From 8f6876d6ba247b81a8997acfb0cfe665bb970675 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 11:02:25 +1000 Subject: [PATCH 1/4] chore(deps-dev): bump shell-quote from 1.8.3 to 1.8.4 in /datafusion/wasmtest/datafusion-wasm-app (#22856) Bumps [shell-quote](https://github.com/ljharb/shell-quote) from 1.8.3 to 1.8.4.
Changelog

Sourced from shell-quote's changelog.

v1.8.4 - 2026-05-22

Commits

Commits

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=shell-quote&package-manager=npm_and_yarn&previous-version=1.8.3&new-version=1.8.4)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores) Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`. [//]: # (dependabot-automerge-start) [//]: # (dependabot-automerge-end) ---
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself) You can disable automated security fix PRs for this repo from the [Security Alerts page](https://github.com/apache/datafusion/network/alerts).
Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .../wasmtest/datafusion-wasm-app/package-lock.json | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json index 3e255bdd3c5e2..c476ea76347ab 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json @@ -3431,11 +3431,10 @@ } }, "node_modules/shell-quote": { - "version": "1.8.3", - "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.8.3.tgz", - "integrity": "sha512-ObmnIF4hXNg1BqhnHmgbDETF8dLPCggZWBjkQfhZpbszZnYur5DUljTcCHii5LC3J5E0yeO/1LIMyH+UvHQgyw==", + "version": "1.8.4", + "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.8.4.tgz", + "integrity": "sha512-VsC6n6vz1ihYYyZZwX7YZSF5l5x36ca17OC+a69h94YqB7X6XLwf+5MOgynYir2SLFUbl8gIYvBo8K8RoNQ6bQ==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" }, @@ -6756,9 +6755,9 @@ "dev": true }, "shell-quote": { - "version": "1.8.3", - "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.8.3.tgz", - "integrity": "sha512-ObmnIF4hXNg1BqhnHmgbDETF8dLPCggZWBjkQfhZpbszZnYur5DUljTcCHii5LC3J5E0yeO/1LIMyH+UvHQgyw==", + "version": "1.8.4", + "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.8.4.tgz", + "integrity": "sha512-VsC6n6vz1ihYYyZZwX7YZSF5l5x36ca17OC+a69h94YqB7X6XLwf+5MOgynYir2SLFUbl8gIYvBo8K8RoNQ6bQ==", "dev": true }, "side-channel": { From 40a64546c8ed25a00c8100d2fb4f44890e22fc25 Mon Sep 17 00:00:00 2001 From: Yongting You <2010youy01@gmail.com> Date: Wed, 10 Jun 2026 09:09:41 +0800 Subject: [PATCH 2/4] refactor(hash-aggr): Forward port the soft limit optimization to the new hash aggregation impl (#22824) ## Which issue does this PR close? - part of https://github.com/apache/datafusion/issues/22710 ## Rationale for this change Part of rewriting hash aggregation into several dedicated streams. In the first step https://github.com/apache/datafusion/pull/22729, `PartialHashAggregateStream` and `FinalHashAggregateStream` has been split from the old `GroupsHashAggregateStream`, but both stream only have basic implementation, no optimizations and extra features like spilling. \* it's incremental migration, so old impl won't change, we plan to delete it once migration is finished This PR forward ports the below optimization to the new implementation: - https://github.com/apache/datafusion/pull/8038 The optimizer part don't have to move, ported changes are only inside aggregate operator. ## What changes are included in this PR? Extends `PartialHashAggregateStream` and `FinalHashAggregateStream` to apply the optimization. See code comment at `datafusion/physical-plan/src/aggregates/hash_aggregate.rs` for the background. ## Are these changes tested? Yes, the original test in https://github.com/apache/datafusion/pull/8038 is only at `ExecutionPlan` level, they're still passing after the change. This PR added new test coverage: check `explain analyze` to ensure the implementation actually respects this soft limit at runtime. ## Are there any user-facing changes? --------- Co-authored-by: Martin Grigorov --- .../limited_distinct_aggregation.rs | 117 +++++++++++++++++- .../src/aggregates/hash_aggregate.rs | 96 ++++++++++++-- .../src/aggregates/hash_table.rs | 4 + .../physical-plan/src/aggregates/mod.rs | 106 +++++++++++++++- 4 files changed, 312 insertions(+), 11 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs index c523b4a752a82..323dfd6183306 100644 --- a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs +++ b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs @@ -36,7 +36,7 @@ use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::{ ExecutionPlan, aggregates::{AggregateExec, AggregateMode}, - collect, + collect, displayable, limit::{GlobalLimitExec, LocalLimitExec}, }; @@ -104,6 +104,121 @@ async fn test_partial_final() -> Result<()> { Ok(()) } +// Ensure operator respect the soft limit and stops early: `AggregateExec`'s +// `output_rows` metric should be smaller than then total distinct group count. +#[tokio::test] +async fn limited_distinct_aggregate_stream_respects_soft_limit() -> Result<()> { + // Snapshot for an aggregate operator node from `EXPLAIN ANALYZE`. + // + // Example: In an `EXPLAIN ANALYZE` output + // ```txt + // AggregateExec: mode=partial, limit=10, metrics=[output_rows=100, ...] + // ``` + // we get: + // ```txt + // AggregateRuntimeMetric { + // mode: Partial, + // limit: Some(10), + // output_rows: 100, + // } + // ``` + #[derive(Debug)] + struct AggregateRuntimeMetric { + mode: AggregateMode, + limit: Option, + output_rows: usize, + } + + fn collect_aggregate_runtime_metrics( + plan: &Arc, + metrics: &mut Vec, + ) { + if let Some(agg) = plan.downcast_ref::() { + let output_rows = agg + .metrics() + .and_then(|metrics| metrics.aggregate_by_name().output_rows()) + .expect("AggregateExec should record output_rows after execution"); + + metrics.push(AggregateRuntimeMetric { + mode: *agg.mode(), + limit: agg.limit_options().map(|config| config.limit()), + output_rows, + }); + } + + for child in plan.children() { + collect_aggregate_runtime_metrics(child, metrics); + } + } + + fn aggregate_runtime_metrics( + plan: &Arc, + ) -> Vec { + let mut metrics = vec![]; + collect_aggregate_runtime_metrics(plan, &mut metrics); + metrics + } + + let cfg = SessionConfig::new() + .with_target_partitions(2) + .with_batch_size(10) + .set_bool("datafusion.execution.enable_migration_aggregate", true); + let ctx = SessionContext::new_with_config(cfg); + + let dataframe = ctx + .sql( + "SELECT DISTINCT value % 100000 AS v \ + FROM generate_series(1000000) \ + LIMIT 10", + ) + .await?; + let plan = dataframe.create_physical_plan().await?; + let formatted_plan = displayable(plan.as_ref()).indent(false).to_string(); + assert!( + formatted_plan.contains("AggregateExec: mode=Partial"), + "expected a partial aggregate in plan:\n{formatted_plan}" + ); + assert!( + formatted_plan.contains("AggregateExec: mode=FinalPartitioned"), + "expected a final partitioned aggregate in plan:\n{formatted_plan}" + ); + + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await?; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 10 + ); + + let metrics = aggregate_runtime_metrics(&plan); + let partial = metrics + .iter() + .find(|metric| metric.mode == AggregateMode::Partial) + .expect("expected partial aggregate metrics"); + let final_aggregate = metrics + .iter() + .find(|metric| { + matches!( + metric.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) + }) + .expect("expected final aggregate metrics"); + + assert_eq!(partial.limit, Some(10)); + assert_eq!(final_aggregate.limit, Some(10)); + + assert!( + partial.output_rows <= 100, + "partial aggregate should stop before emitting all distinct groups: {metrics:?}" + ); + assert!( + final_aggregate.output_rows <= 100, + "final aggregate should stop before emitting all distinct groups: {metrics:?}" + ); + + Ok(()) +} + #[tokio::test] async fn test_single_local() -> Result<()> { let source = mock_data()?; diff --git a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs index f25299631a92c..0c8593efd05bb 100644 --- a/datafusion/physical-plan/src/aggregates/hash_aggregate.rs +++ b/datafusion/physical-plan/src/aggregates/hash_aggregate.rs @@ -60,6 +60,33 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metric /// ## Final Stage Behavior /// Input: partial states /// Output: results for all groups (e.g. for avg(x), it's avg(x) calculated from the state) +/// +/// # Optimization: DISTINCT LIMIT Soft Limit +/// +/// This optimization applies to both [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`] +/// +/// Unordered distinct queries such as: +/// +/// ```sql +/// SELECT DISTINCT x FROM t LIMIT 10; +/// ``` +/// +/// are optimized into a two-stage aggregate like: +/// +/// ```txt +/// LimitExec, limit=10 +/// --AggregateExec(Final), group_by=[x], aggr=[], soft_limit=10 +/// ---- RepartitionExec, partitioning=hash(x) +/// ------ AggregateExec(Partial), group_by=[x], aggr=[], soft_limit=10 +/// -------- Scan(t) +/// ``` +/// +/// After each input batch, the stream checks whether the soft limit has been +/// reached. If so, it emits the accumulated groups and stops reading input. +/// +/// This operator does not guarantee an exact limit because a single batch can +/// cross the threshold. The downstream limit operator enforces the exact result +/// size. pub(crate) struct PartialHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, @@ -78,6 +105,12 @@ pub(crate) struct PartialHashAggregateStream { /// Tracks partial aggregation row reduction, matching `GroupedHashAggregateStream`. reduction_factor: metrics::RatioMetrics, + + /// Optional soft limit on the number of groups to accumulate before output. + /// + /// Invariant: when this is `Some(..)`, the accumulators inside `hash_table` must + /// be empty. See struct comments for details. + group_values_soft_limit: Option, } /// Hash aggregation uses a 2-stage (partial and final) hash aggregation, this stream @@ -99,6 +132,9 @@ pub(crate) struct FinalHashAggregateStream { /// Memory reservation for group keys and accumulators. reservation: MemoryReservation, + + /// See comments for the same variable in [`PartialHashAggregateStream`] + group_values_soft_limit: Option, } impl PartialHashAggregateStream { @@ -139,8 +175,21 @@ impl PartialHashAggregateStream { baseline_metrics, reservation, reduction_factor, + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), }) } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit(&self) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= self.hash_table.building_group_count()) + } + + fn start_output(&mut self) -> Result<()> { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + self.hash_table.start_output() + } } impl Stream for PartialHashAggregateStream { @@ -169,6 +218,18 @@ impl Stream for PartialHashAggregateStream { return Poll::Ready(Some(Err(e))); } + if self.hit_soft_group_limit() { + let timer = elapsed_compute.timer(); + let result = self.start_output(); + timer.done(); + + if let Err(e) = result { + return Poll::Ready(Some(Err(e))); + } + + continue; + } + // TODO: impl memory-limited aggr, when OOM directly send // partial state to final aggregate stage if let Err(e) = @@ -181,11 +242,8 @@ impl Stream for PartialHashAggregateStream { return Poll::Ready(Some(Err(e))); } Poll::Ready(None) => { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - let timer = elapsed_compute.timer(); - let result = self.hash_table.start_output(); + let result = self.start_output(); timer.done(); if let Err(e) = result { @@ -262,8 +320,21 @@ impl FinalHashAggregateStream { hash_table, baseline_metrics, reservation, + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), }) } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit(&self) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= self.hash_table.building_group_count()) + } + + fn start_output(&mut self) -> Result<()> { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + self.hash_table.start_output() + } } impl Stream for FinalHashAggregateStream { @@ -291,6 +362,18 @@ impl Stream for FinalHashAggregateStream { return Poll::Ready(Some(Err(e))); } + if self.hit_soft_group_limit() { + let timer = elapsed_compute.timer(); + let result = self.start_output(); + timer.done(); + + if let Err(e) = result { + return Poll::Ready(Some(Err(e))); + } + + continue; + } + if let Err(e) = self.reservation.try_resize(self.hash_table.memory_size()) { @@ -301,11 +384,8 @@ impl Stream for FinalHashAggregateStream { return Poll::Ready(Some(Err(e))); } Poll::Ready(None) => { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - let timer = elapsed_compute.timer(); - let result = self.hash_table.start_output(); + let result = self.start_output(); timer.done(); if let Err(e) = result { diff --git a/datafusion/physical-plan/src/aggregates/hash_table.rs b/datafusion/physical-plan/src/aggregates/hash_table.rs index 278689d23f264..87f16d0eebe6f 100644 --- a/datafusion/physical-plan/src/aggregates/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/hash_table.rs @@ -342,6 +342,10 @@ impl AggregateHashTable { } } + pub(super) fn building_group_count(&self) -> usize { + self.state.building().group_values.len() + } + pub(super) fn is_building(&self) -> bool { matches!(self.state, AggregateHashTableState::Building(_)) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 5be65f862c5c0..a5f1621812561 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1038,10 +1038,10 @@ impl AggregateExec { } self.mode == AggregateMode::Partial - && self.limit_options.is_none() && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() && self.group_by.is_single() + && self.limit_options_supported_by_hash_stream() } fn should_use_final_hash_stream(&self, context: &TaskContext) -> bool { @@ -1053,12 +1053,17 @@ impl AggregateExec { matches!( self.mode, AggregateMode::Final | AggregateMode::FinalPartitioned - ) && self.limit_options.is_none() + ) && self.limit_options_supported_by_hash_stream() && self.input_order_mode == InputOrderMode::Linear && !self.group_by.is_true_no_grouping() && self.group_by.is_single() } + /// See comments in `PartialHashAggregateStream` limit optimization section + fn limit_options_supported_by_hash_stream(&self) -> bool { + self.limit_options.is_none() || self.is_unordered_unfiltered_group_by_distinct() + } + /// Finds the DataType and SortDirection for this Aggregate, if there is one pub fn get_minmax_desc(&self) -> Option<(FieldRef, bool)> { let agg_expr = self.aggr_expr.iter().exactly_one().ok()?; @@ -3157,6 +3162,103 @@ mod tests { Ok(()) } + #[tokio::test] + async fn limited_distinct_aggregate_uses_migrated_hash_streams() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![1, 2, 1]))], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![3, 4]))], + )?, + ]; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + + let partial_input = TestMemoryExec::try_new_exec( + std::slice::from_ref(&input_batches), + Arc::clone(&schema), + None, + )?; + let partial_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + partial_input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(2))), + ); + + let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(partial_stream, StreamType::PartialHash(_))); + let stream: SendableRecordBatchStream = partial_stream.into(); + let partial_output = collect(stream).await?; + assert_eq!( + partial_output + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 2 + ); + assert_snapshot!(batches_to_sort_string(&partial_output), @r" ++---+ +| a | ++---+ +| 1 | +| 2 | ++---+ +"); + + let final_input = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + let final_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![], + vec![], + final_input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(2))), + ); + + let final_stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(final_stream, StreamType::FinalHash(_))); + let stream: SendableRecordBatchStream = final_stream.into(); + let final_output = collect(stream).await?; + assert_eq!( + final_output + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 2 + ); + assert_snapshot!(batches_to_sort_string(&final_output), @r" ++---+ +| a | ++---+ +| 1 | +| 2 | ++---+ +"); + + Ok(()) + } + #[tokio::test] async fn test_drop_cancel_without_groups() -> Result<()> { let task_ctx = Arc::new(TaskContext::default()); From d77a02d662b27852cc247153d9edcad5a92905d4 Mon Sep 17 00:00:00 2001 From: WeblWabl Date: Tue, 9 Jun 2026 21:52:27 -0500 Subject: [PATCH 3/4] feat: Add From> trait for Precision enum (#22792) This commit adds a `From>` trait for `Precision`. Porting from `influxdb` to upstream: https://github.com/influxdata/influxdb/blob/70335b158808881c2f5a9ef27cd4bcbe4944686a/core/datafusion_util/src/lib.rs#L153-L163 --- datafusion/common/src/stats.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs index 320fd43751025..a64d5e00ee6df 100644 --- a/datafusion/common/src/stats.rs +++ b/datafusion/common/src/stats.rs @@ -318,6 +318,12 @@ impl Precision { } } +impl From> for Precision { + fn from(option: Option) -> Self { + option.map_or(Precision::Absent, Precision::Exact) + } +} + impl Debug for Precision { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { From d23321d0d9626264f9e6b3192a51e9649c590872 Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Wed, 10 Jun 2026 01:35:19 -0400 Subject: [PATCH 4/4] Add logical range partitioning representation (#22777) ## Which issue does this PR close? - Closes #22778. - Related: #21992, #22395. - Needed by #22657. ## Rationale for this change Declared scan output partitioning should use logical partitioning metadata, not physical partitioning types. This adds logical range partitioning so range-partitioned sources can declare their layout at the logical layer. ## What changes are included in this PR? - Add logical `Partitioning::Range` and `RangePartitioning`. - Move `SplitPoint` and shared split-point validation to `datafusion-common`. - Wire logical range partitioning through expression traversal, rewrites, and display. - Keep planning, logical proto, and Substrait support explicitly unsupported for now. ## Are these changes tested? Yes. Unit tests added ## Are there any user-facing changes? Yes. This adds public logical range partitioning API. No breaking API changes. --- datafusion/common/src/lib.rs | 2 + datafusion/common/src/partitioning.rs | 104 +++++++ datafusion/core/src/physical_planner.rs | 66 +++- datafusion/expr/src/logical_plan/display.rs | 17 ++ datafusion/expr/src/logical_plan/mod.rs | 6 +- datafusion/expr/src/logical_plan/plan.rs | 283 +++++++++++++++++- datafusion/expr/src/logical_plan/tree_node.rs | 15 + datafusion/physical-expr/src/lib.rs | 3 +- datafusion/physical-expr/src/partitioning.rs | 142 +-------- datafusion/proto/src/logical_plan/mod.rs | 5 + .../logical_plan/producer/rel/exchange_rel.rs | 10 + 11 files changed, 508 insertions(+), 145 deletions(-) create mode 100644 datafusion/common/src/partitioning.rs diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index e865c548bb554..2f6d9848b6e55 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -30,6 +30,7 @@ mod dfschema; mod functional_dependencies; mod join_type; mod param_value; +mod partitioning; mod schema_reference; mod table_reference; mod unnest; @@ -92,6 +93,7 @@ pub use join_type::{JoinConstraint, JoinSide, JoinType}; pub use nested_struct::cast_column; pub use null_equality::NullEquality; pub use param_value::ParamValues; +pub use partitioning::{SplitPoint, validate_range_split_points}; pub use scalar::{ScalarType, ScalarValue}; pub use schema_reference::SchemaReference; pub use spans::{Location, Span, Spans}; diff --git a/datafusion/common/src/partitioning.rs b/datafusion/common/src/partitioning.rs new file mode 100644 index 0000000000000..8a7212c2e3089 --- /dev/null +++ b/datafusion/common/src/partitioning.rs @@ -0,0 +1,104 @@ +// 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 crate::utils::compare_rows; +use crate::{Result, ScalarValue, error::_plan_err}; +use arrow::compute::SortOptions; +use std::cmp::Ordering; +use std::fmt::{self, Display}; + +/// A boundary between adjacent range partitions. +/// +/// A split point is a tuple with one [`ScalarValue`] per partitioning +/// expression. Split points are interpreted lexicographically according to the +/// ordering of the range partitioning that owns them. +/// +/// `N` split points define `N + 1` partitions: +/// +/// ```text +/// partition 0: key < split_points[0] +/// partition 1: split_points[0] <= key < split_points[1] +/// ... +/// partition N - 1: split_points[N - 2] <= key < split_points[N - 1] +/// partition N: split_points[N - 1] <= key +/// ``` +/// +/// Values equal to split point `i` belong to partition `i + 1`, so interior +/// partitions are lower-inclusive and upper-exclusive. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct SplitPoint { + values: Vec, +} + +impl SplitPoint { + /// Creates a new split point from its tuple values. + pub fn new(values: Vec) -> Self { + Self { values } + } + + /// Returns the tuple values for this split point. + pub fn values(&self) -> &[ScalarValue] { + &self.values + } +} + +impl Display for SplitPoint { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let values = self + .values + .iter() + .map(ToString::to_string) + .collect::>() + .join(", "); + write!(f, "({values})") + } +} + +/// Validates that split points match the ordering width and are strictly +/// ordered according to the provided sort options. +pub fn validate_range_split_points( + split_points: &[SplitPoint], + sort_options: &[SortOptions], +) -> Result<()> { + let width = sort_options.len(); + for (idx, split_point) in split_points.iter().enumerate() { + let split_point_width = split_point.values().len(); + if split_point_width != width { + return _plan_err!( + "Range partitioning split point {idx} has width {split_point_width}, but ordering has width {width}" + ); + } + } + + for (idx, split_points) in split_points.windows(2).enumerate() { + if compare_rows( + split_points[0].values(), + split_points[1].values(), + sort_options, + )? != Ordering::Less + { + return _plan_err!( + "Range partitioning split points must be strictly ordered: split point {idx} ({}) must be less than split point {} ({})", + split_points[0], + idx + 1, + split_points[1] + ); + } + } + + Ok(()) +} diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index dd741ee6ff12e..190a08da12222 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -98,7 +98,7 @@ use datafusion_physical_expr::aggregate::{ }; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::{ - LexOrdering, PhysicalSortExpr, create_physical_sort_exprs, + LexOrdering, PhysicalSortExpr, RangePartitioning, create_physical_sort_exprs, }; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::empty::EmptyExec; @@ -1264,6 +1264,22 @@ impl DefaultPhysicalPlanner { .collect::>>()?; Partitioning::Hash(runtime_expr, *n) } + LogicalPartitioning::Range(range) => { + let sort_exprs = create_physical_sort_exprs( + range.ordering(), + input_dfschema, + execution_props, + )?; + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!( + "Range repartitioning requires non-empty ordering" + ) + })?; + Partitioning::Range(RangePartitioning::try_new( + ordering, + range.split_points().to_vec(), + )?) + } LogicalPartitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning" @@ -3245,8 +3261,8 @@ mod tests { use arrow_schema::{FieldRef, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::{ - DFSchemaRef, ScalarValue, TableReference, ToDFSchema as _, assert_batches_eq, - assert_contains, + DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _, + assert_batches_eq, assert_contains, }; use datafusion_execution::TaskContext; use datafusion_execution::runtime_env::RuntimeEnv; @@ -3255,8 +3271,8 @@ mod tests { use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder, - Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, - WindowFunctionDefinition, col, lit, + RangePartitioning, Signature, TableSource, UserDefinedLogicalNodeCore, + Volatility, WindowFunctionDefinition, col, lit, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; @@ -3304,6 +3320,46 @@ mod tests { Field::new(name, DataType::Int64, nullable) } + #[tokio::test] + async fn logical_range_repartition_plans_output_partitioning() -> Result<()> { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )])?; + let table = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]])?); + let source = Arc::new(DefaultTableSource::new(table)); + let logical_plan = LogicalPlanBuilder::scan("test", source, None)? + .repartition(LogicalPartitioning::Range(RangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?))? + .build()?; + + let planner = DefaultPhysicalPlanner::default(); + let physical_plan = planner + .create_initial_plan(&logical_plan, &make_session_state()) + .await?; + let repartition = physical_plan + .as_ref() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "expected RepartitionExec, got {}", + physical_plan.name() + ) + })?; + let Partitioning::Range(range) = repartition.partitioning() else { + return internal_err!( + "expected Range target partitioning, got {:?}", + repartition.partitioning() + ); + }; + assert_eq!(range.partition_count(), 2); + assert_eq!(physical_plan.output_partitioning().partition_count(), 2); + + Ok(()) + } + #[test] fn test_create_window_expr_unwraps_alias_with_metadata() -> Result<()> { use std::collections::HashMap; diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 58c7feb616179..27b86a6d8cdd5 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -515,6 +515,23 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Partitioning Key": hash_expr }) } + Partitioning::Range(range) => { + let range_expr: Vec = + range.ordering().iter().map(|e| format!("{e}")).collect(); + let split_points: Vec = range + .split_points() + .iter() + .map(|e| format!("{e}")) + .collect(); + + json!({ + "Node Type": "Repartition", + "Partitioning Scheme": "Range", + "Partition Count": range.partition_count(), + "Partitioning Key": range_expr, + "Split Points": split_points + }) + } Partitioning::DistributeBy(expr) => { let dist_by_expr: Vec = expr.iter().map(|e| format!("{e}")).collect(); diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index 5087b25178ab6..e0e51d7e470c3 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -41,9 +41,9 @@ pub use plan::{ Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, - RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery, - SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, Unnest, Values, - Window, projection_schema, + RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, + Subquery, SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, + Unnest, Values, Window, projection_schema, }; pub use statement::{ Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement, diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index b8843953865d2..3608c81878d17 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -51,6 +51,7 @@ use crate::{ }; use crate::statistics::StatisticsRequest; +use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datafusion_common::cse::{NormalizeEq, Normalizeable}; use datafusion_common::format::{ExplainAnalyzeCategories, ExplainFormat, MetricType}; @@ -61,10 +62,12 @@ use datafusion_common::tree_node::{ use datafusion_common::{ Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Dependency, FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues, Result, - ScalarValue, Spans, TableReference, UnnestOptions, aggregate_functional_dependencies, - assert_eq_or_internal_err, assert_or_internal_err, internal_err, plan_err, + ScalarValue, Spans, SplitPoint, TableReference, UnnestOptions, + aggregate_functional_dependencies, assert_eq_or_internal_err, assert_or_internal_err, + internal_err, plan_err, validate_range_split_points, }; use indexmap::IndexSet; +use itertools::Itertools as _; // backwards compatibility use crate::display::PgJsonVisitor; @@ -869,6 +872,32 @@ impl LogicalPlan { input: Arc::new(input), })) } + Partitioning::Range(range) => { + if expr.len() != range.ordering().len() { + return internal_err!( + "Incorrect number of expressions for Range partitioning" + ); + } + let input = self.only_input(inputs)?; + let ordering = range + .ordering() + .iter() + .zip(expr) + .map(|(sort_expr, expr)| SortExpr { + expr, + asc: sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }) + .collect(); + let range = RangePartitioning::try_new( + ordering, + range.split_points().to_vec(), + )?; + Ok(LogicalPlan::Repartition(Repartition { + partitioning_scheme: Partitioning::Range(range), + input: Arc::new(input), + })) + } Partitioning::DistributeBy(_) => { let input = self.only_input(inputs)?; Ok(LogicalPlan::Repartition(Repartition { @@ -2101,6 +2130,9 @@ impl LogicalPlan { n ) } + Partitioning::Range(range) => { + write!(f, "Repartition: {range}") + } Partitioning::DistributeBy(expr) => { let dist_by_expr: Vec = expr.iter().map(|e| format!("{e}")).collect(); @@ -4397,11 +4429,16 @@ impl Debug for Subquery { } } -/// Logical partitioning schemes supported by [`LogicalPlan::Repartition`] +/// Logical partitioning schemes. /// -/// See [`Partitioning`] for more details on partitioning +/// A scheme can describe either requested repartitioning in +/// [`LogicalPlan::Repartition`] or a partitioning property declared by a source. +/// Some schemes are only valid as metadata until planner support is added. /// -/// [`Partitioning`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/enum.Partitioning.html# +/// For physical execution partitioning, see +/// [`datafusion_physical_expr::Partitioning`]. +/// +/// [`datafusion_physical_expr::Partitioning`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/enum.Partitioning.html# #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub enum Partitioning { /// Allocate batches using a round-robin algorithm and the specified number of partitions @@ -4409,10 +4446,118 @@ pub enum Partitioning { /// Allocate rows based on a hash of one of more expressions and the specified number /// of partitions. Hash(Vec, usize), + /// Partition rows by ranges. + /// See [`RangePartitioning`] for the logical contract. + Range(RangePartitioning), /// The DISTRIBUTE BY clause is used to repartition the data based on the input expressions DistributeBy(Vec), } +impl Partitioning { + /// Return the number of partitions, if known. + pub fn partition_count(&self) -> Option { + match self { + Self::RoundRobinBatch(partition_count) | Self::Hash(_, partition_count) => { + Some(*partition_count) + } + Self::Range(range) => Some(range.partition_count()), + Self::DistributeBy(_) => None, + } + } +} + +/// Logical range partitioning. +/// +/// [`RangePartitioning`] describes an ordered logical key space with split points. +/// +/// - `ordering` defines the partitioning key and ordering using logical +/// [`SortExpr`]s. +/// - `split_points` define the boundaries between adjacent partitions. +/// +/// Comparisons use the lexicographic order defined by `ordering`, +/// including `ASC`/`DESC` and null ordering. Split points must be ordered +/// according to that ordering, and each split point must have one value per +/// ordering expression. See [`SplitPoint`] for the shared boundary contract. +/// +/// The expressions are resolved against the declaring plan's schema. This +/// constructor does not validate split point value types against the resolved +/// expression types. Like other user-specified data properties such as +/// sortedness, if a source declares range partitioning, it is responsible for +/// placing each row in the partition described by the split points. DataFusion +/// will not validate this is upheld. +/// +/// NOTE: Range-aware optimizer and execution behavior will be introduced +/// incrementally. See +/// . +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct RangePartitioning { + /// Ordered logical partitioning key. + ordering: Vec, + /// Boundaries between adjacent partitions. + split_points: Vec, +} + +impl RangePartitioning { + /// Creates logical range partitioning metadata and validates split point + /// shape and ordering. + pub fn try_new( + ordering: Vec, + split_points: Vec, + ) -> Result { + if ordering.is_empty() { + return plan_err!("Range partitioning requires non-empty ordering"); + } + + validate_range_split_points(&split_points, &logical_sort_options(&ordering))?; + + Ok(Self { + ordering, + split_points, + }) + } + + /// Return the number of partitions. + pub fn partition_count(&self) -> usize { + self.split_points.len() + 1 + } + + /// Returns the ordering that defines the range key. + pub fn ordering(&self) -> &[SortExpr] { + &self.ordering + } + + /// Returns the ordered split points between partitions. + pub fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } +} + +fn logical_sort_options(ordering: &[SortExpr]) -> Vec { + ordering + .iter() + .map(|sort_expr| SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }) + .collect() +} + +impl Display for RangePartitioning { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let ordering = self.ordering().iter().map(ToString::to_string).join(", "); + let split_points = self + .split_points() + .iter() + .map(ToString::to_string) + .join(", "); + write!( + f, + "Range([{ordering}], [{split_points}], {})", + self.partition_count() + ) + } +} + /// Represent the unnesting operation on a list column, such as the recursion depth and /// the output column name after unnesting /// @@ -4789,6 +4934,134 @@ mod tests { ]) } + fn i32_split_point(value: i32) -> SplitPoint { + SplitPoint::new(vec![ScalarValue::Int32(Some(value))]) + } + + fn null_i32_split_point() -> SplitPoint { + SplitPoint::new(vec![ScalarValue::Int32(None)]) + } + + #[test] + fn logical_range_partitioning_validates_shape() { + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10), i32_split_point(20)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + + let range = RangePartitioning::try_new( + vec![col("id").sort(false, true)], + vec![i32_split_point(20), i32_split_point(10)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + + let err = RangePartitioning::try_new(vec![], vec![]).unwrap_err(); + assert!(err.to_string().contains("non-empty ordering")); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true), col("salary").sort(true, true)], + vec![i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split point 0 has width 1, but ordering has width 2") + ); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(20), i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split points must be strictly ordered") + ); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10), i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split points must be strictly ordered") + ); + + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![null_i32_split_point(), i32_split_point(10)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + } + + #[test] + fn logical_partitioning_reports_known_partition_count() -> Result<()> { + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10)], + )?; + + assert_eq!(Partitioning::RoundRobinBatch(4).partition_count(), Some(4)); + assert_eq!( + Partitioning::Hash(vec![col("id")], 8).partition_count(), + Some(8) + ); + assert_eq!(Partitioning::Range(range).partition_count(), Some(2)); + assert_eq!( + Partitioning::DistributeBy(vec![col("id")]).partition_count(), + None + ); + + Ok(()) + } + + #[test] + fn logical_range_partitioning_participates_in_expression_rewrite() -> Result<()> { + let input = + table_scan(Some("employee_csv"), &employee_schema(), None)?.build()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(input), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10)], + )?), + }); + + let mut visited_exprs = vec![]; + plan.apply_expressions(|expr| { + visited_exprs.push(expr.to_string()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited_exprs, vec!["id"]); + + let plan = plan + .map_expressions(|expr| { + if expr == col("id") { + Ok(Transformed::yes(col("salary"))) + } else { + Ok(Transformed::no(expr)) + } + })? + .data; + + let LogicalPlan::Repartition(Repartition { + partitioning_scheme: Partitioning::Range(range), + .. + }) = plan + else { + unreachable!("expected range repartition"); + }; + assert_eq!(range.ordering()[0].expr, col("salary")); + assert_eq!(range.partition_count(), 2); + + Ok(()) + } + fn display_plan() -> Result { let plan1 = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3]))? .build()?; diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index e0cdec9e2c088..cba2dac24b610 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -37,6 +37,7 @@ //! * [`LogicalPlan::with_new_exprs`]: Create a new plan with different expressions //! * [`LogicalPlan::expressions`]: Return a copy of the plan's expressions +use crate::logical_plan::plan::RangePartitioning; use crate::{ Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit, @@ -427,6 +428,7 @@ impl LogicalPlan { Partitioning::Hash(expr, _) | Partitioning::DistributeBy(expr) => { expr.apply_elements(f) } + Partitioning::Range(range) => range.ordering().to_vec().apply_elements(f), Partitioning::RoundRobinBatch(_) => Ok(TreeNodeRecursion::Continue), }, LogicalPlan::Window(Window { window_expr, .. }) => { @@ -532,6 +534,19 @@ impl LogicalPlan { Partitioning::DistributeBy(expr) => expr .map_elements(f)? .update_data(Partitioning::DistributeBy), + Partitioning::Range(range) => { + let split_points = range.split_points().to_vec(); + range + .ordering() + .to_vec() + .map_elements(f)? + .map_data(|ordering| { + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) + })? + } Partitioning::RoundRobinBatch(_) => Transformed::no(partitioning_scheme), } .update_data(|partitioning_scheme| { diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index c82d1c64dd0d9..b55bd70bdf185 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -58,11 +58,12 @@ pub mod execution_props { pub use aggregate::groups_accumulator::{GroupsAccumulatorAdapter, NullState}; pub use analysis::{AnalysisContext, ExprBoundaries, analyze}; +pub use datafusion_common::SplitPoint; pub use equivalence::{ AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union, }; pub use expressions::{DynamicFilterTracker, DynamicFilterTracking}; -pub use partitioning::{Distribution, Partitioning, RangePartitioning, SplitPoint}; +pub use partitioning::{Distribution, Partitioning, RangePartitioning}; pub use physical_expr::{ add_offset_to_expr, add_offset_to_physical_sort_exprs, create_lex_ordering, create_ordering, create_physical_sort_expr, create_physical_sort_exprs, diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 6009cd995e18c..2e0aaaf3fb4b7 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -21,10 +21,10 @@ use crate::{ EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, expressions::UnKnownColumn, physical_exprs_equal, }; -use datafusion_common::{Result, ScalarValue, plan_err}; +pub use datafusion_common::SplitPoint; +use datafusion_common::{Result, validate_range_split_points}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; -use std::cmp::Ordering; use std::fmt; use std::fmt::Display; use std::sync::Arc; @@ -156,20 +156,7 @@ impl Display for Partitioning { /// Comparisons use the lexicographic order defined by `ordering`, including /// `ASC`/`DESC` and null ordering. Split points must be strictly ordered /// according to that ordering, and each split point must have one value per -/// ordering expression. -/// -/// `N` split points define `N + 1` partitions: -/// -/// ```text -/// partition 0: key < split_points[0] -/// partition 1: split_points[0] <= key < split_points[1] -/// ... -/// partition N - 1: split_points[N - 2] <= key < split_points[N - 1] -/// partition N: split_points[N - 1] <= key -/// ``` -/// -/// Values equal to split point `i` belong to partition `i + 1`, so interior -/// partitions are lower-inclusive and upper-exclusive. +/// ordering expression. See [`SplitPoint`] for the shared boundary convention. /// /// Like other user-specified data properties such as sortedness, if a source /// declares range partitioning, it is responsible for placing each row in the @@ -217,39 +204,6 @@ pub struct RangePartitioning { split_points: Vec, } -/// A boundary between adjacent range partitions. -/// -/// A split point is a tuple with one [`ScalarValue`] per sort expression in the -/// parent [`RangePartitioning`] ordering. -#[derive(Debug, Clone, PartialEq)] -pub struct SplitPoint { - values: Vec, -} - -impl SplitPoint { - /// Creates a new split point from its tuple values. - pub fn new(values: Vec) -> Self { - Self { values } - } - - /// Returns the tuple values for this split point. - pub fn values(&self) -> &[ScalarValue] { - &self.values - } -} - -impl Display for SplitPoint { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let values = self - .values - .iter() - .map(ToString::to_string) - .collect::>() - .join(", "); - write!(f, "({values})") - } -} - impl RangePartitioning { /// Creates range partitioning metadata without validating split points. /// @@ -265,7 +219,13 @@ impl RangePartitioning { /// Creates range partitioning metadata and validates split point shape and /// ordering. pub fn try_new(ordering: LexOrdering, split_points: Vec) -> Result { - validate_range_split_points(&ordering, &split_points)?; + validate_range_split_points( + &split_points, + &ordering + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(), + )?; Ok(Self::new(ordering, split_points)) } @@ -384,86 +344,6 @@ fn format_range_split_points(split_points: &[SplitPoint]) -> String { .join(", ") } -fn validate_range_split_points( - ordering: &LexOrdering, - split_points: &[SplitPoint], -) -> Result<()> { - let width = ordering.len(); - for (idx, split_point) in split_points.iter().enumerate() { - let split_point_width = split_point.values.len(); - if split_point_width != width { - return plan_err!( - "Range partitioning split point {idx} has width {split_point_width}, but ordering has width {width}" - ); - } - } - - for (idx, split_points) in split_points.windows(2).enumerate() { - if compare_split_points(ordering, &split_points[0], &split_points[1])? - != Ordering::Less - { - return plan_err!( - "Range partitioning split points must be strictly ordered: split point {idx} ({}) must be less than split point {} ({})", - split_points[0], - idx + 1, - split_points[1] - ); - } - } - - Ok(()) -} - -fn compare_split_points( - ordering: &LexOrdering, - left: &SplitPoint, - right: &SplitPoint, -) -> Result { - for ((left_value, right_value), sort_expr) in - left.values.iter().zip(&right.values).zip(ordering.iter()) - { - let value_ordering = - compare_scalar_values_for_sort(left_value, right_value, sort_expr)?; - if value_ordering != Ordering::Equal { - return Ok(value_ordering); - } - } - - Ok(Ordering::Equal) -} - -fn compare_scalar_values_for_sort( - left: &ScalarValue, - right: &ScalarValue, - sort_expr: &PhysicalSortExpr, -) -> Result { - match (left.is_null(), right.is_null()) { - (true, true) => Ok(Ordering::Equal), - (true, false) => Ok(if sort_expr.options.nulls_first { - Ordering::Less - } else { - Ordering::Greater - }), - (false, true) => Ok(if sort_expr.options.nulls_first { - Ordering::Greater - } else { - Ordering::Less - }), - (false, false) => { - let Some(ordering) = left.partial_cmp(right) else { - return plan_err!( - "Range partitioning split point values are not comparable: {left:?} and {right:?}" - ); - }; - Ok(if sort_expr.options.descending { - ordering.reverse() - } else { - ordering - }) - } - } -} - fn equivalent_exprs( left: &[Arc], right: &[Arc], @@ -754,7 +634,7 @@ mod tests { use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; - use datafusion_common::Result; + use datafusion_common::{Result, ScalarValue}; struct PartitioningTestFixture { schema: SchemaRef, diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index b691441e95a97..35c2e76d880b9 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -1755,6 +1755,11 @@ impl AsLogicalPlan for LogicalPlanNode { Partitioning::RoundRobinBatch(partition_count) => { PartitionMethod::RoundRobin(*partition_count as u64) } + Partitioning::Range(_) => { + // TODO: Support range repartition protobuf serialization. + // Tracked by https://github.com/apache/datafusion/issues/22787 + return not_impl_err!("Range repartition"); + } Partitioning::DistributeBy(_) => { return not_impl_err!("DistributeBy"); } diff --git a/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs index 50c4b3da86cbe..1b9e91c7c475a 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs @@ -32,6 +32,11 @@ pub fn from_repartition( let partition_count = match repartition.partitioning_scheme { Partitioning::RoundRobinBatch(num) => num, Partitioning::Hash(_, num) => num, + Partitioning::Range(_) => { + // TODO: Support range repartitioning in Substrait exchange output. + // Tracked by https://github.com/apache/datafusion/issues/22788 + return not_impl_err!("Substrait does not support Range repartitioning"); + } Partitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning" @@ -50,6 +55,11 @@ pub fn from_repartition( .collect::>>()?; ExchangeKind::ScatterByFields(ScatterFields { fields }) } + Partitioning::Range(_) => { + // TODO: Support range repartitioning in Substrait exchange output. + // Tracked by https://github.com/apache/datafusion/issues/22788 + return not_impl_err!("Substrait does not support Range repartitioning"); + } Partitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning"