| # 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 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.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.data.variant_shredding import ( |
| VariantSchema, |
| assemble_shredded_column, |
| build_variant_schema, |
| is_shredded_variant, |
| ) |
| from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader |
| from pypaimon.schema.data_types import ( |
| ArrayType, |
| AtomicType, |
| DataField, |
| MapType, |
| MultisetType, |
| PyarrowFieldParser, |
| RowType, |
| ) |
| from pypaimon.table.special_fields import SpecialFields |
| |
| |
| _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, str, str], 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: |
| 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 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 |
| |
| |
| _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, str, str], |
| 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) |
| + sys.getsizeof(key[0]) |
| + sys.getsizeof(key[1]) |
| + sys.getsizeof(key[2]) |
| + 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 _file_format_dataset(file_io: FileIO, file_format: str, file_path: str, |
| cache_max_size: int): |
| file_path_for_pyarrow = file_io.to_filesystem_path(file_path) |
| filesystem = file_io.filesystem |
| |
| def load(): |
| return ds.dataset( |
| file_path_for_pyarrow, format=file_format, filesystem=filesystem) |
| |
| key = (_FilesystemIdentity(filesystem), file_format, file_path_for_pyarrow) |
| 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)) |
| |
| |
| 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. |
| """ |
| |
| 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): |
| self._predicate_field_names = predicate_field_names or set() |
| file_path_for_pyarrow = file_io.to_filesystem_path(file_path) |
| 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) |
| 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 |
| 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_shredding_enabled = ( |
| options is None or options.variant_shredding_enabled()) |
| self._variant_schema_cache: Dict[pa.DataType, VariantSchema] = {} |
| |
| self._bounded_variant_read = ( |
| self._file_format == 'parquet' and self._has_projected_variant()) |
| if has_nested_path and not self._bounded_variant_read: |
| existing_set = set(self.existing_fields) |
| columns_dict = {} |
| 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 |
| 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 |
| if (self._bounded_variant_read |
| or self._selected_parquet_row_groups is not None): |
| import pyarrow.parquet as pq |
| # ParquetFile(filesystem=...) is unavailable in PyArrow 6. |
| self._parquet_file = pq.ParquetFile( |
| file_io.filesystem.open_input_file(file_path_for_pyarrow)) |
| if self._exhausted: |
| self._raw_batches = iter(()) |
| elif self._parquet_file is not None: |
| self._raw_batches = self._iter_row_group_batches() |
| else: |
| reader = self.dataset.scanner( |
| columns=self._scan_columns, |
| filter=self._scan_filter, |
| batch_size=self._scan_batch_size, |
| ).to_reader() |
| self._raw_batches = self._iter_reader_batches(reader) |
| |
| 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(): |
| for batch in self._parquet_file.iter_batches( |
| row_groups=[row_group], |
| columns=columns, |
| batch_size=self._scan_batch_size): |
| if self._has_nested_path: |
| batches = [batch] |
| if self._scan_filter is not None: |
| table = ds.dataset( |
| pa.Table.from_batches([batch]) |
| ).scanner(filter=self._scan_filter).to_table() |
| batches = table.to_batches() |
| for filtered in 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 _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) |
| 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._variant_shredding_enabled: |
| 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_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): |
| 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 |
| |
| |
| 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 _zero_column_batch(num_rows: int) -> RecordBatch: |
| """Build a zero-column batch without losing its logical row count.""" |
| empty_struct = pa.Array.from_buffers( |
| pa.struct([]), num_rows, [None], children=[]) |
| return pa.RecordBatch.from_struct_array(empty_struct) |
| |
| |
| 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, |
| ) |
| raw_offsets = offsets_array.to_pylist() |
| start = raw_offsets[0] |
| end = raw_offsets[-1] |
| offsets = [value - start for value in raw_offsets] |
| for index, is_null in enumerate(column.is_null().to_pylist()): |
| if is_null: |
| offsets[index] = None |
| return pa.array(offsets, type=offsets_array.type), start, end |
