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 450a993e02 [python][ray] Fix HASH_FIXED primary-key writer ownership
(#8541)
450a993e02 is described below
commit 450a993e02b36c40b9e40b7e6ad2b692d12ebe7a
Author: QuakeWang <[email protected]>
AuthorDate: Mon Jul 13 10:36:19 2026 +0800
[python][ray] Fix HASH_FIXED primary-key writer ownership (#8541)
Ray `map_groups` only keeps a `(partition, bucket)` group intact while
the UDF runs. Its output may be split before `PaimonDatasink`, so sink
options that prevent operator fusion can create multiple independent
writers for one bucket with overlapping sequence numbers.
This change writes each HASH_FIXED primary-key group inside the
`map_groups` UDF, returns serialized commit messages, and commits them
on the driver. Writer concurrency and remote options are applied
directly to the group tasks.
---
docs/docs/pypaimon/ray-data.md | 34 +++---
paimon-python/pypaimon/ray/ray_paimon.py | 38 +++----
paimon-python/pypaimon/ray/shuffle.py | 20 ++--
.../pypaimon/tests/ray_repartition_test.py | 101 +++++++++++++++++
paimon-python/pypaimon/write/ray_datasink.py | 120 +++++++++++++++++++++
paimon-python/pypaimon/write/table_write.py | 17 ++-
6 files changed, 278 insertions(+), 52 deletions(-)
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index a9bd84d38d..88c0a7143a 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -211,7 +211,9 @@ write_paimon(
**HASH_FIXED pre-clustering:**
HASH_FIXED rows are always assigned to the correct Paimon bucket by
-the writer. Pre-clustering is only a file-count optimization.
+the writer. For append-only tables, pre-clustering is only a file-count
+optimization. Primary-key tables additionally require one writer per
+`(partition_keys..., bucket)` group to generate ordered sequence numbers.
By default, `write_paimon` writes append-only HASH_FIXED tables
without pre-clustering. This avoids Ray `groupby().map_groups()`
@@ -220,9 +222,9 @@ node.
HASH_FIXED primary-key tables reject the default/off mode. Direct Ray
writes can send the same bucket to multiple writer tasks, and those
-writers can allocate overlapping sequence numbers. Use the explicit
-`map_groups` mode until a bounded pre-clustering strategy preserves
-per-bucket sequence ordering.
+writers can allocate overlapping sequence numbers. The explicit
+`map_groups` mode avoids this by running one writer for each complete
+`(partition_keys..., bucket)` group.
If every `(partition_keys..., bucket)` group fits in memory on a
single Ray node, you can opt in to the legacy small-file optimization:
@@ -237,11 +239,12 @@ write_paimon(
```
`hash_fixed_precluster="map_groups"` groups rows by
-`(partition_keys..., bucket)` before writing so each group lands in a
-single Ray task. This can reduce file count and keeps HASH_FIXED
-primary-key sequence generation per bucket in one writer task, but it
-inherits Ray's `map_groups()` memory bound. Large append-only buckets
-or hot append-only partitions should use the default mode or
+`(partition_keys..., bucket)`. For primary-key tables, the Paimon writer
+runs inside that `map_groups()` task and returns serialized commit
+messages for the driver to commit. Ray output block splitting therefore
+cannot create multiple writers for the same group. The mode inherits
+Ray's `map_groups()` memory bound. Large append-only buckets or hot
+append-only partitions should use the default mode or
`hash_fixed_precluster="off"`.
For non-HASH_FIXED append-only tables, the dataset is written as-is.
@@ -258,15 +261,18 @@ overlapping buckets or sequence numbers for those modes.
- `catalog_options`: kwargs forwarded to `CatalogFactory.create()`.
- `overwrite`: if `True`, overwrite existing data in the table.
- `concurrency`: optional max number of Ray write tasks to run concurrently.
+ For HASH_FIXED primary-key `map_groups` writes, this limits the group
+ writer tasks.
- `ray_remote_args`: optional kwargs passed to `ray.remote()` in write tasks
- (e.g. `{"num_cpus": 2}`).
+ (e.g. `{"num_cpus": 2}`). For HASH_FIXED primary-key `map_groups`
+ writes, these options apply to the group writer tasks.
- `hash_fixed_precluster`: HASH_FIXED pre-clustering mode. `"auto"` and
`"off"` write append-only HASH_FIXED tables directly and reject
HASH_FIXED primary-key tables. `"map_groups"` enables the legacy
- small-file optimization for HASH_FIXED primary-key tables and requires
- each `(partition_keys..., bucket)` group to fit in memory on one Ray
- node. This option does not enable Ray writes for HASH_DYNAMIC or
- CROSS_PARTITION primary-key tables.
+ small-file optimization for append-only tables and runs one writer per
+ HASH_FIXED primary-key group. Each `(partition_keys..., bucket)` group
+ must fit in memory on one Ray node. This option does not enable Ray
+ writes for HASH_DYNAMIC or CROSS_PARTITION primary-key tables.
### `TableWrite.write_ray()` (lower-level)
diff --git a/paimon-python/pypaimon/ray/ray_paimon.py
b/paimon-python/pypaimon/ray/ray_paimon.py
index b375e3950f..3dcf40616d 100644
--- a/paimon-python/pypaimon/ray/ray_paimon.py
+++ b/paimon-python/pypaimon/ray/ray_paimon.py
@@ -292,12 +292,12 @@ def write_paimon(
"""Write a Ray Dataset to a Paimon table.
HASH_FIXED rows are assigned to the correct bucket by the Paimon
- writer. Optional pre-clustering is only a file-count optimization.
- The legacy ``map_groups`` pre-clustering mode materializes each
- ``(partition_keys..., bucket)`` group on one Ray node and should
- only be used when every group fits in memory. HASH_DYNAMIC and
- CROSS_PARTITION primary-key Ray writes are rejected because Ray
- write tasks create independent Paimon writers.
+ writer. For primary-key tables, ``map_groups`` writes each complete
+ ``(partition_keys..., bucket)`` group in one Ray task and returns
+ commit messages to the driver. It should only be used when every
+ group fits in memory. HASH_DYNAMIC and CROSS_PARTITION primary-key
+ Ray writes are rejected because Ray write tasks create independent
+ Paimon writers.
Args:
dataset: The Ray Dataset to write.
@@ -309,26 +309,22 @@ def write_paimon(
hash_fixed_precluster: HASH_FIXED pre-clustering mode. ``"auto"``
and ``"off"`` write append-only HASH_FIXED tables directly
and reject HASH_FIXED primary-key tables. ``"map_groups"``
- preserves the legacy small-file optimization and its single
- group memory bound for HASH_FIXED primary-key tables.
+ writes each HASH_FIXED primary-key group in one task and
+ preserves the legacy single-group memory bound.
"""
_require_ray_data()
from pypaimon.catalog.catalog_factory import CatalogFactory
- from pypaimon.ray.shuffle import maybe_apply_repartition
- from pypaimon.write.ray_datasink import PaimonDatasink
+ from pypaimon.write.ray_datasink import write_paimon_dataset
catalog = CatalogFactory.create(catalog_options)
table = catalog.get_table(table_identifier)
- dataset = maybe_apply_repartition(dataset, table, hash_fixed_precluster)
-
- datasink = PaimonDatasink(table, overwrite=overwrite)
-
- write_kwargs = {}
- if ray_remote_args is not None:
- write_kwargs["ray_remote_args"] = ray_remote_args
- if concurrency is not None:
- write_kwargs["concurrency"] = concurrency
-
- dataset.write_datasink(datasink, **write_kwargs)
+ write_paimon_dataset(
+ dataset,
+ table,
+ overwrite=overwrite,
+ concurrency=concurrency,
+ ray_remote_args=ray_remote_args,
+ hash_fixed_precluster=hash_fixed_precluster,
+ )
diff --git a/paimon-python/pypaimon/ray/shuffle.py
b/paimon-python/pypaimon/ray/shuffle.py
index 5079bf6215..eef29bdfc5 100644
--- a/paimon-python/pypaimon/ray/shuffle.py
+++ b/paimon-python/pypaimon/ray/shuffle.py
@@ -20,9 +20,10 @@
The legacy ``map_groups`` strategy groups rows by
``(partition_keys..., bucket)`` so every distinct group lands in a
-single Ray task. This can reduce file count, but Ray requires each
-``map_groups`` group to fit in memory on one node. Keep that strategy
-behind an explicit opt-in.
+single Ray task. Primary-key writes consume the complete group in that
+task; append-only writes use the regrouped rows as a file-count
+optimization. Ray requires each ``map_groups`` group to fit in memory
+on one node, so keep that strategy behind an explicit opt-in.
For append-only tables in any other bucket mode the dataset is returned
unchanged.
@@ -117,6 +118,15 @@ def maybe_apply_repartition(
)
return dataset
+ grouped, bucket_col = _group_by_partition_bucket(dataset, table)
+ regrouped = grouped.map_groups(_identity_batch, batch_format="pyarrow")
+ return regrouped.drop_columns([bucket_col])
+
+
+def _group_by_partition_bucket(
+ dataset: "ray.data.Dataset",
+ table: "Table",
+):
partition_keys = list(table.table_schema.partition_keys or [])
extractor = table.create_row_key_extractor()
col_names = set(f.name for f in table.table_schema.fields)
@@ -127,9 +137,7 @@ def maybe_apply_repartition(
bucket_udf, batch_format="pyarrow", zero_copy_batch=True,
)
group_keys: List[str] = partition_keys + [bucket_col]
- grouped = ds_with_bucket.groupby(group_keys)
- regrouped = grouped.map_groups(_identity_batch, batch_format="pyarrow")
- return regrouped.drop_columns([bucket_col])
+ return ds_with_bucket.groupby(group_keys), bucket_col
def _identity_batch(batch: pa.Table) -> pa.Table:
diff --git a/paimon-python/pypaimon/tests/ray_repartition_test.py
b/paimon-python/pypaimon/tests/ray_repartition_test.py
index 0d7dd568c0..8229d9467c 100644
--- a/paimon-python/pypaimon/tests/ray_repartition_test.py
+++ b/paimon-python/pypaimon/tests/ray_repartition_test.py
@@ -30,6 +30,9 @@ explicitly selected. These tests cover:
from the sink-visible schema.
* explicit ``map_groups`` mode can produce one file per
(partition, bucket) on the small test dataset.
+ * primary-key group writers remain single-writer when Ray splits the
+ UDF output or writer task options prevent downstream operator fusion.
+ * grouped primary-key writes preserve overwrite and empty-overwrite behavior.
* regression: a table whose schema already contains a column named
``__paimon_bucket__`` still works (collision-safe column name).
* non-HASH_FIXED append-only tables pass through unchanged.
@@ -364,6 +367,104 @@ class RayShuffleTest(unittest.TestCase):
# 4 buckets × 1 file each.
self.assertEqual(len(files), 4)
+ def test_primary_key_group_writer_does_not_depend_on_fusion(self):
+ """Output block splitting must not create independent bucket
writers."""
+ from ray.data import DataContext
+
+ from pypaimon.ray import write_paimon
+
+ pa_schema = pa.schema([
+ pa.field('id', pa.int32(), nullable=False),
+ ('value', pa.string()),
+ ])
+ table_name = 'test_pk_group_writer_no_fusion'
+ identifier = self._make_table(
+ table_name, pa_schema,
+ primary_keys=['id'], options={'bucket': '1'},
+ )
+ rows = pa.Table.from_pydict({
+ 'id': [1] * 800,
+ 'value': [f'{i:04d}-' + 'x' * 4091 for i in range(800)],
+ }, schema=pa_schema)
+
+ context = DataContext.get_current()
+ previous_target = context.target_max_block_size
+ context.target_max_block_size = 64 * 1024
+ try:
+ write_paimon(
+ ray.data.from_arrow(rows),
+ identifier,
+ self.catalog_options,
+ concurrency=4,
+ ray_remote_args={'num_cpus': 0.5},
+ hash_fixed_precluster='map_groups',
+ )
+ finally:
+ context.target_max_block_size = previous_target
+
+ table =
CatalogFactory.create(self.catalog_options).get_table(identifier)
+ data_files = [
+ data_file
+ for split in table.new_read_builder().new_scan().plan().splits()
+ for data_file in split.files
+ ]
+ self.assertEqual(len(data_files), 1)
+ self.assertEqual(
+ (data_files[0].min_sequence_number,
+ data_files[0].max_sequence_number),
+ (1, 801),
+ )
+
+ result = self._read_table(identifier)
+ self.assertEqual(len(result), 1)
+ self.assertTrue(result.iloc[0]['value'].startswith('0799-'))
+
+ def test_primary_key_group_writer_overwrite(self):
+ from pypaimon.ray import write_paimon
+
+ pa_schema = pa.schema([
+ pa.field('id', pa.int32(), nullable=False),
+ ('value', pa.string()),
+ ])
+ identifier = self._make_table(
+ 'test_pk_group_writer_overwrite', pa_schema,
+ primary_keys=['id'], options={'bucket': '1'},
+ )
+
+ initial = pa.Table.from_pydict(
+ {'id': [1, 2], 'value': ['a', 'b']}, schema=pa_schema,
+ )
+ write_paimon(
+ ray.data.from_arrow(initial), identifier, self.catalog_options,
+ hash_fixed_precluster='map_groups',
+ )
+
+ replacement = pa.Table.from_pydict(
+ {'id': [3], 'value': ['c']}, schema=pa_schema,
+ )
+ write_paimon(
+ ray.data.from_arrow(replacement),
+ identifier,
+ self.catalog_options,
+ overwrite=True,
+ concurrency=2,
+ ray_remote_args={'num_cpus': 0.5},
+ hash_fixed_precluster='map_groups',
+ )
+ self.assertEqual(set(self._read_table(identifier)['id']), {3})
+
+ empty = pa.Table.from_batches([], schema=pa_schema)
+ write_paimon(
+ ray.data.from_arrow(empty),
+ identifier,
+ self.catalog_options,
+ overwrite=True,
+ concurrency=2,
+ ray_remote_args={'num_cpus': 0.5},
+ hash_fixed_precluster='map_groups',
+ )
+ self.assertEqual(len(self._read_table(identifier)), 0)
+
def test_fixed_bucket_with_colliding_column_name(self):
"""A table that has a column named ``__paimon_bucket__`` must
still work — the helper picks a collision-free transient
diff --git a/paimon-python/pypaimon/write/ray_datasink.py
b/paimon-python/pypaimon/write/ray_datasink.py
index 18b7f024f0..6f7a336190 100644
--- a/paimon-python/pypaimon/write/ray_datasink.py
+++ b/paimon-python/pypaimon/write/ray_datasink.py
@@ -241,3 +241,123 @@ class PaimonDatasink(_DatasinkBase):
)
finally:
self._pending_commit_messages = []
+
+
+def write_paimon_dataset(
+ dataset,
+ table,
+ *,
+ overwrite: bool = False,
+ static_partition: Optional[Dict[str, Any]] = None,
+ concurrency: Optional[int] = None,
+ ray_remote_args: Optional[Dict[str, Any]] = None,
+ hash_fixed_precluster: str = "auto",
+) -> None:
+ """Write a Ray Dataset through the safe path for the table's bucket
mode."""
+ from pypaimon.ray.shuffle import (
+ HASH_FIXED_PRECLUSTER_MAP_GROUPS,
+ maybe_apply_repartition,
+ )
+ from pypaimon.table.bucket_mode import BucketMode
+
+ if (
+ hash_fixed_precluster == HASH_FIXED_PRECLUSTER_MAP_GROUPS
+ and table.bucket_mode() == BucketMode.HASH_FIXED
+ and getattr(table, "is_primary_key_table", False)
+ ):
+ _write_primary_key_groups(
+ dataset,
+ table,
+ overwrite=overwrite,
+ static_partition=static_partition,
+ concurrency=concurrency,
+ ray_remote_args=ray_remote_args,
+ )
+ return
+
+ dataset = maybe_apply_repartition(dataset, table, hash_fixed_precluster)
+ dataset.write_datasink(
+ PaimonDatasink(
+ table,
+ overwrite=overwrite,
+ static_partition=static_partition,
+ ),
+ concurrency=concurrency,
+ ray_remote_args=ray_remote_args,
+ )
+
+
+def _write_primary_key_groups(
+ dataset,
+ table,
+ *,
+ overwrite: bool,
+ static_partition: Optional[Dict[str, Any]],
+ concurrency: Optional[int],
+ ray_remote_args: Optional[Dict[str, Any]],
+) -> None:
+ import inspect
+ import pickle
+
+ from pypaimon.ray.shuffle import (
+ _coerce_large_string_types,
+ _group_by_partition_bucket,
+ )
+
+ grouped, bucket_col = _group_by_partition_bucket(dataset, table)
+ message_col = "__paimon_commit_messages__"
+ captured_table = table
+
+ # Keep the writer inside the group UDF. Ray may split the UDF output
+ # into multiple blocks, so only serialized commit messages leave it.
+ def _write_group(group: pa.Table) -> pa.Table:
+ if group.num_rows == 0:
+ return pa.table({message_col: pa.array([], type=pa.binary())})
+
+ rows = _coerce_large_string_types(
+ group.drop_columns([bucket_col])
+ )
+ worker_sink = PaimonDatasink(
+ captured_table,
+ overwrite=overwrite,
+ static_partition=static_partition,
+ )
+ commit_messages = worker_sink.write([rows], None)
+ return pa.table({
+ message_col: pa.array(
+ [pickle.dumps(commit_messages)], type=pa.binary()
+ )
+ })
+
+ map_kwargs = {"batch_format": "pyarrow"}
+ if concurrency is not None:
+ concurrency_param = inspect.signature(
+ grouped.map_groups
+ ).parameters.get("concurrency")
+ if (
+ concurrency_param is not None
+ and concurrency_param.kind != inspect.Parameter.VAR_KEYWORD
+ ):
+ map_kwargs["concurrency"] = concurrency
+ else:
+ from ray.data._internal.compute import TaskPoolStrategy
+ map_kwargs["compute"] = TaskPoolStrategy(size=concurrency)
+ if ray_remote_args:
+ map_kwargs.update(ray_remote_args)
+
+ messages = grouped.map_groups(_write_group, **map_kwargs)
+ coordinator = PaimonDatasink(
+ table,
+ overwrite=overwrite,
+ static_partition=static_partition,
+ )
+ coordinator.on_write_start()
+ try:
+ write_returns = []
+ for batch in messages.iter_batches(batch_format="pyarrow"):
+ for blob in batch.column(message_col).to_pylist():
+ write_returns.append(pickle.loads(blob))
+ coordinator.on_write_complete(write_returns)
+ except Exception as error:
+ coordinator.on_write_failed(error)
+ raise
diff --git a/paimon-python/pypaimon/write/table_write.py
b/paimon-python/pypaimon/write/table_write.py
index 4fb4b8043c..629383409f 100644
--- a/paimon-python/pypaimon/write/table_write.py
+++ b/paimon-python/pypaimon/write/table_write.py
@@ -120,31 +120,26 @@ class TableWrite:
hash_fixed_precluster: HASH_FIXED pre-clustering mode. ``"auto"``
and ``"off"`` write append-only HASH_FIXED tables directly
and reject HASH_FIXED primary-key tables. ``"map_groups"``
- preserves the legacy small-file optimization and its single
- group memory bound for HASH_FIXED primary-key tables.
+ writes each HASH_FIXED primary-key group in one task and
+ preserves the legacy single-group memory bound.
static_partition: Optional partition spec to overwrite. When set,
the Ray write runs in overwrite mode for this partition and
overrides any builder-level partition spec.
"""
- from pypaimon.ray.shuffle import maybe_apply_repartition
- from pypaimon.write.ray_datasink import PaimonDatasink
-
- dataset = maybe_apply_repartition(
- dataset, self.table, hash_fixed_precluster)
+ from pypaimon.write.ray_datasink import write_paimon_dataset
overwrite_partition = self.static_partition
if static_partition is not None:
overwrite_partition = static_partition
- datasink = PaimonDatasink(
+ write_paimon_dataset(
+ dataset,
self.table,
overwrite=overwrite,
static_partition=overwrite_partition,
- )
- dataset.write_datasink(
- datasink,
concurrency=concurrency,
ray_remote_args=ray_remote_args,
+ hash_fixed_precluster=hash_fixed_precluster,
)
def close(self):
