diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index c1887c537c75..1921a9764974 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -321,6 +321,11 @@ table_write.close() 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 7f56169eb46d..9b08d6014f21 100644 --- a/docs/docs/pypaimon/ray-data.md +++ b/docs/docs/pypaimon/ray-data.md @@ -248,8 +248,13 @@ 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. 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 @@ -261,11 +266,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/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/table/file_store_table.py b/paimon-python/pypaimon/table/file_store_table.py index 4bc9b1c04859..ac88bf5c12dc 100644 --- a/paimon-python/pypaimon/table/file_store_table.py +++ b/paimon-python/pypaimon/table/file_store_table.py @@ -428,6 +428,12 @@ def new_stream_read_builder(self) -> 'StreamReadBuilder': def new_batch_write_builder(self) -> BatchWriteBuilder: 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/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/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/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/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 d078304d2770..20a11cf6bc2f 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,10 +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, partition_keys=['user_id'], primary_keys=['user_id', 'dt'], - options={'bucket': -2}) - self.catalog.create_table('default.test_postpone', schema, False) - table = self.catalog.get_table('default.test_postpone') + table = self._create_postpone_table( + 'default.test_postpone', + pa_schema=self.pa_schema, + partition_keys=['user_id'], + primary_keys=['user_id', 'dt'], + options={ + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) data = { 'user_id': [1, 2, 3, 4], 'item_id': [1001, 1002, 1003, 1004], @@ -438,8 +494,16 @@ 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, + PostponeFixedBucketWriteBuilder, + ) + self.assertIsInstance( + write_builder, PostponeFixedBucketWriteBuilder) + self.assertIsInstance( + table_write, PostponeFixedBucketBatchTableWrite) table_commit = write_builder.new_commit() table_write.write_arrow(expect) commit_messages = table_write.prepare_commit() @@ -451,18 +515,731 @@ 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_file_store_write_validates_runtime_bucket_count(self): + from pypaimon.write.file_store_write import ( + PostponeFixedBucketFileStoreWrite, + ) + + 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'], + ) + 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_write_builder_keeps_postpone_mode(self): + table = self._create_postpone_table( + 'default.test_postpone_default_builder', + pa_schema=self.pa_schema, + partition_keys=['user_id'], + primary_keys=['user_id', 'dt'], + ) + expected = pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['a'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema) + + self._commit_arrow(table, expected) + + self.assertEqual( + 1, + len(glob.glob( + self.warehouse + + "/default.db/test_postpone_default_builder/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): + table = self._create_postpone_table( + 'default.test_postpone_size_inference', + pa_schema=self.postpone_pa_schema, + partition_keys=['dt'], + primary_keys=['id', 'dt'], + options={ + 'postpone.target-size-per-bucket': '100 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + data = pa.Table.from_pydict({ + 'id': [1, 2], + 'dt': ['small', 'large'], + 'value': ['x', 'x' * 500], + }, schema=self.postpone_pa_schema) + + with self._postpone_write(table) as (write, commit): + 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) + + self.assertEqual( + [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), + ]) + table = self._create_postpone_table( + 'default.test_postpone_java_size_fixture', + pa_schema, + primary_keys=['id'], + options={ + 'postpone.target-size-per-bucket': '20 kb', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + 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_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 + + 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 ( + PostponeFixedBucketRowKeyExtractor, + ) + + pa_schema = pa.schema([ + pa.field('key', pa.string(), nullable=False), + pa.field('value', pa.string(), nullable=False), + ]) + table = self._create_postpone_table( + 'default.test_postpone_java_bucket_fixture', + pa_schema, + primary_keys=['key'], + ) + 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), + ]) + table = self._create_postpone_table( + 'default.test_postpone_bucket_function_' + bucket_function, + pa_schema, + primary_keys=['id'], + options={'bucket-function.type': bucket_function}, + ) + + 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): + 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={ + 'postpone.target-size-per-bucket': '100 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + + with self._postpone_write(table) as (write, commit): + write.write_arrow_batch(pa.RecordBatch.from_pydict({ + 'id': [1], + 'dt': ['p'], + 'value': ['x'], + }, schema=self.postpone_pa_schema)) + write.write_arrow_batch(pa.RecordBatch.from_pydict({ + 'id': [2], + 'dt': ['p'], + 'value': ['x' * 500], + }, schema=self.postpone_pa_schema)) + messages = write.prepare_commit() + self.assertEqual({3}, {message.total_buckets for message in messages}) + commit.commit(messages) + + self.assertEqual( + [1, 2], self._read_sorted(table, 'id').column('id').to_pylist() + ) + + def test_postpone_batch_prefers_target_row_num(self): + 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={ + 'postpone.target-row-num-per-bucket': 1, + 'postpone.target-size-per-bucket': '1 gb', + 'postpone.batch-write-fixed-bucket.max-parallelism': 8, + }, + ) + + 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=self.postpone_pa_schema)) + messages = write.prepare_commit() + 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 + + 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={ + 'postpone.target-row-num-per-bucket': 1, + 'postpone.target-size-per-bucket': '1 gb', + 'postpone.batch-write-fixed-bucket.max-parallelism': 8, + }, + ) + + with self._postpone_write(table) as (write, commit): + 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=self.postpone_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) + + 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, + ) + + table = self._create_postpone_table( + 'default.test_postpone_existing_row_count', + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={ + 'postpone.target-row-num-per-bucket': 2, + 'postpone.batch-write-fixed-bucket.max-parallelism': 8, + }, + ) + 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 + + 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={ + 'postpone.target-size-per-bucket': '100 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }, + ) + 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() + try: + small_write.write_arrow(pa.Table.from_pydict({ + 'id': [1], 'dt': ['p'], 'value': ['x'], + }, schema=self.postpone_pa_schema)) + large_write.write_arrow(pa.Table.from_pydict({ + 'id': [2], 'dt': ['p'], 'value': ['x' * 500], + }, 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}) + 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): + table = self._create_postpone_table( + 'default.test_postpone_reuse', + pa_schema=self.pa_schema, + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={ + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) + expected = pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['a'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema) + + self._commit_arrow(table, expected, fixed_bucket=True) + + 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_postpone_fixed_bucket_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()), + ]) + table = self._create_postpone_table( + 'default.test_postpone_typed_partition_reuse', + pa_schema, + partition_keys=['part', 'day'], + primary_keys=['id', 'part', 'day'], + options={ + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) + identifier = 'default.test_postpone_typed_partition_reuse' + day = datetime.date(2026, 8, 1) + + self._commit_arrow(table, pa.Table.from_pydict({ + 'id': [1], + 'part': [7], + 'day': [day], + 'value': ['first'], + }, 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, + }) + with self._postpone_write(reopened) as (second_write, second_commit): + 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) + + 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): + identifier = 'default.test_postpone_legacy_partition_migration' + legacy = self._create_postpone_table( + identifier, + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={'postpone.batch-write-fixed-bucket': False}, + ) + self._commit_arrow(legacy, pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['legacy'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema)) + + 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, + }) + with self._postpone_write(fixed) as (fixed_write, fixed_commit): + 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) + + 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_overwrite_updates_catalog_bucket_count(self): + table = self._create_postpone_table( + 'default.test_postpone_overwrite_bucket_statistics', + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + ) + self._commit_arrow(table, pa.Table.from_pydict({ + 'user_id': [1], + 'item_id': [1001], + 'behavior': ['legacy'], + 'dt': ['p1'], + }, schema=self.pk_pa_schema)) + + fixed = table.copy({ + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }) + captured_statistics = [] + 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) + + commit.file_store_commit.snapshot_commit.commit = capture_statistics + commit.commit(messages) + + 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 + + table_two_buckets = self._create_postpone_table( + 'default.test_postpone_concurrent_bucket_num', + partition_keys=['dt'], + primary_keys=['user_id', 'dt'], + options={ + 'postpone.target-size-per-bucket': '1 b', + 'postpone.batch-write-fixed-bucket.max-parallelism': 2, + }, + ) + table_three_buckets = table_two_buckets.copy({ + 'postpone.batch-write-fixed-bucket.max-parallelism': 3, + }) + + 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() + 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}) + + 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() + 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'], - 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..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 @@ -211,6 +213,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, @@ -377,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: @@ -404,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): @@ -438,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: @@ -458,6 +471,10 @@ 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: raise RuntimeError(error_msg) from retry_result.exception else: @@ -475,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() @@ -491,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: + if not commit_result_may_be_uncertain: raise CommitConflictError(str(conflict)) from conflict raise conflict @@ -543,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, @@ -559,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: + if not commit_result_may_be_uncertain: raise CommitConflictError( str(conflict_exception) ) from conflict_exception @@ -684,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 " @@ -709,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, @@ -721,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 = ( @@ -839,12 +862,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 +881,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 +899,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 @@ -994,6 +1027,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/file_store_write.py b/paimon-python/pypaimon/write/file_store_write.py index 0cdd84664c4e..9a1dd10721a8 100644 --- a/paimon-python/pypaimon/write/file_store_write.py +++ b/paimon-python/pypaimon/write/file_store_write.py @@ -44,12 +44,16 @@ def __init__(self, table, commit_user): 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() + 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}" @@ -63,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) @@ -91,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() @@ -286,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 @@ -295,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.""" @@ -304,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) @@ -330,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 new file mode 100644 index 000000000000..7ef2ab9e2b0b --- /dev/null +++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py @@ -0,0 +1,204 @@ +# 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.table.bucket_mode import BucketMode +from pypaimon.write.commit_message import CommitMessage +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 ( + require_columns, + row_to_named_values, + 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): + """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 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 + 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) + size = self._planner.input_partition_stats(arrow_row)[partition][1] + self._pending_inputs.append( + ("row", (row, partition, size))) + + 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 _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..297d55d28bcb --- /dev/null +++ b/paimon-python/pypaimon/write/postpone_bucket.py @@ -0,0 +1,345 @@ +# 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 + +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 _parse_type_precision_scale + + +def _ceil_div(dividend, divisor): + return (dividend + divisor - 1) // divisor + + +def _round_to_word(size): + return _ceil_div(size, 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: + """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()) + ) + 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() + ) + 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_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) + ) + 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_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) + ] + else: + partitions = [()] * data.num_rows + + stats = {} + 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): + if partition in self._known_num_buckets: + continue + data_size = 0 + if collect_size: + values = [column[row].as_py() for column in columns] + 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 + + 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, + _ceil_div( + total_rows, self.target_row_num_per_bucket), + ) + else: + estimated_size = data_size + if postpone_rows and row_count: + estimated_size = _ceil_div( + data_size * (row_count + postpone_rows), row_count) + num_buckets = max( + 1, + _ceil_div( + estimated_size, 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..3d20f5fad382 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}") @@ -114,10 +118,16 @@ 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() table_schema = self.table.table_schema @@ -256,14 +266,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 +350,49 @@ 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) + items = list(stats.items()) + return pa.table({ + partition_col: pa.array( + [pickle.dumps(partition) for partition, _ in items], + type=pa.binary(), + ), + rows_col: pa.array( + [value[0] for _, value in items], type=pa.int64() + ), + size_col: pa.array( + [value[1] for _, value in items], 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 +401,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 +412,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 +431,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..f97933fb8331 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,68 @@ 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() + ) + ) + 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 + 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_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: diff --git a/paimon-python/pypaimon/write/table_write.py b/paimon-python/pypaimon/write/table_write.py index c1276e72d2cc..162fe36f047f 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) @@ -70,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, @@ -156,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) @@ -170,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 @@ -245,8 +267,13 @@ 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 options forwarded by ``write_ray`` to the Ray writer.""" + return {} + def close(self): try: self.file_store_write.close()