#49 Fix: Python native read with PyArrow (#53)

Author: chenghuichen

pypaimon/py4j/java_implementation.py

@@ -82,8 +82,12 @@ def new_read_builder(self) -> 'ReadBuilder':
             primary_keys = None
         else:
             primary_keys = [str(key) for key in self._j_table.primaryKeys()]
+        if self._j_table.partitionKeys().isEmpty():
+            partition_keys = None
+        else:
+            partition_keys = [str(key) for key in self._j_table.partitionKeys()]
         return ReadBuilder(j_read_builder, self._j_table.rowType(), self._catalog_options,
-                           primary_keys)
+                           primary_keys, partition_keys)
 
     def new_batch_write_builder(self) -> 'BatchWriteBuilder':
         java_utils.check_batch_write(self._j_table)
@@ -93,11 +97,12 @@ def new_batch_write_builder(self) -> 'BatchWriteBuilder':
 
 class ReadBuilder(read_builder.ReadBuilder):
 
-    def __init__(self, j_read_builder, j_row_type, catalog_options: dict, primary_keys: List[str]):
+    def __init__(self, j_read_builder, j_row_type, catalog_options: dict, primary_keys: List[str], partition_keys: List[str]):
         self._j_read_builder = j_read_builder
         self._j_row_type = j_row_type
         self._catalog_options = catalog_options
         self._primary_keys = primary_keys
+        self._partition_keys = partition_keys
         self._predicate = None
         self._projection = None
 
@@ -128,7 +133,7 @@ def new_scan(self) -> 'TableScan':
     def new_read(self) -> 'TableRead':
         j_table_read = self._j_read_builder.newRead().executeFilter()
         return TableRead(j_table_read, self._j_read_builder.readType(), self._catalog_options,
-                         self._predicate, self._projection, self._primary_keys)
+                         self._predicate, self._projection, self._primary_keys, self._partition_keys)
 
     def new_predicate_builder(self) -> 'PredicateBuilder':
         return PredicateBuilder(self._j_row_type)
@@ -203,14 +208,15 @@ def file_paths(self) -> List[str]:
 class TableRead(table_read.TableRead):
 
     def __init__(self, j_table_read, j_read_type, catalog_options, predicate, projection,
-                 primary_keys: List[str]):
+                 primary_keys: List[str], partition_keys: List[str]):
         self._j_table_read = j_table_read
         self._j_read_type = j_read_type
         self._catalog_options = catalog_options
 
         self._predicate = predicate
         self._projection = projection
         self._primary_keys = primary_keys
