From 4c389878f6ca9151852b04093a2268f267d46391 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 3 Nov 2026 17:15:21 +0800 Subject: [PATCH 1/4] [datafusion] Expose range partitioning for table partition keys --- .../datafusion/src/physical_plan/mod.rs | 1 + .../src/physical_plan/range_partitioning.rs | 401 +++++++++++++++++ .../datafusion/src/physical_plan/scan.rs | 19 +- .../datafusion/tests/range_partitioning.rs | 421 ++++++++++++++++++ docs/src/sql.md | 23 +- 5 files changed, 860 insertions(+), 5 deletions(-) create mode 100644 crates/integrations/datafusion/src/physical_plan/range_partitioning.rs create mode 100644 crates/integrations/datafusion/tests/range_partitioning.rs diff --git a/crates/integrations/datafusion/src/physical_plan/mod.rs b/crates/integrations/datafusion/src/physical_plan/mod.rs index e0e3c8878..1ccab2349 100644 --- a/crates/integrations/datafusion/src/physical_plan/mod.rs +++ b/crates/integrations/datafusion/src/physical_plan/mod.rs @@ -16,6 +16,7 @@ // under the License. mod audit_log; +mod range_partitioning; pub(crate) mod scan; mod search_score; pub(crate) mod sink; diff --git a/crates/integrations/datafusion/src/physical_plan/range_partitioning.rs b/crates/integrations/datafusion/src/physical_plan/range_partitioning.rs new file mode 100644 index 000000000..47008d409 --- /dev/null +++ b/crates/integrations/datafusion/src/physical_plan/range_partitioning.rs @@ -0,0 +1,401 @@ +// 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. + +use std::sync::Arc; + +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::{DataType, Schema, TimeUnit}; +use datafusion::arrow::row::{RowConverter, SortField}; +use datafusion::common::{ScalarValue, SplitPoint}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning}; +use paimon::spec::{BinaryRow, DataField}; +use paimon::table::Table; +use paimon::DataSplit; + +use super::scan::datum_to_scalar; + +/// Group whole splits by their typed partition keys, then combine adjacent key +/// groups into at most the original number of execution partitions. Equal keys +/// must never straddle a boundary: DataFusion can use this declaration to omit +/// the exchange before a final aggregate or a co-partitioned join. +/// +/// This describes distribution only. Files and rows within a range need not be +/// sorted. Failure to prove any part of the contract leaves the original scan +/// assignments and UnknownPartitioning intact. +pub(super) fn range_partitioned_splits( + table: &Table, + schema: &Schema, + read_type: &[DataField], + partitions: &[Arc<[DataSplit]>], +) -> Option<(Vec>, Partitioning)> { + let fields = table.schema().partition_fields(); + if fields.is_empty() + || fields.len() != table.schema().partition_keys().len() + || partitions.len() < 2 + || read_type.len() != schema.fields().len() + || table.schema().core_options().query_auth_enabled() + { + return None; + } + + // Use the actual projected column positions. Variant extraction and omitted + // partition keys must not be mistaken for an unchanged partition column. + let indices = fields + .iter() + .map(|field| { + let index = read_type.iter().position(|read| read == field)?; + let output = schema.field(index); + (output.name() == field.name() && supported_type(output.data_type())).then_some(index) + }) + .collect::>>()?; + let options = SortOptions::new(false, false); // ASC NULLS LAST + let ordering = LexOrdering::new(indices.iter().map(|&index| { + PhysicalSortExpr::new( + Arc::new(Column::new(schema.field(index).name(), index)), + options, + ) + }))?; + + let splits = partitions + .iter() + .flat_map(|splits| splits.iter()) + .collect::>(); + let keys = splits + .iter() + .map(|split| { + let row = split.partition(); + if split.has_query_auth_rules() + || row.arity() != fields.len() as i32 + || row.data().len() < BinaryRow::cal_fix_part_size_in_bytes(row.arity()) as usize + { + return None; + } + fields + .iter() + .zip(&indices) + .enumerate() + .map(|(pos, (field, &index))| { + let data_type = schema.field(index).data_type(); + match row.get_datum(pos, field.data_type()).ok()? { + Some(value) => datum_to_scalar(value, data_type), + None => ScalarValue::try_from(data_type).ok(), + } + }) + .collect::>>() + }) + .collect::>>()?; + if keys.is_empty() { + return None; + } + + // Arrow's row encoding gives exactly the comparison semantics used by the + // engine, including typed numbers, UTF-8 strings, compound keys and NULLs. + let converter = RowConverter::new( + indices + .iter() + .map(|&index| { + SortField::new_with_options(schema.field(index).data_type().clone(), options) + }) + .collect(), + ) + .ok()?; + let columns = (0..fields.len()) + .map(|index| ScalarValue::iter_to_array(keys.iter().map(|key| key[index].clone())).ok()) + .collect::>>()?; + let rows = converter.convert_columns(&columns).ok()?; + let mut sorted = (0..splits.len()).collect::>(); + sorted.sort_by(|&left, &right| rows.row(left).cmp(&rows.row(right))); + let distinct = 1 + sorted + .windows(2) + .filter(|pair| rows.row(pair[0]) != rows.row(pair[1])) + .count(); + let count = distinct.min(partitions.len()); + if count < 2 { + return None; + } + + let mut grouped = vec![Vec::new(); count]; + let mut split_points = Vec::with_capacity(count - 1); + let mut key_index = 0; + for (position, &index) in sorted.iter().enumerate() { + if position > 0 && rows.row(sorted[position - 1]) != rows.row(index) { + key_index += 1; + } + let partition = key_index * count / distinct; + if partition > 0 && grouped[partition].is_empty() { + // The first key of each upper partition is the inclusive boundary. + split_points.push(SplitPoint::new(keys[index].clone())); + } + grouped[partition].push(splits[index].clone()); + } + let range = RangePartitioning::try_new(ordering, split_points).ok()?; + Some(( + grouped.into_iter().map(Arc::from).collect(), + Partitioning::Range(range), + )) +} + +fn supported_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Utf8 + | DataType::Utf8View + | DataType::LargeUtf8 + | DataType::Date32 + | DataType::Time32(TimeUnit::Millisecond) + | DataType::Decimal128(_, _) + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use paimon::catalog::Identifier; + use paimon::io::FileIOBuilder; + use paimon::spec::{ + DataType as PaimonType, Datum, IntType, Schema as PaimonSchema, TableSchema, TimeType, + }; + use paimon::DataSplitBuilder; + + fn table(data_type: PaimonType, query_auth: bool) -> Table { + let schema = PaimonSchema::builder() + .column("pt", data_type) + .partition_keys(["pt"]) + .option("query-auth.enabled", query_auth.to_string()) + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("file").build().unwrap(), + Identifier::new("test", "ranges"), + "file:/tmp/ranges".to_string(), + TableSchema::new(0, &schema), + None, + ) + } + + fn split(row: BinaryRow, bucket: i32) -> DataSplit { + DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(row) + .with_bucket(bucket) + .with_bucket_path("file:/tmp/ranges/bucket".to_string()) + .with_data_files(vec![]) + .build() + .unwrap() + } + + fn partitions(data_type: &PaimonType, values: &[Option]) -> Vec> { + values + .iter() + .enumerate() + .map(|(i, value)| { + Arc::from(vec![split( + BinaryRow::from_datums(&[(value.as_ref(), data_type)]), + i as i32, + )]) + }) + .collect() + } + + #[test] + fn time_partitions_follow_millisecond_order_including_null() { + // TIME cannot currently be created through SQLContext's DDL parser. + let data_type = PaimonType::Time(TimeType::new(3).unwrap()); + let table = table(data_type.clone(), false); + let schema = crate::table::datafusion_arrow_schema(table.schema().fields(), true).unwrap(); + let partitions = partitions( + &data_type, + &[Some(Datum::Time(86_399_999)), None, Some(Datum::Time(1))], + ); + let (grouped, partitioning) = + range_partitioned_splits(&table, &schema, table.schema().fields(), &partitions) + .unwrap(); + assert_eq!( + grouped + .iter() + .map(|group| group[0].bucket()) + .collect::>(), + vec![2, 0, 1] + ); + let Partitioning::Range(range) = partitioning else { + panic!("expected ranges") + }; + assert_eq!( + range.split_points()[0].values(), + &[ScalarValue::Time32Millisecond(Some(86_399_999))] + ); + assert_eq!( + range.split_points()[1].values(), + &[ScalarValue::Time32Millisecond(None)] + ); + } + + #[test] + fn missing_or_malformed_metadata_never_declares_ranges() { + let data_type = PaimonType::Int(IntType::new()); + let table = table(data_type.clone(), false); + let schema = crate::table::datafusion_arrow_schema(table.schema().fields(), true).unwrap(); + for row in [ + BinaryRow::new(0), + BinaryRow::new(1), + BinaryRow::from_bytes(1, vec![0; 9]), + ] { + let partitions = vec![ + Arc::from(vec![split(row, 0)]), + Arc::from(vec![split( + BinaryRow::from_datums(&[(Some(&Datum::Int(2)), &data_type)]), + 1, + )]), + ]; + assert!(range_partitioned_splits( + &table, + &schema, + table.schema().fields(), + &partitions + ) + .is_none()); + } + } + + #[test] + fn uneven_key_groups_keep_every_split_in_exactly_one_range() { + let data_type = PaimonType::Int(IntType::new()); + let table = table(data_type.clone(), false); + let schema = crate::table::datafusion_arrow_schema(table.schema().fields(), true).unwrap(); + let splits = partitions( + &data_type, + &[5, 1, 4, 2, 3, 2, 4].map(|value| Some(Datum::Int(value))), + ); + let partitions = crate::table::bucket_round_robin( + splits + .iter() + .flat_map(|group| group.iter().cloned()) + .collect(), + 3, + ) + .into_iter() + .map(Arc::from) + .collect::>(); + let (grouped, partitioning) = + range_partitioned_splits(&table, &schema, table.schema().fields(), &partitions) + .unwrap(); + let buckets = grouped + .iter() + .map(|group| { + let mut buckets = group.iter().map(DataSplit::bucket).collect::>(); + buckets.sort(); + buckets + }) + .collect::>(); + assert_eq!(buckets, vec![vec![1, 3, 5], vec![2, 4, 6], vec![0]]); + let Partitioning::Range(range) = partitioning else { + panic!("expected ranges") + }; + assert_eq!( + range + .split_points() + .iter() + .map(|point| point.values().to_vec()) + .collect::>(), + vec![ + vec![ScalarValue::Int32(Some(3))], + vec![ScalarValue::Int32(Some(5))] + ] + ); + } + + #[test] + fn binary_keys_do_not_infer_ranges_from_signed_byte_order() { + let data_type = PaimonType::VarBinary(paimon::spec::VarBinaryType::new(8).unwrap()); + let table = table(data_type.clone(), false); + let schema = crate::table::datafusion_arrow_schema(table.schema().fields(), true).unwrap(); + let partitions = partitions( + &data_type, + &[ + Some(Datum::Bytes(vec![0x80])), + Some(Datum::Bytes(vec![0x7f])), + ], + ); + assert!( + range_partitioned_splits(&table, &schema, table.schema().fields(), &partitions) + .is_none() + ); + } + + #[test] + fn authorization_and_single_key_groups_keep_original_assignments() { + let data_type = PaimonType::Int(IntType::new()); + let unrestricted = table(data_type.clone(), false); + let restricted = table(data_type.clone(), true); + let schema = + crate::table::datafusion_arrow_schema(unrestricted.schema().fields(), true).unwrap(); + let mixed = partitions(&data_type, &[Some(Datum::Int(1)), Some(Datum::Int(2))]); + assert!(range_partitioned_splits( + &restricted, + &schema, + restricted.schema().fields(), + &mixed + ) + .is_none()); + // Transported rules must also disable range claims when the table + // handle itself has no query-auth option. + let mut serialized = serde_json::to_value(&mixed[0][0]).unwrap(); + serialized["query_auth_response"] = serde_json::json!({"columnMasking": {"pt": "{}"}}); + let masked: DataSplit = serde_json::from_value(serialized).unwrap(); + assert!(masked.has_query_auth_rules()); + let masked = vec![Arc::from(vec![masked]), Arc::clone(&mixed[1])]; + assert!(range_partitioned_splits( + &unrestricted, + &schema, + unrestricted.schema().fields(), + &masked + ) + .is_none()); + for values in [ + vec![Some(Datum::Int(1)), Some(Datum::Int(1))], + vec![None, None], + ] { + let partitions = partitions(&data_type, &values); + assert!(range_partitioned_splits( + &unrestricted, + &schema, + unrestricted.schema().fields(), + &partitions + ) + .is_none()); + } + let single = vec![Arc::from( + mixed + .iter() + .flat_map(|group| group.iter().cloned()) + .collect::>(), + )]; + assert!(range_partitioned_splits( + &unrestricted, + &schema, + unrestricted.schema().fields(), + &single + ) + .is_none()); + } +} diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 80e59aa57..2ee74f597 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -939,9 +939,23 @@ impl PaimonTableScan { case_sensitive: bool, parquet_read_budget: Arc, ) -> Self { + let partition_count = planned_partitions.len(); + let (planned_partitions, partitioning) = + super::range_partitioning::range_partitioned_splits( + &table, + &schema, + &read_type, + &planned_partitions, + ) + .unwrap_or_else(|| { + ( + planned_partitions, + Partitioning::UnknownPartitioning(partition_count), + ) + }); let plan_properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(schema.clone()), - Partitioning::UnknownPartitioning(planned_partitions.len()), + partitioning, EmissionType::Incremental, Boundedness::Bounded, )); @@ -1350,6 +1364,9 @@ impl DisplayAs for PaimonTableScan { ", partitions={}, splits={total_splits}, files={total_files}", self.planned_partitions.len() )?; + if let Partitioning::Range(range) = self.plan_properties.output_partitioning() { + write!(f, ", partitioning={range}")?; + } } let columns = self diff --git a/crates/integrations/datafusion/tests/range_partitioning.rs b/crates/integrations/datafusion/tests/range_partitioning.rs new file mode 100644 index 000000000..0ec13a6d2 --- /dev/null +++ b/crates/integrations/datafusion/tests/range_partitioning.rs @@ -0,0 +1,421 @@ +// 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. + +mod common; + +use std::sync::Arc; + +use arrow_array::{Array, Int32Array}; +use datafusion::common::ScalarValue; +use datafusion::physical_plan::{ + collect_partitioned, displayable, ExecutionPlan, ExecutionPlanProperties, Partitioning, +}; +use paimon_datafusion::SQLContext; + +fn scan(plan: &Arc) -> Arc { + fn find(plan: &Arc) -> Option> { + if matches!(plan.name(), "PaimonTableScan" | "PaimonAuditLogScan") { + return Some(Arc::clone(plan)); + } + plan.children().into_iter().find_map(find) + } + find(plan).expect("Paimon scan") +} + +async fn plan(ctx: &SQLContext, query: &str) -> Arc { + ctx.sql(query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap() +} + +#[tokio::test] +async fn integer_ranges_keep_equal_keys_together_and_avoid_aggregation_shuffle() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + common::exec( + &ctx, + "CREATE TABLE paimon.test_db.ranges (id INT, pt INT, value INT) PARTITIONED BY (pt)", + ) + .await; + // Two commits put equal keys in distinct files/splits. Numeric order differs + // from directory-name order, and NULL must occupy the final range. + for values in [ + "(1, 10, 6), (2, -10, 2), (3, 2, 4), (4, NULL, 5)", + "(5, 2, 40), (6, 10, 60), (7, -10, 20), (8, NULL, 50)", + ] { + common::exec( + &ctx, + &format!("INSERT INTO paimon.test_db.ranges VALUES {values}"), + ) + .await; + } + + let scan = scan(&plan(&ctx, "SELECT id, pt, value FROM paimon.test_db.ranges").await); + let Partitioning::Range(range) = scan.output_partitioning() else { + panic!( + "expected range partitioning, got {:?}", + scan.output_partitioning() + ); + }; + assert_eq!(range.partition_count(), 2); + assert_eq!(range.ordering()[0].expr.to_string(), "pt@1"); + assert_eq!( + range.split_points()[0].values(), + &[ScalarValue::Int32(Some(10))] + ); + assert!( + scan.output_ordering().is_none(), + "ranges do not imply sorted rows" + ); + let partitions = collect_partitioned(scan, ctx.ctx().task_ctx()) + .await + .unwrap(); + let mut actual = Vec::new(); + for batches in partitions { + let mut keys = Vec::new(); + for batch in batches { + let pt = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + keys.extend((0..pt.len()).map(|i| (!pt.is_null(i)).then(|| pt.value(i)))); + } + keys.sort_by_key(|key| (key.is_none(), *key)); + actual.push(keys); + } + assert_eq!( + actual, + vec![ + vec![Some(-10), Some(-10), Some(2), Some(2)], + vec![Some(10), Some(10), None, None] + ] + ); + + let query = "SELECT CAST(coalesce(pt, -100) AS INT) AS id, CAST(sum(value) AS INT) AS value FROM paimon.test_db.ranges GROUP BY pt"; + let physical = plan(&ctx, query).await; + let text = displayable(physical.as_ref()).indent(true).to_string(); + assert!(!text.contains("RepartitionExec"), "{text}"); + assert_eq!( + common::collect_id_value(&ctx, query).await, + vec![(-100, 55), (-10, 22), (2, 44), (10, 66)] + ); + for (predicate, expected) in [ + ("pt < 10", vec![(2, 2), (3, 4), (5, 40), (7, 20)]), + ("pt >= 10", vec![(1, 6), (6, 60)]), + ("pt IS NULL", vec![(4, 5), (8, 50)]), + ] { + assert_eq!( + common::collect_id_value( + &ctx, + &format!("SELECT id, value FROM paimon.test_db.ranges WHERE {predicate}") + ) + .await, + expected + ); + } + // Local sorting is still required, even though the partitions have ranges. + let sorted = plan( + &ctx, + "SELECT id, pt FROM paimon.test_db.ranges ORDER BY pt, id", + ) + .await; + assert!(displayable(sorted.as_ref()) + .indent(true) + .to_string() + .contains("SortExec")); +} + +#[tokio::test] +async fn joins_align_different_ranges_after_primary_key_updates() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + common::exec(&ctx, "CREATE TABLE paimon.test_db.pk_ranges (id INT NOT NULL, pt INT NOT NULL, value INT, PRIMARY KEY (pt, id)) PARTITIONED BY (pt) WITH ('bucket' = '2')").await; + common::exec(&ctx, "INSERT INTO paimon.test_db.pk_ranges VALUES (1, 1, 10), (2, 2, 20), (3, 3, 30), (4, 4, 40)").await; + common::exec( + &ctx, + "INSERT INTO paimon.test_db.pk_ranges VALUES (2, 2, 200)", + ) + .await; + common::exec( + &ctx, + "CREATE TABLE paimon.test_db.other_ranges (pt INT, value INT) PARTITIONED BY (pt)", + ) + .await; + common::exec(&ctx, "INSERT INTO paimon.test_db.other_ranges VALUES (1, 100), (2, 200), (3, 300), (4, 400), (5, 500), (6, 600)").await; + + let left = scan(&plan(&ctx, "SELECT * FROM paimon.test_db.pk_ranges").await); + let right = scan(&plan(&ctx, "SELECT * FROM paimon.test_db.other_ranges").await); + let (Partitioning::Range(left), Partitioning::Range(right)) = + (left.output_partitioning(), right.output_partitioning()) + else { + panic!("both inputs must have ranges"); + }; + assert_eq!( + left.split_points()[0].values(), + &[ScalarValue::Int32(Some(3))] + ); + assert_eq!( + right.split_points()[0].values(), + &[ScalarValue::Int32(Some(4))] + ); + for prefer_hash in [true, false] { + common::exec( + &ctx, + &format!("SET datafusion.optimizer.prefer_hash_join = {prefer_hash}"), + ) + .await; + common::exec( + &ctx, + "SET datafusion.optimizer.hash_join_single_partition_threshold = 0", + ) + .await; + common::exec( + &ctx, + "SET datafusion.optimizer.hash_join_single_partition_threshold_rows = 0", + ) + .await; + let query = "SELECT l.id, l.value + r.value AS value FROM paimon.test_db.pk_ranges l JOIN paimon.test_db.other_ranges r ON l.pt = r.pt"; + let physical = plan(&ctx, query).await; + let text = displayable(physical.as_ref()).indent(true).to_string(); + let join = if prefer_hash { + "HashJoinExec" + } else { + "SortMergeJoinExec" + }; + assert!(text.contains(join), "{text}"); + assert!( + text.contains("RepartitionExec"), + "different boundaries need alignment: {text}" + ); + assert_eq!( + common::collect_id_value(&ctx, query).await, + vec![(1, 110), (2, 400), (3, 330), (4, 440)] + ); + } +} + +#[tokio::test] +async fn compound_date_and_string_keys_preserve_lexicographic_ranges() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 8").await; + common::exec(&ctx, "CREATE TABLE paimon.test_db.compound (id INT, dt DATE, region STRING) PARTITIONED BY (dt, region)").await; + common::exec(&ctx, "INSERT INTO paimon.test_db.compound VALUES (4, DATE '2026-01-01', 'δΈ­'), (2, DATE '1969-12-31', 'z'), (1, DATE '1969-12-31', 'a'), (3, DATE '2026-01-01', NULL)").await; + let physical = plan(&ctx, "SELECT region, id, dt FROM paimon.test_db.compound").await; + let scan = scan(&physical); + let Partitioning::Range(range) = scan.output_partitioning() else { + panic!("expected compound ranges: {:?}", scan.output_partitioning()); + }; + assert_eq!(range.partition_count(), 4); + assert_eq!(range.ordering()[0].expr.to_string(), "dt@1"); + assert_eq!(range.ordering()[1].expr.to_string(), "region@2"); + assert_eq!( + range.split_points()[0].values(), + &[ + ScalarValue::Date32(Some(-1)), + ScalarValue::Utf8View(Some("z".into())) + ] + ); + let Partitioning::Range(projected) = physical.output_partitioning() else { + panic!("projection must preserve ranges"); + }; + assert_eq!(projected.ordering()[0].expr.to_string(), "dt@2"); + assert_eq!(projected.ordering()[1].expr.to_string(), "region@0"); + let partitions = collect_partitioned(physical, ctx.ctx().task_ctx()) + .await + .unwrap(); + let ids: Vec> = partitions + .into_iter() + .map(|batches| { + batches + .into_iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect() + }) + .collect(); + assert_eq!(ids, vec![vec![1], vec![2], vec![4], vec![3]]); + let query = "SELECT sum(id) FROM paimon.test_db.compound GROUP BY dt, region"; + let physical = plan(&ctx, query).await; + let text = displayable(physical.as_ref()).indent(true).to_string(); + assert!(!text.contains("RepartitionExec"), "{text}"); + assert_eq!(common::row_count(&ctx, query).await, 4); +} + +#[tokio::test] +async fn supported_scalar_partition_types_use_typed_order_and_nulls_last() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 8").await; + for (name, data_type, low, high) in [ + ("bools", "BOOLEAN", "false", "true"), + ("tiny", "TINYINT", "-10", "2"), + ("small", "SMALLINT", "-10", "2"), + ("big", "BIGINT", "-100000000000", "200000000000"), + ( + "decimal", + "DECIMAL(20,3)", + "-12345678901234567.890", + "2.001", + ), + ] { + common::exec( + &ctx, + &format!( + "CREATE TABLE paimon.test_db.{name} (id INT, pt {data_type}) PARTITIONED BY (pt)" + ), + ) + .await; + common::exec( + &ctx, + &format!("INSERT INTO paimon.test_db.{name} VALUES (2, {high}), (3, NULL), (1, {low})"), + ) + .await; + let scan = scan(&plan(&ctx, &format!("SELECT id, pt FROM paimon.test_db.{name}")).await); + assert!( + matches!(scan.output_partitioning(), Partitioning::Range(_)), + "{name}: {:?}", + scan.output_partitioning() + ); + let partitions = collect_partitioned(scan, ctx.ctx().task_ctx()) + .await + .unwrap(); + let ids: Vec> = partitions + .into_iter() + .map(|batches| { + batches + .into_iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect() + }) + .collect(); + assert_eq!(ids, vec![vec![1], vec![2], vec![3]], "{name}"); + } +} + +#[tokio::test] +async fn variant_pushdown_and_audit_scans_preserve_partition_ranges() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + common::exec(&ctx, r#"CREATE TABLE paimon.test_db.variants (pt INT, payload VARIANT) PARTITIONED BY (pt) WITH ('variant.shreddingSchema' = '{"type":"ROW","fields":[{"name":"payload","type":{"type":"ROW","fields":[{"name":"value","type":"INT"}]}}]}')"#).await; + common::exec(&ctx, r#"INSERT INTO paimon.test_db.variants SELECT 2, parse_json('{"value":20}') UNION ALL SELECT 1, parse_json('{"value":10}')"#).await; + let query = "SELECT pt AS id, CAST(variant_get(payload, '$.value', 'int') AS INT) AS value FROM paimon.test_db.variants"; + let physical = plan(&ctx, query).await; + let text = displayable(physical.as_ref()).indent(true).to_string(); + assert!(text.contains("PushedVariants"), "{text}"); + assert!(text.contains("partitioning=Range"), "{text}"); + assert!(matches!( + scan(&physical).output_partitioning(), + Partitioning::Range(_) + )); + assert_eq!( + common::collect_id_value(&ctx, query).await, + vec![(1, 10), (2, 20)] + ); + + let audit = "SELECT pt FROM paimon.test_db.variants$audit_log"; + let scan = scan(&plan(&ctx, audit).await); + assert_eq!(scan.name(), "PaimonAuditLogScan"); + assert!(matches!(scan.output_partitioning(), Partitioning::Range(_))); + assert_eq!(common::row_count(&ctx, audit).await, 2); +} + +#[tokio::test] +async fn unsupported_or_unavailable_partition_keys_keep_unknown_partitioning() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + for (name, definition, partitioning, values, columns) in [ + ("plain", "id INT, pt INT", "", "(1, 10), (2, 20)", "id, pt"), + ( + "floats", + "id INT, pt DOUBLE", + "PARTITIONED BY (pt)", + "(1, -1.0), (2, 2.0)", + "id, pt", + ), + ( + "timestamps", + "id INT, pt TIMESTAMP", + "PARTITIONED BY (pt)", + "(1, TIMESTAMP '1969-12-31 23:59:59'), (2, TIMESTAMP '2026-01-01 00:00:00')", + "id, pt", + ), + ( + "projection", + "id INT, pt INT", + "PARTITIONED BY (pt)", + "(1, 10), (2, 20)", + "id", + ), + ( + "partial_key", + "id INT, pt INT, region STRING", + "PARTITIONED BY (pt, region)", + "(1, 10, 'a'), (2, 20, 'b')", + "id, pt", + ), + ] { + common::exec( + &ctx, + &format!("CREATE TABLE paimon.test_db.{name} ({definition}) {partitioning}"), + ) + .await; + common::exec( + &ctx, + &format!("INSERT INTO paimon.test_db.{name} VALUES {values}"), + ) + .await; + let query = format!("SELECT {columns} FROM paimon.test_db.{name}"); + let scan = scan(&plan(&ctx, &query).await); + assert!( + matches!( + scan.output_partitioning(), + Partitioning::UnknownPartitioning(_) + ), + "{name}: {:?}", + scan.output_partitioning() + ); + assert_eq!(common::row_count(&ctx, &query).await, 2); + } + common::exec( + &ctx, + "CREATE TABLE paimon.test_db.empty_ranges (pt INT) PARTITIONED BY (pt)", + ) + .await; + let scan = scan(&plan(&ctx, "SELECT pt FROM paimon.test_db.empty_ranges").await); + assert!(matches!( + scan.output_partitioning(), + Partitioning::UnknownPartitioning(1) + )); +} diff --git a/docs/src/sql.md b/docs/src/sql.md index 156f52895..f60fc112c 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -78,10 +78,25 @@ on the `RuntimeEnv` passed to `SQLContext::builder`. Paimon `MERGE INTO` continues to use `SQLContext`'s Paimon write implementation. DataFusion's generic `TableProvider::merge_into` hook, `file_row_index()` and -`input_file_name()` for Paimon scans, range-partition declarations, and custom -query-planner FFI are not provided by this upgrade. In particular, Paimon's -logical row IDs are not physical file row offsets, and ordinary scan splits -do not guarantee disjoint sorted ranges. +`input_file_name()` for Paimon scans, and custom query-planner FFI are not +provided by this upgrade. Paimon's logical row IDs are not physical file row +offsets. + +Scans declare DataFusion `Range` partitioning when every table partition key is +present unchanged in the scan output and uses a supported type: Boolean, +integers, strings, Date, millisecond Time, or Decimal128. Composite keys use +lexicographic order with `ASC NULLS LAST`. Whole splits with equal partition +keys stay together; adjacent key groups are combined into at most the configured +`datafusion.execution.target_partitions`. This can avoid exchanges for compatible +aggregations and joins, but does not declare that rows within a range are sorted. +Parallelism may be lower than the target when there are fewer distinct keys. + +Unpartitioned tables, empty scans, scans with only one execution partition or one +distinct key, missing/projected-out keys, unsupported types (including floating +point, binary, and timestamp keys), unreadable partition metadata, and scans with +query authorization or attached authorization rules retain `UnknownPartitioning` +and their original split assignments. File min/max statistics and hash buckets +alone are not used to infer ranges. ## Registering Catalog From ec02b94fc2725232edeace90145922a0122d9294 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 3 Nov 2026 17:39:22 +0800 Subject: [PATCH 2/4] [datafusion] Keep mutable partition columns out of range metadata --- .../src/physical_plan/range_partitioning.rs | 19 ++- .../datafusion/tests/range_partitioning.rs | 131 ++++++++++++++++++ docs/src/sql.md | 6 +- 3 files changed, 153 insertions(+), 3 deletions(-) diff --git a/crates/integrations/datafusion/src/physical_plan/range_partitioning.rs b/crates/integrations/datafusion/src/physical_plan/range_partitioning.rs index 47008d409..05b90c116 100644 --- a/crates/integrations/datafusion/src/physical_plan/range_partitioning.rs +++ b/crates/integrations/datafusion/src/physical_plan/range_partitioning.rs @@ -44,11 +44,28 @@ pub(super) fn range_partitioned_splits( partitions: &[Arc<[DataSplit]>], ) -> Option<(Vec>, Partitioning)> { let fields = table.schema().partition_fields(); + let options = table.schema().core_options(); if fields.is_empty() || fields.len() != table.schema().partition_keys().len() || partitions.len() < 2 || read_type.len() != schema.fields().len() - || table.schema().core_options().query_auth_enabled() + || options.query_auth_enabled() + { + return None; + } + + // Native PK value aggregation can change non-key partition columns. Those + // values can remain in files even after the merge engine is changed, so the + // current engine alone cannot prove range membership. Full PK columns are + // immutable; format-table readers reconstruct partition columns from splits. + let primary_keys = table.schema().primary_keys(); + if !options.is_format_table() + && !primary_keys.is_empty() + && table + .schema() + .partition_keys() + .iter() + .any(|key| !primary_keys.contains(key)) { return None; } diff --git a/crates/integrations/datafusion/tests/range_partitioning.rs b/crates/integrations/datafusion/tests/range_partitioning.rs index 0ec13a6d2..555442b98 100644 --- a/crates/integrations/datafusion/tests/range_partitioning.rs +++ b/crates/integrations/datafusion/tests/range_partitioning.rs @@ -143,6 +143,137 @@ async fn integer_ranges_keep_equal_keys_together_and_avoid_aggregation_shuffle() .contains("SortExec")); } +#[tokio::test] +async fn aggregation_of_non_primary_partition_keys_keeps_unknown_distribution() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + common::exec(&ctx, "CREATE TABLE paimon.test_db.aggregated_ranges (id INT NOT NULL, pt INT, value INT, PRIMARY KEY (id)) PARTITIONED BY (pt) WITH ('bucket' = '-1', 'merge-engine' = 'aggregation', 'fields.pt.aggregate-function' = 'sum', 'fields.value.aggregate-function' = 'sum')").await; + common::exec( + &ctx, + "INSERT INTO paimon.test_db.aggregated_ranges VALUES (1, 1, 10), (2, 2, 30)", + ) + .await; + common::exec( + &ctx, + "INSERT INTO paimon.test_db.aggregated_ranges VALUES (1, 1, 20)", + ) + .await; + let query = "SELECT pt AS id, CAST(sum(value) AS INT) AS value FROM paimon.test_db.aggregated_ranges GROUP BY pt"; + assert_eq!(common::collect_id_value(&ctx, query).await, vec![(2, 60)]); + let physical = plan( + &ctx, + "SELECT id, pt, value FROM paimon.test_db.aggregated_ranges", + ) + .await; + assert!(matches!( + scan(&physical).output_partitioning(), + Partitioning::UnknownPartitioning(_) + )); +} + +#[tokio::test] +async fn partial_update_of_non_primary_partition_keys_keeps_unknown_distribution() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + common::exec(&ctx, "CREATE TABLE paimon.test_db.partial_ranges (id INT NOT NULL, pt INT, ver INT, value INT, PRIMARY KEY (id)) PARTITIONED BY (pt) WITH ('bucket' = '-1', 'merge-engine' = 'partial-update', 'fields.ver.sequence-group' = 'pt,value', 'fields.pt.aggregate-function' = 'sum', 'fields.value.aggregate-function' = 'sum')").await; + common::exec( + &ctx, + "INSERT INTO paimon.test_db.partial_ranges VALUES (1, 1, 1, 10), (2, 2, 1, 30)", + ) + .await; + common::exec( + &ctx, + "INSERT INTO paimon.test_db.partial_ranges VALUES (1, 1, 2, 20)", + ) + .await; + let query = "SELECT pt AS id, CAST(sum(value) AS INT) AS value FROM paimon.test_db.partial_ranges GROUP BY pt"; + assert_eq!(common::collect_id_value(&ctx, query).await, vec![(2, 60)]); + let physical = plan( + &ctx, + "SELECT id, pt, value FROM paimon.test_db.partial_ranges", + ) + .await; + assert!(matches!( + scan(&physical).output_partitioning(), + Partitioning::UnknownPartitioning(_) + )); +} + +#[tokio::test] +async fn historically_aggregated_non_primary_partition_keys_keep_unknown_distribution() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + common::exec(&ctx, "CREATE TABLE paimon.test_db.historical_ranges (id INT NOT NULL, pt INT, value INT, PRIMARY KEY (id)) PARTITIONED BY (pt) WITH ('bucket' = '-1', 'merge-engine' = 'aggregation', 'fields.pt.aggregate-function' = 'sum', 'fields.value.aggregate-function' = 'sum')").await; + // Writer-side aggregation persists pt=2 in a file assigned to partition 1. + common::exec( + &ctx, + "INSERT INTO paimon.test_db.historical_ranges VALUES (1, 1, 10), (1, 1, 20), (2, 2, 30)", + ) + .await; + common::exec( + &ctx, + "ALTER TABLE paimon.test_db.historical_ranges SET ('merge-engine' = 'deduplicate')", + ) + .await; + let query = "SELECT pt AS id, CAST(sum(value) AS INT) AS value FROM paimon.test_db.historical_ranges GROUP BY pt"; + assert_eq!(common::collect_id_value(&ctx, query).await, vec![(2, 60)]); + let physical = plan( + &ctx, + "SELECT id, pt, value FROM paimon.test_db.historical_ranges", + ) + .await; + assert!(matches!( + scan(&physical).output_partitioning(), + Partitioning::UnknownPartitioning(_) + )); +} + +#[tokio::test] +async fn merge_engines_preserve_ranges_for_partition_columns_in_the_primary_key() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + for (name, engine, extra) in [ + ( + "aggregation_full_pk", + "aggregation", + ", 'fields.pt.aggregate-function' = 'sum'", + ), + ( + "partial_full_pk", + "partial-update", + ", 'fields.ver.sequence-group' = 'value'", + ), + ] { + common::exec(&ctx, &format!("CREATE TABLE paimon.test_db.{name} (id INT NOT NULL, pt INT NOT NULL, ver INT, value INT, PRIMARY KEY (pt, id)) PARTITIONED BY (pt) WITH ('bucket' = '1', 'merge-engine' = '{engine}', 'fields.value.aggregate-function' = 'sum'{extra})")).await; + common::exec( + &ctx, + &format!("INSERT INTO paimon.test_db.{name} VALUES (1, 1, 1, 10), (2, 2, 1, 30)"), + ) + .await; + common::exec( + &ctx, + &format!("INSERT INTO paimon.test_db.{name} VALUES (1, 1, 2, 20)"), + ) + .await; + let query = format!("SELECT pt AS id, CAST(sum(value) AS INT) AS value FROM paimon.test_db.{name} GROUP BY pt"); + let physical = plan(&ctx, &query).await; + let text = displayable(physical.as_ref()).indent(true).to_string(); + assert!( + matches!( + scan(&physical).output_partitioning(), + Partitioning::Range(_) + ), + "{text}" + ); + assert!(!text.contains("RepartitionExec"), "{text}"); + assert_eq!( + common::collect_id_value(&ctx, &query).await, + vec![(1, 30), (2, 30)], + "{engine}" + ); + } +} + #[tokio::test] async fn joins_align_different_ranges_after_primary_key_updates() { let (_tmp, ctx) = common::setup_sql_context().await; diff --git a/docs/src/sql.md b/docs/src/sql.md index f60fc112c..e15ac8f75 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -95,8 +95,10 @@ Unpartitioned tables, empty scans, scans with only one execution partition or on distinct key, missing/projected-out keys, unsupported types (including floating point, binary, and timestamp keys), unreadable partition metadata, and scans with query authorization or attached authorization rules retain `UnknownPartitioning` -and their original split assignments. File min/max statistics and hash buckets -alone are not used to infer ranges. +and their original split assignments. Native primary-key tables also retain +`UnknownPartitioning` unless every partition column belongs to the primary key: +current or historical value aggregation can change a non-key partition column. +File min/max statistics and hash buckets alone are not used to infer ranges. ## Registering Catalog From c3bab53fa202401d87fdb0c9f7cccb9bcd41305d Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 3 Nov 2026 18:35:47 +0800 Subject: [PATCH 3/4] [datafusion] Cover compatible joins and format-table ranges --- .../datafusion/tests/range_partitioning.rs | 224 +++++++++++++++++- 1 file changed, 223 insertions(+), 1 deletion(-) diff --git a/crates/integrations/datafusion/tests/range_partitioning.rs b/crates/integrations/datafusion/tests/range_partitioning.rs index 555442b98..0d5096a5a 100644 --- a/crates/integrations/datafusion/tests/range_partitioning.rs +++ b/crates/integrations/datafusion/tests/range_partitioning.rs @@ -19,12 +19,15 @@ mod common; use std::sync::Arc; -use arrow_array::{Array, Int32Array}; +use arrow_array::{Array, Int32Array, RecordBatch}; +use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; use datafusion::common::ScalarValue; use datafusion::physical_plan::{ collect_partitioned, displayable, ExecutionPlan, ExecutionPlanProperties, Partitioning, }; use paimon_datafusion::SQLContext; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::arrow::ArrowWriter; fn scan(plan: &Arc) -> Arc { fn find(plan: &Arc) -> Option> { @@ -343,6 +346,225 @@ async fn joins_align_different_ranges_after_primary_key_updates() { } } +#[tokio::test] +async fn joins_with_identical_ranges_avoid_exchanges() { + let (_tmp, ctx) = common::setup_sql_context().await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + common::exec(&ctx, "CREATE TABLE paimon.test_db.compatible_left (id INT, pt INT, value INT) PARTITIONED BY (pt)").await; + common::exec( + &ctx, + "CREATE TABLE paimon.test_db.compatible_right (pt INT, value INT) PARTITIONED BY (pt)", + ) + .await; + common::exec(&ctx, "INSERT INTO paimon.test_db.compatible_left VALUES (1, 1, 10), (2, 2, 20), (3, 3, 30), (4, 4, 40), (5, 3, 50)").await; + common::exec(&ctx, "INSERT INTO paimon.test_db.compatible_right VALUES (1, 100), (2, 200), (3, 300), (4, 400), (3, 301)").await; + + let left = scan(&plan(&ctx, "SELECT * FROM paimon.test_db.compatible_left").await); + let right = scan(&plan(&ctx, "SELECT * FROM paimon.test_db.compatible_right").await); + let (Partitioning::Range(left), Partitioning::Range(right)) = + (left.output_partitioning(), right.output_partitioning()) + else { + panic!("both inputs must have ranges"); + }; + assert_eq!(left.partition_count(), 2); + assert_eq!(right.partition_count(), 2); + assert_eq!(left.split_points(), right.split_points()); + assert_eq!( + left.split_points()[0].values(), + &[ScalarValue::Int32(Some(3))] + ); + + common::exec( + &ctx, + "SET datafusion.optimizer.hash_join_single_partition_threshold = 0", + ) + .await; + common::exec( + &ctx, + "SET datafusion.optimizer.hash_join_single_partition_threshold_rows = 0", + ) + .await; + for prefer_hash in [true, false] { + common::exec( + &ctx, + &format!("SET datafusion.optimizer.prefer_hash_join = {prefer_hash}"), + ) + .await; + let query = "SELECT l.id, l.value + r.value AS value FROM paimon.test_db.compatible_left l JOIN paimon.test_db.compatible_right r ON l.pt = r.pt"; + let physical = plan(&ctx, query).await; + let text = displayable(physical.as_ref()).indent(true).to_string(); + let join = if prefer_hash { + "HashJoinExec" + } else { + "SortMergeJoinExec" + }; + assert!(text.contains(join), "{text}"); + assert!( + !text.contains("RepartitionExec"), + "compatible ranges should avoid exchanges: {text}" + ); + if !prefer_hash { + assert!( + text.contains("SortExec"), + "unordered range inputs still need local sorts: {text}" + ); + } + let mut rows = common::collect_id_value(&ctx, query).await; + rows.sort_unstable(); + assert_eq!( + rows, + vec![ + (1, 110), + (2, 220), + (3, 330), + (3, 331), + (4, 440), + (5, 350), + (5, 351) + ], + "{join}" + ); + } +} + +#[tokio::test] +async fn format_table_restores_partition_columns_into_declared_ranges() { + use paimon::catalog::{Catalog, Identifier}; + use paimon::spec::{DataType, IntType, Schema}; + use paimon::{CatalogOptions, FileSystemCatalog, Options}; + + let tmp = tempfile::TempDir::new().unwrap(); + let mut options = Options::new(); + // Match the file:/ paths returned by listing, as in the format-table fixtures. + options.set( + CatalogOptions::WAREHOUSE, + format!("file:{}", tmp.path().display()), + ); + let catalog = Arc::new(FileSystemCatalog::new(options).unwrap()); + catalog + .create_database("test_db", false, Default::default()) + .await + .unwrap(); + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("pt", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .partition_keys(["pt"]) + .option("type", "format-table") + .option("file.format", "parquet") + .build() + .unwrap(); + catalog + .create_table(&Identifier::new("test_db", "format_ranges"), schema, false) + .await + .unwrap(); + + // Physical files hold only data columns; pt must come from split metadata. + let physical_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowDataType::Int32, true), + Field::new("value", ArrowDataType::Int32, true), + ])); + for (pt, ids, values) in [ + (10, vec![4, 5], vec![40, 50]), + (-10, vec![1, 2], vec![10, 20]), + (2, vec![3], vec![30]), + ] { + let directory = tmp.path().join(format!("test_db.db/format_ranges/pt={pt}")); + std::fs::create_dir_all(&directory).unwrap(); + let path = directory.join("part-0.parquet"); + let batch = RecordBatch::try_new( + Arc::clone(&physical_schema), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(Int32Array::from(values)), + ], + ) + .unwrap(); + let mut writer = ArrowWriter::try_new( + std::fs::File::create(&path).unwrap(), + Arc::clone(&physical_schema), + None, + ) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let reader = + ParquetRecordBatchReaderBuilder::try_new(std::fs::File::open(&path).unwrap()).unwrap(); + assert!(reader.schema().field_with_name("pt").is_err()); + } + + let ctx = common::create_sql_context(catalog).await; + common::exec(&ctx, "SET datafusion.execution.target_partitions = 2").await; + let physical = plan( + &ctx, + "SELECT id, pt, value FROM paimon.test_db.format_ranges", + ) + .await; + let scan = scan(&physical); + let Partitioning::Range(range) = scan.output_partitioning() else { + panic!( + "format table must declare ranges: {:?}", + scan.output_partitioning() + ); + }; + assert_eq!(range.partition_count(), 2); + assert_eq!( + range.split_points()[0].values(), + &[ScalarValue::Int32(Some(10))] + ); + let mut actual = Vec::new(); + for batches in collect_partitioned(scan, ctx.ctx().task_ctx()) + .await + .unwrap() + { + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let keys = batch + .column_by_name("pt") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column_by_name("value") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + rows.extend( + (0..batch.num_rows()).map(|i| (keys.value(i), ids.value(i), values.value(i))), + ); + } + rows.sort_unstable(); + actual.push(rows); + } + assert_eq!( + actual, + vec![ + vec![(-10, 1, 10), (-10, 2, 20), (2, 3, 30)], + vec![(10, 4, 40), (10, 5, 50)] + ] + ); + + let query = "SELECT pt AS id, CAST(sum(value) AS INT) AS value FROM paimon.test_db.format_ranges GROUP BY pt"; + let physical = plan(&ctx, query).await; + let text = displayable(physical.as_ref()).indent(true).to_string(); + assert!( + !text.contains("RepartitionExec"), + "format-table ranges should avoid an exchange: {text}" + ); + assert_eq!( + common::collect_id_value(&ctx, query).await, + vec![(-10, 30), (2, 30), (10, 90)] + ); +} + #[tokio::test] async fn compound_date_and_string_keys_preserve_lexicographic_ranges() { let (_tmp, ctx) = common::setup_sql_context().await; From 0fec7699f8eaa5f12178573178937871f7788d33 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Tue, 3 Nov 2026 19:40:29 +0800 Subject: [PATCH 4/4] [test] Stabilize cache lock test with inherited descriptors --- crates/paimon/src/io/cache/shared.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/crates/paimon/src/io/cache/shared.rs b/crates/paimon/src/io/cache/shared.rs index 4ced74492..60333270c 100644 --- a/crates/paimon/src/io/cache/shared.rs +++ b/crates/paimon/src/io/cache/shared.rs @@ -1172,12 +1172,23 @@ mod tests { let reverse = [collision, b, a]; let guards = first.load_locks(&forward).await; assert_eq!(guards.len(), 2); + // Model a concurrent fork/dup keeping the locked descriptions alive. + let inherited: Vec<_> = guards + .iter() + .map(|guard| guard.try_clone().unwrap()) + .collect(); let contended = tokio::time::timeout(Duration::from_secs(1), second.load_locks(&reverse)) .await .expect("busy stripes must not block another batch"); assert!(contended.is_empty()); + // Closing only our handles need not release locks inherited by a fork. + // Explicitly unlock before checking non-blocking reacquisition. + for guard in &guards { + guard.unlock().unwrap(); + } drop(guards); assert_eq!(second.load_locks(&reverse).await.len(), 2); + drop(inherited); } #[tokio::test]