From d4f63ae749ed994b7d9800a7048d3358d7bbbdad Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Fri, 24 Jul 2026 04:53:17 -0700 Subject: [PATCH 1/3] [python] Read VARIANT Parquet in bounded batches PyArrow dataset scans can fail when a nested VARIANT child produces multiple chunks. Read projected Parquet VARIANT columns with ParquetFile.iter_batches while preserving row-group pruning, predicates, and projection. --- .../read/reader/format_pyarrow_reader.py | 196 ++++++++++++------ paimon-python/pypaimon/read/split_read.py | 6 +- .../format_pyarrow_variant_row_group_test.py | 184 ++++++++++++++++ 3 files changed, 326 insertions(+), 60 deletions(-) create mode 100644 paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py index efaf8e0546b1..30e13188c679 100644 --- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py +++ b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Set import pyarrow as pa import pyarrow.dataset as ds @@ -30,7 +30,8 @@ is_shredded_variant, ) from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader -from pypaimon.schema.data_types import DataField, PyarrowFieldParser +from pypaimon.schema.data_types import (AtomicType, DataField, + PyarrowFieldParser) from pypaimon.table.special_fields import SpecialFields @@ -48,7 +49,9 @@ 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): + nested_name_paths: Optional[List[List[str]]] = None, + predicate_field_names: Optional[Set[str]] = None): + self._predicate_field_names = predicate_field_names or set() file_path_for_pyarrow = file_io.to_filesystem_path(file_path) self.dataset = ds.dataset(file_path_for_pyarrow, format=file_format, filesystem=file_io.filesystem) self._file_format = file_format @@ -93,69 +96,140 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, for f, path in zip(read_fields, nested_name_paths): if f.name in existing_set: columns_dict[f.name] = ds.field(*path) - self.reader = self.dataset.scanner( - columns=columns_dict, - filter=push_down_predicate, - batch_size=batch_size - ).to_reader() + self._scan_columns = columns_dict else: # Only pass existing fields to PyArrow scanner to avoid errors - self.reader = self.dataset.scanner( - columns=self.existing_fields, - filter=push_down_predicate, - batch_size=batch_size - ).to_reader() + 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 ) - def read_arrow_batch(self) -> Optional[RecordBatch]: + # Read projected VARIANT columns in bounded batches. + self._parquet_file = None + if (self._file_format == 'parquet' + and not has_nested_path + and self._has_projected_variant()): + 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._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 isinstance(f.type, AtomicType) + and f.type.type == 'VARIANT' + 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._scan_filter is None: + yield 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 not self.existing_fields: + return None + 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 _surviving_row_group_ids(self): + total = self._parquet_file.num_row_groups + if self._scan_filter is None: + return range(total) try: - batch = self.reader.read_next_batch() - - if self._file_format == 'orc' and self._output_schema is not None: - batch = self._cast_orc_time_columns(batch) - - if self._shredded_schemas: - 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 - ] - - # Reconstruct the batch with all fields in the correct order - all_columns = [] - out_fields = [] - for field_name in self._read_field_names: - if field_name in self.existing_fields: - # Get the column from the existing batch - column_idx = self.existing_fields.index(field_name) - all_columns.append(batch.column(column_idx)) - out_fields.append(batch.schema.field(column_idx)) - else: - # Get the column from missing fields - 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)) - # Create a new RecordBatch with all columns - return pa.RecordBatch.from_arrays(all_columns, schema=pa.schema(out_fields)) - - except StopIteration: + 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]: + 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._shredded_schemas: + 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: """Replace shredded VARIANT columns with standard struct.""" @@ -197,8 +271,12 @@ def _cast_orc_time_columns(self, batch): return batch def close(self): - if self.reader is not None: - self.reader = None + 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: diff --git a/paimon-python/pypaimon/read/split_read.py b/paimon-python/pypaimon/read/split_read.py index 6ced11f0c127..71e6b623f4c0 100644 --- a/paimon-python/pypaimon/read/split_read.py +++ b/paimon-python/pypaimon/read/split_read.py @@ -330,11 +330,15 @@ def file_reader_supplier(self, file: DataFileMeta, for_merge_read: bool, [nested_path_by_name[f.name] for f in ordered_read_fields] if has_nested else None ) + predicate_fields = ( + predicate_field_names(self.push_down_predicate) + if self.push_down_predicate else set()) format_reader = FormatPyArrowReader( self.table.file_io, file_format, file_path, ordered_read_fields, read_arrow_predicate, batch_size=batch_size, options=self.table.options, - nested_name_paths=ordered_nested_paths) + nested_name_paths=ordered_nested_paths, + predicate_field_names=predicate_fields) elif file_format == CoreOptions.FILE_FORMAT_ROW: if has_nested: raise NotImplementedError( diff --git a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py new file mode 100644 index 000000000000..8e012ca4f32b --- /dev/null +++ b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py @@ -0,0 +1,184 @@ +# 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 inspect +import os +import shutil +import tempfile +import unittest + +import pyarrow as pa +import pyarrow.dataset as ds +import pyarrow.fs as pafs +import pyarrow.parquet as pq + +from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader +from pypaimon.schema.data_types import AtomicType, DataField + +_VARIANT_TYPE = pa.struct([ + pa.field("value", pa.binary(), nullable=False), + pa.field("metadata", pa.binary(), nullable=False), +]) + + +class _LocalFileIO: + filesystem = pafs.LocalFileSystem() + + def to_filesystem_path(self, path): + return path + + +def _drain(reader): + rows = 0 + columns = None + content_keys = set() + while True: + batch = reader.read_arrow_batch() + if batch is None: + break + rows += batch.num_rows + columns = batch.schema.names + if "content_key" in columns: + content_keys |= set( + batch.column(columns.index("content_key")).to_pylist()) + return rows, columns, content_keys + + +class VariantRowGroupReaderTest(unittest.TestCase): + + def setUp(self): + self.n = 2000 + content_key = [ + "robot_pose_raw" if i % 2 == 0 else "imu_raw" + for i in range(self.n) + ] + payload = [ + {"value": b"v%d" % i, "metadata": b"m"} + for i in range(self.n) + ] + table = pa.table({ + "content_key": pa.array(content_key), + "payload": pa.array(payload, type=_VARIANT_TYPE), + }) + self.tmp = tempfile.mkdtemp() + self.path = os.path.join(self.tmp, "topics.parquet") + pq.write_table(table, self.path, row_group_size=1000) + self.assertEqual(2, pq.ParquetFile(self.path).num_row_groups) + self.read_fields = [ + DataField(0, "content_key", AtomicType("STRING")), + DataField(1, "payload", AtomicType("VARIANT")), + ] + + def _reader(self, read_fields, predicate=None, predicate_field_names=None): + return FormatPyArrowReader( + _LocalFileIO(), "parquet", self.path, read_fields, + predicate, batch_size=256, + predicate_field_names=predicate_field_names) + + def _large_dictionary_variant(self): + if "store_schema" not in inspect.signature(pq.write_table).parameters: + self.skipTest("PyArrow does not support store_schema") + row_count = 15_340 + values = pa.DictionaryArray.from_arrays( + pa.array([0] * row_count, type=pa.int32()), + pa.array([b"x" * 140_000], type=pa.binary())) + payload = pa.StructArray.from_arrays( + [values, pa.array([b"m"] * row_count)], + names=["value", "metadata"]) + path = os.path.join(self.tmp, "large-dictionary-variant.parquet") + pq.write_table( + pa.table({"payload": payload}), path, + use_dictionary=True, compression="zstd", store_schema=False) + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [DataField(0, "payload", AtomicType("VARIANT"))], + None, batch_size=128) + return pq.ParquetFile(path).num_row_groups, _drain(reader) + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_filter_only_column_not_in_projection(self): + reader = self._reader( + [DataField(1, "payload", AtomicType("VARIANT"))], + predicate=ds.field("content_key") == "robot_pose_raw", + predicate_field_names={"content_key"}) + rows, columns, _ = _drain(reader) + self.assertEqual(self.n // 2, rows) + self.assertEqual(["payload"], columns) + + def test_reads_all_rows_across_row_groups(self): + rows, columns, _ = _drain(self._reader(self.read_fields)) + self.assertEqual(self.n, rows) + self.assertEqual(["content_key", "payload"], columns) + + def test_reads_large_dictionary_variant_in_single_row_group(self): + row_groups, (rows, columns, _) = self._large_dictionary_variant() + self.assertEqual(1, row_groups) + self.assertEqual(15_340, rows) + self.assertEqual(["payload"], columns) + + def test_predicate_returns_only_matching_rows(self): + predicate = ds.field("content_key") == "robot_pose_raw" + rows, _, content_keys = _drain( + self._reader(self.read_fields, predicate)) + self.assertEqual(self.n // 2, rows) + self.assertEqual({"robot_pose_raw"}, content_keys) + + def test_projection_returns_only_requested_columns(self): + rows, columns, _ = _drain( + self._reader([DataField(1, "payload", AtomicType("VARIANT"))])) + self.assertEqual(self.n, rows) + self.assertEqual(["payload"], columns) + + def test_single_row_group_scalar_read_uses_fast_path(self): + rows, columns, _ = _drain( + self._reader([DataField(0, "content_key", AtomicType("STRING"))])) + self.assertEqual(self.n, rows) + self.assertEqual(["content_key"], columns) + + def test_row_group_pruning_by_statistics(self): + rows_per_group = 1000 + content_key = [] + for group in range(8): + content_key += ( + ["match" if group == 3 else "other"] * rows_per_group) + total = len(content_key) + payload = [ + {"value": b"v%d" % i, "metadata": b"m"} + for i in range(total) + ] + path = os.path.join(self.tmp, "clustered.parquet") + pq.write_table( + pa.table({ + "content_key": pa.array(content_key), + "payload": pa.array(payload, type=_VARIANT_TYPE), + }), + path, row_group_size=rows_per_group) + self.assertEqual(8, pq.ParquetFile(path).num_row_groups) + + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, self.read_fields, + ds.field("content_key") == "match", batch_size=512) + self.assertEqual([3], list(reader._surviving_row_group_ids())) + rows, _, keys = _drain(reader) + self.assertEqual(rows_per_group, rows) + self.assertEqual({"match"}, keys) + + +if __name__ == "__main__": + unittest.main() From 705a6f393df8d6e01a87fdd51be28d02d449d394 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 25 Jul 2026 23:02:37 +0800 Subject: [PATCH 2/3] [python] Test VARIANT projection order --- .../tests/format_pyarrow_variant_row_group_test.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py index 8e012ca4f32b..4acd82b503e9 100644 --- a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py +++ b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py @@ -145,6 +145,17 @@ def test_projection_returns_only_requested_columns(self): self.assertEqual(self.n, rows) self.assertEqual(["payload"], columns) + def test_projection_preserves_requested_order(self): + reader = self._reader([ + DataField(1, "payload", AtomicType("VARIANT")), + DataField(0, "content_key", AtomicType("STRING")), + ]) + batch = reader.read_arrow_batch() + self.assertEqual(["payload", "content_key"], batch.schema.names) + self.assertEqual( + {"value": b"v0", "metadata": b"m"}, batch.column(0)[0].as_py()) + self.assertEqual("robot_pose_raw", batch.column(1)[0].as_py()) + def test_single_row_group_scalar_read_uses_fast_path(self): rows, columns, _ = _drain( self._reader([DataField(0, "content_key", AtomicType("STRING"))])) From 685f4101fbb4d33461053fbcdcb22f9c8d73bfa1 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 25 Jul 2026 23:47:29 +0800 Subject: [PATCH 3/3] [python] Fix dotted and nested VARIANT projections --- .../pypaimon/data/variant_shredding.py | 2 +- .../read/reader/format_pyarrow_reader.py | 247 ++++++++++++++++-- .../format_pyarrow_variant_row_group_test.py | 243 ++++++++++++++++- 3 files changed, 462 insertions(+), 30 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_shredding.py b/paimon-python/pypaimon/data/variant_shredding.py index f14946ba330e..72a4508dd32a 100644 --- a/paimon-python/pypaimon/data/variant_shredding.py +++ b/paimon-python/pypaimon/data/variant_shredding.py @@ -182,7 +182,7 @@ def _parse_typed_value_field(schema: VariantSchema, tv_type: pa.DataType) -> Var if pa.types.is_struct(tv_type): object_fields: List[ObjectField] = [] for j in range(tv_type.num_fields): - sub_f = tv_type.field(j) + sub_f = tv_type[j] if pa.types.is_struct(sub_f.type): sub_schema = build_variant_schema(sub_f.type) else: diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py index 30e13188c679..956ca39f467a 100644 --- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py +++ b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py @@ -30,8 +30,15 @@ is_shredded_variant, ) from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader -from pypaimon.schema.data_types import (AtomicType, DataField, - PyarrowFieldParser) +from pypaimon.schema.data_types import ( + ArrayType, + AtomicType, + DataField, + MapType, + MultisetType, + PyarrowFieldParser, + RowType, +) from pypaimon.table.special_fields import SpecialFields @@ -65,6 +72,7 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, 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: @@ -80,23 +88,21 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, 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._shredded_schemas: Dict[str, VariantSchema] = {} - if options is None or options.variant_shredding_enabled(): - top_level_names = set(file_schema.names) - for name in self.existing_fields: - if name not in top_level_names: - continue - field_type = file_schema.field(name).type - if is_shredded_variant(field_type): - self._shredded_schemas[name] = build_variant_schema(field_type) + self._variant_shredding_enabled = ( + options is None or options.variant_shredding_enabled()) + self._variant_schema_cache: Dict[pa.DataType, VariantSchema] = {} - if has_nested_path: + 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 @@ -109,9 +115,7 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, # Read projected VARIANT columns in bounded batches. self._parquet_file = None - if (self._file_format == 'parquet' - and not has_nested_path - and self._has_projected_variant()): + if self._bounded_variant_read: import pyarrow.parquet as pq # ParquetFile(filesystem=...) is unavailable in PyArrow 6. self._parquet_file = pq.ParquetFile( @@ -129,8 +133,7 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, def _has_projected_variant(self) -> bool: return any( f.name in self.existing_fields - and isinstance(f.type, AtomicType) - and f.type.type == 'VARIANT' + and _contains_variant(f.type) for f in self.read_fields) @staticmethod @@ -148,8 +151,20 @@ def _iter_row_group_batches(self): 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 batch + yield self._select_existing_fields(batch) continue table = ds.dataset( pa.Table.from_batches([batch]) @@ -161,9 +176,14 @@ def _iter_row_group_batches(self): yield out def _row_group_read_columns(self): - if not self.existing_fields: - return None - columns = list(self.existing_fields) + 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: @@ -171,6 +191,35 @@ def _row_group_read_columns(self): columns.append(name) return columns + def _select_existing_fields(self, batch): + 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) + 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: @@ -196,7 +245,7 @@ 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._shredded_schemas: + if self._variant_shredding_enabled: batch = self._assemble_shredded_variants(batch) if not self.missing_fields: @@ -232,15 +281,19 @@ def _type_for_missing(name: str) -> pa.DataType: all_columns, schema=pa.schema(out_fields)) def _assemble_shredded_variants(self, batch: pa.RecordBatch) -> pa.RecordBatch: - """Replace shredded VARIANT columns with standard struct.""" 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): - if f.name in self._shredded_schemas: - schema = self._shredded_schemas[f.name] - new_col = assemble_shredded_column(columns[i], schema) + 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 @@ -292,5 +345,145 @@ def _path_exists_in_arrow_schema(schema: pa.Schema, path: List[str]) -> bool: idx = current_type.get_field_index(name) if idx < 0: return False - current_type = current_type.field(idx).type + current_type = current_type[idx].type return True + + +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 diff --git a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py index 4acd82b503e9..bb0f06bc6ee3 100644 --- a/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py +++ b/paimon-python/pypaimon/tests/format_pyarrow_variant_row_group_test.py @@ -16,6 +16,7 @@ # under the License. import inspect +import json import os import shutil import tempfile @@ -26,8 +27,20 @@ import pyarrow.fs as pafs import pyarrow.parquet as pq +from pypaimon.data.generic_variant import GenericVariant +from pypaimon.data.variant_shredding import ( + parse_shredding_schema_option, + shredding_schema_to_arrow_type, + shred_variant_column, +) from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader -from pypaimon.schema.data_types import AtomicType, DataField +from pypaimon.schema.data_types import ( + ArrayType, + AtomicType, + DataField, + MapType, + RowType, +) _VARIANT_TYPE = pa.struct([ pa.field("value", pa.binary(), nullable=False), @@ -89,7 +102,7 @@ def _reader(self, read_fields, predicate=None, predicate_field_names=None): predicate, batch_size=256, predicate_field_names=predicate_field_names) - def _large_dictionary_variant(self): + def _large_dictionary_payload(self): if "store_schema" not in inspect.signature(pq.write_table).parameters: self.skipTest("PyArrow does not support store_schema") row_count = 15_340 @@ -99,6 +112,10 @@ def _large_dictionary_variant(self): payload = pa.StructArray.from_arrays( [values, pa.array([b"m"] * row_count)], names=["value", "metadata"]) + return row_count, payload + + def _large_dictionary_variant(self): + row_count, payload = self._large_dictionary_payload() path = os.path.join(self.tmp, "large-dictionary-variant.parquet") pq.write_table( pa.table({"payload": payload}), path, @@ -109,6 +126,31 @@ def _large_dictionary_variant(self): None, batch_size=128) return pq.ParquetFile(path).num_row_groups, _drain(reader) + def _shredded_variant_payload(self): + schema = json.dumps({ + "type": "ROW", + "fields": [{ + "id": 0, + "name": "v", + "type": { + "type": "ROW", + "fields": [{ + "id": 1, + "name": "age", + "type": "BIGINT", + }], + }, + }], + }) + obj_fields = parse_shredding_schema_option(schema)["v"] + target_type = shredding_schema_to_arrow_type(obj_fields) + variants = GenericVariant.to_arrow_array([ + GenericVariant.from_python({"age": 1, "extra": "x"}), + None, + GenericVariant.from_python({"age": 3, "extra": "z"}), + ]) + return shred_variant_column(variants, obj_fields, target_type) + def tearDown(self): shutil.rmtree(self.tmp, ignore_errors=True) @@ -156,6 +198,203 @@ def test_projection_preserves_requested_order(self): {"value": b"v0", "metadata": b"m"}, batch.column(0)[0].as_py()) self.assertEqual("robot_pose_raw", batch.column(1)[0].as_py()) + def test_dotted_top_level_name_does_not_match_nested_path(self): + path = os.path.join(self.tmp, "dotted-name.parquet") + pq.write_table(pa.table({ + "a": pa.array( + [{"b": "nested"}], + type=pa.struct([pa.field("b", pa.string())])), + "a.b": pa.array(["top"]), + "payload": pa.array( + [{"value": b"v", "metadata": b"m"}], type=_VARIANT_TYPE), + }), path) + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [ + DataField(1, "a.b", AtomicType("STRING")), + DataField(2, "payload", AtomicType("VARIANT")), + ], + None, batch_size=128) + batch = reader.read_arrow_batch() + self.assertEqual(["a.b", "payload"], batch.schema.names) + self.assertEqual("top", batch.column(0)[0].as_py()) + self.assertEqual( + {"value": b"v", "metadata": b"m"}, batch.column(1)[0].as_py()) + + def test_reads_variant_nested_in_container_types(self): + row_count, payload = self._large_dictionary_payload() + offsets = pa.array(range(row_count + 1), type=pa.int32()) + variant_type = AtomicType("VARIANT") + cases = [ + ( + "row", + pa.StructArray.from_arrays([payload], names=["v"]), + RowType(True, [DataField(1, "v", variant_type)]), + ), + ( + "array", + pa.ListArray.from_arrays(offsets, payload), + ArrayType(True, variant_type), + ), + ( + "map", + pa.MapArray.from_arrays( + offsets, pa.array(["k"] * row_count), payload), + MapType(True, AtomicType("STRING", False), variant_type), + ), + ] + for name, column, data_type in cases: + with self.subTest(name=name): + path = os.path.join(self.tmp, "nested-{}.parquet".format(name)) + pq.write_table( + pa.table({name: column}), path, + use_dictionary=True, compression="zstd", + store_schema=False) + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [DataField(0, name, data_type)], None, batch_size=128) + rows, columns, _ = _drain(reader) + self.assertEqual(row_count, rows) + self.assertEqual([name], columns) + + def test_nested_variant_paths_use_bounded_reader(self): + path = os.path.join(self.tmp, "small-nested-variant.parquet") + payload = pa.array([ + {"value": b"v0", "metadata": b"m"}, + {"value": b"v1", "metadata": b"m"}, + {"value": b"v2", "metadata": b"m"}, + ], type=_VARIANT_TYPE) + offsets = pa.array([0, 1, 2, 3], type=pa.int32()) + pq.write_table(pa.table({ + "row": pa.StructArray.from_arrays( + [payload], names=["v"], + mask=pa.array([False, False, True])), + "array": pa.ListArray.from_arrays(offsets, payload), + "map": pa.MapArray.from_arrays( + offsets, pa.array(["k", "k", "k"]), payload), + "kind": pa.array(["keep", "drop", "keep"]), + }), path) + + variant_type = AtomicType("VARIANT") + cases = [ + ("row", RowType(True, [DataField(1, "v", variant_type)])), + ("array", ArrayType(True, variant_type)), + ("map", MapType( + True, AtomicType("STRING", False), variant_type)), + ] + for name, data_type in cases: + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [DataField(0, name, data_type)], None, batch_size=128) + self.assertIsNotNone(reader._parquet_file) + self.assertEqual(3, _drain(reader)[0]) + + nested_reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [DataField(1, "row_v", AtomicType("VARIANT"))], + ds.field("kind") == "keep", batch_size=128, + nested_name_paths=[["row", "v"]], + predicate_field_names={"kind"}) + self.assertIsNotNone(nested_reader._parquet_file) + batch = nested_reader.read_arrow_batch() + self.assertEqual(["row_v"], batch.schema.names) + self.assertEqual([ + {"value": b"v0", "metadata": b"m"}, None, + ], batch.column(0).to_pylist()) + + def test_reads_nested_variant_projection(self): + row_count, payload = self._large_dictionary_payload() + path = os.path.join(self.tmp, "nested-projection.parquet") + pq.write_table( + pa.table({ + "row": pa.StructArray.from_arrays([payload], names=["v"]), + "kind": pa.array( + ["keep" if i % 2 == 0 else "drop" + for i in range(row_count)]), + }), + path, use_dictionary=True, compression="zstd", store_schema=False) + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [DataField(0, "row_v", AtomicType("VARIANT"))], + ds.field("kind") == "keep", batch_size=128, + nested_name_paths=[["row", "v"]], + predicate_field_names={"kind"}) + rows, columns, _ = _drain(reader) + self.assertEqual((row_count + 1) // 2, rows) + self.assertEqual(["row_v"], columns) + + def test_assembles_shredded_variant_in_nested_types(self): + shredded = self._shredded_variant_payload() + values = shredded.to_pylist() + path = os.path.join(self.tmp, "nested-shredded-variant.parquet") + pq.write_table(pa.table({ + "row": pa.array( + [{"v": values[0]}, {"v": values[1]}, None, {"v": values[2]}], + type=pa.struct([pa.field("v", shredded.type)])), + "array": pa.array( + [[values[0]], [values[1]], None, [values[2]]], + type=pa.list_(pa.field("item", shredded.type))), + "map": pa.array( + [[("k", values[0])], [("k", values[1])], None, + [("k", values[2])]], + type=pa.map_( + pa.field("key", pa.string(), nullable=False), + pa.field("value", shredded.type))), + "kind": pa.array(["keep", "drop", "keep", "keep"]), + }), path) + + variant_type = AtomicType("VARIANT") + data_types = { + "row": RowType(True, [DataField(1, "v", variant_type)]), + "array": ArrayType(True, variant_type), + "map": MapType( + True, AtomicType("STRING", False), variant_type), + } + outputs = {} + for name, data_type in data_types.items(): + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [DataField(0, name, data_type)], None, batch_size=128) + outputs[name] = reader.read_arrow_batch().column(0) + + self.assertEqual(_VARIANT_TYPE, outputs["row"].type[0].type) + self.assertEqual(_VARIANT_TYPE, outputs["array"].type.value_type) + self.assertEqual(_VARIANT_TYPE, outputs["map"].type.item_type) + + def decode(value): + if value is None: + return None + return GenericVariant.from_arrow_struct(value).to_python() + + row_values = outputs["row"].to_pylist() + self.assertEqual({"age": 1, "extra": "x"}, decode(row_values[0]["v"])) + self.assertIsNone(row_values[1]["v"]) + self.assertIsNone(row_values[2]) + + array_values = outputs["array"].to_pylist() + self.assertEqual({"age": 1, "extra": "x"}, decode(array_values[0][0])) + self.assertIsNone(array_values[1][0]) + self.assertIsNone(array_values[2]) + + map_values = outputs["map"].to_pylist() + self.assertEqual( + {"age": 1, "extra": "x"}, decode(map_values[0][0][1])) + self.assertIsNone(map_values[1][0][1]) + self.assertIsNone(map_values[2]) + + nested_reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [DataField(1, "row_v", AtomicType("VARIANT"))], + ds.field("kind") == "keep", batch_size=128, + nested_name_paths=[["row", "v"]], + predicate_field_names={"kind"}) + nested = nested_reader.read_arrow_batch().column(0) + self.assertEqual(_VARIANT_TYPE, nested.type) + nested_values = nested.to_pylist() + self.assertEqual({"age": 1, "extra": "x"}, decode(nested_values[0])) + self.assertIsNone(nested_values[1]) + self.assertEqual({"age": 3, "extra": "z"}, decode(nested_values[2])) + def test_single_row_group_scalar_read_uses_fast_path(self): rows, columns, _ = _drain( self._reader([DataField(0, "content_key", AtomicType("STRING"))]))