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/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 { 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/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/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/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()); 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" 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": {