From 93649b8e3cc77d4f6ea2ba10b5af083d37c7009f Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 23 Oct 2026 15:56:35 +0800 Subject: [PATCH 1/2] [core] Add schema-based bucket routing and postpone input planning --- bindings/python/README.md | 31 + bindings/python/project-description.md | 8 + .../python/pypaimon_rust/datafusion.pyi | 24 + bindings/python/src/context.rs | 2 + bindings/python/src/lib.rs | 1 + bindings/python/src/schema.rs | 2 +- bindings/python/src/write.rs | 5 +- bindings/python/src/write_routing.rs | 208 +++ crates/paimon/src/spec/binary_row.rs | 107 +- .../src/table/authorized_write_tests.rs | 2 +- crates/paimon/src/table/branch_write_tests.rs | 2 +- crates/paimon/src/table/bucket_router.rs | 547 ++++++++ crates/paimon/src/table/format_table_scan.rs | 76 +- crates/paimon/src/table/mod.rs | 5 + crates/paimon/src/table/partition_filter.rs | 54 +- .../paimon/src/table/postpone_bucket_plan.rs | 108 +- .../src/table/postpone_bucket_planner.rs | 1149 +++++++++++++++-- .../src/table/postpone_fixed_bucket_router.rs | 128 +- .../src/table/postpone_fixed_bucket_write.rs | 12 +- .../postpone_fixed_bucket_write_builder.rs | 14 +- .../paimon/tests/nullable_primary_key_test.rs | 163 ++- 21 files changed, 2371 insertions(+), 277 deletions(-) create mode 100644 bindings/python/src/write_routing.rs create mode 100644 crates/paimon/src/table/bucket_router.rs diff --git a/bindings/python/README.md b/bindings/python/README.md index c12c8f3f6..981474f78 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -336,6 +336,37 @@ Spark's temporary-file measurement, automatic rescaling or compaction. Deletion-vector tables are rejected by this writer. A failed write cannot be prepared, and close deletes only unprepared output. +Engines can compute buckets before creating writers. `BucketRouter` uses the +table schema's bucket keys and Java hashing rules, and returns an Arrow `int32` +array in input order. It accepts reordered columns and ignores unused payload +fields. It performs no catalog access or file writes: + +```python +from pypaimon_rust.datafusion import BucketRouter, PostponeBucketPlanner + +buckets = BucketRouter(table.schema()).route(batch) # Fixed-bucket table +``` + +For POSTPONE inputs, resolve the known layouts through the write builder, then +collect statistics across the complete distributed input before inferring new +counts. The planner retains no input batches: + +```python +builder = table.new_postpone_fixed_bucket_write_builder() +known = builder.resolve_bucket_plan() +planner = PostponeBucketPlanner(table.schema(), known) +statistics = planner.input_statistics(batch) +plan = planner.plan(statistics) +buckets = BucketRouter(table.schema()).with_bucket_plan(plan).route(batch) +``` + +Each worker may call `input_statistics` independently. Concatenate its results +before calling `plan`; duplicate partition rows are summed before inference. +Statistics contain partition columns followed by `row_count: int64` and +`data_size: int64`. These columns are interpreted by position, so partition names +may coincide with statistics names. Existing layouts remain unchanged. The +shared result can configure every worker with `with_bucket_plan(plan)`. + Configure overwrite on the batch builder: ```python diff --git a/bindings/python/project-description.md b/bindings/python/project-description.md index 870373f34..f1256e9c7 100644 --- a/bindings/python/project-description.md +++ b/bindings/python/project-description.md @@ -85,6 +85,14 @@ updates, including incremental row-ID updaters, expose the same chainable setter `with_commit_user` is available on batch, stream, and postpone fixed bucket write builders, and the identity applies to both writers and committers from that builder. +`BucketRouter(table.schema()).route(batch)` assigns fixed buckets in Rust and +returns an Arrow array without creating a writer. POSTPONE engines can use +`PostponeBucketPlanner` to collect mergeable per-partition statistics, infer a +shared plan, and configure that router with `with_bucket_plan(plan)`. Collect +statistics across all input batches before planning; the planner stores schema +and counts only. Existing layouts are resolved separately by the write builder +using write authorization, without requiring SELECT. + ### Typed MERGE Actions Batch and stream `merge_into` accept ordered `WhenMatched` and `WhenNotMatched` diff --git a/bindings/python/python/pypaimon_rust/datafusion.pyi b/bindings/python/python/pypaimon_rust/datafusion.pyi index fb50bd68d..9be54bfaf 100644 --- a/bindings/python/python/pypaimon_rust/datafusion.pyi +++ b/bindings/python/python/pypaimon_rust/datafusion.pyi @@ -881,6 +881,30 @@ class BatchWriteBuilder: def new_update(self) -> BatchTableUpdate: ... def new_commit(self) -> BatchTableCommit: ... +class BucketRouter: + """Pure input-order bucket assignment using the table schema's Java rules. + + No catalog access or file writes occur. POSTPONE routing requires a known + plan or a configured default; an explicit plan replaces that default. + """ + def __init__(self, schema: TableSchema) -> None: ... + def with_bucket_plan(self, plan: pyarrow.RecordBatch) -> BucketRouter: ... + def route(self, batch: pyarrow.RecordBatch) -> pyarrow.Int32Array: ... + +class PostponeBucketPlanner: + """Pure per-partition statistics and bucket-count inference. + + Statistics contain partition columns in schema order followed by + row_count: int64 and data_size: int64, interpreted positionally. Combine + statistics for each input partition before calling plan(). Input batches + are never retained. Known bucket counts are preserved. + """ + def __init__(self, schema: TableSchema, + known_plan: Optional[pyarrow.RecordBatch] = None) -> None: ... + def current_plan(self) -> pyarrow.RecordBatch: ... + def input_statistics(self, batch: pyarrow.RecordBatch) -> pyarrow.RecordBatch: ... + def plan(self, statistics: pyarrow.RecordBatch) -> pyarrow.RecordBatch: ... + class PostponeFixedBucketWriteBuilder: """Real-bucket batch writes for a bucket=-2 primary-key table. diff --git a/bindings/python/src/context.rs b/bindings/python/src/context.rs index 75af55b72..c11790612 100644 --- a/bindings/python/src/context.rs +++ b/bindings/python/src/context.rs @@ -1039,6 +1039,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> this.add_class::()?; this.add_class::()?; this.add_class::()?; + this.add_class::()?; + this.add_class::()?; this.add_class::()?; this.add_class::()?; this.add_class::()?; diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 69f081bb0..e22064dc0 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -47,6 +47,7 @@ mod variant; mod vector_search; mod vindex; mod write; +mod write_routing; // ---- #285: observability ---- mod partition; mod snapshot; diff --git a/bindings/python/src/schema.rs b/bindings/python/src/schema.rs index a40eeb062..055a51e36 100644 --- a/bindings/python/src/schema.rs +++ b/bindings/python/src/schema.rs @@ -22,7 +22,7 @@ use pyo3::prelude::*; #[pyclass(name = "TableSchema", module = "pypaimon_rust.datafusion")] pub struct PyTableSchema { - inner: TableSchema, + pub(crate) inner: TableSchema, } impl PyTableSchema { diff --git a/bindings/python/src/write.rs b/bindings/python/src/write.rs index af66d51ed..7a8b0b91a 100644 --- a/bindings/python/src/write.rs +++ b/bindings/python/src/write.rs @@ -427,7 +427,7 @@ impl PyPostponeFixedBucketWriteBuilder { let plan = py .detach(|| runtime().block_on(builder.resolve_bucket_plan())) .map_err(to_py_err)?; - plan.to_arrow(table) + plan.to_arrow(table.schema()) .map_err(to_py_err)? .to_pyarrow(py) .map(Bound::unbind) @@ -454,7 +454,8 @@ impl PyPostponeFixedBucketWriteBuilder { ) -> PyResult> { let plan = RecordBatch::from_pyarrow_bound(plan)?; slf.plan = Some( - PostponeBucketPlan::from_arrow(&slf.batch.context.table, &plan).map_err(to_py_err)?, + PostponeBucketPlan::from_arrow(slf.batch.context.table.schema(), &plan) + .map_err(to_py_err)?, ); Ok(slf) } diff --git a/bindings/python/src/write_routing.rs b/bindings/python/src/write_routing.rs new file mode 100644 index 000000000..0eb20ceb8 --- /dev/null +++ b/bindings/python/src/write_routing.rs @@ -0,0 +1,208 @@ +// 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. + +//! 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::schema::PyTableSchema; +use arrow::array::Array; +use arrow::pyarrow::{FromPyArrow, ToPyArrow}; +use arrow::record_batch::RecordBatch; +use paimon::spec::TableSchema; +use paimon::table::{BucketRouter, PostponeBucketPlan, PostponeBucketPlanner}; +use pyo3::prelude::*; + +#[pyclass(name = "BucketRouter", module = "pypaimon_rust.datafusion")] +pub struct PyBucketRouter { + schema: TableSchema, + inner: BucketRouter, +} + +#[pymethods] +impl PyBucketRouter { + #[new] + fn new(schema: &PyTableSchema) -> PyResult { + let inner = BucketRouter::new(&schema.inner).map_err(to_py_err)?; + Ok(Self { + schema: schema.inner.clone(), + inner, + }) + } + + fn with_bucket_plan<'py>( + mut slf: PyRefMut<'py, Self>, + plan: &Bound<'_, PyAny>, + ) -> PyResult> { + let batch = RecordBatch::from_pyarrow_bound(plan)?; + let plan = PostponeBucketPlan::from_arrow(&slf.schema, &batch).map_err(to_py_err)?; + slf.inner = slf + .inner + .clone() + .with_bucket_plan(plan) + .map_err(to_py_err)?; + Ok(slf) + } + + fn route(&self, py: Python<'_>, batch: &Bound<'_, PyAny>) -> PyResult> { + let batch = write_batch_from_pyarrow(batch)?; + let buckets = py.detach(|| self.inner.route(&batch)).map_err(to_py_err)?; + buckets.to_data().to_pyarrow(py).map(Bound::unbind) + } +} + +#[pyclass(name = "PostponeBucketPlanner", module = "pypaimon_rust.datafusion")] +pub struct PyPostponeBucketPlanner { + schema: TableSchema, + inner: PostponeBucketPlanner, +} + +#[pymethods] +impl PyPostponeBucketPlanner { + #[new] + #[pyo3(signature = (schema, known_plan=None))] + fn new(schema: &PyTableSchema, known_plan: Option<&Bound<'_, PyAny>>) -> PyResult { + let known = match known_plan { + Some(plan) => { + let batch = RecordBatch::from_pyarrow_bound(plan)?; + PostponeBucketPlan::from_arrow(&schema.inner, &batch).map_err(to_py_err)? + } + None => PostponeBucketPlan::default(), + }; + let inner = PostponeBucketPlanner::new(&schema.inner, known).map_err(to_py_err)?; + Ok(Self { + schema: schema.inner.clone(), + inner, + }) + } + + fn current_plan(&self, py: Python<'_>) -> PyResult> { + self.inner + .current_plan() + .to_arrow(&self.schema) + .map_err(to_py_err)? + .to_pyarrow(py) + .map(Bound::unbind) + } + + fn input_statistics(&self, py: Python<'_>, batch: &Bound<'_, PyAny>) -> PyResult> { + let batch = write_batch_from_pyarrow(batch)?; + py.detach(|| self.inner.input_statistics(&batch)) + .map_err(to_py_err)? + .to_pyarrow(py) + .map(Bound::unbind) + } + + fn plan(&mut self, py: Python<'_>, statistics: &Bound<'_, PyAny>) -> PyResult> { + let statistics = RecordBatch::from_pyarrow_bound(statistics)?; + py.detach(|| self.inner.plan(&statistics)) + .map_err(to_py_err)? + .to_arrow(&self.schema) + .map_err(to_py_err)? + .to_pyarrow(py) + .map(Bound::unbind) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use paimon::spec::{DataType, IntType, Schema, VarCharType}; + use pyo3::ffi::c_str; + use pyo3::types::{PyDict, PyModule}; + + #[test] + fn routing_and_planning_use_typed_arrow_inputs_without_table_access() { + Python::attach(|py| { + let package = PyModule::new(py, "pypaimon_rust").unwrap(); + crate::pypaimon_rust(py, &package).unwrap(); + py.import("sys") + .unwrap() + .getattr("modules") + .unwrap() + .set_item("pypaimon_rust", &package) + .unwrap(); + let locals = PyDict::new(py); + locals + .set_item("sdk", package.getattr("datafusion").unwrap()) + .unwrap(); + for (name, bucket, function) in [("fixed", "5", "mod"), ("postpone", "-2", "default")] { + 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", bucket) + .option("bucket-key", "id") + .option("bucket-function.type", function) + .option("postpone.target-row-num-per-bucket", "2") + .option("postpone.batch-write-fixed-bucket.max-parallelism", "8") + .build() + .unwrap(); + locals + .set_item( + name, + Py::new(py, PyTableSchema::new(TableSchema::new(0, &schema))).unwrap(), + ) + .unwrap(); + } + py.run( + c_str!(r#" +import pyarrow as pa +batch = pa.RecordBatch.from_arrays([ + pa.array([-17, -1, 0, 17], type=pa.int32()), + pa.array(['a', 'a', 'a', 'b']), +], names=['id', 'pt']) +router = sdk.BucketRouter(fixed) +buckets = router.route(batch) +assert isinstance(buckets, pa.Int32Array) +assert buckets.to_pylist() == [3, 4, 0, 2] +assert buckets.null_count == 0 + +planner = sdk.PostponeBucketPlanner(postpone) +assert planner.current_plan().schema.names == ['pt', 'total_buckets'] +assert planner.current_plan().num_rows == 0 +statistics = planner.input_statistics(batch) +assert statistics.schema.names == ['pt', 'row_count', 'data_size'] +rows = dict(zip(statistics.column(0).to_pylist(), statistics.column(1).to_pylist())) +assert rows == {'a': 3, 'b': 1} +assert statistics.column(2).to_pylist() == [0] * statistics.num_rows +plan = planner.plan(statistics) +assert dict(zip(plan.column(0).to_pylist(), plan.column(1).to_pylist())) == {'a': 2, 'b': 1} +postpone_router = sdk.BucketRouter(postpone) +assert postpone_router.with_bucket_plan(plan) is postpone_router +assert all(0 <= bucket < (2 if partition == 'a' else 1) + for partition, bucket in zip(batch.column(1).to_pylist(), postpone_router.route(batch).to_pylist())) +known = sdk.PostponeBucketPlanner(postpone, plan) +assert known.input_statistics(batch).num_rows == 0 +bad = pa.RecordBatch.from_arrays([pa.array(['c']), pa.array([-1], type=pa.int64()), + pa.array([0], type=pa.int64())], names=['pt', 'row_count', 'data_size']) +try: + planner.plan(bad) +except ValueError: + pass +else: + raise AssertionError('negative statistics must be rejected') +assert planner.current_plan().equals(plan) +"#), + Some(&locals), + Some(&locals), + ).unwrap(); + }); + } +} diff --git a/crates/paimon/src/spec/binary_row.rs b/crates/paimon/src/spec/binary_row.rs index 98cf25a1d..b72779c38 100644 --- a/crates/paimon/src/spec/binary_row.rs +++ b/crates/paimon/src/spec/binary_row.rs @@ -410,7 +410,7 @@ impl BinaryRow { let field = &fields[field_idx]; match extract_datum_from_arrow(batch, row_idx, field_idx, field.data_type())? { Some(datum) => builder.write_datum(pos, &datum, field.data_type()), - None => builder.set_null_at(pos), + None => builder.write_null(pos, field.data_type()), } } Ok(builder.build()) @@ -425,7 +425,7 @@ impl BinaryRow { for (pos, (datum_opt, data_type)) in datums.iter().enumerate() { match datum_opt { Some(datum) => builder.write_datum(pos, datum, data_type), - None => builder.set_null_at(pos), + None => builder.write_null(pos, data_type), } } @@ -476,6 +476,22 @@ impl BinaryRowBuilder { self.data[offset..offset + 8].fill(0); } + /// Java RowHelper writes NULL non-compact decimals and timestamps through + /// their value writers. They retain a zeroed variable slot and its offset + /// for later updates, and that storage participates in the row hash. + pub fn write_null(&mut self, pos: usize, data_type: &DataType) { + self.set_null_at(pos); + let reserved = match data_type { + DataType::Decimal(decimal) if decimal.precision() > 18 => 16, + DataType::Timestamp(timestamp) if timestamp.precision() > 3 => 8, + DataType::LocalZonedTimestamp(timestamp) if timestamp.precision() > 3 => 8, + _ => return, + }; + let variable_offset = self.data.len(); + self.data.resize(variable_offset + reserved, 0); + self.write_long(pos, ((variable_offset as u64) << 32) as i64); + } + pub fn write_boolean(&mut self, pos: usize, value: bool) { let offset = self.field_offset(pos); self.data[offset] = u8::from(value); @@ -718,7 +734,7 @@ pub fn datums_to_binary_row(datums: &[(&Option, &DataType)]) -> Vec { builder.write_datum(pos, datum, data_type); } None => { - builder.set_null_at(pos); + builder.write_null(pos, data_type); } } } @@ -1539,7 +1555,7 @@ fn write_typed_value( pos: usize, row_idx: usize, typed_col: &TypedColumn, - _data_type: &DataType, + data_type: &DataType, ) -> crate::Result<()> { use arrow_array::Array; match typed_col { @@ -1647,7 +1663,7 @@ fn write_typed_value( } TypedColumn::Decimal128(arr, precision, _scale) => { if arr.is_null(row_idx) { - builder.set_null_at(pos); + builder.write_null(pos, data_type); } else { let unscaled = arr.value(row_idx); if *precision <= 18 { @@ -1717,7 +1733,7 @@ fn write_typed_value( } TypedColumn::TimestampUs(arr) => { if arr.is_null(row_idx) { - builder.set_null_at(pos); + builder.write_null(pos, data_type); } else { let (millis, nanos) = timestamp_parts_from_micros(arr.value(row_idx)); builder.write_timestamp_non_compact(pos, millis, nanos); @@ -1725,7 +1741,7 @@ fn write_typed_value( } TypedColumn::TimestampNs(arr) => { if arr.is_null(row_idx) { - builder.set_null_at(pos); + builder.write_null(pos, data_type); } else { let (millis, nanos) = timestamp_parts_from_nanos(arr.value(row_idx)); builder.write_timestamp_non_compact(pos, millis, nanos); @@ -2368,6 +2384,83 @@ mod tests { assert_eq!(nano, nano_of_milli); } + #[test] + fn nullable_non_compact_fields_match_java_serializer_bytes_and_hashes() { + use crate::spec::{ + DataField, DecimalType, LocalZonedTimestampType, TimestampType, VarCharType, + }; + use arrow_array::{new_null_array, StringArray}; + use std::sync::Arc; + + // Current Java InternalRowSerializer, via RowHelper, for + // GenericRow.of(NULL, BinaryString.fromString("abcdefghijk")). + let compact = + "000100000000000000000000000000000b000000180000006162636465666768696a6b0000000000"; + let decimal = "000100000000000000000000180000000b00000028000000000000000000000000000000000000006162636465666768696a6b0000000000"; + let timestamp = "000100000000000000000000180000000b0000002000000000000000000000006162636465666768696a6b0000000000"; + let mut cases = Vec::new(); + for precision in [18, 19, 20] { + cases.push(( + DataType::Decimal(DecimalType::new(precision, 2).unwrap()), + if precision <= 18 { compact } else { decimal }, + if precision <= 18 { + 1_425_286_886 + } else { + 1_558_048_601 + }, + )); + } + for precision in [3, 4, 9] { + let bytes = if precision <= 3 { compact } else { timestamp }; + let hash = if precision <= 3 { + 1_425_286_886 + } else { + -684_602_128 + }; + cases.push(( + DataType::Timestamp(TimestampType::new(precision).unwrap()), + bytes, + hash, + )); + cases.push(( + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(precision).unwrap()), + bytes, + hash, + )); + } + let string_type = DataType::VarChar(VarCharType::string_type()); + let text = Some(Datum::String("abcdefghijk".into())); + let null = None; + for (data_type, bytes, hash) in cases { + let fields = vec![ + DataField::new(0, "key".into(), data_type.clone()), + DataField::new(1, "next".into(), string_type.clone()), + ]; + let schema = crate::arrow::build_target_arrow_schema(&fields).unwrap(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + new_null_array(schema.field(0).data_type(), 1), + Arc::new(StringArray::from(vec!["abcdefghijk"])), + ], + ) + .unwrap(); + let expected = hex::decode(bytes).unwrap(); + let scalar = BinaryRow::from_arrow(&batch, 0, &[0, 1], &fields).unwrap(); + let batched = batch_build_binary_rows(&batch, &[0, 1], &fields).unwrap(); + let datum = + BinaryRow::from_datums(&[(None, &data_type), (text.as_ref(), &string_type)]); + let serialized = datums_to_binary_row(&[(&null, &data_type), (&text, &string_type)]); + let serialized = BinaryRow::from_serialized_bytes(&serialized).unwrap(); + for row in [&scalar, &batched[0], &datum, &serialized] { + assert_eq!(row.data(), expected, "{data_type:?}"); + assert_eq!(row.hash_code(), hash, "{data_type:?}"); + assert!(row.is_null_at(0)); + assert_eq!(row.get_string(1).unwrap(), "abcdefghijk"); + } + } + } + #[test] fn test_batch_vs_per_row_equivalence() { use arrow_array::{Int32Array, StringArray, Time32MillisecondArray}; diff --git a/crates/paimon/src/table/authorized_write_tests.rs b/crates/paimon/src/table/authorized_write_tests.rs index c78a7b15d..babcd4705 100644 --- a/crates/paimon/src/table/authorized_write_tests.rs +++ b/crates/paimon/src/table/authorized_write_tests.rs @@ -320,7 +320,7 @@ async fn rest_postpone_resolve_bucket_plan_requires_write_identity_without_selec .resolve_bucket_plan() .await .unwrap() - .to_arrow(&table) + .to_arrow(table.schema()) .unwrap(); assert_eq!(exported.num_rows(), 2); assert_eq!(exported.schema().field(0).name(), "p"); diff --git a/crates/paimon/src/table/branch_write_tests.rs b/crates/paimon/src/table/branch_write_tests.rs index 0adbcee16..115a35574 100644 --- a/crates/paimon/src/table/branch_write_tests.rs +++ b/crates/paimon/src/table/branch_write_tests.rs @@ -786,7 +786,7 @@ async fn fixed_bucket_write(table: &Table, value: i32, buckets: i32, overwrite: Arc::new(Int32Array::from(vec![buckets])) as arrow_array::ArrayRef, )]) .unwrap(); - let plan = PostponeBucketPlan::from_arrow(table, &plan_batch).unwrap(); + let plan = PostponeBucketPlan::from_arrow(table.schema(), &plan_batch).unwrap(); let builder = table .new_postpone_fixed_bucket_write_builder() .unwrap() diff --git a/crates/paimon/src/table/bucket_router.rs b/crates/paimon/src/table/bucket_router.rs new file mode 100644 index 000000000..af1f581cd --- /dev/null +++ b/crates/paimon/src/table/bucket_router.rs @@ -0,0 +1,547 @@ +// 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. + +//! Bucket assignment for inputs distributed before constructing file writers. + +use super::bucket_function::{batch_bucket_ids, validate_bucket_function}; +use super::postpone_bucket_plan::data_invalid; +use super::postpone_fixed_bucket_router::validate_postpone_fixed_bucket_schema; +use super::table_write::take_rows; +use super::write_batch_normalize::normalize_write_array; +use crate::spec::{ + batch_to_serialized_bytes, BucketFunctionType, CoreOptions, DataField, TableSchema, + POSTPONE_BUCKET, +}; +use crate::table::PostponeBucketPlan; +use crate::Result; +use arrow_array::{Int32Array, RecordBatch}; +use arrow_schema::SchemaRef; +use std::collections::{HashMap, HashSet}; + +/// Compute the bucket of every input row without opening files or authorizing a table. +/// +/// This is the same bucket computation used by fixed-bucket writers and by +/// Java's `FixedBucketRowKeyExtractor` and `PostponeBucketRouter`. Inputs can +/// omit payload fields and reorder columns, but must contain each partition +/// and bucket key exactly once. Output positions retain the input row order. +#[derive(Debug, Clone)] +pub struct BucketRouter { + fields: Vec, + arrow_schema: SchemaRef, + partition_indices: Vec, + bucket_indices: Vec, + bucket_function: BucketFunctionType, + total_buckets: i32, + plan: PostponeBucketPlan, +} + +impl BucketRouter { + /// Construct an owned router for a fixed-bucket or postpone-bucket schema. + /// + /// Postpone routing uses the configured default until an explicit plan is + /// supplied. Without either, routing an unknown input partition fails. + pub fn new(schema: &TableSchema) -> Result { + let options = CoreOptions::new(schema.options()); + let total_buckets = options.bucket(); + if options.is_format_table() || (total_buckets <= 0 && total_buckets != POSTPONE_BUCKET) { + return Err(crate::Error::Unsupported { + message: "Bucket routing requires a Paimon fixed-bucket or postpone-bucket table" + .into(), + }); + } + if total_buckets == POSTPONE_BUCKET { + validate_postpone_fixed_bucket_schema(schema)?; + } + + let bucket_keys = schema.bucket_keys(); + let required = schema + .partition_keys() + .iter() + .chain(&bucket_keys) + .map(String::as_str) + .collect::>(); + let fields = schema + .fields() + .iter() + .filter(|field| required.contains(field.name())) + .cloned() + .collect::>(); + 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::>>() + }; + let partition_indices = indices(schema.partition_keys())?; + let bucket_indices = indices(&bucket_keys)?; + let bucket_function = options.bucket_function_type()?; + let bucket_fields = bucket_indices + .iter() + .map(|&index| fields[index].clone()) + .collect::>(); + validate_bucket_function(bucket_function, &bucket_fields)?; + let plan = if total_buckets == POSTPONE_BUCKET { + options + .postpone_default_bucket_num()? + .map(PostponeBucketPlan::with_default_bucket_num) + .unwrap_or_default() + } else { + PostponeBucketPlan::default() + }; + Ok(Self { + arrow_schema: crate::arrow::build_target_arrow_schema(&fields)?, + fields, + partition_indices, + bucket_indices, + bucket_function, + total_buckets, + plan, + }) + } + + /// Replace postpone routing metadata, including any configured default. + pub fn with_bucket_plan(mut self, plan: PostponeBucketPlan) -> Result { + if self.total_buckets != POSTPONE_BUCKET { + return Err(data_invalid( + "A bucket plan requires a postpone-bucket schema", + )); + } + self.plan = plan; + Ok(self) + } + + /// Return non-null Int32 bucket IDs in input row order. + pub fn route(&self, batch: &RecordBatch) -> Result { + Ok(self.route_with_partitions(batch)?.1) + } + + pub(super) fn plan(&self) -> &PostponeBucketPlan { + &self.plan + } + + pub(super) fn plan_mut(&mut self) -> &mut PostponeBucketPlan { + &mut self.plan + } + + pub(super) fn route_with_partitions( + &self, + batch: &RecordBatch, + ) -> Result<(Vec>, 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::>>()?; + // 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 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)?; + return Ok((partitions, Int32Array::from(buckets))); + } + + let mut groups: HashMap> = HashMap::new(); + for (row, partition) in partitions.iter().enumerate() { + let count = self.plan.total_buckets(partition).ok_or_else(|| { + data_invalid("Postpone bucket plan does not contain an input partition") + })?; + groups.entry(count).or_default().push(row); + } + let mut buckets = vec![0; batch.num_rows()]; + for (count, rows) in groups { + let group = if rows.len() == input.num_rows() { + input.clone() + } else { + take_rows(&input, &rows)? + }; + for (row, bucket) in rows.into_iter().zip(self.route_count(&group, count)?) { + buckets[row] = bucket; + } + } + Ok((partitions, Int32Array::from(buckets))) + } + + fn route_count(&self, batch: &RecordBatch, count: i32) -> Result> { + if count == 1 || (self.total_buckets == POSTPONE_BUCKET && self.bucket_indices.is_empty()) { + return Ok(vec![0; batch.num_rows()]); + } + batch_bucket_ids( + batch, + &self.bucket_indices, + &self.fields, + self.bucket_function, + count, + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{ + BigIntType, DataType, DateType, DecimalType, IntType, Schema, TimestampType, VarCharType, + }; + use arrow_array::{ + types::Int8Type, Array, ArrayRef, Date32Array, Decimal128Array, DictionaryArray, + Int64Array, Int8Array, LargeStringArray, StringArray, StructArray, + TimestampNanosecondArray, + }; + use arrow_schema::{DataType as ArrowType, Field, Schema as ArrowSchema}; + use std::sync::Arc; + + fn batch(columns: Vec<(&str, ArrayRef)>) -> RecordBatch { + let fields = columns + .iter() + .map(|(name, array)| Field::new(*name, array.data_type().clone(), true)) + .collect::>(); + RecordBatch::try_new( + Arc::new(ArrowSchema::new(fields)), + columns.into_iter().map(|(_, array)| array).collect(), + ) + .unwrap() + } + + fn scalar_schema(data_type: DataType, function: &str, count: i32) -> TableSchema { + TableSchema::new( + 0, + &Schema::builder() + .column("id", data_type) + .option("bucket-key", "id") + .option("bucket-function.type", function) + .option("bucket", count.to_string()) + .build() + .unwrap(), + ) + } + + fn check_java_scalar(data_type: DataType, values: ArrayRef, expected: &[i32]) { + let router = BucketRouter::new(&scalar_schema(data_type, "default", 100)).unwrap(); + assert_eq!( + router.route(&batch(vec![("id", values)])).unwrap().values(), + expected + ); + } + + #[test] + fn fixed_java_reference_keys_ignore_payload_and_input_order() { + // FixedBucketRowKeyExtractorTest.testBucket: row (5, 6, 7), 100 buckets. + // A nested payload is deliberately unrelated to any routing field. + let nested = Arc::new(StructArray::from(vec![( + Arc::new(Field::new("child", ArrowType::Int64, true)), + Arc::new(Int64Array::from(vec![999])) as ArrayRef, + )])); + let input = batch(vec![ + ("c", Arc::new(Int32Array::from(vec![7]))), + ("payload", nested), + ("b", Arc::new(Int32Array::from(vec![6]))), + ("a", Arc::new(Int32Array::from(vec![5]))), + ]); + for (keys, expected) in [("a", 96), ("a,b", 27), ("a,b,c", 40)] { + let schema = TableSchema::new( + 0, + &Schema::builder() + .column("a", DataType::Int(IntType::new())) + .column("b", DataType::Int(IntType::new())) + .column("c", DataType::Int(IntType::new())) + .option("bucket", "100") + .option("bucket-key", keys) + .build() + .unwrap(), + ); + let router = BucketRouter::new(&schema).unwrap(); + assert_eq!(router.route(&input).unwrap().values(), &[expected]); + if keys == "a" { + assert_eq!( + router + .route(&input.project(&[3]).unwrap()) + .unwrap() + .values(), + &[expected] + ); + } + assert!(router.route(&input.slice(0, 0)).unwrap().is_empty()); + } + } + + #[test] + fn default_scalar_values_match_current_java_binary_row_writers() { + // Generated with current Java BinaryRow/BinaryRowWriter and + // Math.abs(row.hashCode() % 100), including variable NULL slots. + check_java_scalar( + DataType::Int(IntType::new()), + Arc::new(Int32Array::from(vec![i32::MIN, -3, 0, 5, i32::MAX])), + &[19, 15, 99, 96, 21], + ); + check_java_scalar( + DataType::VarChar(VarCharType::string_type()), + Arc::new(LargeStringArray::from(vec![ + None, + Some(""), + Some("hello"), + Some("abcdefghijk"), + Some("你好"), + ])), + &[44, 19, 46, 23, 58], + ); + check_java_scalar( + DataType::Decimal(DecimalType::new(20, 2).unwrap()), + Arc::new( + Decimal128Array::from(vec![None, Some(12340), Some(-9999999), Some(0)]) + .with_precision_and_scale(20, 2) + .unwrap(), + ), + &[60, 60, 44, 85], + ); + check_java_scalar( + DataType::Date(DateType::new()), + Arc::new(Date32Array::from(vec![-3, 0, 5])), + &[15, 99, 96], + ); + check_java_scalar( + DataType::Timestamp(TimestampType::new(9).unwrap()), + Arc::new(TimestampNanosecondArray::from(vec![ + -1_000_001, + -1, + 0, + 1, + 1_000_001, + i64::MAX, + ])), + &[38, 85, 54, 32, 55, 39], + ); + } + + #[test] + fn fixed_partition_only_rows_hash_the_empty_bucket_key() { + let schema = TableSchema::new( + 0, + &Schema::builder() + .column("pt", DataType::VarChar(VarCharType::string_type())) + .partition_keys(["pt"]) + .option("bucket", "100") + .build() + .unwrap(), + ); + let router = BucketRouter::new(&schema).unwrap(); + // Current Java BinaryRow(0).hashCode() is -1670924195. Fixed routing + // still hashes that key when every table field is a partition field. + assert_eq!( + router + .route(&batch(vec![( + "pt", + Arc::new(StringArray::from(vec![Some("a"), None])), + )])) + .unwrap() + .values(), + &[95, 95] + ); + } + + #[test] + fn nullable_decimal_primary_key_uses_java_bucket() { + let schema = TableSchema::new( + 0, + &Schema::builder() + .column("id", DataType::Decimal(DecimalType::new(20, 2).unwrap())) + .primary_key(["id"]) + .option("primary-key.nullable", "true") + .option("bucket", "100") + .build() + .unwrap(), + ); + let input = batch(vec![( + "id", + Arc::new( + Decimal128Array::from(vec![None, Some(12340)]) + .with_precision_and_scale(20, 2) + .unwrap(), + ), + )]); + assert_eq!( + BucketRouter::new(&schema) + .unwrap() + .route(&input) + .unwrap() + .values(), + &[60, 60] + ); + } + + #[test] + fn mod_and_hive_preserve_nulls_and_integer_boundaries() { + let input = batch(vec![( + "id", + Arc::new(Int64Array::from(vec![ + None, + Some(i64::MIN), + Some(-3), + Some(0), + Some(i64::MAX), + ])), + )]); + let router = BucketRouter::new(&scalar_schema( + DataType::BigInt(BigIntType::new()), + "mod", + 5, + )) + .unwrap(); + assert_eq!(router.route(&input).unwrap().values(), &[0, 2, 2, 0, 2]); + // HiveBucketFunction uses the positive low 31 bits of Long.hashCode. + let router = BucketRouter::new(&scalar_schema( + DataType::BigInt(BigIntType::new()), + "hive", + 5, + )) + .unwrap(); + assert_eq!(router.route(&input).unwrap().values(), &[0, 0, 2, 0, 0]); + } + + #[test] + fn postpone_routes_interleaved_partitions_and_replaces_default() { + let schema = TableSchema::new( + 0, + &Schema::builder() + .column( + "total_buckets", + DataType::VarChar(VarCharType::string_type()), + ) + .column("id", DataType::Int(IntType::new())) + .partition_keys(["total_buckets"]) + .primary_key(["total_buckets", "id"]) + .option("primary-key.nullable", "true") + .option("bucket", "-2") + .option("postpone.default-bucket-num", "1") + .build() + .unwrap(), + ); + let input = batch(vec![ + ( + "id", + Arc::new(Int32Array::from(vec![5, -3, i32::MIN, 0, i32::MAX])), + ), + ( + "total_buckets", + Arc::new(StringArray::from(vec![ + Some("a"), + Some("b"), + None, + Some("a"), + Some("b"), + ])), + ), + ]); + let router = BucketRouter::new(&schema).unwrap(); + assert_eq!(router.route(&input).unwrap().values(), &[0; 5]); + // Both partition and final count field are called total_buckets; + // the plan codec uses their positions to retain that valid schema. + let plan = PostponeBucketPlan::from_arrow( + &schema, + &batch(vec![ + ( + "total_buckets", + Arc::new(StringArray::from(vec![Some("a"), Some("b"), None])), + ), + ("total_buckets", Arc::new(Int32Array::from(vec![3, 5, 7]))), + ]), + ) + .unwrap(); + let router = router.with_bucket_plan(plan).unwrap(); + // Java BinaryRow hashes: 260196596, -528982515, 916225219, + // -300363099 and -1125657321, reduced by each partition's count. + assert_eq!(router.route(&input).unwrap().values(), &[2, 0, 0, 0, 1]); + let empty = router + .clone() + .with_bucket_plan(PostponeBucketPlan::default()) + .unwrap(); + assert!(empty + .route(&input) + .unwrap_err() + .to_string() + .contains("does not contain")); + } + + #[test] + fn routing_rejects_ambiguous_missing_and_incompatible_columns() { + let schema = scalar_schema(DataType::Int(IntType::new()), "default", 2); + let router = BucketRouter::new(&schema).unwrap(); + let ints: ArrayRef = Arc::new(Int32Array::from(vec![1])); + assert!(router + .route(&batch(vec![("id", ints.clone()), ("id", ints.clone())])) + .unwrap_err() + .to_string() + .contains("Duplicate")); + assert!(router + .route(&batch(vec![("other", ints)])) + .unwrap_err() + .to_string() + .contains("Missing routing column 'id'")); + assert!(router + .route(&batch(vec![("id", Arc::new(Int64Array::from(vec![1])))])) + .is_err()); + let strings = Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![0]), + Arc::new(StringArray::from(vec!["value"])), + ) + .unwrap(), + ); + let string_router = BucketRouter::new(&scalar_schema( + DataType::VarChar(VarCharType::string_type()), + "default", + 2, + )) + .unwrap(); + assert!(string_router.route(&batch(vec![("id", strings)])).is_err()); + assert!(router + .with_bucket_plan(PostponeBucketPlan::default()) + .is_err()); + for count in [-2, -1, 0] { + let mut options = schema.options().clone(); + options.insert("bucket".into(), count.to_string()); + assert!(BucketRouter::new(&schema.copy_with_replaced_options(options)).is_err()); + } + assert!(BucketRouter::new(&scalar_schema( + DataType::VarChar(VarCharType::string_type()), + "mod", + 2 + )) + .is_err()); + } +} diff --git a/crates/paimon/src/table/format_table_scan.rs b/crates/paimon/src/table/format_table_scan.rs index 054b0df2e..79c2790d0 100644 --- a/crates/paimon/src/table/format_table_scan.rs +++ b/crates/paimon/src/table/format_table_scan.rs @@ -812,7 +812,7 @@ fn partition_row_from_path( let mut builder = BinaryRowBuilder::new(partition_fields.len() as i32); for (idx, value) in values.iter().enumerate() { if value == default_partition_name { - builder.set_null_at(idx); + builder.write_null(idx, partition_fields[idx].data_type()); continue; } let Some(datum) = parse_format_partition_value(value, partition_fields[idx].data_type()) @@ -837,7 +837,7 @@ fn partition_row_from_catalog_spec( source: None, })?; if value == default_partition_name { - builder.set_null_at(index); + builder.write_null(index, field.data_type()); continue; } let datum = parse_format_partition_value(value, field.data_type()).ok_or_else(|| { @@ -935,6 +935,78 @@ mod tests { use bytes::Bytes; use std::collections::HashSet; + #[test] + fn null_partition_reconstruction_uses_java_non_compact_slots() { + use crate::spec::{DecimalType, LocalZonedTimestampType, TimestampType}; + + // Java InternalRowSerializer for (NULL, "abcdefghijk"). Reconstructing + // a directory or catalog spec must match the exact partition filter. + let decimal = "000100000000000000000000180000000b00000028000000000000000000000000000000000000006162636465666768696a6b0000000000"; + let timestamp = "000100000000000000000000180000000b0000002000000000000000000000006162636465666768696a6b0000000000"; + for (data_type, expected) in [ + (DataType::Decimal(DecimalType::new(19, 2).unwrap()), decimal), + ( + DataType::Timestamp(TimestampType::new(4).unwrap()), + timestamp, + ), + ( + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(4).unwrap()), + timestamp, + ), + ] { + let fields = vec![ + DataField::new(0, "pt".into(), data_type), + DataField::new( + 1, + "code".into(), + DataType::VarChar(VarCharType::string_type()), + ), + ]; + let keys = vec!["pt".into(), "code".into()]; + let spec = HashMap::from([ + ("pt".into(), "DEFAULT".into()), + ("code".into(), "abcdefghijk".into()), + ]); + let catalog = + partition_row_from_catalog_spec(&spec, &fields, &keys, "DEFAULT").unwrap(); + let named = partition_row_from_path( + "memory:/format", + "memory:/format/pt=DEFAULT/code=abcdefghijk", + &fields, + &keys, + "DEFAULT", + false, + ) + .unwrap() + .unwrap(); + let values = partition_row_from_path( + "memory:/format", + "memory:/format/DEFAULT/abcdefghijk", + &fields, + &keys, + "DEFAULT", + true, + ) + .unwrap() + .unwrap(); + let expected = hex::decode(expected).unwrap(); + for row in [&catalog, &named, &values] { + assert_eq!(row.data(), expected); + let builder = crate::spec::PredicateBuilder::new(&fields); + let filter = PartitionFilter::from_predicate( + Predicate::and(vec![ + builder.is_null("pt").unwrap(), + builder + .equal("code", Datum::String("abcdefghijk".into())) + .unwrap(), + ]), + &fields, + ); + assert!(filter.matches_entry(&row.to_serialized_bytes()).unwrap()); + } + } + } + /// A Parquet format table in memory, partitioned by the given string keys. fn format_table(location: &str, partition_keys: &[&str], options: &[(&str, &str)]) -> Table { let mut builder = Schema::builder(); diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 0e1193fa2..51999afff 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -41,6 +41,7 @@ mod bucket_assigner_dynamic; mod bucket_assigner_fixed; mod bucket_filter; mod bucket_function; +mod bucket_router; mod chunk_shuffle; mod commit_callback; mod commit_message; @@ -206,6 +207,7 @@ pub use audit_log_table::AuditLogTable; pub use batch_vector_search_builder::BatchVectorSearchBuilder; pub use blob_resolver::{BlobReader, BlobStream}; pub use branch_manager::BranchManager; +pub use bucket_router::BucketRouter; pub use commit_callback::{CommitCallback, CommitCallbackContext}; pub use commit_message::{CommitMessage, COMMIT_MESSAGE_SERIALIZER_VERSION}; pub use consumer_manager::ConsumerManager; @@ -243,6 +245,9 @@ pub use partition_spec::parse_partition_spec_value; pub use partition_stat::PartitionStat; pub use pk_vector_bucket_split::{BucketVectorPayload, BucketVectorSearchSplit}; pub use postpone_bucket_plan::{PostponeBucketPlan, POSTPONE_BUCKET_PLAN_TOTAL_BUCKETS_FIELD}; +pub use postpone_bucket_planner::{ + PostponeBucketPlanner, POSTPONE_STATISTICS_DATA_SIZE_FIELD, POSTPONE_STATISTICS_ROW_COUNT_FIELD, +}; pub use postpone_fixed_bucket_write::{ PostponeFixedBucketTableCommit, PostponeFixedBucketTableWrite, }; diff --git a/crates/paimon/src/table/partition_filter.rs b/crates/paimon/src/table/partition_filter.rs index ad54f931e..2dd914328 100644 --- a/crates/paimon/src/table/partition_filter.rs +++ b/crates/paimon/src/table/partition_filter.rs @@ -93,7 +93,7 @@ impl PartitionFilter { Some(datum) => { builder.write_datum(i, datum, partition_fields[i].data_type()); } - None => builder.set_null_at(i), + None => builder.write_null(i, partition_fields[i].data_type()), } } let row = builder.build(); @@ -859,6 +859,58 @@ mod tests { assert!(!filter.matches_entry(&builder.build_serialized()).unwrap()); } + #[test] + fn null_non_compact_partition_candidates_match_java_bytes() { + use crate::spec::{DecimalType, LocalZonedTimestampType, TimestampType}; + + // Current Java InternalRowSerializer for (NULL, "abcdefghijk"). + let compact = + "000100000000000000000000000000000b000000180000006162636465666768696a6b0000000000"; + let decimal = "000100000000000000000000180000000b00000028000000000000000000000000000000000000006162636465666768696a6b0000000000"; + let timestamp = "000100000000000000000000180000000b0000002000000000000000000000006162636465666768696a6b0000000000"; + let mut cases = Vec::new(); + for precision in [18, 19] { + cases.push(( + DataType::Decimal(DecimalType::new(precision, 2).unwrap()), + if precision <= 18 { compact } else { decimal }, + )); + } + for precision in [3, 4] { + let bytes = if precision <= 3 { compact } else { timestamp }; + cases.push(( + DataType::Timestamp(TimestampType::new(precision).unwrap()), + bytes, + )); + cases.push(( + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(precision).unwrap()), + bytes, + )); + } + for (data_type, body) in cases { + let fields = vec![ + DataField::new(0, "pt".into(), data_type.clone()), + DataField::new( + 1, + "code".into(), + DataType::VarChar(VarCharType::string_type()), + ), + ]; + let pb = PredicateBuilder::new(&fields); + let filter = PartitionFilter::from_predicate( + Predicate::and(vec![ + pb.is_null("pt").unwrap(), + pb.equal("code", Datum::String("abcdefghijk".into())) + .unwrap(), + ]), + &fields, + ); + assert!(matches!(filter, PartitionFilter::PartitionSet { .. })); + let mut bytes = 2_i32.to_be_bytes().to_vec(); + bytes.extend(hex::decode(body).unwrap()); + assert!(filter.matches_entry(&bytes).unwrap(), "{data_type:?}"); + } + } + #[test] fn test_decode_failure_in_predicate_mode_fails_open() { let fields = partition_fields_dt(); diff --git a/crates/paimon/src/table/postpone_bucket_plan.rs b/crates/paimon/src/table/postpone_bucket_plan.rs index ff0896507..a94193d8d 100644 --- a/crates/paimon/src/table/postpone_bucket_plan.rs +++ b/crates/paimon/src/table/postpone_bucket_plan.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::spec::{batch_to_serialized_bytes, BinaryRow, Snapshot}; +use crate::spec::{batch_to_serialized_bytes, BinaryRow, DataField, Snapshot, TableSchema}; use crate::table::{Table, TableScan}; use crate::Result; use arrow_array::{new_empty_array, Array, ArrayRef, Int32Array, RecordBatch}; @@ -45,8 +45,8 @@ impl PostponeBucketPlan { default_bucket_num: None, } } - pub fn from_arrow(table: &Table, batch: &RecordBatch) -> Result { - let partition_fields = table.schema().partition_fields(); + pub fn from_arrow(schema: &TableSchema, batch: &RecordBatch) -> Result { + let partition_fields = schema.partition_fields(); let partition_count = partition_fields.len(); if batch.num_columns() != partition_count + 1 { return Err(data_invalid(format!( @@ -119,16 +119,13 @@ impl PostponeBucketPlan { } } } - Ok(Self { - bucket_counts: Arc::new(bucket_counts), - default_bucket_num: None, - }) + Ok(Self::from_counts(bucket_counts)) } /// Export known partitions using the same schema accepted by `from_arrow`. /// A configured default is not an existing partition and is not exported. - pub fn to_arrow(&self, table: &Table) -> Result { - let partition_fields = table.schema().partition_fields(); + pub fn to_arrow(&self, schema: &TableSchema) -> Result { + let partition_fields = schema.partition_fields(); let mut fields = crate::arrow::build_target_arrow_schema(&partition_fields)? .fields() .to_vec(); @@ -141,37 +138,9 @@ impl PostponeBucketPlan { entries.sort_unstable_by_key(|(partition, _)| *partition); let partitions = entries .iter() - .map(|(bytes, _)| { - let row = BinaryRow::from_serialized_bytes(bytes)?; - if row.arity() as usize != partition_fields.len() { - return Err(data_invalid( - "Postpone bucket plan partition arity does not match the table schema", - )); - } - Ok(row) - }) - .collect::>>()?; - let mut columns = Vec::::with_capacity(fields.len()); - for (index, field) in partition_fields.iter().enumerate() { - let arrays = partitions - .iter() - .map(|partition| { - crate::arrow::partition::partition_array(partition, index, field.data_type(), 1) - }) - .collect::>>()?; - let column = if arrays.is_empty() { - new_empty_array(fields[index].data_type()) - } else { - arrow_select::concat::concat( - &arrays - .iter() - .map(|array| array.as_ref()) - .collect::>(), - ) - .map_err(|error| data_invalid(error.to_string()))? - }; - columns.push(column); - } + .map(|(bytes, _)| bytes.as_slice()) + .collect::>(); + let mut columns = serialized_partition_columns(&partition_fields, &partitions)?; columns.push(Arc::new(Int32Array::from_iter_values( entries.iter().map(|(_, count)| **count), ))); @@ -237,6 +206,57 @@ impl PostponeBucketPlan { .copied() .or(self.default_bucket_num) } + + pub(super) fn explicit_bucket_counts(&self) -> &HashMap, i32> { + &self.bucket_counts + } + + pub(super) fn insert(&mut self, partition: Vec, count: i32) { + Arc::make_mut(&mut self.bucket_counts).insert(partition, count); + } +} + +/// Decode the common partition prefix of plans and input statistics. +pub(super) fn serialized_partition_columns( + fields: &[DataField], + partitions: &[&[u8]], +) -> Result> { + let rows = partitions + .iter() + .map(|bytes| { + let row = BinaryRow::from_serialized_bytes(bytes)?; + if row.arity() as usize != fields.len() { + return Err(data_invalid( + "Postpone partition arity does not match the table schema", + )); + } + Ok(row) + }) + .collect::>>()?; + let schema = crate::arrow::build_target_arrow_schema(fields)?; + fields + .iter() + .enumerate() + .map(|(index, field)| { + let arrays = rows + .iter() + .map(|row| { + crate::arrow::partition::partition_array(row, index, field.data_type(), 1) + }) + .collect::>>()?; + if arrays.is_empty() { + Ok(new_empty_array(schema.field(index).data_type())) + } else { + arrow_select::concat::concat( + &arrays + .iter() + .map(|array| array.as_ref()) + .collect::>(), + ) + .map_err(|error| data_invalid(error.to_string())) + } + }) + .collect() } #[cfg(test)] @@ -284,21 +304,21 @@ mod tests { ], ) .unwrap(); - let plan = PostponeBucketPlan::from_arrow(&table, &batch).unwrap(); - let exported = plan.to_arrow(&table).unwrap(); + let plan = PostponeBucketPlan::from_arrow(table.schema(), &batch).unwrap(); + let exported = plan.to_arrow(table.schema()).unwrap(); assert_eq!(exported.schema().field(0).name(), "pt"); assert_eq!(exported.schema().field(0).data_type(), &DataType::Utf8); assert_eq!(exported.schema().field(1), batch.schema().field(1)); assert_eq!(exported.num_rows(), 3); assert_eq!( - PostponeBucketPlan::from_arrow(&table, &exported) + PostponeBucketPlan::from_arrow(table.schema(), &exported) .unwrap() .bucket_counts, plan.bucket_counts ); let expected_schema = exported.schema(); let default = PostponeBucketPlan::with_default_bucket_num(7); - let exported = default.to_arrow(&table).unwrap(); + let exported = default.to_arrow(table.schema()).unwrap(); assert_eq!(exported.schema(), expected_schema); assert_eq!(exported.num_rows(), 0); assert_eq!(default.total_buckets(&[]), Some(7)); diff --git a/crates/paimon/src/table/postpone_bucket_planner.rs b/crates/paimon/src/table/postpone_bucket_planner.rs index 5efac0250..e259ff6b5 100644 --- a/crates/paimon/src/table/postpone_bucket_planner.rs +++ b/crates/paimon/src/table/postpone_bucket_planner.rs @@ -15,118 +15,421 @@ // specific language governing permissions and limitations // under the License. -//! Snapshot-bound sizing for fixed-bucket batch writes. +//! Pure input statistics and bucket inference for postpone batch writes. -use super::postpone_bucket_plan::data_invalid; -use crate::spec::{batch_to_serialized_bytes, BinaryRow, CoreOptions, DataField, DataType}; -use crate::table::{PostponeBucketPlan, Table, TableScan}; +use super::postpone_bucket_plan::{data_invalid, serialized_partition_columns}; +use crate::spec::{ + batch_to_serialized_bytes, BinaryRow, BucketFunctionType, CoreOptions, DataField, DataType, + TableSchema, POSTPONE_BUCKET, +}; +use crate::table::{PostponeBucketPlan, Table}; use crate::Result; use arrow_array::{ - Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, MapArray, RecordBatch, - StructArray, + Array, ArrayRef, FixedSizeListArray, Int64Array, LargeListArray, ListArray, MapArray, + RecordBatch, StructArray, }; -use arrow_schema::{DataType as ArrowType, Schema}; -use std::collections::HashMap; +use arrow_schema::{DataType as ArrowType, Field, Schema}; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; -pub(super) struct PostponeBucketPlanner { +pub const POSTPONE_STATISTICS_ROW_COUNT_FIELD: &str = "row_count"; +pub const POSTPONE_STATISTICS_DATA_SIZE_FIELD: &str = "data_size"; + +type PartitionStatistics = HashMap, (u128, u128)>; + +/// Owns schema and bucket counts without a table, runtime or retained input batches. +/// Workers can collect statistics independently and the caller can concatenate +/// them before planning. Existing partition layouts always take precedence. +pub struct PostponeBucketPlanner { fields: Vec, - partitions: Vec, - known: HashMap, i32>, - input: HashMap, (u128, u128)>, + partition_fields: Vec, + known: PostponeBucketPlan, + default_bucket_num: Option, target_rows: Option, target_size: u128, maximum: u128, - batches: Vec, } impl PostponeBucketPlanner { - pub(super) fn new(table: &Table) -> Result { - let options = CoreOptions::new(table.schema().options()); - let target_rows = options - .postpone_target_row_num_per_bucket()? - .map(|n| n as u128); - Ok(Self { - fields: table.schema().fields().to_vec(), - partitions: table - .schema() - .partition_keys() - .iter() - .map(|name| { - table - .schema() - .fields() - .iter() - .position(|f| f.name() == name) - .unwrap() - }) - .collect(), - known: HashMap::new(), - input: HashMap::new(), - target_rows, - // Java ignores byte sizing when a row target is supplied. - target_size: if target_rows.is_none() { + pub fn new(schema: &TableSchema, known: PostponeBucketPlan) -> Result { + let options = CoreOptions::new(schema.options()); + if options.bucket() != POSTPONE_BUCKET { + return Err(data_invalid( + "Postpone bucket planning requires bucket = -2", + )); + } + if options.bucket_function_type()? != BucketFunctionType::Default { + return Err(data_invalid( + "Postpone bucket planning only supports bucket-function.type=default", + )); + } + let partition_fields = schema.partition_fields(); + if partition_fields.len() != schema.partition_keys().len() { + return Err(data_invalid( + "Postpone partition fields are missing from the table schema", + )); + } + let mut names = HashSet::new(); + if schema + .fields() + .iter() + .any(|field| !names.insert(field.name())) + { + return Err(data_invalid( + "Postpone table schema contains duplicate field names", + )); + } + if known + .explicit_bucket_counts() + .values() + .any(|count| *count <= 0) + { + return Err(data_invalid( + "Postpone known bucket counts must be positive", + )); + } + let partitions = known + .explicit_bucket_counts() + .keys() + .map(Vec::as_slice) + .collect::>(); + serialized_partition_columns(&partition_fields, &partitions)?; + let default_bucket_num = options.postpone_default_bucket_num()?; + // Java and Python ignore sizing options when an exact default is set, + // and ignore the byte target when rows determine the count. + let (target_rows, target_size, maximum) = if default_bucket_num.is_some() { + (None, 1, 1) + } else { + let target_rows = options + .postpone_target_row_num_per_bucket()? + .map(|n| n as u128); + let target_size = if target_rows.is_none() { options.postpone_target_size_per_bucket()? as u128 } else { 1 - }, - maximum: options.postpone_batch_write_fixed_bucket_max_parallelism()? as u128, - batches: Vec::new(), + }; + ( + target_rows, + target_size, + options.postpone_batch_write_fixed_bucket_max_parallelism()? as u128, + ) + }; + Ok(Self { + fields: schema.fields().to_vec(), + partition_fields, + known, + default_bucket_num, + target_rows, + target_size, + maximum, }) } - pub(super) async fn load_metadata(&mut self, table: &Table, snapshot_id: i64) -> Result<()> { - if snapshot_id == 0 { - return Ok(()); + pub fn current_plan(&self) -> PostponeBucketPlan { + self.known.clone() + } + + /// Return partition columns followed positionally by Int64 `row_count` + /// and `data_size`. Partition names may equal either statistics name. + /// Only the supplied columns contribute to Java BinaryRow sizing. + pub fn input_statistics(&self, batch: &RecordBatch) -> Result { + let stats = self.collect_input(batch)?; + let mut entries = stats.into_iter().collect::>(); + entries.sort_unstable_by(|left, right| left.0.cmp(&right.0)); + let partitions = entries + .iter() + .map(|(partition, _)| partition.as_slice()) + .collect::>(); + let mut columns = serialized_partition_columns(&self.partition_fields, &partitions)?; + for component in [0, 1] { + let values = entries + .iter() + .map(|(_, (rows, bytes))| { + i64::try_from(if component == 0 { *rows } else { *bytes }).map_err(|_| { + data_invalid("Postpone input statistics exceed Int64 capacity") + }) + }) + .collect::>>()?; + columns.push(Arc::new(Int64Array::from(values))); } - let snapshot = table.snapshot_manager().get_snapshot(snapshot_id).await?; - let entries = TableScan::new(table, None, vec![], None, None, None) - .with_scan_all_files() - .plan_manifest_entries(&snapshot) - .await?; - for entry in entries { - let partition = entry.partition().to_vec(); - if entry.bucket() >= 0 { - let count = entry.total_buckets(); - if count <= 0 || entry.bucket() >= count { - return Err(data_invalid( - "Invalid total bucket count in existing postpone files", - )); - } - if self - .known - .insert(partition, count) - .is_some_and(|previous| previous != count) - { - return Err(data_invalid("Existing postpone files contain conflicting total bucket counts for one partition")); + RecordBatch::try_new(self.statistics_schema()?, columns) + .map_err(|error| data_invalid(error.to_string())) + } + + /// Aggregate worker statistics before inferring each new partition once. + /// A rejected batch leaves all existing and inferred counts unchanged. + pub fn plan(&mut self, statistics: &RecordBatch) -> Result { + let expected = self.statistics_schema()?; + if statistics.num_columns() != expected.fields().len() + || !expected + .fields() + .iter() + .zip(statistics.schema().fields()) + .all(|(expected, actual)| { + expected.name() == actual.name() && expected.data_type() == actual.data_type() + }) + { + return Err(data_invalid("Postpone statistics must contain the partition columns followed by Int64 row_count and data_size")); + } + self.validate_partition_nulls( + statistics, + &(0..self.partition_fields.len()).collect::>(), + )?; + let rows = statistics + .column(self.partition_fields.len()) + .as_any() + .downcast_ref::() + .unwrap(); + let bytes = statistics + .column(self.partition_fields.len() + 1) + .as_any() + .downcast_ref::() + .unwrap(); + let partitions = batch_to_serialized_bytes( + statistics, + &(0..self.partition_fields.len()).collect::>(), + &self.partition_fields, + )?; + let mut totals = PartitionStatistics::new(); + for (row, partition) in partitions.into_iter().enumerate() { + if rows.is_null(row) + || bytes.is_null(row) + || rows.value(row) < 0 + || bytes.value(row) < 0 + { + return Err(data_invalid(format!("Postpone row_count and data_size must be non-null and nonnegative at row {row}"))); + } + if self.known.total_buckets(&partition).is_none() { + add_statistics( + &mut totals, + partition, + rows.value(row) as u128, + bytes.value(row) as u128, + )?; + } + } + self.apply_statistics(totals); + Ok(self.current_plan()) + } + + fn statistics_schema(&self) -> Result> { + let mut fields = crate::arrow::build_target_arrow_schema(&self.partition_fields)? + .fields() + .to_vec(); + fields.extend([ + Arc::new(Field::new( + POSTPONE_STATISTICS_ROW_COUNT_FIELD, + ArrowType::Int64, + false, + )), + Arc::new(Field::new( + POSTPONE_STATISTICS_DATA_SIZE_FIELD, + ArrowType::Int64, + false, + )), + ]); + Ok(Arc::new(Schema::new(fields))) + } + + fn input_fields(&self, batch: &RecordBatch) -> Result> { + let mut names = HashSet::new(); + batch + .schema() + .fields() + .iter() + .map(|field| { + if !names.insert(field.name().clone()) { + return Err(data_invalid(format!( + "Duplicate postpone input column '{}'", + field.name() + ))); } + self.fields + .iter() + .find(|candidate| candidate.name() == field.name()) + .cloned() + .ok_or_else(|| { + data_invalid(format!("Unknown postpone input column '{}'", field.name())) + }) + }) + .collect() + } + + fn input_partitions(&self, batch: &RecordBatch, fields: &[DataField]) -> Result>> { + let indices = self + .partition_fields + .iter() + .map(|field| { + fields + .iter() + .position(|candidate| candidate.name() == field.name()) + .ok_or_else(|| { + data_invalid(format!( + "Missing postpone partition column '{}'", + field.name() + )) + }) + }) + .collect::>>()?; + self.validate_partition_nulls(batch, &indices)?; + // Writers accept both Arrow BINARY layouts. The BinaryRow encoder + // expects variable binary, while sizing still uses the original input. + let mut columns = batch.columns().to_vec(); + let mut arrow_fields = batch.schema().fields().to_vec(); + let mut normalized = false; + for &index in &indices { + if matches!(columns[index].data_type(), ArrowType::FixedSizeBinary(_)) { + let expected = crate::arrow::build_target_arrow_schema(&[fields[index].clone()])?; + columns[index] = super::write_batch_normalize::normalize_write_array( + &columns[index], + expected.field(0).data_type(), + )?; + arrow_fields[index] = Arc::new( + arrow_fields[index] + .as_ref() + .clone() + .with_data_type(columns[index].data_type().clone()), + ); + normalized = true; + } + } + if normalized { + let normalized = RecordBatch::try_new(Arc::new(Schema::new(arrow_fields)), columns) + .map_err(|error| data_invalid(error.to_string()))?; + batch_to_serialized_bytes(&normalized, &indices, fields) + } else { + batch_to_serialized_bytes(batch, &indices, fields) + } + } + + fn validate_partition_nulls(&self, batch: &RecordBatch, indices: &[usize]) -> Result<()> { + for (field, index) in self.partition_fields.iter().zip(indices) { + if !field.data_type().is_nullable() && batch.column(*index).null_count() != 0 { + return Err(data_invalid(format!( + "Postpone partition column '{}' is NOT NULL but contains null values", + field.name() + ))); } } Ok(()) } + fn collect_input(&self, batch: &RecordBatch) -> Result { + let fields = self.input_fields(batch)?; + let partitions = self.input_partitions(batch, &fields)?; + let mut stats = PartitionStatistics::new(); + for (row, partition) in partitions.into_iter().enumerate() { + if self.known.total_buckets(&partition).is_some() { + continue; + } + let size = if self.default_bucket_num.is_none() && self.target_rows.is_none() { + row_size(batch, row, &fields)? as u128 + } else { + 0 + }; + add_statistics(&mut stats, partition, 1, size)?; + } + Ok(stats) + } + + fn apply_statistics(&mut self, statistics: PartitionStatistics) { + for (partition, (rows, size)) in statistics { + if self.known.total_buckets(&partition).is_some() { + continue; + } + let count = self.default_bucket_num.unwrap_or_else(|| { + let required = match self.target_rows { + Some(target) => rows.div_ceil(target), + None => size.div_ceil(self.target_size), + } + .max(1); + // Capping before rounding avoids overflow while preserving + // Java's round-up-then-cap result, including non-power-of-two caps. + if required >= self.maximum { + self.maximum as i32 + } else { + required.next_power_of_two().min(self.maximum) as i32 + } + }); + self.known.insert(partition, count); + } + } +} + +fn add_statistics( + stats: &mut PartitionStatistics, + partition: Vec, + rows: u128, + size: u128, +) -> Result<()> { + let current = stats.entry(partition).or_default(); + current.0 = current + .0 + .checked_add(rows) + .ok_or_else(|| data_invalid("Postpone row count overflow"))?; + current.1 = current + .1 + .checked_add(size) + .ok_or_else(|| data_invalid("Postpone data size overflow"))?; + Ok(()) +} + +/// The automatic writer separately owns batches pending their final layout. +pub(super) struct BufferedPostponeBucketPlanner { + planner: PostponeBucketPlanner, + input: PartitionStatistics, + batches: Vec, +} + +impl BufferedPostponeBucketPlanner { + pub(super) fn new(table: &Table) -> Result { + Ok(Self { + planner: PostponeBucketPlanner::new(table.schema(), PostponeBucketPlan::default())?, + input: PartitionStatistics::new(), + batches: Vec::new(), + }) + } + + pub(super) async fn load_metadata(&mut self, table: &Table, snapshot_id: i64) -> Result<()> { + self.planner + .known + .load_existing_buckets(table, snapshot_id) + .await + } + pub(super) fn known_plan(&self) -> PostponeBucketPlan { - PostponeBucketPlan::from_counts(self.known.clone()) + self.planner.current_plan() } pub(super) fn add(&mut self, batch: RecordBatch) -> Result> { - let partitions = batch_to_serialized_bytes(&batch, &self.partitions, &self.fields)?; + // Normalization appends _VALUE_KIND to PK batches. Statistics describe + // the logical row while the buffered write retains the event column. + let logical_indices = batch + .schema() + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| { + (field.name() != crate::spec::VALUE_KIND_FIELD_NAME).then_some(index) + }) + .collect::>(); + let logical = batch + .project(&logical_indices) + .map_err(|error| data_invalid(error.to_string()))?; + let fields = self.planner.input_fields(&logical)?; + let partitions = self.planner.input_partitions(&logical, &fields)?; + let stats = self.planner.collect_input(&logical)?; + for (partition, (rows, size)) in stats { + add_statistics(&mut self.input, partition, rows, size)?; + } let mut known_rows = Vec::new(); let mut unknown_rows = Vec::new(); for (row, partition) in partitions.into_iter().enumerate() { - if self.known.contains_key(&partition) { + if self.planner.known.total_buckets(&partition).is_some() { known_rows.push(row); - continue; - } - unknown_rows.push(row); - let size = if self.target_rows.is_none() { - row_size(&batch, row, &self.fields)? } else { - 0 - }; - let stats = self.input.entry(partition).or_default(); - stats.0 += 1; - stats.1 += size as u128; + unknown_rows.push(row); + } } if unknown_rows.is_empty() { return Ok(Some(batch)); @@ -142,27 +445,23 @@ impl PostponeBucketPlanner { } pub(super) fn finish(mut self) -> (PostponeBucketPlan, Vec) { - for (partition, (rows, size)) in self.input { - let required = match self.target_rows { - Some(target) => rows.div_ceil(target), - None => size.div_ceil(self.target_size), - }; - // PostponeUtils rounds inferred counts up before applying the cap. - let count = required.max(1).next_power_of_two().min(self.maximum) as i32; - self.known.insert(partition, count); - } - (PostponeBucketPlan::from_counts(self.known), self.batches) + self.planner.apply_statistics(self.input); + (self.planner.current_plan(), self.batches) } } -fn round_word(size: usize) -> usize { - size.div_ceil(8) * 8 +fn checked_size(size: Option) -> Result { + size.ok_or_else(|| data_invalid("Postpone BinaryRow size overflow")) +} + +fn round_word(size: usize) -> Result { + Ok(checked_size(size.checked_add(7))? / 8 * 8) } fn row_size(batch: &RecordBatch, row: usize, fields: &[DataField]) -> Result { let mut size = BinaryRow::cal_fix_part_size_in_bytes(fields.len() as i32) as usize; for (column, field) in batch.columns().iter().zip(fields) { - size += variable_size(column, row, field.data_type())?; + size = checked_size(size.checked_add(variable_size(column, row, field.data_type())?))?; } Ok(size) } @@ -197,10 +496,13 @@ fn array_width(data_type: &DataType) -> usize { fn array_size(values: &ArrayRef, data_type: &DataType) -> Result { let count = values.len(); - let mut size = round_word(4 + count.div_ceil(32) * 4 + count * array_width(data_type)); + let nulls = checked_size(count.div_ceil(32).checked_mul(4))?; + let values_size = checked_size(count.checked_mul(array_width(data_type)))?; + let fixed = checked_size(checked_size(4usize.checked_add(nulls))?.checked_add(values_size))?; + let mut size = round_word(fixed)?; for row in 0..count { if !values.is_null(row) { - size += variable_size(values, row, data_type)?; + size = checked_size(size.checked_add(variable_size(values, row, data_type)?))?; } } Ok(size) @@ -228,7 +530,11 @@ fn variable_size(array: &ArrayRef, row: usize, data_type: &DataType) -> Result { let length = binary_length(array, row)?; - Ok(if length <= 7 { 0 } else { round_word(length) }) + if length <= 7 { + Ok(0) + } else { + round_word(length) + } } DataType::Row(row_type) => { let values = array @@ -237,10 +543,29 @@ fn variable_size(array: &ArrayRef, row: usize, data_type: &DataType) -> Result { let values = if let Some(values) = array.as_any().downcast_ref::() { @@ -254,7 +579,7 @@ fn variable_size(array: &ArrayRef, row: usize, data_type: &DataType) -> Result map_size(array, row, map.key_type(), map.value_type()), DataType::Multiset(multiset) => map_size( @@ -268,9 +593,13 @@ fn variable_size(array: &ArrayRef, row: usize, data_type: &DataType) -> Result() .ok_or_else(|| data_invalid("Expected VECTOR input while sizing postpone data"))?; - Ok(round_word( - 4 + values.value(row).len() * array_width(vector.element_type()), - )) + let payload = checked_size( + values + .value(row) + .len() + .checked_mul(array_width(vector.element_type())), + )?; + round_word(checked_size(4usize.checked_add(payload))?) } _ => { // Reuse the existing Java BinaryRow encoder for scalar layouts, @@ -311,19 +640,588 @@ fn map_size(array: &ArrayRef, row: usize, key: &DataType, value: &DataType) -> R .downcast_ref::() .ok_or_else(|| data_invalid("Expected MAP input while sizing postpone data"))? .value(row); - Ok(round_word( - 4 + array_size(entries.column(0), key)? + array_size(entries.column(1), value)?, - )) + let keys = array_size(entries.column(0), key)?; + let values = array_size(entries.column(1), value)?; + round_word(checked_size( + checked_size(4usize.checked_add(keys))?.checked_add(values), + )?) } #[cfg(test)] mod tests { use super::*; - use crate::spec::{BlobType, DecimalType}; + use crate::spec::{ + ArrayType, BinaryType, BlobType, DecimalType, FloatType, IntType, MapType, RowType, + TimestampType, VarCharType, VectorType, + }; use crate::table::table_write::tests::{ make_batch, setup_dirs, test_file_io, test_postpone_pk_table, }; - use arrow_array::{Decimal128Array, LargeBinaryArray}; + use arrow_array::builder::{Int32Builder, MapBuilder, StringBuilder}; + use arrow_array::{ + BinaryArray, Decimal128Array, FixedSizeBinaryArray, Float32Array, Int32Array, + LargeBinaryArray, StringArray, TimestampMicrosecondArray, + }; + + fn planner_schema(options: &[(&str, &str)]) -> TableSchema { + let mut builder = crate::spec::Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("pt", DataType::VarChar(VarCharType::string_type())) + .column("payload", DataType::VarChar(VarCharType::string_type())) + .primary_key(["pt", "id"]) + .partition_keys(["pt"]) + .option("primary-key.nullable", "true") + .option("bucket", "-2"); + for (key, value) in options { + builder = builder.option(*key, *value); + } + TableSchema::new(0, &builder.build().unwrap()) + } + + fn input(partitions: Vec>, payloads: Vec>) -> RecordBatch { + // Deliberately reorder and project out id: the full schema is not the input layout. + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("payload", ArrowType::Utf8, true), + Field::new("pt", ArrowType::Utf8, true), + ])), + vec![ + Arc::new(StringArray::from(payloads)), + Arc::new(StringArray::from(partitions)), + ], + ) + .unwrap() + } + + fn statistics( + schema: &TableSchema, + partitions: Vec>, + rows: Vec>, + sizes: Vec>, + ) -> RecordBatch { + let mut fields = crate::arrow::build_target_arrow_schema(&schema.partition_fields()) + .unwrap() + .fields() + .to_vec(); + fields.extend([ + Arc::new(Field::new( + POSTPONE_STATISTICS_ROW_COUNT_FIELD, + ArrowType::Int64, + true, + )), + Arc::new(Field::new( + POSTPONE_STATISTICS_DATA_SIZE_FIELD, + ArrowType::Int64, + true, + )), + ]); + RecordBatch::try_new( + Arc::new(Schema::new(fields)), + vec![ + Arc::new(StringArray::from(partitions)), + Arc::new(Int64Array::from(rows)), + Arc::new(Int64Array::from(sizes)), + ], + ) + .unwrap() + } + + fn counts(schema: &TableSchema, plan: &PostponeBucketPlan) -> HashMap, i32> { + let batch = plan.to_arrow(schema).unwrap(); + let partitions = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let counts = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|row| { + ( + if partitions.is_null(row) { + None + } else { + Some(partitions.value(row).to_string()) + }, + counts.value(row), + ) + }) + .collect() + } + + #[test] + fn input_statistics_size_only_actual_columns_and_leave_plan_unchanged() { + let schema = planner_schema(&[("postpone.target-size-per-bucket", "20 b")]); + let known = PostponeBucketPlan::from_arrow( + &schema, + &RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("pt", ArrowType::Utf8, true), + Field::new("total_buckets", ArrowType::Int32, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["old"])), + Arc::new(Int32Array::from(vec![3])), + ], + ) + .unwrap(), + ) + .unwrap(); + let mut planner = PostponeBucketPlanner::new(&schema, known).unwrap(); + let batch = input( + vec![Some("new"), Some("old"), Some("new"), None], + vec![ + Some("12345678"), + Some("ignored bytes"), + Some("1234567"), + None, + ], + ); + let stats = planner.input_statistics(&batch).unwrap(); + let pts = stats + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let rows = stats + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let bytes = stats + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let values = (0..stats.num_rows()) + .map(|row| { + ( + if pts.is_null(row) { + None + } else { + Some(pts.value(row).to_string()) + }, + (rows.value(row), bytes.value(row)), + ) + }) + .collect::>(); + // Two-field BinaryRows have 24 fixed bytes; eight-byte text needs another word. + assert_eq!( + values, + HashMap::from([(Some("new".into()), (2, 56)), (None, (1, 24))]) + ); + assert_eq!( + counts(&schema, &planner.current_plan()), + HashMap::from([(Some("old".into()), 3)]) + ); + assert_eq!( + counts(&schema, &planner.plan(&stats).unwrap()), + HashMap::from([(Some("old".into()), 3), (Some("new".into()), 4), (None, 2),]) + ); + assert_eq!(planner.input_statistics(&batch).unwrap().num_rows(), 0); + } + + #[test] + fn merged_worker_statistics_are_inferred_once_and_keep_known_counts() { + let schema = planner_schema(&[("postpone.target-row-num-per-bucket", "2")]); + let mut planner = + PostponeBucketPlanner::new(&schema, PostponeBucketPlan::default()).unwrap(); + let worker = planner + .input_statistics(&input( + vec![Some("p"), Some("p"), None], + vec![None, None, None], + )) + .unwrap(); + let merged = + arrow_select::concat::concat_batches(&worker.schema(), [&worker, &worker, &worker]) + .unwrap(); + let first = planner.plan(&merged).unwrap(); + assert_eq!( + counts(&schema, &first), + HashMap::from([(Some("p".into()), 4), (None, 2)]) + ); + let repeated = statistics( + &schema, + vec![Some("p"), Some("q")], + vec![Some(100), Some(0)], + vec![Some(0), Some(0)], + ); + assert_eq!( + counts(&schema, &planner.plan(&repeated).unwrap()), + HashMap::from([(Some("p".into()), 4), (None, 2), (Some("q".into()), 1),]) + ); + assert_eq!( + counts(&schema, &first).len(), + 2, + "previous plans own their counts" + ); + } + + #[test] + fn fixed_binary_partition_statistics_preserve_slices_nulls_and_routing() { + let schema = crate::spec::Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("pt", DataType::Binary(BinaryType::new(8).unwrap())) + .column("payload", DataType::VarChar(VarCharType::string_type())) + .primary_key(["pt", "id"]) + .partition_keys(["pt"]) + .option("primary-key.nullable", "true") + .option("bucket", "-2") + .option("postpone.target-size-per-bucket", "32 b") + .build() + .unwrap(); + let schema = TableSchema::new(0, &schema); + let partition = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [ + Some(b"discard!".as_slice()), + None, + Some(b"12345678".as_slice()), + Some(b"12345678".as_slice()), + Some(b"outside!".as_slice()), + ] + .into_iter(), + 8, + ) + .unwrap(); + let input = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("payload", ArrowType::Utf8, true), + Field::new("pt", ArrowType::FixedSizeBinary(8), true), + Field::new("id", ArrowType::Int32, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["123456789"; 5])), + Arc::new(partition), + Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4])), + ], + ) + .unwrap() + .slice(1, 3); + let canonical = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("payload", ArrowType::Utf8, true), + Field::new("pt", ArrowType::Binary, true), + Field::new("id", ArrowType::Int32, false), + ])), + vec![ + input.column(0).clone(), + Arc::new(BinaryArray::from(vec![ + None, + Some(b"12345678".as_slice()), + Some(b"12345678".as_slice()), + ])), + input.column(2).clone(), + ], + ) + .unwrap(); + let mut planner = + PostponeBucketPlanner::new(&schema, PostponeBucketPlan::default()).unwrap(); + let actual = planner.input_statistics(&input).unwrap(); + assert_eq!(actual, planner.input_statistics(&canonical).unwrap()); + let partitions = actual + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let sizes = actual + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let expected_sizes = (0..actual.num_rows()) + .map(|row| { + ( + if partitions.is_null(row) { + None + } else { + Some(partitions.value(row).to_vec()) + }, + sizes.value(row), + ) + }) + .collect::>(); + // Three logical columns: 32 fixed bytes plus 16 payload bytes; + // only non-NULL eight-byte partition values need an additional word. + assert_eq!( + expected_sizes, + HashMap::from([(None, 48), (Some(b"12345678".to_vec()), 112)]) + ); + let plan = planner.plan(&actual).unwrap(); + let exported = plan.to_arrow(&schema).unwrap(); + let counts = exported + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let pts = exported + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..exported.num_rows() { + assert_eq!(counts.value(row), if pts.is_null(row) { 2 } else { 4 }); + } + let router = crate::table::BucketRouter::new(&schema) + .unwrap() + .with_bucket_plan(plan) + .unwrap(); + assert_eq!( + router.route(&input).unwrap(), + router.route(&canonical).unwrap() + ); + } + + #[test] + fn inference_preserves_default_and_row_precedence_and_caps_large_aggregates() { + let schema = planner_schema(&[ + ("postpone.default-bucket-num", "7"), + ( + "postpone.batch-write-fixed-bucket.max-parallelism", + "invalid", + ), + ("postpone.target-row-num-per-bucket", "0"), + ("postpone.target-size-per-bucket", "invalid"), + ]); + let mut planner = + PostponeBucketPlanner::new(&schema, PostponeBucketPlan::default()).unwrap(); + let stats = planner + .input_statistics(&input(vec![Some("a")], vec![Some("long payload")])) + .unwrap(); + assert_eq!( + stats + .column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 0 + ); + assert_eq!( + counts(&schema, &planner.plan(&stats).unwrap()), + HashMap::from([(Some("a".into()), 7)]) + ); + + let schema = planner_schema(&[ + ("postpone.target-row-num-per-bucket", "2"), + ("postpone.target-size-per-bucket", "invalid"), + ("postpone.batch-write-fixed-bucket.max-parallelism", "3"), + ]); + let mut planner = + PostponeBucketPlanner::new(&schema, PostponeBucketPlan::default()).unwrap(); + let stats = statistics( + &schema, + vec![Some("a"), Some("b"), Some("large"), Some("large")], + vec![Some(2), Some(5), Some(i64::MAX), Some(i64::MAX)], + vec![Some(0); 4], + ); + assert_eq!( + counts(&schema, &planner.plan(&stats).unwrap()), + HashMap::from([ + (Some("a".into()), 1), + (Some("b".into()), 3), + (Some("large".into()), 3), + ]) + ); + assert!(round_word(usize::MAX).is_err()); + let mut totals = HashMap::from([(vec![], (u128::MAX, 0))]); + assert!(add_statistics(&mut totals, vec![], 1, 0).is_err()); + } + + #[test] + fn statistics_suffixes_are_positional_when_partition_names_collide() { + let schema = crate::spec::Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("row_count", DataType::Int(IntType::new())) + .column("data_size", DataType::Int(IntType::new())) + .primary_key(["id", "row_count", "data_size"]) + .partition_keys(["row_count", "data_size"]) + .option("bucket", "-2") + .option("postpone.target-row-num-per-bucket", "1") + .build() + .unwrap(); + let schema = TableSchema::new(0, &schema); + let mut planner = + PostponeBucketPlanner::new(&schema, PostponeBucketPlan::default()).unwrap(); + let input = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("data_size", ArrowType::Int32, false), + Field::new("row_count", ArrowType::Int32, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![20, 20, 20])), + Arc::new(Int32Array::from(vec![10, 10, 10])), + ], + ) + .unwrap(); + let stats = planner.input_statistics(&input).unwrap(); + assert_eq!( + stats + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + vec!["row_count", "data_size", "row_count", "data_size"] + ); + assert_eq!( + stats + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 10 + ); + let plan = planner.plan(&stats).unwrap().to_arrow(&schema).unwrap(); + assert_eq!( + plan.column(2) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 4 + ); + } + + #[test] + fn invalid_input_and_statistics_are_rejected_before_mutating_counts() { + let schema = planner_schema(&[("postpone.target-row-num-per-bucket", "1")]); + let mut planner = + PostponeBucketPlanner::new(&schema, PostponeBucketPlan::default()).unwrap(); + let batch = input(vec![Some("a")], vec![None]); + for names in [["unknown", "pt"], ["pt", "pt"], ["payload", "id"]] { + let invalid = RecordBatch::try_new( + Arc::new(Schema::new( + names + .map(|name| Field::new(name, ArrowType::Utf8, true)) + .to_vec(), + )), + batch.columns().to_vec(), + ) + .unwrap(); + assert!(planner.input_statistics(&invalid).is_err()); + } + for (rows, sizes) in [ + (vec![Some(1), Some(-1)], vec![Some(0), Some(0)]), + (vec![Some(1), None], vec![Some(0), Some(0)]), + (vec![Some(1), Some(1)], vec![Some(0), Some(-1)]), + (vec![Some(1), Some(1)], vec![Some(0), None]), + ] { + let stats = statistics(&schema, vec![Some("a"), Some("b")], rows, sizes); + assert!(planner.plan(&stats).is_err()); + assert_eq!( + planner.current_plan().to_arrow(&schema).unwrap().num_rows(), + 0 + ); + } + let conflict = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("pt", ArrowType::Utf8, true), + Field::new("total_buckets", ArrowType::Int32, false), + ])), + vec![ + Arc::new(StringArray::from(vec!["a", "a"])), + Arc::new(Int32Array::from(vec![2, 3])), + ], + ) + .unwrap(); + assert!(PostponeBucketPlan::from_arrow(&schema, &conflict).is_err()); + } + + #[test] + fn complex_input_size_matches_java_binary_layout() { + let string = DataType::VarChar(VarCharType::string_type()); + let integer = DataType::Int(IntType::new()); + let nested = DataType::Row(RowType::new(vec![ + DataField::new(20, "tiny".into(), integer.clone()), + DataField::new(21, "label".into(), string.clone()), + ])); + let schema = crate::spec::Schema::builder() + .column("id", integer.clone()) + .column("payload", string.clone()) + .column("blob", DataType::Blob(BlobType::new())) + .column( + "decimal", + DataType::Decimal(DecimalType::new(38, 0).unwrap()), + ) + .column("ts", DataType::Timestamp(TimestampType::new(6).unwrap())) + .column("items", DataType::Array(ArrayType::new(string.clone()))) + .column("pairs", DataType::Map(MapType::new(string, integer))) + .column("nested", nested) + .column( + "vector", + DataType::Vector(VectorType::new(3, DataType::Float(FloatType::new())).unwrap()), + ) + .primary_key(["id"]) + .option("bucket", "-2") + .build() + .unwrap(); + let schema = TableSchema::new(0, &schema); + let mut list = arrow_array::builder::ListBuilder::new(StringBuilder::new()); + list.values().append_value("1234567"); + list.values().append_value("12345678"); + list.values().append_null(); + list.append(true); + let mut map = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + map.keys().append_value("12345678"); + map.values().append_value(1); + map.append(true).unwrap(); + let nested = StructArray::from(vec![ + ( + Arc::new(Field::new("label", ArrowType::Utf8, true)), + Arc::new(StringArray::from(vec!["123456789"])) as ArrayRef, + ), + ( + Arc::new(Field::new("tiny", ArrowType::Int32, true)), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + ), + ]); + let vector = FixedSizeListArray::try_new( + Arc::new(Field::new("item", ArrowType::Float32, true)), + 3, + Arc::new(Float32Array::from(vec![1., 2., 3.])), + None, + ) + .unwrap(); + let columns: Vec = vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec!["12345678"])), + Arc::new(LargeBinaryArray::from(vec![Some(b"123456789".as_slice())])), + Arc::new( + Decimal128Array::from(vec![None]) + .with_precision_and_scale(38, 0) + .unwrap(), + ), + Arc::new(TimestampMicrosecondArray::from(vec![None])), + Arc::new(list.finish()), + Arc::new(map.finish()), + Arc::new(nested), + Arc::new(vector), + ]; + let fields = schema + .fields() + .iter() + .zip(&columns) + .map(|(field, column)| Field::new(field.name(), column.data_type().clone(), true)) + .collect::>(); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + let planner = PostponeBucketPlanner::new(&schema, PostponeBucketPlan::default()).unwrap(); + let statistics = planner.input_statistics(&batch).unwrap(); + // Fixed row: 80; string/blob: 8+16; NULL decimal/timestamp: 16+8; + // array/map/nested row/vector: 40+48+40+16. Same layout as Java/Python. + assert_eq!( + statistics + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 272 + ); + } #[test] fn binary_row_sizing_handles_blob_and_null_array_values() { @@ -350,6 +1248,38 @@ mod tests { assert_eq!(array_size(&values, &decimal).unwrap(), 32); } + #[test] + fn buffered_statistics_preserve_normalized_primary_key_events() { + let table = test_postpone_pk_table(&test_file_io(), "memory:/buffered_postpone_events") + .copy_with_options(HashMap::from([( + "postpone.target-row-num-per-bucket".into(), + "1".into(), + )])); + let write = table.new_write_builder().new_write().unwrap(); + let normalized = write + .normalize_write_batch(&make_batch(vec![1, 2], vec![10, 20])) + .unwrap() + .unwrap(); + let expected_kinds = normalized + .column_by_name(crate::spec::VALUE_KIND_FIELD_NAME) + .unwrap() + .clone(); + let mut planner = BufferedPostponeBucketPlanner::new(&table).unwrap(); + assert!(planner.add(normalized).unwrap().is_none()); + let (plan, batches) = planner.finish(); + assert_eq!( + plan.total_buckets(&crate::spec::EMPTY_SERIALIZED_ROW), + Some(2) + ); + assert_eq!(batches.len(), 1); + assert_eq!( + batches[0] + .column_by_name(crate::spec::VALUE_KIND_FIELD_NAME) + .unwrap(), + &expected_kinds + ); + } + #[tokio::test] async fn inference_excludes_old_pending_files_and_streams_known_partitions() { let io = test_file_io(); @@ -373,7 +1303,7 @@ mod tests { .commit(writer.prepare_commit().await.unwrap()) .await .unwrap(); - let mut planner = PostponeBucketPlanner::new(&table).unwrap(); + let mut planner = BufferedPostponeBucketPlanner::new(&table).unwrap(); planner.load_metadata(&table, 1).await.unwrap(); planner.add(make_batch(vec![5], vec![50])).unwrap(); let (plan, batches) = planner.finish(); @@ -382,8 +1312,9 @@ mod tests { Some(1) ); assert_eq!(batches.len(), 1); - let mut planner = PostponeBucketPlanner::new(&table).unwrap(); + let mut planner = BufferedPostponeBucketPlanner::new(&table).unwrap(); planner + .planner .known .insert(crate::spec::EMPTY_SERIALIZED_ROW.clone(), 3); let known = planner @@ -409,7 +1340,7 @@ mod tests { ), ]), ); - let mut planner = PostponeBucketPlanner::new(&table).unwrap(); + let mut planner = BufferedPostponeBucketPlanner::new(&table).unwrap(); planner .add(make_batch(vec![1, 2, 3], vec![10, 20, 30])) .unwrap(); @@ -422,7 +1353,7 @@ mod tests { ("postpone.target-row-num-per-bucket".into(), "2".into()), ("postpone.target-size-per-bucket".into(), "invalid".into()), ])); - let mut planner = PostponeBucketPlanner::new(&table).unwrap(); + let mut planner = BufferedPostponeBucketPlanner::new(&table).unwrap(); planner .add(make_batch(vec![1, 2, 3], vec![10, 20, 30])) .unwrap(); diff --git a/crates/paimon/src/table/postpone_fixed_bucket_router.rs b/crates/paimon/src/table/postpone_fixed_bucket_router.rs index fe116c7f1..debe0079d 100644 --- a/crates/paimon/src/table/postpone_fixed_bucket_router.rs +++ b/crates/paimon/src/table/postpone_fixed_bucket_router.rs @@ -15,14 +15,9 @@ // specific language governing permissions and limitations // under the License. -use super::postpone_bucket_plan::data_invalid; -use crate::spec::{ - batch_to_serialized_bytes, BucketFunctionType, CoreOptions, DataField, EMPTY_SERIALIZED_ROW, - POSTPONE_BUCKET, -}; -use crate::table::bucket_function::{batch_bucket_ids, validate_bucket_function}; +use crate::spec::{CoreOptions, TableSchema, POSTPONE_BUCKET}; use crate::table::table_write::take_rows; -use crate::table::{PostponeBucketPlan, Table}; +use crate::table::{BucketRouter, PostponeBucketPlan, Table}; use crate::Result; use arrow_array::RecordBatch; use std::collections::HashMap; @@ -35,37 +30,18 @@ pub(crate) struct PostponeBucketBatch { } pub(crate) struct PostponeFixedBucketRouter { - fields: Vec, - partition_field_indices: Vec, - bucket_key_indices: Vec, - bucket_function_type: BucketFunctionType, - plan: PostponeBucketPlan, + router: BucketRouter, } impl PostponeFixedBucketRouter { pub(super) fn set_plan(&mut self, plan: PostponeBucketPlan) { - self.plan = plan; + *self.router.plan_mut() = plan; } pub(crate) fn new(table: &Table, plan: PostponeBucketPlan) -> Result { validate_postpone_fixed_bucket_table(table)?; - let schema = table.schema(); - let options = CoreOptions::new(schema.options()); - let bucket_key_indices = field_indices(schema.fields(), &schema.bucket_keys()); - let bucket_key_fields = bucket_key_indices - .iter() - .map(|&index| schema.fields()[index].clone()) - .collect::>(); - let bucket_function_type = options.bucket_function_type()?; - if !bucket_key_fields.is_empty() { - validate_bucket_function(bucket_function_type, &bucket_key_fields)?; - } Ok(Self { - fields: schema.fields().to_vec(), - partition_field_indices: field_indices(schema.fields(), schema.partition_keys()), - bucket_key_indices, - bucket_function_type, - plan, + router: BucketRouter::new(table.schema())?.with_bucket_plan(plan)?, }) } @@ -74,56 +50,49 @@ impl PostponeFixedBucketRouter { table: &Table, snapshot_id: i64, ) -> Result<()> { - self.plan.load_existing_buckets(table, snapshot_id).await + self.router + .plan_mut() + .load_existing_buckets(table, snapshot_id) + .await } pub(crate) fn route(&self, batch: &RecordBatch) -> Result> { - let mut output = Vec::new(); - for (partition, batch) in - partition_batches(batch, &self.partition_field_indices, &self.fields)? - { - let total_buckets = self.plan.total_buckets(&partition).ok_or_else(|| { - data_invalid("Postpone bucket plan does not contain an input partition") - })?; - let buckets = if total_buckets <= 1 || self.bucket_key_indices.is_empty() { - vec![0; batch.num_rows()] - } else { - batch_bucket_ids( - &batch, - &self.bucket_key_indices, - &self.fields, - self.bucket_function_type, - total_buckets, - )? - }; - let mut groups: HashMap> = HashMap::new(); - for (row, bucket) in buckets.into_iter().enumerate() { - groups.entry(bucket).or_default().push(row); - } - for (bucket, rows) in groups { - output.push(PostponeBucketBatch { - partition: partition.clone(), - bucket, - batch: take_rows(&batch, &rows)?, - }); - } + let (partitions, buckets) = self.router.route_with_partitions(batch)?; + let mut groups: HashMap<(Vec, i32), Vec> = HashMap::new(); + for (row, partition) in partitions.into_iter().enumerate() { + groups + .entry((partition, buckets.value(row))) + .or_default() + .push(row); } - Ok(output) + groups + .into_iter() + .map(|((partition, bucket), rows)| { + Ok(PostponeBucketBatch { + partition, + bucket, + batch: take_rows(batch, &rows)?, + }) + }) + .collect() } pub(crate) fn total_buckets(&self, partition: &[u8]) -> Option { - self.plan.total_buckets(partition) + self.router.plan().total_buckets(partition) } } pub(crate) fn validate_postpone_fixed_bucket_table(table: &Table) -> Result<()> { - let schema = table.schema(); - let bucket = CoreOptions::new(schema.options()).bucket(); - if table.is_format_table() || bucket != POSTPONE_BUCKET || schema.primary_keys().is_empty() { + validate_postpone_fixed_bucket_schema(table.schema()) +} + +pub(super) fn validate_postpone_fixed_bucket_schema(schema: &TableSchema) -> Result<()> { + let options = CoreOptions::new(schema.options()); + let bucket = options.bucket(); + if options.is_format_table() || bucket != POSTPONE_BUCKET || schema.primary_keys().is_empty() { return Err(crate::Error::Unsupported { message: format!( - "Postpone fixed-bucket writes require a Paimon primary-key table with bucket=-2, but table '{}' has bucket={bucket}", - table.identifier().full_name() + "Postpone fixed-bucket writes require a Paimon primary-key table with bucket=-2, but schema has bucket={bucket}" ), }); } @@ -139,30 +108,3 @@ pub(crate) fn validate_postpone_fixed_bucket_table(table: &Table) -> Result<()> } Ok(()) } - -fn field_indices(fields: &[DataField], names: &[String]) -> Vec { - names - .iter() - .filter_map(|name| fields.iter().position(|field| field.name() == name)) - .collect() -} - -fn partition_batches( - batch: &RecordBatch, - partition_field_indices: &[usize], - fields: &[DataField], -) -> Result, RecordBatch)>> { - let partitions = if partition_field_indices.is_empty() { - vec![EMPTY_SERIALIZED_ROW.clone(); batch.num_rows()] - } else { - batch_to_serialized_bytes(batch, partition_field_indices, fields)? - }; - let mut groups: HashMap, Vec> = HashMap::new(); - for (row, partition) in partitions.into_iter().enumerate() { - groups.entry(partition).or_default().push(row); - } - groups - .into_iter() - .map(|(partition, rows)| Ok((partition, take_rows(batch, &rows)?))) - .collect() -} diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write.rs b/crates/paimon/src/table/postpone_fixed_bucket_write.rs index a46730c0f..28d958706 100644 --- a/crates/paimon/src/table/postpone_fixed_bucket_write.rs +++ b/crates/paimon/src/table/postpone_fixed_bucket_write.rs @@ -16,7 +16,7 @@ // under the License. use super::postpone_bucket_plan::data_invalid; -use super::postpone_bucket_planner::PostponeBucketPlanner; +use super::postpone_bucket_planner::BufferedPostponeBucketPlanner; use super::postpone_fixed_bucket_router::{ validate_postpone_fixed_bucket_table, PostponeFixedBucketRouter, }; @@ -35,7 +35,7 @@ pub struct PostponeFixedBucketTableWrite { overwrite: bool, check_from_snapshot: Option, prepare_started: bool, - planner: Option, + planner: Option, } impl PostponeFixedBucketTableWrite { @@ -80,7 +80,7 @@ impl PostponeFixedBucketTableWrite { ) -> Result { validate_postpone_fixed_bucket_table(table)?; let planner = if plan.is_none() { - Some(PostponeBucketPlanner::new(table)?) + Some(BufferedPostponeBucketPlanner::new(table)?) } else { None }; @@ -332,7 +332,7 @@ mod tests { total_buckets: Vec, ) -> PostponeBucketPlan { PostponeBucketPlan::from_arrow( - table, + table.schema(), &make_partition_bucket_plan_batch(partitions, total_buckets), ) .unwrap() @@ -348,7 +348,7 @@ mod tests { vec![Arc::new(Int32Array::from(vec![total_buckets]))], ) .unwrap(); - PostponeBucketPlan::from_arrow(table, &batch).unwrap() + PostponeBucketPlan::from_arrow(table.schema(), &batch).unwrap() } fn cross_partition_postpone_table(file_io: &FileIO, table_path: &str) -> Table { @@ -1000,7 +1000,7 @@ mod tests { let table = test_postpone_partitioned_table(&file_io, table_path); let error = PostponeBucketPlan::from_arrow( - &table, + table.schema(), &make_partition_bucket_plan_batch(vec!["p"], vec![0]), ) .unwrap_err(); diff --git a/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs index 21ebc7141..56d368cc2 100644 --- a/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs +++ b/crates/paimon/src/table/postpone_fixed_bucket_write_builder.rs @@ -145,7 +145,7 @@ mod tests { vec![Arc::new(Int32Array::from(vec![total_buckets]))], ) .unwrap(); - PostponeBucketPlan::from_arrow(table, &batch).unwrap() + PostponeBucketPlan::from_arrow(table.schema(), &batch).unwrap() } #[tokio::test] @@ -163,7 +163,7 @@ mod tests { .resolve_bucket_plan() .await .unwrap() - .to_arrow(&table) + .to_arrow(table.schema()) .unwrap() .num_rows(), 0 @@ -174,7 +174,7 @@ mod tests { .resolve_bucket_plan() .await .unwrap() - .to_arrow(&table) + .to_arrow(table.schema()) .unwrap(); assert_eq!(exported.num_rows(), 1); assert_eq!( @@ -194,7 +194,7 @@ mod tests { .resolve_bucket_plan() .await .unwrap() - .to_arrow(&table) + .to_arrow(table.schema()) .unwrap() .num_rows(), 0 @@ -225,7 +225,7 @@ mod tests { .resolve_bucket_plan() .await .unwrap() - .to_arrow(&table) + .to_arrow(table.schema()) .unwrap() .num_rows(), 0 @@ -254,7 +254,7 @@ mod tests { .resolve_bucket_plan() .await .unwrap() - .to_arrow(&historical) + .to_arrow(historical.schema()) .unwrap(); assert_eq!(exported.num_rows(), 1); assert_eq!( @@ -274,7 +274,7 @@ mod tests { .resolve_bucket_plan() .await .unwrap() - .to_arrow(&historical) + .to_arrow(historical.schema()) .unwrap() .num_rows(), 0 diff --git a/crates/paimon/tests/nullable_primary_key_test.rs b/crates/paimon/tests/nullable_primary_key_test.rs index c56e09885..a5342a8d2 100644 --- a/crates/paimon/tests/nullable_primary_key_test.rs +++ b/crates/paimon/tests/nullable_primary_key_test.rs @@ -17,16 +17,173 @@ mod common; -use arrow_array::{Array, ArrayRef, Int32Array, Int8Array, RecordBatch}; +use arrow_array::{ + new_null_array, Array, ArrayRef, Decimal128Array, Int32Array, Int8Array, RecordBatch, + StringArray, +}; use common::incremental_helpers::{memory_table, persist_table_schema, setup_dirs, write_batch}; use futures::TryStreamExt; -use paimon::spec::{DataType, IntType, Predicate, PredicateBuilder, Schema, TableSchema}; -use paimon::table::{MergeAssignment, MergeSource, Table, WhenMatched, WhenNotMatched}; +use paimon::spec::{ + DataType, DecimalType, IntType, LocalZonedTimestampType, Predicate, PredicateBuilder, Schema, + TableSchema, TimestampType, VarCharType, +}; +use paimon::table::{ + BucketRouter, MergeAssignment, MergeSource, Table, WhenMatched, WhenNotMatched, +}; use std::collections::HashMap; use std::sync::Arc; type Row = (Option, Option, Option); +#[tokio::test] +async fn null_non_compact_partitions_survive_exact_read_filters() { + let cases = [ + ( + "decimal18", + DataType::Decimal(DecimalType::new(18, 2).unwrap()), + ), + ( + "decimal19", + DataType::Decimal(DecimalType::new(19, 2).unwrap()), + ), + ( + "timestamp3", + DataType::Timestamp(TimestampType::new(3).unwrap()), + ), + ( + "timestamp4", + DataType::Timestamp(TimestampType::new(4).unwrap()), + ), + ( + "ltz3", + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(3).unwrap()), + ), + ( + "ltz4", + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(4).unwrap()), + ), + ]; + for (name, data_type) in cases { + let schema = Schema::builder() + .column("pt", data_type) + .column("code", DataType::VarChar(VarCharType::string_type())) + .column("id", DataType::Int(IntType::new())) + .partition_keys(["pt", "code"]) + .option("bucket", "1") + .option("bucket-key", "id") + .build() + .unwrap(); + let path = format!("memory:/nullable/partition-{name}"); + let (io, table) = memory_table(&path, TableSchema::new(0, &schema)); + setup_dirs(&io, &path).await; + persist_table_schema(&io, &path, table.schema()).await; + let arrow_schema = + paimon::arrow::build_target_arrow_schema(table.schema().fields()).unwrap(); + let batch = RecordBatch::try_new( + arrow_schema.clone(), + vec![ + new_null_array(arrow_schema.field(0).data_type(), 2), + Arc::new(StringArray::from(vec!["abcdefghijk", "other"])), + Arc::new(Int32Array::from(vec![1, 2])), + ], + ) + .unwrap(); + write_batch(&table, &batch).await; + let pb = PredicateBuilder::new(table.schema().fields()); + let mut builder = table.new_read_builder(); + builder.with_filter(Predicate::and(vec![ + pb.is_null("pt").unwrap(), + pb.equal("code", paimon::spec::Datum::String("abcdefghijk".into())) + .unwrap(), + ])); + let plan = builder.new_scan().plan().await.unwrap(); + assert_eq!(plan.splits().len(), 1, "{name}"); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let ids = batches + .iter() + .flat_map(|batch| { + batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!(ids, vec![1], "{name}"); + } +} + +#[tokio::test] +async fn nullable_decimal_key_uses_java_bucket_across_writer_restarts() { + let schema = Schema::builder() + .column("id", DataType::Decimal(DecimalType::new(20, 2).unwrap())) + .column("value", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("primary-key.nullable", "true") + .option("bucket", "100") + .build() + .unwrap(); + let path = "memory:/nullable/decimal-bucket"; + let (io, table) = memory_table(path, TableSchema::new(0, &schema)); + setup_dirs(&io, path).await; + persist_table_schema(&io, path, table.schema()).await; + let router = BucketRouter::new(table.schema()).unwrap(); + for value in [10, 20] { + let batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new( + Decimal128Array::from(vec![None]) + .with_precision_and_scale(20, 2) + .unwrap(), + ) as ArrayRef, + ), + ("value", Arc::new(Int32Array::from(vec![value])) as ArrayRef), + ]) + .unwrap(); + assert_eq!(router.route(&batch).unwrap().value(0), 60); + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer.write_arrow_batch(&batch).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].bucket, 60); + builder.new_commit().commit(messages).await.unwrap(); + } + let read = table.new_read_builder(); + let plan = read.new_scan().with_scan_all_files().plan().await.unwrap(); + let batches: Vec = read + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + let result = batches.iter().find(|batch| batch.num_rows() > 0).unwrap(); + assert!(result.column(0).is_null(0)); + assert_eq!( + result + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 20 + ); +} + async fn table(name: &str, options: &[(&str, &str)], primary_key: bool) -> Table { let mut schema = Schema::builder() // Engines may declare keys NOT NULL before applying the table option. From 1bf78423e22694258840d7b47b592f6084fe4245 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Fri, 23 Oct 2026 17:16:33 +0800 Subject: [PATCH 2/2] [core] Fix bucket plan binding and typed null partition identity --- bindings/c/src/write.rs | 2 +- crates/paimon/src/table/key_normalization.rs | 40 ++++--- crates/paimon/src/table/table_upsert.rs | 116 +++++++++++-------- crates/paimon/src/table/table_write.rs | 22 ++-- 4 files changed, 104 insertions(+), 76 deletions(-) diff --git a/bindings/c/src/write.rs b/bindings/c/src/write.rs index 3282e0d5e..dfb0dddbe 100644 --- a/bindings/c/src/write.rs +++ b/bindings/c/src/write.rs @@ -354,7 +354,7 @@ pub unsafe extern "C" fn paimon_postpone_fixed_bucket_write_builder_with_bucket_ Ok(batch) => batch, Err(error) => return error, }; - match PostponeBucketPlan::from_arrow(&state.table, &batch) { + match PostponeBucketPlan::from_arrow(state.table.schema(), &batch) { Ok(plan) => { state.bucket_plan = Some(plan); ptr::null_mut() diff --git a/crates/paimon/src/table/key_normalization.rs b/crates/paimon/src/table/key_normalization.rs index 8622f8688..27b85d94e 100644 --- a/crates/paimon/src/table/key_normalization.rs +++ b/crates/paimon/src/table/key_normalization.rs @@ -29,8 +29,8 @@ mod numeric; #[cfg(test)] mod encoded_tests; -// Java may reserve variable-length storage for NULL timestamps/decimals. -// Compare logical partitions using one encoding, independent of that padding. +// Re-encode logical partitions with the same typed NULL slots as Java and +// Arrow input serialization, independently of their original storage order. pub(super) fn canonical_partition_key( row: &BinaryRow, fields: &[DataField], @@ -39,7 +39,7 @@ pub(super) fn canonical_partition_key( for (index, field) in fields.iter().enumerate() { match extract_datum(row, index, field.data_type())? { Some(value) => builder.write_datum(index, &value, field.data_type()), - None => builder.set_null_at(index), + None => builder.write_null(index, field.data_type()), } } Ok(builder.build().to_serialized_bytes()) @@ -385,7 +385,7 @@ mod tests { } #[test] - fn partition_identity_ignores_java_null_variable_storage() { + fn partition_identity_reencodes_java_null_variable_storage() { let fields = vec![ DataField::new( 0, @@ -404,24 +404,26 @@ mod tests { ), ]; let mut builder = BinaryRowBuilder::new(3); - builder.set_null_at(0); - builder.set_null_at(1); + builder.write_null(0, fields[0].data_type()); + builder.write_null(1, fields[1].data_type()); builder.write_string(2, "after-null-fields"); let canonical = builder.build(); - let mut java = canonical.data()[..32].to_vec(); - // BinaryWriter.writeTimestamp(NULL, 6) reserves 8 bytes; - // writeDecimal(NULL, 20) reserves another 16, even though both are NULL. - java.extend_from_slice(&[0; 24]); - java.extend_from_slice(&canonical.data()[32..]); - java[8..16].copy_from_slice(&(32_u64 << 32).to_le_bytes()); - java[16..24].copy_from_slice(&(40_u64 << 32).to_le_bytes()); - let offset_and_size = u64::from_le_bytes(canonical.data()[24..32].try_into().unwrap()); - java[24..32].copy_from_slice(&(offset_and_size + (24_u64 << 32)).to_le_bytes()); - let java = BinaryRow::from_bytes(3, java); - assert_ne!(java.to_serialized_bytes(), canonical.to_serialized_bytes()); - assert_eq!( - canonical_partition_key(&java, &fields).unwrap(), + // Physical offsets can differ when variable fields are written in a + // different order. Both encodings retain Java's typed NULL storage. + let mut reordered = BinaryRowBuilder::new(3); + reordered.write_string(2, "after-null-fields"); + reordered.write_null(1, fields[1].data_type()); + reordered.write_null(0, fields[0].data_type()); + let reordered = reordered.build(); + assert_ne!( + reordered.to_serialized_bytes(), canonical.to_serialized_bytes() ); + for row in [&canonical, &reordered] { + assert_eq!( + canonical_partition_key(row, &fields).unwrap(), + canonical.to_serialized_bytes() + ); + } } } diff --git a/crates/paimon/src/table/table_upsert.rs b/crates/paimon/src/table/table_upsert.rs index c78abe560..7ccc886b9 100644 --- a/crates/paimon/src/table/table_upsert.rs +++ b/crates/paimon/src/table/table_upsert.rs @@ -1739,57 +1739,75 @@ mod tests { } #[tokio::test] - async fn upsert_finds_java_padded_null_timestamp_partition() { - let schema = PaimonSchema::builder() - .column( - "part", + async fn upsert_finds_java_null_non_compact_partitions() { + // Current Java BinaryRowWriter bodies for a single NULL DECIMAL(20,2) + // or TIMESTAMP/LTZ(6). The ordinary writer must produce these bytes; + // upsert must then match the same partition instead of appending a row. + let decimal = "0001000000000000000000001000000000000000000000000000000000000000"; + let timestamp = "000100000000000000000000100000000000000000000000"; + for (name, data_type, expected) in [ + ( + "decimal", + PaimonType::Decimal(crate::spec::DecimalType::new(20, 2).unwrap()), + decimal, + ), + ( + "timestamp", PaimonType::Timestamp(TimestampType::new(6).unwrap()), - ) - .column("key", PaimonType::Int(IntType::new())) - .column("value", PaimonType::Int(IntType::new())) - .partition_keys(["part"]) - .option("data-evolution.enabled", "true") - .option("row-tracking.enabled", "true") - .build() - .unwrap(); - let table = Table::new( - FileIOBuilder::new("memory").build().unwrap(), - Identifier::new("db", "padded_null_partition"), - "memory:/upsert_padded_null_partition".into(), - TableSchema::new(0, &schema), - None, - ); - let batch = |value| { - RecordBatch::try_from_iter([ - ( - "part", - Arc::new(arrow_array::TimestampMicrosecondArray::from(vec![ - None::, - ])) as ArrayRef, + timestamp, + ), + ( + "ltz", + PaimonType::LocalZonedTimestamp( + crate::spec::LocalZonedTimestampType::new(6).unwrap(), ), - ("key", Arc::new(Int32Array::from(vec![7])) as ArrayRef), - ("value", Arc::new(Int32Array::from(vec![value])) as ArrayRef), - ]) - .unwrap() - }; - let mut write = table.new_write_builder().new_write().unwrap(); - write.write_arrow_batch(&batch(10)).await.unwrap(); - let mut messages = write.prepare_commit().await.unwrap(); - for message in &mut messages { - let row = BinaryRow::from_serialized_bytes(&message.partition).unwrap(); - let mut padded = row.data().to_vec(); - // Java reserves a variable slot for NULL TIMESTAMP(6), although - // the null bit makes its offset irrelevant to the logical value. - padded.extend_from_slice(&[0; 8]); - padded[8..16].copy_from_slice(&(16_u64 << 32).to_le_bytes()); - message.partition = BinaryRow::from_bytes(1, padded).to_serialized_bytes(); + timestamp, + ), + ] { + let schema = PaimonSchema::builder() + .column("part", data_type) + .column("key", PaimonType::Int(IntType::new())) + .column("value", PaimonType::Int(IntType::new())) + .partition_keys(["part"]) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("db", name), + format!("memory:/upsert_null_partition_{name}"), + TableSchema::new(0, &schema), + None, + ); + let arrow_schema = + crate::arrow::build_target_arrow_schema(table.schema().fields()).unwrap(); + let batch = |value| { + RecordBatch::try_new( + arrow_schema.clone(), + vec![ + arrow_array::new_null_array(arrow_schema.field(0).data_type(), 1), + Arc::new(Int32Array::from(vec![7])), + Arc::new(Int32Array::from(vec![value])), + ], + ) + .unwrap() + }; + let mut write = table.new_write_builder().new_write().unwrap(); + write.write_arrow_batch(&batch(10)).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + assert_eq!(messages.len(), 1); + let partition = BinaryRow::from_serialized_bytes(&messages[0].partition).unwrap(); + assert_eq!(partition.data(), hex::decode(expected).unwrap(), "{name}"); + TableCommit::new(table.clone(), "seed".into()) + .commit(messages) + .await + .unwrap(); + write.close().await; + for value in [42, 43] { + commit_upsert(&table, vec![batch(value)], vec!["key".into()]).await; + assert_eq!(values_by_row_id(&table).await, vec![(0, value)], "{name}"); + } } - TableCommit::new(table.clone(), "seed".into()) - .commit(messages) - .await - .unwrap(); - write.close().await; - commit_upsert(&table, vec![batch(42)], vec!["key".into()]).await; - assert_eq!(values_by_row_id(&table).await, vec![(0, 42)]); } } diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 1e8c2217d..71e20b9af 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -4407,10 +4407,20 @@ pub(in crate::table) mod tests { async fn test_non_compact_null_bucket_key() { let file_io = test_file_io(); - let bucket_cols = ["d", "ltz", "ntz"]; + // Current Java InternalRowSerializer(type).toBinaryRow(GenericRow.of(null)): + // DECIMAL(38,18) reserves 16 variable bytes; TIMESTAMP/LTZ(6) reserve 8. + // Keep independent golden hashes and buckets instead of deriving the + // expectation from an untyped NULL, which omits those reserved slots. + let bucket_cases = [ + ("d", -1_380_280_760, 8), + ("ltz", -340_882_661, 5), + ("ntz", -340_882_661, 5), + ]; let total_buckets = 16; - for bucket_col in &bucket_cols { + for (column_index, &(bucket_col, java_hash, expected_bucket)) in + bucket_cases.iter().enumerate() + { let table_path = format!("memory:/test_null_bk_{bucket_col}"); setup_dirs(&file_io, &table_path).await; @@ -4423,7 +4433,7 @@ pub(in crate::table) mod tests { .column("ntz", DataType::Timestamp(TimestampType::new(6).unwrap())) .column("k", DataType::Int(IntType::new())) .option("bucket", total_buckets.to_string()) - .option("bucket-key", *bucket_col) + .option("bucket-key", bucket_col) .build() .unwrap(); let table_schema = TableSchema::new(0, &schema); @@ -4480,10 +4490,8 @@ pub(in crate::table) mod tests { .unwrap(); let bucket = batch_output.buckets[0]; - // Expected: BinaryRow with 1 field, null at pos 0 - let mut builder = BinaryRowBuilder::new(1); - builder.set_null_at(0); - let expected_bucket = (builder.build().hash_code() % total_buckets).wrapping_abs(); + let key = BinaryRow::from_arrow(&batch, 0, &[column_index], &fields).unwrap(); + assert_eq!(key.hash_code(), java_hash, "NULL bucket-key '{bucket_col}'"); assert_eq!( bucket, expected_bucket,