| # 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. |
| |
| from typing import Callable, Optional, Set |
| |
| import pyarrow |
| from pyarrow import RecordBatch |
| |
| from pypaimon.common.options.core_options import CoreOptions |
| from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader |
| from pypaimon.table.row.blob import Blob, BlobViewStruct |
| |
| |
| class BlobInlineConvertReader(RecordBatchReader): |
| """Resolves BlobView and BlobDescriptor fields in record batches. |
| |
| Processing is split into two clear stages: |
| Stage 1 (BlobView resolution): If view fields exist, use a lightweight |
| prescan reader (only projecting view columns) to collect |
| BlobViewStructs, bulk-preload their descriptors, then read |
| full data from the main reader and replace view field values |
| with descriptor bytes or real blob data according to the |
| blob-as-descriptor option. |
| Stage 2 (BlobDescriptor resolution): Controlled by blob-as-descriptor option. |
| If false, resolve BlobDescriptor bytes from descriptor fields |
| into real blob data bytes. BlobView fields are already resolved |
| in Stage 1 with the upstream table FileIO. |
| If true, return as-is. |
| """ |
| |
| def __init__(self, inner: RecordBatchReader, table, |
| prescan_reader_factory: Optional[Callable[[Set[str]], RecordBatchReader]] = None, |
| blob_parallelism: int = 1): |
| """ |
| Args: |
| inner: The main data reader (reads all columns). |
| table: The table instance. |
| prescan_reader_factory: Optional factory that creates a lightweight |
| reader projecting only the specified field names. Used for |
| prescan to collect BlobViewStructs without reading all columns. |
| Signature: (field_names: Set[str]) -> RecordBatchReader |
| blob_parallelism: number of threads for concurrent blob reads. |
| """ |
| self._inner = inner |
| self._table = table |
| self._prescan_reader_factory = prescan_reader_factory |
| self._blob_parallelism = blob_parallelism |
| self._adopt_metadata(inner) |
| # Preserve original BlobViewStruct bytes when resolve disabled: skip both |
| # view resolution (Stage 1) and descriptor-to-data resolution (Stage 2). |
| resolve_enabled = CoreOptions.blob_view_resolve_enabled( |
| table.options) and self._table.catalog_environment.catalog_loader is not None |
| self._view_fields = CoreOptions.blob_view_fields(table.options) if resolve_enabled else set() |
| self._descriptor_fields = CoreOptions.blob_descriptor_fields(table.options) |
| self._blob_as_descriptor = CoreOptions.blob_as_descriptor(table.options) |
| self._prescan_done = False |
| self._blob_view_lookup = None |
| |
| def read_arrow_batch(self) -> Optional[RecordBatch]: |
| # Align with Java: only enter blob view resolution when catalog_loader is available |
| # If catalog_loader is None, skip both Stage 1 (view resolution) and Stage 2 (descriptor resolution) |
| if self._view_fields and not self._prescan_done: |
| self._prescan_view_structs() |
| |
| batch = self._inner.read_arrow_batch() |
| if batch is None: |
| return None |
| # Resolve view fields using the preloaded lookup |
| view_file_ios = {} |
| if self._view_fields and self._blob_view_lookup is not None: |
| batch, view_file_ios = self._resolve_view_fields(batch, self._blob_view_lookup) |
| # Resolve BlobDescriptor -> real bytes (if blob-as-descriptor=false) |
| return self._resolve_descriptor_fields(batch, view_file_ios) |
| |
| # ------------------------------------------------------------------ |
| # Stage 1: BlobView prescan (lightweight, only reads view columns) |
| # ------------------------------------------------------------------ |
| |
| def _prescan_view_structs(self): |
| """Use a lightweight prescan reader (projecting only view columns) to |
| collect all BlobViewStructs and bulk-preload their descriptors.""" |
| from pypaimon.table.row.blob import BlobViewStruct |
| from pypaimon.utils.blob_view_lookup import BlobViewLookup |
| |
| all_view_structs = [] |
| |
| prescan_reader = self._prescan_reader_factory(self._view_fields) |
| try: |
| while True: |
| batch = prescan_reader.read_arrow_batch() |
| if batch is None: |
| break |
| for field_name in self._view_fields: |
| if field_name not in batch.schema.names: |
| continue |
| for value in batch.column(field_name).to_pylist(): |
| value = self._normalize_blob_to_bytes(value) |
| if value is None: |
| continue |
| if isinstance(value, bytes) and BlobViewStruct.is_blob_view_struct(value): |
| all_view_structs.append(BlobViewStruct.deserialize(value)) |
| else: |
| raise ValueError( |
| f"Expected BlobViewStruct bytes in view field '{field_name}', " |
| f"but got non-BlobViewStruct bytes." |
| ) |
| finally: |
| prescan_reader.close() |
| |
| # Bulk-preload BlobViewStruct -> BlobDescriptor mapping |
| if all_view_structs: |
| self._blob_view_lookup = BlobViewLookup(self._table) |
| self._blob_view_lookup.preload(all_view_structs) |
| self._prescan_done = True |
| |
| def _resolve_view_fields(self, batch, blob_view_lookup): |
| """Replace BlobViewStruct bytes in view fields with descriptor bytes.""" |
| view_file_ios = {} |
| for field_name in self._view_fields: |
| if field_name not in batch.schema.names: |
| continue |
| values = [self._normalize_blob_to_bytes(v) for v in batch.column(field_name).to_pylist()] |
| converted_values = [] |
| field_file_ios = [] |
| for value in values: |
| if value is None or not ( |
| isinstance(value, bytes) and BlobViewStruct.is_blob_view_struct(value)): |
| converted_values.append(value) |
| field_file_ios.append(None) |
| continue |
| |
| view_struct = BlobViewStruct.deserialize(value) |
| if blob_view_lookup.resolve_to_null(view_struct): |
| converted_values.append(None) |
| field_file_ios.append(None) |
| else: |
| descriptor = blob_view_lookup.resolve_descriptor(view_struct) |
| converted_values.append(descriptor.serialize()) |
| file_io = blob_view_lookup.resolve_file_io(view_struct) |
| field_file_ios.append(file_io) |
| |
| column_idx = batch.schema.names.index(field_name) |
| batch = batch.set_column( |
| column_idx, |
| pyarrow.field(field_name, pyarrow.large_binary(), nullable=True), |
| pyarrow.array(converted_values, type=pyarrow.large_binary()), |
| ) |
| view_file_ios[field_name] = field_file_ios |
| return batch, view_file_ios |
| |
| # ------------------------------------------------------------------ |
| # Stage 2: BlobData resolution (unified exit) |
| # ------------------------------------------------------------------ |
| |
| def _resolve_descriptor_fields(self, batch, view_file_ios=None): |
| if self._blob_as_descriptor: |
| return batch |
| |
| for field_name in self._descriptor_fields: |
| if field_name not in batch.schema.names: |
| continue |
| values = [self._normalize_blob_to_bytes(v) for v in batch.column(field_name).to_pylist()] |
| blobs = [Blob.from_bytes(v, self._table.file_io) for v in values] |
| |
| if self._blob_parallelism > 1: |
| converted_values = self._table.file_io.read_blobs_concurrent( |
| blobs, self._blob_parallelism) |
| else: |
| converted_values = [b.to_data() if b else None for b in blobs] |
| |
| column_idx = batch.schema.names.index(field_name) |
| batch = batch.set_column( |
| column_idx, |
| pyarrow.field(field_name, pyarrow.large_binary(), nullable=True), |
| pyarrow.array(converted_values, type=pyarrow.large_binary()), |
| ) |
| |
| view_file_ios = view_file_ios or {} |
| for field_name in self._view_fields: |
| field_file_ios = view_file_ios.get(field_name) |
| if field_name not in batch.schema.names or field_file_ios is None: |
| continue |
| values = [self._normalize_blob_to_bytes(v) for v in batch.column(field_name).to_pylist()] |
| blobs_by_file_io = {} |
| converted_values = [] |
| |
| for idx, value in enumerate(values): |
| file_io = field_file_ios[idx] or self._table.file_io |
| blob = Blob.from_bytes(value, file_io) |
| if self._blob_parallelism > 1: |
| converted_values.append(None) |
| if blob is not None: |
| blobs_by_file_io.setdefault(file_io, []).append((idx, blob)) |
| else: |
| converted_values.append(blob.to_data() if blob else None) |
| |
| for file_io, indexed_blobs in blobs_by_file_io.items(): |
| blobs = [item[1] for item in indexed_blobs] |
| results = file_io.read_blobs_concurrent(blobs, self._blob_parallelism) |
| for (idx, _), data in zip(indexed_blobs, results): |
| converted_values[idx] = data |
| |
| column_idx = batch.schema.names.index(field_name) |
| batch = batch.set_column( |
| column_idx, |
| pyarrow.field(field_name, pyarrow.large_binary(), nullable=True), |
| pyarrow.array(converted_values, type=pyarrow.large_binary()), |
| ) |
| |
| return batch |
| |
| # ------------------------------------------------------------------ |
| # Utilities |
| # ------------------------------------------------------------------ |
| |
| @staticmethod |
| def _normalize_blob_to_bytes(value): |
| if value is None: |
| return None |
| if hasattr(value, 'as_py'): |
| value = value.as_py() |
| if isinstance(value, str): |
| value = value.encode('utf-8') |
| if isinstance(value, bytearray): |
| value = bytes(value) |
| return value |
| |
| def close(self): |
| self._inner.close() |
