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 14cf97d3f5 [python] Make HASH_FIXED Ray pre-clustering opt-in (#8046)
14cf97d3f5 is described below
commit 14cf97d3f52c83b589fd2192c78857b2e14d83fb
Author: QuakeWang <[email protected]>
AuthorDate: Sun May 31 21:24:39 2026 +0800
[python] Make HASH_FIXED Ray pre-clustering opt-in (#8046)
---
docs/docs/pypaimon/ray-data.md | 67 ++++++++---
paimon-python/pypaimon/ray/ray_paimon.py | 20 +++-
paimon-python/pypaimon/ray/shuffle.py | 66 +++++++----
.../pypaimon/tests/ray_repartition_test.py | 131 +++++++++++++++++----
.../pypaimon/tests/test_ray_shuffle_helper.py | 59 +++++++++-
paimon-python/pypaimon/write/table_write.py | 11 ++
6 files changed, 287 insertions(+), 67 deletions(-)
diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md
index 3ee4db3289..b411561322 100644
--- a/docs/docs/pypaimon/ray-data.md
+++ b/docs/docs/pypaimon/ray-data.md
@@ -205,19 +205,43 @@ write_paimon(
)
```
-**Automatic (partition, bucket) clustering for HASH_FIXED tables:**
+**HASH_FIXED pre-clustering:**
-For HASH_FIXED tables, `write_paimon` automatically clusters rows by
-`(partition_keys..., bucket)` before writing so each (partition,
-bucket) lands in a single Ray task — one writer, one file group. This
-avoids the small-file storm that Ray's default round-robin
-distribution would otherwise produce (`partitions × buckets ×
-ray_tasks` files instead of `partitions × buckets`).
+HASH_FIXED rows are always assigned to the correct Paimon bucket by
+the writer. Pre-clustering is only a file-count optimization.
-Bucket assignment uses the same hash routine the writer uses, so the
-bucket seen by the groupby is byte-equivalent to the one the writer
-would compute. No user configuration is required. For non-HASH_FIXED
-tables the dataset is written as-is.
+By default, `write_paimon` writes append-only HASH_FIXED tables
+without pre-clustering. This avoids Ray `groupby().map_groups()`
+materializing an entire `(partition_keys..., bucket)` group on one Ray
+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.
+
+If every `(partition_keys..., bucket)` group fits in memory on a
+single Ray node, you can opt in to the legacy small-file optimization:
+
+```python
+write_paimon(
+ ray_dataset,
+ "database_name.table_name",
+ catalog_options={"warehouse": "/path/to/warehouse"},
+ hash_fixed_precluster="map_groups",
+)
+```
+
+`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
+`hash_fixed_precluster="off"`.
+
+For non-HASH_FIXED tables the dataset is written as-is.
**Parameters:**
- `dataset`: the Ray Dataset to write.
@@ -227,13 +251,20 @@ tables the dataset is written as-is.
- `concurrency`: optional max number of Ray write tasks to run concurrently.
- `ray_remote_args`: optional kwargs passed to `ray.remote()` in write tasks
(e.g. `{"num_cpus": 2}`).
+- `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 and requires each `(partition_keys..., bucket)`
+ group to fit in memory on one Ray node.
### `TableWrite.write_ray()` (lower-level)
If you have already constructed a `table_write` from a write builder, you can
-hand a Ray Dataset directly to it. `write_ray()` commits through the Ray
-Datasink API, so there is no `prepare_commit` / `commit` step to run for the
-Ray write itself — just close the writer when you are done with it:
+hand a Ray Dataset directly to it. `write_ray()` uses the same HASH_FIXED
+pre-clustering modes and safety checks as the top-level `write_paimon()` API.
+It commits through the Ray Datasink API, so there is no `prepare_commit` /
+`commit` step to run for the Ray write itself — just close the writer when you
+are done with it:
```python
import ray
@@ -248,12 +279,18 @@ table_commit = write_builder.new_commit()
# 2. Write Ray Dataset
ray_dataset = ray.data.read_json("/path/to/data.jsonl")
-table_write.write_ray(ray_dataset, overwrite=False, concurrency=2)
+table_write.write_ray(
+ ray_dataset,
+ overwrite=False,
+ concurrency=2,
+ hash_fixed_precluster="auto",
+)
# Parameters:
# - dataset: Ray Dataset to write
# - overwrite: Whether to overwrite existing data (default: False)
# - concurrency: Optional max number of concurrent Ray tasks
# - ray_remote_args: Optional kwargs passed to ray.remote() (e.g.,
{"num_cpus": 2})
+# - hash_fixed_precluster: Same HASH_FIXED modes as write_paimon()
# 3. Commit data (required for write_pandas/write_arrow/write_arrow_batch only)
commit_messages = table_write.prepare_commit()
diff --git a/paimon-python/pypaimon/ray/ray_paimon.py
b/paimon-python/pypaimon/ray/ray_paimon.py
index 86505097d8..e2924dcda6 100644
--- a/paimon-python/pypaimon/ray/ray_paimon.py
+++ b/paimon-python/pypaimon/ray/ray_paimon.py
@@ -114,14 +114,17 @@ def write_paimon(
overwrite: bool = False,
concurrency: Optional[int] = None,
ray_remote_args: Optional[Dict[str, Any]] = None,
+ hash_fixed_precluster: str = "auto",
) -> None:
"""Write a Ray Dataset to a Paimon table.
- For HASH_FIXED tables, rows are automatically clustered by
- ``(partition_keys..., bucket)`` before writing so that each
- (partition, bucket) lands in a single Ray task. This avoids the
- small-file storm that Ray's default round-robin distribution would
- otherwise produce. No user configuration is required.
+ 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_FIXED primary-key
+ tables require ``map_groups`` until Ray writes have a bounded
+ strategy that preserves per-bucket sequence ordering.
Args:
dataset: The Ray Dataset to write.
@@ -130,6 +133,11 @@ def write_paimon(
overwrite: If ``True``, overwrite existing data in the table.
concurrency: Optional max number of Ray write tasks to run
concurrently.
ray_remote_args: Optional kwargs passed to ``ray.remote`` in write
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"``
+ preserves the legacy small-file optimization and its single
+ group memory bound.
"""
from pypaimon.catalog.catalog_factory import CatalogFactory
from pypaimon.ray.shuffle import maybe_apply_repartition
@@ -138,7 +146,7 @@ def write_paimon(
catalog = CatalogFactory.create(catalog_options)
table = catalog.get_table(table_identifier)
- dataset = maybe_apply_repartition(dataset, table)
+ dataset = maybe_apply_repartition(dataset, table, hash_fixed_precluster)
datasink = PaimonDatasink(table, overwrite=overwrite)
diff --git a/paimon-python/pypaimon/ray/shuffle.py
b/paimon-python/pypaimon/ray/shuffle.py
index b17f7a7ab1..c7a1c5b85b 100644
--- a/paimon-python/pypaimon/ray/shuffle.py
+++ b/paimon-python/pypaimon/ray/shuffle.py
@@ -16,21 +16,13 @@
# limitations under the License.
################################################################################
-"""Pre-repartition a Ray Dataset by (partition, bucket) before writing
-to a Paimon table.
-
-Without this, Ray's default round-robin block distribution scatters rows
-that share the same (partition, bucket) across many Ray tasks. Each
-task then opens its own writer and emits its own data file, producing
-``partitions × buckets × ray_tasks`` files instead of the
-``partitions × buckets`` the writer would naturally produce.
-
-For HASH_FIXED tables we group rows by ``(partition_keys..., bucket)``
-so every distinct group lands in a single Ray task. ``bucket`` is
-computed using the same ``FixedBucketRowKeyExtractor`` the writer
-uses, so the bucket assignment seen by the groupby is byte-equivalent
-to the writer's. HASH_FIXED writes are always pre-clustered; no user
-opt-in is required.
+"""Optional pre-clustering for Ray writes to HASH_FIXED Paimon tables.
+
+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.
For any other bucket mode the dataset is returned unchanged.
"""
@@ -51,6 +43,14 @@ if TYPE_CHECKING:
# runtime by ``_pick_bucket_col_name`` so user tables that happen to
# contain a column with this name still work correctly.
BUCKET_KEY_COL = "__paimon_bucket__"
+HASH_FIXED_PRECLUSTER_AUTO = "auto"
+HASH_FIXED_PRECLUSTER_OFF = "off"
+HASH_FIXED_PRECLUSTER_MAP_GROUPS = "map_groups"
+HASH_FIXED_PRECLUSTER_MODES = frozenset([
+ HASH_FIXED_PRECLUSTER_AUTO,
+ HASH_FIXED_PRECLUSTER_OFF,
+ HASH_FIXED_PRECLUSTER_MAP_GROUPS,
+])
def _pick_bucket_col_name(existing_names) -> str:
@@ -67,16 +67,42 @@ def _pick_bucket_col_name(existing_names) -> str:
def maybe_apply_repartition(
dataset: "ray.data.Dataset",
table: "Table",
+ hash_fixed_precluster: str = HASH_FIXED_PRECLUSTER_AUTO,
) -> "ray.data.Dataset":
- """Cluster rows by ``(partition_keys..., bucket)`` for HASH_FIXED tables.
-
- For any other bucket mode the dataset is returned unchanged.
- HASH_FIXED writes are always pre-clustered, with no user opt-in
- required.
+ """Optionally cluster rows for HASH_FIXED tables.
+
+ ``auto`` currently behaves like ``off`` for append-only tables
+ because the old ``map_groups`` strategy materializes each
+ ``(partition, bucket)`` group on one Ray node. For primary-key
+ tables, direct writes are rejected because multiple Ray tasks can
+ write the same bucket with overlapping sequence numbers. Use
+ ``map_groups`` only when both bounds are acceptable for the
+ workload.
"""
+ if hash_fixed_precluster not in HASH_FIXED_PRECLUSTER_MODES:
+ raise ValueError(
+ "hash_fixed_precluster must be one of {}, got {!r}".format(
+ sorted(HASH_FIXED_PRECLUSTER_MODES),
+ hash_fixed_precluster,
+ )
+ )
+
if table.bucket_mode() != BucketMode.HASH_FIXED:
return dataset
+ if hash_fixed_precluster in (
+ HASH_FIXED_PRECLUSTER_AUTO,
+ HASH_FIXED_PRECLUSTER_OFF,
+ ):
+ if getattr(table, "is_primary_key_table", False):
+ raise ValueError(
+ "HASH_FIXED primary-key Ray writes require "
+ "hash_fixed_precluster='map_groups'. Direct writes can "
+ "create overlapping sequence numbers when multiple Ray "
+ "tasks write the same bucket."
+ )
+ return dataset
+
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)
diff --git a/paimon-python/pypaimon/tests/ray_repartition_test.py
b/paimon-python/pypaimon/tests/ray_repartition_test.py
index b66b014b4f..01048e850b 100644
--- a/paimon-python/pypaimon/tests/ray_repartition_test.py
+++ b/paimon-python/pypaimon/tests/ray_repartition_test.py
@@ -16,18 +16,20 @@
# limitations under the License.
################################################################################
-"""End-to-end tests for HASH_FIXED auto-clustering on ``write_paimon``.
+"""End-to-end tests for HASH_FIXED Ray writes.
-For HASH_FIXED tables, ``write_paimon`` automatically pre-clusters rows
-by ``(partition_keys..., bucket)`` (matching Spark/Flink). These tests
-cover:
+For append-only HASH_FIXED tables, ``write_paimon`` writes rows to the
+correct bucket by default without pre-clustering. HASH_FIXED
+primary-key tables fail fast unless the legacy ``map_groups`` mode is
+explicitly selected. These tests cover:
- * roundtrip correctness on a HASH_FIXED PK table.
+ * default roundtrip correctness on an append-only HASH_FIXED table.
+ * default fail-fast behaviour on a HASH_FIXED PK table.
* roundtrip correctness on a partitioned HASH_FIXED PK table.
- * the transient bucket column is stripped from the sink-visible
- schema.
- * the output is one file per (partition, bucket) — i.e. the
- small-file storm is eliminated.
+ * explicit ``map_groups`` mode strips the transient bucket column
+ from the sink-visible schema.
+ * explicit ``map_groups`` mode can produce one file per
+ (partition, bucket) on the small test dataset.
* regression: a table whose schema already contains a column named
``__paimon_bucket__`` still works (collision-safe column name).
* non-HASH_FIXED tables (BUCKET_UNAWARE etc.) pass through unchanged.
@@ -114,19 +116,18 @@ class RayShuffleTest(unittest.TestCase):
))
return files
- # ----- HASH_FIXED auto-clustering -----
+ # ----- HASH_FIXED writes -----
- def test_fixed_bucket_roundtrip(self):
+ def test_append_only_fixed_bucket_roundtrip(self):
from pypaimon.ray import write_paimon
pa_schema = pa.schema([
- pa.field('id', pa.int32(), nullable=False),
+ ('id', pa.int32()),
('name', pa.string()),
])
- table_name = 'test_fixed_bucket_roundtrip'
+ table_name = 'test_append_only_fixed_bucket_roundtrip'
identifier = self._make_table(
- table_name, pa_schema,
- primary_keys=['id'], options={'bucket': '4'},
+ table_name, pa_schema, options={'bucket': '4'},
)
rows = pa.Table.from_pydict(
@@ -141,6 +142,54 @@ class RayShuffleTest(unittest.TestCase):
self.assertEqual(set(result['id']), set(range(40)))
self.assertNotIn('__paimon_bucket__', result.columns)
+ def test_primary_key_fixed_bucket_default_fails_fast(self):
+ from pypaimon.ray import write_paimon
+
+ pa_schema = pa.schema([
+ pa.field('id', pa.int32(), nullable=False),
+ ('name', pa.string()),
+ ])
+ table_name = 'test_pk_fixed_bucket_default_fails_fast'
+ identifier = self._make_table(
+ table_name, pa_schema,
+ primary_keys=['id'], options={'bucket': '4'},
+ )
+
+ rows = pa.Table.from_pydict(
+ {'id': list(range(40)), 'name': [f'v{i}' for i in range(40)]},
+ schema=pa_schema,
+ )
+ ds = ray.data.from_arrow(rows).repartition(4)
+
+ with self.assertRaisesRegex(ValueError, "HASH_FIXED primary-key"):
+ write_paimon(ds, identifier, self.catalog_options)
+
+ def test_table_write_ray_primary_key_fixed_bucket_default_fails_fast(self):
+ pa_schema = pa.schema([
+ pa.field('id', pa.int32(), nullable=False),
+ ('name', pa.string()),
+ ])
+ table_name = 'test_table_write_ray_pk_default_fails_fast'
+ identifier = self._make_table(
+ table_name, pa_schema,
+ primary_keys=['id'], options={'bucket': '4'},
+ )
+
+ rows = pa.Table.from_pydict(
+ {'id': list(range(40)), 'name': [f'v{i}' for i in range(40)]},
+ schema=pa_schema,
+ )
+ ds = ray.data.from_arrow(rows).repartition(4)
+
+ catalog = CatalogFactory.create(self.catalog_options)
+ table = catalog.get_table(identifier)
+ writer = table.new_batch_write_builder().new_write()
+ try:
+ with self.assertRaisesRegex(ValueError, "HASH_FIXED primary-key"):
+ writer.write_ray(ds)
+ finally:
+ writer.close()
+
def test_partitioned_fixed_bucket_roundtrip(self):
"""Partitioned table — confirms the post-groupby schema does not
end up with duplicated partition-key or bucket columns."""
@@ -164,16 +213,50 @@ class RayShuffleTest(unittest.TestCase):
'value': list(range(20)),
}, schema=pa_schema)
ds = ray.data.from_arrow(rows).repartition(4)
- write_paimon(ds, identifier, self.catalog_options)
+ write_paimon(
+ ds,
+ identifier,
+ self.catalog_options,
+ hash_fixed_precluster="map_groups",
+ )
result = self._read_table(identifier)
self.assertEqual(set(result.columns), {'id', 'dt', 'value'})
self.assertEqual(len(result), 20)
self.assertEqual(set(result['dt']), {'2026-01-01', '2026-01-02'})
+ def
test_table_write_ray_primary_key_fixed_bucket_map_groups_roundtrip(self):
+ pa_schema = pa.schema([
+ pa.field('id', pa.int32(), nullable=False),
+ ('name', pa.string()),
+ ])
+ table_name = 'test_table_write_ray_pk_map_groups'
+ identifier = self._make_table(
+ table_name, pa_schema,
+ primary_keys=['id'], options={'bucket': '4'},
+ )
+
+ rows = pa.Table.from_pydict(
+ {'id': list(range(40)), 'name': [f'v{i}' for i in range(40)]},
+ schema=pa_schema,
+ )
+ ds = ray.data.from_arrow(rows).repartition(4)
+
+ catalog = CatalogFactory.create(self.catalog_options)
+ table = catalog.get_table(identifier)
+ writer = table.new_batch_write_builder().new_write()
+ try:
+ writer.write_ray(ds, hash_fixed_precluster="map_groups")
+ finally:
+ writer.close()
+
+ result = self._read_table(identifier)
+ self.assertEqual(len(result), 40)
+ self.assertEqual(set(result['id']), set(range(40)))
+
def test_fixed_bucket_writes_one_file_per_bucket(self):
- """With multiple input blocks, auto-clustering collapses per-task
- files into per-bucket files."""
+ """With multiple input blocks, explicit map_groups clustering
+ collapses per-task files into per-bucket files."""
from pypaimon.ray import write_paimon
pa_schema = pa.schema([
@@ -190,11 +273,12 @@ class RayShuffleTest(unittest.TestCase):
primary_keys=['id'], options={'bucket': '4'},
)
- # Materialise 4 input blocks. Without auto-clustering each task
- # would emit one file per bucket it touched (up to 16 files).
+ # Materialise 4 input blocks. Without the explicit map_groups
+ # mode, each task would emit one file per bucket it touched.
write_paimon(
ray.data.from_arrow(rows).repartition(4),
identifier, self.catalog_options,
+ hash_fixed_precluster="map_groups",
)
files = self._count_data_files('test_one_file_per_bucket')
@@ -223,7 +307,12 @@ class RayShuffleTest(unittest.TestCase):
schema=pa_schema,
)
ds = ray.data.from_arrow(rows).repartition(2)
- write_paimon(ds, identifier, self.catalog_options)
+ write_paimon(
+ ds,
+ identifier,
+ self.catalog_options,
+ hash_fixed_precluster="map_groups",
+ )
result = self._read_table(identifier)
self.assertEqual(len(result), 10)
diff --git a/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
b/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
index a849f2788e..eb974dd79a 100644
--- a/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
+++ b/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
@@ -135,12 +135,13 @@ class CoerceLargeStringTypesTest(unittest.TestCase):
class BucketModeDispatchTest(unittest.TestCase):
- """``maybe_apply_repartition`` clusters HASH_FIXED tables and
- returns other bucket modes unchanged."""
+ """``maybe_apply_repartition`` only clusters HASH_FIXED tables when
+ the legacy ``map_groups`` mode is explicitly selected."""
def _make_table(self, bucket_mode):
table = MagicMock()
table.bucket_mode.return_value = bucket_mode
+ table.is_primary_key_table = False
return table
def test_bucket_unaware_returns_dataset_unchanged(self):
@@ -161,7 +162,48 @@ class BucketModeDispatchTest(unittest.TestCase):
self.assertIs(maybe_apply_repartition(dataset, table), dataset)
- def test_hash_fixed_runs_map_batches_groupby_chain(self):
+ def test_hash_fixed_default_returns_dataset_unchanged(self):
+ dataset = MagicMock(name="dataset")
+ table = MagicMock()
+ table.bucket_mode.return_value = BucketMode.HASH_FIXED
+ table.is_primary_key_table = False
+
+ self.assertIs(maybe_apply_repartition(dataset, table), dataset)
+ dataset.map_batches.assert_not_called()
+
+ def test_hash_fixed_off_returns_dataset_unchanged(self):
+ dataset = MagicMock(name="dataset")
+ table = MagicMock()
+ table.bucket_mode.return_value = BucketMode.HASH_FIXED
+ table.is_primary_key_table = False
+
+ self.assertIs(
+ maybe_apply_repartition(dataset, table, "off"),
+ dataset,
+ )
+ dataset.map_batches.assert_not_called()
+
+ def test_hash_fixed_primary_key_default_raises_value_error(self):
+ dataset = MagicMock(name="dataset")
+ table = MagicMock()
+ table.bucket_mode.return_value = BucketMode.HASH_FIXED
+ table.is_primary_key_table = True
+
+ with self.assertRaises(ValueError):
+ maybe_apply_repartition(dataset, table)
+ dataset.map_batches.assert_not_called()
+
+ def test_hash_fixed_primary_key_off_raises_value_error(self):
+ dataset = MagicMock(name="dataset")
+ table = MagicMock()
+ table.bucket_mode.return_value = BucketMode.HASH_FIXED
+ table.is_primary_key_table = True
+
+ with self.assertRaises(ValueError):
+ maybe_apply_repartition(dataset, table, "off")
+ dataset.map_batches.assert_not_called()
+
+ def test_hash_fixed_map_groups_runs_map_batches_groupby_chain(self):
dataset = MagicMock(name="dataset")
dataset.map_batches.return_value.groupby.return_value \
.map_groups.return_value.drop_columns.return_value = "clustered"
@@ -173,7 +215,7 @@ class BucketModeDispatchTest(unittest.TestCase):
type("F", (), {"name": "value"})(),
]
- out = maybe_apply_repartition(dataset, table)
+ out = maybe_apply_repartition(dataset, table, "map_groups")
self.assertEqual(out, "clustered")
# The helper appends a transient bucket column, groups by it,
@@ -199,12 +241,19 @@ class BucketModeDispatchTest(unittest.TestCase):
type("F", (), {"name": "dt"})(),
]
- maybe_apply_repartition(dataset, table)
+ maybe_apply_repartition(dataset, table, "map_groups")
group_call = dataset.map_batches.return_value.groupby.call_args
passed_keys = group_call.args[0]
self.assertEqual(passed_keys, ["dt", BUCKET_KEY_COL])
+ def test_invalid_precluster_mode_raises_value_error(self):
+ dataset = object()
+ table = self._make_table(BucketMode.HASH_FIXED)
+
+ with self.assertRaises(ValueError):
+ maybe_apply_repartition(dataset, table, "hash_shuffle")
+
if __name__ == "__main__":
unittest.main()
diff --git a/paimon-python/pypaimon/write/table_write.py
b/paimon-python/pypaimon/write/table_write.py
index 80ef5a3572..9336fff028 100644
--- a/paimon-python/pypaimon/write/table_write.py
+++ b/paimon-python/pypaimon/write/table_write.py
@@ -77,6 +77,7 @@ class TableWrite:
overwrite: bool = False,
concurrency: Optional[int] = None,
ray_remote_args: Optional[Dict[str, Any]] = None,
+ hash_fixed_precluster: str = "auto",
) -> None:
"""
Write a Ray Dataset to Paimon table.
@@ -89,8 +90,18 @@ class TableWrite:
By default, dynamically decided based on available resources.
ray_remote_args: Optional kwargs passed to :func:`ray.remote` in
write tasks.
For example, ``{"num_cpus": 2, "max_retries": 3}``.
+ 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.
"""
+ 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)
+
datasink = PaimonDatasink(self.table, overwrite=overwrite)
dataset.write_datasink(
datasink,
