| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| import base64 | |
| import binascii | |
| import os | |
| import sys | |
| import threading | |
| from collections import OrderedDict, deque | |
| from concurrent.futures import Future | |
| from typing import Any, Callable, Deque, Dict, Iterator, List, Optional, Set, Tuple | |
| import pyarrow as pa | |
| import pyarrow.compute as pc | |
| import pyarrow.dataset as ds | |
| from pyarrow import RecordBatch | |
| from pypaimon.common.file_io import FileIO | |
| from pypaimon.common.options.config import CatalogOptions | |
| from pypaimon.common.options.core_options import CoreOptions | |
| from pypaimon.filesystem.pyarrow_file_io import _pyarrow_lt_7 | |
| from pypaimon.data.map_shared_shredding import ( | |
| is_map_selected_keys_field, | |
| is_shared_shredding, | |
| map_selected_keys, | |
| parse_shared_shredding_metadata, | |
| parse_shared_shredding_selection_metadata, | |
| ) | |
| from pypaimon.execution.python.read.map_shared_shredding import ( | |
| assemble_normal_map_selected_keys, | |
| assemble_shared_shredding_map, | |
| assemble_shared_shredding_selected_keys, | |
| shared_shredding_selected_paths, | |
| ) | |
| from pypaimon.data.variant_shredding import ( | |
| VariantSchema, | |
| build_variant_schema, | |
| is_shredded_variant, | |
| ) | |
| from pypaimon.execution.python.read.variant_shredding import assemble_shredded_column | |
| from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader | |
| from pypaimon.schema.arrow_schema import parquet_storage_schema | |
| from pypaimon.schema.data_types import ( | |
| ArrayType, | |
| AtomicType, | |
| DataField, | |
| MapType, | |
| MultisetType, | |
| PyarrowFieldParser, | |
| RowType, | |
| ) | |
| from pypaimon.table.special_fields import SpecialFields | |
| from pypaimon.utils.arrow_utils import zero_column_batch as _zero_column_batch | |
| _DEFAULT_FILE_FORMAT_METADATA_CACHE_MAX_SIZE = 50 * 1024 * 1024 | |
| _FILE_FORMAT_METADATA_CACHE_MAX_ENTRIES = 4096 | |
| _FILE_FORMAT_METADATA_CACHE_MIN_ENTRY_SIZE = 8 * 1024 | |
| _FILE_FORMAT_METADATA_CACHE_CONTAINER_OVERHEAD = 256 | |
| class _FilesystemIdentity: | |
| def __init__(self, filesystem): | |
| self.filesystem = filesystem | |
| def __hash__(self): | |
| return id(self.filesystem) | |
| def __eq__(self, other): | |
| return ( | |
| isinstance(other, _FilesystemIdentity) | |
| and self.filesystem is other.filesystem | |
| ) | |
| class _FileFormatDatasetCache: | |
| def __init__( | |
| self, | |
| max_size: int, | |
| max_entries: int = _FILE_FORMAT_METADATA_CACHE_MAX_ENTRIES): | |
| self.max_size = max_size | |
| self.max_entries = max_entries | |
| self.estimated_size = 0 | |
| self._entries = OrderedDict() | |
| self._loads = {} | |
| self._lock = threading.Lock() | |
| def get_or_load(self, key: Tuple[Any, ...], loader: Callable[[], Any], | |
| size_estimator: Callable[[Any], Optional[int]]): | |
| with self._lock: | |
| entry = self._entries.get(key) | |
| if entry is not None: | |
| self._entries.move_to_end(key) | |
| return entry[0] | |
| future = self._loads.get(key) | |
| if future is None: | |
| future = Future() | |
| self._loads[key] = future | |
| should_load = True | |
| else: | |
| should_load = False | |
| if not should_load: | |
| return future.result() | |
| try: | |
| dataset = loader() | |
| estimated_size = size_estimator(dataset) | |
| except BaseException as exception: | |
| future.set_exception(exception) | |
| with self._lock: | |
| self._loads.pop(key, None) | |
| raise | |
| with self._lock: | |
| if (estimated_size is not None | |
| and estimated_size <= self.max_size): | |
| estimated_size = max(1, estimated_size) | |
| self._entries[key] = (dataset, estimated_size) | |
| self.estimated_size += estimated_size | |
| self._entries.move_to_end(key) | |
| self._evict() | |
| future.set_result(dataset) | |
| with self._lock: | |
| self._loads.pop(key, None) | |
| return dataset | |
| def get(self, key): | |
| with self._lock: | |
| entry = self._entries.get(key) | |
| if entry is None: | |
| return None | |
| self._entries.move_to_end(key) | |
| return entry[0] | |
| def put(self, key, dataset, estimated_size): | |
| estimated_size = max(1, estimated_size) | |
| with self._lock: | |
| if estimated_size > self.max_size or key in self._entries: | |
| return | |
| self._entries[key] = (dataset, estimated_size) | |
| self.estimated_size += estimated_size | |
| self._evict() | |
| def resize(self, max_size: int): | |
| with self._lock: | |
| self.max_size = max_size | |
| self._evict() | |
| def _evict(self): | |
| while ( | |
| self.estimated_size > self.max_size | |
| or len(self._entries) > self.max_entries): | |
| _, (_, evicted_size) = self._entries.popitem(last=False) | |
| self.estimated_size -= evicted_size | |
| class _DecodedRowGroupCache: | |
| def __init__(self, max_size: int): | |
| self._cache = _FileFormatDatasetCache(max_size) | |
| def iter_or_load(self, key, loader): | |
| cached = self._cache.get(key) | |
| if cached is not None: | |
| yield from cached | |
| return | |
| batches = [] | |
| size = 0 | |
| for batch in loader(): | |
| # Yield while loading so one oversized row group is never | |
| # materialized in full merely to discover that it cannot fit. | |
| if batches is not None: | |
| size += batch.nbytes | |
| if size <= self._cache.max_size: | |
| batches.append(batch) | |
| else: | |
| batches = None | |
| yield batch | |
| if batches is not None: | |
| self._cache.put(key, batches, size) | |
| _FILE_FORMAT_DATASET_CACHE = None | |
| _FILE_FORMAT_DATASET_CACHE_LOCK = threading.Lock() | |
| _FILE_FORMAT_DATASET_CACHE_PID = os.getpid() | |
| def _ensure_file_format_dataset_cache_process(): | |
| global _FILE_FORMAT_DATASET_CACHE | |
| global _FILE_FORMAT_DATASET_CACHE_LOCK | |
| global _FILE_FORMAT_DATASET_CACHE_PID | |
| current_pid = os.getpid() | |
| if current_pid != _FILE_FORMAT_DATASET_CACHE_PID: | |
| _FILE_FORMAT_DATASET_CACHE = None | |
| _FILE_FORMAT_DATASET_CACHE_LOCK = threading.Lock() | |
| _FILE_FORMAT_DATASET_CACHE_PID = current_pid | |
| def _file_format_dataset_cache(max_size: int) -> _FileFormatDatasetCache: | |
| global _FILE_FORMAT_DATASET_CACHE | |
| _ensure_file_format_dataset_cache_process() | |
| with _FILE_FORMAT_DATASET_CACHE_LOCK: | |
| if _FILE_FORMAT_DATASET_CACHE is None: | |
| _FILE_FORMAT_DATASET_CACHE = _FileFormatDatasetCache(max_size) | |
| else: | |
| _FILE_FORMAT_DATASET_CACHE.resize(max_size) | |
| return _FILE_FORMAT_DATASET_CACHE | |
| def _reset_file_format_dataset_cache(): | |
| global _FILE_FORMAT_DATASET_CACHE | |
| _ensure_file_format_dataset_cache_process() | |
| with _FILE_FORMAT_DATASET_CACHE_LOCK: | |
| _FILE_FORMAT_DATASET_CACHE = None | |
| def _estimate_file_format_dataset_size(dataset, file_format: str) -> Optional[int]: | |
| try: | |
| if file_format == 'parquet': | |
| footer_size = 0 | |
| for fragment in dataset.get_fragments(): | |
| metadata = fragment.metadata | |
| if metadata is not None: | |
| footer_size += int(metadata.serialized_size) | |
| if footer_size > 0: | |
| return footer_size | |
| return int(dataset.schema.serialize().size) | |
| except Exception: | |
| return None | |
| def _estimate_file_format_cache_entry_size( | |
| key: Tuple[Any, ...], | |
| dataset, | |
| file_format: str) -> Optional[int]: | |
| metadata_size = _estimate_file_format_dataset_size(dataset, file_format) | |
| if metadata_size is None: | |
| return None | |
| # PyArrow does not expose the native size retained by Dataset and Fragment | |
| # objects. Account for all visible Python objects and apply a conservative | |
| # floor so tiny files cannot turn a byte-bounded cache into an effectively | |
| # unbounded object cache. | |
| visible_size = ( | |
| metadata_size | |
| + sys.getsizeof(key) | |
| + sum(sys.getsizeof(part) for part in key) | |
| + sys.getsizeof(dataset) | |
| + sys.getsizeof((dataset, metadata_size)) | |
| + _FILE_FORMAT_METADATA_CACHE_CONTAINER_OVERHEAD | |
| ) | |
| return max(_FILE_FORMAT_METADATA_CACHE_MIN_ENTRY_SIZE, visible_size) | |
| def _file_format_metadata_cache_max_size(file_io: FileIO) -> int: | |
| properties = getattr(file_io, 'properties', None) | |
| if properties is None: | |
| return _DEFAULT_FILE_FORMAT_METADATA_CACHE_MAX_SIZE | |
| return properties.get( | |
| CatalogOptions.FILE_FORMAT_METADATA_CACHE_MAX_SIZE).get_bytes() | |
| def _uses_block_cache(file_io: FileIO, file_path: str) -> bool: | |
| from pypaimon.filesystem.caching_file_io import CachingFileIO | |
| return isinstance(file_io, CachingFileIO) and file_io._is_cacheable(file_path) | |
| def _open_pyarrow_input_file(file_io: FileIO, file_path: str, file_size=None): | |
| if _uses_block_cache(file_io, file_path): | |
| return pa.PythonFile( | |
| file_io._new_input_stream_with_known_size(file_path, file_size), mode='r') | |
| return file_io.filesystem.open_input_file(file_io.to_filesystem_path(file_path)) | |
| def _file_format_dataset(file_io: FileIO, file_format: str, file_path: str, | |
| cache_max_size: int, | |
| file_size: Optional[int] = None): | |
| file_path_for_pyarrow = file_io.to_filesystem_path(file_path) | |
| filesystem = file_io.filesystem | |
| known_size = file_size if file_size is not None and file_size > 0 else None | |
| # PyArrow's Python FileSystemHandler API opens files by path, so its | |
| # FileInfo overload cannot forward file_size to handlers such as Jindo. | |
| # Let capable handlers consume the size from Paimon's immutable file | |
| # metadata before PyArrow opens the file. | |
| handler = getattr(filesystem, "handler", None) | |
| register_file_size = getattr(handler, "register_file_size", None) | |
| if known_size is not None and register_file_size is not None: | |
| register_file_size(file_path_for_pyarrow, known_size) | |
| if _uses_block_cache(file_io, file_path): | |
| from pypaimon.filesystem.caching_file_system import CachedFileSystemHandler | |
| import pyarrow.fs as pafs | |
| filesystem = pafs.PyFileSystem(CachedFileSystemHandler(file_io, file_path, known_size)) | |
| def load(): | |
| if file_format == 'parquet': | |
| parquet_format = ds.ParquetFileFormat() | |
| fragment_options = {} | |
| if known_size is not None and not _pyarrow_lt_7(): | |
| fragment_options["file_size"] = known_size | |
| try: | |
| fragment = parquet_format.make_fragment( | |
| file_path_for_pyarrow, filesystem=filesystem, | |
| **fragment_options) | |
| except TypeError as error: | |
| # PyArrow 7-12 expose make_fragment but do not accept the | |
| # file_size hint. The handler already received the size. | |
| if not fragment_options or 'file_size' not in str(error): | |
| raise | |
| fragment = parquet_format.make_fragment( | |
| file_path_for_pyarrow, filesystem=filesystem) | |
| # Reuse this fragment's footer for schema discovery and scanning. | |
| return ds.FileSystemDataset( | |
| [fragment], fragment.physical_schema, parquet_format, filesystem) | |
| return ds.dataset( | |
| file_path_for_pyarrow, format=file_format, filesystem=filesystem) | |
| key = ( | |
| _FilesystemIdentity(file_io if _uses_block_cache(file_io, file_path) else filesystem), | |
| file_format, file_path_for_pyarrow, | |
| known_size) | |
| if cache_max_size <= 0: | |
| _reset_file_format_dataset_cache() | |
| return load() | |
| return _file_format_dataset_cache(cache_max_size).get_or_load( | |
| key, | |
| load, | |
| lambda dataset: _estimate_file_format_cache_entry_size( | |
| key, dataset, file_format)) | |
| def _orc_schema_with_field_metadata(file_io: FileIO, file_path: str, | |
| fallback: pa.Schema) -> pa.Schema: | |
| """Read Paimon's Arrow schema from ORC user metadata when necessary.""" | |
| if any(is_shared_shredding(field) for field in fallback): | |
| return fallback | |
| import pyarrow.orc as orc | |
| source = _open_pyarrow_input_file(file_io, file_path) | |
| try: | |
| metadata = orc.ORCFile(source).metadata | |
| arrow_schema = metadata.get(b"ARROW:schema") | |
| if arrow_schema is None: | |
| arrow_schema = metadata.get("ARROW:schema") | |
| if arrow_schema is None: | |
| return fallback | |
| if isinstance(arrow_schema, str): | |
| arrow_schema = arrow_schema.encode("latin-1") | |
| try: | |
| arrow_schema = base64.b64decode(arrow_schema, validate=True) | |
| except binascii.Error: | |
| pass | |
| return pa.ipc.read_schema(pa.BufferReader(arrow_schema)) | |
| finally: | |
| source.close() | |
| class FormatPyArrowReader(RecordBatchReader): | |
| """ | |
| A Format Reader that reads record batch from a Parquet or ORC file using PyArrow, | |
| and filters it based on the provided predicate and projection. | |
| When a VARIANT column is stored in the shredded Parquet format (a struct with | |
| ``metadata``, ``value``, and ``typed_value`` fields), this reader transparently | |
| reconstructs the standard ``struct<value: binary, metadata: binary>`` representation. | |
| It also restores shared-shredding MAP columns from their physical struct layout. | |
| """ | |
| def __init__(self, file_io: FileIO, file_format: str, file_path: str, | |
| read_fields: List[DataField], | |
| push_down_predicate: Any, batch_size: int = 1024, | |
| options: CoreOptions = None, | |
| nested_name_paths: Optional[List[List[str]]] = None, | |
| predicate_field_names: Optional[Set[str]] = None, | |
| row_indices: Optional[List[int]] = None, | |
| row_ranges: Optional[List[Tuple[int, int]]] = None, | |
| row_group_cache: Optional[_DecodedRowGroupCache] = None, | |
| file_size: Optional[int] = None): | |
| from pypaimon.filesystem.resolving_file_io import ResolvingFileIO | |
| if isinstance(file_io, ResolvingFileIO): | |
| file_io = file_io._get_fileio(file_path) | |
| self._predicate_field_names = predicate_field_names or set() | |
| file_path_for_pyarrow = file_io.to_filesystem_path(file_path) | |
| self._row_group_cache = row_group_cache | |
| self._row_group_cache_filesystem = _FilesystemIdentity( | |
| file_io if _uses_block_cache(file_io, file_path) else file_io.filesystem) | |
| self._row_group_cache_path = file_path_for_pyarrow | |
| cache_max_size = _file_format_metadata_cache_max_size(file_io) | |
| self.dataset = _file_format_dataset( | |
| file_io, file_format, file_path, cache_max_size, file_size) | |
| self._range_slicer = None | |
| self._selected_parquet_row_groups = None | |
| self._exhausted = False | |
| if row_indices is not None and row_ranges is not None: | |
| raise ValueError( | |
| "row_indices and row_ranges cannot both be provided") | |
| row_selection_supplied = ( | |
| row_indices is not None or row_ranges is not None) | |
| if row_selection_supplied and file_format == 'parquet': | |
| if push_down_predicate is not None: | |
| raise ValueError( | |
| "row selections cannot be combined with a scanner-level " | |
| "push-down predicate because filtering shifts row " | |
| "positions") | |
| runs = ( | |
| _normalize_runs(row_ranges) | |
| if row_ranges is not None | |
| else _to_runs(row_indices) | |
| ) | |
| if not runs: | |
| self._exhausted = True | |
| else: | |
| fragment = next(iter(self.dataset.get_fragments()), None) | |
| row_group_fragments = ( | |
| list(fragment.split_by_row_group()) | |
| if fragment is not None else []) | |
| selected_infos = [] | |
| selected_ids = [] | |
| offset = 0 | |
| run_index = 0 | |
| for row_group_fragment in row_group_fragments: | |
| row_group = row_group_fragment.row_groups[0] | |
| row_count = row_group.num_rows | |
| lower, upper = offset, offset + row_count - 1 | |
| while ( | |
| run_index < len(runs) | |
| and runs[run_index][1] < lower): | |
| run_index += 1 | |
| if (run_index < len(runs) | |
| and runs[run_index][0] <= upper): | |
| selected_infos.append((offset, row_count)) | |
| selected_ids.append(row_group.id) | |
| offset += row_count | |
| if not selected_ids: | |
| self._exhausted = True | |
| else: | |
| self._selected_parquet_row_groups = selected_ids | |
| self._range_slicer = _RowRunSlicer( | |
| selected_infos, runs) | |
| self._file_format = file_format | |
| self.read_fields = read_fields | |
| self._read_field_names = [f.name for f in read_fields] | |
| if nested_name_paths is not None and len(nested_name_paths) != len(read_fields): | |
| raise ValueError( | |
| "nested_name_paths length {} does not match read_fields length {}".format( | |
| len(nested_name_paths), len(read_fields))) | |
| self._nested_name_paths = nested_name_paths | |
| has_nested_path = bool( | |
| nested_name_paths and any(len(p) > 1 for p in nested_name_paths)) | |
| self._has_nested_path = has_nested_path | |
| file_schema = self.dataset.schema | |
| physical_parquet_schema = parquet_storage_schema(file_schema) if file_format == 'parquet' else file_schema | |
| has_fixed_size_lists = not file_schema.equals(physical_parquet_schema) | |
| has_logical_map = any( | |
| isinstance(field.type, MapType) | |
| or is_map_selected_keys_field(field) | |
| for field in read_fields) | |
| metadata_schema = ( | |
| _orc_schema_with_field_metadata(file_io, file_path, file_schema) | |
| if file_format == 'orc' and has_logical_map else file_schema) | |
| if has_nested_path: | |
| self.existing_fields = [] | |
| self.missing_fields = [] | |
| for f, path in zip(read_fields, nested_name_paths): | |
| if _path_exists_in_arrow_schema(file_schema, path): | |
| self.existing_fields.append(f.name) | |
| else: | |
| self.missing_fields.append(f.name) | |
| else: | |
| file_schema_names = set(file_schema.names) | |
| self.existing_fields = [f.name for f in read_fields if f.name in file_schema_names] | |
| self.missing_fields = [f.name for f in read_fields if f.name not in file_schema_names] | |
| self._variant_schema_cache: Dict[pa.DataType, VariantSchema] = {} | |
| self._shared_shredding_maps = {} | |
| self._selected_key_maps = {} | |
| self._selected_shared_map_paths = {} | |
| logical_maps_by_source = {} | |
| selected_maps_by_source = {} | |
| if nested_name_paths is None: | |
| source_names = [field.name for field in read_fields] | |
| else: | |
| source_names = [ | |
| path[0] if len(path) == 1 else None | |
| for path in nested_name_paths | |
| ] | |
| for logical_field, source_name in zip(read_fields, source_names): | |
| if (source_name is not None | |
| and isinstance(logical_field.type, MapType)): | |
| logical_maps_by_source.setdefault(source_name, []).append( | |
| logical_field) | |
| elif (source_name is not None | |
| and is_map_selected_keys_field(logical_field)): | |
| selected_maps_by_source[source_name] = logical_field | |
| for field in metadata_schema: | |
| logical_fields = logical_maps_by_source.get(field.name, []) | |
| if logical_fields and is_shared_shredding(field): | |
| metadata = parse_shared_shredding_metadata(field) | |
| for logical_field in logical_fields: | |
| logical_arrow_type = PyarrowFieldParser.from_paimon_type( | |
| logical_field.type) | |
| self._shared_shredding_maps[logical_field.name] = ( | |
| logical_arrow_type, metadata) | |
| selected_field = selected_maps_by_source.get(field.name) | |
| if selected_field is not None: | |
| keys = map_selected_keys(selected_field.description) | |
| value_type = PyarrowFieldParser.from_paimon_type( | |
| selected_field.type.fields[0].type) | |
| selection_metadata = None | |
| if is_shared_shredding(field): | |
| selection_metadata = \ | |
| parse_shared_shredding_selection_metadata(field) | |
| self._selected_shared_map_paths[field.name] = \ | |
| shared_shredding_selected_paths( | |
| field.name, keys, selection_metadata) | |
| self._selected_key_maps[field.name] = ( | |
| keys, value_type, selection_metadata) | |
| self._bounded_variant_read = ( | |
| self._file_format == 'parquet' and self._has_projected_variant()) | |
| self._select_nested_after_scan = False | |
| if has_nested_path and not self._bounded_variant_read: | |
| existing_set = set(self.existing_fields) | |
| columns_dict = {} | |
| try: | |
| for f, path in zip(read_fields, nested_name_paths): | |
| if f.name in existing_set: | |
| columns_dict[f.name] = ds.field(*path) | |
| self._scan_columns = columns_dict | |
| except TypeError: | |
| # PyArrow 6 only accepts one field name and cannot build a | |
| # nested FieldRef. Read the required top-level columns and | |
| # extract their children after scanning instead. | |
| self._scan_columns = [] | |
| for f, path in zip(read_fields, nested_name_paths): | |
| if (f.name in existing_set | |
| and path[0] not in self._scan_columns): | |
| self._scan_columns.append(path[0]) | |
| self._select_nested_after_scan = True | |
| elif has_nested_path: | |
| self._scan_columns = None | |
| else: | |
| # Only pass existing fields to PyArrow scanner to avoid errors | |
| self._scan_columns = self.existing_fields | |
| self._scan_filter = push_down_predicate | |
| self._scan_batch_size = batch_size | |
| self._output_schema = ( | |
| PyarrowFieldParser.from_paimon_schema(read_fields) if read_fields else None | |
| ) | |
| # Read projected VARIANT columns in bounded batches. | |
| self._parquet_file = None | |
| self._parquet_source = None | |
| self._page_index_reader = None | |
| self._orc_file = None | |
| self._orc_source = None | |
| if (self._bounded_variant_read | |
| or has_fixed_size_lists | |
| or self._selected_parquet_row_groups is not None | |
| or (file_format == 'parquet' | |
| and self._selected_shared_map_paths)): | |
| import pyarrow.parquet as pq | |
| # ParquetFile(filesystem=...) is unavailable in PyArrow 6. | |
| self._parquet_source = _open_pyarrow_input_file(file_io, file_path, file_size) | |
| try: | |
| metadata = None | |
| if has_fixed_size_lists: | |
| # Native/older writers can retain fixed-size lists in | |
| # ARROW:schema. Preserve every physical footer field, | |
| # including INT96/DECIMAL encodings and page-index offsets. | |
| from pypaimon.execution.python.read.reader.parquet_page_index_reader import ( | |
| parquet_metadata_with_arrow_schema, | |
| ) | |
| fragment = next(iter(self.dataset.get_fragments())) | |
| metadata = parquet_metadata_with_arrow_schema(fragment.metadata, physical_parquet_schema) | |
| self._parquet_file = pq.ParquetFile(self._parquet_source, metadata=metadata) | |
| if (self._selected_parquet_row_groups is not None | |
| and options is not None | |
| and options.parquet_column_index_enabled() | |
| and self._row_group_cache is None): | |
| from pypaimon.execution.python.read.reader.parquet_page_index_reader import ( | |
| ParquetPageIndexReader, | |
| ) | |
| self._page_index_reader = ParquetPageIndexReader.create( | |
| self._parquet_source, self._parquet_file, | |
| self._row_group_read_columns(), | |
| self._selected_parquet_row_groups, batch_size) | |
| except BaseException: | |
| self._parquet_source.close() | |
| raise | |
| if file_format == 'orc' and self._selected_shared_map_paths: | |
| import pyarrow.orc as orc | |
| self._orc_source = _open_pyarrow_input_file(file_io, file_path) | |
| self._orc_file = orc.ORCFile(self._orc_source) | |
| if self._exhausted: | |
| self._raw_batches = iter(()) | |
| elif self._page_index_reader is not None: | |
| # Page selection already preserves original row positions. Slice | |
| # fallback row groups here too, before mixing the two streams. | |
| self._range_slicer = None | |
| self._raw_batches = self._iter_page_index_batches(selected_infos, runs) | |
| elif self._parquet_file is not None: | |
| self._raw_batches = self._iter_row_group_batches() | |
| elif self._orc_file is not None: | |
| self._raw_batches = self._iter_orc_batches() | |
| else: | |
| reader = self.dataset.scanner( | |
| columns=self._scan_columns, | |
| filter=self._scan_filter, | |
| batch_size=self._scan_batch_size, | |
| ).to_reader() | |
| raw_batches = self._iter_reader_batches(reader) | |
| if self._select_nested_after_scan: | |
| raw_batches = ( | |
| self._select_nested_fields(batch) | |
| for batch in raw_batches) | |
| self._raw_batches = raw_batches | |
| def _has_projected_variant(self) -> bool: | |
| return any( | |
| f.name in self.existing_fields | |
| and _contains_variant(f.type) | |
| for f in self.read_fields) | |
| @staticmethod | |
| def _iter_reader_batches(reader): | |
| while True: | |
| try: | |
| yield reader.read_next_batch() | |
| except StopIteration: | |
| return | |
| def _iter_row_group_batches(self): | |
| columns = self._row_group_read_columns() | |
| for row_group in self._surviving_row_group_ids(): | |
| if (self._row_group_cache is not None | |
| and self._selected_parquet_row_groups is not None): | |
| key = ( | |
| self._row_group_cache_filesystem, | |
| self._row_group_cache_path, | |
| row_group, | |
| tuple(columns), | |
| self._scan_batch_size, | |
| ) | |
| batches = self._row_group_cache.iter_or_load( | |
| key, | |
| lambda: self._read_parquet_row_group_batches( | |
| row_group, columns), | |
| ) | |
| else: | |
| batches = self._read_parquet_row_group_batches( | |
| row_group, columns) | |
| for batch in batches: | |
| if self._has_nested_path: | |
| filtered_batches = [batch] | |
| if self._scan_filter is not None: | |
| table = ds.dataset( | |
| pa.Table.from_batches([batch]) | |
| ).scanner(filter=self._scan_filter).to_table() | |
| filtered_batches = table.to_batches() | |
| for filtered in filtered_batches: | |
| out = self._select_nested_fields(filtered) | |
| if out.num_rows: | |
| yield out | |
| continue | |
| if self._scan_filter is None: | |
| yield self._select_existing_fields(batch) | |
| continue | |
| table = ds.dataset( | |
| pa.Table.from_batches([batch]) | |
| ).scanner(filter=self._scan_filter).to_table() | |
| if self.existing_fields: | |
| table = table.select(self.existing_fields) | |
| for out in table.to_batches(): | |
| if out.num_rows: | |
| yield out | |
| def _iter_page_index_batches(self, selected_infos, runs): | |
| select = self._select_nested_fields if self._has_nested_path else self._select_existing_fields | |
| run_index = 0 | |
| for group, (offset, count) in zip( | |
| self._selected_parquet_row_groups, selected_infos): | |
| while run_index < len(runs) and runs[run_index][1] < offset: | |
| run_index += 1 | |
| local_runs = [] | |
| position = run_index | |
| while position < len(runs) and runs[position][0] < offset + count: | |
| lower, upper = runs[position] | |
| local_runs.append((max(0, lower - offset), | |
| min(count - 1, upper - offset))) | |
| position += 1 | |
| batches = self._page_index_reader.read_row_group(group, local_runs) | |
| if batches is None: | |
| raw = self._read_parquet_row_group_batches( | |
| group, self._row_group_read_columns()) | |
| slicer = _RowRunSlicer([(0, count)], local_runs) | |
| while True: | |
| batch = slicer.next_batch(raw) | |
| if batch is None: | |
| break | |
| yield select(batch) | |
| else: | |
| try: | |
| for batch in batches: | |
| yield select(batch) | |
| finally: | |
| batches.close() | |
| def _read_parquet_row_group_batches(self, row_group, columns): | |
| return self._parquet_file.iter_batches( | |
| row_groups=[row_group], | |
| columns=columns, | |
| batch_size=self._scan_batch_size, | |
| ) | |
| def _iter_orc_batches(self): | |
| columns = self._physical_read_columns() | |
| for stripe in range(self._orc_file.nstripes): | |
| batch = self._orc_file.read_stripe(stripe, columns=columns) | |
| for offset in range(0, batch.num_rows, self._scan_batch_size): | |
| current = batch.slice(offset, self._scan_batch_size) | |
| if self._scan_filter is not None: | |
| table = ds.dataset(pa.Table.from_batches([current])).scanner( | |
| filter=self._scan_filter).to_table() | |
| batches = table.to_batches() | |
| else: | |
| batches = [current] | |
| for filtered in batches: | |
| if filtered.num_rows: | |
| yield self._select_existing_fields(filtered) | |
| def _row_group_read_columns(self): | |
| if self._has_nested_path: | |
| existing = set(self.existing_fields) | |
| columns = [] | |
| for field, path in zip(self.read_fields, self._nested_name_paths): | |
| if field.name in existing and path[0] not in columns: | |
| columns.append(path[0]) | |
| else: | |
| columns = list(self.existing_fields) | |
| return self._physical_read_columns(columns) | |
| def _physical_read_columns(self, columns=None): | |
| if columns is None: | |
| columns = list(self.existing_fields) | |
| else: | |
| columns = list(columns) | |
| for field_name, paths in self._selected_shared_map_paths.items(): | |
| columns = [name for name in columns if name != field_name] | |
| for path in paths: | |
| if path not in columns: | |
| columns.append(path) | |
| if self._scan_filter is not None: | |
| file_names = set(self.dataset.schema.names) | |
| for name in self._predicate_field_names: | |
| if name in file_names and name not in columns: | |
| columns.append(name) | |
| return columns | |
| def _select_existing_fields(self, batch): | |
| if not self.existing_fields: | |
| return _zero_column_batch(batch.num_rows) | |
| columns = [] | |
| fields = [] | |
| for name in self.existing_fields: | |
| index = batch.schema.get_field_index(name) | |
| if index < 0: | |
| raise KeyError("Field not found in batch: {}".format(name)) | |
| columns.append(batch.column(index)) | |
| fields.append(batch.schema.field(index)) | |
| return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) | |
| def _select_nested_fields(self, batch): | |
| columns = [] | |
| names = [] | |
| existing = set(self.existing_fields) | |
| for field, path in zip(self.read_fields, self._nested_name_paths): | |
| if field.name not in existing: | |
| continue | |
| index = batch.schema.get_field_index(path[0]) | |
| if index < 0: | |
| raise KeyError("Field not found in batch: {}".format(path[0])) | |
| column = batch.column(index) | |
| for name in path[1:]: | |
| index = column.type.get_field_index(name) | |
| column = column.flatten()[index] | |
| columns.append(column) | |
| names.append(field.name) | |
| if not columns: | |
| return _zero_column_batch(batch.num_rows) | |
| return pa.RecordBatch.from_arrays(columns, names=names) | |
| def _surviving_row_group_ids(self): | |
| total = self._parquet_file.num_row_groups | |
| if self._scan_filter is None: | |
| if self._selected_parquet_row_groups is not None: | |
| return self._selected_parquet_row_groups | |
| return range(total) | |
| try: | |
| ids = set() | |
| for fragment in self.dataset.get_fragments( | |
| filter=self._scan_filter): | |
| for row_group in fragment.split_by_row_group( | |
| self._scan_filter): | |
| ids.update(info.id for info in row_group.row_groups) | |
| return sorted(ids) | |
| except Exception: | |
| return range(total) | |
| def read_arrow_batch(self) -> Optional[RecordBatch]: | |
| if self._range_slicer is not None: | |
| batch = self._range_slicer.next_batch(self._raw_batches) | |
| else: | |
| batch = next(self._raw_batches, None) | |
| if batch is None: | |
| return None | |
| return self._post_process_batch(batch) | |
| def _post_process_batch(self, batch: RecordBatch) -> RecordBatch: | |
| if self._file_format == 'orc' and self._output_schema is not None: | |
| batch = self._cast_orc_time_columns(batch) | |
| if self._shared_shredding_maps: | |
| batch = self._assemble_shared_shredding_maps(batch) | |
| if self._selected_key_maps: | |
| batch = self._assemble_selected_key_maps(batch) | |
| batch = self._assemble_shredded_variants(batch) | |
| if not self.missing_fields: | |
| return batch | |
| def _type_for_missing(name: str) -> pa.DataType: | |
| if self._output_schema is not None: | |
| idx = self._output_schema.get_field_index(name) | |
| if idx >= 0: | |
| return self._output_schema.field(idx).type | |
| return pa.null() | |
| missing_columns = [ | |
| pa.nulls(batch.num_rows, type=_type_for_missing(name)) | |
| for name in self.missing_fields | |
| ] | |
| all_columns = [] | |
| out_fields = [] | |
| for field_name in self._read_field_names: | |
| if field_name in self.existing_fields: | |
| column_idx = self.existing_fields.index(field_name) | |
| all_columns.append(batch.column(column_idx)) | |
| out_fields.append(batch.schema.field(column_idx)) | |
| else: | |
| column_idx = self.missing_fields.index(field_name) | |
| col_type = _type_for_missing(field_name) | |
| all_columns.append(missing_columns[column_idx]) | |
| nullable = not SpecialFields.is_system_field(field_name) | |
| out_fields.append( | |
| pa.field(field_name, col_type, nullable=nullable)) | |
| return pa.RecordBatch.from_arrays( | |
| all_columns, schema=pa.schema(out_fields)) | |
| def _assemble_shared_shredding_maps( | |
| self, batch: pa.RecordBatch) -> pa.RecordBatch: | |
| columns = list(batch.columns) | |
| fields = list(batch.schema) | |
| changed = False | |
| for index, field in enumerate(fields): | |
| shared = self._shared_shredding_maps.get(field.name) | |
| if shared is None: | |
| continue | |
| map_type, (name_by_id, num_columns) = shared | |
| columns[index] = assemble_shared_shredding_map( | |
| columns[index], map_type, name_by_id, num_columns) | |
| fields[index] = pa.field( | |
| field.name, map_type, nullable=field.nullable) | |
| changed = True | |
| if not changed: | |
| return batch | |
| return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) | |
| def _assemble_selected_key_maps( | |
| self, batch: pa.RecordBatch) -> pa.RecordBatch: | |
| columns = list(batch.columns) | |
| fields = list(batch.schema) | |
| changed = False | |
| for index, field in enumerate(fields): | |
| selected = self._selected_key_maps.get(field.name) | |
| if selected is None: | |
| continue | |
| keys, value_type, metadata = selected | |
| if metadata is None: | |
| columns[index] = assemble_normal_map_selected_keys( | |
| columns[index], keys, value_type) | |
| else: | |
| columns[index] = assemble_shared_shredding_selected_keys( | |
| columns[index], keys, value_type, metadata) | |
| fields[index] = pa.field( | |
| field.name, columns[index].type, nullable=field.nullable) | |
| changed = True | |
| if not changed: | |
| return batch | |
| return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) | |
| def _assemble_shredded_variants(self, batch: pa.RecordBatch) -> pa.RecordBatch: | |
| changed = False | |
| columns = list(batch.columns) | |
| fields = list(batch.schema) | |
| logical_types = {field.name: field.type for field in self.read_fields} | |
| for i, f in enumerate(fields): | |
| logical_type = logical_types.get(f.name) | |
| if logical_type is not None: | |
| new_col, column_changed = _assemble_variant_column( | |
| columns[i], logical_type, self._variant_schema_cache) | |
| else: | |
| new_col, column_changed = columns[i], False | |
| if column_changed: | |
| columns[i] = new_col | |
| fields[i] = pa.field(f.name, new_col.type, nullable=f.nullable) | |
| changed = True | |
| if not changed: | |
| return batch | |
| return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) | |
| def _cast_orc_time_columns(self, batch): | |
| """Cast int32 TIME columns back to time32('ms') when reading ORC. | |
| """ | |
| columns = [] | |
| fields = [] | |
| changed = False | |
| for i, name in enumerate(batch.schema.names): | |
| col = batch.column(i) | |
| idx = self._output_schema.get_field_index(name) | |
| if idx >= 0 and pa.types.is_int32(col.type) \ | |
| and pa.types.is_time(self._output_schema.field(idx).type): | |
| col = col.cast(self._output_schema.field(idx).type) | |
| fields.append(self._output_schema.field(idx)) | |
| changed = True | |
| else: | |
| fields.append(batch.schema.field(i)) | |
| columns.append(col) | |
| if changed: | |
| return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) | |
| return batch | |
| def close(self): | |
| close_batches = getattr(self._raw_batches, 'close', None) | |
| if close_batches is not None: | |
| close_batches() | |
| self._raw_batches = None | |
| if self._parquet_file is not None: | |
| close = getattr(self._parquet_file, 'close', None) | |
| if close is not None: | |
| close() | |
| self._parquet_file = None | |
| if self._parquet_source is not None: | |
| self._parquet_source.close() | |
| self._parquet_source = None | |
| self._page_index_reader = None | |
| if self._orc_source is not None: | |
| self._orc_source.close() | |
| self._orc_source = None | |
| self._orc_file = None | |
| def _path_exists_in_arrow_schema(schema: pa.Schema, path: List[str]) -> bool: | |
| """Check whether a name path is fully resolvable in the given schema.""" | |
| if not path: | |
| return False | |
| if path[0] not in schema.names: | |
| return False | |
| current_type = schema.field(path[0]).type | |
| for name in path[1:]: | |
| if not pa.types.is_struct(current_type): | |
| return False | |
| idx = current_type.get_field_index(name) | |
| if idx < 0: | |
| return False | |
| current_type = current_type[idx].type | |
| return True | |
| def _to_runs(row_indices: List[int]) -> List[Tuple[int, int]]: | |
| """Collapse row indices into sorted, distinct, inclusive runs.""" | |
| if not row_indices: | |
| return [] | |
| sorted_indices = sorted(set(row_indices)) | |
| runs = [] | |
| start = previous = sorted_indices[0] | |
| for index in sorted_indices[1:]: | |
| if index == previous + 1: | |
| previous = index | |
| continue | |
| runs.append((start, previous)) | |
| start = previous = index | |
| runs.append((start, previous)) | |
| return runs | |
| def _normalize_runs( | |
| row_ranges: List[Tuple[int, int]]) -> List[Tuple[int, int]]: | |
| """Sort and merge inclusive row ranges without expanding their rows.""" | |
| if not row_ranges: | |
| return [] | |
| ranges = sorted(row_ranges) | |
| merged = [] | |
| for lower, upper in ranges: | |
| if lower > upper: | |
| raise ValueError( | |
| "Invalid row range: {} > {}".format(lower, upper)) | |
| if merged and lower <= merged[-1][1] + 1: | |
| merged[-1] = (merged[-1][0], max(merged[-1][1], upper)) | |
| else: | |
| merged.append((lower, upper)) | |
| return merged | |
| class _RowRunSlicer: | |
| """Slice selected Parquet row groups down to requested file-local rows.""" | |
| def __init__( | |
| self, | |
| selected_infos: List[Tuple[int, int]], | |
| runs: List[Tuple[int, int]]): | |
| self._segments = [] | |
| concatenated_offset = 0 | |
| for file_offset, row_count in selected_infos: | |
| self._segments.append(( | |
| concatenated_offset, | |
| concatenated_offset + row_count, | |
| file_offset, | |
| )) | |
| concatenated_offset += row_count | |
| self._runs = [(lower, upper + 1) for lower, upper in runs] | |
| self._stream_offset = 0 | |
| self._segment_index = 0 | |
| self._run_index = 0 | |
| self._pending: Deque[RecordBatch] = deque() | |
| def next_batch( | |
| self, batches: Iterator[RecordBatch]) -> Optional[RecordBatch]: | |
| while not self._pending: | |
| batch = next(batches, None) | |
| if batch is None: | |
| return None | |
| self._slice_batch(batch) | |
| return self._pending.popleft() | |
| def _slice_batch(self, batch: RecordBatch) -> None: | |
| batch_start = self._stream_offset | |
| batch_end = batch_start + batch.num_rows | |
| self._stream_offset = batch_end | |
| position = batch_start | |
| while position < batch_end: | |
| while ( | |
| self._segment_index < len(self._segments) | |
| and position >= self._segments[ | |
| self._segment_index][1]): | |
| self._segment_index += 1 | |
| if self._segment_index >= len(self._segments): | |
| return | |
| segment_start, segment_end, file_start = self._segments[ | |
| self._segment_index] | |
| part_end = min(batch_end, segment_end) | |
| local_start = file_start + position - segment_start | |
| local_end = file_start + part_end - segment_start | |
| while ( | |
| self._run_index < len(self._runs) | |
| and self._runs[self._run_index][1] <= local_start): | |
| self._run_index += 1 | |
| run_index = self._run_index | |
| while ( | |
| run_index < len(self._runs) | |
| and self._runs[run_index][0] < local_end): | |
| run_start, run_end = self._runs[run_index] | |
| lower = max(local_start, run_start) | |
| upper = min(local_end, run_end) | |
| if lower < upper: | |
| offset = ( | |
| position - batch_start + lower - local_start) | |
| self._pending.append( | |
| batch.slice(offset, upper - lower)) | |
| if run_end <= local_end: | |
| run_index += 1 | |
| else: | |
| break | |
| self._run_index = run_index | |
| position = part_end | |
| def _contains_variant(data_type) -> bool: | |
| if isinstance(data_type, AtomicType): | |
| return data_type.type.upper() == 'VARIANT' | |
| if isinstance(data_type, (ArrayType, MultisetType)): | |
| return _contains_variant(data_type.element) | |
| if isinstance(data_type, MapType): | |
| return (_contains_variant(data_type.key) | |
| or _contains_variant(data_type.value)) | |
| if isinstance(data_type, RowType): | |
| return any(_contains_variant(field.type) for field in data_type.fields) | |
| return False | |
| def _assemble_variant_column(column, data_type, schema_cache): | |
| if isinstance(data_type, AtomicType): | |
| if (data_type.type.upper() != 'VARIANT' | |
| or not is_shredded_variant(column.type)): | |
| return column, False | |
| schema = schema_cache.get(column.type) | |
| if schema is None: | |
| schema = build_variant_schema(column.type) | |
| schema_cache[column.type] = schema | |
| return assemble_shredded_column(column, schema), True | |
| if isinstance(data_type, RowType) and pa.types.is_struct(column.type): | |
| logical_fields = {field.name: field.type for field in data_type.fields} | |
| columns = [] | |
| fields = [] | |
| changed = False | |
| for index, arrow_field in enumerate(column.type): | |
| child = column.field(index) | |
| logical_type = logical_fields.get(arrow_field.name) | |
| if logical_type is not None: | |
| child, child_changed = _assemble_variant_column( | |
| child, logical_type, schema_cache) | |
| changed = changed or child_changed | |
| columns.append(child) | |
| fields.append(pa.field( | |
| arrow_field.name, | |
| child.type, | |
| nullable=arrow_field.nullable, | |
| metadata=arrow_field.metadata, | |
| )) | |
| if changed: | |
| mask = column.is_null() if column.null_count else None | |
| return pa.StructArray.from_arrays( | |
| columns, fields=fields, mask=mask), True | |
| return column, False | |
| if (isinstance(data_type, (ArrayType, MultisetType)) | |
| and (pa.types.is_list(column.type) | |
| or pa.types.is_large_list(column.type))): | |
| offsets, start, end = _normalized_offsets(column) | |
| values = column.values.slice(start, end - start) | |
| values, changed = _assemble_variant_column( | |
| values, data_type.element, schema_cache) | |
| if not changed: | |
| return column, False | |
| if pa.types.is_large_list(column.type): | |
| result = pa.LargeListArray.from_arrays(offsets, values) | |
| list_type = pa.large_list(pa.field( | |
| column.type.value_field.name, | |
| values.type, | |
| nullable=column.type.value_field.nullable, | |
| metadata=column.type.value_field.metadata, | |
| )) | |
| else: | |
| result = pa.ListArray.from_arrays(offsets, values) | |
| list_type = pa.list_(pa.field( | |
| column.type.value_field.name, | |
| values.type, | |
| nullable=column.type.value_field.nullable, | |
| metadata=column.type.value_field.metadata, | |
| )) | |
| return pa.Array.from_buffers( | |
| list_type, | |
| len(result), | |
| result.buffers()[:2], | |
| null_count=result.null_count, | |
| children=[values], | |
| ), True | |
| if isinstance(data_type, MapType) and pa.types.is_map(column.type): | |
| offsets, start, end = _normalized_offsets(column) | |
| keys = column.keys.slice(start, end - start) | |
| items = column.items.slice(start, end - start) | |
| keys, key_changed = _assemble_variant_column( | |
| keys, data_type.key, schema_cache) | |
| items, item_changed = _assemble_variant_column( | |
| items, data_type.value, schema_cache) | |
| if not key_changed and not item_changed: | |
| return column, False | |
| result = pa.MapArray.from_arrays(offsets, keys, items) | |
| map_type = pa.map_( | |
| pa.field( | |
| column.type.key_field.name, | |
| keys.type, | |
| nullable=False, | |
| metadata=column.type.key_field.metadata, | |
| ), | |
| pa.field( | |
| column.type.item_field.name, | |
| items.type, | |
| nullable=column.type.item_field.nullable, | |
| metadata=column.type.item_field.metadata, | |
| ), | |
| keys_sorted=getattr(column.type, 'keys_sorted', False), | |
| ) | |
| entries = pa.StructArray.from_arrays( | |
| [keys, items], fields=[map_type.key_field, map_type.item_field]) | |
| return pa.Array.from_buffers( | |
| map_type, | |
| len(result), | |
| result.buffers()[:2], | |
| null_count=result.null_count, | |
| children=[entries], | |
| ), True | |
| return column, False | |
| def _normalized_offsets(column): | |
| offsets_array = getattr(column, 'offsets', None) | |
| if offsets_array is None: | |
| offsets_array = pa.Array.from_buffers( | |
| pa.int32(), | |
| len(column) + 1, | |
| [None, column.buffers()[1]], | |
| offset=column.offset, | |
| ) | |
| start = offsets_array[0].as_py() | |
| end = offsets_array[-1].as_py() | |
| offsets = pc.subtract(offsets_array, pa.scalar(start, type=offsets_array.type)) | |
| if column.null_count: | |
| mask = pa.concat_arrays([column.is_null(), pa.array([False])]) | |
| offsets = pc.if_else(mask, pa.scalar(None, type=offsets_array.type), offsets) | |
| return offsets, start, end |