| # 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 re | |
| import struct | |
| from datetime import datetime, timedelta | |
| from decimal import Decimal | |
| from typing import Any, List, Optional | |
| import pyarrow as pa | |
| import pyarrow.dataset as ds | |
| from pyarrow import RecordBatch | |
| from pypaimon.common.file_io import FileIO, supports_pread, pread | |
| from pypaimon.common.predicate import Predicate | |
| from pypaimon.manifest.schema.simple_stats import SimpleStats | |
| from pypaimon.read.push_down_utils import rewrite_predicate_indices | |
| from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader | |
| from pypaimon.schema.data_types import DataField, PyarrowFieldParser | |
| from pypaimon.table.row.generic_row import GenericRow | |
| from pypaimon.table.special_fields import SpecialFields | |
| class FormatMosaicReader(RecordBatchReader): | |
| def __init__(self, file_io: FileIO, file_path: str, read_fields: List[DataField], | |
| push_down_predicate: Any, batch_size: int = 1024, | |
| row_group_predicate: Optional[Predicate] = None): | |
| from mosaic import MosaicReader | |
| self._read_fields = read_fields | |
| self._read_field_names = [f.name for f in read_fields] | |
| self._batch_size = batch_size | |
| stream = file_io.new_input_stream(file_path) | |
| file_length = file_io.get_file_size(file_path) | |
| if supports_pread(stream): | |
| self._stream = stream | |
| def read_at(offset, length): | |
| return pread(stream, length, offset) | |
| else: | |
| self._stream = None | |
| file_data = stream.read() | |
| stream.close() | |
| file_length = len(file_data) | |
| def read_at(offset, length): | |
| return file_data[offset:offset + length] | |
| self._reader = MosaicReader.from_input_file(read_at, file_length) | |
| file_schema_names = set(f.name for f in self._reader.schema) | |
| 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] | |
| if self.existing_fields: | |
| self._reader.project(self.existing_fields) | |
| if self.missing_fields: | |
| output_schema = PyarrowFieldParser.from_paimon_schema(read_fields) | |
| self._missing_out_fields = [] | |
| for name in self.missing_fields: | |
| idx = output_schema.get_field_index(name) | |
| col_type = output_schema.field(idx).type if idx >= 0 else pa.null() | |
| nullable = not SpecialFields.is_system_field(name) | |
| self._missing_out_fields.append(pa.field(name, col_type, nullable=nullable)) | |
| self._current_rg = 0 | |
| self._num_row_groups = self._reader.num_row_groups | |
| self._current_batches = None | |
| self._predicate = push_down_predicate | |
| self._row_group_predicate = None | |
| if row_group_predicate is not None: | |
| try: | |
| self._row_group_predicate = rewrite_predicate_indices( | |
| row_group_predicate, read_fields) | |
| except ValueError: | |
| pass | |
| def _next_row_group_batches(self): | |
| while self._current_rg < self._num_row_groups: | |
| if not self._matches_row_group(self._current_rg): | |
| self._current_rg += 1 | |
| continue | |
| batch = self._reader.read_row_group(self._current_rg) | |
| self._current_rg += 1 | |
| if batch.num_rows == 0: | |
| continue | |
| batch = self._fill_missing_fields(batch) | |
| if self._predicate is not None: | |
| dataset = ds.InMemoryDataset(pa.Table.from_batches([batch])) | |
| scanner = dataset.scanner(filter=self._predicate, batch_size=self._batch_size) | |
| return scanner.to_reader() | |
| else: | |
| return iter(pa.Table.from_batches([batch]).to_batches( | |
| max_chunksize=self._batch_size)) | |
| return None | |
| def _matches_row_group(self, row_group_index: int) -> bool: | |
| if self._row_group_predicate is None: | |
| return True | |
| try: | |
| row_count = self._reader.row_group_num_rows(row_group_index) | |
| stats_map = self._reader.get_row_group_statistics(row_group_index) | |
| except Exception: | |
| return True | |
| if not stats_map: | |
| return True | |
| try: | |
| stats = self._to_simple_stats(stats_map) | |
| return self._row_group_predicate.test_by_simple_stats(stats, row_count) | |
| except Exception: | |
| return True | |
| def _to_simple_stats(self, stats_map) -> SimpleStats: | |
| min_values = [None] * len(self._read_fields) | |
| max_values = [None] * len(self._read_fields) | |
| null_counts = [None] * len(self._read_fields) | |
| for i, field in enumerate(self._read_fields): | |
| stats = stats_map.get(field.name) | |
| if stats is None: | |
| continue | |
| null_counts[i] = stats.null_count | |
| if stats.has_min_max: | |
| min_values[i] = _convert_stats_value(stats.min, field.type) | |
| max_values[i] = _convert_stats_value( | |
| stats.max, field.type, round_up_timestamp=True) | |
| return SimpleStats( | |
| GenericRow(min_values, self._read_fields), | |
| GenericRow(max_values, self._read_fields), | |
| null_counts) | |
| def _fill_missing_fields(self, batch: RecordBatch) -> RecordBatch: | |
| if not self.missing_fields: | |
| return batch | |
| all_columns = [] | |
| out_fields = [] | |
| for field_name in self._read_field_names: | |
| if field_name in self.existing_fields: | |
| col_idx = self.existing_fields.index(field_name) | |
| all_columns.append(batch.column(col_idx)) | |
| out_fields.append(batch.schema.field(col_idx)) | |
| else: | |
| miss_idx = self.missing_fields.index(field_name) | |
| out_field = self._missing_out_fields[miss_idx] | |
| all_columns.append(pa.nulls(batch.num_rows, type=out_field.type)) | |
| out_fields.append(out_field) | |
| return pa.RecordBatch.from_arrays(all_columns, schema=pa.schema(out_fields)) | |
| def read_arrow_batch(self) -> Optional[RecordBatch]: | |
| while True: | |
| if self._current_batches is not None: | |
| try: | |
| if hasattr(self._current_batches, 'read_next_batch'): | |
| return self._current_batches.read_next_batch() | |
| else: | |
| return next(self._current_batches) | |
| except StopIteration: | |
| self._current_batches = None | |
| self._current_batches = self._next_row_group_batches() | |
| if self._current_batches is None: | |
| return None | |
| def close(self): | |
| if self._stream is not None: | |
| self._stream.close() | |
| self._stream = None | |
| self._reader = None | |
| self._current_batches = None | |
| _EPOCH = datetime(1970, 1, 1) | |
| def _convert_stats_value( | |
| value: Optional[bytes], data_type, round_up_timestamp: bool = False) -> Any: | |
| if value is None: | |
| return None | |
| type_name = str(data_type).upper().strip() | |
| type_name = re.sub(r"\s+NOT\s+NULL$", "", type_name) | |
| type_name = re.sub(r"\s+NULL$", "", type_name) | |
| base_type = type_name.split("(", 1)[0].strip() | |
| if base_type in ("CHAR", "VARCHAR", "STRING"): | |
| return value.decode("utf-8") | |
| if base_type in ("BINARY", "VARBINARY", "BYTES", "BLOB"): | |
| return value | |
| if len(value) == 0: | |
| return None | |
| if base_type == "BOOLEAN": | |
| return value[0] != 0 | |
| if base_type == "TINYINT": | |
| return struct.unpack(">b", value[:1])[0] | |
| if base_type == "SMALLINT": | |
| return struct.unpack(">h", value)[0] | |
| if base_type in ("INT", "INTEGER", "DATE", "TIME"): | |
| return struct.unpack(">i", value)[0] | |
| if base_type == "BIGINT": | |
| return struct.unpack(">q", value)[0] | |
| if base_type == "FLOAT": | |
| return struct.unpack(">f", value)[0] | |
| if base_type == "DOUBLE": | |
| return struct.unpack(">d", value)[0] | |
| if base_type in ("DECIMAL", "NUMERIC"): | |
| scale = _type_scale(type_name) | |
| return Decimal(int.from_bytes(value, byteorder="big", signed=True)).scaleb(-scale) | |
| if base_type in ("TIMESTAMP", "TIMESTAMP_LTZ"): | |
| precision = _type_precision(type_name, default=6) | |
| if precision <= 3: | |
| return _EPOCH + timedelta(milliseconds=struct.unpack(">q", value)[0]) | |
| if precision <= 6: | |
| return _EPOCH + timedelta(microseconds=struct.unpack(">q", value)[0]) | |
| millis, nanos_of_milli = struct.unpack(">qi", value) | |
| microseconds = nanos_of_milli // 1000 | |
| if round_up_timestamp and nanos_of_milli % 1000 != 0: | |
| microseconds += 1 | |
| return _EPOCH + timedelta( | |
| milliseconds=millis, | |
| microseconds=microseconds) | |
| return None | |
| def _type_precision(type_name: str, default: int = 0) -> int: | |
| params = _type_params(type_name) | |
| if not params: | |
| return default | |
| return int(params[0]) | |
| def _type_scale(type_name: str) -> int: | |
| params = _type_params(type_name) | |
| if len(params) < 2: | |
| return 0 | |
| return int(params[1]) | |
| def _type_params(type_name: str) -> List[str]: | |
| match = re.search(r"\(([^)]*)\)", type_name) | |
| if match is None: | |
| return [] | |
| return [p.strip() for p in match.group(1).split(",") if p.strip()] |