diff --git a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java index bd5911f54eaa..63ae166ab0c3 100644 --- a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -33,6 +33,7 @@ import org.apache.paimon.data.GenericRow; import org.apache.paimon.data.InternalArray; import org.apache.paimon.data.InternalRow; +import org.apache.paimon.data.serializer.InternalRowSerializer; import org.apache.paimon.data.variant.GenericVariant; import org.apache.paimon.deletionvectors.BitmapDeletionVector; import org.apache.paimon.deletionvectors.DeletionVector; @@ -42,6 +43,7 @@ import org.apache.paimon.fs.Path; import org.apache.paimon.fs.local.LocalFileIO; import org.apache.paimon.globalindex.sorted.SortedGlobalIndexBuilder; +import org.apache.paimon.index.HashBucketAssigner; import org.apache.paimon.index.IndexFileMeta; import org.apache.paimon.io.CompactIncrement; import org.apache.paimon.io.DataFileMeta; @@ -335,6 +337,35 @@ public void testJavaWriteReadPkTable() throws Exception { } } + @Test + @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") + public void testJavaWriteDynamicBucketHashIndex() throws Exception { + Identifier identifier = identifier("dynamic_hash_java_to_python"); + Schema schema = + Schema.newBuilder() + .column("key1", DataTypes.STRING()) + .column("key2", DataTypes.BIGINT()) + .column("value", DataTypes.STRING()) + .primaryKey("key1", "key2") + .option("bucket", "-1") + .option("dynamic-bucket.target-row-num", "1") + .option("file.format", "parquet") + .build(); + catalog.createTable(identifier, schema, true); + FileStoreTable table = (FileStoreTable) catalog.getTable(identifier); + + try (StreamTableWrite write = table.newWrite(commitUser); + InnerTableCommit commit = table.newCommit(commitUser)) { + GenericRow row = + GenericRow.of( + BinaryString.fromString("hello-java"), + 42L, + BinaryString.fromString("java-old")); + write.write(row, assignDynamicBucket(table, row)); + commit.commit(0, write.prepareCommit(true, 0)); + } + } + @Test @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") public void testPKDeletionVectorWrite() throws Exception { @@ -525,6 +556,51 @@ public void testReadPkTable() throws Exception { } } + @Test + @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") + public void testReadPythonDynamicBucketHashIndex() throws Exception { + Identifier identifier = identifier("dynamic_hash_python_to_java"); + FileStoreTable table = (FileStoreTable) catalog.getTable(identifier); + + try (StreamTableWrite write = table.newWrite(commitUser); + InnerTableCommit commit = table.newCommit(commitUser)) { + GenericRow row = + GenericRow.of( + BinaryString.fromString("hello-java"), + 42L, + BinaryString.fromString("java-new")); + write.write(row, assignDynamicBucket(table, row)); + commit.commit(1, write.prepareCommit(true, 1)); + } + + List result = + getResult( + table.newRead(), + table.newScan().plan().splits(), + row -> DataFormatTestUtil.toStringNoRowKind(row, table.rowType())); + assertThat(result) + .containsExactlyInAnyOrder( + "hello-java, 42, java-new", "python-only, 7, python-only"); + } + + private int assignDynamicBucket(FileStoreTable table, InternalRow row) { + InternalRowSerializer serializer = + new InternalRowSerializer(DataTypes.STRING(), DataTypes.BIGINT()); + int keyHash = + serializer.toBinaryRow(GenericRow.of(row.getString(0), row.getLong(1))).hashCode(); + HashBucketAssigner assigner = + new HashBucketAssigner( + table.snapshotManager(), + commitUser, + table.store().newIndexFileHandler(), + 1, + 1, + 0, + 1, + -1); + return assigner.assign(BinaryRow.EMPTY_ROW, keyHash); + } + @Test @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") public void testBtreeIndexWrite() throws Exception { diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh index 6423b59e049e..b89bf6770fd8 100755 --- a/paimon-python/dev/run_mixed_tests.sh +++ b/paimon-python/dev/run_mixed_tests.sh @@ -98,6 +98,7 @@ run_batched_java_write_tests() { local result=0 local core_tests="org.apache.paimon.JavaPyE2ETest#testJavaWriteReadPkTable" + core_tests="${core_tests}+testJavaWriteDynamicBucketHashIndex" core_tests="${core_tests}+testPKDeletionVectorWrite" core_tests="${core_tests}+testBtreeIndexWrite" core_tests="${core_tests}+testBtreeRawFallbackWrite" @@ -184,7 +185,7 @@ run_java_write_test() { echo "Running Maven test for JavaPyE2ETest.testJavaWriteReadPkTable (Parquet/Orc/Avro)..." echo "Note: Maven may download dependencies on first run, this may take a while..." local parquet_result=0 - if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testJavaWriteReadPkTable -pl paimon-core -Drun.e2e.tests=true; then + if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testJavaWriteReadPkTable+testJavaWriteDynamicBucketHashIndex -pl paimon-core -Drun.e2e.tests=true; then echo -e "${GREEN}✓ Java write Parquet/Orc/Avro test completed successfully${NC}" else echo -e "${RED}✗ Java write Parquet/Orc/Avro test failed${NC}" @@ -222,7 +223,7 @@ run_python_read_test() { # Run the parameterized Python test method (runs for both Parquet/Orc/Avro and Lance) echo "Running Python test for JavaPyReadWriteTest.test_read_pk_table..." - if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest -k "test_read_pk_table" -v; then + if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest -k "test_read_pk_table or test_read_java_dynamic_bucket_hash_index" -v; then echo -e "${GREEN}✓ Python test completed successfully${NC}" # source deactivate return 0 @@ -241,7 +242,7 @@ run_python_write_test() { # Run the parameterized Python test method for writing data (pk table, includes bucket num assertion) echo "Running Python test for JavaPyReadWriteTest (test_py_write_read_pk_table)..." - if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest -k "test_py_write_read_pk_table" -v; then + if python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest -k "test_py_write_read_pk_table or test_py_write_dynamic_bucket_hash_index" -v; then echo -e "${GREEN}✓ Python write test completed successfully${NC}" return 0 else @@ -260,7 +261,7 @@ run_java_read_test() { echo "Running Maven test for JavaPyE2ETest.testReadPkTable (Java Read Parquet/Orc/Avro)..." echo "Note: Maven may download dependencies on first run, this may take a while..." local parquet_result=0 - if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testReadPkTable -pl paimon-core -Drun.e2e.tests=true -Dpython.version="$PYTHON_VERSION"; then + if mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testReadPkTable+testReadPythonDynamicBucketHashIndex -pl paimon-core -Drun.e2e.tests=true -Dpython.version="$PYTHON_VERSION"; then echo -e "${GREEN}✓ Java read Parquet/Orc/Avro test completed successfully${NC}" else echo -e "${RED}✗ Java read Parquet/Orc/Avro test failed${NC}" diff --git a/paimon-python/pypaimon/daft/daft_datasink.py b/paimon-python/pypaimon/daft/daft_datasink.py index 747ccd2bc4d8..c21767b0d2e9 100644 --- a/paimon-python/pypaimon/daft/daft_datasink.py +++ b/paimon-python/pypaimon/daft/daft_datasink.py @@ -19,6 +19,7 @@ from __future__ import annotations import logging +import pickle from typing import TYPE_CHECKING, Any import pyarrow as pa @@ -36,6 +37,9 @@ logger = logging.getLogger(__name__) _PaimonIdentifier = tuple[str, str, str | None] +COMMIT_MESSAGES_COLUMN = "__paimon_commit_messages__" +_WORKER_TABLE_CACHE: dict[tuple[Any, ...], FileStoreTable] = {} +_WORKER_TABLE_CACHE_SIZE = 32 def _options_to_dict(options: Any) -> dict[str, Any]: @@ -122,17 +126,39 @@ def _load_table( ) -class PaimonDataSink(DataSink[list[Any]]): - """DataSink for writing data to an Apache Paimon table. +def _load_worker_table( + catalog_options: dict[str, Any], + table_identifier: _PaimonIdentifier | None, + table_path: str | None, + schema_token: tuple[int, int], +) -> FileStoreTable: + """Reuse immutable table metadata within a Daft worker process.""" + cache_key = ( + tuple(sorted((str(key), repr(value)) + for key, value in catalog_options.items())), + table_identifier, + table_path, + schema_token, + ) + cached = _WORKER_TABLE_CACHE.get(cache_key) + if cached is not None: + return cached + + loaded = _load_table(catalog_options, table_identifier, table_path) + if len(_WORKER_TABLE_CACHE) >= _WORKER_TABLE_CACHE_SIZE: + _WORKER_TABLE_CACHE.pop(next(iter(_WORKER_TABLE_CACHE))) + _WORKER_TABLE_CACHE[cache_key] = loaded + return loaded - Delegates all file I/O and commit logic to pypaimon's BatchTableWrite / - BatchTableCommit APIs. Each `write()` call (one per parallel worker) opens - an independent BatchTableWrite, accumulates micro-partitions, then calls - `prepare_commit()` to obtain CommitMessage objects. `finalize()` gathers - all CommitMessages from all workers and performs a single atomic commit. - """ - def __init__(self, table: FileStoreTable, mode: str = "append") -> None: +class PaimonDataSink(DataSink[list[Any]]): + """DataSink for writing data to an Apache Paimon table.""" + + def __init__( + self, + table: FileStoreTable, + mode: str = "append", + ) -> None: if mode not in ("append", "overwrite"): raise ValueError(f"Only 'append' or 'overwrite' mode is supported, got: {mode!r}") self._mode = mode @@ -264,6 +290,22 @@ def finalize(self, write_results: list[WriteResult[list[Any]]]) -> MicroPartitio all_commit_messages = [msg for wr in write_results for msg in wr.result] table_commit = self._write_builder.new_commit() + non_empty_workers = sum( + any(not message.is_empty() for message in result.result) + for result in write_results + ) + primary_key_error = self._primary_key_write_error(non_empty_workers) + if primary_key_error is not None: + try: + table_commit.abort(all_commit_messages) + except Exception: + logger.warning( + "Failed to abort uncommitted direct Daft PK files.", + exc_info=True, + ) + finally: + table_commit.close() + raise ValueError(primary_key_error) try: table_commit.commit(all_commit_messages) finally: @@ -280,3 +322,445 @@ def finalize(self, write_results: list[WriteResult[list[Any]]]) -> MicroPartitio "file_name": pa.array([f.file_name for f in all_files], type=pa.string()), } ) + + def _primary_key_write_error( + self, non_empty_workers: int + ) -> str | None: + if not self._table.is_primary_key_table: + return None + + from pypaimon.table.bucket_mode import BucketMode + + bucket_mode = self._table.bucket_mode() + if bucket_mode == BucketMode.POSTPONE_MODE: + return None + if bucket_mode in ( + BucketMode.HASH_FIXED, + BucketMode.HASH_DYNAMIC, + ) and non_empty_workers <= 1: + return None + if bucket_mode in (BucketMode.HASH_FIXED, BucketMode.HASH_DYNAMIC): + reason = "require a single non-empty Daft write task" + elif bucket_mode == BucketMode.CROSS_PARTITION: + reason = ( + "require a persistent global primary-key index, which " + "PyPaimon does not yet support" + ) + else: + reason = f"are unsafe for {bucket_mode.name}" + return ( + f"Direct PaimonDataSink primary-key writes {reason}. Use " + "pypaimon.daft.write_paimon or PaimonTable.append/overwrite " + "for coordinated input." + ) + + +class PaimonCommitDataSink(PaimonDataSink): + """Commit files already produced by partition/bucket group UDFs.""" + + def __init__( + self, + table: FileStoreTable, + mode: str = "append", + commit_messages_column: str = COMMIT_MESSAGES_COLUMN, + ) -> None: + self._commit_messages_column = commit_messages_column + super().__init__(table, mode) + + def __getstate__(self) -> dict[str, Any]: + state = super().__getstate__() + state["_commit_messages_column"] = self._commit_messages_column + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self._commit_messages_column = state.pop( + "_commit_messages_column", COMMIT_MESSAGES_COLUMN + ) + super().__setstate__(state) + + @property + def commit_user(self) -> str | None: + return self._commit_user + + def name(self) -> str: + return "Paimon Commit" + + def _primary_key_write_error( + self, non_empty_workers: int + ) -> str | None: + return None + + def write( + self, micropartitions: Iterator[MicroPartition] + ) -> Iterator[WriteResult[list[Any]]]: + commit_messages = [] + total_bytes = 0 + for mp in micropartitions: + for rb in mp.get_record_batches(): + batch = rb.to_arrow_record_batch() + column_index = batch.schema.get_field_index( + self._commit_messages_column + ) + if column_index < 0: + raise ValueError( + "Missing internal column " + f"{self._commit_messages_column!r}" + ) + for payload in batch.column(column_index).to_pylist(): + if payload is None: + continue + total_bytes += len(payload) + decoded = pickle.loads(payload) + if not isinstance(decoded, list): + raise TypeError( + "Invalid Paimon group-write result: expected a list " + f"of commit messages, got {type(decoded).__name__}" + ) + commit_messages.extend(decoded) + + yield WriteResult( + result=commit_messages, + bytes_written=total_bytes, + rows_written=sum( + file.row_count + for message in commit_messages + for file in message.new_files + ), + ) + + +def _series_to_arrow_table(columns, schema: pa.Schema) -> pa.Table: + arrays = [column.to_arrow() for column in columns] + table = pa.Table.from_arrays(arrays, names=schema.names) + if table.schema != schema: + table = table.cast(schema) + return table + + +def make_fixed_bucket_udf(table: FileStoreTable): + """Create a Daft batch UDF which computes Paimon's fixed bucket.""" + import daft + from daft import DataType, Series + + from pypaimon.schema.data_types import PyarrowFieldParser + from pypaimon.write.row_key_extractor import FixedBucketRowKeyExtractor + + schema = PyarrowFieldParser.from_paimon_schema(table.fields) + extractor = FixedBucketRowKeyExtractor(table.table_schema) + + @daft.func.batch(return_dtype=DataType.int32()) + def _fixed_bucket(*columns) -> Series: + arrow_table = _series_to_arrow_table(columns, schema) + buckets = [] + for batch in arrow_table.to_batches(): + _, batch_buckets = extractor.extract_partition_bucket_batch(batch) + buckets.extend(batch_buckets) + return Series.from_arrow(pa.array(buckets, type=pa.int32())) + + return _fixed_bucket + + +def make_dynamic_routing_udf( + table: FileStoreTable, + num_channels: int, + num_assigners: int, + assigner_column: str, + key_hash_column: str, +): + """Calculate each dynamic key hash once and route it to an assigner.""" + import daft + from daft import DataType, Series + + from pypaimon.index.dynamic_bucket import compute_assigner, to_signed_int32 + from pypaimon.schema.data_types import PyarrowFieldParser + from pypaimon.write.row_key_extractor import DynamicBucketRowKeyExtractor + + schema = PyarrowFieldParser.from_paimon_schema(table.fields) + extractor = DynamicBucketRowKeyExtractor(table.table_schema) + return_fields = { + assigner_column: DataType.int32(), + key_hash_column: DataType.int32(), + } + + @daft.func.batch( + return_dtype=DataType.struct(return_fields), + unnest=True, + ) + def _dynamic_routing(*columns) -> Series: + arrow_table = _series_to_arrow_table(columns, schema) + assigners = [] + key_hashes = [] + for batch in arrow_table.to_batches(): + _, partition_hashes, batch_key_hashes = ( + extractor.extract_hashes_batch(batch) + ) + assigners.extend([ + compute_assigner( + partition_hash, key_hash, num_channels, num_assigners + ) + for partition_hash, key_hash in zip( + partition_hashes, batch_key_hashes + ) + ]) + key_hashes.extend( + to_signed_int32(value) for value in batch_key_hashes + ) + return Series.from_arrow( + pa.StructArray.from_arrays( + [ + pa.array(assigners, type=pa.int32()), + pa.array(key_hashes, type=pa.int32()), + ], + names=[assigner_column, key_hash_column], + ) + ) + + return _dynamic_routing + + +def make_dynamic_bucket_assignment_udf( + table: FileStoreTable, + num_channels: int, + num_assigners: int, + bucket_column: str, + key_hash_column: str, + new_mapping_column: str, + base_snapshot_column: str, + ignore_existing: bool = False, + base_snapshot_id: int | None = None, +): + """Assign persistent buckets inside one partition/assigner group.""" + import daft + from daft import DataType, Series + + from pypaimon.schema.data_types import PyarrowFieldParser + from pypaimon.write.row_key_extractor import DynamicBucketRowKeyExtractor + + catalog_options = _extract_catalog_options(table) + table_identifier = _extract_identifier(table) + table_path_value = getattr(table, "table_path", None) + table_path = str(table_path_value) if table_path_value is not None else None + schema_token = ( + table.table_schema.id, + table.table_schema.time_millis, + ) + schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data_column_count = len(schema) + partition_names = set(table.partition_keys) + output_fields = [ + field for field in schema if field.name not in partition_names + ] + return_fields = { + field.name: DataType.from_arrow_type(field.type) + for field in output_fields + } + return_fields[key_hash_column] = DataType.int32() + return_fields[bucket_column] = DataType.int32() + return_fields[new_mapping_column] = DataType.bool() + return_fields[base_snapshot_column] = DataType.int64() + + @daft.func.batch( + return_dtype=DataType.struct(return_fields), + unnest=True, + ) + def _assign_dynamic_bucket(*columns) -> Series: + assigner_values = columns[data_column_count].to_pylist() + key_hash_values = columns[data_column_count + 1].to_pylist() + if not assigner_values: + return Series.from_arrow( + pa.array([], type=pa.struct([ + pa.field(field.name, field.type) for field in output_fields + ] + [ + pa.field(key_hash_column, pa.int32()), + pa.field(bucket_column, pa.int32()), + pa.field(new_mapping_column, pa.bool_()), + pa.field(base_snapshot_column, pa.int64()), + ])) + ) + assign_id = assigner_values[0] + if any(value != assign_id for value in assigner_values): + raise RuntimeError( + "A dynamic-bucket group contained multiple assigners" + ) + + arrow_table = _series_to_arrow_table( + columns[:data_column_count], schema + ).combine_chunks() + worker_table = _load_worker_table( + catalog_options, + table_identifier, + table_path, + schema_token, + ) + extractor = DynamicBucketRowKeyExtractor( + worker_table.table_schema, + table=worker_table, + num_channels=num_channels, + num_assigners=num_assigners, + assign_id=assign_id, + ignore_existing=ignore_existing, + maintain_index=False, + base_snapshot_id=base_snapshot_id, + ) + buckets = [] + new_mappings = [] + offset = 0 + for batch in arrow_table.to_batches(): + batch_hashes = key_hash_values[offset:offset + batch.num_rows] + _, batch_buckets, batch_new_mappings = ( + extractor.extract_partition_bucket_status_from_hashes_batch( + batch, batch_hashes + ) + ) + buckets.extend(batch_buckets) + new_mappings.extend(batch_new_mappings) + offset += batch.num_rows + arrays = [ + arrow_table.column(field.name).combine_chunks() + for field in output_fields + ] + arrays.append(pa.array(key_hash_values, type=pa.int32())) + arrays.append(pa.array(buckets, type=pa.int32())) + arrays.append(pa.array(new_mappings, type=pa.bool_())) + arrays.append(pa.array( + [extractor.base_snapshot_id] * arrow_table.num_rows, + type=pa.int64(), + )) + return Series.from_arrow( + pa.StructArray.from_arrays( + arrays, + names=[field.name for field in output_fields] + [ + key_hash_column, + bucket_column, + new_mapping_column, + base_snapshot_column, + ], + ) + ) + + return _assign_dynamic_bucket + + +def make_group_write_udf( + table: FileStoreTable, + mode: str, + commit_user: str | None, + *, + precomputed_bucket: bool = False, + base_snapshot_id: int | None = None, +): + """Create a UDF which owns one complete Paimon write group.""" + import daft + from daft import DataType, Series + + from pypaimon.schema.data_types import PyarrowFieldParser + from pypaimon.table.bucket_mode import BucketMode + + catalog_options = _extract_catalog_options(table) + table_identifier = _extract_identifier(table) + table_path_value = getattr(table, "table_path", None) + table_path = str(table_path_value) if table_path_value is not None else None + schema_token = ( + table.table_schema.id, + table.table_schema.time_millis, + ) + schema = PyarrowFieldParser.from_paimon_schema(table.fields) + data_column_count = len(schema) + dynamic_bucket = table.bucket_mode() == BucketMode.HASH_DYNAMIC + + @daft.func.batch(return_dtype=DataType.binary()) + def _write_group(*columns) -> Series: + data_columns = columns[:data_column_count] + bucket = None + key_hash_values = None + new_mapping_values = None + group_base_snapshot_id = base_snapshot_id + if precomputed_bucket: + bucket_values = columns[data_column_count].to_pylist() + if not bucket_values: + return Series.from_arrow(pa.array([], type=pa.binary())) + bucket = bucket_values[0] + if any(value != bucket for value in bucket_values): + raise RuntimeError( + "A Paimon write group contained multiple buckets" + ) + if dynamic_bucket: + key_hash_values = columns[data_column_count + 1].to_pylist() + new_mapping_values = columns[data_column_count + 2].to_pylist() + base_snapshot_values = columns[ + data_column_count + 3 + ].to_pylist() + observed_base_snapshot_id = base_snapshot_values[0] + if any( + value != observed_base_snapshot_id + for value in base_snapshot_values + ): + raise RuntimeError( + "A dynamic-bucket group contained multiple base " + "snapshots" + ) + if group_base_snapshot_id is None: + group_base_snapshot_id = observed_base_snapshot_id + elif observed_base_snapshot_id != group_base_snapshot_id: + raise RuntimeError( + "A dynamic-bucket group did not match the driver " + "base snapshot" + ) + + arrow_table = _series_to_arrow_table(data_columns, schema) + worker_table = _load_worker_table( + catalog_options, + table_identifier, + table_path, + schema_token, + ) + write_builder = worker_table.new_batch_write_builder() + if commit_user is not None: + write_builder.commit_user = commit_user + if mode == "overwrite": + write_builder.overwrite({}) + table_write = write_builder.new_write() + if precomputed_bucket and dynamic_bucket: + table_write.with_dynamic_bucket_index( + ignore_existing=mode == "overwrite", + base_snapshot_id=group_base_snapshot_id, + ) + try: + if precomputed_bucket: + offset = 0 + for batch in arrow_table.to_batches(): + batch_hashes = None + batch_new_mappings = None + if key_hash_values is not None: + batch_hashes = key_hash_values[ + offset:offset + batch.num_rows + ] + batch_new_mappings = new_mapping_values[ + offset:offset + batch.num_rows + ] + table_write.write_arrow_batch_to_bucket( + batch, + bucket, + batch_hashes, + batch_new_mappings, + ) + offset += batch.num_rows + else: + table_write.write_arrow(arrow_table) + commit_messages = table_write.prepare_commit() + table_write.close() + except Exception: + try: + table_write.abort() + except Exception: + logger.warning( + "Failed to abort Daft Paimon group write.", + exc_info=True, + ) + raise + + return Series.from_arrow( + pa.array([pickle.dumps(commit_messages)], type=pa.binary()) + ) + + return _write_group diff --git a/paimon-python/pypaimon/daft/daft_paimon.py b/paimon-python/pypaimon/daft/daft_paimon.py index e34b35b0d178..b294ede3cdcb 100644 --- a/paimon-python/pypaimon/daft/daft_paimon.py +++ b/paimon-python/pypaimon/daft/daft_paimon.py @@ -229,10 +229,203 @@ def _write_table( table: FileStoreTable, mode: str = "append", ) -> "daft.DataFrame": - """Write a Daft DataFrame to a Paimon table object.""" - from pypaimon.daft.daft_datasink import PaimonDataSink + """Write a Daft DataFrame to a Paimon table object. + + Primary-key HASH modes materialize each ``(partition, bucket)`` group in + one Daft worker to guarantee a single Paimon writer owner. A large fixed + bucket can therefore become a single-worker memory/CPU hot spot. Dynamic + buckets are normally bounded by ``dynamic-bucket.target-row-num``. Each + group also returns its commit-message list in one pickled binary cell, so + groups producing many files can create large task results. + """ + from pypaimon.daft.daft_datasink import ( + COMMIT_MESSAGES_COLUMN, + PaimonCommitDataSink, + PaimonDataSink, + make_dynamic_bucket_assignment_udf, + make_dynamic_routing_udf, + make_fixed_bucket_udf, + make_group_write_udf, + ) + from pypaimon.table.bucket_mode import BucketMode + from pypaimon.index.dynamic_bucket import SHORT_MAX_VALUE + + if not table.is_primary_key_table: + return df.write_sink(PaimonDataSink(table, mode)) - return df.write_sink(PaimonDataSink(table, mode)) + _validate_write_column_names(df, table) + bucket_mode = table.bucket_mode() + if bucket_mode == BucketMode.POSTPONE_MODE: + return df.write_sink(PaimonDataSink(table, mode)) + if bucket_mode == BucketMode.CROSS_PARTITION: + raise ValueError( + "CROSS_PARTITION primary-key Daft writes require a persistent " + "global primary-key index, which PyPaimon does not yet support" + ) + + import daft + + data_columns = [daft.col(name) for name in table.field_names] + commit_column = _pick_internal_column( + df.column_names, COMMIT_MESSAGES_COLUMN + ) + commit_sink = PaimonCommitDataSink(table, mode, commit_column) + + if bucket_mode == BucketMode.HASH_FIXED: + bucket_column = _pick_internal_column( + df.column_names + [commit_column], "__paimon_bucket__" + ) + bucket_udf = make_fixed_bucket_udf(table) + with_bucket = df.with_column( + bucket_column, + bucket_udf(*data_columns), + ) + group_keys = list(table.partition_keys) + [bucket_column] + group_write = make_group_write_udf( + table, + mode, + commit_sink.commit_user, + precomputed_bucket=True, + ) + prepared = with_bucket.groupby(*group_keys).map_groups( + group_write( + *data_columns, + daft.col(bucket_column), + ).alias(commit_column) + ) + return prepared.write_sink(commit_sink) + + if bucket_mode == BucketMode.HASH_DYNAMIC: + latest_snapshot = table.snapshot_manager().get_latest_snapshot() + base_snapshot_id = ( + latest_snapshot.id if latest_snapshot is not None else 0 + ) + input_partitions = df.num_partitions() or 1 + max_buckets = table.options.dynamic_bucket_max_buckets() + # Bucket ids are stored as Java shorts and Paimon allocates ids in + # [0, Short.MAX_VALUE), so every assigner must own at least one id. + num_assigners = min(max(1, input_partitions), SHORT_MAX_VALUE) + if max_buckets > 0: + num_assigners = min(num_assigners, max_buckets) + num_channels = num_assigners + + assigner_column = _pick_internal_column( + df.column_names + [commit_column], "__paimon_assigner__" + ) + bucket_column = _pick_internal_column( + df.column_names + [commit_column, assigner_column], + "__paimon_bucket__", + ) + key_hash_column = _pick_internal_column( + df.column_names + [ + commit_column, + assigner_column, + bucket_column, + ], + "__paimon_key_hash__", + ) + new_mapping_column = _pick_internal_column( + df.column_names + [ + commit_column, + assigner_column, + bucket_column, + key_hash_column, + ], + "__paimon_new_mapping__", + ) + base_snapshot_column = _pick_internal_column( + df.column_names + [ + commit_column, + assigner_column, + bucket_column, + key_hash_column, + new_mapping_column, + ], + "__paimon_base_snapshot__", + ) + routing_udf = make_dynamic_routing_udf( + table, + num_channels, + num_assigners, + assigner_column, + key_hash_column, + ) + with_routing = df.select( + *data_columns, + routing_udf(*data_columns), + ) + bucket_assignment_udf = make_dynamic_bucket_assignment_udf( + table, + num_channels, + num_assigners, + bucket_column, + key_hash_column, + new_mapping_column, + base_snapshot_column, + ignore_existing=mode == "overwrite", + base_snapshot_id=base_snapshot_id, + ) + assigner_group_keys = list(table.partition_keys) + [assigner_column] + with_bucket = with_routing.groupby(*assigner_group_keys).map_groups( + bucket_assignment_udf( + *data_columns, + daft.col(assigner_column), + daft.col(key_hash_column), + ) + ) + # After a parallelism change, keys handled by different current + # assigners can restore to the same historical bucket. Regroup by the + # final bucket so exactly one file writer owns it in this commit. + group_write = make_group_write_udf( + table, + mode, + commit_sink.commit_user, + precomputed_bucket=True, + base_snapshot_id=base_snapshot_id, + ) + group_keys = list(table.partition_keys) + [bucket_column] + prepared = with_bucket.groupby(*group_keys).map_groups( + group_write( + *data_columns, + daft.col(bucket_column), + daft.col(key_hash_column), + daft.col(new_mapping_column), + daft.col(base_snapshot_column), + ).alias(commit_column) + ) + return prepared.write_sink(commit_sink) + + raise ValueError( + f"Unsupported primary-key bucket mode for Daft writes: {bucket_mode.name}" + ) + + +def _validate_write_column_names(df: "daft.DataFrame", table: FileStoreTable) -> None: + input_names = list(df.column_names) + target_names = list(table.field_names) + if len(set(input_names)) != len(input_names): + raise ValueError( + f"Cannot write to Paimon with duplicate input field names: {input_names}" + ) + missing = [name for name in target_names if name not in input_names] + extra = [name for name in input_names if name not in target_names] + if missing or extra: + details = [] + if missing: + details.append(f"missing fields: {missing}") + if extra: + details.append(f"extra fields: {extra}") + raise ValueError(f"Paimon write schema mismatch: {'; '.join(details)}") + + +def _pick_internal_column(existing_names, preferred: str) -> str: + existing = set(existing_names) + if preferred not in existing: + return preferred + suffix = 1 + while f"{preferred}_{suffix}" in existing: + suffix += 1 + return f"{preferred}_{suffix}" def read_paimon( diff --git a/paimon-python/pypaimon/index/dynamic_bucket.py b/paimon-python/pypaimon/index/dynamic_bucket.py new file mode 100644 index 000000000000..c5bbd08b7dfa --- /dev/null +++ b/paimon-python/pypaimon/index/dynamic_bucket.py @@ -0,0 +1,515 @@ +################################################################################ +# 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. +################################################################################ + +"""Persistent hash assignment and index maintenance for dynamic buckets.""" + +import random +import struct +from dataclasses import dataclass +from typing import Dict, Iterator, List, Optional, Set, Tuple + +from pypaimon.index.index_file_handler import IndexFileHandler +from pypaimon.index.index_file_meta import IndexFileMeta +from pypaimon.manifest.index_manifest_entry import IndexManifestEntry +from pypaimon.table.row.generic_row import GenericRow + +HASH_INDEX = "HASH" +_ADD = 0 +_DELETE = 1 +SHORT_MAX_VALUE = 32767 +_SNAPSHOT_UNSET = object() + + +def to_signed_int32(value: int) -> int: + """Return the Java ``int`` representation of ``value``.""" + value &= 0xFFFFFFFF + return value - 0x100000000 if value >= 0x80000000 else value + + +def _java_remainder(value: int, divisor: int) -> int: + remainder = abs(value) % divisor + return -remainder if value < 0 else remainder + + +def compute_assigner( + partition_hash: int, + key_hash: int, + num_channels: int, + num_assigners: int, +) -> int: + """Mirror Java ``BucketAssigner.computeAssigner``.""" + if num_channels <= 0 or num_assigners <= 0: + raise ValueError("num_channels and num_assigners must be positive") + start = abs(_java_remainder(to_signed_int32(partition_hash), num_channels)) + assigner = abs(_java_remainder(to_signed_int32(key_hash), num_assigners)) + return (start + assigner) % num_channels + + +def is_my_bucket(bucket: int, num_assigners: int, assign_id: int) -> bool: + return bucket % num_assigners == assign_id % num_assigners + + +def _iter_hashes(table, entry: IndexManifestEntry) -> Iterator[int]: + meta = entry.index_file + path = meta.external_path + if path is None: + # Files without an external path always live in the table's index + # directory, even if global-index.external-path was configured later. + index_path = table.path_factory().global_index_path_factory().index_path() + path = f"{index_path}/{meta.file_name}" + with table.file_io.new_input_stream(path) as stream: + remainder = b"" + while True: + payload = stream.read(64 * 1024) + if not payload: + break + payload = remainder + payload + aligned_size = len(payload) - len(payload) % 4 + for value in struct.iter_unpack(">i", payload[:aligned_size]): + yield value[0] + remainder = payload[aligned_size:] + if remainder: + raise RuntimeError( + f"Corrupt dynamic-bucket HASH index {path}: " + "expected a multiple of 4 bytes" + ) + + +def _read_hashes(table, entry: IndexManifestEntry) -> Set[int]: + return set(_iter_hashes(table, entry)) + + +class _PartitionIndex: + def __init__( + self, + hash_to_bucket: Dict[int, int], + bucket_information: Dict[int, int], + target_bucket_row_number: int, + ) -> None: + self.hash_to_bucket = hash_to_bucket + self.non_full_bucket_information = bucket_information + self.total_bucket_set = set(bucket_information) + self.total_bucket_list = list(bucket_information) + self.target_bucket_row_number = target_bucket_row_number + + def assign( + self, + key_hash: int, + bucket_filter, + max_buckets_num: int, + max_bucket_id: int, + ) -> int: + return self.assign_with_status( + key_hash, bucket_filter, max_buckets_num, max_bucket_id + )[0] + + def assign_with_status( + self, + key_hash: int, + bucket_filter, + max_buckets_num: int, + max_bucket_id: int, + ) -> Tuple[int, bool]: + if key_hash in self.hash_to_bucket: + return self.hash_to_bucket[key_hash], False + + for bucket, row_count in list(self.non_full_bucket_information.items()): + if row_count < self.target_bucket_row_number: + self.non_full_bucket_information[bucket] = row_count + 1 + self.hash_to_bucket[key_hash] = bucket + return bucket, True + del self.non_full_bucket_information[bucket] + + global_max_bucket_id = ( + SHORT_MAX_VALUE if max_buckets_num == -1 else max_buckets_num + ) - 1 + if not self.total_bucket_set or max_bucket_id < global_max_bucket_id: + for bucket in range(global_max_bucket_id + 1): + if bucket_filter(bucket) and bucket not in self.total_bucket_set: + self.non_full_bucket_information[bucket] = 1 + self.total_bucket_set.add(bucket) + self.total_bucket_list.append(bucket) + self.hash_to_bucket[key_hash] = bucket + return bucket, True + if max_buckets_num == -1: + raise RuntimeError( + "No dynamic bucket id remains below Java Short.MAX_VALUE. " + "Increase dynamic-bucket.target-row-num." + ) + + if not self.total_bucket_list: + raise RuntimeError( + "No dynamic bucket is available for this assigner. Reduce the " + "assigner parallelism or increase dynamic-bucket.max-buckets." + ) + bucket = random.choice(self.total_bucket_list) + self.hash_to_bucket[key_hash] = bucket + return bucket, True + + +class HashBucketAssigner: + """Restore dynamic-bucket mappings from HASH indexes before assigning.""" + + def __init__( + self, + table, + num_channels: int, + num_assigners: int, + assign_id: int, + target_bucket_row_number: int, + max_buckets_num: int, + ignore_existing: bool = False, + snapshot=_SNAPSHOT_UNSET, + ) -> None: + if not 0 <= assign_id < num_channels: + raise ValueError( + f"assign_id must be in [0, {num_channels}), got {assign_id}" + ) + self.table = table + self.num_channels = num_channels + self.num_assigners = num_assigners + self.assign_id = assign_id + self.target_bucket_row_number = target_bucket_row_number + self.max_buckets_num = max_buckets_num + self.ignore_existing = ignore_existing + self.snapshot = ( + table.snapshot_manager().get_latest_snapshot() + if snapshot is _SNAPSHOT_UNSET + else snapshot + ) + self.max_bucket_id = 0 + self._partition_indexes: Dict[Tuple, _PartitionIndex] = {} + + def assign(self, partition: Tuple, partition_hash: int, key_hash: int) -> int: + return self.assign_with_status(partition, partition_hash, key_hash)[0] + + def assign_with_status( + self, partition: Tuple, partition_hash: int, key_hash: int + ) -> Tuple[int, bool]: + result = self.assign_batch( + [partition], [partition_hash], [key_hash] + ) + return result[0] + + def assign_batch( + self, + partitions: List[Tuple], + partition_hashes: List[int], + key_hashes: List[int], + ) -> List[Tuple[int, bool]]: + if not ( + len(partitions) == len(partition_hashes) == len(key_hashes) + ): + raise ValueError( + "Partition, partition-hash, and key-hash counts must match" + ) + requested_by_partition: Dict[Tuple, Set[int]] = {} + normalized_hashes = [to_signed_int32(value) for value in key_hashes] + for partition, partition_hash, key_hash in zip( + partitions, partition_hashes, normalized_hashes + ): + self._validate_assigner(partition_hash, key_hash) + requested_by_partition.setdefault(partition, set()).add(key_hash) + + for partition, requested in requested_by_partition.items(): + index = self._partition_indexes.get(partition) + if index is None: + index = self._load_partition( + partition, + requested, + ) + self._partition_indexes[partition] = index + else: + missing = requested.difference(index.hash_to_bucket) + if missing: + self._restore_requested_hashes( + partition, + missing, + index.hash_to_bucket, + ) + + results = [] + for partition, key_hash in zip(partitions, normalized_hashes): + index = self._partition_indexes[partition] + bucket, is_new = index.assign_with_status( + key_hash, + lambda value: is_my_bucket( + value, self.num_assigners, self.assign_id + ), + self.max_buckets_num, + self.max_bucket_id, + ) + self.max_bucket_id = max(self.max_bucket_id, bucket) + results.append((bucket, is_new)) + return results + + def _validate_assigner(self, partition_hash: int, key_hash: int) -> None: + expected_assigner = compute_assigner( + partition_hash, + key_hash, + self.num_channels, + self.num_assigners, + ) + if expected_assigner != self.assign_id: + raise ValueError( + f"Record assigner {expected_assigner} does not match writer " + f"assigner {self.assign_id}" + ) + + def _load_partition( + self, partition: Tuple, requested_hashes: Set[int] + ) -> _PartitionIndex: + if self.ignore_existing: + return _PartitionIndex({}, {}, self.target_bucket_row_number) + snapshot = self.snapshot + if snapshot is None: + return _PartitionIndex({}, {}, self.target_bucket_row_number) + entries = IndexFileHandler(self.table).scan( + snapshot, + lambda entry: entry.index_file.index_type == HASH_INDEX + and tuple(entry.partition.values) == partition, + ) + entries_by_bucket = {} + for entry in entries: + previous = entries_by_bucket.get(entry.bucket) + if previous is not None: + raise RuntimeError( + "Found multiple dynamic-bucket HASH indexes for partition " + f"{partition}, bucket {entry.bucket}: " + f"{[previous.index_file.file_name, entry.index_file.file_name]}" + ) + entries_by_bucket[entry.bucket] = entry + + data_buckets = self._active_data_buckets(partition, snapshot) + missing_buckets = data_buckets.difference(entries_by_bucket) + if missing_buckets: + raise RuntimeError( + f"Dynamic-bucket partition {partition} has data files but no " + "complete HASH index for buckets " + f"{sorted(missing_buckets)}. Rewrite the partition before " + "performing incremental writes." + ) + + hash_to_bucket: Dict[int, int] = {} + bucket_information: Dict[int, int] = {} + for entry in entries: + self.max_bucket_id = max(self.max_bucket_id, entry.bucket) + if is_my_bucket( + entry.bucket, self.num_assigners, self.assign_id + ): + bucket_information[entry.bucket] = entry.index_file.row_count + remaining = set(requested_hashes) + for entry in entries: + if not remaining: + break + for key_hash in _iter_hashes(self.table, entry): + if key_hash in remaining: + hash_to_bucket[key_hash] = entry.bucket + remaining.remove(key_hash) + if not remaining: + break + return _PartitionIndex( + hash_to_bucket, + bucket_information, + self.target_bucket_row_number, + ) + + def _restore_requested_hashes( + self, + partition: Tuple, + requested_hashes: Set[int], + hash_to_bucket: Dict[int, int], + ) -> None: + if self.ignore_existing or not requested_hashes: + return + snapshot = self.snapshot + if snapshot is None: + return + entries = IndexFileHandler(self.table).scan( + snapshot, + lambda entry: entry.index_file.index_type == HASH_INDEX + and tuple(entry.partition.values) == partition, + ) + remaining = set(requested_hashes) + for entry in entries: + for key_hash in _iter_hashes(self.table, entry): + if key_hash not in remaining: + continue + hash_to_bucket[key_hash] = entry.bucket + remaining.remove(key_hash) + if not remaining: + return + + def _active_data_buckets(self, partition: Tuple, snapshot) -> Set[int]: + if snapshot is None: + return set() + from pypaimon.common.options.core_options import CoreOptions + + scan_table = self.table.copy( + {CoreOptions.SCAN_SNAPSHOT_ID.key(): str(snapshot.id)} + ) + read_builder = scan_table.new_read_builder() + if scan_table.partition_keys: + predicate_builder = read_builder.new_predicate_builder() + predicates = [ + predicate_builder.equal(key, value) + for key, value in zip(scan_table.partition_keys, partition) + ] + read_builder.with_filter( + predicate_builder.and_predicates(predicates) + ) + return { + split.bucket + for split in read_builder.new_scan().plan_for_write().splits() + if split.files + } + + +@dataclass +class DynamicBucketIndexChanges: + additions: List[IndexManifestEntry] + deletions: List[IndexManifestEntry] + + +class DynamicBucketIndexMaintainer: + """Maintain one complete HASH index file for every modified bucket.""" + + def __init__( + self, + table, + ignore_existing: bool = False, + snapshot=_SNAPSHOT_UNSET, + ) -> None: + self.table = table + self.ignore_existing = ignore_existing + self.snapshot = ( + table.snapshot_manager().get_latest_snapshot() + if snapshot is _SNAPSHOT_UNSET + else snapshot + ) + self.base_snapshot_id = ( + self.snapshot.id if self.snapshot is not None else 0 + ) + self._states: Dict[ + Tuple[Tuple, int], Tuple[Set[int], Optional[IndexManifestEntry], bool] + ] = {} + self._new_paths: List[str] = [] + + def notify_new_record( + self, partition: Tuple, bucket: int, key_hash: int + ) -> None: + key = (partition, bucket) + state = self._states.get(key) + if state is None: + hashes, old_entry = self._load_bucket(partition, bucket) + state = (hashes, old_entry, False) + hashes, old_entry, modified = state + previous_size = len(hashes) + hashes.add(to_signed_int32(key_hash)) + self._states[key] = (hashes, old_entry, modified or len(hashes) != previous_size) + + def prepare_commit(self) -> Dict[Tuple[Tuple, int], DynamicBucketIndexChanges]: + changes: Dict[Tuple[Tuple, int], DynamicBucketIndexChanges] = {} + for key, (hashes, old_entry, modified) in self._states.items(): + if not modified: + continue + partition, bucket = key + new_entry = self._write_index(partition, bucket, hashes) + changes[key] = DynamicBucketIndexChanges( + additions=[new_entry], + deletions=[ + IndexManifestEntry( + kind=_DELETE, + partition=old_entry.partition, + bucket=old_entry.bucket, + index_file=old_entry.index_file, + ) + ] if old_entry is not None else [], + ) + self._states[key] = (hashes, new_entry, False) + return changes + + def release_prepared(self) -> None: + """Release files after their commit messages leave writer ownership.""" + self._new_paths.clear() + + def abort(self) -> None: + for path in self._new_paths: + self.table.file_io.delete_quietly(path) + self._new_paths.clear() + + def _load_bucket( + self, partition: Tuple, bucket: int + ) -> Tuple[Set[int], Optional[IndexManifestEntry]]: + """Load a complete bucket index for full-file replacement. + + This matches Java's HASH index lifecycle. A modified bucket therefore + incurs O(bucket size) read memory and IO until an incremental index + format is introduced. + """ + if self.ignore_existing: + return set(), None + snapshot = self.snapshot + if snapshot is None: + return set(), None + entries = IndexFileHandler(self.table).scan( + snapshot, + lambda entry: entry.index_file.index_type == HASH_INDEX + and tuple(entry.partition.values) == partition + and entry.bucket == bucket, + ) + if len(entries) > 1: + raise RuntimeError( + "Found multiple dynamic-bucket HASH indexes for partition " + f"{partition}, bucket {bucket}: " + f"{[entry.index_file.file_name for entry in entries]}" + ) + if not entries: + return set(), None + return _read_hashes(self.table, entries[0]), entries[0] + + def _write_index( + self, partition: Tuple, bucket: int, hashes: Set[int] + ) -> IndexManifestEntry: + path_factory = self.table.path_factory().global_index_path_factory() + self.table.file_io.check_or_mkdirs(path_factory.global_index_root_path()) + path = path_factory.new_path() + payload = b"".join(struct.pack(">i", value) for value in sorted(hashes)) + try: + with self.table.file_io.new_output_stream(path) as stream: + stream.write(payload) + except Exception: + self.table.file_io.delete_quietly(path) + raise + self._new_paths.append(path) + + meta = IndexFileMeta( + index_type=HASH_INDEX, + file_name=path.rsplit("/", 1)[-1], + file_size=self.table.file_io.get_file_size(path), + row_count=len(hashes), + external_path=path if path_factory.is_external_path() else None, + ) + return IndexManifestEntry( + kind=_ADD, + partition=GenericRow( + list(partition), self.table.partition_keys_fields + ), + bucket=bucket, + index_file=meta, + ) diff --git a/paimon-python/pypaimon/manifest/index_manifest_file.py b/paimon-python/pypaimon/manifest/index_manifest_file.py index 2f711d722e7c..3c1c9448bac7 100644 --- a/paimon-python/pypaimon/manifest/index_manifest_file.py +++ b/paimon-python/pypaimon/manifest/index_manifest_file.py @@ -85,6 +85,7 @@ class IndexManifestFile: DELETION_VECTORS_INDEX = "DELETION_VECTORS" + HASH_INDEX = _HASH_INDEX def __init__(self, table): from pypaimon.table.file_store_table import FileStoreTable @@ -254,6 +255,7 @@ def combine_changes( delete_names = {e.index_file.file_name for e in deletes} survivors = [e for e in previous if e.index_file.file_name not in delete_names] _validate_retained_global_index_files(survivors, adds) + _validate_hash_index_changes(survivors, adds) combined = survivors + adds if not combined: return None @@ -351,6 +353,35 @@ def _validate_retained_global_index_files( ) +def _validate_hash_index_changes( + retained_entries: List[IndexManifestEntry], + added_entries: List[IndexManifestEntry], +) -> None: + """Keep the one-HASH-file-per-bucket manifest invariant.""" + by_bucket = { + (tuple(entry.partition.values), entry.bucket): entry + for entry in retained_entries + if entry.index_file.index_type == _HASH_INDEX + } + for added in added_entries: + if added.kind != _ADD or added.index_file.index_type != _HASH_INDEX: + continue + key = (tuple(added.partition.values), added.bucket) + previous = by_bucket.get(key) + if previous is not None: + raise RuntimeError( + "Trying to add HASH index file {} for partition {}, bucket {}, " + "but HASH index file {} still exists. Remove the previous file " + "first.".format( + added.index_file.file_name, + key[0], + key[1], + previous.index_file.file_name, + ) + ) + by_bucket[key] = added + + def _is_global_index(index_type: str) -> bool: return index_type not in (IndexManifestFile.DELETION_VECTORS_INDEX, _HASH_INDEX) diff --git a/paimon-python/pypaimon/table/file_store_table.py b/paimon-python/pypaimon/table/file_store_table.py index 8edbc1402156..4b89de1a5375 100644 --- a/paimon-python/pypaimon/table/file_store_table.py +++ b/paimon-python/pypaimon/table/file_store_table.py @@ -476,7 +476,9 @@ def drop_global_index(self, index_column, index_type: str = "btree", dry_run=dry_run, ) - def create_row_key_extractor(self) -> RowKeyExtractor: + def create_row_key_extractor( + self, ignore_existing: bool = False + ) -> RowKeyExtractor: bucket_mode = self.bucket_mode() if bucket_mode == BucketMode.HASH_FIXED: return FixedBucketRowKeyExtractor(self.table_schema) @@ -484,8 +486,17 @@ def create_row_key_extractor(self) -> RowKeyExtractor: return UnawareBucketRowKeyExtractor(self.table_schema) elif bucket_mode == BucketMode.POSTPONE_MODE: return PostponeBucketRowKeyExtractor(self.table_schema) - elif bucket_mode == BucketMode.HASH_DYNAMIC or bucket_mode == BucketMode.CROSS_PARTITION: - return DynamicBucketRowKeyExtractor(self.table_schema) + elif bucket_mode == BucketMode.HASH_DYNAMIC: + return DynamicBucketRowKeyExtractor( + self.table_schema, + table=self, + ignore_existing=ignore_existing, + ) + elif bucket_mode == BucketMode.CROSS_PARTITION: + raise ValueError( + "CROSS_PARTITION primary-key writes require a persistent " + "global primary-key index, which PyPaimon does not yet support" + ) else: raise ValueError(f"Unsupported bucket mode: {bucket_mode}") 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 new file mode 100644 index 000000000000..7c327a3a664a --- /dev/null +++ b/paimon-python/pypaimon/tests/daft/daft_pk_distributed_write_test.py @@ -0,0 +1,623 @@ +################################################################################ +# 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. +################################################################################ + +"""Correctness tests for distributed Daft writes to primary-key tables.""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +import textwrap +from unittest.mock import Mock + +import pyarrow as pa +import pytest + +pypaimon = pytest.importorskip("pypaimon") +daft = pytest.importorskip("daft") + +from daft.recordbatch.micropartition import MicroPartition +from pypaimon.daft import daft_datasink +from pypaimon.daft.daft_paimon import _read_table, _write_table +from pypaimon.daft.daft_datasink import PaimonDataSink +from pypaimon.index.index_file_handler import IndexFileHandler + + +@pytest.fixture +def catalog(tmp_path): + result = pypaimon.CatalogFactory.create({"warehouse": str(tmp_path)}) + result.create_database("test_db", ignore_if_exists=True) + return result + + +def _create_pk_table(catalog, name, *, partitioned=False, options=None): + fields = [ + pa.field("id", pa.int64()), + pa.field("value", pa.string()), + ] + partition_keys = [] + primary_keys = ["id"] + if partitioned: + fields.append(pa.field("dt", pa.string())) + partition_keys = ["dt"] + primary_keys.append("dt") + + schema = pypaimon.Schema.from_pyarrow_schema( + pa.schema(fields), + partition_keys=partition_keys, + primary_keys=primary_keys, + options={"file.format": "parquet", **(options or {})}, + ) + identifier = f"test_db.{name}" + catalog.create_table(identifier, schema, ignore_if_exists=False) + return catalog.get_table(identifier) + + +def test_hash_fixed_pk_uses_one_writer_per_partition_bucket(catalog): + table = _create_pk_table( + catalog, + "fixed_pk", + partitioned=True, + options={"bucket": "1"}, + ) + data = daft.from_pydict( + { + "id": [1, 1], + "value": ["old", "new"], + "dt": ["2026-07-20", "2026-07-20"], + } + ).into_partitions(2) + + summary = _write_table(data, table).to_pydict() + + # One complete (partition, bucket) group must be consumed by one + # BatchTableWrite. Two independent writers would produce two files with + # overlapping sequence-number ranges for this primary key. + assert len(summary["file_name"]) == 1 + assert sum(summary["rows"]) == 1 + + +def test_bare_datasink_supports_single_partition_primary_key_write(catalog): + table = _create_pk_table( + catalog, + "bare_sink_pk", + options={"bucket": "1"}, + ) + + summary = daft.from_pydict( + {"id": [1, 2], "value": ["a", "b"]} + ).write_sink(PaimonDataSink(table)).to_pydict() + + assert sum(summary["rows"]) == 2 + assert _read_table(table).sort("id").to_pydict() == { + "id": [1, 2], + "value": ["a", "b"], + } + + +def test_bare_datasink_rejects_multiple_primary_key_write_tasks(catalog): + table = _create_pk_table( + catalog, + "bare_sink_distributed_pk", + options={"bucket": "1"}, + ) + sink = PaimonDataSink(table) + results = [] + for value in (1, 2): + micropartition = MicroPartition.from_arrow( + pa.table({"id": [value], "value": [f"v-{value}"]}) + ) + results.extend(sink.write(iter([micropartition]))) + + with pytest.raises( + ValueError, match="require a single non-empty Daft write task" + ): + sink.finalize(results) + + assert table.snapshot_manager().get_latest_snapshot() is None + + +def test_bare_datasink_supports_single_task_dynamic_bucket(catalog): + table = _create_pk_table( + catalog, + "bare_sink_dynamic_pk", + options={"bucket": "-1"}, + ) + sink = PaimonDataSink(table) + micropartition = MicroPartition.from_arrow( + pa.table({"id": [1], "value": ["v-1"]}) + ) + results = list(sink.write(iter([micropartition]))) + + summary = sink.finalize(results).to_pydict() + + assert sum(summary["rows"]) == 1 + assert _read_table(table).to_pydict() == { + "id": [1], + "value": ["v-1"], + } + + +def test_bare_datasink_rejects_multi_task_dynamic_bucket(catalog): + table = _create_pk_table( + catalog, + "bare_sink_distributed_dynamic_pk", + options={"bucket": "-1"}, + ) + sink = PaimonDataSink(table) + results = [] + for value in (1, 2): + micropartition = MicroPartition.from_arrow( + pa.table({"id": [value], "value": [f"v-{value}"]}) + ) + results.extend(sink.write(iter([micropartition]))) + + with pytest.raises( + ValueError, match="require a single non-empty Daft write task" + ): + sink.finalize(results) + + assert table.snapshot_manager().get_latest_snapshot() is None + + +def test_commit_exception_does_not_abort_possibly_committed_files(catalog): + table = _create_pk_table( + catalog, + "uncertain_commit_pk", + options={"bucket": "1"}, + ) + sink = PaimonDataSink(table) + micropartition = MicroPartition.from_arrow( + pa.table({"id": [1], "value": ["committed"]}) + ) + write_results = list(sink.write(iter([micropartition]))) + + real_commit = sink._write_builder.new_commit() + + def commit_then_raise(messages): + real_commit.commit(messages) + raise RuntimeError("callback failed after snapshot commit") + + uncertain_commit = Mock() + uncertain_commit.commit.side_effect = commit_then_raise + uncertain_commit.close.side_effect = real_commit.close + sink._write_builder.new_commit = Mock(return_value=uncertain_commit) + + with pytest.raises(RuntimeError, match="callback failed"): + sink.finalize(write_results) + + uncertain_commit.abort.assert_not_called() + uncertain_commit.close.assert_called_once_with() + assert _read_table(table).to_pydict() == { + "id": [1], + "value": ["committed"], + } + + +def test_postpone_mode_primary_key_write_preserves_existing_path(catalog): + table = _create_pk_table( + catalog, + "postpone_pk", + options={"bucket": "-2"}, + ) + + summary = _write_table( + daft.from_pydict( + {"id": [1, 2], "value": ["a", "b"]} + ).into_partitions(2), + table, + ).to_pydict() + + assert sum(summary["rows"]) == 2 + assert table.snapshot_manager().get_latest_snapshot() is not None + + +def test_group_udfs_reuse_worker_table_metadata(catalog, monkeypatch): + table = _create_pk_table( + catalog, + "worker_table_cache_pk", + options={"bucket": "1"}, + ) + catalog_options = daft_datasink._extract_catalog_options(table) + identifier = daft_datasink._extract_identifier(table) + schema_token = ( + table.table_schema.id, + table.table_schema.time_millis, + ) + loader = Mock(wraps=daft_datasink._load_table) + monkeypatch.setattr(daft_datasink, "_load_table", loader) + daft_datasink._WORKER_TABLE_CACHE.clear() + + first = daft_datasink._load_worker_table( + catalog_options, identifier, table.table_path, schema_token + ) + second = daft_datasink._load_worker_table( + catalog_options, identifier, table.table_path, schema_token + ) + + assert first is second + loader.assert_called_once_with( + catalog_options, identifier, table.table_path + ) + daft_datasink._WORKER_TABLE_CACHE.clear() + + +def test_hash_dynamic_pk_restores_bucket_mapping_across_commits(catalog): + table = _create_pk_table( + catalog, + "dynamic_pk", + options={ + "bucket": "-1", + "dynamic-bucket.target-row-num": "1", + }, + ) + + _write_table( + daft.from_pydict({"id": [1], "value": ["old"]}), + table, + ) + _write_table( + daft.from_pydict( + { + # A fresh task-local SimpleHashBucketAssigner puts id=2 in + # bucket 0 and incorrectly moves the existing id=1 to bucket 1. + "id": [2, 1], + "value": ["other", "new"], + } + ), + table, + ) + + result = _read_table(table).sort("id").to_pydict() + assert result == {"id": [1, 2], "value": ["new", "other"]} + + snapshot = table.snapshot_manager().get_latest_snapshot() + hash_indexes = [ + entry + for entry in IndexFileHandler(table).scan(snapshot) + if entry.index_file.index_type == "HASH" + ] + assert hash_indexes + assert len({(tuple(e.partition.values), e.bucket) for e in hash_indexes}) == len( + hash_indexes + ) + + +def test_hash_dynamic_pins_one_driver_snapshot_for_all_group_udfs( + catalog, monkeypatch +): + table = _create_pk_table( + catalog, + "dynamic_pinned_snapshot", + options={"bucket": "-1"}, + ) + _write_table( + daft.from_pydict({"id": [1], "value": ["old"]}), table + ).to_pydict() + expected_snapshot_id = table.snapshot_manager().get_latest_snapshot().id + assignment_factory = Mock( + wraps=daft_datasink.make_dynamic_bucket_assignment_udf + ) + group_write_factory = Mock(wraps=daft_datasink.make_group_write_udf) + monkeypatch.setattr( + daft_datasink, + "make_dynamic_bucket_assignment_udf", + assignment_factory, + ) + monkeypatch.setattr( + daft_datasink, "make_group_write_udf", group_write_factory + ) + + _write_table( + daft.from_pydict({"id": [1], "value": ["new"]}), table + ).to_pydict() + + assert ( + assignment_factory.call_args.kwargs["base_snapshot_id"] + == expected_snapshot_id + ) + assert ( + group_write_factory.call_args.kwargs["base_snapshot_id"] + == expected_snapshot_id + ) + + +def test_hash_dynamic_partitioned_pk_preserves_partition_columns(catalog): + table = _create_pk_table( + catalog, + "dynamic_partitioned_pk", + partitioned=True, + options={ + "bucket": "-1", + "dynamic-bucket.target-row-num": "1", + }, + ) + + _write_table( + daft.from_pydict( + { + "id": [1, 1], + "value": ["old-a", "old-b"], + "dt": ["a", "b"], + } + ), + table, + ) + _write_table( + daft.from_pydict( + { + "id": [2, 1, 2, 1], + "value": ["other-a", "new-a", "other-b", "new-b"], + "dt": ["a", "a", "b", "b"], + } + ).into_partitions(2), + table, + ) + + rows = _read_table(table).to_pydict() + assert set(zip(rows["id"], rows["value"], rows["dt"])) == { + (1, "new-a", "a"), + (2, "other-a", "a"), + (1, "new-b", "b"), + (2, "other-b", "b"), + } + + +def test_hash_dynamic_pk_overwrite_replaces_hash_index(catalog): + table = _create_pk_table( + catalog, + "dynamic_pk_overwrite", + options={ + "bucket": "-1", + "dynamic-bucket.target-row-num": "10", + }, + ) + _write_table( + daft.from_pydict({"id": [1, 2], "value": ["old-1", "old-2"]}), + table, + ) + + _write_table( + daft.from_pydict({"id": [3], "value": ["first"]}), + table, + mode="overwrite", + ) + snapshot = table.snapshot_manager().get_latest_snapshot() + hash_indexes = [ + entry + for entry in IndexFileHandler(table).scan(snapshot) + if entry.index_file.index_type == "HASH" + ] + assert sum(entry.index_file.row_count for entry in hash_indexes) == 1 + + _write_table( + daft.from_pydict( + {"id": [4, 3], "value": ["other", "updated"]} + ), + table, + ) + assert _read_table(table).sort("id").to_pydict() == { + "id": [3, 4], + "value": ["updated", "other"], + } + + +def test_cross_partition_pk_write_requires_global_index(catalog): + schema = pypaimon.Schema.from_pyarrow_schema( + pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("value", pa.string()), + pa.field("dt", pa.string()), + ] + ), + partition_keys=["dt"], + primary_keys=["id"], + options={"bucket": "-1", "file.format": "parquet"}, + ) + catalog.create_table("test_db.cross_partition_pk", schema, False) + table = catalog.get_table("test_db.cross_partition_pk") + data = daft.from_pydict({ + "id": [1, 2], + "value": ["a", "b"], + "dt": ["2026-07-20", "2026-07-21"], + }).into_partitions(2) + + with pytest.raises( + ValueError, match="CROSS_PARTITION.*global primary-key index" + ): + _write_table(data, table) + + +@pytest.mark.skipif( + importlib.util.find_spec("ray") is None, + reason="Distributed Daft correctness test requires Ray", +) +def test_pk_distributed_write_on_ray(): + """Exercise real multi-task ownership and compatible PK bucket modes.""" + python_root = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..") + ) + script = textwrap.dedent( + r''' + import os + import shutil + import tempfile + + os.environ.setdefault("RAY_TMPDIR", "/tmp") + + import daft + import pyarrow as pa + import ray + from daft import runners + from pypaimon import CatalogFactory, Schema + from pypaimon.daft.daft_paimon import _read_table, _write_table + from pypaimon.write.row_key_extractor import DynamicBucketRowKeyExtractor + + root = tempfile.mkdtemp(prefix="paimon-daft-pk-ray-") + try: + catalog = CatalogFactory.create({"warehouse": root}) + catalog.create_database("test_db", False) + ray.init(num_cpus=2, include_dashboard=False) + runners.set_runner_ray() + + fixed_schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field("id", pa.int64()), + pa.field("value", pa.string()), + pa.field("dt", pa.string()), + ]), + partition_keys=["dt"], + primary_keys=["id", "dt"], + options={"bucket": "1", "file.format": "parquet"}, + ) + catalog.create_table("test_db.fixed_pk", fixed_schema, False) + fixed = catalog.get_table("test_db.fixed_pk") + fixed_summary = _write_table( + daft.from_pydict({ + "id": [1, 1], + "value": ["old", "new"], + "dt": ["2026-07-20", "2026-07-20"], + }).into_partitions(2), + fixed, + ).to_pydict() + assert len(fixed_summary["file_name"]) == 1 + assert sum(fixed_summary["rows"]) == 1 + + postpone_schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field("id", pa.int64()), + pa.field("value", pa.string()), + ]), + primary_keys=["id"], + options={"bucket": "-2", "file.format": "parquet"}, + ) + catalog.create_table( + "test_db.postpone_pk", postpone_schema, False + ) + postpone = catalog.get_table("test_db.postpone_pk") + postpone_summary = _write_table( + daft.from_pydict({ + "id": [1, 2, 3, 4], + "value": ["a", "b", "c", "d"], + }).into_partitions(2), + postpone, + ).to_pydict() + assert sum(postpone_summary["rows"]) == 4 + + cross_partition_schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field("id", pa.int64()), + pa.field("value", pa.string()), + pa.field("dt", pa.string()), + ]), + partition_keys=["dt"], + primary_keys=["id"], + options={"bucket": "-1", "file.format": "parquet"}, + ) + catalog.create_table( + "test_db.cross_partition_pk", + cross_partition_schema, + False, + ) + cross_partition = catalog.get_table( + "test_db.cross_partition_pk" + ) + try: + _write_table( + daft.from_pydict({ + "id": [1, 1], + "value": ["old", "new"], + "dt": ["p1", "p2"], + }).into_partitions(2), + cross_partition, + ) + except ValueError as error: + assert "global primary-key index" in str(error) + else: + raise AssertionError("CROSS_PARTITION write was not rejected") + + dynamic_schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field("id", pa.int64()), + pa.field("value", pa.string()), + ]), + primary_keys=["id"], + options={ + "bucket": "-1", + "dynamic-bucket.target-row-num": "100", + "file.format": "parquet", + }, + ) + catalog.create_table("test_db.dynamic_pk", dynamic_schema, False) + dynamic = catalog.get_table("test_db.dynamic_pk") + _write_table( + daft.from_pydict({ + "id": list(range(1, 9)), + "value": [f"old-{i}" for i in range(1, 9)], + }), + dynamic, + ) + assigner_probe = pa.RecordBatch.from_pydict({ + "id": list(range(1, 9)), + "value": [f"new-{i}" for i in range(1, 9)], + }) + assigners = DynamicBucketRowKeyExtractor( + dynamic.table_schema + ).extract_assigners_batch(assigner_probe, 2, 2) + assert set(assigners) == {0, 1} + dynamic_summary = _write_table( + daft.from_pydict({ + "id": list(range(1, 9)), + "value": [f"new-{i}" for i in range(1, 9)], + }).into_partitions(2), + dynamic, + ).to_pydict() + # The first commit used one assigner and placed every key in + # bucket 0. The second uses two assigners; both can resolve old + # keys to bucket 0, but the final bucket shuffle must still leave + # exactly one writer for it. + assert len(dynamic_summary["file_name"]) == 1 + actual = _read_table(dynamic).sort("id").to_pydict() + assert actual == { + "id": list(range(1, 9)), + "value": [f"new-{i}" for i in range(1, 9)], + } + finally: + if ray.is_initialized(): + ray.shutdown() + shutil.rmtree(root, ignore_errors=True) + ''' + ) + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [python_root, env.get("PYTHONPATH", "")] + ).rstrip(os.pathsep) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + env=env, + text=True, + timeout=180, + ) + assert result.returncode == 0, ( + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" + ) diff --git a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py index ac9f44cc22fc..8b51f73431f8 100644 --- a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py +++ b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py @@ -67,6 +67,79 @@ def setUpClass(cls): }) cls.catalog.create_database('default', True) + def test_read_java_dynamic_bucket_hash_index(self): + table = self.catalog.get_table( + 'default.dynamic_hash_java_to_python' + ) + read_builder = table.new_read_builder() + initial = read_builder.new_read().to_arrow( + read_builder.new_scan().plan().splits() + ) + self.assertEqual( + { + 'key1': ['hello-java'], + 'key2': [42], + 'value': ['java-old'], + }, + initial.to_pydict(), + ) + + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(pa.table({ + 'key1': ['python-only', 'hello-java'], + 'key2': pa.array([7, 42], type=pa.int64()), + 'value': ['python-only', 'python-new'], + })) + commit = builder.new_commit() + commit.commit(writer.prepare_commit()) + writer.close() + commit.close() + + read_builder = table.new_read_builder() + result = table_sort_by(read_builder.new_read().to_arrow( + read_builder.new_scan().plan().splits() + ), 'key1') + self.assertEqual( + { + 'key1': ['hello-java', 'python-only'], + 'key2': [42, 7], + 'value': ['python-new', 'python-only'], + }, + result.to_pydict(), + ) + + def test_py_write_dynamic_bucket_hash_index(self): + table_name = 'default.dynamic_hash_python_to_java' + self.catalog.drop_table(table_name, True) + schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field('key1', pa.string()), + pa.field('key2', pa.int64()), + pa.field('value', pa.string()), + ]), + primary_keys=['key1', 'key2'], + options={ + 'bucket': '-1', + 'dynamic-bucket.target-row-num': '1', + 'file.format': 'parquet', + }, + ) + self.catalog.create_table(table_name, schema, False) + table = self.catalog.get_table(table_name) + + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(pa.table({ + 'key1': ['hello-java', 'python-only'], + 'key2': pa.array([42, 7], type=pa.int64()), + 'value': ['python-old', 'python-only'], + })) + commit = builder.new_commit() + commit.commit(writer.prepare_commit()) + writer.close() + commit.close() + @parameterized.expand(get_file_format_params()) def test_py_write_read_append_table(self, file_format): pa_schema = pa.schema([ diff --git a/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py new file mode 100644 index 000000000000..30d2dcaa8580 --- /dev/null +++ b/paimon-python/pypaimon/tests/write/dynamic_bucket_test.py @@ -0,0 +1,669 @@ +################################################################################ +# 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. +################################################################################ + +import datetime +import tempfile +import unittest +from decimal import Decimal +from unittest.mock import Mock, patch + +import pyarrow as pa + +from pypaimon import CatalogFactory, Schema +from pypaimon.index.dynamic_bucket import ( + HashBucketAssigner, + _PartitionIndex, + _iter_hashes, + compute_assigner, + to_signed_int32, +) +from pypaimon.index.index_file_handler import IndexFileHandler +from pypaimon.index.index_file_meta import IndexFileMeta +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.row_key_extractor import DynamicBucketRowKeyExtractor + + +class DynamicBucketTest(unittest.TestCase): + + @staticmethod + def _create_table( + root, name, target_row_num=100, max_buckets=None + ): + catalog = CatalogFactory.create({'warehouse': root}) + catalog.create_database('default', True) + options = { + 'bucket': '-1', + 'dynamic-bucket.target-row-num': str(target_row_num), + 'file.format': 'parquet', + } + if max_buckets is not None: + options['dynamic-bucket.max-buckets'] = str(max_buckets) + schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field('id', pa.int64()), + pa.field('value', pa.string()), + ]), + primary_keys=['id'], + options=options, + ) + catalog.create_table(f'default.{name}', schema, False) + return catalog.get_table(f'default.{name}') + + @staticmethod + def _prepare_indexed_write(table, ids): + builder = table.new_batch_write_builder() + writer = builder.new_write().with_dynamic_bucket_index() + writer.write_arrow_batch(pa.RecordBatch.from_pydict({ + 'id': ids, + 'value': [f'v-{value}' for value in ids], + })) + return writer, builder.new_commit(), writer.prepare_commit() + + @staticmethod + def _hash_indexes(table): + snapshot = table.snapshot_manager().get_latest_snapshot() + return [ + entry for entry in IndexFileHandler(table).scan(snapshot) + if entry.index_file.index_type == 'HASH' + ] + + @staticmethod + def _commit_arrow(table, ids, values): + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(pa.table({'id': ids, 'value': values})) + messages = writer.prepare_commit() + commit = builder.new_commit() + commit.commit(messages) + writer.close() + commit.close() + return messages + + @staticmethod + def _read_arrow(table): + builder = table.new_read_builder() + return builder.new_read().to_arrow(builder.new_scan().plan().splits()) + + def test_compute_assigner_matches_java(self): + max_int = 2 ** 31 - 1 + self.assertEqual(compute_assigner(max_int, 0, 5, 5), 2) + self.assertEqual(compute_assigner(max_int, 1, 5, 5), 3) + self.assertEqual(compute_assigner(max_int, 2, 5, 5), 4) + self.assertEqual(compute_assigner(max_int, 3, 5, 5), 0) + self.assertEqual(compute_assigner(2, 0, 5, 3), 2) + self.assertEqual(compute_assigner(2, 1, 5, 3), 3) + self.assertEqual(compute_assigner(2, 2, 5, 3), 4) + self.assertEqual(compute_assigner(2, 3, 5, 3), 2) + self.assertEqual(compute_assigner(3, 1, 5, 1), 3) + self.assertEqual(compute_assigner(3, 2, 5, 1), 3) + min_int = -(2 ** 31) + self.assertEqual(compute_assigner(min_int, 0, 5, 5), 3) + self.assertEqual(compute_assigner(2, min_int, 5, 5), 0) + + def test_binary_row_hash_matches_java_for_bucket_key_types(self): + # Generated with Java InternalRowSerializer.toBinaryRow(...).hashCode(). + cases = [ + ( + 'inline string', + ('hello',), + [DataField(0, 'key', AtomicType('STRING'))], + 243722546, + ), + ( + 'variable string', + ('hello-java',), + [DataField(0, 'key', AtomicType('STRING'))], + -201703277, + ), + ( + 'composite bucket key', + ('hello-java', 42), + [ + DataField(0, 'key1', AtomicType('STRING')), + DataField(1, 'key2', AtomicType('BIGINT')), + ], + -2066620165, + ), + ( + 'compact decimal', + (Decimal('12345.67'),), + [DataField(0, 'key', AtomicType('DECIMAL(10, 2)'))], + 754928256, + ), + ( + 'variable decimal', + (Decimal('12345678901234567890123.45'),), + [DataField(0, 'key', AtomicType('DECIMAL(25, 2)'))], + 1388205002, + ), + ( + 'compact timestamp', + (datetime.datetime(2026, 1, 2, 3, 4, 5, 123000),), + [DataField(0, 'key', AtomicType('TIMESTAMP(3)'))], + -1766746798, + ), + ( + 'variable timestamp', + (datetime.datetime(2026, 1, 2, 3, 4, 5, 123456),), + [DataField(0, 'key', AtomicType('TIMESTAMP(6)'))], + 1245041971, + ), + ( + 'inline binary', + (bytes([0, 1, 255, 127]),), + [DataField(0, 'key', AtomicType('BYTES'))], + 586821318, + ), + ( + 'variable binary', + (bytes(range(10)),), + [DataField(0, 'key', AtomicType('BYTES'))], + 1822312655, + ), + ] + + for name, values, fields, java_hash in cases: + with self.subTest(name=name): + actual = DynamicBucketRowKeyExtractor._binary_row_hash_code( + values, fields + ) + self.assertEqual(java_hash, to_signed_int32(actual)) + + def test_unbounded_bucket_id_matches_java_short_limit(self): + index = _PartitionIndex({}, {}, 1) + + bucket = index.assign( + key_hash=1, + bucket_filter=lambda candidate: candidate == 32766, + max_buckets_num=-1, + max_bucket_id=32765, + ) + + self.assertEqual(bucket, 32766) + + def test_assigner_rejects_record_owned_by_another_writer(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'wrong_assigner') + assigner = HashBucketAssigner( + table=table, + num_channels=2, + num_assigners=2, + assign_id=0, + target_bucket_row_number=100, + max_buckets_num=-1, + ) + + with self.assertRaisesRegex( + ValueError, 'Record assigner 1 does not match writer assigner 0' + ): + assigner.assign((), partition_hash=0, key_hash=1) + + def test_max_buckets_rejects_assigner_without_bucket(self): + index = _PartitionIndex({}, {}, 1) + + with self.assertRaisesRegex( + RuntimeError, 'No dynamic bucket is available for this assigner' + ): + index.assign( + key_hash=1, + bucket_filter=lambda _: False, + max_buckets_num=1, + max_bucket_id=0, + ) + + def test_unbounded_buckets_rejects_java_short_id_exhaustion(self): + index = _PartitionIndex({}, {}, 1) + + with self.assertRaisesRegex( + RuntimeError, 'No dynamic bucket id remains below Java Short.MAX_VALUE' + ): + index.assign( + key_hash=1, + bucket_filter=lambda _: False, + max_buckets_num=-1, + max_bucket_id=32766, + ) + + def test_corrupt_hash_index_rejects_trailing_bytes(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'corrupt_index') + path = f'{root}/corrupt-hash-index' + with table.file_io.new_output_stream(path) as stream: + stream.write(b'\x00\x00\x00') + entry = IndexManifestEntry( + kind=0, + partition=GenericRow([], []), + bucket=0, + index_file=IndexFileMeta( + index_type='HASH', + file_name='corrupt-hash-index', + file_size=3, + row_count=1, + external_path=path, + ), + ) + + with self.assertRaisesRegex( + RuntimeError, 'expected a multiple of 4 bytes' + ): + list(_iter_hashes(table, entry)) + + def test_regular_dynamic_writer_uses_persistent_index(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'regular_dynamic') + writer = table.new_batch_write_builder().new_write() + + self.assertIs(table, writer.row_key_extractor._table) + writer.write_arrow_batch(pa.RecordBatch.from_pydict({ + 'id': [1], + 'value': ['v-1'], + })) + messages = writer.prepare_commit() + + self.assertTrue(any(message.index_adds for message in messages)) + self.assertFalse(any(message.index_deletes for message in messages)) + + def test_regular_dynamic_writer_restores_mapping_across_commits(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'regular_upsert', target_row_num=1) + self._commit_arrow(table, [1], ['old']) + self._commit_arrow(table, [2, 1], ['other', 'new']) + + result = self._read_arrow(table).sort_by('id').to_pydict() + self.assertEqual({'id': [1, 2], 'value': ['new', 'other']}, result) + + def test_regular_dynamic_writer_retains_only_requested_index_hashes(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'bounded_restore') + self._commit_arrow( + table, list(range(10)), [f'old-{value}' for value in range(10)] + ) + + writer = table.new_batch_write_builder().new_write() + writer.write_arrow(pa.table({'id': [3], 'value': ['new-3']})) + + partition_index = writer.row_key_extractor._assigner._partition_indexes[()] + self.assertEqual(1, len(partition_index.hash_to_bucket)) + self.assertEqual({}, writer.row_key_extractor._index_maintainer._states) + + def test_legacy_dynamic_data_without_hash_index_fails_fast(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'legacy_no_index') + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.row_key_extractor = DynamicBucketRowKeyExtractor( + table.table_schema + ) + writer.write_arrow(pa.table({'id': [1], 'value': ['old']})) + messages = writer.prepare_commit() + builder.new_commit().commit(messages) + writer.close() + + new_writer = table.new_batch_write_builder().new_write() + with self.assertRaisesRegex( + RuntimeError, 'has data files but no complete HASH index' + ): + new_writer.write_arrow(pa.table({'id': [1], 'value': ['new']})) + + def test_cross_partition_write_requires_global_index(self): + with tempfile.TemporaryDirectory() as root: + catalog = CatalogFactory.create({'warehouse': root}) + catalog.create_database('default', True) + schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field('id', pa.int64()), + pa.field('value', pa.string()), + pa.field('dt', pa.string()), + ]), + partition_keys=['dt'], + primary_keys=['id'], + options={'bucket': '-1'}, + ) + catalog.create_table('default.cross_partition', schema, False) + table = catalog.get_table('default.cross_partition') + + with self.assertRaisesRegex( + ValueError, 'CROSS_PARTITION.*global primary-key index' + ): + table.new_batch_write_builder().new_write() + + def test_batch_writer_abort_after_prepare_deletes_hash_index(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'abort_prepared') + writer = table.new_batch_write_builder().new_write() + writer.write_arrow(pa.table({'id': [1], 'value': ['v-1']})) + messages = writer.prepare_commit() + index_path = messages[0].index_adds[0].index_file.external_path + if index_path is None: + index_path = ( + table.path_factory().global_index_path_factory() + .to_path(messages[0].index_adds[0].index_file.file_name) + ) + self.assertTrue(table.file_io.exists(index_path)) + + writer.abort() + + self.assertFalse(table.file_io.exists(index_path)) + + def test_stream_writer_releases_prepared_hash_index_ownership(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'stream_prepared') + builder = table.new_stream_write_builder() + writer = builder.new_write() + commit = builder.new_commit() + writer.write_arrow(pa.table({'id': [1], 'value': ['v-1']})) + messages = writer.prepare_commit(1) + index_path = messages[0].index_adds[0].index_file.external_path + if index_path is None: + index_path = ( + table.path_factory().global_index_path_factory() + .to_path(messages[0].index_adds[0].index_file.file_name) + ) + self.assertEqual( + [], writer.row_key_extractor._index_maintainer._new_paths + ) + + commit.commit(messages, 1) + writer.close() + + self.assertTrue(table.file_io.exists(index_path)) + self.assertEqual( + {'id': [1], 'value': ['v-1']}, + self._read_arrow(table).to_pydict(), + ) + commit.close() + + def test_regular_dynamic_extractor_skips_partition_hash(self): + with tempfile.TemporaryDirectory() as root: + catalog = CatalogFactory.create({'warehouse': root}) + catalog.create_database('default', True) + schema = Schema.from_pyarrow_schema( + pa.schema([ + pa.field('id', pa.int64()), + pa.field('value', pa.string()), + pa.field('dt', pa.string()), + ]), + partition_keys=['dt'], + primary_keys=['id', 'dt'], + options={ + 'bucket': '-1', + 'dynamic-bucket.target-row-num': '100', + }, + ) + catalog.create_table('default.partitioned', schema, False) + table = catalog.get_table('default.partitioned') + extractor = DynamicBucketRowKeyExtractor(table.table_schema) + hash_code = Mock(wraps=extractor._binary_row_hash_code) + extractor._binary_row_hash_code = hash_code + + extractor.extract_partition_bucket_batch( + pa.RecordBatch.from_pydict({ + 'id': [1, 2], + 'value': ['a', 'b'], + 'dt': ['p', 'p'], + }) + ) + + self.assertEqual(2, hash_code.call_count) + self.assertNotIn( + ('p',), + [call.args[0] for call in hash_code.call_args_list], + ) + + def test_concurrent_initial_hash_index_add_conflicts(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'concurrent_initial') + writer1, commit1, messages1 = self._prepare_indexed_write( + table, [1] + ) + writer2, commit2, messages2 = self._prepare_indexed_write( + table, [2] + ) + + commit1.commit(messages1) + stale_data_paths = [ + file.file_path for message in messages2 for file in message.new_files + ] + stale_index_paths = [ + entry.index_file.external_path + or table.path_factory().global_index_path_factory().to_path( + entry.index_file.file_name + ) + for message in messages2 for entry in message.index_adds + ] + with self.assertRaisesRegex( + RuntimeError, 'HASH index assignment conflict' + ): + commit2.commit(messages2) + + self.assertTrue(all( + not table.file_io.exists(path) + for path in stale_data_paths + stale_index_paths + )) + + indexes = self._hash_indexes(table) + self.assertEqual(1, len(indexes)) + self.assertEqual(1, indexes[0].index_file.row_count) + writer1.close() + writer2.close() + commit1.close() + commit2.close() + + def test_concurrent_hash_index_replacement_conflicts(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'concurrent_replace') + seed_writer, seed_commit, seed_messages = ( + self._prepare_indexed_write(table, [1]) + ) + seed_commit.commit(seed_messages) + seed_writer.close() + seed_commit.close() + + writer1, commit1, messages1 = self._prepare_indexed_write( + table, [2] + ) + writer2, commit2, messages2 = self._prepare_indexed_write( + table, [3] + ) + old_index = messages1[0].index_deletes[0].index_file.file_name + self.assertEqual( + old_index, + messages2[0].index_deletes[0].index_file.file_name, + ) + + commit1.commit(messages1) + with self.assertRaisesRegex( + RuntimeError, 'HASH index assignment conflict' + ): + commit2.commit(messages2) + + indexes = self._hash_indexes(table) + self.assertEqual(1, len(indexes)) + self.assertEqual(2, indexes[0].index_file.row_count) + writer1.close() + writer2.close() + commit1.close() + commit2.close() + + def test_concurrent_disjoint_bucket_replacements_conflict(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table( + root, + 'concurrent_disjoint', + target_row_num=1, + max_buckets=2, + ) + seed_writer, seed_commit, seed_messages = ( + self._prepare_indexed_write(table, [1, 2]) + ) + seed_commit.commit(seed_messages) + seed_writer.close() + seed_commit.close() + + with patch( + 'pypaimon.index.dynamic_bucket.random.choice', + return_value=0, + ): + writer1, commit1, messages1 = self._prepare_indexed_write( + table, [3] + ) + with patch( + 'pypaimon.index.dynamic_bucket.random.choice', + return_value=1, + ): + writer2, commit2, messages2 = self._prepare_indexed_write( + table, [3] + ) + + self.assertNotEqual(messages1[0].bucket, messages2[0].bucket) + commit1.commit(messages1) + stale_paths = [ + file.file_path + for message in messages2 + for file in message.new_files + ] + [ + entry.index_file.external_path + or table.path_factory().global_index_path_factory().to_path( + entry.index_file.file_name + ) + for message in messages2 + for entry in message.index_adds + ] + + with self.assertRaisesRegex( + RuntimeError, 'assigned from snapshot.*latest snapshot' + ): + commit2.commit(messages2) + + self.assertTrue(all( + not table.file_io.exists(path) for path in stale_paths + )) + writer1.close() + writer2.close() + commit1.close() + commit2.close() + + def test_data_only_upsert_conflicts_after_overwrite_remaps_key(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table( + root, 'data_only_overwrite_conflict', target_row_num=1 + ) + self._commit_arrow(table, [1, 2], ['one', 'two']) + + stale_builder = table.new_batch_write_builder() + stale_writer = stale_builder.new_write() + stale_writer.write_arrow( + pa.table({'id': [2], 'value': ['stale-upsert']}) + ) + stale_messages = stale_writer.prepare_commit() + stale_commit = stale_builder.new_commit() + self.assertFalse(any( + message.index_adds or message.index_deletes + for message in stale_messages + )) + + overwrite_builder = table.new_batch_write_builder().overwrite({}) + overwrite_writer = overwrite_builder.new_write() + overwrite_writer.write_arrow( + pa.table({'id': [2], 'value': ['overwrite']}) + ) + overwrite_messages = overwrite_writer.prepare_commit() + overwrite_commit = overwrite_builder.new_commit() + overwrite_commit.commit(overwrite_messages) + overwrite_writer.close() + overwrite_commit.close() + + with self.assertRaisesRegex( + RuntimeError, 'HASH index assignment conflict' + ): + stale_commit.commit(stale_messages) + + self.assertEqual( + {'id': [2], 'value': ['overwrite']}, + self._read_arrow(table).to_pydict(), + ) + stale_writer.close() + stale_commit.close() + + def test_retry_then_hash_index_conflict_preserves_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]) + prepared_paths = [ + file.file_path + for message in messages + for file in message.new_files + ] + [ + entry.index_file.external_path + or table.path_factory().global_index_path_factory().to_path( + entry.index_file.file_name + ) + for message in messages + for entry in message.index_adds + ] + original_snapshot_commit = commit.file_store_commit.snapshot_commit + calls = 0 + + def lose_first_compare_and_set(snapshot, statistics): + nonlocal calls + calls += 1 + concurrent_writer, concurrent_commit, concurrent_messages = ( + self._prepare_indexed_write(table, [2]) + ) + concurrent_commit.commit(concurrent_messages) + concurrent_writer.close() + concurrent_commit.close() + return False + + with patch.object( + original_snapshot_commit, + 'commit', + side_effect=lose_first_compare_and_set, + ), patch.object( + commit.file_store_commit, + '_commit_retry_wait', + ): + with self.assertRaisesRegex( + RuntimeError, 'HASH index assignment conflict' + ): + commit.commit(messages) + + self.assertEqual(1, calls) + self.assertTrue(all( + table.file_io.exists(path) for path in prepared_paths + )) + writer.close() + commit.close() + + def test_invalid_dynamic_bucket_key_reports_schema_error(self): + with tempfile.TemporaryDirectory() as root: + table = self._create_table(root, 'invalid_bucket_key') + options = dict(table.table_schema.options) + options['bucket-key'] = 'missing_column' + invalid_schema = table.table_schema.copy(options) + + with self.assertRaisesRegex( + ValueError, 'bucket-key references unknown columns' + ): + DynamicBucketRowKeyExtractor(invalid_schema) + + +if __name__ == '__main__': + unittest.main() diff --git a/paimon-python/pypaimon/write/commit/conflict_detection.py b/paimon-python/pypaimon/write/commit/conflict_detection.py index 30615ec893a5..3534960c1c39 100644 --- a/paimon-python/pypaimon/write/commit/conflict_detection.py +++ b/paimon-python/pypaimon/write/commit/conflict_detection.py @@ -149,6 +149,10 @@ def __init__(self, range_, field_ids): self.field_ids = field_ids +class CommitConflictError(RuntimeError): + """A deterministic pre-snapshot conflict which is safe to abort.""" + + class ConflictDetection: """Detects conflicts between base and delta files during commit.""" @@ -177,6 +181,13 @@ def has_row_id_check_from_snapshot(self): def has_global_index_additions(index_entries=None): return bool(ConflictDetection.global_index_file_additions(index_entries)) + @staticmethod + def has_hash_index_changes(index_entries=None): + return any( + entry.index_file.index_type == IndexManifestFile.HASH_INDEX + for entry in (index_entries or []) + ) + def check_conflicts( self, latest_snapshot, @@ -214,6 +225,11 @@ def check_conflicts( if conflict is not None: return conflict + conflict = self.check_hash_index_conflicts( + latest_snapshot, delta_index_entries) + if conflict is not None: + return conflict + if commit_kind != "COMPACT": next_row_id = latest_snapshot.next_row_id if latest_snapshot else None conflict = self.check_row_id_existence( @@ -232,6 +248,78 @@ def check_conflicts( return self.check_row_id_from_snapshot(latest_snapshot, delta_entries) + def check_hash_index_conflicts( + self, latest_snapshot, delta_index_entries=None): + """Detect stale full-file replacements of dynamic-bucket HASH indexes.""" + hash_entries = [ + entry for entry in (delta_index_entries or []) + if entry.index_file.index_type == IndexManifestFile.HASH_INDEX + ] + if not hash_entries: + return None + + delete_entries = [entry for entry in hash_entries if entry.kind == 1] + add_entries = [entry for entry in hash_entries if entry.kind == 0] + delete_names = { + entry.index_file.file_name for entry in delete_entries + } + + current_entries = [] + if latest_snapshot is not None and latest_snapshot.index_manifest is not None: + current_entries = [ + entry for entry in IndexManifestFile(self.table).read( + latest_snapshot.index_manifest) + if entry.kind == 0 + and entry.index_file.index_type == IndexManifestFile.HASH_INDEX + ] + + current_names = { + entry.index_file.file_name for entry in current_entries + } + for delete in delete_entries: + if delete.index_file.file_name not in current_names: + return RuntimeError( + "HASH index conflict detected: index file {} is not " + "present in the latest snapshot.".format( + delete.index_file.file_name + ) + ) + + additions_by_bucket = {} + for add in add_entries: + key = (tuple(add.partition.values), add.bucket) + previous_add = additions_by_bucket.get(key) + if previous_add is not None: + return RuntimeError( + "HASH index conflict detected: multiple index files {} " + "and {} were added for partition {}, bucket {} in one " + "commit.".format( + previous_add.index_file.file_name, + add.index_file.file_name, + key[0], + key[1], + ) + ) + additions_by_bucket[key] = add + + retained = [ + entry for entry in current_entries + if entry.index_file.file_name not in delete_names + and tuple(entry.partition.values) == key[0] + and entry.bucket == key[1] + ] + if retained: + return RuntimeError( + "HASH index conflict detected: partition {}, bucket {} " + "already has newer index file {}.".format( + key[0], + key[1], + retained[0].index_file.file_name, + ) + ) + + return None + def check_deletion_vector_index_conflicts(self, latest_snapshot, delta_index_entries=None, diff --git a/paimon-python/pypaimon/write/commit_message.py b/paimon-python/pypaimon/write/commit_message.py index c0d4e469ec16..3076b014b16d 100644 --- a/paimon-python/pypaimon/write/commit_message.py +++ b/paimon-python/pypaimon/write/commit_message.py @@ -34,6 +34,7 @@ class CommitMessage: index_adds: List['IndexManifestEntry'] = field(default_factory=list) index_deletes: List['IndexManifestEntry'] = field(default_factory=list) changelog_files: List[DataFileMeta] = field(default_factory=list) + hash_index_base_snapshot: 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 1147d7385866..801771bf14ac 100644 --- a/paimon-python/pypaimon/write/file_store_commit.py +++ b/paimon-python/pypaimon/write/file_store_commit.py @@ -38,7 +38,10 @@ from pypaimon.table.row.offset_row import OffsetRow from pypaimon.write.commit.commit_rollback import CommitRollback from pypaimon.write.commit.commit_scanner import CommitScanner -from pypaimon.write.commit.conflict_detection import ConflictDetection +from pypaimon.write.commit.conflict_detection import ( + CommitConflictError, + ConflictDetection, +) from pypaimon.write.commit.overwrite_changes_provider import OverwriteChangesProvider from pypaimon.table.special_fields import SpecialFields from pypaimon.write.commit_callback import CommitCallback, CommitCallbackContext @@ -150,6 +153,9 @@ def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): for msg in commit_messages: index_deletes.extend(msg.index_deletes) index_adds.extend(msg.index_adds) + hash_index_base_snapshot = self._hash_index_base_snapshot( + commit_messages + ) if not index_deletes: from pypaimon.write.global_index_update_checker import ( @@ -185,6 +191,9 @@ def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): allow_rollback = True if self.conflict_detection.has_global_index_additions(index_adds): detect_conflicts = True + if self.conflict_detection.has_hash_index_changes( + index_adds + index_deletes): + detect_conflicts = True self._try_commit(commit_kind=commit_kind, commit_identifier=commit_identifier, @@ -193,7 +202,8 @@ def commit(self, commit_messages: List[CommitMessage], commit_identifier: int): detect_conflicts=detect_conflicts, allow_rollback=allow_rollback, index_deletes=index_deletes, - index_adds=index_adds) + index_adds=index_adds, + hash_index_base_snapshot=hash_index_base_snapshot) def overwrite(self, overwrite_partition, commit_messages: List[CommitMessage], commit_identifier: int): """Commit the given commit messages in overwrite mode.""" @@ -216,8 +226,20 @@ def overwrite(self, overwrite_partition, commit_messages: List[CommitMessage], c partition_filter = self._create_static_partition_filter(overwrite_partition, commit_messages) changelog_entries = self._collect_changelog_entries(commit_messages) + index_adds = [ + entry for message in commit_messages for entry in message.index_adds + ] + index_deletes = [ + entry for message in commit_messages for entry in message.index_deletes + ] + hash_index_base_snapshot = self._hash_index_base_snapshot( + commit_messages + ) if not skip_overwrite: + index_deletes = self._overwrite_hash_index_deletes( + partition_filter, index_deletes + ) provider = self._overwrite_changes_provider(partition_filter, commit_messages) self._try_commit( commit_kind="OVERWRITE", @@ -226,8 +248,52 @@ def overwrite(self, overwrite_partition, commit_messages: List[CommitMessage], c changelog_entries=changelog_entries, detect_conflicts=True, allow_rollback=False, + index_deletes=index_deletes, + index_adds=index_adds, + hash_index_base_snapshot=hash_index_base_snapshot, ) + @staticmethod + def _hash_index_base_snapshot( + commit_messages: List[CommitMessage], + ) -> Optional[int]: + # Include data-only dynamic-bucket upserts. Their existing mappings + # are stable across append commits, but a concurrent overwrite may + # rebuild the HASH index and move a key to another bucket. + base_snapshots = [ + getattr(message, "hash_index_base_snapshot", None) + for message in commit_messages + if getattr(message, "hash_index_base_snapshot", None) is not None + ] + return min(base_snapshots) if base_snapshots else None + + def _overwrite_hash_index_deletes(self, partition_filter, deletes): + """Delete HASH indexes for every partition replaced by overwrite.""" + from pypaimon.index.dynamic_bucket import HASH_INDEX + from pypaimon.index.index_file_handler import IndexFileHandler + from pypaimon.manifest.index_manifest_entry import IndexManifestEntry + from pypaimon.table.bucket_mode import BucketMode + + if self.table.bucket_mode() != BucketMode.HASH_DYNAMIC: + return deletes + + by_file_name = {entry.index_file.file_name: entry for entry in deletes} + snapshot = self.snapshot_manager.get_latest_snapshot() + for entry in IndexFileHandler(self.table).scan(snapshot): + if entry.index_file.index_type != HASH_INDEX: + continue + if partition_filter is not None and not partition_filter.test( + entry.partition + ): + continue + by_file_name[entry.index_file.file_name] = IndexManifestEntry( + kind=1, + partition=entry.partition, + bucket=entry.bucket, + index_file=entry.index_file, + ) + return list(by_file_name.values()) + def drop_partitions(self, partitions: List[Dict[str, str]], commit_identifier: int) -> None: if not partitions: raise ValueError("Partitions list cannot be empty.") @@ -289,7 +355,8 @@ def truncate_table(self, commit_identifier: int) -> None: def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, detect_conflicts=False, allow_rollback=False, index_deletes=None, - index_adds=None, changelog_entries=None): + index_adds=None, changelog_entries=None, + hash_index_base_snapshot=None): retry_count = 0 retry_result = None @@ -314,6 +381,7 @@ def _try_commit(self, commit_kind, commit_identifier, commit_entries_plan, allow_rollback=allow_rollback, index_deletes=index_deletes, index_adds=index_adds, + hash_index_base_snapshot=hash_index_base_snapshot, ) if result.is_success(): @@ -369,11 +437,27 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str detect_conflicts: bool = False, allow_rollback: bool = False, index_deletes=None, - index_adds=None) -> CommitResult: + index_adds=None, + hash_index_base_snapshot=None) -> CommitResult: start_millis = int(time.time() * 1000) if self._is_duplicate_commit(retry_result, latest_snapshot, commit_identifier, commit_kind): return SuccessResult() + latest_snapshot_id = latest_snapshot.id if latest_snapshot else 0 + if ( + hash_index_base_snapshot is not None + and latest_snapshot_id != hash_index_base_snapshot + ): + conflict = RuntimeError( + "HASH index assignment conflict detected: assigned from " + "snapshot {}, but the latest snapshot is {}.".format( + hash_index_base_snapshot, latest_snapshot_id + ) + ) + if retry_result is None: + raise CommitConflictError(str(conflict)) from conflict + raise conflict + unique_id = uuid.uuid4() base_manifest_list = f"manifest-list-{unique_id}-0" delta_manifest_list = f"manifest-list-{unique_id}-1" @@ -388,9 +472,10 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str # Base entries for conflict detection. On retry, reuse the previous # attempt's base + read only the incremental changes (mirrors Java). base_data_files = None - if detect_conflicts and latest_snapshot is not None: + if detect_conflicts: incremental = None - if (retry_result is not None + if (latest_snapshot is not None + and retry_result is not None and retry_result.latest_snapshot is not None and retry_result.base_data_files is not None): incremental = self.commit_scanner.read_incremental_changes( @@ -403,11 +488,13 @@ def _try_commit_once(self, retry_result: Optional[RetryResult], commit_kind: str if incremental: base_data_files.extend(incremental) base_data_files = FileEntry.merge_entries(base_data_files) - else: + elif latest_snapshot is not None: # First attempt, or incremental could not be built (missing # snapshot): scan the changed partitions in full. base_data_files = self.commit_scanner.read_all_entries_from_changed_partitions( latest_snapshot, commit_entries, index_entries) + else: + base_data_files = [] conflict_exception = self.conflict_detection.check_conflicts( latest_snapshot, @@ -423,6 +510,13 @@ 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: + raise CommitConflictError( + str(conflict_exception) + ) from conflict_exception + # A previous attempt may have committed despite returning an + # error. Preserve the generic, uncertain-result semantics so + # callers do not delete files which a snapshot may reference. raise conflict_exception # Apply row tracking logic after conflict detection (matches Java ordering) diff --git a/paimon-python/pypaimon/write/row_key_extractor.py b/paimon-python/pypaimon/write/row_key_extractor.py index b74ec5f03c7a..5436d5f140cc 100644 --- a/paimon-python/pypaimon/write/row_key_extractor.py +++ b/paimon-python/pypaimon/write/row_key_extractor.py @@ -19,11 +19,12 @@ import random import struct from abc import ABC, abstractmethod -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Optional, Tuple import pyarrow as pa from pypaimon.common.options.core_options import CoreOptions +from pypaimon.index.dynamic_bucket import SHORT_MAX_VALUE, is_my_bucket from pypaimon.schema.table_schema import TableSchema from pypaimon.table.bucket_mode import BucketMode from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer @@ -90,6 +91,10 @@ def extract_partition_bucket_batch(self, data: pa.RecordBatch) -> Tuple[List[Tup buckets = self._extract_buckets_batch(data) return partitions, buckets + def extract_partitions_batch(self, data: pa.RecordBatch) -> List[Tuple]: + """Return partition tuples without calculating bucket hashes.""" + return self._extract_partitions_batch(data) + def extract_partition_bucket_row( self, values_by_name: Dict[str, Any]) -> Tuple[Tuple, int]: partition = tuple( @@ -118,6 +123,14 @@ def _extract_partitions_batch(self, data: pa.RecordBatch) -> List[Tuple]: return partitions + @staticmethod + def _binary_row_hash_code(values: Tuple, fields: List) -> int: + return _hash_bytes_by_words( + GenericRowSerializer.to_bytes( + GenericRow(list(values), fields, RowKind.INSERT) + )[4:] + ) + @abstractmethod def _extract_buckets_batch(self, table: pa.RecordBatch) -> List[int]: """Extract bucket numbers for all rows. Must be implemented by subclasses.""" @@ -148,7 +161,10 @@ def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: columns = [data.column(i) for i in self.bucket_key_indices] return [ _bucket_from_hash( - self._binary_row_hash_code(tuple(col[row_idx].as_py() for col in columns)), + self._binary_row_hash_code( + tuple(col[row_idx].as_py() for col in columns), + self._bucket_key_fields, + ), self.num_buckets, ) for row_idx in range(data.num_rows) @@ -157,16 +173,12 @@ def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: return _bucket_from_hash( self._binary_row_hash_code( - tuple(values_by_name[name] for name in self.bucket_keys) + tuple(values_by_name[name] for name in self.bucket_keys), + self._bucket_key_fields, ), self.num_buckets, ) - def _binary_row_hash_code(self, row_values: Tuple) -> int: - row = GenericRow(list(row_values), self._bucket_key_fields, RowKind.INSERT) - serialized = GenericRowSerializer.to_bytes(row) - return _hash_bytes_by_words(serialized[4:]) - class UnawareBucketRowKeyExtractor(RowKeyExtractor): """Extractor for unaware bucket mode (bucket = -1, no primary keys).""" @@ -185,13 +197,6 @@ def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: return 0 -_SHORT_MAX_VALUE = 32767 - - -def _is_my_bucket(bucket: int, num_assigners: int, assign_id: int) -> bool: - return bucket % num_assigners == assign_id % num_assigners - - def _pick_randomly(bucket_list: List[int]) -> int: return random.choice(bucket_list) @@ -242,8 +247,8 @@ def assign( def _load_new_bucket( self, max_buckets_num: int, num_assigners: int, assign_id: int ) -> None: - for i in range(_SHORT_MAX_VALUE): - if _is_my_bucket(i, num_assigners, assign_id) and ( + for i in range(SHORT_MAX_VALUE): + if is_my_bucket(i, num_assigners, assign_id) and ( i not in self.bucket_information ): if max_buckets_num == -1 or i <= max_buckets_num - 1: @@ -282,74 +287,293 @@ def assign(self, partition: Tuple, hash_value: int) -> int: class DynamicBucketRowKeyExtractor(RowKeyExtractor): - """ - Row key extractor for dynamic bucket mode - Ensures bucket configuration is set to -1 and prevents bucket extraction - """ + """Extract dynamic buckets and maintain their persistent hash mapping.""" - def __init__(self, table_schema: 'TableSchema'): + def __init__( + self, + table_schema: 'TableSchema', + table=None, + num_channels: int = 1, + num_assigners: int = 1, + assign_id: int = 0, + ignore_existing: bool = False, + maintain_index: bool = True, + base_snapshot_id: Optional[int] = None, + ): super().__init__(table_schema) num_buckets = int(table_schema.options.get(CoreOptions.BUCKET.key(), -1)) - if num_buckets != -1: raise ValueError( - f"Only 'bucket' = '-1' is allowed for 'DynamicBucketRowKeyExtractor', but found: {num_buckets}" + "Only 'bucket' = '-1' is allowed for " + f"'DynamicBucketRowKeyExtractor', but found: {num_buckets}" ) + opts = CoreOptions.from_dict(table_schema.options) - self._assigner = SimpleHashBucketAssigner( - num_assigners=1, - assign_id=0, - target_bucket_row_number=opts.dynamic_bucket_target_row_num(), - max_buckets_num=opts.dynamic_bucket_max_buckets(), - ) - # TODO: extract bucket key init logic to base class (shared with FixedBucketRowKeyExtractor) - bucket_key_option = opts.bucket_key() - if bucket_key_option and bucket_key_option.strip(): - self.bucket_keys = [k.strip() for k in bucket_key_option.split(',')] - else: - self.bucket_keys = [ - pk for pk in table_schema.primary_keys - if pk not in table_schema.partition_keys - ] + self._table = table + self.base_snapshot_id = 0 + target_bucket_row_number = opts.dynamic_bucket_target_row_num() + max_buckets_num = opts.dynamic_bucket_max_buckets() + + self.bucket_keys = table_schema.bucket_keys self.bucket_key_indices = self._get_field_indices(self.bucket_keys) - field_map = {f.name: f for f in table_schema.fields} - self._bucket_key_fields = [ - field_map[name] for name in self.bucket_keys if name in field_map + self._bucket_key_fields = table_schema.logical_bucket_key_fields + self._partition_fields = [ + table_schema.fields[index] for index in self.partition_indices ] - def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: + if table is None: + self._assigner = SimpleHashBucketAssigner( + num_assigners=num_assigners, + assign_id=assign_id, + target_bucket_row_number=target_bucket_row_number, + max_buckets_num=max_buckets_num, + ) + self._index_maintainer = None + else: + from pypaimon.index.dynamic_bucket import ( + DynamicBucketIndexMaintainer, + HashBucketAssigner, + ) + + if base_snapshot_id is None: + snapshot = table.snapshot_manager().get_latest_snapshot() + elif base_snapshot_id == 0: + snapshot = None + else: + snapshot = table.snapshot_manager().get_snapshot_by_id( + base_snapshot_id + ) + self.base_snapshot_id = snapshot.id if snapshot is not None else 0 + self._assigner = HashBucketAssigner( + table=table, + num_channels=num_channels, + num_assigners=num_assigners, + assign_id=assign_id, + target_bucket_row_number=target_bucket_row_number, + max_buckets_num=max_buckets_num, + ignore_existing=ignore_existing, + snapshot=snapshot, + ) + self._index_maintainer = ( + DynamicBucketIndexMaintainer( + table, + ignore_existing=ignore_existing, + snapshot=snapshot, + ) + if maintain_index + else None + ) + + def extract_hashes_batch( + self, data: pa.RecordBatch + ) -> Tuple[List[Tuple], List[int], List[int]]: + """Return partitions, BinaryRow partition hashes, and key hashes.""" partitions = self._extract_partitions_batch(data) - columns = [data.column(i) for i in self.bucket_key_indices] - buckets = [] - for row_idx in range(data.num_rows): - key_hash = _hash_bytes_by_words( - GenericRowSerializer.to_bytes( - GenericRow( - [columns[j][row_idx].as_py() for j in range(len(columns))], - self._bucket_key_fields, - RowKind.INSERT, - ) - )[4:] + key_hashes = self._extract_key_hashes_batch(data) + partition_hash_cache = {} + partition_hashes = [] + for partition in partitions: + if partition not in partition_hash_cache: + partition_hash_cache[partition] = self._binary_row_hash_code( + partition, self._partition_fields + ) + partition_hashes.append(partition_hash_cache[partition]) + return partitions, partition_hashes, key_hashes + + def _extract_key_hashes_batch(self, data: pa.RecordBatch) -> List[int]: + key_columns = [data.column(i) for i in self.bucket_key_indices] + return [ + self._binary_row_hash_code( + tuple(column[row_idx].as_py() for column in key_columns), + self._bucket_key_fields, + ) + for row_idx in range(data.num_rows) + ] + + def extract_assigners_batch( + self, data: pa.RecordBatch, num_channels: int, num_assigners: int + ) -> List[int]: + from pypaimon.index.dynamic_bucket import compute_assigner + + _, partition_hashes, key_hashes = self.extract_hashes_batch(data) + return [ + compute_assigner( + partition_hash, key_hash, num_channels, num_assigners ) - buckets.append( - self._assigner.assign(partitions[row_idx], key_hash)) + for partition_hash, key_hash in zip(partition_hashes, key_hashes) + ] + + def extract_partition_bucket_from_hashes_batch( + self, data: pa.RecordBatch, key_hashes: List[int] + ) -> Tuple[List[Tuple], List[int]]: + """Assign buckets using key hashes calculated by an upstream stage.""" + partitions, buckets, _ = ( + self.extract_partition_bucket_status_from_hashes_batch( + data, key_hashes + ) + ) + return partitions, buckets + + def extract_partition_bucket_status_from_hashes_batch( + self, data: pa.RecordBatch, key_hashes: List[int] + ) -> Tuple[List[Tuple], List[int], List[bool]]: + """Assign carried hashes and report whether each mapping is new.""" + from pypaimon.index.dynamic_bucket import to_signed_int32 + + if len(key_hashes) != data.num_rows: + raise ValueError( + "Precomputed key hash count {} does not match row count {}".format( + len(key_hashes), data.num_rows + ) + ) + partitions = self._extract_partitions_batch(data) + normalized_hashes = [to_signed_int32(value) for value in key_hashes] + if self._table is None: + buckets = [ + self._assigner.assign(partition, key_hash) + for partition, key_hash in zip(partitions, normalized_hashes) + ] + return partitions, buckets, [True] * len(buckets) + + partition_hash_cache = {} + partition_hashes = [] + for partition in partitions: + partition_hash = partition_hash_cache.get(partition) + if partition_hash is None: + partition_hash = self._binary_row_hash_code( + partition, self._partition_fields + ) + partition_hash_cache[partition] = partition_hash + partition_hashes.append(partition_hash) + assignments = self._assigner.assign_batch( + partitions, partition_hashes, normalized_hashes + ) + buckets = [assignment[0] for assignment in assignments] + new_mappings = [assignment[1] for assignment in assignments] + if self._index_maintainer is not None: + for partition, bucket, key_hash, is_new in zip( + partitions, buckets, normalized_hashes, new_mappings + ): + if is_new: + self._index_maintainer.notify_new_record( + partition, bucket, key_hash + ) + return partitions, buckets, new_mappings + + def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]: + if self._table is None: + partitions = self._extract_partitions_batch(data) + key_hashes = self._extract_key_hashes_batch(data) + return [ + self._assigner.assign(partition, key_hash) + for partition, key_hash in zip(partitions, key_hashes) + ] + + partitions, partition_hashes, key_hashes = self.extract_hashes_batch(data) + assignments = self._assigner.assign_batch( + partitions, partition_hashes, key_hashes + ) + buckets = [assignment[0] for assignment in assignments] + if self._index_maintainer is not None: + for partition, key_hash, (bucket, is_new) in zip( + partitions, key_hashes, assignments + ): + if is_new: + self._index_maintainer.notify_new_record( + partition, bucket, key_hash + ) return buckets def _extract_bucket_row(self, values_by_name: Dict[str, Any]) -> int: - key_hash = _hash_bytes_by_words( - GenericRowSerializer.to_bytes( - GenericRow( - [values_by_name[name] for name in self.bucket_keys], - self._bucket_key_fields, - RowKind.INSERT, - ) - )[4:] + key_hash = self._binary_row_hash_code( + tuple(values_by_name[name] for name in self.bucket_keys), + self._bucket_key_fields, ) partition = tuple( values_by_name[self.table_schema.fields[i].name] for i in self.partition_indices ) - return self._assigner.assign(partition, key_hash) + if self._table is None: + return self._assigner.assign(partition, key_hash) + partition_hash = self._binary_row_hash_code( + partition, self._partition_fields + ) + bucket, is_new = self._assigner.assign_with_status( + partition, partition_hash, key_hash + ) + if self._index_maintainer is not None and is_new: + self._index_maintainer.notify_new_record(partition, bucket, key_hash) + return bucket + + def prepare_commit(self): + if self._index_maintainer is None: + return {} + return self._index_maintainer.prepare_commit() + + def notify_precomputed_bucket_batch( + self, data: pa.RecordBatch, bucket: int + ) -> Optional[Tuple]: + """Maintain HASH indexes for rows assigned by an upstream coordinator.""" + partitions, _, key_hashes = self.extract_hashes_batch(data) + return self.notify_precomputed_bucket_hashes_batch( + data, bucket, key_hashes, partitions=partitions + ) + + def notify_precomputed_bucket_hashes_batch( + self, + data: pa.RecordBatch, + bucket: int, + key_hashes: List[int], + partitions: Optional[List[Tuple]] = None, + new_mappings: Optional[List[bool]] = None, + ) -> Optional[Tuple]: + """Maintain HASH indexes using hashes carried through the shuffle.""" + from pypaimon.index.dynamic_bucket import to_signed_int32 + + if self._index_maintainer is None: + raise RuntimeError( + "Precomputed dynamic buckets require a persistent table extractor" + ) + if len(key_hashes) != data.num_rows: + raise ValueError( + "Precomputed key hash count {} does not match row count {}".format( + len(key_hashes), data.num_rows + ) + ) + if partitions is None: + partitions = self._extract_partitions_batch(data) + if new_mappings is not None and len(new_mappings) != data.num_rows: + raise ValueError( + "Precomputed new-mapping count {} does not match row count {}".format( + len(new_mappings), data.num_rows + ) + ) + if not partitions: + return None + partition = tuple(partitions[0]) + if new_mappings is None: + new_mappings = [True] * data.num_rows + for actual_partition, key_hash, is_new in zip( + partitions, key_hashes, new_mappings + ): + if tuple(actual_partition) != tuple(partition): + raise RuntimeError( + "A precomputed dynamic-bucket group contained multiple " + f"partitions: expected {partition}, got {actual_partition}" + ) + if is_new: + self._index_maintainer.notify_new_record( + partition, bucket, to_signed_int32(key_hash) + ) + return partition + + def release_prepared(self) -> None: + if self._index_maintainer is not None: + self._index_maintainer.release_prepared() + + def abort(self) -> None: + if self._index_maintainer is not None: + self._index_maintainer.abort() class PostponeBucketRowKeyExtractor(RowKeyExtractor): diff --git a/paimon-python/pypaimon/write/table_commit.py b/paimon-python/pypaimon/write/table_commit.py index 4a758cb17662..2f4d8e60a0ce 100644 --- a/paimon-python/pypaimon/write/table_commit.py +++ b/paimon-python/pypaimon/write/table_commit.py @@ -23,6 +23,7 @@ logger = logging.getLogger(__name__) from pypaimon.write.commit_callback import CommitCallback from pypaimon.write.commit_message import CommitMessage +from pypaimon.write.commit.conflict_detection import CommitConflictError from pypaimon.write.file_store_commit import FileStoreCommit @@ -63,30 +64,44 @@ def add_commit_callback(self, callback: CommitCallback) -> None: def _commit(self, commit_messages: List[CommitMessage], commit_identifier: int = BATCH_COMMIT_IDENTIFIER): non_empty_messages = [msg for msg in commit_messages if not msg.is_empty()] - if self.overwrite_partition is not None: - # Always call overwrite() even with empty messages, so that - # FileStoreCommit.overwrite can handle the empty case properly - # (e.g. static overwrite with empty data should delete the partition). - logger.info( - "Committing overwrite to table %s, %d non-empty messages", - self.table.identifier, len(non_empty_messages) - ) - self.file_store_commit.overwrite( - overwrite_partition=self.overwrite_partition, - commit_messages=non_empty_messages, - commit_identifier=commit_identifier - ) - else: - if not non_empty_messages: - return - logger.info( - "Committing table %s, %d non-empty messages", - self.table.identifier, len(non_empty_messages) - ) - self.file_store_commit.commit( - commit_messages=non_empty_messages, - commit_identifier=commit_identifier - ) + try: + if self.overwrite_partition is not None: + # Always call overwrite() even with empty messages, so that + # FileStoreCommit.overwrite can handle the empty case properly + # (e.g. static overwrite with empty data should delete the partition). + logger.info( + "Committing overwrite to table %s, %d non-empty messages", + self.table.identifier, len(non_empty_messages) + ) + self.file_store_commit.overwrite( + overwrite_partition=self.overwrite_partition, + commit_messages=non_empty_messages, + commit_identifier=commit_identifier + ) + else: + if not non_empty_messages: + return + logger.info( + "Committing table %s, %d non-empty messages", + self.table.identifier, len(non_empty_messages) + ) + self.file_store_commit.commit( + commit_messages=non_empty_messages, + 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. + try: + self.file_store_commit.abort(non_empty_messages) + except Exception: + logger.warning( + "Failed to abort files after a deterministic commit " + "conflict.", + exc_info=True, + ) + raise def abort(self, commit_messages: List[CommitMessage]): self.file_store_commit.abort(commit_messages) diff --git a/paimon-python/pypaimon/write/table_write.py b/paimon-python/pypaimon/write/table_write.py index bfcb4c4f3a0d..c1276e72d2cc 100644 --- a/paimon-python/pypaimon/write/table_write.py +++ b/paimon-python/pypaimon/write/table_write.py @@ -38,7 +38,9 @@ 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() + 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 @@ -70,6 +72,92 @@ def write_arrow_batch(self, data: pa.RecordBatch): sub_table = pa.compute.take(data, indices_array) self.file_store_write.write(partition, bucket, sub_table) + def with_dynamic_bucket_index( + self, + ignore_existing: bool = False, + base_snapshot_id: Optional[int] = None, + ): + """Enable persistent HASH-index maintenance for coordinated writes.""" + from pypaimon.table.bucket_mode import BucketMode + from pypaimon.write.row_key_extractor import DynamicBucketRowKeyExtractor + + if self.table.bucket_mode() != BucketMode.HASH_DYNAMIC: + raise ValueError( + "Dynamic bucket index maintenance is only valid for " + "HASH_DYNAMIC tables" + ) + if self.file_store_write.data_writers: + raise RuntimeError( + "Dynamic bucket index maintenance must be enabled before writing" + ) + self.row_key_extractor = DynamicBucketRowKeyExtractor( + self.table.table_schema, + table=self.table, + ignore_existing=ignore_existing, + base_snapshot_id=base_snapshot_id, + ) + return self + + def write_arrow_batch_to_bucket( + self, + data: pa.RecordBatch, + bucket: int, + key_hashes: Optional[List[int]] = None, + new_mappings: Optional[List[bool]] = None, + ): + """Write one complete group whose bucket was computed upstream.""" + from pypaimon.table.bucket_mode import BucketMode + from pypaimon.write.row_key_extractor import DynamicBucketRowKeyExtractor + + bucket_mode = self.table.bucket_mode() + if bucket_mode not in (BucketMode.HASH_FIXED, BucketMode.HASH_DYNAMIC): + raise ValueError( + "Precomputed bucket writes are only valid for HASH_FIXED or " + "HASH_DYNAMIC tables" + ) + if not isinstance(self.row_key_extractor, DynamicBucketRowKeyExtractor): + if bucket_mode == BucketMode.HASH_DYNAMIC: + raise RuntimeError("Dynamic bucket extractor is not configured") + self._validate_pyarrow_schema(data.schema) + if bucket_mode == BucketMode.HASH_DYNAMIC: + if key_hashes is None: + partition = self.row_key_extractor.notify_precomputed_bucket_batch( + data, bucket + ) + else: + partition = ( + self.row_key_extractor + .notify_precomputed_bucket_hashes_batch( + data, + bucket, + key_hashes, + new_mappings=new_mappings, + ) + ) + else: + if key_hashes is not None: + raise ValueError( + "Precomputed key hashes are only valid for HASH_DYNAMIC tables" + ) + if new_mappings is not None: + raise ValueError( + "Precomputed new-mapping flags are only valid for " + "HASH_DYNAMIC tables" + ) + partitions = self.row_key_extractor.extract_partitions_batch(data) + if not partitions: + return + partition = tuple(partitions[0]) + for actual_partition in partitions[1:]: + if tuple(actual_partition) != partition: + raise RuntimeError( + "A precomputed fixed-bucket group contained multiple " + f"partitions: expected {partition}, got {actual_partition}" + ) + if partition is None: + return + self.file_store_write.write(partition, bucket, data) + def write_row(self, row): values_by_name = row_to_named_values(row, self.table.table_schema.fields) column_names = ( @@ -160,10 +248,57 @@ def write_ray( ) def close(self): - self.file_store_write.close() + try: + self.file_store_write.close() + finally: + self._release_prepared_indexes() def abort(self): - self.file_store_write.abort() + try: + self.file_store_write.abort() + finally: + abort = getattr(self.row_key_extractor, "abort", None) + if abort is not None: + abort() + + def _prepare_commit(self, commit_identifier) -> List[CommitMessage]: + commit_messages = self.file_store_write.prepare_commit(commit_identifier) + prepare_indexes = getattr(self.row_key_extractor, "prepare_commit", None) + if prepare_indexes is None: + return commit_messages + + index_changes = prepare_indexes() + base_snapshot_id = getattr( + self.row_key_extractor, "base_snapshot_id", None + ) + messages_by_bucket = { + (tuple(message.partition), message.bucket): message + for message in commit_messages + } + for (partition, bucket), changes in index_changes.items(): + message = messages_by_bucket.get((partition, bucket)) + if message is None: + message = CommitMessage( + partition=partition, + bucket=bucket, + new_files=[], + ) + commit_messages.append(message) + messages_by_bucket[(partition, bucket)] = message + message.index_adds.extend(changes.additions) + message.index_deletes.extend(changes.deletions) + if base_snapshot_id is not None: + # Data-only upserts must participate too. A concurrent overwrite + # can rebuild the HASH index and move an existing key, making a + # stale data file unsafe even when this writer added no mapping. + for message in commit_messages: + message.hash_index_base_snapshot = base_snapshot_id + return commit_messages + + def _release_prepared_indexes(self) -> None: + release = getattr(self.row_key_extractor, "release_prepared", None) + if release is not None: + release() def _validate_pyarrow_schema(self, data_schema: pa.Schema): if self._is_compatible_pyarrow_schema(data_schema, self.table_pyarrow_schema): @@ -218,10 +353,12 @@ def prepare_commit(self) -> List[CommitMessage]: if self.batch_committed: raise RuntimeError("BatchTableWrite only supports one-time committing.") self.batch_committed = True - return self.file_store_write.prepare_commit(BATCH_COMMIT_IDENTIFIER) + return self._prepare_commit(BATCH_COMMIT_IDENTIFIER) class StreamTableWrite(TableWrite): def prepare_commit(self, commit_identifier) -> List[CommitMessage]: - return self.file_store_write.prepare_commit(commit_identifier) + messages = self._prepare_commit(commit_identifier) + self._release_prepared_indexes() + return messages