This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 4595e4b4 Support distributed dynamic bucket input execution (#1208)
4595e4b4 is described below
commit 4595e4b410b0c3e7ebd34b3b068d895bc2d51b0f
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 23 19:33:47 2026 +0800
Support distributed dynamic bucket input execution (#1208)
---
bindings/python/README.md | 24 +
.../python/python/pypaimon_rust/datafusion.pyi | 31 +
bindings/python/src/context.rs | 2 +
bindings/python/src/write_routing.rs | 226 +++++++-
crates/paimon/src/table/bucket_assigner.rs | 4 +-
crates/paimon/src/table/bucket_router.rs | 64 +-
crates/paimon/src/table/dynamic_bucket_assigner.rs | 642 +++++++++++++++++++++
.../src/table/dynamic_bucket_assigner_router.rs | 464 +++++++++++++++
...assigner_dynamic.rs => dynamic_bucket_index.rs} | 339 ++++++++++-
crates/paimon/src/table/mod.rs | 8 +-
crates/paimon/src/table/table_write.rs | 4 +-
11 files changed, 1748 insertions(+), 60 deletions(-)
diff --git a/bindings/python/README.md b/bindings/python/README.md
index 981474f7..fc9cc1d2 100644
--- a/bindings/python/README.md
+++ b/bindings/python/README.md
@@ -367,6 +367,30 @@ Statistics contain partition columns followed by
`row_count: int64` and
may coincide with statistics names. Existing layouts remain unchanged. The
shared result can configure every worker with `with_bucket_plan(plan)`.
+For HASH_DYNAMIC primary-key inputs (`bucket=-1`), route rows to assigners
+before allocating buckets:
+
+```python
+from pypaimon_rust.datafusion import DynamicBucketAssigner,
DynamicBucketAssignerRouter
+
+routing = DynamicBucketAssignerRouter(table.schema(), num_channels=4,
num_assigners=4).route(batch)
+# Group rows and key_hash by routing's assigner column, then on that worker:
+assigner = DynamicBucketAssigner(table, num_channels=4, num_assigners=4,
assign_id=worker_id,
+ restore_snapshot_id=base_snapshot_id)
+buckets = assigner.assign(worker_batch, worker_key_hashes)
+```
+
+The assigner captures the latest write snapshot when none is supplied; zero
+selects an empty base. `overwrite=True` starts with empty assignments. Every
worker
+must use the same baseline, available as `assigner.restore_snapshot_id`.
+Configure the eventual file writers with that restore snapshot and the same
+overwrite mode.
+Assignment reads HASH indexes without creating files. Regroup its output by
+partition and bucket before writing, so one writer publishes the complete
+HASH index for each bucket. Key hashes can be omitted for Core to compute
+them; when supplied, the batch only needs partition columns. Existing keys
+retain their original buckets even when assigner parallelism changes.
+
Configure overwrite on the batch builder:
```python
diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi
b/bindings/python/python/pypaimon_rust/datafusion.pyi
index 9be54bfa..073f583d 100644
--- a/bindings/python/python/pypaimon_rust/datafusion.pyi
+++ b/bindings/python/python/pypaimon_rust/datafusion.pyi
@@ -891,6 +891,37 @@ class BucketRouter:
def with_bucket_plan(self, plan: pyarrow.RecordBatch) -> BucketRouter: ...
def route(self, batch: pyarrow.RecordBatch) -> pyarrow.Int32Array: ...
+class DynamicBucketAssignerRouter:
+ """Pure Java HASH_DYNAMIC routing before grouping input by assigner."""
+ def __init__(self, schema: TableSchema, num_channels: int, num_assigners:
int) -> None: ...
+ def route(self, batch: pyarrow.RecordBatch) -> pyarrow.RecordBatch:
+ """Return non-null assigner and key_hash Int32 columns in input
order."""
+ ...
+
+class DynamicBucketAssigner:
+ """Stateful bucket assignment for distributed HASH_DYNAMIC inputs.
+
+ Existing mappings remain stable across assigner parallelism changes.
+ New keys use buckets owned by this assigner. Input must already be
+ grouped by the DynamicBucketAssignerRouter's assigner column.
+ """
+ def __init__(self, table: Table, num_channels: int, num_assigners: int,
assign_id: int,
+ *, restore_snapshot_id: Optional[int] = None, overwrite: bool
= False) -> None:
+ """Capture one write-authorized restore point; None is latest and zero
is empty.
+
+ Overwrite starts empty. This stage creates no data or index files;
+ group its output by partition and bucket before writing.
+ """
+ ...
+ @property
+ def restore_snapshot_id(self) -> int:
+ """Captured restore snapshot; zero denotes an empty base."""
+ ...
+ def assign(self, batch: pyarrow.RecordBatch,
+ key_hashes: Optional[pyarrow.Int32Array] = None) ->
pyarrow.Int32Array:
+ """Return buckets in input order, optionally reusing routed key
hashes."""
+ ...
+
class PostponeBucketPlanner:
"""Pure per-partition statistics and bucket-count inference.
diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs
index c1179061..10a952fd 100644
--- a/bindings/python/src/context.rs
+++ b/bindings/python/src/context.rs
@@ -1040,6 +1040,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_,
PyModule>) -> PyResult<()>
this.add_class::<crate::write::PyBatchWriteBuilder>()?;
this.add_class::<crate::write::PyPostponeFixedBucketWriteBuilder>()?;
this.add_class::<crate::write_routing::PyBucketRouter>()?;
+ this.add_class::<crate::write_routing::PyDynamicBucketAssignerRouter>()?;
+ this.add_class::<crate::write_routing::PyDynamicBucketAssigner>()?;
this.add_class::<crate::write_routing::PyPostponeBucketPlanner>()?;
this.add_class::<crate::write::PyStreamWriteBuilder>()?;
this.add_class::<crate::write::PyBatchTableWrite>()?;
diff --git a/bindings/python/src/write_routing.rs
b/bindings/python/src/write_routing.rs
index 0eb20ceb..f053051c 100644
--- a/bindings/python/src/write_routing.rs
+++ b/bindings/python/src/write_routing.rs
@@ -17,15 +17,23 @@
//! Arrow transport for pure Core input routing and bucket planning.
-use crate::arrow_input::write_batch_from_pyarrow;
-use crate::error::to_py_err;
+use crate::arrow_input::{array_from_pyarrow, write_batch_from_pyarrow};
+use crate::error::{to_catalog_py_err, to_py_err};
use crate::schema::PyTableSchema;
-use arrow::array::Array;
+use crate::table::PyTable;
+use arrow::array::{Array, Int32Array};
+use arrow::datatypes::{DataType, Field, Schema};
use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use arrow::record_batch::RecordBatch;
use paimon::spec::TableSchema;
-use paimon::table::{BucketRouter, PostponeBucketPlan, PostponeBucketPlanner};
+use paimon::table::{
+ BucketRouter, DynamicBucketAssigner, DynamicBucketAssignerRouter,
PostponeBucketPlan,
+ PostponeBucketPlanner,
+};
+use paimon_datafusion::runtime::runtime;
+use pyo3::exceptions::PyTypeError;
use pyo3::prelude::*;
+use std::sync::Arc;
#[pyclass(name = "BucketRouter", module = "pypaimon_rust.datafusion")]
pub struct PyBucketRouter {
@@ -65,6 +73,107 @@ impl PyBucketRouter {
}
}
+#[pyclass(
+ name = "DynamicBucketAssignerRouter",
+ module = "pypaimon_rust.datafusion"
+)]
+pub struct PyDynamicBucketAssignerRouter {
+ inner: DynamicBucketAssignerRouter,
+}
+
+#[pymethods]
+impl PyDynamicBucketAssignerRouter {
+ #[new]
+ fn new(schema: &PyTableSchema, num_channels: i32, num_assigners: i32) ->
PyResult<Self> {
+ Ok(Self {
+ inner: DynamicBucketAssignerRouter::new(&schema.inner,
num_channels, num_assigners)
+ .map_err(to_py_err)?,
+ })
+ }
+
+ fn route(&self, py: Python<'_>, batch: &Bound<'_, PyAny>) ->
PyResult<Py<PyAny>> {
+ let batch = write_batch_from_pyarrow(batch)?;
+ let routing = py.detach(||
self.inner.route(&batch)).map_err(to_py_err)?;
+ RecordBatch::try_new(
+ Arc::new(Schema::new(vec![
+ Field::new("assigner", DataType::Int32, false),
+ Field::new("key_hash", DataType::Int32, false),
+ ])),
+ vec![Arc::new(routing.assigners), Arc::new(routing.key_hashes)],
+ )
+ .map_err(|error| PyTypeError::new_err(error.to_string()))?
+ .to_pyarrow(py)
+ .map(Bound::unbind)
+ }
+}
+
+#[pyclass(name = "DynamicBucketAssigner", module = "pypaimon_rust.datafusion")]
+pub struct PyDynamicBucketAssigner {
+ inner: DynamicBucketAssigner,
+}
+
+#[pymethods]
+impl PyDynamicBucketAssigner {
+ #[new]
+ #[pyo3(signature = (table, num_channels, num_assigners, assign_id, *,
restore_snapshot_id=None, overwrite=false))]
+ fn new(
+ py: Python<'_>,
+ table: &PyTable,
+ num_channels: i32,
+ num_assigners: i32,
+ assign_id: i32,
+ restore_snapshot_id: Option<i64>,
+ overwrite: bool,
+ ) -> PyResult<Self> {
+ let builder =
+ DynamicBucketAssigner::builder(&table.inner, num_channels,
num_assigners, assign_id);
+ let builder = match restore_snapshot_id {
+ Some(id) => builder
+ .with_restore_snapshot(id)
+ .map_err(to_catalog_py_err)?,
+ None => builder,
+ };
+ let builder = if overwrite {
+ builder.with_overwrite()
+ } else {
+ builder
+ };
+ let inner = py
+ .detach(|| runtime().block_on(builder.build()))
+ .map_err(to_catalog_py_err)?;
+ Ok(Self { inner })
+ }
+
+ #[getter]
+ fn restore_snapshot_id(&self) -> i64 {
+ self.inner.restore_snapshot_id()
+ }
+
+ #[pyo3(signature = (batch, key_hashes=None))]
+ fn assign(
+ &mut self,
+ py: Python<'_>,
+ batch: &Bound<'_, PyAny>,
+ key_hashes: Option<&Bound<'_, PyAny>>,
+ ) -> PyResult<Py<PyAny>> {
+ let batch = write_batch_from_pyarrow(batch)?;
+ let hashes = key_hashes.map(array_from_pyarrow).transpose()?;
+ let hashes = hashes
+ .as_ref()
+ .map(|hashes| {
+ hashes
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .ok_or_else(|| PyTypeError::new_err("key_hashes must be an
Int32Array"))
+ })
+ .transpose()?;
+ let buckets = py
+ .detach(|| runtime().block_on(self.inner.assign(&batch, hashes)))
+ .map_err(to_catalog_py_err)?;
+ buckets.to_data().to_pyarrow(py).map(Bound::unbind)
+ }
+}
+
#[pyclass(name = "PostponeBucketPlanner", module = "pypaimon_rust.datafusion")]
pub struct PyPostponeBucketPlanner {
schema: TableSchema,
@@ -125,6 +234,115 @@ mod tests {
use pyo3::ffi::c_str;
use pyo3::types::{PyDict, PyModule};
+ #[test]
+ fn dynamic_routing_and_assignment_preserve_arrow_rows_and_write_baseline()
{
+ use paimon::catalog::Identifier;
+ use paimon::io::FileIOBuilder;
+ use paimon::table::Table;
+
+ Python::attach(|py| {
+ let schema = Schema::builder()
+ .column("pt", DataType::VarChar(VarCharType::string_type()))
+ .column("id", DataType::Int(IntType::new()))
+ .column("payload",
DataType::VarChar(VarCharType::string_type()))
+ .partition_keys(["pt"])
+ .primary_key(["pt", "id"])
+ .option("bucket", "-1")
+ .option("dynamic-bucket.target-row-num", "1")
+ .build()
+ .unwrap();
+ let table = Arc::new(Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("default", "dynamic_arrow_binding"),
+ "memory:/dynamic_arrow_binding".into(),
+ TableSchema::new(0, &schema),
+ None,
+ ));
+ let locals = PyDict::new(py);
+ locals
+ .set_item(
+ "table",
+ Py::new(py, crate::table::PyTable::new(table)).unwrap(),
+ )
+ .unwrap();
+ locals
+ .set_item(
+ "DynamicBucketAssignerRouter",
+ py.get_type::<PyDynamicBucketAssignerRouter>(),
+ )
+ .unwrap();
+ locals
+ .set_item(
+ "DynamicBucketAssigner",
+ py.get_type::<PyDynamicBucketAssigner>(),
+ )
+ .unwrap();
+ py.run(
+ c_str!(
+ r#"
+import pyarrow as pa
+router = DynamicBucketAssignerRouter(table.schema(), 1, 1)
+batch = pa.record_batch([
+ pa.array(['unused', 'a', 'a', 'a', 'unused']),
+ pa.array([0, 10, 20, 10, 0], type=pa.int32()),
+ pa.array(['x', 'v1', 'v2', 'v3', 'x']),
+], names=['pt', 'id', 'payload']).slice(1, 3)
+routing = router.route(batch)
+assert routing.schema.names == ['assigner', 'key_hash']
+assert routing.column(0).to_pylist() == [0, 0, 0]
+hashes = routing.column(1)
+assert hashes.type == pa.int32() and hashes.null_count == 0
+assert hashes[0] == hashes[2] and hashes[0] != hashes[1]
+
+builder = table.new_batch_write_builder().with_restore_snapshot(0)
+assert not hasattr(builder, 'new_dynamic_bucket_assigner')
+assigner = DynamicBucketAssigner(table, 1, 1, 0, restore_snapshot_id=0)
+assert assigner.restore_snapshot_id == 0
+try:
+ DynamicBucketAssigner(table, 1, 1, 0, restore_snapshot_id=-1)
+except ValueError:
+ pass
+else:
+ raise AssertionError('negative restore snapshots must be rejected')
+for invalid in [pa.array([None, 1, 2], type=pa.int32()), hashes.slice(0, 2),
+ pa.array([1, 2, 3], type=pa.int64())]:
+ try:
+ assigner.assign(batch, invalid)
+ except (ValueError, TypeError):
+ pass
+ else:
+ raise AssertionError('invalid hash transport must fail before
assignment')
+buckets = assigner.assign(batch, hashes)
+assert buckets.to_pylist() == [0, 1, 0]
+assert assigner.assign(batch.slice(1), hashes.slice(1)).to_pylist() == [1, 0]
+assert DynamicBucketAssigner(table, 1, 1, 0,
restore_snapshot_id=0).assign(batch).equals(buckets)
+
+writer = builder.new_write()
+writer.write_arrow(batch.take(pa.array([0, 2])), bucket=0)
+writer.write_arrow(batch.slice(1, 1), bucket=1)
+messages = writer.prepare_commit()
+builder.new_commit().commit(messages)
+
+restored = DynamicBucketAssigner(table, 1, 1, 0)
+assert restored.restore_snapshot_id == 1
+assert restored.assign(batch, hashes).equals(buckets)
+# A captured empty base stays empty after publication of a newer snapshot.
+assert assigner.restore_snapshot_id == 0
+empty = DynamicBucketAssigner(table, 1, 1, 0, restore_snapshot_id=0)
+assert empty.restore_snapshot_id == 0
+assert empty.assign(batch.slice(1, 1)).to_pylist() == [0]
+overwritten = DynamicBucketAssigner(table, 1, 1, 0, overwrite=True)
+assert overwritten.restore_snapshot_id == 0
+assert overwritten.assign(batch.slice(1, 1)).to_pylist() == [0]
+"#
+ ),
+ Some(&locals),
+ Some(&locals),
+ )
+ .unwrap();
+ });
+ }
+
#[test]
fn routing_and_planning_use_typed_arrow_inputs_without_table_access() {
Python::attach(|py| {
diff --git a/crates/paimon/src/table/bucket_assigner.rs
b/crates/paimon/src/table/bucket_assigner.rs
index 64abf810..ca42f352 100644
--- a/crates/paimon/src/table/bucket_assigner.rs
+++ b/crates/paimon/src/table/bucket_assigner.rs
@@ -24,8 +24,8 @@ use crate::io::FileIO;
use crate::spec::{DataField, IndexFileMeta};
use crate::table::bucket_assigner_constant::ConstantBucketAssigner;
use crate::table::bucket_assigner_cross::CrossPartitionAssigner;
-use crate::table::bucket_assigner_dynamic::DynamicBucketAssigner;
use crate::table::bucket_assigner_fixed::FixedBucketAssigner;
+use crate::table::dynamic_bucket_index::DynamicBucketIndex;
use crate::Result;
use arrow_array::RecordBatch;
use std::collections::HashMap;
@@ -66,7 +66,7 @@ pub(crate) trait BucketAssigner: Send {
pub(crate) enum BucketAssignerEnum {
Constant(ConstantBucketAssigner),
Fixed(FixedBucketAssigner),
- Dynamic(Box<DynamicBucketAssigner>),
+ Dynamic(Box<DynamicBucketIndex>),
CrossPartition(Box<CrossPartitionAssigner>),
}
diff --git a/crates/paimon/src/table/bucket_router.rs
b/crates/paimon/src/table/bucket_router.rs
index af1f581c..38b8c0c1 100644
--- a/crates/paimon/src/table/bucket_router.rs
+++ b/crates/paimon/src/table/bucket_router.rs
@@ -28,7 +28,7 @@ use crate::spec::{
};
use crate::table::PostponeBucketPlan;
use crate::Result;
-use arrow_array::{Int32Array, RecordBatch};
+use arrow_array::{Int32Array, RecordBatch, RecordBatchOptions};
use arrow_schema::SchemaRef;
use std::collections::{HashMap, HashSet};
@@ -146,31 +146,7 @@ impl BucketRouter {
&self,
batch: &RecordBatch,
) -> Result<(Vec<Vec<u8>>, Int32Array)> {
- let input_schema = batch.schema();
- let mut names = HashSet::new();
- for field in input_schema.fields() {
- if !names.insert(field.name()) {
- return Err(data_invalid(format!(
- "Duplicate input column '{}' in bucket routing",
- field.name()
- )));
- }
- }
- let columns = self
- .fields
- .iter()
- .zip(self.arrow_schema.fields())
- .map(|(field, expected)| {
- let index = input_schema.index_of(field.name()).map_err(|_| {
- data_invalid(format!("Missing routing column '{}'",
field.name()))
- })?;
- normalize_write_array(batch.column(index),
expected.data_type())
- })
- .collect::<Result<Vec<_>>>()?;
- // Project by name before applying indices resolved from the table
- // schema. Unused payloads do not participate in routing validation.
- let input = RecordBatch::try_new(self.arrow_schema.clone(), columns)
- .map_err(|error| data_invalid(error.to_string()))?;
+ let input = project_routing_input(batch, &self.fields,
&self.arrow_schema)?;
let partitions = batch_to_serialized_bytes(&input,
&self.partition_indices, &self.fields)?;
if self.total_buckets > 0 {
let buckets = self.route_count(&input, self.total_buckets)?;
@@ -212,6 +188,42 @@ impl BucketRouter {
}
}
+/// Resolve required columns once, keeping unrelated payloads out of
validation.
+pub(super) fn project_routing_input(
+ batch: &RecordBatch,
+ fields: &[DataField],
+ arrow_schema: &SchemaRef,
+) -> Result<RecordBatch> {
+ let input_schema = batch.schema();
+ let mut names = HashSet::new();
+ for field in input_schema.fields() {
+ if !names.insert(field.name()) {
+ return Err(data_invalid(format!(
+ "Duplicate input column '{}' in bucket routing",
+ field.name()
+ )));
+ }
+ }
+ let columns = fields
+ .iter()
+ .zip(arrow_schema.fields())
+ .map(|(field, expected)| {
+ let index = input_schema
+ .index_of(field.name())
+ .map_err(|_| data_invalid(format!("Missing routing column
'{}'", field.name())))?;
+ normalize_write_array(batch.column(index), expected.data_type())
+ })
+ .collect::<Result<Vec<_>>>()?;
+ // Empty projections still represent the original number of rows, for
+ // example when upstream hashes route an unpartitioned input.
+ RecordBatch::try_new_with_options(
+ arrow_schema.clone(),
+ columns,
+ &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
+ )
+ .map_err(|error| data_invalid(error.to_string()))
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/crates/paimon/src/table/dynamic_bucket_assigner.rs
b/crates/paimon/src/table/dynamic_bucket_assigner.rs
new file mode 100644
index 00000000..c4979ba4
--- /dev/null
+++ b/crates/paimon/src/table/dynamic_bucket_assigner.rs
@@ -0,0 +1,642 @@
+// 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.
+
+//! Snapshot-bound dynamic bucket assignment before distributed file writing.
+
+use std::collections::{HashMap, HashSet};
+
+use arrow_array::{Int32Array, RecordBatch};
+
+use super::dynamic_bucket_assigner_router::DynamicBucketAssignerRouter;
+use super::dynamic_bucket_index::DynamicBucketIndex;
+use super::write_builder::ensure_table_write_authorized;
+use crate::spec::{CoreOptions, PartitionComputer};
+use crate::table::Table;
+use crate::Result;
+
+/// Assign dynamic buckets using the same key ownership and bucket allocation
+/// as Java's HashBucketAssigner. This stage produces no files or commit
+/// messages. The final table writer maintains complete HASH index files after
+/// inputs have been regrouped by their assigned partition and bucket.
+pub struct DynamicBucketAssigner {
+ router: DynamicBucketAssignerRouter,
+ index: DynamicBucketIndex,
+ num_assigners: i32,
+ assign_id: i32,
+ restore_snapshot_id: i64,
+}
+
+/// Configuration for a distributed dynamic-bucket assigner.
+///
+/// The restore snapshot is captured by [`Self::build`]. It is independent of
+/// writer configuration and must also be supplied to the eventual file
writers.
+pub struct DynamicBucketAssignerBuilder<'a> {
+ table: &'a Table,
+ num_channels: i32,
+ num_assigners: i32,
+ assign_id: i32,
+ restore_snapshot_id: Option<i64>,
+ overwrite: bool,
+}
+
+impl DynamicBucketAssignerBuilder<'_> {
+ /// Restore a specific snapshot. Zero denotes an empty restore state;
+ /// without this option, [`Self::build`] captures the latest snapshot.
+ pub fn with_restore_snapshot(mut self, snapshot_id: i64) -> Result<Self> {
+ if snapshot_id < 0 {
+ return Err(crate::Error::DataInvalid {
+ message: "Restore snapshot id must not be negative".into(),
+ source: None,
+ });
+ }
+ self.restore_snapshot_id = Some(snapshot_id);
+ Ok(self)
+ }
+
+ /// Assign against empty state, ignoring any configured restore snapshot.
+ pub fn with_overwrite(mut self) -> Self {
+ self.overwrite = true;
+ self
+ }
+
+ /// Authorize and capture the immutable restore point used by this
assigner.
+ pub async fn build(self) -> Result<DynamicBucketAssigner> {
+ DynamicBucketAssigner::new(
+ self.table,
+ self.num_channels,
+ self.num_assigners,
+ self.assign_id,
+ self.restore_snapshot_id,
+ self.overwrite,
+ )
+ .await
+ }
+}
+
+impl DynamicBucketAssigner {
+ /// Configure the dynamic-bucket assignment stage separately from file
writing.
+ pub fn builder(
+ table: &Table,
+ num_channels: i32,
+ num_assigners: i32,
+ assign_id: i32,
+ ) -> DynamicBucketAssignerBuilder<'_> {
+ DynamicBucketAssignerBuilder {
+ table,
+ num_channels,
+ num_assigners,
+ assign_id,
+ restore_snapshot_id: None,
+ overwrite: false,
+ }
+ }
+
+ async fn new(
+ table: &Table,
+ num_channels: i32,
+ num_assigners: i32,
+ assign_id: i32,
+ restore_snapshot_id: Option<i64>,
+ overwrite: bool,
+ ) -> Result<Self> {
+ ensure_table_write_authorized(table)?;
+ let router = DynamicBucketAssignerRouter::new(table.schema(),
num_channels, num_assigners)?;
+ if !(0..num_channels).contains(&assign_id) {
+ return Err(crate::Error::DataInvalid {
+ message: format!("assign_id must be in [0, {num_channels}),
got {assign_id}"),
+ source: None,
+ });
+ }
+ let schema = table.schema();
+ let fields = schema.fields();
+ let indices = |names: &[String]| {
+ names
+ .iter()
+ .map(|name| {
+ fields
+ .iter()
+ .position(|field| field.name() == name)
+ .unwrap()
+ })
+ .collect::<Vec<_>>()
+ };
+ let options = CoreOptions::new(schema.options());
+ let partition_computer = PartitionComputer::new(
+ schema.partition_keys(),
+ fields,
+ options.partition_default_name(),
+ options.legacy_partition_name(),
+ )?;
+ let mut index = DynamicBucketIndex::new(
+ table.clone(),
+ indices(schema.partition_keys()),
+ indices(&schema.trimmed_primary_keys()),
+ overwrite,
+ partition_computer,
+ )?;
+ let manager = table.snapshot_manager();
+ let snapshot = match (overwrite, restore_snapshot_id) {
+ (true, _) | (_, Some(0)) => None,
+ (_, Some(id)) => Some(manager.get_snapshot(id).await?),
+ (_, None) => manager.get_latest_snapshot().await?,
+ };
+ let restore_snapshot_id = snapshot.as_ref().map_or(0, |snapshot|
snapshot.id());
+ index.pin_restore_snapshot(snapshot);
+ Ok(Self {
+ router,
+ index,
+ num_assigners,
+ assign_id,
+ restore_snapshot_id,
+ })
+ }
+
+ /// Fixed restore point shared with the eventual file writers. Zero means
+ /// empty restore state, including overwrite, and must not resolve latest.
+ pub fn restore_snapshot_id(&self) -> i64 {
+ self.restore_snapshot_id
+ }
+
+ /// Return non-null bucket IDs in input order. Supplied hashes avoid
+ /// reading or normalizing primary-key columns again; only partition
+ /// columns are then required. Validate the entire input's owner and hash
+ /// arrays before restoring any index state or assigning any new key.
+ pub async fn assign(
+ &mut self,
+ batch: &RecordBatch,
+ key_hashes: Option<&Int32Array>,
+ ) -> Result<Int32Array> {
+ let (partitions, routing) = self.router.route_with_partitions(batch,
key_hashes)?;
+ for owner in routing.assigners.values() {
+ if *owner != self.assign_id {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Record assigner {owner} does not match writer
assigner {}",
+ self.assign_id
+ ),
+ source: None,
+ });
+ }
+ }
+ let mut requested: HashMap<Vec<u8>, HashSet<i32>> = HashMap::new();
+ for (partition, hash) in
partitions.iter().zip(routing.key_hashes.values()) {
+ requested
+ .entry(partition.clone())
+ .or_default()
+ .insert(*hash);
+ }
+ self.index
+ .restore_input_hashes(&requested, self.num_assigners,
self.assign_id)
+ .await?;
+ partitions
+ .iter()
+ .zip(routing.key_hashes.values())
+ .map(|(partition, hash)| {
+ self.index
+ .assign_input_hash(partition, *hash, self.num_assigners,
self.assign_id)
+ })
+ .collect::<Result<Vec<_>>>()
+ .map(Int32Array::from)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::catalog::Identifier;
+ use crate::io::FileIOBuilder;
+ use crate::spec::{DataType, IntType, Schema, TableSchema};
+ use arrow_array::{ArrayRef, RecordBatchOptions};
+ use arrow_schema::Schema as ArrowSchema;
+ use std::sync::Arc;
+
+ fn table(name: &str) -> Table {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("value", DataType::Int(IntType::new()))
+ .primary_key(["id"])
+ .option("bucket", "-1")
+ .option("dynamic-bucket.target-row-num", "1")
+ .option("dynamic-bucket.max-buckets", "4")
+ .build()
+ .unwrap();
+ Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("db", name),
+ format!("memory:/{name}"),
+ TableSchema::new(0, &schema),
+ None,
+ )
+ }
+
+ fn batch(ids: &[i32]) -> RecordBatch {
+ RecordBatch::try_from_iter([
+ ("id", Arc::new(Int32Array::from(ids.to_vec())) as ArrayRef),
+ (
+ "value",
+ Arc::new(Int32Array::from(vec![10; ids.len()])) as ArrayRef,
+ ),
+ ])
+ .unwrap()
+ }
+
+ fn keys_for_owner(table: &Table, owner: i32) -> (Vec<i32>, Int32Array) {
+ let input = batch(&(0..32).collect::<Vec<_>>());
+ let routing = DynamicBucketAssignerRouter::new(table.schema(), 2, 2)
+ .unwrap()
+ .route(&input)
+ .unwrap();
+ let positions = routing
+ .assigners
+ .values()
+ .iter()
+ .enumerate()
+ .filter_map(|(index, assigner)| (*assigner ==
owner).then_some(index))
+ .collect::<Vec<_>>();
+ assert!(positions.len() >= 3);
+ (
+ positions.iter().map(|index| *index as i32).collect(),
+ Int32Array::from_iter_values(
+ positions
+ .iter()
+ .map(|index| routing.key_hashes.value(*index)),
+ ),
+ )
+ }
+
+ async fn seed(table: &Table, ids: &[i32], bucket: i32) {
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ writer
+ .write_arrow(&[batch(ids)], Some(bucket))
+ .await
+ .unwrap();
+ let messages = writer.prepare_commit().await.unwrap();
+ builder.new_commit().commit(messages).await.unwrap();
+ writer.close().await;
+ }
+
+ #[tokio::test]
+ async fn
input_assignment_preserves_owned_holes_foreign_keys_and_later_batches() {
+ let table = table("input_owned_holes");
+ let (keys, hashes) = keys_for_owner(&table, 0);
+ // A changed parallelism can route this historical key to assigner 0
+ // even though it was originally written to foreign bucket 3.
+ seed(&table, &keys[..1], 3).await;
+ let mut assigner = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(assigner.restore_snapshot_id(), 1);
+ let routed = assigner
+ .assign(&batch(&keys[..2]), Some(&hashes.slice(0, 2)))
+ .await
+ .unwrap();
+ assert_eq!(routed.values(), &[3, 0]);
+ // Repeating the historical key does not increment a bucket's count;
+ // the next new key fills unused owned bucket 2 despite global id 3.
+ let routed = assigner
+ .assign(&batch(&[keys[0], keys[2], keys[1]]), None)
+ .await
+ .unwrap();
+ assert_eq!(routed.values(), &[3, 2, 0]);
+ assert_eq!(
+ table
+ .snapshot_manager()
+ .get_latest_snapshot()
+ .await
+ .unwrap()
+ .unwrap()
+ .id(),
+ 1
+ );
+ // Input assignment never writes HASH files or exposes commit messages.
+ let fresh = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(fresh.restore_snapshot_id(), 1);
+ }
+
+ #[tokio::test]
+ async fn
input_assignment_restores_new_requested_hashes_without_resetting_counts() {
+ let table = table("input_requested_hashes");
+ let (keys, _) = keys_for_owner(&table, 0);
+ seed(&table, &keys[..1], 3).await;
+ seed(&table, &keys[1..2], 2).await;
+ let mut assigner = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(
+ assigner
+ .assign(&batch(&keys[..1]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[3]
+ );
+ // This hash was not retained during the first call's restore.
+ assert_eq!(
+ assigner
+ .assign(&batch(&keys[1..2]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[2]
+ );
+ assert_eq!(
+ assigner
+ .assign(&batch(&keys[2..3]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[0]
+ );
+ }
+
+ #[tokio::test]
+ async fn input_assignment_pins_latest_and_empty_restore_points() {
+ let table = table("input_snapshot_pin");
+ let (keys, _) = keys_for_owner(&table, 0);
+ assert!(DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .with_restore_snapshot(-1)
+ .err()
+ .unwrap()
+ .to_string()
+ .contains("must not be negative"));
+ let mut empty = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(empty.restore_snapshot_id(), 0);
+ // Creating the builder does not capture an earlier restore point.
+ let pending = DynamicBucketAssigner::builder(&table, 2, 2, 0);
+ seed(&table, &keys[..1], 3).await;
+ let mut pinned = pending.build().await.unwrap();
+ assert_eq!(pinned.restore_snapshot_id(), 1);
+ seed(&table, &keys[1..2], 2).await;
+ assert_eq!(
+ empty
+ .assign(&batch(&keys[..1]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[0]
+ );
+ assert_eq!(
+ pinned
+ .assign(&batch(&keys[1..2]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[0]
+ );
+ let mut explicit = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .with_restore_snapshot(1)
+ .unwrap()
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(explicit.restore_snapshot_id(), 1);
+ assert_eq!(
+ explicit
+ .assign(&batch(&keys[1..2]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[0]
+ );
+ let mut latest = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(latest.restore_snapshot_id(), 2);
+ assert_eq!(
+ latest
+ .assign(&batch(&keys[1..2]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[2]
+ );
+ let mut explicit_empty = DynamicBucketAssigner::builder(&table, 2, 2,
0)
+ .with_restore_snapshot(0)
+ .unwrap()
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(explicit_empty.restore_snapshot_id(), 0);
+ assert_eq!(
+ explicit_empty
+ .assign(&batch(&keys[..1]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[0]
+ );
+ let mut overwrite = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .with_restore_snapshot(999)
+ .unwrap()
+ .with_overwrite()
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(overwrite.restore_snapshot_id(), 0);
+ assert_eq!(
+ overwrite
+ .assign(&batch(&keys[..1]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[0]
+ );
+ }
+
+ #[tokio::test]
+ async fn
input_assignment_checks_all_owners_and_hashes_before_state_changes() {
+ let table = table("input_validate_before_assignment");
+ let (keys, hashes) = keys_for_owner(&table, 0);
+ let (foreign, _) = keys_for_owner(&table, 1);
+ let mut assigner = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .build()
+ .await
+ .unwrap();
+ let error = assigner
+ .assign(&batch(&[keys[0], foreign[0]]), None)
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("does not match"));
+ let error = assigner
+ .assign(&batch(&keys[..2]), Some(&hashes.slice(0, 1)))
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("length") ||
error.to_string().contains("row"));
+ let error = assigner
+ .assign(
+ &batch(&keys[..2]),
+ Some(&Int32Array::from(vec![Some(hashes.value(0)), None])),
+ )
+ .await
+ .unwrap_err();
+ assert!(error.to_string().contains("null") ||
error.to_string().contains("NULL"));
+ // None of the rejected calls consumed the first bucket with keys[0].
+ assert_eq!(
+ assigner
+ .assign(&batch(&keys[1..2]), None)
+ .await
+ .unwrap()
+ .values(),
+ &[0]
+ );
+ }
+
+ #[tokio::test]
+ async fn prehashed_unpartitioned_input_keeps_zero_column_batch_rows() {
+ let table = table("input_prehashed_zero_columns");
+ let (_, hashes) = keys_for_owner(&table, 0);
+ let input = RecordBatch::try_new_with_options(
+ Arc::new(ArrowSchema::empty()),
+ vec![],
+ &RecordBatchOptions::default().with_row_count(Some(2)),
+ )
+ .unwrap();
+ let mut assigner = DynamicBucketAssigner::builder(&table, 2, 2, 0)
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(
+ assigner
+ .assign(&input, Some(&hashes.slice(0, 2)))
+ .await
+ .unwrap()
+ .values(),
+ &[0, 2]
+ );
+ assert!(assigner.assign(&input, None).await.is_err());
+ assert!(!table
+ .file_io()
+ .exists_dir("memory:/input_prehashed_zero_columns/index")
+ .await
+ .unwrap());
+ }
+
+ #[tokio::test]
+ async fn input_assignment_rejects_data_without_complete_hash_indexes() {
+ let table = table("input_missing_hash_index");
+ let builder = table.new_write_builder();
+ let mut writer = builder.new_write().unwrap();
+ writer.write_arrow(&[batch(&[1])], Some(0)).await.unwrap();
+ let mut messages = writer.prepare_commit().await.unwrap();
+ for message in &mut messages {
+ message.new_index_files.clear();
+ }
+ builder.new_commit().commit(messages).await.unwrap();
+ writer.close().await;
+ let mut assigner = DynamicBucketAssigner::builder(&table, 1, 1, 0)
+ .build()
+ .await
+ .unwrap();
+ assert!(assigner
+ .assign(&batch(&[1]), None)
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("no complete HASH"));
+ let mut overwrite = DynamicBucketAssigner::builder(&table, 1, 1, 0)
+ .with_overwrite()
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(
+ overwrite.assign(&batch(&[1]), None).await.unwrap().values(),
+ &[0]
+ );
+ }
+
+ #[tokio::test]
+ async fn
input_assignment_validates_duplicate_and_corrupt_restored_hash_indexes() {
+ for duplicate in [true, false] {
+ let table = table(if duplicate {
+ "input_duplicate_hash"
+ } else {
+ "input_corrupt_hash"
+ });
+ seed(&table, &[1], 0).await;
+ let manager = table.snapshot_manager();
+ let snapshot =
manager.get_latest_snapshot().await.unwrap().unwrap();
+ let path =
manager.manifest_path(snapshot.index_manifest().unwrap());
+ let mut entries =
crate::spec::IndexManifest::read(table.file_io(), &path)
+ .await
+ .unwrap();
+ assert_eq!(entries.len(), 1);
+ if duplicate {
+ entries.push(entries[0].clone());
+ } else {
+ entries[0].index_file.row_count += 1;
+ }
+ crate::spec::IndexManifest::write(table.file_io(), &path, &entries)
+ .await
+ .unwrap();
+ let mut assigner = DynamicBucketAssigner::builder(&table, 1, 1, 0)
+ .build()
+ .await
+ .unwrap();
+ let error = assigner.assign(&batch(&[1]), None).await.unwrap_err();
+ assert!(error.to_string().contains(if duplicate {
+ "Multiple HASH"
+ } else {
+ "Corrupt HASH"
+ }));
+ }
+ }
+
+ #[tokio::test]
+ async fn input_assignment_uses_write_identity_without_a_select_grant() {
+ let unbound =
table("input_unbound_auth").copy_with_options(HashMap::from([(
+ "query-auth.enabled".into(),
+ "true".into(),
+ )]));
+ assert!(DynamicBucketAssigner::builder(&unbound, 1, 1, 0)
+ .with_restore_snapshot(0)
+ .unwrap()
+ .build()
+ .await
+ .err()
+ .unwrap()
+ .to_string()
+ .contains("catalog-loaded REST"));
+ // This REST identity intentionally has no running query endpoint or
+ // SELECT grant. Empty restore state must not issue a query-auth call.
+ let rest = crate::table::rest_query_auth_table().await;
+ let bound = Table {
+ rest_env: rest.rest_env,
+ ..unbound
+ }
+ .with_query_auth_session();
+ let mut assigner = DynamicBucketAssigner::builder(&bound, 1, 1, 0)
+ .with_restore_snapshot(0)
+ .unwrap()
+ .build()
+ .await
+ .unwrap();
+ assert_eq!(
+ assigner.assign(&batch(&[1]), None).await.unwrap().values(),
+ &[0]
+ );
+ }
+}
diff --git a/crates/paimon/src/table/dynamic_bucket_assigner_router.rs
b/crates/paimon/src/table/dynamic_bucket_assigner_router.rs
new file mode 100644
index 00000000..218bbbc2
--- /dev/null
+++ b/crates/paimon/src/table/dynamic_bucket_assigner_router.rs
@@ -0,0 +1,464 @@
+// 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.
+
+//! Stateless input routing before distributed dynamic bucket assignment.
+
+use super::bucket_router::project_routing_input;
+use super::postpone_bucket_plan::data_invalid;
+use crate::spec::{
+ batch_build_binary_rows, batch_hash_codes, BinaryRow, CoreOptions,
DataField, TableSchema,
+};
+use crate::Result;
+use arrow_array::{Array, Int32Array, RecordBatch};
+use arrow_schema::SchemaRef;
+use std::collections::HashSet;
+
+/// Assigner IDs and Java BinaryRow key hashes in the original input order.
+/// The assigner subsequently uses its persisted HASH index to choose buckets.
+#[derive(Debug, Clone)]
+pub struct DynamicBucketAssignerRouting {
+ pub assigners: Int32Array,
+ pub key_hashes: Int32Array,
+}
+
+/// Hash only partition and primary-key columns without table access or I/O.
+///
+/// Java's `BucketAssigner.computeAssigner` combines the partition hash with
+/// the trimmed primary-key hash. Neither hash is a bucket ID: one assigner
+/// must own the persisted HASH-index state for each routed input.
+#[derive(Debug, Clone)]
+pub struct DynamicBucketAssignerRouter {
+ fields: Vec<DataField>,
+ arrow_schema: SchemaRef,
+ partition_indices: Vec<usize>,
+ key_indices: Vec<usize>,
+ partition_fields: Vec<DataField>,
+ partition_schema: SchemaRef,
+ num_channels: i32,
+ num_assigners: i32,
+}
+
+impl DynamicBucketAssignerRouter {
+ pub fn new(schema: &TableSchema, num_channels: i32, num_assigners: i32) ->
Result<Self> {
+ if num_channels <= 0 || num_assigners <= 0 {
+ return Err(data_invalid(
+ "num_channels and num_assigners must be positive",
+ ));
+ }
+ schema.validate_resolved_structure()?;
+ let options = CoreOptions::new(schema.options());
+ if options.is_format_table()
+ || options.bucket() != -1
+ || schema.primary_keys().is_empty()
+ || schema
+ .partition_keys()
+ .iter()
+ .any(|key| !schema.primary_keys().contains(key))
+ {
+ return Err(crate::Error::Unsupported {
+ message: "Dynamic bucket routing requires a Paimon primary-key
table with bucket=-1 and partition columns included in the primary key".into(),
+ });
+ }
+ if schema.options().contains_key("bucket-key") {
+ return Err(data_invalid(
+ "Cannot define 'bucket-key' in dynamic bucket mode",
+ ));
+ }
+
+ let keys = schema.trimmed_primary_keys();
+ let required = schema
+ .partition_keys()
+ .iter()
+ .chain(&keys)
+ .map(String::as_str)
+ .collect::<HashSet<_>>();
+ let fields = schema
+ .fields()
+ .iter()
+ .filter(|field| required.contains(field.name()))
+ .cloned()
+ .collect::<Vec<_>>();
+ let indices = |names: &[String]| {
+ names
+ .iter()
+ .map(|name| {
+ fields
+ .iter()
+ .position(|field| field.name() == name)
+ .ok_or_else(|| data_invalid(format!("Unknown routing
column '{name}'")))
+ })
+ .collect::<Result<Vec<_>>>()
+ };
+ let partition_indices = indices(schema.partition_keys())?;
+ let key_indices = indices(&keys)?;
+ let partition_fields = schema.partition_fields();
+ Ok(Self {
+ arrow_schema: crate::arrow::build_target_arrow_schema(&fields)?,
+ partition_schema:
crate::arrow::build_target_arrow_schema(&partition_fields)?,
+ fields,
+ partition_fields,
+ partition_indices,
+ key_indices,
+ num_channels,
+ num_assigners,
+ })
+ }
+
+ pub fn route(&self, batch: &RecordBatch) ->
Result<DynamicBucketAssignerRouting> {
+ Ok(self.route_input(batch, None)?.1)
+ }
+
+ /// Prehashed input needs only partition columns. Its payload and key
+ /// arrays are deliberately not read or normalized a second time.
+ pub(super) fn route_with_partitions(
+ &self,
+ batch: &RecordBatch,
+ key_hashes: Option<&Int32Array>,
+ ) -> Result<(Vec<Vec<u8>>, DynamicBucketAssignerRouting)> {
+ let (partitions, routing) = self.route_input(batch, key_hashes)?;
+ Ok((
+ partitions
+ .into_iter()
+ .map(|row| row.to_serialized_bytes())
+ .collect(),
+ routing,
+ ))
+ }
+
+ fn route_input(
+ &self,
+ batch: &RecordBatch,
+ key_hashes: Option<&Int32Array>,
+ ) -> Result<(Vec<BinaryRow>, DynamicBucketAssignerRouting)> {
+ if let Some(hashes) = key_hashes {
+ if hashes.len() != batch.num_rows() || hashes.null_count() != 0 {
+ return Err(data_invalid(
+ "Dynamic key hashes must be non-null and match the input
row count",
+ ));
+ }
+ }
+ let (input, fields, partition_indices, hashes) = if let Some(hashes) =
key_hashes {
+ (
+ project_routing_input(batch, &self.partition_fields,
&self.partition_schema)?,
+ &self.partition_fields,
+ (0..self.partition_fields.len()).collect::<Vec<_>>(),
+ hashes.clone(),
+ )
+ } else {
+ let input = project_routing_input(batch, &self.fields,
&self.arrow_schema)?;
+ let hashes =
+ Int32Array::from(batch_hash_codes(&input, &self.key_indices,
&self.fields)?);
+ (input, &self.fields, self.partition_indices.clone(), hashes)
+ };
+ let partitions = batch_build_binary_rows(&input, &partition_indices,
fields)?;
+ let assigners =
Int32Array::from_iter_values(partitions.iter().zip(hashes.values()).map(
+ |(partition, hash)| {
+ compute_assigner(
+ partition.hash_code(),
+ *hash,
+ self.num_channels,
+ self.num_assigners,
+ )
+ },
+ ));
+ Ok((
+ partitions,
+ DynamicBucketAssignerRouting {
+ assigners,
+ key_hashes: hashes,
+ },
+ ))
+ }
+}
+
+/// Java signed remainder and `int` addition, including overflow before the
+/// final remainder. The router constructor validates both divisors.
+pub(super) fn compute_assigner(
+ partition_hash: i32,
+ key_hash: i32,
+ num_channels: i32,
+ num_assigners: i32,
+) -> i32 {
+ let start = (partition_hash % num_channels).wrapping_abs();
+ let assigner = (key_hash % num_assigners).wrapping_abs();
+ start.wrapping_add(assigner) % num_channels
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::spec::{
+ DataType, DecimalType, IntType, LocalZonedTimestampType, Schema,
TimestampType, VarCharType,
+ };
+ use arrow_array::{
+ ArrayRef, Decimal128Array, LargeStringArray, StringArray,
TimestampMicrosecondArray,
+ };
+ use arrow_schema::{DataType as ArrowType, Field, Schema as ArrowSchema};
+ use std::collections::HashMap;
+ use std::sync::Arc;
+
+ fn schema() -> TableSchema {
+ TableSchema::new(
+ 0,
+ &Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("p", DataType::VarChar(VarCharType::string_type()))
+ .column("label", DataType::VarChar(VarCharType::string_type()))
+ .column("payload",
DataType::VarChar(VarCharType::string_type()))
+ .primary_key(["p", "id", "label"])
+ .partition_keys(["p"])
+ .option("bucket", "-1")
+ .option("primary-key.nullable", "true")
+ .build()
+ .unwrap(),
+ )
+ }
+
+ fn input() -> RecordBatch {
+ RecordBatch::try_from_iter([
+ (
+ "label",
+ Arc::new(LargeStringArray::from(vec![
+ Some("hello"),
+ Some(""),
+ Some("abcdefghijk"),
+ Some("你好"),
+ None,
+ ])) as ArrayRef,
+ ),
+ (
+ "ignored",
+ Arc::new(Int32Array::from(vec![0; 5])) as ArrayRef,
+ ),
+ (
+ "p",
+ Arc::new(StringArray::from(vec![
+ Some("east"),
+ Some("west"),
+ None,
+ Some("abcdefghijk"),
+ Some("east"),
+ ])) as ArrayRef,
+ ),
+ (
+ "id",
+ Arc::new(Int32Array::from(vec![
+ Some(5),
+ Some(-3),
+ None,
+ Some(i32::MIN),
+ Some(i32::MAX),
+ ])) as ArrayRef,
+ ),
+ ])
+ .unwrap()
+ }
+
+ #[test]
+ fn routes_reordered_key_only_rows_using_current_java_hashes() {
+ // Current BinaryRowWriter + BucketAssigner.computeAssigner(7, 3).
+ // The partition columns do not participate in the trimmed key hash.
+ let router = DynamicBucketAssignerRouter::new(&schema(), 7,
3).unwrap();
+ let batch = input();
+ let routed = router.route(&batch).unwrap();
+ assert_eq!(
+ routed.key_hashes.values(),
+ &[
+ -1_451_184_867,
+ 556_362_117,
+ 1_425_286_886,
+ 308_673_824,
+ -1_856_341_548
+ ]
+ );
+ assert_eq!(routed.assigners.values(), &[6, 6, 5, 2, 6]);
+ assert_eq!(routed.assigners.null_count(), 0);
+ assert_eq!(routed.key_hashes.null_count(), 0);
+ // SDK Tables arrive as batches: boundaries must not change ownership.
+ let pieces = [
+ router.route(&batch.slice(0, 2)).unwrap(),
+ router.route(&batch.slice(2, 3)).unwrap(),
+ ];
+ assert_eq!(
+ pieces
+ .iter()
+ .flat_map(|part| part.assigners.values().iter().copied())
+ .collect::<Vec<_>>(),
+ routed.assigners.values().to_vec()
+ );
+ let empty = router.route(&batch.slice(0, 0)).unwrap();
+ assert_eq!(empty.assigners.len(), 0);
+ assert_eq!(empty.key_hashes.len(), 0);
+ }
+
+ #[test]
+ fn prehashed_routing_reads_only_partition_columns() {
+ let router = DynamicBucketAssignerRouter::new(&schema(), 7,
3).unwrap();
+ let batch = input();
+ let expected = router.route(&batch).unwrap();
+ let partition_only = batch.project(&[2]).unwrap();
+ let (partitions, actual) = router
+ .route_with_partitions(&partition_only, Some(&expected.key_hashes))
+ .unwrap();
+ assert_eq!(actual.assigners, expected.assigners);
+ assert_eq!(actual.key_hashes, expected.key_hashes);
+ assert_eq!(partitions.len(), batch.num_rows());
+ assert!(router.route(&partition_only).is_err());
+ for hashes in [Int32Array::from(vec![0]), Int32Array::from(vec![None;
5])] {
+ assert!(router
+ .route_with_partitions(&partition_only, Some(&hashes))
+ .is_err());
+ }
+ }
+
+ #[test]
+ fn unpartitioned_prehashed_rows_retain_their_count() {
+ let schema = TableSchema::new(
+ 0,
+ &Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .primary_key(["id"])
+ .option("bucket", "-1")
+ .build()
+ .unwrap(),
+ );
+ let router = DynamicBucketAssignerRouter::new(&schema, 7, 3).unwrap();
+ let batch = RecordBatch::try_from_iter([(
+ "payload",
+ Arc::new(StringArray::from(vec!["ignored"; 4])) as ArrayRef,
+ )])
+ .unwrap();
+ let hashes = Int32Array::from(vec![i32::MIN, -1, 0, i32::MAX]);
+ let (partitions, routed) = router.route_with_partitions(&batch,
Some(&hashes)).unwrap();
+ // The empty Java BinaryRow hash is -1670924195, whose abs(% 7) is 3.
+ assert_eq!(partitions.len(), 4);
+ assert_eq!(routed.assigners.values(), &[5, 4, 3, 4]);
+ }
+
+ #[test]
+ fn nullable_noncompact_keys_keep_java_variable_slots() {
+ for (data_type, values, expected_hash) in [
+ (
+ DataType::Decimal(DecimalType::new(20, 2).unwrap()),
+ Arc::new(
+ Decimal128Array::from(vec![None::<i128>])
+ .with_precision_and_scale(20, 2)
+ .unwrap(),
+ ) as ArrayRef,
+ -1_380_280_760,
+ ),
+ (
+ DataType::Timestamp(TimestampType::new(6).unwrap()),
+ Arc::new(TimestampMicrosecondArray::from(vec![None::<i64>]))
as ArrayRef,
+ -340_882_661,
+ ),
+ (
+
DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(6).unwrap()),
+
Arc::new(TimestampMicrosecondArray::from(vec![None::<i64>]).with_timezone("UTC"))
+ as ArrayRef,
+ -340_882_661,
+ ),
+ ] {
+ let schema = TableSchema::new(
+ 0,
+ &Schema::builder()
+ .column("id", data_type)
+ .primary_key(["id"])
+ .option("bucket", "-1")
+ .option("primary-key.nullable", "true")
+ .build()
+ .unwrap(),
+ );
+ let router = DynamicBucketAssignerRouter::new(&schema, 7,
3).unwrap();
+ let batch = RecordBatch::try_from_iter([("id", values)]).unwrap();
+ assert_eq!(
+ router.route(&batch).unwrap().key_hashes.values(),
+ &[expected_hash]
+ );
+ }
+ }
+
+ #[test]
+ fn invalid_schemas_parallelism_and_required_inputs_fail_before_routing() {
+ for (channels, assigners) in [(0, 1), (1, 0), (-1, 1), (1, -1)] {
+ assert!(DynamicBucketAssignerRouter::new(&schema(), channels,
assigners).is_err());
+ }
+ for options in [
+ HashMap::from([("bucket".into(), "4".into())]),
+ HashMap::from([("bucket-key".into(), "id".into())]),
+ HashMap::from([("bucket-key".into(), "".into())]),
+ ] {
+ assert!(
+
DynamicBucketAssignerRouter::new(&schema().copy_with_options(options), 7, 3)
+ .is_err()
+ );
+ }
+ let no_primary_key = TableSchema::new(
+ 0,
+ &Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .build()
+ .unwrap(),
+ );
+ assert!(DynamicBucketAssignerRouter::new(&no_primary_key, 7,
3).is_err());
+ let cross_partition = TableSchema::new(
+ 0,
+ &Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("p", DataType::VarChar(VarCharType::string_type()))
+ .primary_key(["id"])
+ .partition_keys(["p"])
+ .build()
+ .unwrap(),
+ );
+ assert!(DynamicBucketAssignerRouter::new(&cross_partition, 7,
3).is_err());
+ let router = DynamicBucketAssignerRouter::new(&schema(), 7,
3).unwrap();
+ let batch = input();
+ assert!(router.route(&batch.project(&[0, 2]).unwrap()).is_err());
+ let duplicate = RecordBatch::try_new(
+ Arc::new(ArrowSchema::new(vec![
+ Field::new("id", ArrowType::Int32, true),
+ Field::new("id", ArrowType::Int32, true),
+ ])),
+ vec![batch.column(3).clone(), batch.column(3).clone()],
+ )
+ .unwrap();
+ assert!(router.route(&duplicate).is_err());
+ let incompatible = RecordBatch::try_from_iter([
+ ("p", batch.column(2).clone()),
+ ("label", batch.column(0).clone()),
+ ("id", Arc::new(StringArray::from(vec!["1"; 5])) as ArrayRef),
+ ])
+ .unwrap();
+ assert!(router.route(&incompatible).is_err());
+ }
+
+ #[test]
+ fn java_assigner_remainder_keeps_negative_hashes_and_int_min() {
+ for (partition, expected) in [
+ (i32::MIN, [4, 3, 2, 3]),
+ (-1, [3, 2, 1, 2]),
+ (0, [2, 1, 0, 1]),
+ (i32::MAX, [3, 2, 1, 2]),
+ ] {
+ for (key, expected) in [i32::MIN, -1, 0,
i32::MAX].into_iter().zip(expected) {
+ assert_eq!(compute_assigner(partition, key, 7, 3), expected);
+ }
+ }
+ assert_eq!(compute_assigner(-1, i32::MIN, 1, 1), 0);
+ }
+}
diff --git a/crates/paimon/src/table/bucket_assigner_dynamic.rs
b/crates/paimon/src/table/dynamic_bucket_index.rs
similarity index 79%
rename from crates/paimon/src/table/bucket_assigner_dynamic.rs
rename to crates/paimon/src/table/dynamic_bucket_index.rs
index d6422b60..316670ff 100644
--- a/crates/paimon/src/table/bucket_assigner_dynamic.rs
+++ b/crates/paimon/src/table/dynamic_bucket_index.rs
@@ -15,12 +15,12 @@
// specific language governing permissions and limitations
// under the License.
-//! Dynamic bucket assigner for PK tables with bucket=-1 where PK includes
partition fields.
+//! Dynamic bucket index state for PK tables whose primary key includes
partition fields.
//!
//! Also contains the per-bucket index maintainer
(`DynamicBucketIndexMaintainer`)
//! and per-partition index (`PartitionIndex`) used by both dynamic and
cross-partition modes.
-use crate::io::FileIO;
+use crate::io::{FileIO, FileRead};
use crate::spec::MAX_DYNAMIC_BUCKETS;
use crate::spec::{
batch_hash_codes, batch_to_serialized_bytes, bucket_path_under, BinaryRow,
CoreOptions,
@@ -44,6 +44,7 @@ use uuid::Uuid;
/// Index type identifier for hash index files, matching Java's
`HashIndexFile.HASH_INDEX`.
const HASH_INDEX: &str = "HASH";
+const HASH_INDEX_READ_CHUNK_SIZE: u64 = 8 * 1024 * 1024;
/// Read/write hash index files.
///
@@ -80,6 +81,49 @@ impl HashIndexFile {
Ok(hashes)
}
+ /// Distributed input assigners retain only hashes requested by their
+ /// current batch. Scan bounded ranges rather than materializing a whole
+ /// partition's index or constructing writer-side index maintainers.
+ async fn read_requested(
+ file_io: &FileIO,
+ path: &str,
+ expected_count: i64,
+ requested: &HashSet<i32>,
+ ) -> Result<HashSet<i32>> {
+ let input = file_io.new_input(path)?;
+ let size = input.metadata().await?.size;
+ let count = u64::try_from(expected_count).ok();
+ if size % 4 != 0 || count.and_then(|count| count.checked_mul(4)) !=
Some(size) {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Corrupt HASH index {path}: expected {expected_count}
hashes, got {size} bytes"
+ ),
+ source: None,
+ });
+ }
+ let reader = input.reader().await?;
+ let mut found = HashSet::new();
+ let mut offset = 0;
+ while offset < size {
+ let end = (offset + HASH_INDEX_READ_CHUNK_SIZE).min(size);
+ let content = reader.read(offset..end).await?;
+ if content.len() as u64 != end - offset {
+ return Err(crate::Error::DataInvalid {
+ message: format!("Short read in HASH index {path}"),
+ source: None,
+ });
+ }
+ for bytes in content.as_chunks::<4>().0 {
+ let hash = i32::from_be_bytes(*bytes);
+ if requested.contains(&hash) {
+ found.insert(hash);
+ }
+ }
+ offset = end;
+ }
+ Ok(found)
+ }
+
/// Write key hashes to a new hash index file, returning its metadata.
async fn write_at(
file_io: &FileIO,
@@ -430,15 +474,36 @@ impl PartitionIndex {
/// 2. Otherwise, find a non-full bucket and assign the hash there.
/// 3. If all buckets are full, create a new bucket.
fn assign(&mut self, hash: i32, max_buckets: i32) -> Result<i32> {
+ let (bucket, new) = self.allocate(hash, max_buckets, |_| true)?;
+ if new {
+ self.bucket_maintainers
+ .entry(bucket)
+ .or_insert_with(|| DynamicBucketIndexMaintainer::new(vec![]))
+ .notify_new_record(hash);
+ }
+ Ok(bucket)
+ }
+
+ /// Java PartitionIndex allocation, shared by writers and distributed
+ /// input assignment. Only writers notify HASH-file maintainers.
+ fn allocate(
+ &mut self,
+ hash: i32,
+ max_buckets: i32,
+ bucket_filter: impl Fn(i32) -> bool,
+ ) -> Result<(i32, bool)> {
// 1. Already assigned
if let Some(&bucket) = self.hash_to_bucket.get(&hash) {
- return Ok(bucket);
+ return Ok((bucket, false));
}
// 2. Find a non-full bucket
let mut full_buckets = Vec::new();
let mut assigned_bucket = None;
for (&bucket, count) in &mut self.non_full_buckets {
+ if !bucket_filter(bucket) {
+ continue;
+ }
if *count < self.target_bucket_row_number {
*count += 1;
self.hash_to_bucket.insert(hash, bucket);
@@ -452,11 +517,7 @@ impl PartitionIndex {
self.non_full_buckets.remove(&b);
}
if let Some(bucket) = assigned_bucket {
- self.bucket_maintainers
- .entry(bucket)
- .or_insert_with(|| DynamicBucketIndexMaintainer::new(vec![]))
- .notify_new_record(hash);
- return Ok(bucket);
+ return Ok((bucket, true));
}
// Java PartitionIndex first allocates an unused id in range. A
configured
@@ -466,7 +527,10 @@ impl PartitionIndex {
} else {
max_buckets
};
- while self.next_bucket_id < limit &&
self.all_buckets.contains(&self.next_bucket_id) {
+ while self.next_bucket_id < limit
+ && (self.all_buckets.contains(&self.next_bucket_id)
+ || !bucket_filter(self.next_bucket_id))
+ {
self.next_bucket_id += 1;
}
let bucket = if self.next_bucket_id < limit {
@@ -485,16 +549,12 @@ impl PartitionIndex {
.bucket_ids
.choose(&mut rand::thread_rng())
.ok_or_else(|| crate::Error::DataInvalid {
- message: "No dynamic bucket is available".to_string(),
+ message: "No dynamic bucket is available for this
assigner".to_string(),
source: None,
})?
};
self.hash_to_bucket.insert(hash, bucket);
- self.bucket_maintainers
- .entry(bucket)
- .or_insert_with(|| DynamicBucketIndexMaintainer::new(vec![]))
- .notify_new_record(hash);
- Ok(bucket)
+ Ok((bucket, true))
}
/// Write hash index files for all modified buckets, returning (bucket,
index_files) pairs.
@@ -529,13 +589,15 @@ impl PartitionIndex {
}
// ---------------------------------------------------------------------------
-// DynamicBucketAssigner
+// DynamicBucketIndex
// ---------------------------------------------------------------------------
-/// Bucket assigner for dynamic bucket mode (bucket=-1) where PK includes
partition fields.
+/// HASH index state for dynamic bucket mode (bucket=-1).
///
-/// Maintains a per-partition `PartitionIndex` that maps key hashes to bucket
ids.
-pub(crate) struct DynamicBucketAssigner {
+/// Restores and maintains per-partition key-to-bucket mappings. Distributed
+/// assigners restore requested mappings; file writers additionally maintain
+/// complete HASH indexes for publication.
+pub(crate) struct DynamicBucketIndex {
partition_field_indices: Vec<usize>,
primary_key_indices: Vec<usize>,
/// Schema fields for BinaryRow extraction.
@@ -561,7 +623,7 @@ pub(crate) struct DynamicBucketAssigner {
index_file_in_data_file_dir: bool,
}
-impl DynamicBucketAssigner {
+impl DynamicBucketIndex {
pub fn new(
table: Table,
partition_field_indices: Vec<usize>,
@@ -602,6 +664,156 @@ impl DynamicBucketAssigner {
self.restore_snapshot_id = Some(snapshot_id);
}
+ /// Input assignment and the eventual file writer must restore the exact
+ /// same snapshot, including an explicitly captured empty table.
+ pub(super) fn pin_restore_snapshot(&mut self, snapshot: Option<Snapshot>) {
+ self.restore_snapshot_id = Some(snapshot.as_ref().map_or(0,
Snapshot::id));
+ self.snapshot = snapshot;
+ }
+
+ /// Restore only mappings requested by this input, without constructing
+ /// writer-side maintainers. Historical mappings may belong to a bucket
+ /// owned by another current assigner after a parallelism change.
+ pub(super) async fn restore_input_hashes(
+ &mut self,
+ requested: &HashMap<Vec<u8>, HashSet<i32>>,
+ num_assigners: i32,
+ assign_id: i32,
+ ) -> Result<()> {
+ let mut missing = Vec::new();
+ for (partition, hashes) in requested {
+ let index = self.partition_indexes.get(partition);
+ let hashes = hashes
+ .iter()
+ .copied()
+ .filter(|hash| index.is_none_or(|index|
!index.hash_to_bucket.contains_key(hash)))
+ .collect::<HashSet<_>>();
+ if !hashes.is_empty() {
+ missing.push((partition, hashes));
+ }
+ }
+ if missing.is_empty() {
+ return Ok(());
+ }
+ self.ensure_index_entries_loaded().await?;
+ // Apply restored maps only after every requested partition was read
+ // successfully. Input owner/shape checks run before this method.
+ let mut restored = Vec::with_capacity(missing.len());
+ for (partition, hashes) in missing {
+ let entries = self
+ .cached_index_entries
+ .as_deref()
+ .unwrap_or(&[])
+ .iter()
+ .filter(|entry| {
+ entry.partition == *partition &&
entry.index_file.index_type == HASH_INDEX
+ })
+ .collect::<Vec<_>>();
+ let mut buckets = HashSet::new();
+ for entry in &entries {
+ if !(0..MAX_DYNAMIC_BUCKETS).contains(&entry.bucket) {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Dynamic bucket id must be between 0 and {}, but
was {}",
+ MAX_DYNAMIC_BUCKETS - 1,
+ entry.bucket
+ ),
+ source: None,
+ });
+ }
+ if !buckets.insert(entry.bucket) {
+ return Err(crate::Error::DataInvalid {
+ message: format!(
+ "Multiple HASH indexes for dynamic bucket {}",
+ entry.bucket
+ ),
+ source: None,
+ });
+ }
+ }
+ let new_partition =
!self.partition_indexes.contains_key(partition);
+ if new_partition {
+ self.validate_data_buckets(
+ partition,
+ &entries
+ .iter()
+ .map(|entry| (*entry).clone())
+ .collect::<Vec<_>>(),
+ )
+ .await?;
+ }
+ let mut index =
PartitionIndex::empty(self.target_bucket_row_number);
+ let partition_path = self.partition_path(partition)?;
+ let options = CoreOptions::new(self.table.schema().options());
+ let layout = HashIndexLayout {
+ table_path: self.table.location().trim_end_matches('/'),
+ partition_path: &partition_path,
+ data_file_path_directory: options.data_file_path_directory(),
+ index_file_in_data_file_dir: self.index_file_in_data_file_dir,
+ };
+ for entry in entries {
+ let path = layout.resolve(
+ entry.bucket,
+ &entry.index_file.file_name,
+ entry.index_file.external_path.as_deref(),
+ );
+ for hash in HashIndexFile::read_requested(
+ self.table.file_io(),
+ &path,
+ entry.index_file.row_count,
+ &hashes,
+ )
+ .await?
+ {
+ if let Some(previous) = index.hash_to_bucket.insert(hash,
entry.bucket) {
+ if previous != entry.bucket {
+ return Err(crate::Error::DataInvalid {
+ message: format!("HASH index assigns hash
{hash} to both buckets {previous} and {}", entry.bucket),
+ source: None,
+ });
+ }
+ }
+ }
+ // Restored foreign buckets only supply historical mappings.
+ // New keys can use or grow buckets owned by this assigner.
+ if new_partition && entry.bucket % num_assigners == assign_id
% num_assigners {
+ index.all_buckets.insert(entry.bucket);
+ index.bucket_ids.push(entry.bucket);
+ if entry.index_file.row_count <
self.target_bucket_row_number {
+ index
+ .non_full_buckets
+ .insert(entry.bucket, entry.index_file.row_count);
+ }
+ }
+ }
+ restored.push((partition.clone(), index));
+ }
+ for (partition, index) in restored {
+ if let Some(existing) = self.partition_indexes.get_mut(&partition)
{
+ existing.hash_to_bucket.extend(index.hash_to_bucket);
+ } else {
+ self.partition_indexes.insert(partition, index);
+ }
+ }
+ Ok(())
+ }
+
+ pub(super) fn assign_input_hash(
+ &mut self,
+ partition: &[u8],
+ hash: i32,
+ num_assigners: i32,
+ assign_id: i32,
+ ) -> Result<i32> {
+ self.partition_indexes
+ .get_mut(partition)
+ .expect("requested input hashes were restored before assignment")
+ .allocate(hash, self.max_buckets, |bucket| {
+ bucket % num_assigners == assign_id % num_assigners
+ })
+ .map(|(bucket, _)| bucket)
+ }
+
/// Java's DynamicBucketIndexMaintainer sees keys after bucket assignment.
/// Keep that same full-file replacement lifecycle for precomputed groups.
pub async fn notify_precomputed_batch(
@@ -661,10 +873,14 @@ impl DynamicBucketAssigner {
return Ok(());
}
let snapshot_manager = self.table.snapshot_manager();
- let latest_snapshot = match self.restore_snapshot_id {
- Some(0) => None,
- Some(id) => Some(snapshot_manager.get_snapshot(id).await?),
- None => snapshot_manager.get_latest_snapshot().await?,
+ let latest_snapshot = if self.snapshot.is_some() {
+ self.snapshot.clone()
+ } else {
+ match self.restore_snapshot_id {
+ Some(0) => None,
+ Some(id) => Some(snapshot_manager.get_snapshot(id).await?),
+ None => snapshot_manager.get_latest_snapshot().await?,
+ }
};
let entries = if let Some(snapshot) = &latest_snapshot {
@@ -769,7 +985,7 @@ impl DynamicBucketAssigner {
}
}
-impl BucketAssigner for DynamicBucketAssigner {
+impl BucketAssigner for DynamicBucketIndex {
async fn assign_batch(
&mut self,
batch: &RecordBatch,
@@ -1005,6 +1221,79 @@ mod tests {
// -- PartitionIndex tests --
+ #[test]
+ fn input_allocation_uses_owned_holes_without_writer_maintainers() {
+ let mut index = PartitionIndex::empty(1);
+ index.hash_to_bucket.insert(10, 3);
+ index.hash_to_bucket.insert(20, 0);
+ index.all_buckets.insert(0);
+ index.bucket_ids.push(0);
+ // Java restores foreign mappings but only keeps owned buckets as
+ // candidates for new keys. Historical max bucket 3 cannot hide id 2.
+ assert_eq!(
+ index.allocate(10, 4, |bucket| bucket % 2 == 0).unwrap(),
+ (3, false)
+ );
+ assert_eq!(
+ index.allocate(30, 4, |bucket| bucket % 2 == 0).unwrap(),
+ (2, true)
+ );
+ let (bucket, new) = index.allocate(40, 4, |bucket| bucket % 2 ==
0).unwrap();
+ assert!(new);
+ assert!([0, 2].contains(&bucket));
+ assert_eq!(
+ index.allocate(40, 4, |bucket| bucket % 2 == 0).unwrap(),
+ (bucket, false)
+ );
+ assert!(index.bucket_maintainers.is_empty());
+ }
+
+ #[tokio::test]
+ async fn
requested_hash_reads_cross_chunks_and_validate_complete_file_size() {
+ let io = crate::io::FileIOBuilder::new("memory").build().unwrap();
+ let count = HASH_INDEX_READ_CHUNK_SIZE as usize / 4 + 3;
+ let hashes = (0..count as i32).collect::<Vec<_>>();
+ let meta = write_hash_index(&io, "memory:/input-hash-chunks/index",
&hashes)
+ .await
+ .unwrap();
+ let path = format!("memory:/input-hash-chunks/index/{}",
meta.file_name);
+ let requested =
+ HashSet::from([0, count as i32 - 4, count as i32 - 3, count as i32
- 1, -1]);
+ let found = HashIndexFile::read_requested(&io, &path, meta.row_count,
&requested)
+ .await
+ .unwrap();
+ assert_eq!(
+ found,
+ requested
+ .difference(&HashSet::from([-1]))
+ .copied()
+ .collect()
+ );
+ assert_eq!(found.len(), 4);
+ assert!(
+ HashIndexFile::read_requested(&io, &path, meta.row_count - 1,
&requested)
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("Corrupt HASH")
+ );
+ let output =
io.new_output("memory:/input-hash-chunks/trailing").unwrap();
+ output
+ .write(bytes::Bytes::from_static(&[0, 0, 0, 1, 2]))
+ .await
+ .unwrap();
+ assert!(HashIndexFile::read_requested(
+ &io,
+ "memory:/input-hash-chunks/trailing",
+ 1,
+ &requested
+ )
+ .await
+ .unwrap_err()
+ .to_string()
+ .contains("Corrupt HASH"));
+ }
+
#[test]
fn test_assign_new_keys() {
let mut index = PartitionIndex::empty(3);
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 51999aff..86b8f54d 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -37,7 +37,6 @@ mod branch_write_tests;
mod bucket_assigner;
mod bucket_assigner_constant;
mod bucket_assigner_cross;
-mod bucket_assigner_dynamic;
mod bucket_assigner_fixed;
mod bucket_filter;
mod bucket_function;
@@ -58,6 +57,9 @@ mod data_file_writer;
mod de_vector_read;
mod de_vector_scan;
mod dedicated_format_file_writer;
+mod dynamic_bucket_assigner;
+mod dynamic_bucket_assigner_router;
+mod dynamic_bucket_index;
mod external_path;
mod format_partition;
mod format_partition_location;
@@ -214,6 +216,10 @@ pub use consumer_manager::ConsumerManager;
pub use cow_writer::{CopyOnWriteMergeWriter, FileInfo};
pub use data_evolution_writer::{DataEvolutionDeleteWriter,
DataEvolutionWriter};
pub use de_vector_scan::PreparedVectorSearchFilter;
+pub use dynamic_bucket_assigner::{DynamicBucketAssigner,
DynamicBucketAssignerBuilder};
+pub use dynamic_bucket_assigner_router::{
+ DynamicBucketAssignerRouter, DynamicBucketAssignerRouting,
+};
pub use format_partition::{
format_partition_value, parse_format_partition_value,
FormatTablePartitionPaths,
};
diff --git a/crates/paimon/src/table/table_write.rs
b/crates/paimon/src/table/table_write.rs
index 71e20b9a..9a18df92 100644
--- a/crates/paimon/src/table/table_write.rs
+++ b/crates/paimon/src/table/table_write.rs
@@ -32,13 +32,13 @@ use crate::spec::{
use crate::table::bucket_assigner::{BatchAssignOutput, BucketAssignerEnum,
PartitionBucketKey};
use crate::table::bucket_assigner_constant::ConstantBucketAssigner;
use crate::table::bucket_assigner_cross::CrossPartitionAssigner;
-use crate::table::bucket_assigner_dynamic::DynamicBucketAssigner;
use crate::table::bucket_assigner_fixed::FixedBucketAssigner;
use crate::table::bucket_function::validate_bucket_function;
use crate::table::commit_message::CommitMessage;
use crate::table::data_file_index_writer::FileIndexOptions;
use crate::table::data_file_writer::DataFileWriter;
use
crate::table::dedicated_format_file_writer::AppendDedicatedFormatFileWriter;
+use crate::table::dynamic_bucket_index::DynamicBucketIndex;
use crate::table::format_table_writer::FormatTableWriter;
use crate::table::kv_file_writer::{KeyValueFileWriter, KeyValueWriteConfig};
use crate::table::partition_filter::PartitionFilter;
@@ -470,7 +470,7 @@ impl TableWrite {
merge_engine,
)?))
} else if is_dynamic_bucket {
- BucketAssignerEnum::Dynamic(Box::new(DynamicBucketAssigner::new(
+ BucketAssignerEnum::Dynamic(Box::new(DynamicBucketIndex::new(
table.clone(),
partition_field_indices,
primary_key_indices.clone(),