+        self._partition_keys = partition_keys
 
         self._arrow_schema = java_utils.to_arrow_schema(j_read_type)
         self._j_bytes_reader = get_gateway().jvm.InvocationUtil.createParallelBytesReader(
@@ -259,7 +265,7 @@ def to_record_generator(self, splits: List['Split']) -> Optional[Iterator[Any]]:
         try:
             j_splits = list(s.to_j_split() for s in splits)
             j_reader = get_gateway().jvm.InvocationUtil.createReader(self._j_table_read, j_splits)
-            converter = ReaderConverter(self._predicate, self._projection, self._primary_keys)
+            converter = ReaderConverter(self._predicate, self._projection, self._primary_keys, self._partition_keys)
             pynative_reader = converter.convert_java_reader(j_reader)
 
             def _record_generator():

pypaimon/pynative/reader/core/columnar_row_iterator.py

@@ -32,7 +32,7 @@ class ColumnarRowIterator(FileRecordIterator[InternalRow]):
 
     def __init__(self, file_path: str, record_batch: pa.RecordBatch):
         self.file_path = file_path
-        self._record_batch = record_batch
+        self.record_batch = record_batch
         self._row = ColumnarRow(record_batch)
 
         self.num_rows = record_batch.num_rows
@@ -58,4 +58,4 @@ def reset(self, next_file_pos: int):
         self.next_file_pos = next_file_pos
 
     def release_batch(self):
-        del self._record_batch
+        del self.record_batch

pypaimon/pynative/reader/data_file_record_reader.py

@@ -16,28 +16,118 @@
 # limitations under the License.
 ################################################################################
 
-from typing import Optional
+from typing import Optional, List, Any
+import pyarrow as pa
 
+from pypaimon.pynative.common.exception import PyNativeNotImplementedError
 from pypaimon.pynative.common.row.internal_row import InternalRow
 from pypaimon.pynative.reader.core.file_record_iterator import FileRecordIterator
 from pypaimon.pynative.reader.core.file_record_reader import FileRecordReader
 from pypaimon.pynative.reader.core.record_reader import RecordReader
+from pypaimon.pynative.reader.core.columnar_row_iterator import ColumnarRowIterator
+
+
+class PartitionInfo:
+    """
+    Partition information about how the row mapping of outer row.
+    """
+
+    def __init__(self, mapping: List[int], partition_values: List[Any]):
+        self.mapping = mapping  # Mapping array similar to Java version
+        self.partition_values = partition_values  # Partition values to be injected
+
+    def size(self) -> int:
+        return len(self.mapping) - 1
+
+    def in_partition_row(self, pos: int) -> bool:
+        return self.mapping[pos] < 0
+
+    def get_real_index(self, pos: int) -> int:
+        return abs(self.mapping[pos]) - 1
+
+    def get_partition_value(self, pos: int) -> Any:
+        real_index = self.get_real_index(pos)
+        return self.partition_values[real_index] if real_index < len(self.partition_values) else None
+
+
+class MappedColumnarRowIterator(ColumnarRowIterator):
+    """
+    ColumnarRowIterator with mapping support for partition and index mapping.
+    """
+
+    def __init__(self, file_path: str, record_batch: pa.RecordBatch,
+                 partition_info: Optional[PartitionInfo] = None,
+                 index_mapping: Optional[List[int]] = None):
+        mapped_batch = self._apply_mappings(record_batch, partition_info, index_mapping)
+        super().__init__(file_path, mapped_batch)
+
+    def _apply_mappings(self, record_batch: pa.RecordBatch,
+                        partition_info: Optional[PartitionInfo],
+                        index_mapping: Optional[List[int]]) -> pa.RecordBatch:
+        arrays = []
+        names = []
+
+        if partition_info is not None:
+            for i in range(partition_info.size()):
+                if partition_info.in_partition_row(i):
+                    partition_value = partition_info.get_partition_value(i)
+                    const_array = pa.array([partition_value] * record_batch.num_rows)
+                    arrays.append(const_array)
+                    names.append(f"partition_field_{i}")
+                else:
+                    real_index = partition_info.get_real_index(i)
+                    if real_index < record_batch.num_columns:
+                        arrays.append(record_batch.column(real_index))
+                        names.append(record_batch.column_names[real_index])
+        else:
+            arrays = [record_batch.column(i) for i in range(record_batch.num_columns)]
+            names = record_batch.column_names[:]
+
+        if index_mapping is not None:
+            mapped_arrays = []
+            mapped_names = []
+            for i, real_index in enumerate(index_mapping):
+                if real_index >= 0 and real_index < len(arrays):
+                    mapped_arrays.append(arrays[real_index])
+                    mapped_names.append(names[real_index] if real_index < len(names) else f"field_{i}")
+                else:
+                    null_array = pa.array([None] * record_batch.num_rows)
+                    mapped_arrays.append(null_array)
+                    mapped_names.append(f"null_field_{i}")
+            arrays = mapped_arrays
+            names = mapped_names
+
+        final_batch = pa.RecordBatch.from_arrays(arrays, names=names)
+        return final_batch
 
 
 class DataFileRecordReader(FileRecordReader[InternalRow]):
     """
     Reads InternalRow from data files.
     """
 
-    def __init__(self, wrapped_reader: RecordReader):
+    def __init__(self, wrapped_reader: RecordReader,
+                 index_mapping: Optional[List[int]] = None,
+                 partition_info: Optional[PartitionInfo] = None):
         self.wrapped_reader = wrapped_reader
+        self.index_mapping = index_mapping
+        self.partition_info = partition_info
 
     def read_batch(self) -> Optional[FileRecordIterator['InternalRow']]:
         iterator = self.wrapped_reader.read_batch()
         if iterator is None:
             return None
 
-        # TODO: Handle partition_info, index_mapping, and cast_mapping
+        if isinstance(iterator, ColumnarRowIterator):
+            if self.partition_info is not None or self.index_mapping is not None:
+                iterator = MappedColumnarRowIterator(
+                    iterator.file_path,
+                    iterator.record_batch,
+                    self.partition_info,
+                    self.index_mapping
+                )
+        else:
+            raise PyNativeNotImplementedError("partition_info & index_mapping for non ColumnarRowIterator")
 
         return iterator
 

pypaimon/pynative/reader/pyarrow_dataset_reader.py

@@ -35,16 +35,9 @@ class PyArrowDatasetReader(FileRecordReader[InternalRow]):
     """
 
     def __init__(self, format, file_path, batch_size, projection,
-                 predicate: Predicate, primary_keys: List[str]):
+                 predicate: Predicate, primary_keys: List[str], fields: List[str]):
+
         if primary_keys is not None:
-            if projection is not None:
-                key_columns = []
-                for pk in primary_keys:
-                    key_column = f"_KEY_{pk}"
-                    if key_column not in projection:
-                        key_columns.append(key_column)
-                system_columns = ["_SEQUENCE_NUMBER", "_VALUE_KIND"]
-                projection = key_columns + system_columns + projection
             # TODO: utilize predicate to improve performance
             predicate = None
 
@@ -54,7 +47,7 @@ def __init__(self, format, file_path, batch_size, projection,
         self._file_path = file_path
         self.dataset = ds.dataset(file_path, format=format)
         self.scanner = self.dataset.scanner(
-            columns=projection,
+            columns=fields,
             filter=predicate,
             batch_size=batch_size
         )

pypaimon/pynative/reader/sort_merge_reader.py

@@ -196,11 +196,18 @@ def release_batch(self):
 
 
 class SortMergeReader:
-    def __init__(self, readers, primary_keys):
+    def __init__(self, readers, primary_keys, partition_keys):
         self.next_batch_readers = list(readers)
         self.merge_function = DeduplicateMergeFunction(False)
 
-        key_columns = [f"_KEY_{pk}" for pk in primary_keys]
+        if partition_keys:
+            trimmed_primary_keys = [pk for pk in primary_keys if pk not in partition_keys]
+            if not trimmed_primary_keys:
+                raise ValueError(f"Primary key constraint {primary_keys} same with partition fields")
+        else:
+            trimmed_primary_keys = primary_keys
+
+        key_columns = [f"_KEY_{pk}" for pk in trimmed_primary_keys]
         key_schema = pa.schema([pa.field(column, pa.string()) for column in key_columns])
         self.user_key_comparator = built_comparator(key_schema)
 

pypaimon/pynative/tests/test_pynative_reader.py

@@ -38,6 +38,12 @@ def setUpClass(cls):
             ('f1', pa.string()),
             ('f2', pa.string())
         ])
+        cls.partition_pk_pa_schema = pa.schema([
+            ('user_id', pa.int32(), False),
+            ('item_id', pa.int32()),
+            ('behavior', pa.string()),
+            ('dt', pa.string(), False)
+        ])
         cls._expected_full_data = pd.DataFrame({
             'f0': [1, 2, 3, 4, 5, 6, 7, 8],
             'f1': ['a', 'b', 'c', None, 'e', 'f', 'g', 'h'],
@@ -201,7 +207,7 @@ def testPkParquetReaderWithMinHeap(self):
         actual = self._read_test_table(read_builder)
         self.assertEqual(actual, self.expected_full_pk)
 
-    def testPkOrcReader(self):
+    def skip_testPkOrcReader(self):
         schema = Schema(self.pk_pa_schema, primary_keys=['f0'], options={
             'bucket': '1',
             'file.format': 'orc'
@@ -214,7 +220,7 @@ def testPkOrcReader(self):
         actual = self._read_test_table(read_builder)
         self.assertEqual(actual, self.expected_full_pk)
 
-    def testPkAvroReader(self):
+    def skip_testPkAvroReader(self):
         schema = Schema(self.pk_pa_schema, primary_keys=['f0'], options={
             'bucket': '1',
             'file.format': 'avro'
@@ -263,6 +269,51 @@ def testPkReaderWithProjection(self):
         expected = self.expected_full_pk.select(['f0', 'f2'])
         self.assertEqual(actual, expected)
 
+    def testPartitionPkParquetReader(self):
+        schema = Schema(self.partition_pk_pa_schema,
+                        partition_keys=['dt'],
+                        primary_keys=['dt', 'user_id'],
+                        options={
+                            'bucket': '2'
+                        })
+        self.catalog.create_table('default.test_partition_pk_parquet', schema, False)
+        table = self.catalog.get_table('default.test_partition_pk_parquet')
+        self._write_partition_test_table(table)
+
+        read_builder = table.new_read_builder()
+        actual = self._read_test_table(read_builder)
+        expected = pa.Table.from_pandas(
+            pd.DataFrame({
+                'user_id': [1, 2, 3, 4, 5, 7, 8],
+                'item_id': [1, 2, 3, 4, 5, 7, 8],
+                'behavior': ["b-1", "b-2-new", "b-3", None, "b-5", "b-7", None],
+                'dt': ["p-1", "p-1", "p-1", "p-1", "p-2", "p-1", "p-2"]
+            }),
+            schema=self.partition_pk_pa_schema)
+        self.assertEqual(actual.sort_by('user_id'), expected)
+
+    def testPartitionPkParquetReaderWriteOnce(self):
+        schema = Schema(self.partition_pk_pa_schema,
+                        partition_keys=['dt'],
+                        primary_keys=['dt', 'user_id'],
+                        options={
+                            'bucket': '1'
+                        })
+        self.catalog.create_table('default.test_partition_pk_parquet2', schema, False)
+        table = self.catalog.get_table('default.test_partition_pk_parquet2')
+        self._write_partition_test_table(table, write_once=True)
+
+        read_builder = table.new_read_builder()
+        actual = self._read_test_table(read_builder)
+        expected = pa.Table.from_pandas(
+            pd.DataFrame({
+                'user_id': [1, 2, 3, 4],
+                'item_id': [1, 2, 3, 4],
+                'behavior': ['b-1', 'b-2', 'b-3', None],
+                'dt': ['p-1', 'p-1', 'p-1', 'p-1']
+            }), schema=self.partition_pk_pa_schema)
+        self.assertEqual(actual, expected)
+
     def _write_test_table(self, table, for_pk=False):
         write_builder = table.new_batch_write_builder()
 
@@ -301,6 +352,40 @@ def _write_test_table(self, table, for_pk=False):
         table_write.close()
         table_commit.close()
 
+    def _write_partition_test_table(self, table, write_once=False):
+        write_builder = table.new_batch_write_builder()
+
+        table_write = write_builder.new_write()
+        table_commit = write_builder.new_commit()
+        data1 = {
+            'user_id': [1, 2, 3, 4],
+            'item_id': [1, 2, 3, 4],
+            'behavior': ['b-1', 'b-2', 'b-3', None],
+            'dt': ['p-1', 'p-1', 'p-1', 'p-1']
+        }
+        pa_table = pa.Table.from_pydict(data1, schema=self.partition_pk_pa_schema)
+        table_write.write_arrow(pa_table)
+        table_commit.commit(table_write.prepare_commit())
+        table_write.close()
+        table_commit.close()
+
+        if write_once:
+            return
+
+        table_write = write_builder.new_write()
+        table_commit = write_builder.new_commit()
+        data1 = {
+            'user_id': [5, 2, 7, 8],
+            'item_id': [5, 2, 7, 8],
+            'behavior': ['b-5', 'b-2-new', 'b-7', None],
+            'dt': ['p-2', 'p-1', 'p-1', 'p-2']
+        }
+        pa_table = pa.Table.from_pydict(data1, schema=self.partition_pk_pa_schema)
+        table_write.write_arrow(pa_table)
+        table_commit.commit(table_write.prepare_commit())
+        table_write.close()
+        table_commit.close()
+
     def _read_test_table(self, read_builder):
         table_read = read_builder.new_read()
         splits = read_builder.new_scan().plan().splits()