This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 296b2e8ff8 [python] Make write_pandas respect write_cols and add
to_duckdb parallelism (#8636)
296b2e8ff8 is described below
commit 296b2e8ff8f94d9397412733f32996ebcdcc8fca
Author: chaoyang <[email protected]>
AuthorDate: Thu Jul 16 22:11:28 2026 +0800
[python] Make write_pandas respect write_cols and add to_duckdb parallelism
(#8636)
---
paimon-python/pypaimon/read/table_read.py | 9 ++-
paimon-python/pypaimon/tests/reader_base_test.py | 38 +++++++++++++
.../pypaimon/tests/write/table_write_test.py | 65 ++++++++++++++++++++++
paimon-python/pypaimon/write/table_write.py | 9 ++-
4 files changed, 118 insertions(+), 3 deletions(-)
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 89ca539b73..6e641b5096 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -538,11 +538,16 @@ class TableRead:
return arrow_table.to_pandas()
def to_duckdb(self, splits: List[Split], table_name: str,
- connection: Optional["DuckDBPyConnection"] = None) ->
"DuckDBPyConnection":
+ connection: Optional["DuckDBPyConnection"] = None,
+ parallelism: Optional[int] = None) -> "DuckDBPyConnection":
+ """Materialize ``splits`` into an in-memory table registered with
DuckDB.
+
+ See :meth:`to_arrow` for the semantics of ``parallelism``.
+ """
import duckdb
con = connection or duckdb.connect(database=":memory:")
- con.register(table_name, self.to_arrow(splits))
+ con.register(table_name, self.to_arrow(splits,
parallelism=parallelism))
return con
def to_ray(
diff --git a/paimon-python/pypaimon/tests/reader_base_test.py
b/paimon-python/pypaimon/tests/reader_base_test.py
index 12875cabaf..455b7765b8 100644
--- a/paimon-python/pypaimon/tests/reader_base_test.py
+++ b/paimon-python/pypaimon/tests/reader_base_test.py
@@ -364,6 +364,44 @@ class ReaderBasicTest(unittest.TestCase):
expect = pd.DataFrame(self.raw_data)
pd.testing.assert_frame_equal(actual.reset_index(drop=True),
expect.reset_index(drop=True))
+ def test_reader_duckDB_parallelism(self):
+ # A dedicated partitioned table with rows across multiple partitions so
+ # scan planning yields >= 2 splits, exercising the real parallel
fan-out
+ # path (``_should_run_parallel`` requires parallelism >= 2 AND
+ # splits >= 2).
+ schema = Schema.from_pyarrow_schema(self.pa_schema,
partition_keys=['dt'])
+ self.catalog.create_table('default.test_duckdb_parallel', schema,
False)
+ table = self.catalog.get_table('default.test_duckdb_parallel')
+ write_builder = table.new_batch_write_builder()
+ table_write = write_builder.new_write()
+ table_commit = write_builder.new_commit()
+ table_write.write_arrow(pa.Table.from_pydict({
+ 'user_id': [1, 2, 3, 4, 5, 6],
+ 'item_id': [1001, 1002, 1003, 1004, 1005, 1006],
+ 'behavior': ['a', 'b', 'c', 'd', 'e', 'f'],
+ 'dt': ['p1', 'p1', 'p2', 'p2', 'p3', 'p3'],
+ }, schema=self.pa_schema))
+ table_commit.commit(table_write.prepare_commit())
+ table_write.close()
+ table_commit.close()
+
+ read_builder = table.new_read_builder()
+ splits = read_builder.new_scan().plan().splits()
+ self.assertGreaterEqual(len(splits), 2)
+
+ serial = read_builder.new_read().to_duckdb(
+ splits, 'duckdb_serial', parallelism=1)
+ parallel = read_builder.new_read().to_duckdb(
+ splits, 'duckdb_parallel', parallelism=4)
+
+ serial_df = serial.query(
+ "SELECT * FROM duckdb_serial ORDER BY item_id").fetchdf()
+ parallel_df = parallel.query(
+ "SELECT * FROM duckdb_parallel ORDER BY item_id").fetchdf()
+ pd.testing.assert_frame_equal(
+ serial_df.reset_index(drop=True),
+ parallel_df.reset_index(drop=True))
+
def test_mixed_add_and_delete_entries_compute_stats(self):
"""Test record_count calculation with mixed ADD/DELETE entries in same
partition."""
pa_schema = pa.schema([
diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py
b/paimon-python/pypaimon/tests/write/table_write_test.py
index f6eb97c5ee..ee2f75bec4 100644
--- a/paimon-python/pypaimon/tests/write/table_write_test.py
+++ b/paimon-python/pypaimon/tests/write/table_write_test.py
@@ -587,6 +587,71 @@ class TableWriteTest(unittest.TestCase):
self.assertTrue(str(e.exception).startswith(
"Input schema isn't consistent with table schema and write cols."))
+ def test_write_pandas_respects_write_cols(self):
+ import pandas as pd
+
+ pa_schema = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ('score', pa.int64()),
+ ])
+ schema = Schema.from_pyarrow_schema(pa_schema)
+ self.catalog.create_table(
+ 'default.test_write_pandas_write_cols', schema, False)
+ table = self.catalog.get_table(
+ 'default.test_write_pandas_write_cols')
+
+ write_builder = table.new_batch_write_builder()
+ table_write = write_builder.new_write().with_write_type(['id', 'name'])
+ table_commit = write_builder.new_commit()
+ # DataFrame only carries the written subset; missing ``score`` is
+ # padded with null on read.
+ table_write.write_pandas(pd.DataFrame({
+ 'id': [1, 2],
+ 'name': ['a', 'b'],
+ }))
+ table_commit.commit(table_write.prepare_commit())
+ table_write.close()
+ table_commit.close()
+
+ expected = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['a', 'b'],
+ 'score': [None, None],
+ }, schema=pa_schema)
+ actual = self._read_sorted(table, 'id')
+ self.assertEqual(expected, actual)
+
+ def test_write_pandas_full_columns_unchanged(self):
+ import pandas as pd
+
+ pa_schema = pa.schema([
+ ('id', pa.int32()),
+ ('name', pa.string()),
+ ])
+ schema = Schema.from_pyarrow_schema(pa_schema)
+ self.catalog.create_table(
+ 'default.test_write_pandas_full', schema, False)
+ table = self.catalog.get_table('default.test_write_pandas_full')
+
+ write_builder = table.new_batch_write_builder()
+ table_write = write_builder.new_write()
+ table_commit = write_builder.new_commit()
+ table_write.write_pandas(pd.DataFrame({
+ 'id': [1, 2],
+ 'name': ['a', 'b'],
+ }))
+ table_commit.commit(table_write.prepare_commit())
+ table_write.close()
+ table_commit.close()
+
+ expected = pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'name': ['a', 'b'],
+ }, schema=pa_schema)
+ actual = self._read_sorted(table, 'id')
+ self.assertEqual(expected, actual)
+
def test_validate_schema_allows_binary_family_for_write_cols(self):
pa_schema = pa.schema([
('id', pa.int32()),
diff --git a/paimon-python/pypaimon/write/table_write.py
b/paimon-python/pypaimon/write/table_write.py
index 629383409f..ad88295b64 100644
--- a/paimon-python/pypaimon/write/table_write.py
+++ b/paimon-python/pypaimon/write/table_write.py
@@ -75,7 +75,14 @@ class TableWrite:
self.file_store_write.write_row(partition, bucket, row, values_by_name)
def write_pandas(self, dataframe):
- pa_schema =
PyarrowFieldParser.from_paimon_schema(self.table.table_schema.fields)
+ write_cols = self.file_store_write.write_cols
+ if write_cols is not None:
+ # Column-subset write (append-only ``with_write_type``): build the
+ # RecordBatch against the subset schema so the input only needs the
+ # written columns, mirroring the ``write_arrow`` path.
+ pa_schema = self._write_cols_pyarrow_schema(write_cols)
+ else:
+ pa_schema = self.table_pyarrow_schema
record_batch = pa.RecordBatch.from_pandas(dataframe, schema=pa_schema)
return self.write_arrow_batch(record_batch)
