From d8f76e7b35a250771424ceda20449cd7341d5d01 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 3 Aug 2026 01:44:00 -0700 Subject: [PATCH 1/7] [python] Write postpone batches to fixed buckets --- docs/docs/pypaimon/python-api.mdx | 4 + docs/docs/pypaimon/ray-data.md | 16 +- .../pypaimon/common/options/core_options.py | 67 +++ paimon-python/pypaimon/daft/daft_datasink.py | 8 +- paimon-python/pypaimon/ray/ray_paimon.py | 7 +- paimon-python/pypaimon/ray/shuffle.py | 4 +- .../daft/daft_pk_distributed_write_test.py | 4 + .../pypaimon/tests/ray_repartition_test.py | 30 +- paimon-python/pypaimon/tests/ray_sink_test.py | 42 ++ .../tests/write/conflict_detection_test.py | 16 + .../pypaimon/tests/write/table_write_test.py | 558 +++++++++++++++++- .../write/commit/conflict_detection.py | 25 + .../commit/overwrite_changes_provider.py | 7 +- .../pypaimon/write/commit_message.py | 1 + .../pypaimon/write/file_store_commit.py | 20 +- .../pypaimon/write/file_store_write.py | 5 +- .../write/postpone_batch_table_write.py | 165 ++++++ .../pypaimon/write/postpone_bucket.py | 189 ++++++ paimon-python/pypaimon/write/ray_datasink.py | 104 +++- .../pypaimon/write/row_key_extractor.py | 58 +- paimon-python/pypaimon/write/table_write.py | 24 +- paimon-python/pypaimon/write/write_builder.py | 23 + 22 files changed, 1337 insertions(+), 40 deletions(-) create mode 100644 paimon-python/pypaimon/write/postpone_batch_table_write.py create mode 100644 paimon-python/pypaimon/write/postpone_bucket.py diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index c1887c537c75..d6d12e6d2d01 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -321,6 +321,10 @@ table_write.close() table_commit.close() ``` +For postpone-bucket tables (`bucket = -2`) with fixed-bucket batch writes +enabled, inputs for new partitions are buffered until `prepare_commit()`. +Memory grows with buffered input; existing-bucket partitions remain incremental. + By default, the data will be appended to table. If you want to overwrite table, you should use `TableWrite#overwrite` API: diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 7f56169eb46d..9f2b9d353e76 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -248,8 +248,12 @@ 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. -Postpone-bucket primary-key tables (`bucket = -2`) are also written -as-is to the `bucket-postpone` directory. HASH_DYNAMIC and +By default, postpone-bucket primary-key tables (`bucket = -2`) are +written to real buckets and are immediately visible. Existing partitions +reuse their bucket count; new partitions infer one from the configured target +row count or size. Ray groups rows by `(partition, bucket)`, so each group must +fit in one worker's memory. Set `postpone.batch-write-fixed-bucket` to `false` +to write to `bucket-postpone` instead. HASH_DYNAMIC and CROSS_PARTITION primary-key Ray writes are not supported and fail fast, including the default dynamic-bucket primary-key table (`bucket = -1`). Ray write tasks create independent Paimon writers, which can assign @@ -261,11 +265,11 @@ 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. + For grouped HASH_FIXED primary-key and postpone-bucket writes, this + limits the group writer tasks. - `ray_remote_args`: optional kwargs passed to `ray.remote()` in write tasks - (e.g. `{"num_cpus": 2}`). For HASH_FIXED primary-key `map_groups` - writes, these options apply to the group writer tasks. + (e.g. `{"num_cpus": 2}`). For grouped HASH_FIXED primary-key and + postpone-bucket 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 diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index 8ee6016e6465..c51225baec33 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -227,6 +227,46 @@ class CoreOptions: ) ) + POSTPONE_BATCH_WRITE_FIXED_BUCKET: ConfigOption[bool] = ( + ConfigOptions.key("postpone.batch-write-fixed-bucket") + .boolean_type() + .default_value(True) + .with_description( + "Whether to write data into fixed buckets for batch writes to a " + "postpone bucket table." + ) + ) + + POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM: ConfigOption[int] = ( + ConfigOptions.key("postpone.batch-write-fixed-bucket.max-parallelism") + .int_type() + .default_value(2048) + .with_description( + "Maximum bucket number inferred for a postpone batch write." + ) + ) + + POSTPONE_TARGET_ROW_NUM_PER_BUCKET: ConfigOption[int] = ( + ConfigOptions.key("postpone.target-row-num-per-bucket") + .long_type() + .no_default_value() + .with_description( + "Target row number per bucket when batch writing a postpone " + "partition without real bucket data." + ) + ) + + POSTPONE_TARGET_SIZE_PER_BUCKET: ConfigOption[MemorySize] = ( + ConfigOptions.key("postpone.target-size-per-bucket") + .memory_type() + .default_value(MemorySize.parse("1 gb")) + .with_description( + "Target uncompressed input size per bucket when batch writing a " + "postpone partition without real bucket data. This option is " + "ignored when postpone.target-row-num-per-bucket is configured." + ) + ) + SCAN_MANIFEST_PARALLELISM: ConfigOption[int] = ( ConfigOptions.key("scan.manifest.parallelism") .int_type() @@ -1066,6 +1106,33 @@ def dynamic_bucket_target_row_num(self, default=None): def dynamic_bucket_max_buckets(self, default=None): return self.options.get(CoreOptions.DYNAMIC_BUCKET_MAX_BUCKETS, default) + def postpone_batch_write_fixed_bucket(self, default=None): + return self.options.get( + CoreOptions.POSTPONE_BATCH_WRITE_FIXED_BUCKET, default + ) + + def postpone_batch_write_fixed_bucket_max_parallelism(self, default=None): + return self.options.get( + CoreOptions.POSTPONE_BATCH_WRITE_FIXED_BUCKET_MAX_PARALLELISM, + default, + ) + + def postpone_target_row_num_per_bucket(self, default=None): + return self.options.get( + CoreOptions.POSTPONE_TARGET_ROW_NUM_PER_BUCKET, default + ) + + def postpone_target_size_per_bucket(self, default=None): + if default is not None and not isinstance(default, MemorySize): + default = ( + MemorySize.of_bytes(default) + if isinstance(default, int) + else MemorySize.parse(default) + ) + return self.options.get( + CoreOptions.POSTPONE_TARGET_SIZE_PER_BUCKET, default + ).get_bytes() + def scan_manifest_parallelism(self, default=None): return self.options.get(CoreOptions.SCAN_MANIFEST_PARALLELISM, default) diff --git a/paimon-python/pypaimon/daft/daft_datasink.py b/paimon-python/pypaimon/daft/daft_datasink.py index c21767b0d2e9..6065bcc19e57 100644 --- a/paimon-python/pypaimon/daft/daft_datasink.py +++ b/paimon-python/pypaimon/daft/daft_datasink.py @@ -195,9 +195,15 @@ def _init_table(self, table: FileStoreTable) -> None: self._table = table from pypaimon.schema.data_types import PyarrowFieldParser + from pypaimon.table.bucket_mode import BucketMode self._target_schema: pa.Schema = PyarrowFieldParser.from_paimon_schema(table.fields) - self._write_builder = table.new_batch_write_builder() + write_table = table + if table.bucket_mode() == BucketMode.POSTPONE_MODE: + write_table = table.copy({ + "postpone.batch-write-fixed-bucket": "false", + }) + self._write_builder = write_table.new_batch_write_builder() if ( self._commit_user is not None and hasattr(self._write_builder, "commit_user") diff --git a/paimon-python/pypaimon/ray/ray_paimon.py b/paimon-python/pypaimon/ray/ray_paimon.py index 3dcf40616db3..f068fda7efa5 100644 --- a/paimon-python/pypaimon/ray/ray_paimon.py +++ b/paimon-python/pypaimon/ray/ray_paimon.py @@ -295,9 +295,10 @@ def write_paimon( 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. + group fits in memory. Postpone-bucket batch writes resolve existing + bucket counts once on the driver and always use this grouped path. + 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. diff --git a/paimon-python/pypaimon/ray/shuffle.py b/paimon-python/pypaimon/ray/shuffle.py index eef29bdfc558..b15bff5a8b20 100644 --- a/paimon-python/pypaimon/ray/shuffle.py +++ b/paimon-python/pypaimon/ray/shuffle.py @@ -126,9 +126,11 @@ def maybe_apply_repartition( def _group_by_partition_bucket( dataset: "ray.data.Dataset", table: "Table", + extractor=None, ): partition_keys = list(table.table_schema.partition_keys or []) - extractor = table.create_row_key_extractor() + if extractor is None: + extractor = table.create_row_key_extractor() col_names = set(f.name for f in table.table_schema.fields) bucket_col = _pick_bucket_col_name(col_names) bucket_udf = _make_bucket_udf(extractor, bucket_col) diff --git a/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py b/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py index 7c327a3a664a..1c2ff559b909 100644 --- a/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py +++ b/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py @@ -227,6 +227,10 @@ def test_postpone_mode_primary_key_write_preserves_existing_path(catalog): assert sum(summary["rows"]) == 2 assert table.snapshot_manager().get_latest_snapshot() is not None + scanner = table.new_read_builder().new_scan().file_scanner + manifests, _ = scanner.manifest_scanner() + entries = scanner.manifest_file_manager.read_entries_parallel(manifests) + assert {entry.bucket for entry in entries} == {-2} def test_group_udfs_reuse_worker_table_metadata(catalog, monkeypatch): diff --git a/paimon-python/pypaimon/tests/ray_repartition_test.py b/paimon-python/pypaimon/tests/ray_repartition_test.py index 8229d9467ce6..b835657cf9b7 100644 --- a/paimon-python/pypaimon/tests/ray_repartition_test.py +++ b/paimon-python/pypaimon/tests/ray_repartition_test.py @@ -37,7 +37,7 @@ ``__paimon_bucket__`` still works (collision-safe column name). * non-HASH_FIXED append-only tables pass through unchanged. * dynamic-bucket primary-key tables fail fast, while postpone-bucket - primary-key tables pass through. + primary-key tables plan once and group multi-block input by real bucket. """ import glob @@ -241,7 +241,7 @@ def test_table_write_ray_primary_key_dynamic_bucket_default_fails_fast(self): finally: writer.close() - def test_primary_key_postpone_bucket_roundtrip_to_postpone_files(self): + def test_primary_key_postpone_bucket_multi_block_single_writer_per_bucket(self): from pypaimon.ray import write_paimon pa_schema = pa.schema([ @@ -253,24 +253,34 @@ def test_primary_key_postpone_bucket_roundtrip_to_postpone_files(self): identifier = self._make_table( table_name, pa_schema, primary_keys=['id', 'dt'], partition_keys=['dt'], - options={'bucket': '-2'}, + options={ + 'bucket': '-2', + 'postpone.target-size-per-bucket': '1000 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': '2', + }, ) rows = pa.Table.from_pydict({ - 'id': list(range(10)), - 'dt': ['2026-01-01'] * 5 + ['2026-01-02'] * 5, - 'value': list(range(10)), + 'id': list(range(100)), + 'dt': ['2026-01-01'] * 50 + ['2026-01-02'] * 50, + 'value': list(range(100)), }, schema=pa_schema) + dataset = ray.data.from_arrow(rows).repartition(8).materialize() + self.assertEqual(8, dataset.num_blocks()) write_paimon( - ray.data.from_arrow(rows).repartition(2), + dataset, identifier, self.catalog_options, + concurrency=4, ) files = self._count_data_files(table_name) - self.assertGreater(len(files), 0) - self.assertTrue(all('/bucket-postpone/' in path for path in files)) - self.assertEqual(len(self._read_table(identifier)), 0) + self.assertEqual(4, len(files)) + self.assertEqual( + {'bucket-0', 'bucket-1'}, + {os.path.basename(os.path.dirname(path)) for path in files}, + ) + self.assertEqual(len(self._read_table(identifier)), 100) def test_partitioned_fixed_bucket_roundtrip(self): """Partitioned table — confirms the post-groupby schema does not diff --git a/paimon-python/pypaimon/tests/ray_sink_test.py b/paimon-python/pypaimon/tests/ray_sink_test.py index 5dc462304aeb..e6e465f19755 100644 --- a/paimon-python/pypaimon/tests/ray_sink_test.py +++ b/paimon-python/pypaimon/tests/ray_sink_test.py @@ -304,6 +304,48 @@ def test_write(self): mock_write.prepare_commit.assert_called_once() mock_write.abort.assert_called_once() + def test_postpone_worker_uses_driver_bucket_plan_without_manifest_scan(self): + from pypaimon.write.postpone_bucket import ( + PostponeBucketPlan, + PostponeBucketPlanner, + ) + + pa_schema = pa.schema([ + pa.field('id', pa.int64(), nullable=False), + ('name', pa.string()), + ('value', pa.float64()), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + primary_keys=['id'], + options={ + 'bucket': '-2', + }, + ) + identifier = 'test_db.test_postpone_worker_plan' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + datasink = PaimonDatasink( + table, + postpone_bucket_plan=PostponeBucketPlan({(): 2}), + ) + data = pa.Table.from_pydict({ + 'id': list(range(20)), + 'name': ['name-{}'.format(i) for i in range(20)], + 'value': [float(i) for i in range(20)], + }, schema=pa_schema) + + with patch.object( + PostponeBucketPlanner, + '_load_bucket_metadata', + side_effect=AssertionError("worker must not scan manifests"), + ) as load: + messages = datasink.write([data], Mock(spec=TaskContext)) + + load.assert_not_called() + self.assertEqual({0, 1}, {message.bucket for message in messages}) + self.assertEqual({2}, {message.total_buckets for message in messages}) + def test_write_does_not_return_prepared_messages_when_dedicated_close_aborts(self): from pypaimon.write.writer.dedicated_format_writer import DedicatedFormatWriter diff --git a/paimon-python/pypaimon/tests/write/conflict_detection_test.py b/paimon-python/pypaimon/tests/write/conflict_detection_test.py index 62889370f56d..4d02a212523b 100644 --- a/paimon-python/pypaimon/tests/write/conflict_detection_test.py +++ b/paimon-python/pypaimon/tests/write/conflict_detection_test.py @@ -311,6 +311,22 @@ def test_delete_entry_missing_from_base_conflicts(self): self.assertIsNotNone(result) self.assertIn("File deletion conflicts", str(result)) + def test_bucket_num_mismatch_conflicts(self): + detection = self._make_detection() + old = _make_entry("old") + new = _make_entry("new") + new.total_buckets = 2 + + result = detection.check_conflicts( + latest_snapshot=None, + base_entries=[old], + delta_entries=[new], + commit_kind="APPEND", + ) + + self.assertIsNotNone(result) + self.assertIn("Total buckets", str(result)) + class _FakeSnapshot: diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index d078304d2770..e1ce3cd3cafd 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -426,8 +426,16 @@ def test_multi_prepare_commit_pk(self): self.assertEqual(self.pk_expected, actual) def test_postpone_read_write(self): - schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], - options={'bucket': -2}) + schema = Schema.from_pyarrow_schema( + self.pa_schema, + partition_keys=['user_id'], + primary_keys=['user_id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) self.catalog.create_table('default.test_postpone', schema, False) table = self.catalog.get_table('default.test_postpone') data = { @@ -440,6 +448,11 @@ def test_postpone_read_write(self): write_builder = table.new_batch_write_builder() table_write = write_builder.new_write() + from pypaimon.write.postpone_batch_table_write import ( + PostponeFixedBucketBatchTableWrite, + ) + self.assertIsInstance( + table_write, PostponeFixedBucketBatchTableWrite) table_commit = write_builder.new_commit() table_write.write_arrow(expect) commit_messages = table_write.prepare_commit() @@ -451,18 +464,551 @@ def test_postpone_read_write(self): self.assertTrue(os.path.exists(self.warehouse + "/default.db/test_postpone/snapshot/snapshot-1")) self.assertTrue(os.path.exists(self.warehouse + "/default.db/test_postpone/manifest")) self.assertEqual(len(glob.glob(self.warehouse + "/default.db/test_postpone/manifest/*")), 3) - self.assertEqual(len(glob.glob(self.warehouse + "/default.db/test_postpone/user_id=2/bucket-postpone/*.avro")), - 1) + self.assertEqual({2}, {message.total_buckets for message in commit_messages}) + self.assertEqual( + 1, + len(glob.glob( + self.warehouse + + "/default.db/test_postpone/user_id=2/bucket-[01]/*.parquet" + )), + ) read_builder = table.new_read_builder() table_read = read_builder.new_read() splits = read_builder.new_scan().plan().splits() actual = table_read.to_arrow(splits) - self.assertTrue(not actual) + self.assertEqual(expect, actual) + + def test_postpone_batch_fixed_bucket_disabled(self): + schema = Schema.from_pyarrow_schema( + self.pa_schema, + partition_keys=['user_id'], + primary_keys=['user_id', 'dt'], + options={ + 'bucket': -2, + 'postpone.batch-write-fixed-bucket': False, + }, + ) + self.catalog.create_table('default.test_postpone_disabled', schema, False) + table = self.catalog.get_table('default.test_postpone_disabled') + expected = pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['a'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema) + + write_builder = table.new_batch_write_builder() + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + table_write.write_arrow(expected) + table_commit.commit(table_write.prepare_commit()) + table_write.close() + table_commit.close() + + self.assertEqual( + 1, + len(glob.glob( + self.warehouse + + "/default.db/test_postpone_disabled/user_id=1/" + + "bucket-postpone/*.avro" + )), + ) + splits = table.new_read_builder().new_scan().plan().splits() + self.assertTrue(not table.new_read_builder().new_read().to_arrow(splits)) + + def test_postpone_batch_infers_bucket_num_from_input_size(self): + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('dt', pa.string(), nullable=False), + ('value', pa.string()), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '100 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + identifier = 'default.test_postpone_size_inference' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + data = pa.Table.from_pydict({ + 'id': [1, 2], + 'dt': ['small', 'large'], + 'value': ['x', 'x' * 500], + }, schema=pa_schema) + + builder = table.new_batch_write_builder() + write = builder.new_write() + commit = builder.new_commit() + try: + write.write_arrow(data) + messages = write.prepare_commit() + total_buckets = { + tuple(message.partition): message.total_buckets + for message in messages + } + self.assertEqual(1, total_buckets[('small',)]) + self.assertEqual(3, total_buckets[('large',)]) + commit.commit(messages) + finally: + write.close() + commit.close() + + self.assertEqual( + [1, 2], self._read_sorted(table, 'id').column('id').to_pylist() + ) + + def test_postpone_batch_plans_all_record_batches(self): + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('dt', pa.string(), nullable=False), + ('value', pa.string()), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '100 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + identifier = 'default.test_postpone_multi_batch_plan' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + + builder = table.new_batch_write_builder() + write = builder.new_write() + commit = builder.new_commit() + try: + write.write_arrow_batch(pa.RecordBatch.from_pydict({ + 'id': [1], + 'dt': ['p'], + 'value': ['x'], + }, schema=pa_schema)) + write.write_arrow_batch(pa.RecordBatch.from_pydict({ + 'id': [2], + 'dt': ['p'], + 'value': ['x' * 500], + }, schema=pa_schema)) + messages = write.prepare_commit() + self.assertEqual({3}, {message.total_buckets for message in messages}) + commit.commit(messages) + finally: + write.close() + commit.close() + + self.assertEqual( + [1, 2], self._read_sorted(table, 'id').column('id').to_pylist() + ) + + def test_postpone_batch_prefers_target_row_num(self): + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('dt', pa.string(), nullable=False), + ('value', pa.string()), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-row-num-per-bucket': 1, + 'postpone.target-size-per-bucket': '1 gb', + 'postpone.batch-write-fixed-bucket.max-parallelism': 8, + }, + ) + identifier = 'default.test_postpone_row_num_plan' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + + builder = table.new_batch_write_builder() + write = builder.new_write() + commit = builder.new_commit() + try: + write.write_arrow(pa.Table.from_pydict({ + 'id': list(range(8)), + 'dt': ['p'] * 8, + 'value': ['x'] * 8, + }, schema=pa_schema)) + messages = write.prepare_commit() + self.assertEqual({8}, {message.total_buckets for message in messages}) + commit.commit(messages) + finally: + write.close() + commit.close() + + def test_postpone_batch_plans_write_rows_with_arrow_input(self): + from pypaimon.table.row.generic_row import GenericRow + + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('dt', pa.string(), nullable=False), + ('value', pa.string()), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-row-num-per-bucket': 1, + 'postpone.target-size-per-bucket': '1 gb', + 'postpone.batch-write-fixed-bucket.max-parallelism': 8, + }, + ) + identifier = 'default.test_postpone_write_row_plan' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + + builder = table.new_batch_write_builder() + write = builder.new_write() + commit = builder.new_commit() + try: + for row_id in range(8): + write.write_row( + GenericRow([row_id, 'rows', 'x'], table.fields) + ) + write.write_row(GenericRow([8, 'mixed', 'x'], table.fields)) + write.write_arrow(pa.Table.from_pydict({ + 'id': list(range(9, 16)), + 'dt': ['mixed'] * 7, + 'value': ['x'] * 7, + }, schema=pa_schema)) + + messages = write.prepare_commit() + total_buckets = { + tuple(message.partition): message.total_buckets + for message in messages + } + self.assertEqual(8, total_buckets[('rows',)]) + self.assertEqual(8, total_buckets[('mixed',)]) + commit.commit(messages) + finally: + write.close() + commit.close() + + self.assertEqual(16, self._read_sorted(table, 'id').num_rows) + + def test_postpone_target_row_num_counts_existing_postpone_rows(self): + from pypaimon.write.postpone_bucket import ( + PostponeBucketPlanner, + ) + + schema = Schema.from_pyarrow_schema( + self.pk_pa_schema, + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-row-num-per-bucket': 2, + 'postpone.batch-write-fixed-bucket.max-parallelism': 8, + }, + ) + identifier = 'default.test_postpone_existing_row_count' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + append_planner = PostponeBucketPlanner( + table, + known_num_buckets={}, + postpone_row_counts={('p',): 3}, + ) + append_plan = append_planner.plan({('p',): (1, 10)}) + self.assertEqual(2, append_plan.num_buckets(('p',))) + + overwrite_planner = PostponeBucketPlanner( + table, + known_num_buckets={}, + postpone_row_counts={('p',): 3}, + ) + overwrite_plan = overwrite_planner.plan( + {('p',): (1, 10)}, include_postpone_rows=False + ) + self.assertEqual(1, overwrite_plan.num_buckets(('p',))) + + def test_postpone_worker_bucket_plan_mismatch_fails_commit(self): + from pypaimon.write.commit.conflict_detection import CommitConflictError + + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('dt', pa.string(), nullable=False), + ('value', pa.string()), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '100 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + identifier = 'default.test_postpone_worker_plan_mismatch' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + small_builder = table.new_batch_write_builder() + large_builder = table.new_batch_write_builder() + small_write = small_builder.new_write() + large_write = large_builder.new_write() + commit = small_builder.new_commit() + try: + small_write.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'dt': ['p'], 'value': ['x'], + }, schema=pa_schema)) + large_write.write_arrow(pa.Table.from_pydict({ + 'id': [2], 'dt': ['p'], 'value': ['x' * 500], + }, schema=pa_schema)) + small_messages = small_write.prepare_commit() + large_messages = large_write.prepare_commit() + self.assertEqual({1}, {m.total_buckets for m in small_messages}) + self.assertEqual({3}, {m.total_buckets for m in large_messages}) + with self.assertRaisesRegex(CommitConflictError, 'Total buckets'): + commit.commit(small_messages + large_messages) + finally: + small_write.close() + large_write.close() + commit.close() + + def test_postpone_batch_fixed_bucket_reuses_existing_bucket_num(self): + schema = Schema.from_pyarrow_schema( + self.pa_schema, + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) + self.catalog.create_table('default.test_postpone_reuse', schema, False) + table = self.catalog.get_table('default.test_postpone_reuse') + expected = pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['a'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema) + + write_builder = table.new_batch_write_builder() + table_write = write_builder.new_write() + table_commit = write_builder.new_commit() + table_write.write_arrow(expected) + table_commit.commit(table_write.prepare_commit()) + table_write.close() + table_commit.close() + + copied_table = table.copy({ + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }) + from pypaimon.write.postpone_bucket import PostponeBucketPlanner + + planner = PostponeBucketPlanner(copied_table) + plan = planner.plan({('p2',): (1, 10)}) + copied_write = ( + copied_table.new_batch_write_builder() + .with_bucket_plan(plan) + .new_write() + ) + self.assertEqual(2, copied_write.row_key_extractor.num_buckets(('p1',))) + self.assertEqual(3, copied_write.row_key_extractor.num_buckets(('p2',))) + copied_write.close() + + def test_postpone_reuses_bucket_num_for_int_date_partition(self): + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('part', pa.int32(), nullable=False), + pa.field('day', pa.date32(), nullable=False), + ('value', pa.string()), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + partition_keys=['part', 'day'], + primary_keys=['id', 'part', 'day'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) + identifier = 'default.test_postpone_typed_partition_reuse' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + day = datetime.date(2026, 8, 1) + + first_builder = table.new_batch_write_builder() + first_write = first_builder.new_write() + first_commit = first_builder.new_commit() + first_write.write_arrow(pa.Table.from_pydict({ + 'id': [1], + 'part': [7], + 'day': [day], + 'value': ['first'], + }, schema=pa_schema)) + first_commit.commit(first_write.prepare_commit()) + first_write.close() + first_commit.close() + + # Reopen to load the existing bucket count from manifests. + reopened = self.catalog.get_table(identifier).copy({ + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }) + second_builder = reopened.new_batch_write_builder() + second_write = second_builder.new_write() + second_commit = second_builder.new_commit() + try: + second_write.write_arrow(pa.Table.from_pydict({ + 'id': [2, 3], + 'part': [7, 8], + 'day': [day, day], + 'value': ['existing', 'new'], + }, schema=pa_schema)) + messages = second_write.prepare_commit() + total_buckets = { + tuple(message.partition): message.total_buckets + for message in messages + } + self.assertEqual(2, total_buckets[(7, day)]) + self.assertEqual(3, total_buckets[(8, day)]) + second_commit.commit(messages) + finally: + second_write.close() + second_commit.close() + + actual = self._read_sorted(reopened, 'id') + self.assertEqual([1, 2, 3], actual.column('id').to_pylist()) + + def test_postpone_legacy_partition_migrates_to_fixed_bucket(self): + schema = Schema.from_pyarrow_schema( + self.pk_pa_schema, + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={ + 'bucket': -2, + 'postpone.batch-write-fixed-bucket': False, + }, + ) + identifier = 'default.test_postpone_legacy_partition_migration' + self.catalog.create_table(identifier, schema, False) + legacy = self.catalog.get_table(identifier) + legacy_builder = legacy.new_batch_write_builder() + legacy_write = legacy_builder.new_write() + legacy_commit = legacy_builder.new_commit() + legacy_write.write_arrow(pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['legacy'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema)) + legacy_commit.commit(legacy_write.prepare_commit()) + legacy_write.close() + legacy_commit.close() + + fixed = self.catalog.get_table(identifier).copy({ + 'postpone.batch-write-fixed-bucket': True, + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }) + fixed_builder = fixed.new_batch_write_builder() + fixed_write = fixed_builder.new_write() + fixed_commit = fixed_builder.new_commit() + try: + fixed_write.write_arrow(pa.Table.from_pydict({ + 'user_id': [2], + 'item_id': [1002], + 'behavior': ['fixed'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema)) + messages = fixed_write.prepare_commit() + # Legacy -2 files do not define a real bucket count. + self.assertEqual({2}, {m.total_buckets for m in messages}) + fixed_commit.commit(messages) + finally: + fixed_write.close() + fixed_commit.close() + + scanner = fixed.new_read_builder().new_scan().file_scanner + manifests, _ = scanner.manifest_scanner() + entries = scanner.manifest_file_manager.read_entries_parallel( + manifests, drop_stats=False + ) + partition_entries = [ + entry for entry in entries + if tuple(entry.partition.values) == ('p1',) + ] + self.assertEqual({-2, 2}, { + entry.total_buckets for entry in partition_entries + }) + self.assertIn(-2, {entry.bucket for entry in partition_entries}) + self.assertTrue(any(entry.bucket >= 0 for entry in partition_entries)) + + actual = self._read_sorted(fixed, 'user_id') + self.assertEqual([2], actual.column('user_id').to_pylist()) + + def test_postpone_concurrent_new_partition_bucket_num_conflict(self): + from pypaimon.write.commit.conflict_detection import CommitConflictError + + schema = Schema.from_pyarrow_schema( + self.pk_pa_schema, + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) + identifier = 'default.test_postpone_concurrent_bucket_num' + self.catalog.create_table(identifier, schema, False) + table_two_buckets = self.catalog.get_table(identifier) + table_three_buckets = table_two_buckets.copy({ + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }) + + builder_two = table_two_buckets.new_batch_write_builder() + builder_three = table_three_buckets.new_batch_write_builder() + write_two = builder_two.new_write() + write_three = builder_three.new_write() + commit_two = builder_two.new_commit() + commit_three = builder_three.new_commit() + try: + write_two.write_arrow(pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['a'], + 'dt': ['new-partition'], + }, schema=self.pk_pa_schema)) + write_three.write_arrow(pa.Table.from_pydict({ + 'user_id': [2], + 'item_id': [1002], + 'behavior': ['b'], + 'dt': ['new-partition'], + }, schema=self.pk_pa_schema)) + + messages_two = write_two.prepare_commit() + messages_three = write_three.prepare_commit() + self.assertEqual({2}, {m.total_buckets for m in messages_two}) + self.assertEqual({3}, {m.total_buckets for m in messages_three}) + + commit_two.commit(messages_two) + with self.assertRaisesRegex(CommitConflictError, "Total buckets"): + commit_three.commit(messages_three) + finally: + write_two.close() + write_three.close() + commit_two.close() + commit_three.close() def test_data_file_prefix_postpone(self): """Test that generated data file names follow the expected prefix format.""" schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], - options={'bucket': -2}) + options={'bucket': -2, 'postpone.batch-write-fixed-bucket': False}) self.catalog.create_table('default.test_file_prefix_postpone', schema, False) table = self.catalog.get_table('default.test_file_prefix_postpone') diff --git a/paimon-python/pypaimon/write/commit/conflict_detection.py b/paimon-python/pypaimon/write/commit/conflict_detection.py index 37992239593d..e89eadfc0c5d 100644 --- a/paimon-python/pypaimon/write/commit/conflict_detection.py +++ b/paimon-python/pypaimon/write/commit/conflict_detection.py @@ -232,6 +232,10 @@ def check_conflicts( "Trying to delete file {} which is not previously added.".format( entry.file.file_name)) + conflict = self.check_bucket_num_conflicts(merged_entries, commit_kind) + if conflict is not None: + return conflict + conflict = self.check_overwrite_from_snapshot( latest_snapshot, delta_entries, commit_kind) if conflict is not None: @@ -265,6 +269,27 @@ def check_conflicts( return self.check_row_id_from_snapshot(latest_snapshot, delta_entries) + @staticmethod + def check_bucket_num_conflicts(entries, commit_kind): + if commit_kind == "OVERWRITE": + return None + + total_buckets = {} + for entry in entries: + if entry.kind != 0 or entry.total_buckets <= 0: + continue + partition = tuple(entry.partition.values) + previous = total_buckets.get(partition) + if previous is not None and previous != entry.total_buckets: + return RuntimeError( + "Total buckets of partition {} changed from {} to {} " + "without overwrite. Give up committing.".format( + partition, previous, entry.total_buckets + ) + ) + total_buckets[partition] = entry.total_buckets + return None + def check_hash_index_conflicts( self, latest_snapshot, delta_index_entries=None): """Detect stale full-file replacements of dynamic-bucket HASH indexes.""" diff --git a/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py b/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py index d3718454e4bd..39460acc618d 100644 --- a/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py +++ b/paimon-python/pypaimon/write/commit/overwrite_changes_provider.py @@ -132,12 +132,17 @@ def _build_result(self, existing_entries: List[ManifestEntry]) -> List[ManifestE # New files being written by this overwrite. for msg in self.commit_messages: partition = GenericRow(list(msg.partition), self.table.partition_keys_fields) + total_buckets = ( + msg.total_buckets + if msg.total_buckets is not None + else self.table.total_buckets + ) for file in msg.new_files: entries.append(ManifestEntry( kind=0, partition=partition, bucket=msg.bucket, - total_buckets=self.table.total_buckets, + total_buckets=total_buckets, file=file, )) return entries diff --git a/paimon-python/pypaimon/write/commit_message.py b/paimon-python/pypaimon/write/commit_message.py index 3076b014b16d..c170d0d216a8 100644 --- a/paimon-python/pypaimon/write/commit_message.py +++ b/paimon-python/pypaimon/write/commit_message.py @@ -35,6 +35,7 @@ class CommitMessage: index_deletes: List['IndexManifestEntry'] = field(default_factory=list) changelog_files: List[DataFileMeta] = field(default_factory=list) hash_index_base_snapshot: Optional[int] = None + total_buckets: Optional[int] = None def is_empty(self): return ( diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 73f3d22429f5..a1640d5c205c 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -211,6 +211,10 @@ def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): if self.conflict_detection.has_hash_index_changes( index_adds + index_deletes): detect_conflicts = True + if any(message.total_buckets is not None + for message in commit_messages): + # Detect concurrent bucket-count changes in postpone APPENDs. + detect_conflicts = True self._try_commit(commit_kind=commit_kind, commit_identifier=commit_identifier, @@ -839,12 +843,17 @@ def _collect_changelog_entries(self, commit_messages: List[CommitMessage]) -> Li changelog_entries = [] for msg in commit_messages: partition = GenericRow(list(msg.partition), self.table.partition_keys_fields) + total_buckets = ( + msg.total_buckets + if msg.total_buckets is not None + else self.table.total_buckets + ) for file in msg.changelog_files: changelog_entries.append(ManifestEntry( kind=0, partition=partition, bucket=msg.bucket, - total_buckets=self.table.total_buckets, + total_buckets=total_buckets, file=file )) return changelog_entries @@ -853,12 +862,17 @@ def _collect_manifest_entries(self, commit_messages: List[CommitMessage]) -> Lis commit_entries = [] for msg in commit_messages: partition = GenericRow(list(msg.partition), self.table.partition_keys_fields) + total_buckets = ( + msg.total_buckets + if msg.total_buckets is not None + else self.table.total_buckets + ) for file in msg.new_files: commit_entries.append(ManifestEntry( kind=0, partition=partition, bucket=msg.bucket, - total_buckets=self.table.total_buckets, + total_buckets=total_buckets, file=file, )) for file in msg.deleted_files: @@ -866,7 +880,7 @@ def _collect_manifest_entries(self, commit_messages: List[CommitMessage]) -> Lis kind=1, partition=partition, bucket=msg.bucket, - total_buckets=self.table.total_buckets, + total_buckets=total_buckets, file=file, )) return commit_entries diff --git a/paimon-python/pypaimon/write/file_store_write.py b/paimon-python/pypaimon/write/file_store_write.py index 0cdd84664c4e..372e8aee58ab 100644 --- a/paimon-python/pypaimon/write/file_store_write.py +++ b/paimon-python/pypaimon/write/file_store_write.py @@ -39,7 +39,7 @@ class FileStoreWrite: """Base class for file store write operations.""" - def __init__(self, table, commit_user): + def __init__(self, table, commit_user, postpone_fixed_bucket=False): from pypaimon.table.file_store_table import FileStoreTable self.table: FileStoreTable = table @@ -50,7 +50,8 @@ def __init__(self, table, commit_user): self.commit_identifier = 0 self.options = CoreOptions.copy(table.options) self.changelog_producer = self.options.changelog_producer() - if self.table.bucket_mode() == BucketMode.POSTPONE_MODE: + if (self.table.bucket_mode() == BucketMode.POSTPONE_MODE + and not postpone_fixed_bucket): self.options.set(CoreOptions.DATA_FILE_PREFIX, (f"{self.options.data_file_prefix()}-u-{commit_user}" f"-s-{random.randint(0, 2 ** 31 - 2)}-w-")) diff --git a/paimon-python/pypaimon/write/postpone_batch_table_write.py b/paimon-python/pypaimon/write/postpone_batch_table_write.py new file mode 100644 index 000000000000..ccea7327caa0 --- /dev/null +++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py @@ -0,0 +1,165 @@ +# 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. + +from typing import Any, Dict, List, Optional + +import pyarrow as pa + +from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER +from pypaimon.write.commit_message import CommitMessage +from pypaimon.write.file_store_write import FileStoreWrite +from pypaimon.write.postpone_bucket import PostponeBucketPlanner +from pypaimon.write.row_key_extractor import PostponeFixedBucketRowKeyExtractor +from pypaimon.write.row_utils import ( + require_columns, + row_to_named_values, + row_values_to_arrow_table, +) +from pypaimon.write.table_write import BatchTableWrite + + +class PostponeFixedBucketBatchTableWrite(BatchTableWrite): + """Batch writer which plans postpone rows before routing new partitions.""" + + def __init__( + self, + table, + commit_user, + static_partition: Optional[dict] = None, + bucket_plan=None, + ): + self._planner = PostponeBucketPlanner( + table, + known_num_buckets=( + bucket_plan.as_dict() if bucket_plan is not None else None + ), + ) + self._bucket_plan = ( + bucket_plan + if bucket_plan is not None + else self._planner.current_plan() + ) + self._plan_provided = bucket_plan is not None + self._pending_inputs = [] + super().__init__(table, commit_user, static_partition) + + def _create_file_store_write(self, commit_user): + return FileStoreWrite( + self.table, commit_user, postpone_fixed_bucket=True) + + def _create_row_key_extractor(self, static_partition): + return PostponeFixedBucketRowKeyExtractor( + self.table, self._bucket_plan) + + def _buffer_input(self, data) -> bool: + if self._plan_provided: + return False + return any( + not self._bucket_plan.contains(partition) + for partition in self._planner.input_partition_stats(data) + ) + + def write_arrow(self, table: pa.Table): + if not self._buffer_input(table): + return super().write_arrow(table) + self._validate_pyarrow_schema(table.schema) + self._pending_inputs.extend( + ("batch", batch) for batch in table.to_batches()) + + def write_arrow_batch(self, data: pa.RecordBatch): + if not self._buffer_input(data): + return super().write_arrow_batch(data) + self._validate_pyarrow_schema(data.schema) + self._pending_inputs.append(("batch", data)) + + def write_row(self, row): + if self._plan_provided: + return super().write_row(row) + + values_by_name = row_to_named_values( + row, self.table.table_schema.fields) + column_names = ( + self.file_store_write.write_cols + if self.file_store_write.write_cols is not None + else list(self.table.field_names) + ) + require_columns(values_by_name, column_names, "write_row") + require_columns(values_by_name, self.table.partition_keys, "write_row") + partition = tuple( + values_by_name[key] for key in self.table.partition_keys) + if self._bucket_plan.contains(partition): + return super().write_row(row) + + arrow_row = row_values_to_arrow_table( + values_by_name, self.table.table_schema.fields, column_names) + self._pending_inputs.append( + ("row", (row, partition, arrow_row.nbytes))) + + def _flush_pending_inputs(self): + if not self._pending_inputs: + return + + partition_stats = {} + for input_type, value in self._pending_inputs: + if input_type == "batch": + stats_by_partition = ( + self._planner.input_partition_stats(value)) + else: + _, partition, size = value + stats_by_partition = {partition: (1, size)} + for partition, stats in stats_by_partition.items(): + rows, size = partition_stats.get(partition, (0, 0)) + partition_stats[partition] = ( + rows + stats[0], size + stats[1]) + + self._bucket_plan = self._planner.plan( + partition_stats, + include_postpone_rows=self.static_partition is None, + ) + self.row_key_extractor.with_bucket_plan(self._bucket_plan) + inputs = self._pending_inputs + self._pending_inputs = [] + for input_type, value in inputs: + if input_type == "batch": + super().write_arrow_batch(value) + else: + super().write_row(value[0]) + + def prepare_commit(self) -> List[CommitMessage]: + if self.batch_committed: + raise RuntimeError( + "BatchTableWrite only supports one-time committing.") + self.batch_committed = True + self._flush_pending_inputs() + return self._prepare_commit(BATCH_COMMIT_IDENTIFIER) + + def _prepare_commit(self, commit_identifier) -> List[CommitMessage]: + messages = super()._prepare_commit(commit_identifier) + for message in messages: + message.total_buckets = self.row_key_extractor.num_buckets( + message.partition) + return messages + + def _distributed_write_options(self) -> Dict[str, Any]: + return {"postpone_bucket_planner": self._planner} + + def close(self): + self._pending_inputs = [] + super().close() + + def abort(self): + self._pending_inputs = [] + super().abort() diff --git a/paimon-python/pypaimon/write/postpone_bucket.py b/paimon-python/pypaimon/write/postpone_bucket.py new file mode 100644 index 000000000000..7d7f3627467c --- /dev/null +++ b/paimon-python/pypaimon/write/postpone_bucket.py @@ -0,0 +1,189 @@ +# 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. + +from typing import Dict, Tuple + +import pyarrow as pa + +from pypaimon.table.bucket_mode import BucketMode + + +class PostponeBucketPlan: + """Resolved bucket counts by partition.""" + + def __init__(self, num_buckets: Dict[Tuple, int]): + self._num_buckets = dict(num_buckets) + + def contains(self, partition: Tuple) -> bool: + return tuple(partition) in self._num_buckets + + def num_buckets(self, partition: Tuple) -> int: + partition = tuple(partition) + if partition not in self._num_buckets: + raise ValueError("Missing bucket plan for partition {}".format(partition)) + return self._num_buckets[partition] + + def as_dict(self) -> Dict[Tuple, int]: + return dict(self._num_buckets) + + +class PostponeBucketPlanner: + """Plans fixed bucket counts for postpone batch writes.""" + + def __init__( + self, + table, + known_num_buckets=None, + postpone_row_counts=None, + ): + options = table.options + if options.bucket() != BucketMode.POSTPONE_BUCKET.value: + raise ValueError( + "Postpone fixed bucket writes require bucket = -2, got {}" + .format(options.bucket()) + ) + + self.max_num_buckets = ( + options.postpone_batch_write_fixed_bucket_max_parallelism() + ) + if self.max_num_buckets <= 0: + raise ValueError( + "postpone.batch-write-fixed-bucket.max-parallelism must be " + "positive, got {}".format(self.max_num_buckets) + ) + self.target_row_num_per_bucket = ( + options.postpone_target_row_num_per_bucket() + ) + if self.target_row_num_per_bucket is not None: + if self.target_row_num_per_bucket <= 0: + raise ValueError( + "postpone.target-row-num-per-bucket must be positive, " + "got {}".format(self.target_row_num_per_bucket) + ) + self.target_size_per_bucket = None + else: + self.target_size_per_bucket = ( + options.postpone_target_size_per_bucket() + ) + if self.target_size_per_bucket <= 0: + raise ValueError( + "postpone.target-size-per-bucket must be positive, got " + "{}".format(self.target_size_per_bucket) + ) + + self._partition_indices = [ + table.field_names.index(key) for key in table.partition_keys + ] + if known_num_buckets is None: + known_num_buckets, loaded_postpone_counts = ( + self._load_bucket_metadata(table) + ) + if postpone_row_counts is None: + postpone_row_counts = loaded_postpone_counts + self._known_num_buckets = dict(known_num_buckets) + self._postpone_row_counts = dict(postpone_row_counts or {}) + + @staticmethod + def _load_bucket_metadata(table): + scan = table.new_read_builder().new_scan().file_scanner + manifest_files, _ = scan.manifest_scanner() + entries = scan.manifest_file_manager.read_entries_parallel( + manifest_files, + max_workers=table.options.scan_manifest_parallelism(), + ) + + known = {} + postpone_counts = {} + for entry in entries: + partition = tuple(entry.partition.values) + if entry.bucket == BucketMode.POSTPONE_BUCKET.value: + postpone_counts[partition] = ( + postpone_counts.get(partition, 0) + entry.file.row_count + ) + elif entry.bucket >= 0 and entry.total_buckets > 0: + previous = known.get(partition) + if previous is not None and previous != entry.total_buckets: + raise RuntimeError( + "Partition {} has different total buckets {} and {}" + .format(partition, previous, entry.total_buckets) + ) + known[partition] = entry.total_buckets + return known, postpone_counts + + def current_plan(self) -> PostponeBucketPlan: + return PostponeBucketPlan(self._known_num_buckets) + + def input_partition_stats(self, data) -> Dict[Tuple, Tuple[int, int]]: + if self._partition_indices: + columns = [data.column(i) for i in self._partition_indices] + partitions = [ + tuple(column[row].as_py() for column in columns) + for row in range(data.num_rows) + ] + else: + partitions = [()] * data.num_rows + + row_indices = {} + for index, partition in enumerate(partitions): + row_indices.setdefault(partition, []).append(index) + + stats = {} + for partition, indices in row_indices.items(): + partition_data = ( + data + if len(indices) == data.num_rows + else pa.compute.take(data, pa.array(indices, type=pa.int64())) + ) + stats[partition] = (len(indices), partition_data.nbytes) + return stats + + def plan( + self, + partition_stats: Dict[Tuple, Tuple[int, int]], + include_postpone_rows: bool = True, + ) -> PostponeBucketPlan: + for partition, (row_count, data_size) in partition_stats.items(): + partition = tuple(partition) + if partition in self._known_num_buckets: + continue + postpone_rows = ( + self._postpone_row_counts.get(partition, 0) + if include_postpone_rows + else 0 + ) + if self.target_row_num_per_bucket is not None: + total_rows = row_count + postpone_rows + num_buckets = max( + 1, + (total_rows + self.target_row_num_per_bucket - 1) + // self.target_row_num_per_bucket, + ) + else: + estimated_size = data_size + if postpone_rows and row_count: + estimated_size = ( + data_size * (row_count + postpone_rows) + row_count - 1 + ) // row_count + num_buckets = max( + 1, + (estimated_size + self.target_size_per_bucket - 1) + // self.target_size_per_bucket, + ) + self._known_num_buckets[partition] = min( + num_buckets, self.max_num_buckets + ) + return self.current_plan() diff --git a/paimon-python/pypaimon/write/ray_datasink.py b/paimon-python/pypaimon/write/ray_datasink.py index 37c753b4ca79..3c34bd665170 100644 --- a/paimon-python/pypaimon/write/ray_datasink.py +++ b/paimon-python/pypaimon/write/ray_datasink.py @@ -73,10 +73,12 @@ def __init__( table: "Table", overwrite: bool = False, static_partition: Optional[Dict[str, Any]] = None, + postpone_bucket_plan=None, ): self.table = table self.overwrite = overwrite self.static_partition = static_partition + self._postpone_bucket_plan = postpone_bucket_plan self._table_name = table.identifier.get_full_name() self._writer_builder: Optional["WriteBuilder"] = None self._pending_commit_messages: List["CommitMessage"] = [] @@ -97,6 +99,8 @@ def __setstate__(self, state: dict) -> None: self._table_name = self.table.identifier.get_full_name() if not hasattr(self, 'static_partition'): self.static_partition = None + if not hasattr(self, '_postpone_bucket_plan'): + self._postpone_bucket_plan = None def on_write_start(self, schema=None) -> None: logger.info(f"Starting write job for table {self._table_name}") @@ -118,6 +122,8 @@ def write( if self._is_overwrite(): writer_builder = writer_builder.overwrite(self.static_partition) + if self._postpone_bucket_plan is not None: + writer_builder.with_bucket_plan(self._postpone_bucket_plan) table_write = writer_builder.new_write() table_schema = self.table.table_schema @@ -256,14 +262,63 @@ def write_paimon_dataset( concurrency: Optional[int] = None, ray_remote_args: Optional[Dict[str, Any]] = None, hash_fixed_precluster: str = "auto", + postpone_bucket_planner=None, ) -> 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, + HASH_FIXED_PRECLUSTER_MODES, maybe_apply_repartition, ) from pypaimon.table.bucket_mode import BucketMode + 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.POSTPONE_MODE + and table.options.postpone_batch_write_fixed_bucket() + ): + from pypaimon.write.postpone_bucket import ( + PostponeBucketPlanner, + ) + from pypaimon.write.row_key_extractor import ( + PostponeFixedBucketRowKeyExtractor, + ) + + planner = ( + postpone_bucket_planner + if postpone_bucket_planner is not None + else PostponeBucketPlanner(table) + ) + plan = planner.current_plan() + if table.partition_keys or not plan.contains(()): + dataset, partition_stats = _collect_partition_stats( + dataset, planner + ) + plan = planner.plan( + partition_stats, + include_postpone_rows=not ( + overwrite or static_partition is not None + ), + ) + _write_primary_key_groups( + dataset, + table, + overwrite=overwrite, + static_partition=static_partition, + concurrency=concurrency, + ray_remote_args=ray_remote_args, + bucket_extractor=PostponeFixedBucketRowKeyExtractor(table, plan), + postpone_bucket_plan=plan, + ) + return + if ( hash_fixed_precluster == HASH_FIXED_PRECLUSTER_MAP_GROUPS and table.bucket_mode() == BucketMode.HASH_FIXED @@ -291,6 +346,48 @@ def write_paimon_dataset( ) +def _collect_partition_stats(dataset, planner): + import pickle + + partition_col = "__paimon_partition__" + rows_col = "__paimon_rows__" + size_col = "__paimon_size__" + + def _stats(batch: pa.Table) -> pa.Table: + stats = planner.input_partition_stats(batch) + return pa.table({ + partition_col: pa.array( + [pickle.dumps(partition) for partition in stats], + type=pa.binary(), + ), + rows_col: pa.array( + [value[0] for value in stats.values()], type=pa.int64() + ), + size_col: pa.array( + [value[1] for value in stats.values()], type=pa.int64() + ), + }) + + materialized = dataset.materialize() + stats_dataset = materialized.map_batches( + _stats, batch_format="pyarrow", zero_copy_batch=True + ) + combined = {} + for batch in stats_dataset.iter_batches(batch_format="pyarrow"): + for partition, rows, size in zip( + batch.column(partition_col).to_pylist(), + batch.column(rows_col).to_pylist(), + batch.column(size_col).to_pylist(), + ): + key = pickle.loads(partition) + previous_rows, previous_size = combined.get(key, (0, 0)) + combined[key] = ( + previous_rows + rows, + previous_size + size, + ) + return materialized, combined + + def _write_primary_key_groups( dataset, table, @@ -299,6 +396,8 @@ def _write_primary_key_groups( static_partition: Optional[Dict[str, Any]], concurrency: Optional[int], ray_remote_args: Optional[Dict[str, Any]], + bucket_extractor=None, + postpone_bucket_plan=None, ) -> None: import inspect import pickle @@ -308,7 +407,9 @@ def _write_primary_key_groups( _group_by_partition_bucket, ) - grouped, bucket_col = _group_by_partition_bucket(dataset, table) + grouped, bucket_col = _group_by_partition_bucket( + dataset, table, extractor=bucket_extractor + ) message_col = "__paimon_commit_messages__" captured_table = table @@ -325,6 +426,7 @@ def _write_group(group: pa.Table) -> pa.Table: captured_table, overwrite=overwrite, static_partition=static_partition, + postpone_bucket_plan=postpone_bucket_plan, ) commit_messages = worker_sink.write([rows], None) return pa.table({ diff --git a/paimon-python/pypaimon/write/row_key_extractor.py b/paimon-python/pypaimon/write/row_key_extractor.py index 5436d5f140cc..713023874b27 100644 --- a/paimon-python/pypaimon/write/row_key_extractor.py +++ b/paimon-python/pypaimon/write/row_key_extractor.py @@ -577,7 +577,7 @@ def abort(self) -> None: class PostponeBucketRowKeyExtractor(RowKeyExtractor): - """Extractor for unaware bucket mode (bucket = -1, no primary keys).""" + """Extractor for postpone bucket mode which writes to bucket -2.""" def __init__(self, table_schema: TableSchema): super().__init__(table_schema) @@ -590,3 +590,59 @@ def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: return BucketMode.POSTPONE_BUCKET.value + + +class PostponeFixedBucketRowKeyExtractor(RowKeyExtractor): + """Route postpone batches using a resolved bucket plan.""" + + def __init__(self, table, bucket_plan): + super().__init__(table.table_schema) + if table.options.bucket() != BucketMode.POSTPONE_BUCKET.value: + raise ValueError( + "Postpone fixed bucket writes require bucket = -2, got {}".format( + table.options.bucket() + ) + ) + self.bucket_keys = table.table_schema.bucket_keys + self.bucket_key_indices = self._get_field_indices(self.bucket_keys) + self._bucket_key_fields = table.table_schema.logical_bucket_key_fields + self._bucket_plan = bucket_plan + + def with_bucket_plan(self, bucket_plan) -> None: + self._bucket_plan = bucket_plan + + def num_buckets(self, partition: Tuple) -> int: + return self._bucket_plan.num_buckets(partition) + + def extract_partition_bucket_batch( + self, data: pa.RecordBatch + ) -> Tuple[List[Tuple], List[int]]: + partitions = self._extract_partitions_batch(data) + columns = [data.column(i) for i in self.bucket_key_indices] + buckets = [ + _bucket_from_hash( + self._binary_row_hash_code( + tuple(col[row_idx].as_py() for col in columns), + self._bucket_key_fields, + ), + self.num_buckets(partition), + ) + for row_idx, partition in enumerate(partitions) + ] + return partitions, buckets + + def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: + return self.extract_partition_bucket_batch(data)[1] + + def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: + partition = tuple( + values_by_name[self.table_schema.fields[i].name] + for i in self.partition_indices + ) + return _bucket_from_hash( + self._binary_row_hash_code( + tuple(values_by_name[name] for name in self.bucket_keys), + self._bucket_key_fields, + ), + self.num_buckets(partition), + ) diff --git a/paimon-python/pypaimon/write/table_write.py b/paimon-python/pypaimon/write/table_write.py index c1276e72d2cc..bf23a69ecf99 100644 --- a/paimon-python/pypaimon/write/table_write.py +++ b/paimon-python/pypaimon/write/table_write.py @@ -23,7 +23,10 @@ from pypaimon.schema.data_types import PyarrowFieldParser from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER from pypaimon.table.row.blob import BlobConsumer -from pypaimon.write.row_utils import require_columns, row_to_named_values +from pypaimon.write.row_utils import ( + require_columns, + row_to_named_values, +) from pypaimon.write.commit_message import CommitMessage from pypaimon.write.file_store_write import FileStoreWrite @@ -37,14 +40,21 @@ def __init__(self, table, commit_user, static_partition: Optional[dict] = None): self.table: FileStoreTable = table self.table_pyarrow_schema = PyarrowFieldParser.from_paimon_schema(self.table.table_schema.fields) - self.file_store_write = FileStoreWrite(self.table, commit_user) - self.row_key_extractor = self.table.create_row_key_extractor( - ignore_existing=static_partition is not None - ) self.commit_user = commit_user self.static_partition = static_partition + self.file_store_write = self._create_file_store_write(commit_user) + self.row_key_extractor = self._create_row_key_extractor(static_partition) + + def _create_file_store_write(self, commit_user): + return FileStoreWrite(self.table, commit_user) + + def _create_row_key_extractor(self, static_partition): + return self.table.create_row_key_extractor( + ignore_existing=static_partition is not None + ) def write_arrow(self, table: pa.Table): + self._validate_pyarrow_schema(table.schema) batches_iterator = table.to_batches() for batch in batches_iterator: self.write_arrow_batch(batch) @@ -245,8 +255,12 @@ def write_ray( concurrency=concurrency, ray_remote_args=ray_remote_args, hash_fixed_precluster=hash_fixed_precluster, + **self._distributed_write_options(), ) + def _distributed_write_options(self) -> Dict[str, Any]: + return {} + def close(self): try: self.file_store_write.close() diff --git a/paimon-python/pypaimon/write/write_builder.py b/paimon-python/pypaimon/write/write_builder.py index f7a0459305de..51d689b724cf 100644 --- a/paimon-python/pypaimon/write/write_builder.py +++ b/paimon-python/pypaimon/write/write_builder.py @@ -58,7 +58,30 @@ def _create_commit_user(self): class BatchWriteBuilder(WriteBuilder): + def __init__(self, table): + super().__init__(table) + self.bucket_plan = None + + def with_bucket_plan(self, bucket_plan): + """Use a precomputed partition bucket plan.""" + self.bucket_plan = bucket_plan + return self + def new_write(self) -> BatchTableWrite: + from pypaimon.table.bucket_mode import BucketMode + if (self.table.bucket_mode() == BucketMode.POSTPONE_MODE + and self.table.options.postpone_batch_write_fixed_bucket()): + from pypaimon.write.postpone_batch_table_write import ( + PostponeFixedBucketBatchTableWrite, + ) + return PostponeFixedBucketBatchTableWrite( + self.table, + self.commit_user, + self.static_partition, + self.bucket_plan, + ) + if self.bucket_plan is not None: + raise ValueError("Bucket plans require a postpone-bucket table") return BatchTableWrite(self.table, self.commit_user, self.static_partition) def new_update(self) -> BatchTableUpdate: From 69aa464b148972fd786cd19942f2834638da90a9 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 3 Aug 2026 02:21:58 -0700 Subject: [PATCH 2/7] [python] Align postpone writer with fixed-bucket builder --- .../pypaimon/table/file_store_table.py | 9 +++ .../pypaimon/tests/write/table_write_test.py | 30 +++++++++ .../pypaimon/write/file_store_write.py | 63 +++++++++++++++++-- .../write/postpone_batch_table_write.py | 58 ++++++++++++++--- paimon-python/pypaimon/write/table_write.py | 18 +++++- paimon-python/pypaimon/write/write_builder.py | 23 ------- 6 files changed, 160 insertions(+), 41 deletions(-) diff --git a/paimon-python/pypaimon/table/file_store_table.py b/paimon-python/pypaimon/table/file_store_table.py index 4bc9b1c04859..e3bd4152d07b 100644 --- a/paimon-python/pypaimon/table/file_store_table.py +++ b/paimon-python/pypaimon/table/file_store_table.py @@ -426,8 +426,17 @@ def new_stream_read_builder(self) -> 'StreamReadBuilder': return StreamReadBuilder(self) def new_batch_write_builder(self) -> BatchWriteBuilder: + if (self.bucket_mode() == BucketMode.POSTPONE_MODE + and self.options.postpone_batch_write_fixed_bucket()): + return self.new_postpone_fixed_bucket_write_builder() return BatchWriteBuilder(self) + def new_postpone_fixed_bucket_write_builder(self): + from pypaimon.write.postpone_batch_table_write import ( + PostponeFixedBucketWriteBuilder, + ) + return PostponeFixedBucketWriteBuilder(self) + def new_stream_write_builder(self) -> StreamWriteBuilder: return StreamWriteBuilder(self) diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index e1ce3cd3cafd..1aeb3f56b12a 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -450,7 +450,10 @@ def test_postpone_read_write(self): table_write = write_builder.new_write() from pypaimon.write.postpone_batch_table_write import ( PostponeFixedBucketBatchTableWrite, + PostponeFixedBucketWriteBuilder, ) + self.assertIsInstance( + write_builder, PostponeFixedBucketWriteBuilder) self.assertIsInstance( table_write, PostponeFixedBucketBatchTableWrite) table_commit = write_builder.new_commit() @@ -478,6 +481,33 @@ def test_postpone_read_write(self): actual = table_read.to_arrow(splits) self.assertEqual(expect, actual) + def test_postpone_file_store_write_validates_runtime_bucket_count(self): + from pypaimon.write.file_store_write import ( + PostponeFixedBucketFileStoreWrite, + ) + + schema = Schema.from_pyarrow_schema( + self.pa_schema, + partition_keys=['user_id'], + primary_keys=['user_id', 'dt'], + options={'bucket': -2}, + ) + identifier = 'default.test_postpone_runtime_bucket_validation' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + write = PostponeFixedBucketFileStoreWrite(table, 'test-user') + try: + with self.assertRaisesRegex(ValueError, 'must be positive'): + write.write((1,), 0, self.pk_expected.to_batches()[0], 0) + with self.assertRaisesRegex(ValueError, 'out of range'): + write.write((1,), 2, self.pk_expected.to_batches()[0], 2) + + write._check_runtime_bucket((1,), 0, 2) + with self.assertRaisesRegex(RuntimeError, 'new bucket num 3'): + write._check_runtime_bucket((1,), 0, 3) + finally: + write.abort() + def test_postpone_batch_fixed_bucket_disabled(self): schema = Schema.from_pyarrow_schema( self.pa_schema, diff --git a/paimon-python/pypaimon/write/file_store_write.py b/paimon-python/pypaimon/write/file_store_write.py index 372e8aee58ab..9a1dd10721a8 100644 --- a/paimon-python/pypaimon/write/file_store_write.py +++ b/paimon-python/pypaimon/write/file_store_write.py @@ -39,19 +39,22 @@ class FileStoreWrite: """Base class for file store write operations.""" - def __init__(self, table, commit_user, postpone_fixed_bucket=False): + def __init__(self, table, commit_user): from pypaimon.table.file_store_table import FileStoreTable self.table: FileStoreTable = table self.data_writers: Dict[Tuple, DataWriter] = {} + self._runtime_total_buckets: Dict[Tuple, int] = {} self.max_seq_numbers: dict = {} self.write_cols = None self.blob_consumer = None self.commit_identifier = 0 self.options = CoreOptions.copy(table.options) self.changelog_producer = self.options.changelog_producer() - if (self.table.bucket_mode() == BucketMode.POSTPONE_MODE - and not postpone_fixed_bucket): + self._configure_data_file_prefix(commit_user) + + def _configure_data_file_prefix(self, commit_user): + if self.table.bucket_mode() == BucketMode.POSTPONE_MODE: self.options.set(CoreOptions.DATA_FILE_PREFIX, (f"{self.options.data_file_prefix()}-u-{commit_user}" f"-s-{random.randint(0, 2 ** 31 - 2)}-w-")) @@ -64,14 +67,29 @@ def disable_rolling(self): self.options.set( CoreOptions.TARGET_FILE_ROW_NUM, str(max_value)) - def write(self, partition: Tuple, bucket: int, data: pa.RecordBatch): + def write( + self, + partition: Tuple, + bucket: int, + data: pa.RecordBatch, + total_buckets=None, + ): + self._check_runtime_bucket(partition, bucket, total_buckets) key = (partition, bucket) if key not in self.data_writers: self.data_writers[key] = self._create_data_writer(partition, bucket, self.options) writer = self.data_writers[key] writer.write(data) - def write_row(self, partition: Tuple, bucket: int, row, values_by_name: dict): + def write_row( + self, + partition: Tuple, + bucket: int, + row, + values_by_name: dict, + total_buckets=None, + ): + self._check_runtime_bucket(partition, bucket, total_buckets) key = (partition, bucket) if key not in self.data_writers: self.data_writers[key] = self._create_data_writer(partition, bucket, self.options) @@ -92,6 +110,31 @@ def write_row(self, partition: Tuple, bucket: int, row, values_by_name: dict): ) writer.write(data.to_batches()[0]) + def _check_runtime_bucket(self, partition, bucket, total_buckets): + if total_buckets is None: + return + if (isinstance(total_buckets, bool) + or not isinstance(total_buckets, int) + or total_buckets <= 0): + raise ValueError("Total number of buckets must be positive") + if bucket < 0 or bucket >= total_buckets: + raise ValueError( + "Bucket {} is out of range [0, {})".format( + bucket, total_buckets + ) + ) + + partition = tuple(partition) + previous = self._runtime_total_buckets.get(partition) + if previous is not None and previous != total_buckets: + raise RuntimeError( + "Try to write partition {} with a new bucket num {}, but " + "the previous bucket num is {}.".format( + partition, total_buckets, previous + ) + ) + self._runtime_total_buckets[partition] = total_buckets + def _create_data_writer(self, partition: Tuple, bucket: int, options: CoreOptions) -> DataWriter: row_limit = options.target_file_row_num() max_value = CoreOptions.TARGET_FILE_ROW_NUM.default_value() @@ -287,6 +330,7 @@ def prepare_commit(self, commit_identifier) -> List[CommitMessage]: bucket=bucket, new_files=committed_files, changelog_files=changelog_files, + total_buckets=self._runtime_total_buckets.get(partition), ) commit_messages.append(commit_message) return commit_messages @@ -296,6 +340,7 @@ def close(self): for writer in self.data_writers.values(): writer.close() self.data_writers.clear() + self._runtime_total_buckets.clear() def abort(self): """Abort all data writers and clean up files produced by this write.""" @@ -305,6 +350,7 @@ def abort(self): except Exception as e: logger.warning("Failed to abort data writer.", exc_info=e) self.data_writers.clear() + self._runtime_total_buckets.clear() def _seq_number_stats(self, partition: Tuple) -> Dict[int, int]: buckets = self.max_seq_numbers.get(partition) @@ -331,3 +377,10 @@ def _load_seq_number_stats(self, partition: Tuple) -> dict: if current_seq_num > existing_max: max_seq_numbers[split.bucket] = current_seq_num return max_seq_numbers + + +class PostponeFixedBucketFileStoreWrite(FileStoreWrite): + """File store write with runtime bucket counts for postpone tables.""" + + def _configure_data_file_prefix(self, commit_user): + pass diff --git a/paimon-python/pypaimon/write/postpone_batch_table_write.py b/paimon-python/pypaimon/write/postpone_batch_table_write.py index ccea7327caa0..e0ce213c789b 100644 --- a/paimon-python/pypaimon/write/postpone_batch_table_write.py +++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py @@ -19,8 +19,9 @@ import pyarrow as pa from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER +from pypaimon.table.bucket_mode import BucketMode from pypaimon.write.commit_message import CommitMessage -from pypaimon.write.file_store_write import FileStoreWrite +from pypaimon.write.file_store_write import PostponeFixedBucketFileStoreWrite from pypaimon.write.postpone_bucket import PostponeBucketPlanner from pypaimon.write.row_key_extractor import PostponeFixedBucketRowKeyExtractor from pypaimon.write.row_utils import ( @@ -29,6 +30,32 @@ row_values_to_arrow_table, ) from pypaimon.write.table_write import BatchTableWrite +from pypaimon.write.write_builder import BatchWriteBuilder + + +class PostponeFixedBucketWriteBuilder(BatchWriteBuilder): + """Write builder for fixed-bucket batches on a postpone table.""" + + def __init__(self, table): + if table.bucket_mode() != BucketMode.POSTPONE_MODE: + raise ValueError( + "Postpone fixed-bucket write requires a postpone-bucket table" + ) + super().__init__(table) + self._bucket_plan = None + + def with_bucket_plan(self, bucket_plan): + """Use a precomputed partition bucket plan.""" + self._bucket_plan = bucket_plan + return self + + def new_write(self): + return PostponeFixedBucketBatchTableWrite( + self.table, + self.commit_user, + self.static_partition, + self._bucket_plan, + ) class PostponeFixedBucketBatchTableWrite(BatchTableWrite): @@ -57,13 +84,31 @@ def __init__( super().__init__(table, commit_user, static_partition) def _create_file_store_write(self, commit_user): - return FileStoreWrite( - self.table, commit_user, postpone_fixed_bucket=True) + return PostponeFixedBucketFileStoreWrite(self.table, commit_user) def _create_row_key_extractor(self, static_partition): return PostponeFixedBucketRowKeyExtractor( self.table, self._bucket_plan) + def _write_partition_bucket_batch(self, partition, bucket, data): + self.file_store_write.write( + partition, + bucket, + data, + self.row_key_extractor.num_buckets(partition), + ) + + def _write_partition_bucket_row( + self, partition, bucket, row, values_by_name + ): + self.file_store_write.write_row( + partition, + bucket, + row, + values_by_name, + self.row_key_extractor.num_buckets(partition), + ) + def _buffer_input(self, data) -> bool: if self._plan_provided: return False @@ -146,13 +191,6 @@ def prepare_commit(self) -> List[CommitMessage]: self._flush_pending_inputs() return self._prepare_commit(BATCH_COMMIT_IDENTIFIER) - def _prepare_commit(self, commit_identifier) -> List[CommitMessage]: - messages = super()._prepare_commit(commit_identifier) - for message in messages: - message.total_buckets = self.row_key_extractor.num_buckets( - message.partition) - return messages - def _distributed_write_options(self) -> Dict[str, Any]: return {"postpone_bucket_planner": self._planner} diff --git a/paimon-python/pypaimon/write/table_write.py b/paimon-python/pypaimon/write/table_write.py index bf23a69ecf99..469325a4d408 100644 --- a/paimon-python/pypaimon/write/table_write.py +++ b/paimon-python/pypaimon/write/table_write.py @@ -80,7 +80,10 @@ def write_arrow_batch(self, data: pa.RecordBatch): else: indices_array = pa.array(row_indices, type=pa.int64()) sub_table = pa.compute.take(data, indices_array) - self.file_store_write.write(partition, bucket, sub_table) + self._write_partition_bucket_batch(partition, bucket, sub_table) + + def _write_partition_bucket_batch(self, partition, bucket, data): + self.file_store_write.write(partition, bucket, data) def with_dynamic_bucket_index( self, @@ -166,7 +169,7 @@ def write_arrow_batch_to_bucket( ) if partition is None: return - self.file_store_write.write(partition, bucket, data) + self._write_partition_bucket_batch(partition, bucket, data) def write_row(self, row): values_by_name = row_to_named_values(row, self.table.table_schema.fields) @@ -180,7 +183,16 @@ def write_row(self, row): partition, bucket = ( self.row_key_extractor.extract_partition_bucket_row(values_by_name) ) - self.file_store_write.write_row(partition, bucket, row, values_by_name) + self._write_partition_bucket_row( + partition, bucket, row, values_by_name + ) + + def _write_partition_bucket_row( + self, partition, bucket, row, values_by_name + ): + self.file_store_write.write_row( + partition, bucket, row, values_by_name + ) def write_pandas(self, dataframe): write_cols = self.file_store_write.write_cols diff --git a/paimon-python/pypaimon/write/write_builder.py b/paimon-python/pypaimon/write/write_builder.py index 51d689b724cf..f7a0459305de 100644 --- a/paimon-python/pypaimon/write/write_builder.py +++ b/paimon-python/pypaimon/write/write_builder.py @@ -58,30 +58,7 @@ def _create_commit_user(self): class BatchWriteBuilder(WriteBuilder): - def __init__(self, table): - super().__init__(table) - self.bucket_plan = None - - def with_bucket_plan(self, bucket_plan): - """Use a precomputed partition bucket plan.""" - self.bucket_plan = bucket_plan - return self - def new_write(self) -> BatchTableWrite: - from pypaimon.table.bucket_mode import BucketMode - if (self.table.bucket_mode() == BucketMode.POSTPONE_MODE - and self.table.options.postpone_batch_write_fixed_bucket()): - from pypaimon.write.postpone_batch_table_write import ( - PostponeFixedBucketBatchTableWrite, - ) - return PostponeFixedBucketBatchTableWrite( - self.table, - self.commit_user, - self.static_partition, - self.bucket_plan, - ) - if self.bucket_plan is not None: - raise ValueError("Bucket plans require a postpone-bucket table") return BatchTableWrite(self.table, self.commit_user, self.static_partition) def new_update(self) -> BatchTableUpdate: From acf7e4ef67ca5cbb3cf44847ad9ae66ada6e9acf Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 3 Aug 2026 03:43:13 -0700 Subject: [PATCH 3/7] [python] Keep postpone fixed-bucket writes explicit --- docs/docs/pypaimon/python-api.mdx | 6 +-- paimon-python/pypaimon/daft/daft_datasink.py | 8 +--- .../pypaimon/table/file_store_table.py | 3 -- .../pypaimon/tests/write/table_write_test.py | 44 +++++++++---------- paimon-python/pypaimon/write/ray_datasink.py | 8 +++- 5 files changed, 32 insertions(+), 37 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index d6d12e6d2d01..c7b2aa153e4b 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -321,9 +321,9 @@ table_write.close() table_commit.close() ``` -For postpone-bucket tables (`bucket = -2`) with fixed-bucket batch writes -enabled, inputs for new partitions are buffered until `prepare_commit()`. -Memory grows with buffered input; existing-bucket partitions remain incremental. +`new_batch_write_builder()` keeps postpone-bucket writes in `bucket-postpone`. +Use `new_postpone_fixed_bucket_write_builder()` to write real buckets. New +partitions are buffered until `prepare_commit()`; existing ones are incremental. By default, the data will be appended to table. If you want to overwrite table, you should use `TableWrite#overwrite` API: diff --git a/paimon-python/pypaimon/daft/daft_datasink.py b/paimon-python/pypaimon/daft/daft_datasink.py index 6065bcc19e57..c21767b0d2e9 100644 --- a/paimon-python/pypaimon/daft/daft_datasink.py +++ b/paimon-python/pypaimon/daft/daft_datasink.py @@ -195,15 +195,9 @@ def _init_table(self, table: FileStoreTable) -> None: self._table = table from pypaimon.schema.data_types import PyarrowFieldParser - from pypaimon.table.bucket_mode import BucketMode self._target_schema: pa.Schema = PyarrowFieldParser.from_paimon_schema(table.fields) - write_table = table - if table.bucket_mode() == BucketMode.POSTPONE_MODE: - write_table = table.copy({ - "postpone.batch-write-fixed-bucket": "false", - }) - self._write_builder = write_table.new_batch_write_builder() + self._write_builder = table.new_batch_write_builder() if ( self._commit_user is not None and hasattr(self._write_builder, "commit_user") diff --git a/paimon-python/pypaimon/table/file_store_table.py b/paimon-python/pypaimon/table/file_store_table.py index e3bd4152d07b..ac88bf5c12dc 100644 --- a/paimon-python/pypaimon/table/file_store_table.py +++ b/paimon-python/pypaimon/table/file_store_table.py @@ -426,9 +426,6 @@ def new_stream_read_builder(self) -> 'StreamReadBuilder': return StreamReadBuilder(self) def new_batch_write_builder(self) -> BatchWriteBuilder: - if (self.bucket_mode() == BucketMode.POSTPONE_MODE - and self.options.postpone_batch_write_fixed_bucket()): - return self.new_postpone_fixed_bucket_write_builder() return BatchWriteBuilder(self) def new_postpone_fixed_bucket_write_builder(self): diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index 1aeb3f56b12a..a24e6e1f262f 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -446,7 +446,7 @@ def test_postpone_read_write(self): } expect = pa.Table.from_pydict(data, schema=self.pk_pa_schema) - write_builder = table.new_batch_write_builder() + write_builder = table.new_postpone_fixed_bucket_write_builder() table_write = write_builder.new_write() from pypaimon.write.postpone_batch_table_write import ( PostponeFixedBucketBatchTableWrite, @@ -508,18 +508,16 @@ def test_postpone_file_store_write_validates_runtime_bucket_count(self): finally: write.abort() - def test_postpone_batch_fixed_bucket_disabled(self): + def test_postpone_batch_write_builder_keeps_postpone_mode(self): schema = Schema.from_pyarrow_schema( self.pa_schema, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], - options={ - 'bucket': -2, - 'postpone.batch-write-fixed-bucket': False, - }, + options={'bucket': -2}, ) - self.catalog.create_table('default.test_postpone_disabled', schema, False) - table = self.catalog.get_table('default.test_postpone_disabled') + identifier = 'default.test_postpone_default_builder' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) expected = pa.Table.from_pydict({ 'user_id': [1], 'item_id': [1001], @@ -539,7 +537,7 @@ def test_postpone_batch_fixed_bucket_disabled(self): 1, len(glob.glob( self.warehouse - + "/default.db/test_postpone_disabled/user_id=1/" + + "/default.db/test_postpone_default_builder/user_id=1/" + "bucket-postpone/*.avro" )), ) @@ -571,7 +569,7 @@ def test_postpone_batch_infers_bucket_num_from_input_size(self): 'value': ['x', 'x' * 500], }, schema=pa_schema) - builder = table.new_batch_write_builder() + builder = table.new_postpone_fixed_bucket_write_builder() write = builder.new_write() commit = builder.new_commit() try: @@ -612,7 +610,7 @@ def test_postpone_batch_plans_all_record_batches(self): self.catalog.create_table(identifier, schema, False) table = self.catalog.get_table(identifier) - builder = table.new_batch_write_builder() + builder = table.new_postpone_fixed_bucket_write_builder() write = builder.new_write() commit = builder.new_commit() try: @@ -658,7 +656,7 @@ def test_postpone_batch_prefers_target_row_num(self): self.catalog.create_table(identifier, schema, False) table = self.catalog.get_table(identifier) - builder = table.new_batch_write_builder() + builder = table.new_postpone_fixed_bucket_write_builder() write = builder.new_write() commit = builder.new_commit() try: @@ -697,7 +695,7 @@ def test_postpone_batch_plans_write_rows_with_arrow_input(self): self.catalog.create_table(identifier, schema, False) table = self.catalog.get_table(identifier) - builder = table.new_batch_write_builder() + builder = table.new_postpone_fixed_bucket_write_builder() write = builder.new_write() commit = builder.new_commit() try: @@ -783,8 +781,8 @@ def test_postpone_worker_bucket_plan_mismatch_fails_commit(self): identifier = 'default.test_postpone_worker_plan_mismatch' self.catalog.create_table(identifier, schema, False) table = self.catalog.get_table(identifier) - small_builder = table.new_batch_write_builder() - large_builder = table.new_batch_write_builder() + small_builder = table.new_postpone_fixed_bucket_write_builder() + large_builder = table.new_postpone_fixed_bucket_write_builder() small_write = small_builder.new_write() large_write = large_builder.new_write() commit = small_builder.new_commit() @@ -826,7 +824,7 @@ def test_postpone_batch_fixed_bucket_reuses_existing_bucket_num(self): 'dt': ['p1'], }, schema=self.pk_pa_schema) - write_builder = table.new_batch_write_builder() + write_builder = table.new_postpone_fixed_bucket_write_builder() table_write = write_builder.new_write() table_commit = write_builder.new_commit() table_write.write_arrow(expected) @@ -842,7 +840,7 @@ def test_postpone_batch_fixed_bucket_reuses_existing_bucket_num(self): planner = PostponeBucketPlanner(copied_table) plan = planner.plan({('p2',): (1, 10)}) copied_write = ( - copied_table.new_batch_write_builder() + copied_table.new_postpone_fixed_bucket_write_builder() .with_bucket_plan(plan) .new_write() ) @@ -872,7 +870,7 @@ def test_postpone_reuses_bucket_num_for_int_date_partition(self): table = self.catalog.get_table(identifier) day = datetime.date(2026, 8, 1) - first_builder = table.new_batch_write_builder() + first_builder = table.new_postpone_fixed_bucket_write_builder() first_write = first_builder.new_write() first_commit = first_builder.new_commit() first_write.write_arrow(pa.Table.from_pydict({ @@ -889,7 +887,7 @@ def test_postpone_reuses_bucket_num_for_int_date_partition(self): reopened = self.catalog.get_table(identifier).copy({ 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }) - second_builder = reopened.new_batch_write_builder() + second_builder = reopened.new_postpone_fixed_bucket_write_builder() second_write = second_builder.new_write() second_commit = second_builder.new_commit() try: @@ -945,7 +943,7 @@ def test_postpone_legacy_partition_migrates_to_fixed_bucket(self): 'postpone.target-size-per-bucket': '1 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 2, }) - fixed_builder = fixed.new_batch_write_builder() + fixed_builder = fixed.new_postpone_fixed_bucket_write_builder() fixed_write = fixed_builder.new_write() fixed_commit = fixed_builder.new_commit() try: @@ -1001,8 +999,10 @@ def test_postpone_concurrent_new_partition_bucket_num_conflict(self): 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }) - builder_two = table_two_buckets.new_batch_write_builder() - builder_three = table_three_buckets.new_batch_write_builder() + builder_two = ( + table_two_buckets.new_postpone_fixed_bucket_write_builder()) + builder_three = ( + table_three_buckets.new_postpone_fixed_bucket_write_builder()) write_two = builder_two.new_write() write_three = builder_three.new_write() commit_two = builder_two.new_commit() diff --git a/paimon-python/pypaimon/write/ray_datasink.py b/paimon-python/pypaimon/write/ray_datasink.py index 3c34bd665170..7e7c8c1f7f95 100644 --- a/paimon-python/pypaimon/write/ray_datasink.py +++ b/paimon-python/pypaimon/write/ray_datasink.py @@ -118,10 +118,14 @@ def write( table_write = None try: - writer_builder = self.table.new_batch_write_builder() + writer_builder = ( + self.table.new_postpone_fixed_bucket_write_builder() + if self._postpone_bucket_plan is not None + else self.table.new_batch_write_builder() + ) if self._is_overwrite(): writer_builder = writer_builder.overwrite(self.static_partition) - + if self._postpone_bucket_plan is not None: writer_builder.with_bucket_plan(self._postpone_bucket_plan) table_write = writer_builder.new_write() From 07e51908dff5e40c48fbc35c5a4ea1a70dfb79bc Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 3 Aug 2026 04:59:52 -0700 Subject: [PATCH 4/7] [python] Align postpone planning and commit cleanup --- docs/docs/pypaimon/python-api.mdx | 1 + docs/docs/pypaimon/ray-data.md | 5 +- .../pypaimon/tests/file_store_commit_test.py | 29 +++ .../tests/write/dynamic_bucket_test.py | 7 +- .../pypaimon/tests/write/table_write_test.py | 170 +++++++++++++++++- .../pypaimon/write/file_store_commit.py | 8 +- .../write/postpone_batch_table_write.py | 3 +- .../pypaimon/write/postpone_bucket.py | 45 +++-- .../pypaimon/write/row_key_extractor.py | 9 + paimon-python/pypaimon/write/table_commit.py | 5 +- 10 files changed, 252 insertions(+), 30 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index c7b2aa153e4b..1921a9764974 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -324,6 +324,7 @@ table_commit.close() `new_batch_write_builder()` keeps postpone-bucket writes in `bucket-postpone`. Use `new_postpone_fixed_bucket_write_builder()` to write real buckets. New partitions are buffered until `prepare_commit()`; existing ones are incremental. +This builder currently supports `bucket-function.type=default` only. By default, the data will be appended to table. If you want to overwrite table, you should use `TableWrite#overwrite` API: diff --git a/docs/docs/pypaimon/ray-data.md b/docs/docs/pypaimon/ray-data.md index 9f2b9d353e76..9b08d6014f21 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -252,8 +252,9 @@ By default, postpone-bucket primary-key tables (`bucket = -2`) are written to real buckets and are immediately visible. Existing partitions reuse their bucket count; new partitions infer one from the configured target row count or size. Ray groups rows by `(partition, bucket)`, so each group must -fit in one worker's memory. Set `postpone.batch-write-fixed-bucket` to `false` -to write to `bucket-postpone` instead. HASH_DYNAMIC and +fit in one worker's memory. This path supports `bucket-function.type=default` +only. Set `postpone.batch-write-fixed-bucket` to `false` to write to +`bucket-postpone` instead. HASH_DYNAMIC and CROSS_PARTITION primary-key Ray writes are not supported and fail fast, including the default dynamic-bucket primary-key table (`bucket = -1`). Ray write tasks create independent Paimon writers, which can assign diff --git a/paimon-python/pypaimon/tests/file_store_commit_test.py b/paimon-python/pypaimon/tests/file_store_commit_test.py index 88350bfd5f01..f8397041144a 100644 --- a/paimon-python/pypaimon/tests/file_store_commit_test.py +++ b/paimon-python/pypaimon/tests/file_store_commit_test.py @@ -175,6 +175,35 @@ def test_generate_partition_statistics_multiple_files_same_partition( expected_time = file_meta_2.creation_time_epoch_millis() self.assertEqual(stat.last_file_creation_time, expected_time) + def test_partition_statistics_use_replacement_bucket_count( + self, mock_manifest_list_manager, mock_manifest_file_manager): + file_store_commit = self._create_file_store_commit() + file_meta = Mock(row_count=1, file_size=10, creation_time=None) + + for old_buckets, new_buckets in [(-2, 2), (2, 3)]: + with self.subTest(old=old_buckets, new=new_buckets): + partition = GenericRow(['2024-01-15', 'us-east-1'], None) + entries = [ + ManifestEntry( + kind=1, + partition=partition, + bucket=0, + total_buckets=old_buckets, + file=file_meta, + ), + ManifestEntry( + kind=0, + partition=partition, + bucket=0, + total_buckets=new_buckets, + file=file_meta, + ), + ] + + statistics = ( + file_store_commit._generate_partition_statistics(entries)) + self.assertEqual(new_buckets, statistics[0].total_buckets) + def test_generate_partition_statistics_multiple_partitions( self, mock_manifest_list_manager, mock_manifest_file_manager): """Test partition statistics generation with multiple different partitions.""" diff --git a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py index 61fece493462..f7714f5e9994 100644 --- a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py +++ b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py @@ -37,6 +37,7 @@ from pypaimon.manifest.index_manifest_entry import IndexManifestEntry from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.table.row.generic_row import GenericRow +from pypaimon.write.commit.conflict_detection import CommitConflictError from pypaimon.write.row_key_extractor import DynamicBucketRowKeyExtractor @@ -602,7 +603,7 @@ def test_data_only_upsert_conflicts_after_overwrite_remaps_key(self): stale_writer.close() stale_commit.close() - def test_retry_then_hash_index_conflict_preserves_prepared_files(self): + def test_retry_then_hash_index_conflict_aborts_prepared_files(self): with tempfile.TemporaryDirectory() as root: table = self._create_table(root, 'retry_hash_conflict') writer, commit, messages = self._prepare_indexed_write(table, [1]) @@ -642,13 +643,13 @@ def lose_first_compare_and_set( '_commit_retry_wait', ): with self.assertRaisesRegex( - RuntimeError, 'HASH index assignment conflict' + CommitConflictError, 'HASH index assignment conflict' ): commit.commit(messages) self.assertEqual(1, calls) self.assertTrue(all( - table.file_io.exists(path) for path in prepared_paths + not table.file_io.exists(path) for path in prepared_paths )) writer.close() commit.close() diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index a24e6e1f262f..78788866fda5 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -590,6 +590,90 @@ def test_postpone_batch_infers_bucket_num_from_input_size(self): [1, 2], self._read_sorted(table, 'id').column('id').to_pylist() ) + def test_postpone_size_inference_matches_java_binary_row(self): + from pypaimon.write.postpone_bucket import PostponeBucketPlanner + + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('key', pa.string(), nullable=False), + pa.field('value', pa.string(), nullable=False), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + primary_keys=['id'], + options={ + 'bucket': -2, + 'postpone.target-size-per-bucket': '20 kb', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + identifier = 'default.test_postpone_java_size_fixture' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + data = pa.RecordBatch.from_pydict({ + 'id': list(range(1000)), + 'key': ['k'] * 1000, + 'value': ['v'] * 1000, + }, schema=pa_schema) + + planner = PostponeBucketPlanner( + table, known_num_buckets={}, postpone_row_counts={}) + stats = planner.input_partition_stats(data) + self.assertEqual((1000, 32000), stats[()]) + self.assertEqual(2, planner.plan(stats).num_buckets(())) + + def test_postpone_default_bucket_function_matches_java(self): + from pypaimon.write.postpone_bucket import PostponeBucketPlan + from pypaimon.write.row_key_extractor import ( + PostponeFixedBucketRowKeyExtractor, + ) + + pa_schema = pa.schema([ + pa.field('key', pa.string(), nullable=False), + pa.field('value', pa.string(), nullable=False), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + primary_keys=['key'], + options={'bucket': -2}, + ) + identifier = 'default.test_postpone_java_bucket_fixture' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + extractor = PostponeFixedBucketRowKeyExtractor( + table, PostponeBucketPlan({(): 4})) + data = pa.RecordBatch.from_pydict({ + 'key': ['hello-java'], + 'value': ['v'], + }, schema=pa_schema) + + # Java BinaryRow hash -201703277 maps to bucket 1 of 4. + self.assertEqual([1], extractor.extract_partition_bucket_batch(data)[1]) + + @parameterized.expand([('mod',), ('hive',)]) + def test_postpone_rejects_unsupported_bucket_function( + self, bucket_function): + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('value', pa.string(), nullable=False), + ]) + schema = Schema.from_pyarrow_schema( + pa_schema, + primary_keys=['id'], + options={ + 'bucket': -2, + 'bucket-function.type': bucket_function, + }, + ) + identifier = ( + 'default.test_postpone_bucket_function_' + bucket_function) + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + + with self.assertRaisesRegex( + ValueError, 'only support bucket-function.type=default'): + table.new_postpone_fixed_bucket_write_builder().new_write() + def test_postpone_batch_plans_all_record_batches(self): pa_schema = pa.schema([ pa.field('id', pa.int32(), nullable=False), @@ -979,6 +1063,66 @@ def test_postpone_legacy_partition_migrates_to_fixed_bucket(self): actual = self._read_sorted(fixed, 'user_id') self.assertEqual([2], actual.column('user_id').to_pylist()) + def test_postpone_overwrite_updates_catalog_bucket_count(self): + schema = Schema.from_pyarrow_schema( + self.pk_pa_schema, + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={'bucket': -2}, + ) + identifier = 'default.test_postpone_overwrite_bucket_statistics' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + + legacy_builder = table.new_batch_write_builder() + legacy_write = legacy_builder.new_write() + legacy_commit = legacy_builder.new_commit() + legacy_write.write_arrow(pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['legacy'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema)) + legacy_commit.commit(legacy_write.prepare_commit()) + legacy_write.close() + legacy_commit.close() + + fixed = table.copy({ + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }) + builder = ( + fixed.new_postpone_fixed_bucket_write_builder() + .overwrite({'dt': 'p1'}) + ) + write = builder.new_write() + commit = builder.new_commit() + write.write_arrow(pa.Table.from_pydict({ + 'user_id': [2], + 'item_id': [1002], + 'behavior': ['fixed'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema)) + messages = write.prepare_commit() + captured_statistics = [] + real_commit = commit.file_store_commit.snapshot_commit.commit + + def capture_statistics(base_snapshot_uuid, snapshot, statistics): + captured_statistics.extend(statistics) + return real_commit(base_snapshot_uuid, snapshot, statistics) + + commit.file_store_commit.snapshot_commit.commit = capture_statistics + try: + commit.commit(messages) + finally: + write.close() + commit.close() + + self.assertEqual(2, captured_statistics[0].total_buckets) + self.assertEqual( + [2], self._read_sorted(fixed, 'user_id').column('user_id').to_pylist() + ) + def test_postpone_concurrent_new_partition_bucket_num_conflict(self): from pypaimon.write.commit.conflict_detection import CommitConflictError @@ -1026,9 +1170,33 @@ def test_postpone_concurrent_new_partition_bucket_num_conflict(self): self.assertEqual({2}, {m.total_buckets for m in messages_two}) self.assertEqual({3}, {m.total_buckets for m in messages_three}) - commit_two.commit(messages_two) + losing_paths = [ + file.external_path or file.file_path + for message in messages_three + for file in message.new_files + ] + self.assertTrue(all( + table_three_buckets.file_io.exists(path) + for path in losing_paths + )) + concurrent_commit = {'done': False} + + def fail_cas_after_concurrent_commit(*_): + if not concurrent_commit['done']: + concurrent_commit['done'] = True + commit_two.commit(messages_two) + return False + raise AssertionError('Bucket conflict should precede another CAS') + + commit_three.file_store_commit.snapshot_commit.commit = ( + fail_cas_after_concurrent_commit) with self.assertRaisesRegex(CommitConflictError, "Total buckets"): commit_three.commit(messages_three) + self.assertTrue(concurrent_commit['done']) + self.assertTrue(all( + not table_three_buckets.file_io.exists(path) + for path in losing_paths + )) finally: write_two.close() write_three.close() diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index a1640d5c205c..7b69aa4768c1 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -462,6 +462,8 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, f"after {elapsed_ms} millis with {retry_count} retries, " f"there maybe exist commit conflicts between multiple jobs." ) + if retry_result is not None and retry_result.exception is None: + raise CommitConflictError(error_msg) if retry_result is not None and retry_result.exception: raise RuntimeError(error_msg) from retry_result.exception else: @@ -495,7 +497,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str hash_index_base_snapshot, latest_snapshot_id ) ) - if retry_result is None: + if retry_result is None or retry_result.exception is None: raise CommitConflictError(str(conflict)) from conflict raise conflict @@ -563,7 +565,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str # Rolled back: base/snapshot no longer valid; next attempt # re-scans from scratch (matches Java RollbackRetryResult). return RetryResult(None, conflict_exception) - if retry_result is None: + if retry_result is None or retry_result.exception is None: raise CommitConflictError( str(conflict_exception) ) from conflict_exception @@ -1008,6 +1010,8 @@ def _generate_partition_statistics(self, commit_entries: List[ManifestEntry]) -> 'last_file_creation_time': 0, 'total_buckets': entry.total_buckets } + partition_stats[partition_key]['total_buckets'] = ( + entry.total_buckets) # Following Java implementation: PartitionEntry.fromDataFile() file_meta = entry.file diff --git a/paimon-python/pypaimon/write/postpone_batch_table_write.py b/paimon-python/pypaimon/write/postpone_batch_table_write.py index e0ce213c789b..7ef2ab9e2b0b 100644 --- a/paimon-python/pypaimon/write/postpone_batch_table_write.py +++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py @@ -150,8 +150,9 @@ def write_row(self, row): arrow_row = row_values_to_arrow_table( values_by_name, self.table.table_schema.fields, column_names) + size = self._planner.input_partition_stats(arrow_row)[partition][1] self._pending_inputs.append( - ("row", (row, partition, arrow_row.nbytes))) + ("row", (row, partition, size))) def _flush_pending_inputs(self): if not self._pending_inputs: diff --git a/paimon-python/pypaimon/write/postpone_bucket.py b/paimon-python/pypaimon/write/postpone_bucket.py index 7d7f3627467c..5b6855fae324 100644 --- a/paimon-python/pypaimon/write/postpone_bucket.py +++ b/paimon-python/pypaimon/write/postpone_bucket.py @@ -17,9 +17,9 @@ from typing import Dict, Tuple -import pyarrow as pa - from pypaimon.table.bucket_mode import BucketMode +from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer +from pypaimon.table.row.internal_row import RowKind class PostponeBucketPlan: @@ -56,6 +56,15 @@ def __init__( "Postpone fixed bucket writes require bucket = -2, got {}" .format(options.bucket()) ) + bucket_function = str( + table.table_schema.options.get("bucket-function.type", "default") + ).strip().lower() + if bucket_function != "default": + raise ValueError( + "Postpone fixed bucket writes only support " + "bucket-function.type=default, got {}" + .format(bucket_function) + ) self.max_num_buckets = ( options.postpone_batch_write_fixed_bucket_max_parallelism() @@ -85,9 +94,8 @@ def __init__( "{}".format(self.target_size_per_bucket) ) - self._partition_indices = [ - table.field_names.index(key) for key in table.partition_keys - ] + self._partition_keys = list(table.partition_keys) + self._field_dict = dict(table.field_dict) if known_num_buckets is None: known_num_buckets, loaded_postpone_counts = ( self._load_bucket_metadata(table) @@ -128,8 +136,8 @@ def current_plan(self) -> PostponeBucketPlan: return PostponeBucketPlan(self._known_num_buckets) def input_partition_stats(self, data) -> Dict[Tuple, Tuple[int, int]]: - if self._partition_indices: - columns = [data.column(i) for i in self._partition_indices] + if self._partition_keys: + columns = [data.column(key) for key in self._partition_keys] partitions = [ tuple(column[row].as_py() for column in columns) for row in range(data.num_rows) @@ -137,18 +145,19 @@ def input_partition_stats(self, data) -> Dict[Tuple, Tuple[int, int]]: else: partitions = [()] * data.num_rows - row_indices = {} - for index, partition in enumerate(partitions): - row_indices.setdefault(partition, []).append(index) - stats = {} - for partition, indices in row_indices.items(): - partition_data = ( - data - if len(indices) == data.num_rows - else pa.compute.take(data, pa.array(indices, type=pa.int64())) - ) - stats[partition] = (len(indices), partition_data.nbytes) + fields = [self._field_dict[name] for name in data.schema.names] + columns = [data.column(i) for i in range(len(fields))] + collect_size = self.target_row_num_per_bucket is None + for row, partition in enumerate(partitions): + data_size = 0 + if collect_size: + values = [column[row].as_py() for column in columns] + data_size = len(GenericRowSerializer.to_bytes( + GenericRow(values, fields, RowKind.INSERT) + )) - 4 + row_count, total_size = stats.get(partition, (0, 0)) + stats[partition] = (row_count + 1, total_size + data_size) return stats def plan( diff --git a/paimon-python/pypaimon/write/row_key_extractor.py b/paimon-python/pypaimon/write/row_key_extractor.py index 713023874b27..f97933fb8331 100644 --- a/paimon-python/pypaimon/write/row_key_extractor.py +++ b/paimon-python/pypaimon/write/row_key_extractor.py @@ -603,6 +603,15 @@ def __init__(self, table, bucket_plan): table.options.bucket() ) ) + bucket_function = str( + table.table_schema.options.get("bucket-function.type", "default") + ).strip().lower() + if bucket_function != "default": + raise ValueError( + "Postpone fixed bucket writes only support " + "bucket-function.type=default, got {}" + .format(bucket_function) + ) self.bucket_keys = table.table_schema.bucket_keys self.bucket_key_indices = self._get_field_indices(self.bucket_keys) self._bucket_key_fields = table.table_schema.logical_bucket_key_fields diff --git a/paimon-python/pypaimon/write/table_commit.py b/paimon-python/pypaimon/write/table_commit.py index 2f4d8e60a0ce..a70c52ebd348 100644 --- a/paimon-python/pypaimon/write/table_commit.py +++ b/paimon-python/pypaimon/write/table_commit.py @@ -90,9 +90,8 @@ def _commit(self, commit_messages: List[CommitMessage], commit_identifier: int = commit_identifier=commit_identifier ) except CommitConflictError: - # Conflict detection runs before manifest and snapshot creation, so - # these files are known to be uncommitted. Generic commit failures - # are intentionally not aborted because their success is uncertain. + # These files are known to be uncommitted. Generic commit failures + # remain untouched because their success is uncertain. try: self.file_store_commit.abort(non_empty_messages) except Exception: From a80974c8c6e1bd530c7347d40f1b57c9ab2ebfe4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 3 Aug 2026 05:35:47 -0700 Subject: [PATCH 5/7] [python] Simplify postpone writer tests --- .../pypaimon/tests/write/table_write_test.py | 368 +++++++----------- 1 file changed, 131 insertions(+), 237 deletions(-) diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index 78788866fda5..a056e940588d 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -15,13 +15,13 @@ # specific language governing permissions and limitations # under the License. -import glob import datetime +import glob import os import shutil - import tempfile import unittest +from contextlib import contextmanager from unittest.mock import Mock, patch from pypaimon import CatalogFactory, Schema @@ -57,6 +57,11 @@ def setUpClass(cls): ('behavior', pa.string()), pa.field('dt', pa.string(), nullable=False) ]) + cls.postpone_pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + pa.field('dt', pa.string(), nullable=False), + ('value', pa.string()), + ]) cls.expected = pa.Table.from_pydict({ 'user_id': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010], @@ -91,6 +96,51 @@ def _read_sorted(table, sort_keys): return read_builder.new_read().to_arrow( read_builder.new_scan().plan().splits()).sort_by(sort_keys) + def _create_postpone_table( + self, identifier, pa_schema=None, partition_keys=None, + primary_keys=None, options=None): + options = dict(options or {}) + options['bucket'] = -2 + schema = Schema.from_pyarrow_schema( + pa_schema if pa_schema is not None else self.pk_pa_schema, + partition_keys=partition_keys or [], + primary_keys=primary_keys or [], + options=options, + ) + self.catalog.create_table(identifier, schema, False) + return self.catalog.get_table(identifier) + + @staticmethod + def _commit_arrow(table, data, fixed_bucket=False): + builder = ( + table.new_postpone_fixed_bucket_write_builder() + if fixed_bucket else table.new_batch_write_builder() + ) + write = builder.new_write() + commit = builder.new_commit() + try: + write.write_arrow(data) + messages = write.prepare_commit() + commit.commit(messages) + return messages + finally: + write.close() + commit.close() + + @staticmethod + @contextmanager + def _postpone_write(table, overwrite=None): + builder = table.new_postpone_fixed_bucket_write_builder() + if overwrite is not None: + builder.overwrite(overwrite) + write = builder.new_write() + commit = builder.new_commit() + try: + yield write, commit + finally: + write.close() + commit.close() + @staticmethod def _mock_table_write(partitions, buckets): table_write = object.__new__(TableWrite) @@ -426,18 +476,16 @@ def test_multi_prepare_commit_pk(self): self.assertEqual(self.pk_expected, actual) def test_postpone_read_write(self): - schema = Schema.from_pyarrow_schema( - self.pa_schema, + table = self._create_postpone_table( + 'default.test_postpone', + pa_schema=self.pa_schema, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '1 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 2, }, ) - self.catalog.create_table('default.test_postpone', schema, False) - table = self.catalog.get_table('default.test_postpone') data = { 'user_id': [1, 2, 3, 4], 'item_id': [1001, 1002, 1003, 1004], @@ -486,15 +534,12 @@ def test_postpone_file_store_write_validates_runtime_bucket_count(self): PostponeFixedBucketFileStoreWrite, ) - schema = Schema.from_pyarrow_schema( - self.pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_runtime_bucket_validation', + pa_schema=self.pa_schema, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], - options={'bucket': -2}, ) - identifier = 'default.test_postpone_runtime_bucket_validation' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) write = PostponeFixedBucketFileStoreWrite(table, 'test-user') try: with self.assertRaisesRegex(ValueError, 'must be positive'): @@ -509,15 +554,12 @@ def test_postpone_file_store_write_validates_runtime_bucket_count(self): write.abort() def test_postpone_batch_write_builder_keeps_postpone_mode(self): - schema = Schema.from_pyarrow_schema( - self.pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_default_builder', + pa_schema=self.pa_schema, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], - options={'bucket': -2}, ) - identifier = 'default.test_postpone_default_builder' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) expected = pa.Table.from_pydict({ 'user_id': [1], 'item_id': [1001], @@ -525,13 +567,7 @@ def test_postpone_batch_write_builder_keeps_postpone_mode(self): 'dt': ['p1'], }, schema=self.pk_pa_schema) - write_builder = table.new_batch_write_builder() - table_write = write_builder.new_write() - table_commit = write_builder.new_commit() - table_write.write_arrow(expected) - table_commit.commit(table_write.prepare_commit()) - table_write.close() - table_commit.close() + self._commit_arrow(table, expected) self.assertEqual( 1, @@ -545,34 +581,23 @@ def test_postpone_batch_write_builder_keeps_postpone_mode(self): self.assertTrue(not table.new_read_builder().new_read().to_arrow(splits)) def test_postpone_batch_infers_bucket_num_from_input_size(self): - pa_schema = pa.schema([ - pa.field('id', pa.int32(), nullable=False), - pa.field('dt', pa.string(), nullable=False), - ('value', pa.string()), - ]) - schema = Schema.from_pyarrow_schema( - pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_size_inference', + pa_schema=self.postpone_pa_schema, partition_keys=['dt'], primary_keys=['id', 'dt'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '100 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }, ) - identifier = 'default.test_postpone_size_inference' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) data = pa.Table.from_pydict({ 'id': [1, 2], 'dt': ['small', 'large'], 'value': ['x', 'x' * 500], - }, schema=pa_schema) + }, schema=self.postpone_pa_schema) - builder = table.new_postpone_fixed_bucket_write_builder() - write = builder.new_write() - commit = builder.new_commit() - try: + with self._postpone_write(table) as (write, commit): write.write_arrow(data) messages = write.prepare_commit() total_buckets = { @@ -582,9 +607,6 @@ def test_postpone_batch_infers_bucket_num_from_input_size(self): self.assertEqual(1, total_buckets[('small',)]) self.assertEqual(3, total_buckets[('large',)]) commit.commit(messages) - finally: - write.close() - commit.close() self.assertEqual( [1, 2], self._read_sorted(table, 'id').column('id').to_pylist() @@ -598,18 +620,15 @@ def test_postpone_size_inference_matches_java_binary_row(self): pa.field('key', pa.string(), nullable=False), pa.field('value', pa.string(), nullable=False), ]) - schema = Schema.from_pyarrow_schema( + table = self._create_postpone_table( + 'default.test_postpone_java_size_fixture', pa_schema, primary_keys=['id'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '20 kb', 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }, ) - identifier = 'default.test_postpone_java_size_fixture' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) data = pa.RecordBatch.from_pydict({ 'id': list(range(1000)), 'key': ['k'] * 1000, @@ -632,14 +651,11 @@ def test_postpone_default_bucket_function_matches_java(self): pa.field('key', pa.string(), nullable=False), pa.field('value', pa.string(), nullable=False), ]) - schema = Schema.from_pyarrow_schema( + table = self._create_postpone_table( + 'default.test_postpone_java_bucket_fixture', pa_schema, primary_keys=['key'], - options={'bucket': -2}, ) - identifier = 'default.test_postpone_java_bucket_fixture' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) extractor = PostponeFixedBucketRowKeyExtractor( table, PostponeBucketPlan({(): 4})) data = pa.RecordBatch.from_pydict({ @@ -657,132 +673,87 @@ def test_postpone_rejects_unsupported_bucket_function( pa.field('id', pa.int32(), nullable=False), pa.field('value', pa.string(), nullable=False), ]) - schema = Schema.from_pyarrow_schema( + table = self._create_postpone_table( + 'default.test_postpone_bucket_function_' + bucket_function, pa_schema, primary_keys=['id'], - options={ - 'bucket': -2, - 'bucket-function.type': bucket_function, - }, + options={'bucket-function.type': bucket_function}, ) - identifier = ( - 'default.test_postpone_bucket_function_' + bucket_function) - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) with self.assertRaisesRegex( ValueError, 'only support bucket-function.type=default'): table.new_postpone_fixed_bucket_write_builder().new_write() def test_postpone_batch_plans_all_record_batches(self): - pa_schema = pa.schema([ - pa.field('id', pa.int32(), nullable=False), - pa.field('dt', pa.string(), nullable=False), - ('value', pa.string()), - ]) - schema = Schema.from_pyarrow_schema( - pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_multi_batch_plan', + pa_schema=self.postpone_pa_schema, partition_keys=['dt'], primary_keys=['id', 'dt'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '100 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }, ) - identifier = 'default.test_postpone_multi_batch_plan' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) - builder = table.new_postpone_fixed_bucket_write_builder() - write = builder.new_write() - commit = builder.new_commit() - try: + with self._postpone_write(table) as (write, commit): write.write_arrow_batch(pa.RecordBatch.from_pydict({ 'id': [1], 'dt': ['p'], 'value': ['x'], - }, schema=pa_schema)) + }, schema=self.postpone_pa_schema)) write.write_arrow_batch(pa.RecordBatch.from_pydict({ 'id': [2], 'dt': ['p'], 'value': ['x' * 500], - }, schema=pa_schema)) + }, schema=self.postpone_pa_schema)) messages = write.prepare_commit() self.assertEqual({3}, {message.total_buckets for message in messages}) commit.commit(messages) - finally: - write.close() - commit.close() self.assertEqual( [1, 2], self._read_sorted(table, 'id').column('id').to_pylist() ) def test_postpone_batch_prefers_target_row_num(self): - pa_schema = pa.schema([ - pa.field('id', pa.int32(), nullable=False), - pa.field('dt', pa.string(), nullable=False), - ('value', pa.string()), - ]) - schema = Schema.from_pyarrow_schema( - pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_row_num_plan', + pa_schema=self.postpone_pa_schema, partition_keys=['dt'], primary_keys=['id', 'dt'], options={ - 'bucket': -2, 'postpone.target-row-num-per-bucket': 1, 'postpone.target-size-per-bucket': '1 gb', 'postpone.batch-write-fixed-bucket.max-parallelism': 8, }, ) - identifier = 'default.test_postpone_row_num_plan' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) - builder = table.new_postpone_fixed_bucket_write_builder() - write = builder.new_write() - commit = builder.new_commit() - try: + with self._postpone_write(table) as (write, commit): write.write_arrow(pa.Table.from_pydict({ 'id': list(range(8)), 'dt': ['p'] * 8, 'value': ['x'] * 8, - }, schema=pa_schema)) + }, schema=self.postpone_pa_schema)) messages = write.prepare_commit() self.assertEqual({8}, {message.total_buckets for message in messages}) commit.commit(messages) - finally: - write.close() - commit.close() def test_postpone_batch_plans_write_rows_with_arrow_input(self): from pypaimon.table.row.generic_row import GenericRow - pa_schema = pa.schema([ - pa.field('id', pa.int32(), nullable=False), - pa.field('dt', pa.string(), nullable=False), - ('value', pa.string()), - ]) - schema = Schema.from_pyarrow_schema( - pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_write_row_plan', + pa_schema=self.postpone_pa_schema, partition_keys=['dt'], primary_keys=['id', 'dt'], options={ - 'bucket': -2, 'postpone.target-row-num-per-bucket': 1, 'postpone.target-size-per-bucket': '1 gb', 'postpone.batch-write-fixed-bucket.max-parallelism': 8, }, ) - identifier = 'default.test_postpone_write_row_plan' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) - builder = table.new_postpone_fixed_bucket_write_builder() - write = builder.new_write() - commit = builder.new_commit() - try: + with self._postpone_write(table) as (write, commit): for row_id in range(8): write.write_row( GenericRow([row_id, 'rows', 'x'], table.fields) @@ -792,7 +763,7 @@ def test_postpone_batch_plans_write_rows_with_arrow_input(self): 'id': list(range(9, 16)), 'dt': ['mixed'] * 7, 'value': ['x'] * 7, - }, schema=pa_schema)) + }, schema=self.postpone_pa_schema)) messages = write.prepare_commit() total_buckets = { @@ -802,9 +773,6 @@ def test_postpone_batch_plans_write_rows_with_arrow_input(self): self.assertEqual(8, total_buckets[('rows',)]) self.assertEqual(8, total_buckets[('mixed',)]) commit.commit(messages) - finally: - write.close() - commit.close() self.assertEqual(16, self._read_sorted(table, 'id').num_rows) @@ -813,19 +781,15 @@ def test_postpone_target_row_num_counts_existing_postpone_rows(self): PostponeBucketPlanner, ) - schema = Schema.from_pyarrow_schema( - self.pk_pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_existing_row_count', partition_keys=['dt'], primary_keys=['user_id', 'dt'], options={ - 'bucket': -2, 'postpone.target-row-num-per-bucket': 2, 'postpone.batch-write-fixed-bucket.max-parallelism': 8, }, ) - identifier = 'default.test_postpone_existing_row_count' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) append_planner = PostponeBucketPlanner( table, known_num_buckets={}, @@ -847,24 +811,16 @@ def test_postpone_target_row_num_counts_existing_postpone_rows(self): def test_postpone_worker_bucket_plan_mismatch_fails_commit(self): from pypaimon.write.commit.conflict_detection import CommitConflictError - pa_schema = pa.schema([ - pa.field('id', pa.int32(), nullable=False), - pa.field('dt', pa.string(), nullable=False), - ('value', pa.string()), - ]) - schema = Schema.from_pyarrow_schema( - pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_worker_plan_mismatch', + pa_schema=self.postpone_pa_schema, partition_keys=['dt'], primary_keys=['id', 'dt'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '100 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }, ) - identifier = 'default.test_postpone_worker_plan_mismatch' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) small_builder = table.new_postpone_fixed_bucket_write_builder() large_builder = table.new_postpone_fixed_bucket_write_builder() small_write = small_builder.new_write() @@ -873,10 +829,10 @@ def test_postpone_worker_bucket_plan_mismatch_fails_commit(self): try: small_write.write_arrow(pa.Table.from_pydict({ 'id': [1], 'dt': ['p'], 'value': ['x'], - }, schema=pa_schema)) + }, schema=self.postpone_pa_schema)) large_write.write_arrow(pa.Table.from_pydict({ 'id': [2], 'dt': ['p'], 'value': ['x' * 500], - }, schema=pa_schema)) + }, schema=self.postpone_pa_schema)) small_messages = small_write.prepare_commit() large_messages = large_write.prepare_commit() self.assertEqual({1}, {m.total_buckets for m in small_messages}) @@ -889,18 +845,16 @@ def test_postpone_worker_bucket_plan_mismatch_fails_commit(self): commit.close() def test_postpone_batch_fixed_bucket_reuses_existing_bucket_num(self): - schema = Schema.from_pyarrow_schema( - self.pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_reuse', + pa_schema=self.pa_schema, partition_keys=['dt'], primary_keys=['user_id', 'dt'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '1 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 2, }, ) - self.catalog.create_table('default.test_postpone_reuse', schema, False) - table = self.catalog.get_table('default.test_postpone_reuse') expected = pa.Table.from_pydict({ 'user_id': [1], 'item_id': [1001], @@ -908,13 +862,7 @@ def test_postpone_batch_fixed_bucket_reuses_existing_bucket_num(self): 'dt': ['p1'], }, schema=self.pk_pa_schema) - write_builder = table.new_postpone_fixed_bucket_write_builder() - table_write = write_builder.new_write() - table_commit = write_builder.new_commit() - table_write.write_arrow(expected) - table_commit.commit(table_write.prepare_commit()) - table_write.close() - table_commit.close() + self._commit_arrow(table, expected, fixed_bucket=True) copied_table = table.copy({ 'postpone.batch-write-fixed-bucket.max-parallelism': 3, @@ -939,42 +887,31 @@ def test_postpone_reuses_bucket_num_for_int_date_partition(self): pa.field('day', pa.date32(), nullable=False), ('value', pa.string()), ]) - schema = Schema.from_pyarrow_schema( + table = self._create_postpone_table( + 'default.test_postpone_typed_partition_reuse', pa_schema, partition_keys=['part', 'day'], primary_keys=['id', 'part', 'day'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '1 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 2, }, ) identifier = 'default.test_postpone_typed_partition_reuse' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) day = datetime.date(2026, 8, 1) - first_builder = table.new_postpone_fixed_bucket_write_builder() - first_write = first_builder.new_write() - first_commit = first_builder.new_commit() - first_write.write_arrow(pa.Table.from_pydict({ + self._commit_arrow(table, pa.Table.from_pydict({ 'id': [1], 'part': [7], 'day': [day], 'value': ['first'], - }, schema=pa_schema)) - first_commit.commit(first_write.prepare_commit()) - first_write.close() - first_commit.close() + }, schema=pa_schema), fixed_bucket=True) # Reopen to load the existing bucket count from manifests. reopened = self.catalog.get_table(identifier).copy({ 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }) - second_builder = reopened.new_postpone_fixed_bucket_write_builder() - second_write = second_builder.new_write() - second_commit = second_builder.new_commit() - try: + with self._postpone_write(reopened) as (second_write, second_commit): second_write.write_arrow(pa.Table.from_pydict({ 'id': [2, 3], 'part': [7, 8], @@ -989,48 +926,31 @@ def test_postpone_reuses_bucket_num_for_int_date_partition(self): self.assertEqual(2, total_buckets[(7, day)]) self.assertEqual(3, total_buckets[(8, day)]) second_commit.commit(messages) - finally: - second_write.close() - second_commit.close() actual = self._read_sorted(reopened, 'id') self.assertEqual([1, 2, 3], actual.column('id').to_pylist()) def test_postpone_legacy_partition_migrates_to_fixed_bucket(self): - schema = Schema.from_pyarrow_schema( - self.pk_pa_schema, + identifier = 'default.test_postpone_legacy_partition_migration' + legacy = self._create_postpone_table( + identifier, partition_keys=['dt'], primary_keys=['user_id', 'dt'], - options={ - 'bucket': -2, - 'postpone.batch-write-fixed-bucket': False, - }, + options={'postpone.batch-write-fixed-bucket': False}, ) - identifier = 'default.test_postpone_legacy_partition_migration' - self.catalog.create_table(identifier, schema, False) - legacy = self.catalog.get_table(identifier) - legacy_builder = legacy.new_batch_write_builder() - legacy_write = legacy_builder.new_write() - legacy_commit = legacy_builder.new_commit() - legacy_write.write_arrow(pa.Table.from_pydict({ + self._commit_arrow(legacy, pa.Table.from_pydict({ 'user_id': [1], 'item_id': [1001], 'behavior': ['legacy'], 'dt': ['p1'], }, schema=self.pk_pa_schema)) - legacy_commit.commit(legacy_write.prepare_commit()) - legacy_write.close() - legacy_commit.close() fixed = self.catalog.get_table(identifier).copy({ 'postpone.batch-write-fixed-bucket': True, 'postpone.target-size-per-bucket': '1 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 2, }) - fixed_builder = fixed.new_postpone_fixed_bucket_write_builder() - fixed_write = fixed_builder.new_write() - fixed_commit = fixed_builder.new_commit() - try: + with self._postpone_write(fixed) as (fixed_write, fixed_commit): fixed_write.write_arrow(pa.Table.from_pydict({ 'user_id': [2], 'item_id': [1002], @@ -1041,9 +961,6 @@ def test_postpone_legacy_partition_migrates_to_fixed_bucket(self): # Legacy -2 files do not define a real bucket count. self.assertEqual({2}, {m.total_buckets for m in messages}) fixed_commit.commit(messages) - finally: - fixed_write.close() - fixed_commit.close() scanner = fixed.new_read_builder().new_scan().file_scanner manifests, _ = scanner.manifest_scanner() @@ -1064,59 +981,40 @@ def test_postpone_legacy_partition_migrates_to_fixed_bucket(self): self.assertEqual([2], actual.column('user_id').to_pylist()) def test_postpone_overwrite_updates_catalog_bucket_count(self): - schema = Schema.from_pyarrow_schema( - self.pk_pa_schema, + table = self._create_postpone_table( + 'default.test_postpone_overwrite_bucket_statistics', partition_keys=['dt'], primary_keys=['user_id', 'dt'], - options={'bucket': -2}, ) - identifier = 'default.test_postpone_overwrite_bucket_statistics' - self.catalog.create_table(identifier, schema, False) - table = self.catalog.get_table(identifier) - - legacy_builder = table.new_batch_write_builder() - legacy_write = legacy_builder.new_write() - legacy_commit = legacy_builder.new_commit() - legacy_write.write_arrow(pa.Table.from_pydict({ + self._commit_arrow(table, pa.Table.from_pydict({ 'user_id': [1], 'item_id': [1001], 'behavior': ['legacy'], 'dt': ['p1'], }, schema=self.pk_pa_schema)) - legacy_commit.commit(legacy_write.prepare_commit()) - legacy_write.close() - legacy_commit.close() fixed = table.copy({ 'postpone.target-size-per-bucket': '1 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 2, }) - builder = ( - fixed.new_postpone_fixed_bucket_write_builder() - .overwrite({'dt': 'p1'}) - ) - write = builder.new_write() - commit = builder.new_commit() - write.write_arrow(pa.Table.from_pydict({ - 'user_id': [2], - 'item_id': [1002], - 'behavior': ['fixed'], - 'dt': ['p1'], - }, schema=self.pk_pa_schema)) - messages = write.prepare_commit() captured_statistics = [] - real_commit = commit.file_store_commit.snapshot_commit.commit + with self._postpone_write( + fixed, overwrite={'dt': 'p1'}) as (write, commit): + write.write_arrow(pa.Table.from_pydict({ + 'user_id': [2], + 'item_id': [1002], + 'behavior': ['fixed'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema)) + messages = write.prepare_commit() + real_commit = commit.file_store_commit.snapshot_commit.commit - def capture_statistics(base_snapshot_uuid, snapshot, statistics): - captured_statistics.extend(statistics) - return real_commit(base_snapshot_uuid, snapshot, statistics) + def capture_statistics(base_snapshot_uuid, snapshot, statistics): + captured_statistics.extend(statistics) + return real_commit(base_snapshot_uuid, snapshot, statistics) - commit.file_store_commit.snapshot_commit.commit = capture_statistics - try: + commit.file_store_commit.snapshot_commit.commit = capture_statistics commit.commit(messages) - finally: - write.close() - commit.close() self.assertEqual(2, captured_statistics[0].total_buckets) self.assertEqual( @@ -1126,19 +1024,15 @@ def capture_statistics(base_snapshot_uuid, snapshot, statistics): def test_postpone_concurrent_new_partition_bucket_num_conflict(self): from pypaimon.write.commit.conflict_detection import CommitConflictError - schema = Schema.from_pyarrow_schema( - self.pk_pa_schema, + table_two_buckets = self._create_postpone_table( + 'default.test_postpone_concurrent_bucket_num', partition_keys=['dt'], primary_keys=['user_id', 'dt'], options={ - 'bucket': -2, 'postpone.target-size-per-bucket': '1 b', 'postpone.batch-write-fixed-bucket.max-parallelism': 2, }, ) - identifier = 'default.test_postpone_concurrent_bucket_num' - self.catalog.create_table(identifier, schema, False) - table_two_buckets = self.catalog.get_table(identifier) table_three_buckets = table_two_buckets.copy({ 'postpone.batch-write-fixed-bucket.max-parallelism': 3, }) From 63ba1c371dee69f1a1db0f0a3b43c7c6e0fc689d Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 3 Aug 2026 07:17:56 -0700 Subject: [PATCH 6/7] [python] Address postpone fixed bucket review --- .../pypaimon/tests/write/table_write_test.py | 111 +++++++++++++ .../pypaimon/write/file_store_commit.py | 33 +++- .../pypaimon/write/postpone_bucket.py | 152 +++++++++++++++++- 3 files changed, 283 insertions(+), 13 deletions(-) diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index a056e940588d..08076e19b6f5 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -641,6 +641,43 @@ def test_postpone_size_inference_matches_java_binary_row(self): self.assertEqual((1000, 32000), stats[()]) self.assertEqual(2, planner.plan(stats).num_buckets(())) + def test_postpone_size_inference_supports_java_type_surface(self): + from pypaimon.write.postpone_bucket import PostponeBucketPlanner + + variant_type = pa.struct([ + pa.field('value', pa.binary(), nullable=False), + pa.field('metadata', pa.binary(), nullable=False), + ]) + pa_schema = pa.schema([ + pa.field('id', pa.int32(), nullable=False), + ('items', pa.list_(pa.int32())), + ('attributes', pa.map_(pa.string(), pa.int32())), + ('nested', pa.struct([('label', pa.string())])), + ('embedding', pa.list_(pa.float32(), 2)), + ('payload', variant_type), + ('event_time', pa.timestamp('us', tz='UTC')), + ]) + table = self._create_postpone_table( + 'default.test_postpone_java_type_surface', + pa_schema, + primary_keys=['id'], + ) + data = pa.RecordBatch.from_pydict({ + 'id': [1], + 'items': [[1, 2, 3]], + 'attributes': [[('a', 1)]], + 'nested': [{'label': 'x'}], + 'embedding': [[1.0, 2.0]], + 'payload': [{'value': b'\x00', 'metadata': b'\x01'}], + 'event_time': [datetime.datetime( + 2026, 1, 1, tzinfo=datetime.timezone.utc + )], + }, schema=pa_schema) + + planner = PostponeBucketPlanner( + table, known_num_buckets={}, postpone_row_counts={}) + self.assertEqual((1, 176), planner.input_partition_stats(data)[()]) + def test_postpone_default_bucket_function_matches_java(self): from pypaimon.write.postpone_bucket import PostponeBucketPlan from pypaimon.write.row_key_extractor import ( @@ -1097,6 +1134,80 @@ def fail_cas_after_concurrent_commit(*_): commit_two.close() commit_three.close() + def test_uncertain_commit_then_cas_failure_keeps_files(self): + table = self._create_postpone_table( + 'default.test_uncertain_commit_then_cas_failure', + pa_schema=self.postpone_pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + ) + builder = table.new_postpone_fixed_bucket_write_builder() + write = builder.new_write() + commit = builder.new_commit() + try: + write.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'dt': ['p'], 'value': ['v'], + }, schema=self.postpone_pa_schema)) + messages = write.prepare_commit() + data_paths = [ + file.external_path or file.file_path + for message in messages + for file in message.new_files + ] + uncertain_error = TimeoutError('lost commit response') + file_store_commit = commit.file_store_commit + file_store_commit.commit_max_retries = 1 + snapshot_commit = file_store_commit.snapshot_commit + real_commit = snapshot_commit.commit + attempts = 0 + + def uncertain_then_cas_failure(base_uuid, snapshot, statistics): + nonlocal attempts + attempts += 1 + if attempts == 1: + self.assertTrue(real_commit(base_uuid, snapshot, statistics)) + self._commit_arrow( + table, + pa.Table.from_pydict({ + 'id': [2], 'dt': ['p'], 'value': ['v2'], + }, schema=self.postpone_pa_schema), + fixed_bucket=True, + ) + raise uncertain_error + return False + + real_get_snapshot = file_store_commit.snapshot_manager.get_snapshot_by_id + + def hide_first_snapshot(snapshot_id): + return None if snapshot_id == 1 else real_get_snapshot(snapshot_id) + + # Keep the retry on the CAS path after snapshot 1 becomes unavailable. + with patch.object( + snapshot_commit, + 'commit', + side_effect=uncertain_then_cas_failure, + ), patch.object( + file_store_commit.snapshot_manager, + 'get_snapshot_by_id', + side_effect=hide_first_snapshot, + ), patch.object( + file_store_commit.conflict_detection, + 'check_conflicts', + return_value=None, + ), patch.object(file_store_commit, '_commit_retry_wait'): + with self.assertRaises(RuntimeError) as context: + commit.commit(messages) + + self.assertIs(uncertain_error, context.exception.__cause__) + self.assertEqual(2, attempts) + self.assertTrue(all(table.file_io.exists(path) for path in data_paths)) + self.assertEqual( + [1, 2], self._read_sorted(table, 'id').column('id').to_pylist() + ) + finally: + write.close() + commit.close() + def test_data_file_prefix_postpone(self): """Test that generated data file names follow the expected prefix format.""" schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], diff --git a/paimon-python/pypaimon/write/file_store_commit.py b/paimon-python/pypaimon/write/file_store_commit.py index 7b69aa4768c1..405052b8f556 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -73,9 +73,11 @@ def is_success(self) -> bool: class RetryResult(CommitResult): def __init__(self, latest_snapshot, exception: Optional[Exception] = None, - base_data_files: Optional[List[ManifestEntry]] = None): + base_data_files: Optional[List[ManifestEntry]] = None, + commit_result_may_be_uncertain: bool = False): self.latest_snapshot = latest_snapshot self.exception = exception + self.commit_result_may_be_uncertain = commit_result_may_be_uncertain # Base entries as of latest_snapshot, carried so the next attempt reuses # them and reads only the incremental changes. self.base_data_files = base_data_files @@ -381,6 +383,8 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, retry_count = 0 retry_result = None + commit_result_may_be_uncertain = False + uncertain_commit_exception = None rewritten_commit_entries = None start_time_ms = int(time.time() * 1000) while True: @@ -408,6 +412,7 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, index_deletes=index_deletes, index_adds=index_adds, hash_index_base_snapshot=hash_index_base_snapshot, + commit_result_may_be_uncertain=commit_result_may_be_uncertain, ) if isinstance(result, RewriteResult): @@ -442,6 +447,10 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, break else: retry_result = result + if result.commit_result_may_be_uncertain: + commit_result_may_be_uncertain = True + if uncertain_commit_exception is None: + uncertain_commit_exception = result.exception elapsed_ms = int(time.time() * 1000) - start_time_ms if elapsed_ms > self.commit_timeout or retry_count >= self.commit_max_retries: @@ -462,6 +471,8 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, f"after {elapsed_ms} millis with {retry_count} retries, " f"there maybe exist commit conflicts between multiple jobs." ) + if commit_result_may_be_uncertain: + raise RuntimeError(error_msg) from uncertain_commit_exception if retry_result is not None and retry_result.exception is None: raise CommitConflictError(error_msg) if retry_result is not None and retry_result.exception: @@ -481,7 +492,8 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str allow_rollback: bool = False, index_deletes=None, index_adds=None, - hash_index_base_snapshot=None) -> CommitResult: + hash_index_base_snapshot=None, + commit_result_may_be_uncertain: bool = False) -> CommitResult: start_millis = int(time.time() * 1000) if self._is_duplicate_commit(retry_result, latest_snapshot, commit_identifier, commit_kind): return SuccessResult() @@ -497,7 +509,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str hash_index_base_snapshot, latest_snapshot_id ) ) - if retry_result is None or retry_result.exception is None: + if not commit_result_may_be_uncertain: raise CommitConflictError(str(conflict)) from conflict raise conflict @@ -549,7 +561,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str if conflict_exception is not None: rewrite_result = self._try_rewrite_row_id_conflict( - retry_result, + commit_result_may_be_uncertain, conflict_exception, latest_snapshot, base_data_files, @@ -565,7 +577,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str # Rolled back: base/snapshot no longer valid; next attempt # re-scans from scratch (matches Java RollbackRetryResult). return RetryResult(None, conflict_exception) - if retry_result is None or retry_result.exception is None: + if not commit_result_may_be_uncertain: raise CommitConflictError( str(conflict_exception) ) from conflict_exception @@ -690,7 +702,12 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str except Exception as e: # Commit exception, not sure about the situation and should not clean up the files logger.warning("Retry commit for exception.", exc_info=True) - return RetryResult(latest_snapshot, e, base_data_files=base_data_files) + return RetryResult( + latest_snapshot, + e, + base_data_files=base_data_files, + commit_result_may_be_uncertain=True, + ) logger.info( "Successfully commit snapshot %d to table %s by user %s " @@ -715,7 +732,7 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str def _try_rewrite_row_id_conflict( self, - retry_result, + commit_result_may_be_uncertain, conflict_exception, latest_snapshot, base_data_files, @@ -727,7 +744,7 @@ def _try_rewrite_row_id_conflict( return None if commit_kind != "APPEND" or changelog_entries: return None - if retry_result is not None and retry_result.exception is not None: + if commit_result_may_be_uncertain: return None non_compaction_conflict = ( diff --git a/paimon-python/pypaimon/write/postpone_bucket.py b/paimon-python/pypaimon/write/postpone_bucket.py index 5b6855fae324..89c6c99499d4 100644 --- a/paimon-python/pypaimon/write/postpone_bucket.py +++ b/paimon-python/pypaimon/write/postpone_bucket.py @@ -17,9 +17,153 @@ from typing import Dict, Tuple +from pypaimon.schema.data_types import ( + ArrayType, + AtomicType, + MapType, + MultisetType, + RowType, + VectorType, +) from pypaimon.table.bucket_mode import BucketMode -from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer -from pypaimon.table.row.internal_row import RowKind +from pypaimon.table.row.generic_row import _parse_type_precision_scale + + +def _round_to_word(size): + return (size + 7) // 8 * 8 + + +class _BinaryRowSizeEstimator: + """Calculates the Java internal binary size without serializing rows.""" + + @classmethod + def row_size(cls, values, fields): + size = ((len(fields) + 71) // 64) * 8 + len(fields) * 8 + for value, field in zip(values, fields): + size += cls._variable_size(value, field.type) + return size + + @classmethod + def _variable_size(cls, value, data_type): + if isinstance(data_type, AtomicType): + return cls._atomic_variable_size(value, data_type) + if value is None: + return 0 + if isinstance(data_type, ArrayType): + return _round_to_word(cls._array_size(value, data_type.element)) + if isinstance(data_type, VectorType): + return _round_to_word( + 4 + data_type.length * cls._primitive_width(data_type.element) + ) + if isinstance(data_type, MapType): + keys, values = cls._map_values(value) + return _round_to_word( + 4 + + cls._array_size(keys, data_type.key) + + cls._array_size(values, data_type.value) + ) + if isinstance(data_type, MultisetType): + keys, values = cls._map_values(value) + return _round_to_word( + 4 + + cls._array_size(keys, data_type.element) + + cls._array_size(values, AtomicType("INT", False)) + ) + if isinstance(data_type, RowType): + return _round_to_word( + cls.row_size(cls._row_values(value, data_type), data_type.fields) + ) + raise ValueError("Unsupported data type: {}".format(data_type)) + + @classmethod + def _atomic_variable_size(cls, value, data_type): + type_name = data_type.type.upper() + if type_name.startswith(("DECIMAL", "NUMERIC")): + precision, _ = _parse_type_precision_scale(data_type) + return 16 if precision > 18 else 0 + if type_name.startswith("TIMESTAMP"): + precision, _ = _parse_type_precision_scale(data_type) + return 8 if precision > 3 else 0 + if value is None: + return 0 + if type_name == "VARIANT": + value_bytes, metadata = cls._variant_bytes(value) + return _round_to_word(4 + len(value_bytes) + len(metadata)) + if type_name == "BLOB": + value = value if isinstance(value, (bytes, bytearray)) else value.to_data() + return cls._binary_size(value) + if type_name.startswith(("CHAR", "VARCHAR", "STRING")): + return cls._binary_size(str(value).encode("utf-8")) + if type_name.startswith(("BINARY", "VARBINARY", "BYTES")): + return cls._binary_size(value) + return 0 + + @classmethod + def _array_size(cls, values, element_type): + values = list(values) + header_size = 4 + ((len(values) + 31) // 32) * 4 + size = _round_to_word( + header_size + len(values) * cls._fixed_width(element_type) + ) + return size + sum( + cls._variable_size(value, element_type) + for value in values + if value is not None + ) + + @staticmethod + def _binary_size(value): + length = len(bytes(value)) + return 0 if length <= 7 else _round_to_word(length) + + @staticmethod + def _variant_bytes(value): + if isinstance(value, dict): + return bytes(value["value"]), bytes(value["metadata"]) + return bytes(value.value()), bytes(value.metadata()) + + @staticmethod + def _map_values(value): + items = value.items() if isinstance(value, dict) else value + items = list(items) + return [item[0] for item in items], [item[1] for item in items] + + @staticmethod + def _row_values(value, row_type): + if isinstance(value, dict): + return [value[field.name] for field in row_type.fields] + if hasattr(value, "values"): + return value.values + return list(value) + + @classmethod + def _fixed_width(cls, data_type): + if isinstance(data_type, AtomicType): + type_name = data_type.type.upper() + if type_name in ("BOOLEAN", "BOOL", "TINYINT", "BYTE"): + return 1 + if type_name in ("SMALLINT", "SHORT"): + return 2 + if type_name in ("INT", "INTEGER", "FLOAT", "REAL", "DATE", "TIME") \ + or type_name.startswith("TIME("): + return 4 + return 8 + if isinstance(data_type, (ArrayType, MapType, MultisetType, RowType)): + return 8 + raise ValueError("Unsupported array element type: {}".format(data_type)) + + @staticmethod + def _primitive_width(data_type): + type_name = data_type.type.upper() + if type_name in ("BOOLEAN", "TINYINT"): + return 1 + if type_name == "SMALLINT": + return 2 + if type_name in ("INT", "INTEGER", "FLOAT"): + return 4 + if type_name in ("BIGINT", "DOUBLE"): + return 8 + raise ValueError("Unsupported vector element type: {}".format(data_type)) class PostponeBucketPlan: @@ -153,9 +297,7 @@ def input_partition_stats(self, data) -> Dict[Tuple, Tuple[int, int]]: data_size = 0 if collect_size: values = [column[row].as_py() for column in columns] - data_size = len(GenericRowSerializer.to_bytes( - GenericRow(values, fields, RowKind.INSERT) - )) - 4 + data_size = _BinaryRowSizeEstimator.row_size(values, fields) row_count, total_size = stats.get(partition, (0, 0)) stats[partition] = (row_count + 1, total_size + data_size) return stats From 0eb3e456879d13a36bc8b07eb3c494784f66041c Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 3 Aug 2026 08:56:23 -0700 Subject: [PATCH 7/7] [python] Refine postpone bucket planning --- .../pypaimon/tests/write/table_write_test.py | 28 +++++++++++++++++++ .../pypaimon/write/postpone_bucket.py | 21 ++++++++------ paimon-python/pypaimon/write/ray_datasink.py | 7 +++-- paimon-python/pypaimon/write/table_write.py | 1 + 4 files changed, 46 insertions(+), 11 deletions(-) diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index 08076e19b6f5..20a11cf6bc2f 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -641,6 +641,29 @@ def test_postpone_size_inference_matches_java_binary_row(self): self.assertEqual((1000, 32000), stats[()]) self.assertEqual(2, planner.plan(stats).num_buckets(())) + def test_postpone_stats_skip_known_partitions(self): + from pypaimon.write.postpone_bucket import PostponeBucketPlanner + + table = self._create_postpone_table( + 'default.test_postpone_skip_known_stats', + pa_schema=self.postpone_pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + ) + data = pa.RecordBatch.from_pydict({ + 'id': [1, 2], + 'dt': ['known', 'new'], + 'value': ['x', 'y'], + }, schema=self.postpone_pa_schema) + planner = PostponeBucketPlanner( + table, + known_num_buckets={('known',): 2}, + postpone_row_counts={}, + ) + + self.assertEqual( + {('new',)}, set(planner.input_partition_stats(data))) + def test_postpone_size_inference_supports_java_type_surface(self): from pypaimon.write.postpone_bucket import PostponeBucketPlanner @@ -775,6 +798,11 @@ def test_postpone_batch_prefers_target_row_num(self): self.assertEqual({8}, {message.total_buckets for message in messages}) commit.commit(messages) + self.assertEqual( + list(range(8)), + self._read_sorted(table, 'id').column('id').to_pylist(), + ) + def test_postpone_batch_plans_write_rows_with_arrow_input(self): from pypaimon.table.row.generic_row import GenericRow diff --git a/paimon-python/pypaimon/write/postpone_bucket.py b/paimon-python/pypaimon/write/postpone_bucket.py index 89c6c99499d4..297d55d28bcb 100644 --- a/paimon-python/pypaimon/write/postpone_bucket.py +++ b/paimon-python/pypaimon/write/postpone_bucket.py @@ -29,8 +29,12 @@ from pypaimon.table.row.generic_row import _parse_type_precision_scale +def _ceil_div(dividend, divisor): + return (dividend + divisor - 1) // divisor + + def _round_to_word(size): - return (size + 7) // 8 * 8 + return _ceil_div(size, 8) * 8 class _BinaryRowSizeEstimator: @@ -294,6 +298,8 @@ def input_partition_stats(self, data) -> Dict[Tuple, Tuple[int, int]]: columns = [data.column(i) for i in range(len(fields))] collect_size = self.target_row_num_per_bucket is None for row, partition in enumerate(partitions): + if partition in self._known_num_buckets: + continue data_size = 0 if collect_size: values = [column[row].as_py() for column in columns] @@ -320,19 +326,18 @@ def plan( total_rows = row_count + postpone_rows num_buckets = max( 1, - (total_rows + self.target_row_num_per_bucket - 1) - // self.target_row_num_per_bucket, + _ceil_div( + total_rows, self.target_row_num_per_bucket), ) else: estimated_size = data_size if postpone_rows and row_count: - estimated_size = ( - data_size * (row_count + postpone_rows) + row_count - 1 - ) // row_count + estimated_size = _ceil_div( + data_size * (row_count + postpone_rows), row_count) num_buckets = max( 1, - (estimated_size + self.target_size_per_bucket - 1) - // self.target_size_per_bucket, + _ceil_div( + estimated_size, self.target_size_per_bucket), ) self._known_num_buckets[partition] = min( num_buckets, self.max_num_buckets diff --git a/paimon-python/pypaimon/write/ray_datasink.py b/paimon-python/pypaimon/write/ray_datasink.py index 7e7c8c1f7f95..3d20f5fad382 100644 --- a/paimon-python/pypaimon/write/ray_datasink.py +++ b/paimon-python/pypaimon/write/ray_datasink.py @@ -359,16 +359,17 @@ def _collect_partition_stats(dataset, planner): def _stats(batch: pa.Table) -> pa.Table: stats = planner.input_partition_stats(batch) + items = list(stats.items()) return pa.table({ partition_col: pa.array( - [pickle.dumps(partition) for partition in stats], + [pickle.dumps(partition) for partition, _ in items], type=pa.binary(), ), rows_col: pa.array( - [value[0] for value in stats.values()], type=pa.int64() + [value[0] for _, value in items], type=pa.int64() ), size_col: pa.array( - [value[1] for value in stats.values()], type=pa.int64() + [value[1] for _, value in items], type=pa.int64() ), }) diff --git a/paimon-python/pypaimon/write/table_write.py b/paimon-python/pypaimon/write/table_write.py index 469325a4d408..162fe36f047f 100644 --- a/paimon-python/pypaimon/write/table_write.py +++ b/paimon-python/pypaimon/write/table_write.py @@ -271,6 +271,7 @@ def write_ray( ) def _distributed_write_options(self) -> Dict[str, Any]: + """Return options forwarded by ``write_ray`` to the Ray writer.""" return {} def close(self):