diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index 0389b3cb17fe9..edb3294051cde 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -325,7 +325,7 @@ pub fn evaluate_partition_prefix<'a>( } } -fn filter_partitions( +pub fn filter_partitioned_file( pf: PartitionedFile, filters: &[Expr], df_schema: &DFSchema, @@ -447,7 +447,7 @@ pub async fn pruned_partition_list<'a>( )) }) .try_filter_map(move |pf| { - futures::future::ready(filter_partitions(pf, filters, &df_schema)) + futures::future::ready(filter_partitioned_file(pf, filters, &df_schema)) }) .boxed()) } diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 0ab15e05abba1..65444706dd02f 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -21,7 +21,7 @@ use datafusion_common::plan_err; use datafusion_datasource::ListingTableUrl; use datafusion_datasource::file_format::FileFormat; use datafusion_execution::config::SessionConfig; -use datafusion_expr::SortExpr; +use datafusion_expr::{Partitioning, SortExpr}; use futures::StreamExt; use futures::TryStreamExt; use itertools::Itertools; @@ -61,6 +61,46 @@ pub struct ListingOptions { /// multiple equivalent orderings, the outer `Vec` will have a /// single element. pub file_sort_order: Vec>, + /// Declared output partitioning for scans from this table. + /// + /// Expressions are logical expressions over the full table schema. When set, + /// [`ListingTable`](crate::ListingTable) creates one file group per + /// declared output partition. When unset, file grouping uses the scan-time + /// [`SessionConfig::target_partitions`](datafusion_execution::config::SessionConfig::target_partitions). + /// + /// Files are listed in path order, split into whole-file groups across the + /// declared partition count, and then padded with trailing empty groups when + /// needed. DataFusion does not route files by partition values or validate + /// row placement, so callers must ensure file group `i` contains rows for + /// partition `i`. Layouts that require explicit file-to-partition assignment + /// are not supported. + /// + /// For example, range partitioning on column `a` with split points + /// `[10, 20, 30]` declares four output partitions. With three path-ordered + /// files, the trailing partition is preserved as empty: + /// + /// ```text + /// files in path order: f0, f1, f2 + /// + /// file groups: + /// partition 0: [f0] + /// partition 1: [f1] + /// partition 2: [f2] + /// partition 3: [] + /// ``` + /// + /// With five path-ordered files, a partition can contain multiple files: + /// + /// ```text + /// files in path order: f0, f1, f2, f3, f4 + /// + /// file groups: + /// partition 0: [f0, f1] + /// partition 1: [f2, f3] + /// partition 2: [f4] + /// partition 3: [] + /// ``` + pub output_partitioning: Option, } impl ListingOptions { @@ -78,6 +118,7 @@ impl ListingOptions { collect_stat: false, target_partitions: 1, file_sort_order: vec![], + output_partitioning: None, } } @@ -136,6 +177,17 @@ impl ListingOptions { self } + /// Set declared output partitioning. + /// + /// See [`Self::output_partitioning`] for the contract. + pub fn with_output_partitioning( + mut self, + output_partitioning: Option, + ) -> Self { + self.output_partitioning = output_partitioning; + self + } + /// Set `table partition columns` on [`ListingOptions`] and returns self. /// /// "partition columns," used to support [Hive Partitioning], are diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 7ee743a6abe71..3627e560a5096 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -16,14 +16,17 @@ // under the License. use crate::config::SchemaSource; -use crate::helpers::{expr_applicable_for_cols, pruned_partition_list}; +use crate::helpers::{ + expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list, +}; use crate::{ListingOptions, ListingTableConfig}; use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef}; use async_trait::async_trait; use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider}; use datafusion_common::stats::Precision; use datafusion_common::{ - Constraints, SchemaExt, Statistics, internal_datafusion_err, plan_err, project_schema, + Constraints, DFSchema, SchemaExt, Statistics, internal_datafusion_err, plan_err, + project_schema, }; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; @@ -38,8 +41,10 @@ use datafusion_execution::cache::TableScopedPath; use datafusion_execution::cache::cache_manager::FileStatisticsCache; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::ExecutionProps; -use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; -use datafusion_physical_expr::create_lex_ordering; +use datafusion_expr::{ + Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType, +}; +use datafusion_physical_expr::{create_lex_ordering, create_physical_partitioning}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::ExecutionPlan; @@ -433,6 +438,20 @@ fn derive_common_ordering_from_files(file_groups: &[FileGroup]) -> Option datafusion_common::Result { + let files = file_group + .into_inner() + .into_iter() + .map(|file| filter_partitioned_file(file, filters, df_schema)) + .filter_map(Result::transpose) + .collect::>>()?; + Ok(FileGroup::new(files)) +} + // Expressions can be used for partition pruning if they can be evaluated using // only the partition columns and there are partition columns. fn can_be_evaluated_for_partition_pruning( @@ -500,9 +519,19 @@ impl TableProvider for ListingTable { can_be_evaluated_for_partition_pruning(&table_partition_col_names, filter) }); - // We should not limit the number of partitioned files to scan if there are filters and limit - // at the same time. This is because the limit should be applied after the filters are applied. - let statistic_file_limit = if filters.is_empty() { limit } else { None }; + let declared_output_partitioning = self.options.output_partitioning.as_ref(); + + // We should not limit files before assigning declared output partitions + // or before applying non-partition filters. + let statistic_file_limit = + if filters.is_empty() && declared_output_partitioning.is_none() { + limit + } else { + None + }; + let file_group_count = declared_output_partitioning + .and_then(LogicalPartitioning::partition_count) + .unwrap_or_else(|| state.config().target_partitions()); let ListFilesResult { file_groups: mut partitioned_file_lists, @@ -522,17 +551,19 @@ impl TableProvider for ListingTable { state.execution_props(), &partitioned_file_lists, )?; - match state - .config_options() - .execution - .split_file_groups_by_statistics + let split_file_groups_by_statistics = declared_output_partitioning.is_none() + && state + .config_options() + .execution + .split_file_groups_by_statistics; + match split_file_groups_by_statistics .then(|| { output_ordering.first().map(|output_ordering| { FileScanConfig::split_groups_by_statistics_with_target_partitions( &self.table_schema, &partitioned_file_lists, output_ordering, - self.options.target_partitions, + file_group_count, ) }) }) @@ -540,7 +571,7 @@ impl TableProvider for ListingTable { { Some(Err(e)) => log::debug!("failed to split file groups by statistics: {e}"), Some(Ok(new_groups)) => { - if new_groups.len() <= self.options.target_partitions { + if new_groups.len() <= file_group_count { partitioned_file_lists = new_groups; } else { log::debug!( @@ -551,6 +582,41 @@ impl TableProvider for ListingTable { None => {} // no ordering required }; + let output_partitioning = if let Some(output_partitioning) = + declared_output_partitioning + { + let output_partitioning = match output_partitioning { + LogicalPartitioning::RoundRobinBatch(_) => { + return datafusion_common::not_impl_err!( + "RoundRobinBatch output partitioning is not supported for ListingTable" + ); + } + LogicalPartitioning::DistributeBy(_) => { + return datafusion_common::not_impl_err!( + "DistributeBy output partitioning is not supported for ListingTable" + ); + } + LogicalPartitioning::Hash(_, _) | LogicalPartitioning::Range(_) => { + let df_schema = DFSchema::try_from(Arc::clone(&self.table_schema))?; + create_physical_partitioning( + output_partitioning, + &df_schema, + state.execution_props(), + )? + } + }; + let partition_count = output_partitioning.partition_count(); + if partitioned_file_lists.len() != partition_count { + return plan_err!( + "ListingTable output_partitioning has {partition_count} partitions, but the scan has {} file groups", + partitioned_file_lists.len() + ); + } + Some(output_partitioning) + } else { + None + }; + let Some(object_store_url) = self.table_paths.first().map(ListingTableUrl::object_store) else { @@ -560,24 +626,23 @@ impl TableProvider for ListingTable { }; let file_source = self.create_file_source(); + let scan_config = FileScanConfigBuilder::new(object_store_url, file_source) + .with_file_groups(partitioned_file_lists) + .with_constraints(self.constraints.clone()) + .with_statistics(statistics) + .with_projection_indices(projection)? + .with_limit(limit) + .with_output_ordering(output_ordering) + .with_output_partitioning(output_partitioning) + .with_expr_adapter(self.expr_adapter_factory.clone()) + .with_partitioned_by_file_group(partitioned_by_file_group) + .build(); // create the execution plan let plan = self .options .format - .create_physical_plan( - state, - FileScanConfigBuilder::new(object_store_url, file_source) - .with_file_groups(partitioned_file_lists) - .with_constraints(self.constraints.clone()) - .with_statistics(statistics) - .with_projection_indices(projection)? - .with_limit(limit) - .with_output_ordering(output_ordering) - .with_expr_adapter(self.expr_adapter_factory.clone()) - .with_partitioned_by_file_group(partitioned_by_file_group) - .build(), - ) + .create_physical_plan(state, scan_config) .await?; Ok(ScanResult::new(plan)) @@ -689,28 +754,41 @@ impl ListingTable { /// Get the list of files for a scan as well as the file level statistics. /// The list is grouped to let the execution plan know how the files should /// be distributed to different threads / executors. + /// + /// If [`ListingOptions::output_partitioning`] is set, returns one file + /// group per declared partition, including empty trailing groups. pub async fn list_files_for_scan<'a>( &'a self, ctx: &'a dyn Session, filters: &'a [Expr], limit: Option, ) -> datafusion_common::Result { - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? + if let Some(output_partitioning) = self.options.output_partitioning.as_ref() { + self.list_files_for_declared_output_partitioning( + ctx, + output_partitioning, + filters, + ) + .await } else { - return Ok(ListFilesResult { - file_groups: vec![], - statistics: Statistics::new_unknown(&self.file_schema), - grouped_by_partition: false, - }); - }; + self.list_files_for_regular_scan(ctx, filters, limit).await + } + } + + async fn collect_files_for_scan<'a>( + &'a self, + ctx: &'a dyn Session, + store: &'a Arc, + listing_time_filters: &'a [Expr], + file_limit: Option, + ) -> datafusion_common::Result<(FileGroup, bool)> { // list files (with partitions) let file_list = future::try_join_all(self.table_paths.iter().map(|table_path| { pruned_partition_list( ctx, store.as_ref(), table_path, - filters, + listing_time_filters, &self.options.file_extension, &self.options.table_partition_cols, ) @@ -724,7 +802,7 @@ impl ListingTable { .map(|part_file| async { let part_file = part_file?; let (statistics, ordering) = if self.options.collect_stat { - self.do_collect_statistics_and_ordering(ctx, &store, &part_file) + self.do_collect_statistics_and_ordering(ctx, store, &part_file) .await? } else { (Arc::new(Statistics::new_unknown(&self.file_schema)), None) @@ -736,8 +814,34 @@ impl ListingTable { .boxed() .buffer_unordered(ctx.config_options().execution.meta_fetch_concurrency); - let (file_group, inexact_stats) = - get_files_with_limit(files, limit, self.options.collect_stat).await?; + get_files_with_limit(files, file_limit, self.options.collect_stat).await + } + + async fn list_files_for_regular_scan<'a>( + &'a self, + ctx: &'a dyn Session, + filters: &'a [Expr], + limit: Option, + ) -> datafusion_common::Result { + let file_group_count = self.options.target_partitions; + if file_group_count == 0 { + return plan_err!( + "ListingTable requires target_partitions to be greater than zero" + ); + } + + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? + } else { + return Ok(ListFilesResult { + file_groups: vec![], + statistics: Statistics::new_unknown(&self.file_schema), + grouped_by_partition: false, + }); + }; + let (file_group, inexact_stats) = self + .collect_files_for_scan(ctx, &store, filters, limit) + .await?; // Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N // @@ -746,28 +850,102 @@ impl ListingTable { // hash repartitioning for aggregates and joins on partition columns. let threshold = ctx.config_options().optimizer.preserve_file_partitions; - let (file_groups, grouped_by_partition) = if threshold > 0 - && !self.options.table_partition_cols.is_empty() - { - let grouped = - file_group.group_by_partition_values(self.options.target_partitions); - if grouped.len() >= threshold { - (grouped, true) + let (file_groups, grouped_by_partition) = + if threshold > 0 && !self.options.table_partition_cols.is_empty() { + let grouped = file_group.group_by_partition_values(file_group_count); + if grouped.len() >= threshold { + (grouped, true) + } else { + let all_files: Vec<_> = + grouped.into_iter().flat_map(|g| g.into_inner()).collect(); + ( + FileGroup::new(all_files).split_files(file_group_count), + false, + ) + } } else { - let all_files: Vec<_> = - grouped.into_iter().flat_map(|g| g.into_inner()).collect(); - ( - FileGroup::new(all_files).split_files(self.options.target_partitions), - false, - ) - } + (file_group.split_files(file_group_count), false) + }; + + self.list_files_result_from_groups( + ctx, + file_groups, + inexact_stats, + grouped_by_partition, + ) + } + + async fn list_files_for_declared_output_partitioning<'a>( + &'a self, + ctx: &'a dyn Session, + output_partitioning: &LogicalPartitioning, + filters: &'a [Expr], + ) -> datafusion_common::Result { + let Some(file_group_count) = output_partitioning.partition_count() else { + return datafusion_common::not_impl_err!( + "DistributeBy output partitioning is not supported for ListingTable" + ); + }; + if file_group_count == 0 { + return plan_err!( + "ListingTable output_partitioning requires at least one partition" + ); + } + + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? } else { - ( - file_group.split_files(self.options.target_partitions), - false, - ) + return Ok(ListFilesResult { + file_groups: vec![], + statistics: Statistics::new_unknown(&self.file_schema), + grouped_by_partition: false, + }); }; + let (file_group, inexact_stats) = + self.collect_files_for_scan(ctx, &store, &[], None).await?; + let mut file_groups = file_group.split_files(file_group_count); + if !file_groups.is_empty() { + file_groups.resize_with(file_group_count, || FileGroup::new(vec![])); + } + let file_groups = + self.filter_declared_file_groups_by_partition_filters(file_groups, filters)?; + self.list_files_result_from_groups(ctx, file_groups, inexact_stats, false) + } + + fn filter_declared_file_groups_by_partition_filters( + &self, + file_groups: Vec, + filters: &[Expr], + ) -> datafusion_common::Result> { + if filters.is_empty() { + return Ok(file_groups); + } + + let df_schema = DFSchema::from_unqualified_fields( + self.options + .table_partition_cols + .iter() + .map(|(name, data_type)| Field::new(name, data_type.clone(), true)) + .collect(), + Default::default(), + )?; + + file_groups + .into_iter() + .map(|file_group| { + filter_file_group_by_partition_filters(file_group, filters, &df_schema) + }) + .collect::>>() + } + + fn list_files_result_from_groups( + &self, + _ctx: &dyn Session, + file_groups: Vec, + inexact_stats: bool, + grouped_by_partition: bool, + ) -> datafusion_common::Result { let (file_groups, stats) = compute_all_files_statistics( file_groups, self.schema(), diff --git a/datafusion/catalog/src/streaming.rs b/datafusion/catalog/src/streaming.rs index e609877c2b778..5bfecef1fb2ed 100644 --- a/datafusion/catalog/src/streaming.rs +++ b/datafusion/catalog/src/streaming.rs @@ -24,7 +24,10 @@ use async_trait::async_trait; use datafusion_common::{DFSchema, Result, plan_err}; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::equivalence::project_ordering; -use datafusion_physical_expr::{LexOrdering, create_physical_sort_exprs}; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, create_physical_sort_exprs, +}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use log::debug; @@ -38,6 +41,7 @@ pub struct StreamingTable { partitions: Vec>, infinite: bool, sort_order: Vec, + output_partitioning: Option, } impl StreamingTable { @@ -62,6 +66,7 @@ impl StreamingTable { partitions, infinite: false, sort_order: vec![], + output_partitioning: None, }) } @@ -76,6 +81,33 @@ impl StreamingTable { self.sort_order = sort_order; self } + + /// Declares the output partitioning of this streaming table. + /// + /// The partitioning expressions refer to the table schema before scan + /// projection. If a scan projection removes a partitioning expression, the + /// physical plan reports unknown partitioning. + pub fn with_output_partitioning(mut self, output_partitioning: Partitioning) -> Self { + self.output_partitioning = Some(output_partitioning); + self + } + + fn output_partitioning( + &self, + projection: Option<&Vec>, + ) -> Result { + let Some(output_partitioning) = &self.output_partitioning else { + return Ok(Partitioning::UnknownPartitioning(self.partitions.len())); + }; + let Some(projection) = projection else { + return Ok(output_partitioning.clone()); + }; + + let projection_mapping = + ProjectionMapping::from_indices(projection, &self.schema)?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&self.schema)); + Ok(output_partitioning.project(&projection_mapping, &eq_properties)) + } } #[async_trait] @@ -119,13 +151,16 @@ impl TableProvider for StreamingTable { vec![] }; - Ok(Arc::new(StreamingTableExec::try_new( + let exec = StreamingTableExec::try_new( Arc::clone(&self.schema), self.partitions.clone(), projection, LexOrdering::new(physical_sort), self.infinite, limit, - )?)) + )? + .with_output_partitioning(self.output_partitioning(projection)?)?; + + Ok(Arc::new(exec)) } } 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/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index d14ec1f56dce2..0d6231ce27833 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -125,7 +125,7 @@ mod tests { }, }; use arrow::{compute::SortOptions, record_batch::RecordBatch}; - use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_catalog::TableProvider; use datafusion_catalog_listing::{ ListingOptions, ListingTable, ListingTableConfig, SchemaSource, @@ -139,12 +139,17 @@ mod tests { use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::FileFormat; use datafusion_expr::dml::InsertOp; - use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator}; + use datafusion_expr::{ + BinaryExpr, LogicalPlanBuilder, Operator, Partitioning as LogicalPartitioning, + RangePartitioning as LogicalRangePartitioning, + }; use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::expressions::binary; + use datafusion_physical_expr::expressions::{Column, binary}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::empty::EmptyExec; - use datafusion_physical_plan::{ExecutionPlanProperties, collect}; + use datafusion_physical_plan::{ + ExecutionPlanProperties, Partitioning, RangePartitioning, SplitPoint, collect, + }; use std::collections::HashMap; use std::io::Write; use std::sync::Arc; @@ -177,6 +182,21 @@ mod tests { .collect() } + fn listing_table_with_files( + ctx: &SessionContext, + files: &[&str], + table_path: &str, + options: ListingOptions, + schema: Schema, + ) -> Result { + register_test_store(ctx, &files.iter().map(|f| (*f, 10)).collect::>()); + + let config = ListingTableConfig::new(ListingTableUrl::parse(table_path)?) + .with_listing_options(options) + .with_schema(Arc::new(schema)); + ListingTable::try_new(config) + } + #[tokio::test] async fn test_schema_source_tracking_comprehensive() -> Result<()> { let ctx = SessionContext::new(); @@ -1286,6 +1306,236 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_list_files_uses_declared_output_partitioning_count() -> Result<()> { + let files = ["bucket/key-prefix/file0", "bucket/key-prefix/file1"]; + + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + ); + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(LogicalPartitioning::Range( + LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::from(10i32)]), + SplitPoint::new(vec![ScalarValue::from(20i32)]), + SplitPoint::new(vec![ScalarValue::from(30i32)]), + ], + )?, + ))); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new("a", DataType::Int32, false)]), + )?; + + let result = table.list_files_for_scan(&ctx.state(), &[], None).await?; + let group_sizes = result + .file_groups + .iter() + .map(|group| group.len()) + .collect::>(); + + assert_eq!(group_sizes, vec![1, 1, 0, 0]); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_normalizes_split_point_types() -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(123_000_000_000), + None, + )])], + )?); + let expected_output_partitioning = + Partitioning::Range(RangePartitioning::try_new( + LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions::default(), + )]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::TimestampSecond( + Some(123), + None, + )])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]), + )?; + + let scan = table.scan(&ctx.state(), None, &[], None).await?; + assert_eq!(scan.output_partitioning(), &expected_output_partitioning); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_rejects_invalid_split_point_type() + -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "not-an-int".to_string(), + ))])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new("a", DataType::Int32, false)]), + )?; + + let err = table.scan(&ctx.state(), None, &[], None).await.unwrap_err(); + assert_contains!( + err.to_string(), + "Range output partitioning split point 0 value 0 with type Utf8 cannot be represented exactly as ordering expression type Int32" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_rejects_lossy_timestamp_split_point() + -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(123_456), + None, + )])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]), + )?; + + let err = table.scan(&ctx.state(), None, &[], None).await.unwrap_err(); + assert_contains!( + err.to_string(), + "Range output partitioning split point 0 value 0 with type Timestamp(ns) cannot be represented exactly as ordering expression type Timestamp(s)" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partition_filter_preserves_declared_output_partitioning() -> Result<()> + { + let files = ["bucket/test/pid=1/file1", "bucket/test/pid=2/file2"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("pid").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::from(2i32)])], + )?); + let expected_output_partitioning = + Partitioning::Range(RangePartitioning::try_new( + LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("pid", 1)), + SortOptions::default(), + )]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::from(2i32)])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_table_partition_cols(vec![("pid".to_string(), DataType::Int32)]) + .with_output_partitioning(Some(output_partitioning.clone())); + + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/test/", + opt, + Schema::new(vec![Field::new("a", DataType::Boolean, false)]), + )?; + + let unfiltered = table.scan(&ctx.state(), None, &[], None).await?; + assert_eq!( + unfiltered.output_partitioning(), + &expected_output_partitioning + ); + + let filter = Expr::eq(col("pid"), lit(2_i32)); + let file_groups = table + .list_files_for_scan(&ctx.state(), std::slice::from_ref(&filter), None) + .await? + .file_groups + .into_iter() + .map(|group| { + group + .into_inner() + .into_iter() + .map(|file| file.path().to_string()) + .collect::>() + }) + .collect::>(); + assert_eq!( + file_groups, + vec![ + Vec::::new(), + vec!["bucket/test/pid=2/file2".to_string()] + ] + ); + + let filtered = table.scan(&ctx.state(), None, &[filter], None).await?; + assert_eq!( + filtered.output_partitioning(), + &expected_output_partitioning + ); + + Ok(()) + } + #[tokio::test] async fn test_listing_table_prunes_extra_files_in_hive() -> Result<()> { let files = [ diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 43af743caa030..2a41d83aac957 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -32,10 +32,11 @@ use crate::logical_expr::{ Aggregate, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, Values, Window, }; use crate::logical_expr::{ - Expr, LogicalPlan, Partitioning as LogicalPartitioning, PlanType, Repartition, - UserDefinedLogicalNode, + Expr, LogicalPlan, PlanType, Repartition, UserDefinedLogicalNode, +}; +use crate::physical_expr::{ + create_physical_expr, create_physical_exprs, create_physical_partitioning, }; -use crate::physical_expr::{create_physical_expr, create_physical_exprs}; use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use crate::physical_plan::analyze::AnalyzeExec; use crate::physical_plan::explain::ExplainExec; @@ -52,8 +53,8 @@ use crate::physical_plan::union::UnionExec; use crate::physical_plan::unnest::UnnestExec; use crate::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use crate::physical_plan::{ - ExecutionPlan, ExecutionPlanProperties, InputOrderMode, Partitioning, PhysicalExpr, - WindowExpr, displayable, windows, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PhysicalExpr, WindowExpr, + displayable, windows, }; use crate::schema_equivalence::schema_satisfied_by; @@ -1247,25 +1248,11 @@ impl DefaultPhysicalPlanner { }) => { let physical_input = children.one()?; let input_dfschema = input.as_ref().schema(); - let physical_partitioning = match partitioning_scheme { - LogicalPartitioning::RoundRobinBatch(n) => { - Partitioning::RoundRobinBatch(*n) - } - LogicalPartitioning::Hash(expr, n) => { - let runtime_expr = expr - .iter() - .map(|e| { - create_physical_expr(e, input_dfschema, execution_props) - }) - .collect::>>()?; - Partitioning::Hash(runtime_expr, *n) - } - LogicalPartitioning::DistributeBy(_) => { - return not_impl_err!( - "Physical plan does not support DistributeBy partitioning" - ); - } - }; + let physical_partitioning = create_physical_partitioning( + partitioning_scheme, + input_dfschema, + execution_props, + )?; Arc::new(RepartitionExec::try_new( physical_input, physical_partitioning, @@ -3219,8 +3206,8 @@ mod tests { use crate::datasource::MemTable; use crate::datasource::file_format::options::CsvReadOptions; use crate::physical_plan::{ - DisplayAs, DisplayFormatType, PlanProperties, SendableRecordBatchStream, - expressions, + DisplayAs, DisplayFormatType, Partitioning, PlanProperties, + SendableRecordBatchStream, expressions, }; use crate::prelude::{SessionConfig, SessionContext}; use crate::test_util::{scan_empty, scan_empty_with_partitions}; @@ -3231,8 +3218,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; @@ -3241,8 +3228,8 @@ mod tests { use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder, - Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, - WindowFunctionDefinition, col, lit, + Partitioning as LogicalPartitioning, RangePartitioning, Signature, TableSource, + UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, lit, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; @@ -3290,6 +3277,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/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 12abf79041091..54b80644a31a3 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -20,8 +20,8 @@ use std::ops::Deref; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - check_integrity, coalesce_partitions_exec, parquet_exec_with_sort, - parquet_exec_with_stats, repartition_exec, schema, sort_exec, + bounded_window_exec_with_can_repartition, check_integrity, coalesce_partitions_exec, + parquet_exec_with_sort, parquet_exec_with_stats, repartition_exec, schema, sort_exec, sort_exec_with_preserve_partitioning, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; @@ -59,9 +59,13 @@ use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; -use datafusion_physical_expr::Distribution; +use datafusion_physical_expr::{ + Distribution, EquivalenceProperties, RangePartitioning, SplitPoint, +}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion_physical_plan::execution_plan::ExecutionPlan; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, ExecutionPlan, +}; use datafusion_physical_plan::expressions::col; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::JoinOn; @@ -70,7 +74,8 @@ use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, displayable, + DisplayAs, DisplayFormatType, ExecutionPlanProperties, Partitioning, PlanProperties, + displayable, }; use insta::Settings; @@ -263,8 +268,12 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { vec![&self.input] } - fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + ]) } fn maintains_input_order(&self) -> Vec { @@ -325,6 +334,104 @@ fn parquet_exec_multiple_sorted( DataSourceExec::from_data_source(config) } +#[derive(Debug)] +struct PartitionedTestExec { + cache: Arc, +} + +impl PartitionedTestExec { + fn new(output_partitioning: Partitioning) -> Self { + Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema()), + output_partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } +} + +impl DisplayAs for PartitionedTestExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => write!( + f, + "PartitionedTestExec: output_partitioning={}", + self.cache.output_partitioning() + ), + DisplayFormatType::TreeRender => write!(f, ""), + } + } +} + +impl ExecutionPlan for PartitionedTestExec { + fn name(&self) -> &'static str { + "PartitionedTestExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } +} + +fn parquet_exec_with_output_partitioning( + output_partitioning: Partitioning, +) -> Arc { + Arc::new(PartitionedTestExec::new(output_partitioning)) +} + +fn range_partitioning( + key: &str, + split_points: impl IntoIterator, + options: SortOptions, +) -> Result { + let ordering = [PhysicalSortExpr { + expr: col(key, &schema())?, + options, + }] + .into(); + let split_points = split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect(); + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) +} + +fn hash_partitioning(key: &str, partition_count: usize) -> Result { + Ok(Partitioning::Hash( + vec![col(key, &schema())?], + partition_count, + )) +} + fn csv_exec() -> Arc { csv_exec_with_sort(vec![]) } @@ -701,6 +808,553 @@ impl TestConfig { } } +fn partitioned_join_plan( + left_partitioning: Partitioning, + right_partitioning: Partitioning, + join_type: JoinType, +) -> Arc { + let left = parquet_exec_with_output_partitioning(left_partitioning); + let right = parquet_exec_with_output_partitioning(right_partitioning); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &right.schema()).unwrap()) as _, + )]; + hash_join_exec(left, right, &join_on, &join_type) +} + +#[test] +fn inner_range_join_keeps_range_partitioning() -> Result<()> { + let join = partitioned_join_plan( + range_partitioning("a", [10, 20], SortOptions::default())?, + range_partitioning("b", [10, 20], SortOptions::default())?, + JoinType::Inner, + ); + + let plan = TestConfig::default() + .with_query_execution_partitions(3) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + PartitionedTestExec: output_partitioning=Range([b@1 ASC], [(10), (20)], 3) + " + ); + + Ok(()) +} + +#[test] +fn inner_range_join_rehashes_different_bounds() -> Result<()> { + let join = partitioned_join_plan( + range_partitioning("a", [10, 20], SortOptions::default())?, + range_partitioning("b", [15, 20], SortOptions::default())?, + JoinType::Inner, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([b@1 ASC], [(15), (20)], 3) + " + ); + + Ok(()) +} + +#[test] +fn inner_hash_join_rehashes_mismatched_counts() -> Result<()> { + let join = partitioned_join_plan( + hash_partitioning("a", 11)?, + hash_partitioning("b", 12)?, + JoinType::Inner, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=11 + PartitionedTestExec: output_partitioning=Hash([a@0], 11) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=12 + PartitionedTestExec: output_partitioning=Hash([b@1], 12) + " + ); + Ok(()) +} + +#[test] +fn inner_hash_join_rehashes_to_target_count() -> Result<()> { + let join = partitioned_join_plan( + hash_partitioning("a", 3)?, + hash_partitioning("b", 3)?, + JoinType::Inner, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Hash([a@0], 3) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Hash([b@1], 3) + " + ); + Ok(()) +} + +#[test] +fn non_inner_range_join_rehashes() -> Result<()> { + let join = partitioned_join_plan( + range_partitioning("a", [10, 20], SortOptions::default())?, + range_partitioning("b", [10, 20], SortOptions::default())?, + JoinType::Left, + ); + + let plan = TestConfig::default().to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Left, on=[(a@0, b@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + RepartitionExec: partitioning=Hash([b@1], 10), input_partitions=3 + PartitionedTestExec: output_partitioning=Range([b@1 ASC], [(10), (20)], 3) + " + ); + Ok(()) +} + +#[test] +fn range_aggregate_reuses_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let aggregate = + aggregate_exec_with_alias(input, vec![("a".to_string(), "a".to_string())]); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] + AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + " + ); + + Ok(()) +} + +#[test] +fn range_grouping_set_aggregate_rehashes_with_grouping_id() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20], + SortOptions::default(), + )?); + let input_schema = input.schema(); + let group_by = PhysicalGroupBy::new( + vec![ + (col("a", &input_schema)?, "a".to_string()), + (col("b", &input_schema)?, "b".to_string()), + ], + vec![ + (lit(ScalarValue::Int64(None)), "a".to_string()), + (lit(ScalarValue::Int64(None)), "b".to_string()), + ], + vec![vec![false, true], vec![false, false]], + true, + ); + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + input, + Arc::clone(&input_schema), + )?); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + group_by.as_final(), + vec![], + vec![], + Arc::clone(&partial) as _, + partial.schema(), + )?); + + let plan = TestConfig::default() + .with_query_execution_partitions(3) + .to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, __grouping_id@2 as __grouping_id], aggr=[] + RepartitionExec: partitioning=Hash([a@0, b@1, __grouping_id@2], 3), input_partitions=3 + AggregateExec: mode=Partial, gby=[(a@0 as a, NULL as b), (a@0 as a, b@1 as b)], aggr=[] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20)], 3) + " + ); + + Ok(()) +} + +#[test] +fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 30, 40], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4 + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4) + " + ); + + Ok(()) +} + +#[test] +fn range_right_mark_hash_join_reuses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::RightMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(a@0, a@0)] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + " + ); + + Ok(()) +} + +#[test] +fn range_right_semi_hash_join_rehashes_incompatible_sort_options() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [20], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [20], + SortOptions { + descending: true, + nulls_first: true, + }, + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::RightSemi); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(20)], 2) + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + PartitionedTestExec: output_partitioning=Range([a@0 DESC], [(20)], 2) + " + ); + + Ok(()) +} + +#[test] +fn range_window_reuses_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let window = bounded_window_exec_with_can_repartition( + "a", + vec![], + &[col("a", &schema())?], + input, + true, + ); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(window, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + "# + ); + + Ok(()) +} + +#[test] +fn range_window_rehashes_incompatible_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let window = bounded_window_exec_with_can_repartition( + "b", + vec![], + &[col("b", &schema())?], + input, + true, + ); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(window, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([b@1], 4), input_partitions=4 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + "# + ); + + Ok(()) +} + +#[test] +fn range_full_hash_join_reuses_compatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Full); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + " + ); + + Ok(()) +} + +#[test] +fn range_full_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 30, 40], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Full); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4 + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4) + " + ); + + Ok(()) +} + +#[test] +fn range_left_mark_hash_join_reuses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(a@0, a@0)] + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + " + ); + + Ok(()) +} + +#[test] +fn range_left_anti_hash_join_rehashes_incompatible_null_options() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions { + descending: false, + nulls_first: false, + }, + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftAnti); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + PartitionedTestExec: output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4) + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + PartitionedTestExec: output_partitioning=Range([a@0 ASC NULLS LAST], [(10), (20), (30)], 4) + " + ); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 6f88e01059fc9..827a001b59894 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -724,7 +724,7 @@ fn test_output_req_after_projection() -> Result<()> { ] .into(), )), - Distribution::HashPartitioned(vec![ + Distribution::KeyPartitioned(vec![ Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1)), ]), @@ -746,7 +746,7 @@ fn test_output_req_after_projection() -> Result<()> { actual, @r" ProjectionExec: expr=[c@2 as c, a@0 as new_a, b@1 as b] - OutputRequirementExec: order_by=[(b@1, asc), (c@2 + a@0, asc)], dist_by=HashPartitioned[[a@0, b@1]]) + OutputRequirementExec: order_by=[(b@1, asc), (c@2 + a@0, asc)], dist_by=KeyPartitioned[[a@0, b@1]]) DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false " ); @@ -762,7 +762,7 @@ fn test_output_req_after_projection() -> Result<()> { assert_snapshot!( actual, @r" - OutputRequirementExec: order_by=[(b@2, asc), (c@0 + new_a@1, asc)], dist_by=HashPartitioned[[new_a@1, b@2]]) + OutputRequirementExec: order_by=[(b@2, asc), (c@0 + new_a@1, asc)], dist_by=KeyPartitioned[[new_a@1, b@2]]) DataSourceExec: file_groups={1 group: [[x]]}, projection=[c, a@0 as new_a, b], file_type=csv, has_header=false " ); @@ -797,10 +797,12 @@ fn test_output_req_after_projection() -> Result<()> { Arc::new(Column::new("new_a", 1)), Arc::new(Column::new("b", 2)), ]; - if let Distribution::HashPartitioned(vec) = after_optimize + if let Distribution::KeyPartitioned(vec) = after_optimize .downcast_ref::() .unwrap() - .required_input_distribution()[0] + .input_distribution_requirements() + .child_distribution(0) + .unwrap() .clone() { assert!( @@ -809,7 +811,7 @@ fn test_output_req_after_projection() -> Result<()> { .all(|(actual, expected)| actual.eq(&expected)) ); } else { - panic!("Expected HashPartitioned distribution!"); + panic!("Expected KeyPartitioned distribution!"); }; Ok(()) diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index 217570846d56e..184125dcbe180 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -19,8 +19,9 @@ use insta::assert_snapshot; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - bounded_window_exec, global_limit_exec, local_limit_exec, memory_exec, - projection_exec, repartition_exec, sort_exec, sort_expr, sort_expr_options, + bounded_window_exec, bounded_window_exec_with_can_repartition, global_limit_exec, + hash_join_exec, local_limit_exec, memory_exec, projection_exec, repartition_exec, + sort_exec, sort_exec_with_preserve_partitioning, sort_expr, sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; @@ -29,12 +30,13 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; use datafusion::prelude::{CsvReadOptions, SessionContext}; use datafusion_common::config::ConfigOptions; -use datafusion_common::{JoinType, Result, ScalarValue}; -use datafusion_physical_expr::Partitioning; +use datafusion_common::{JoinType, NullEquality, Result, ScalarValue}; use datafusion_physical_expr::expressions::{Literal, col}; +use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; +use datafusion_physical_plan::joins::{StreamJoinPartitionMode, SymmetricHashJoinExec}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::{ExecutionPlan, displayable}; @@ -400,6 +402,146 @@ fn assert_sanity_check(plan: &Arc, is_sane: bool) { ); } +fn range_partitioned_exec( + schema: &SchemaRef, + key: &str, + split_points: impl IntoIterator, +) -> Result> { + let split_points = split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [sort_expr(key, schema)].into(), + split_points, + )?); + RepartitionExec::try_new(memory_exec(schema), partitioning) + .map(|exec| Arc::new(exec) as Arc) +} + +#[test] +fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Inner, + )?; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Inner, + )?; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_partitioned_right_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Right, + )?; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Right, + )?; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_sort_merge_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + let ordering: LexOrdering = [sort_expr("a", &schema)].into(); + + let compatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&compatible_join, true); + + let incompatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering, + range_partitioned_exec(&schema, "a", [20])?, + ), + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_symmetric_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = Arc::new(SymmetricHashJoinExec::try_new( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = Arc::new(SymmetricHashJoinExec::try_new( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + #[tokio::test] /// Tests that plan is valid when the sort requirements are satisfied. async fn test_bounded_window_agg_sort_requirement() -> Result<()> { @@ -458,6 +600,76 @@ async fn test_bounded_window_agg_no_sort_requirement() -> Result<()> { Ok(()) } +#[tokio::test] +/// Tests that a window over a compatible range-partitioned input satisfies +/// the window's key distribution requirement without a hash repartition. +async fn test_bounded_window_agg_range_partitioning() -> Result<()> { + let schema = create_test_schema2(); + let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?; + let ordering: LexOrdering = [sort_expr_options( + "a", + &schema, + SortOptions { + descending: false, + nulls_first: false, + }, + )] + .into(); + let partition_by = vec![col("a", &schema)?]; + let sort = sort_exec_with_preserve_partitioning(ordering, source); + let bw = + bounded_window_exec_with_can_repartition("a", vec![], &partition_by, sort, true); + let plan_str = displayable(bw.as_ref()).indent(true).to_string(); + let actual = plan_str.trim(); + assert_snapshot!( + actual, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + "# + ); + assert_sanity_check(&bw, true); + Ok(()) +} + +#[tokio::test] +/// Tests that a window over an incompatible range-partitioned input fails +/// the window's key distribution requirement. +async fn test_bounded_window_agg_incompatible_range_partitioning() -> Result<()> { + let schema = create_test_schema2(); + let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?; + let ordering: LexOrdering = [sort_expr_options( + "b", + &schema, + SortOptions { + descending: false, + nulls_first: false, + }, + )] + .into(); + let partition_by = vec![col("b", &schema)?]; + let sort = sort_exec_with_preserve_partitioning(ordering, source); + let bw = + bounded_window_exec_with_can_repartition("b", vec![], &partition_by, sort, true); + let plan_str = displayable(bw.as_ref()).indent(true).to_string(); + let actual = plan_str.trim(); + assert_snapshot!( + actual, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + "# + ); + // Range([a]) does not colocate `b` values, so the window's key + // distribution requirement is not satisfied. + assert_sanity_check(&bw, false); + Ok(()) +} + #[tokio::test] /// A valid when a single partition requirement /// is satisfied. diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 09225cb0385a7..1973813215ed6 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -263,6 +263,22 @@ pub fn bounded_window_exec_with_partition( sort_exprs: impl IntoIterator, partition_by: &[Arc], input: Arc, +) -> Arc { + bounded_window_exec_with_can_repartition( + col_name, + sort_exprs, + partition_by, + input, + false, + ) +} + +pub fn bounded_window_exec_with_can_repartition( + col_name: &str, + sort_exprs: impl IntoIterator, + partition_by: &[Arc], + input: Arc, + can_repartition: bool, ) -> Arc { let sort_exprs = sort_exprs.into_iter().collect::>(); let schema = input.schema(); @@ -285,7 +301,7 @@ pub fn bounded_window_exec_with_partition( vec![window_expr], Arc::clone(&input), InputOrderMode::Sorted, - false, + can_repartition, ) .unwrap(), ) diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index e8ff6758ccdd4..b837373632f07 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -708,8 +708,12 @@ impl ExecutionPlan for TopKExec { &self.cache } - fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + ]) } fn children(&self) -> Vec<&Arc> { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 4bf86e17d387d..5e4126793cb1f 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -27,8 +27,7 @@ use crate::{ file_stream::work_source::SharedWorkSource, source::DataSource, statistics::MinMaxStatistics, }; -use arrow::datatypes::FieldRef; -use arrow::datatypes::{DataType, Schema, SchemaRef}; +use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::{ Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, @@ -40,7 +39,7 @@ use datafusion_expr::Operator; use crate::source::OpenArgs; use datafusion_physical_expr::expressions::{BinaryExpr, Column}; -use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::projection::{ProjectionExprs, ProjectionMapping}; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; @@ -205,7 +204,17 @@ pub struct FileScanConfig { /// /// If the number of file partitions > target_partitions, the file partitions will be grouped /// in a round-robin fashion such that number of file partitions = target_partitions. + /// + /// Follow-up: remove this redundant field in favor of + /// `output_partitioning`, see . pub partitioned_by_file_group: bool, + /// Declared physical output partitioning for this scan. + /// + /// Expressions are against the full table schema, before scan projection or + /// filtering. `ListingTable` validates partition count before building the + /// scan, and direct builders with mismatched counts fall back to + /// `UnknownPartitioning`. + pub output_partitioning: Option, } /// A builder for [`FileScanConfig`]'s. @@ -272,6 +281,7 @@ pub struct FileScanConfigBuilder { file_groups: Vec, statistics: Option, output_ordering: Vec, + output_partitioning: Option, file_compression_type: Option, batch_size: Option, expr_adapter_factory: Option>, @@ -295,6 +305,7 @@ impl FileScanConfigBuilder { file_groups: vec![], statistics: None, output_ordering: vec![], + output_partitioning: None, file_compression_type: None, limit: None, preserve_order: false, @@ -461,6 +472,15 @@ impl FileScanConfigBuilder { self } + /// Set declared physical output partitioning for this scan. + pub fn with_output_partitioning( + mut self, + output_partitioning: Option, + ) -> Self { + self.output_partitioning = output_partitioning; + self + } + /// Set the file compression type pub fn with_file_compression_type( mut self, @@ -519,6 +539,7 @@ impl FileScanConfigBuilder { file_groups, statistics, output_ordering, + output_partitioning, file_compression_type, batch_size, expr_adapter_factory: expr_adapter, @@ -548,6 +569,7 @@ impl FileScanConfigBuilder { expr_adapter_factory: expr_adapter, statistics, partitioned_by_file_group, + output_partitioning, } } } @@ -560,6 +582,7 @@ impl From for FileScanConfigBuilder { file_groups: config.file_groups, statistics: Some(config.statistics), output_ordering: config.output_ordering, + output_partitioning: config.output_partitioning, file_compression_type: Some(config.file_compression_type), limit: config.limit, preserve_order: config.preserve_order, @@ -571,6 +594,52 @@ impl From for FileScanConfigBuilder { } } +fn hash_partitioning_from_partition_fields( + schema: &Schema, + partition_cols: &[FieldRef], + partition_count: usize, +) -> Option { + if partition_cols.is_empty() { + return None; + } + + let mut exprs: Vec> = Vec::with_capacity(partition_cols.len()); + for partition_col in partition_cols { + let name = partition_col.name(); + let idx = schema + .fields() + .iter() + .position(|field| field.name() == name)?; + exprs.push(Arc::new(Column::new(name, idx))); + } + + Some(Partitioning::Hash(exprs, partition_count)) +} + +fn project_output_partitioning( + partitioning: &Partitioning, + mapping: &ProjectionMapping, + input_schema: &SchemaRef, + partition_count: usize, +) -> Partitioning { + let input_eq_properties = EquivalenceProperties::new(Arc::clone(input_schema)); + match partitioning { + Partitioning::Hash(exprs, _) => { + let projected_exprs = input_eq_properties + .project_expressions(exprs, mapping) + .collect::>>(); + projected_exprs + .map(|exprs| Partitioning::Hash(exprs, partition_count)) + .unwrap_or_else(|| Partitioning::UnknownPartitioning(partition_count)) + } + Partitioning::Range(_) + | Partitioning::RoundRobinBatch(_) + | Partitioning::UnknownPartitioning(_) => { + partitioning.project(mapping, &input_eq_properties) + } + } +} + impl DataSource for FileScanConfig { fn open( &self, @@ -658,6 +727,10 @@ impl DataSource for FileScanConfig { display_orderings(f, &orderings)?; + if self.output_partitioning.is_some() { + write!(f, ", output_partitioning={}", self.output_partitioning())?; + } + if !self.constraints.is_empty() { write!(f, ", {}", self.constraints)?; } @@ -681,10 +754,9 @@ impl DataSource for FileScanConfig { repartition_file_min_size: usize, output_ordering: Option, ) -> Result>> { - // When files are grouped by partition values, we cannot allow byte-range - // splitting. It would mix rows from different partition values across - // file groups, breaking the Hash partitioning. - if self.partitioned_by_file_group { + // When file groups define output partitioning, repartitioning files + // would invalidate the partition-to-file-group mapping. + if self.output_partitioning.is_some() || self.partitioned_by_file_group { return Ok(None); } @@ -700,13 +772,18 @@ impl DataSource for FileScanConfig { /// Returns the output partitioning for this file scan. /// - /// When `partitioned_by_file_group` is true, this returns `Partitioning::Hash` on - /// the Hive partition columns, allowing the optimizer to skip hash repartitioning - /// for aggregates and joins on those columns. + /// When `output_partitioning` is set, this returns the declared partitioning + /// after applying scan projection. When `partitioned_by_file_group` is true, + /// this returns `Partitioning::Hash` on the Hive partition columns, allowing + /// the optimizer to skip hash repartitioning for aggregates and joins on + /// those columns. + /// + /// If projection or partition count validation fails, this returns + /// `UnknownPartitioning`. /// /// Tradeoffs - /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries with - /// `GROUP BY` or `ORDER BY` on partition columns. + /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries whose + /// required distribution is satisfied by the scan's output partitioning. /// - Cost: Files are grouped by partition values rather than split by byte /// ranges, which may reduce I/O parallelism when partition sizes are uneven. /// For simple aggregations without `ORDER BY`, this cost may outweigh the benefit. @@ -715,39 +792,45 @@ impl DataSource for FileScanConfig { /// - Idea: Could allow byte-range splitting within partition-aware groups, /// preserving I/O parallelism while maintaining partition semantics. fn output_partitioning(&self) -> Partitioning { - if self.partitioned_by_file_group { - let partition_cols = self.table_partition_cols(); - if !partition_cols.is_empty() { - let projected_schema = match self.projected_schema() { - Ok(schema) => schema, - Err(_) => { - debug!( - "Could not get projected schema, falling back to UnknownPartitioning." - ); - return Partitioning::UnknownPartitioning(self.file_groups.len()); - } - }; - - // Build Column expressions for partition columns based on their - // position in the projected schema - let mut exprs: Vec> = Vec::new(); - for partition_col in partition_cols { - if let Some((idx, _)) = projected_schema - .fields() - .iter() - .enumerate() - .find(|(_, f)| f.name() == partition_col.name()) - { - exprs.push(Arc::new(Column::new(partition_col.name(), idx))); - } - } + let Some(output_partitioning) = self.output_partitioning.clone().or_else(|| { + self.partitioned_by_file_group.then(|| { + hash_partitioning_from_partition_fields( + self.file_source.table_schema().table_schema(), + self.table_partition_cols(), + self.file_groups.len(), + ) + })? + }) else { + return Partitioning::UnknownPartitioning(self.file_groups.len()); + }; + if output_partitioning.partition_count() != self.file_groups.len() { + warn!( + "Declared output partitioning has {} partitions, but file scan has {} file groups. Falling back to UnknownPartitioning.", + output_partitioning.partition_count(), + self.file_groups.len() + ); + return Partitioning::UnknownPartitioning(self.file_groups.len()); + } - if exprs.len() == partition_cols.len() { - return Partitioning::Hash(exprs, self.file_groups.len()); + if let Some(projection) = self.file_source.projection() { + let schema = self.file_source.table_schema().table_schema(); + return match projection.projection_mapping(schema) { + Ok(mapping) => project_output_partitioning( + &output_partitioning, + &mapping, + schema, + self.file_groups.len(), + ), + Err(e) => { + debug!( + "Could not project output partitioning, falling back to UnknownPartitioning: {e}" + ); + Partitioning::UnknownPartitioning(self.file_groups.len()) } - } + }; } - Partitioning::UnknownPartitioning(self.file_groups.len()) + + output_partitioning } /// Computes the effective equivalence properties of this file scan, taking @@ -1041,7 +1124,10 @@ impl DataSource for FileScanConfig { /// when file order must be preserved or the file groups define the output /// partitioning needed for the rest of the plan fn create_sibling_state(&self) -> Option> { - if self.preserve_order || self.partitioned_by_file_group { + if self.preserve_order + || self.output_partitioning.is_some() + || self.partitioned_by_file_group + { return None; } @@ -2445,6 +2531,58 @@ mod tests { assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_))); } + #[test] + fn test_declared_output_partitioning_projects_with_scan() { + let file_schema = aggr_test_schema(); + let output_partitioning = + Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4); + + let mut config = config_for_projection( + Arc::clone(&file_schema), + Some(vec![1, 2]), + Statistics::new_unknown(&file_schema), + vec![], + ); + config.file_groups = vec![ + FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]), + ]; + config.output_partitioning = Some(output_partitioning); + + match config.output_partitioning() { + Partitioning::Hash(exprs, num_partitions) => { + assert_eq!(num_partitions, 4); + assert_eq!(exprs.len(), 1); + let column = exprs[0].downcast_ref::().unwrap(); + assert_eq!(column.name(), "c2"); + assert_eq!(column.index(), 0); + } + _ => panic!("Expected Hash partitioning"), + } + + let mut config = config_for_projection( + Arc::clone(&file_schema), + Some(vec![2]), + Statistics::new_unknown(&file_schema), + vec![], + ); + config.file_groups = vec![ + FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]), + ]; + config.output_partitioning = + Some(Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4)); + + assert!(matches!( + config.output_partitioning(), + Partitioning::UnknownPartitioning(4) + )); + } + #[test] fn test_output_partitioning_no_partition_columns() { let file_schema = aggr_test_schema(); diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index e3df1ad6381f4..18ebe80773e8a 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -31,8 +31,9 @@ use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequire use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - PlanProperties, SendableRecordBatchStream, execute_input_stream, + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, Partitioning, PlanProperties, + SendableRecordBatchStream, execute_input_stream, }; use async_trait::async_trait; @@ -189,9 +190,16 @@ impl ExecutionPlan for DataSinkExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { // DataSink is responsible for dynamically partitioning its // own input at execution time, and so requires a single input partition. - vec![Distribution::SinglePartition; self.children().len()] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition; + self.children().len() + ]) } fn required_input_ordering(&self) -> Vec> { 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 c2b01868c97f3..5da6d97d64403 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -41,8 +41,8 @@ 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, ToStringifiedPlan, Union, Unnest, Values, Window, + RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, + Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, projection_schema, }; pub use statement::{ diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c572b202f03ce..a5530af914164 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -50,6 +50,7 @@ use crate::{ WindowFunctionDefinition, build_join_schema, expr_vec_fmt, requalify_sides_if_needed, }; +use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datafusion_common::cse::{NormalizeEq, Normalizeable}; use datafusion_common::format::ExplainFormat; @@ -60,10 +61,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; @@ -864,6 +867,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 { @@ -2091,6 +2120,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(); @@ -4137,11 +4169,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 @@ -4149,10 +4186,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 /// @@ -4490,6 +4635,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 ef9382a57209a..32396d2cd26f2 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, @@ -414,6 +415,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, .. }) => { @@ -519,6 +521,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/ffi/src/physical_expr/partitioning.rs b/datafusion/ffi/src/physical_expr/partitioning.rs index 434b6a097e645..2a9a8528c6c3e 100644 --- a/datafusion/ffi/src/physical_expr/partitioning.rs +++ b/datafusion/ffi/src/physical_expr/partitioning.rs @@ -17,20 +17,35 @@ use std::sync::Arc; -use datafusion_physical_expr::Partitioning; +use datafusion_common::{DataFusionError, ScalarValue, SplitPoint}; +use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, +}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use stabby::vec::Vec as SVec; +use crate::arrow_wrappers::WrappedArray; use crate::physical_expr::FFI_PhysicalExpr; +use crate::physical_expr::sort::FFI_PhysicalSortExpr; + +/// A stable struct for sharing [`RangePartitioning`] across FFI boundaries. +/// See [`RangePartitioning`] for the descriptions of each field. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_RangePartitioning { + split_points: SVec>, + ordering: SVec, +} /// A stable struct for sharing [`Partitioning`] across FFI boundaries. -/// See ['Partitioning'] for the meaning of each variant. +/// See [`Partitioning`] for the meaning of each variant. #[repr(C)] #[derive(Debug)] pub enum FFI_Partitioning { RoundRobinBatch(usize), Hash(SVec, usize), UnknownPartitioning(usize), + Range(FFI_RangePartitioning), } impl From<&Partitioning> for FFI_Partitioning { @@ -45,44 +60,130 @@ impl From<&Partitioning> for FFI_Partitioning { .collect(); Self::Hash(exprs, *size) } + Partitioning::Range(range) => { + // Producer-side conversion should be infallible at ABI boundary + let split_points = range + .split_points() + .iter() + .map(|split_point| { + split_point + .values() + .iter() + .map(|value| { + WrappedArray::try_from(value).expect( + "ScalarValue in RangePartitioning should convert to WrappedArray", + ) + }) + .collect() + }) + .collect(); + let ordering = range + .ordering() + .iter() + .map(FFI_PhysicalSortExpr::from) + .collect(); + Self::Range(FFI_RangePartitioning { + split_points, + ordering, + }) + } Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size), } } } -impl From<&FFI_Partitioning> for Partitioning { - fn from(value: &FFI_Partitioning) -> Self { - match value { +impl TryFrom for Partitioning { + type Error = DataFusionError; + + fn try_from(value: FFI_Partitioning) -> Result { + Ok(match value { FFI_Partitioning::RoundRobinBatch(size) => { - Partitioning::RoundRobinBatch(*size) + Partitioning::RoundRobinBatch(size) } FFI_Partitioning::Hash(exprs, size) => { let exprs = exprs.iter().map(>::from).collect(); - Self::Hash(exprs, *size) + Self::Hash(exprs, size) + } + FFI_Partitioning::Range(range) => { + let split_points = range + .split_points + .into_iter() + .map(|split_point| { + split_point + .into_iter() + .map(ScalarValue::try_from) + .collect::, _>>() + .map(SplitPoint::new) + }) + .collect::, _>>()?; + + let ordering = + LexOrdering::new(range.ordering.iter().map(PhysicalSortExpr::from)) + .ok_or_else(|| { + DataFusionError::Internal( + "FFI Range partitioning ordering must be non-empty" + .to_string(), + ) + })?; + + Self::Range(RangePartitioning::try_new(ordering, split_points)?) } FFI_Partitioning::UnknownPartitioning(size) => { - Self::UnknownPartitioning(*size) + Self::UnknownPartitioning(size) } - } + }) } } #[cfg(test)] mod tests { - use datafusion_physical_expr::Partitioning; - use datafusion_physical_expr::expressions::lit; + use std::sync::Arc; + + use arrow_schema::SortOptions; + use datafusion_common::{Result, ScalarValue, SplitPoint}; + use datafusion_physical_expr::expressions::{Column, lit}; + use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, + }; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use stabby::vec::Vec as SVec; - use crate::physical_expr::partitioning::FFI_Partitioning; + use crate::physical_expr::partitioning::{FFI_Partitioning, FFI_RangePartitioning}; + + fn range_partitioning() -> Result { + let a = Arc::new(Column::new("a", 0)) as Arc; + let b = Arc::new(Column::new("b", 1)) as Arc; + let ordering = LexOrdering::new([ + PhysicalSortExpr::new(a, SortOptions::default()), + PhysicalSortExpr::new(b, SortOptions::new(true, false)), + ]) + .expect("non-empty ordering"); + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int64(Some(10)), + ScalarValue::Utf8(Some("a".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(20)), + ScalarValue::Utf8(Some("b".to_string())), + ]), + ]; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) + } #[test] - fn round_trip_ffi_partitioning() { + fn round_trip_ffi_partitioning() -> Result<()> { for partitioning in [ Partitioning::RoundRobinBatch(10), Partitioning::Hash(vec![lit(1)], 10), Partitioning::UnknownPartitioning(10), + range_partitioning()?, ] { let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); - let returned: Partitioning = (&ffi_partitioning).into(); + let returned: Partitioning = ffi_partitioning.try_into()?; if let Partitioning::UnknownPartitioning(return_size) = returned { let Partitioning::UnknownPartitioning(original_size) = partitioning @@ -94,5 +195,32 @@ mod tests { assert_eq!(partitioning, returned); } } + + Ok(()) + } + + #[test] + fn round_trip_ffi_range_partitioning_compound_key() -> Result<()> { + let partitioning = range_partitioning()?; + + let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); + let returned: Partitioning = ffi_partitioning.try_into()?; + assert_eq!(partitioning, returned); + + Ok(()) + } + + #[test] + fn ffi_range_partitioning_rejects_empty_ordering() { + let ffi_partitioning = FFI_Partitioning::Range(FFI_RangePartitioning { + split_points: SVec::new(), + ordering: SVec::new(), + }); + + let err = Partitioning::try_from(ffi_partitioning).unwrap_err(); + assert!( + err.to_string().contains("ordering must be non-empty"), + "{err}" + ); } } diff --git a/datafusion/ffi/src/plan_properties.rs b/datafusion/ffi/src/plan_properties.rs index b286ee2d7d30c..09ef26af32349 100644 --- a/datafusion/ffi/src/plan_properties.rs +++ b/datafusion/ffi/src/plan_properties.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use datafusion_common::error::{DataFusionError, Result}; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::PlanProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -172,6 +172,7 @@ impl TryFrom for PlanProperties { .unwrap_or_default(); let partitioning = unsafe { (ffi_props.output_partitioning)(&ffi_props) }; + let partitioning = Partitioning::try_from(partitioning)?; let eq_properties = if sort_exprs.is_empty() { EquivalenceProperties::new(Arc::new(schema)) @@ -187,7 +188,7 @@ impl TryFrom for PlanProperties { Ok(PlanProperties::new( eq_properties, - (&partitioning).into(), + partitioning, emission_type, boundedness, )) @@ -260,13 +261,15 @@ impl From for EmissionType { #[cfg(test)] mod tests { + use arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_expr::PhysicalSortExpr; use datafusion::physical_plan::Partitioning; + use datafusion_common::{ScalarValue, SplitPoint}; + use datafusion_physical_expr::{LexOrdering, RangePartitioning}; use super::*; fn create_test_props() -> Result { - use arrow::datatypes::{DataType, Field, Schema}; let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); @@ -282,6 +285,25 @@ mod tests { )) } + fn create_range_test_props() -> Result { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let col = datafusion::physical_plan::expressions::col("a", &schema)?; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(col)]) + .expect("non-empty ordering"); + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), + ]; + let range = RangePartitioning::try_new(ordering, split_points)?; + + Ok(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::Range(range), + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + #[test] fn test_round_trip_ffi_plan_properties() -> Result<()> { let original_props = create_test_props()?; @@ -314,4 +336,22 @@ mod tests { Ok(()) } + + #[test] + fn test_round_trip_ffi_plan_properties_range_partitioning() -> Result<()> { + let original_props = create_range_test_props()?; + + let mut local_props_ptr = FFI_PlanProperties::from(&original_props); + local_props_ptr.library_marker_id = crate::mock_foreign_marker_id; + + let foreign_props: PlanProperties = local_props_ptr.try_into()?; + + assert_eq!( + format!("{:?}", foreign_props.output_partitioning()), + format!("{:?}", original_props.output_partitioning()) + ); + assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}")); + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 848bf81d15979..ec7bf648e22ea 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -55,14 +55,18 @@ 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 partitioning::{Distribution, Partitioning}; +pub use partitioning::{ + Distribution, Partitioning, PartitioningSatisfaction, 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, - physical_exprs_bag_equal, physical_exprs_contains, physical_exprs_equal, + create_ordering, create_physical_partitioning, create_physical_sort_expr, + create_physical_sort_exprs, physical_exprs_bag_equal, physical_exprs_contains, + physical_exprs_equal, }; pub use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index d24c60b63e6bd..14065ab8853ad 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -21,7 +21,10 @@ use crate::{ EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, expressions::UnKnownColumn, physical_exprs_equal, }; +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::fmt; use std::fmt::Display; use std::sync::Arc; @@ -117,6 +120,8 @@ 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 source-declared ranges + Range(RangePartitioning), /// Unknown partitioning scheme with a known number of partitions UnknownPartitioning(usize), } @@ -133,6 +138,7 @@ impl Display for Partitioning { .join(", "); write!(f, "Hash([{phy_exprs_str}], {size})") } + Partitioning::Range(range) => write!(f, "{range}"), Partitioning::UnknownPartitioning(size) => { write!(f, "UnknownPartitioning({size})") } @@ -140,6 +146,191 @@ impl Display for Partitioning { } } +/// Physical range partitioning. +/// +/// [`RangePartitioning`] describes an ordered key space with split points. +/// +/// - `ordering` defines the partitioning key and ordering. +/// - `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 strictly ordered +/// according to that ordering, and each split point must have one value per +/// 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 +/// partition described by the split points. DataFusion will not validate this is +/// upheld. +/// +/// For a single range key: +/// +/// ```text +/// ordering = [date ASC NULLS LAST] +/// split_points = [ +/// (2022-01-01), +/// (2023-01-01), +/// ] +/// +/// partition 0: date before 2022-01-01 +/// partition 1: date between 2022-01-01 (inclusive) and 2023-01-01 (exclusive) +/// partition 2: date at/after 2023-01-01 +/// ``` +/// +/// The same model extends to compound keys. +/// For `ordering = [time ASC, city ASC]`, split points are ordered +/// lexicographically by `(time, city)`: +/// +/// ```text +/// ordering = [time ASC NULLS LAST, city ASC NULLS LAST] +/// split_points = [ +/// (2022, Allston), +/// (2023, Allston), +/// ] +/// +/// partition 0: keys before (2022, Allston) +/// partition 1: keys between (2022, Allston) and (2023, Allston) +/// partition 2: keys at/after (2023, Allston) +/// ``` +/// +/// NOTE: Optimizer and execution behavior for this partitioning is intentionally +/// not implemented and will be introduced incrementally. See +/// . +#[derive(Debug, Clone, PartialEq)] +pub struct RangePartitioning { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Boundaries between adjacent partitions. + split_points: Vec, +} + +impl RangePartitioning { + /// Creates range partitioning metadata without validating split points. + /// + /// Use [`Self::try_new`] to validate the contract documented on + /// [`RangePartitioning`]. + pub fn new(ordering: LexOrdering, split_points: Vec) -> Self { + Self { + ordering, + split_points, + } + } + + /// 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( + &split_points, + &ordering + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(), + )?; + Ok(Self::new(ordering, split_points)) + } + + /// Returns the ordering that defines the range key. + pub fn ordering(&self) -> &LexOrdering { + &self.ordering + } + + /// Returns the ordered split points between partitions. + pub fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Returns the number of partitions. + pub fn partition_count(&self) -> usize { + self.split_points.len() + 1 + } + + /// Calculates the range partitioning after applying the given projection. + /// + /// Returns `None` if any range key cannot be projected or if projection + /// collapses distinct range keys into duplicate output expressions. + fn project( + &self, + mapping: &ProjectionMapping, + input_eq_properties: &EquivalenceProperties, + ) -> Option { + let exprs = self + .ordering + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + let projected_exprs = input_eq_properties + .project_expressions(&exprs, mapping) + .collect::>>()?; + let sort_exprs = self + .ordering + .iter() + .zip(projected_exprs) + .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, sort_expr.options)) + .collect::>(); + let ordering = LexOrdering::new(sort_exprs)?; + if ordering.len() != self.ordering.len() { + return None; + } + + Some(Self { + ordering, + split_points: self.split_points.clone(), + }) + } +} + +impl Display for RangePartitioning { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let split_points = format_range_split_points(&self.split_points); + write!( + f, + "Range([{}], [{}], {})", + self.ordering, + split_points, + self.partition_count() + ) + } +} + +fn format_range_split_points(split_points: &[SplitPoint]) -> String { + split_points + .iter() + .map(ToString::to_string) + .collect::>() + .join(", ") +} + +fn equivalent_exprs( + left: &[Arc], + right: &[Arc], + eq_properties: &EquivalenceProperties, +) -> bool { + if physical_exprs_equal(left, right) { + return true; + } + + let eq_groups = eq_properties.eq_group(); + if eq_groups.is_empty() { + return false; + } + + let normalized_left = normalize_exprs(left, eq_properties); + let normalized_right = normalize_exprs(right, eq_properties); + + physical_exprs_equal(&normalized_left, &normalized_right) +} + +fn normalize_exprs( + exprs: &[Arc], + eq_properties: &EquivalenceProperties, +) -> Vec> { + let eq_groups = eq_properties.eq_group(); + exprs + .iter() + .map(|expr| eq_groups.normalize_expr(Arc::clone(expr))) + .collect() +} + /// Represents how a [`Partitioning`] satisfies a [`Distribution`] requirement. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PartitioningSatisfaction { @@ -167,6 +358,7 @@ impl Partitioning { use Partitioning::*; match self { RoundRobinBatch(n) | Hash(_, n) | UnknownPartitioning(n) => *n, + Range(range) => range.partition_count(), } } @@ -201,6 +393,10 @@ impl Partitioning { /// Returns how this [`Partitioning`] satisfies the partitioning scheme mandated /// by the `required` [`Distribution`]. + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] pub fn satisfaction( &self, required: &Distribution, @@ -212,84 +408,111 @@ impl Partitioning { Distribution::SinglePartition if self.partition_count() == 1 => { PartitioningSatisfaction::Exact } - // When partition count is 1, hash requirement is satisfied. - Distribution::HashPartitioned(_) if self.partition_count() == 1 => { + // When partition count is 1, partitioned requirements are satisfied. + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) + if self.partition_count() == 1 => + { PartitioningSatisfaction::Exact } - Distribution::HashPartitioned(required_exprs) => match self { + Distribution::HashPartitioned(required_exprs) + | Distribution::KeyPartitioned(required_exprs) => match self { // Here we do not check the partition count for hash partitioning and assumes the partition count // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins, // then we need to have the partition count and hash functions validation. - Partitioning::Hash(partition_exprs, _) => { - // Empty hash partitioning is invalid - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - // Fast path: exact match - if physical_exprs_equal(required_exprs, partition_exprs) { - return PartitioningSatisfaction::Exact; - } - - // Normalization path using equivalence groups - let eq_groups = eq_properties.eq_group(); - if !eq_groups.is_empty() { - let normalized_required_exprs = required_exprs - .iter() - .map(|e| eq_groups.normalize_expr(Arc::clone(e))) - .collect::>(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|e| eq_groups.normalize_expr(Arc::clone(e))) - .collect::>(); - if physical_exprs_equal( - &normalized_required_exprs, - &normalized_partition_exprs, - ) { - return PartitioningSatisfaction::Exact; - } - - if allow_subset - && Self::is_subset_partitioning( - &normalized_partition_exprs, - &normalized_required_exprs, - ) - { - return PartitioningSatisfaction::Subset; - } - } else if allow_subset - && Self::is_subset_partitioning(partition_exprs, required_exprs) - { - return PartitioningSatisfaction::Subset; - } - + Partitioning::Hash(partition_exprs, _) => Self::key_satisfaction( + partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ), + Partitioning::Range(range) => { + let partition_exprs = range + .ordering() + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + Self::key_satisfaction( + &partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ) + } + Partitioning::RoundRobinBatch(_) + | Partitioning::UnknownPartitioning(_) => { PartitioningSatisfaction::NotSatisfied } - _ => PartitioningSatisfaction::NotSatisfied, }, - _ => PartitioningSatisfaction::NotSatisfied, + Distribution::SinglePartition => PartitioningSatisfaction::NotSatisfied, } } + fn key_satisfaction( + partition_exprs: &[Arc], + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, + ) -> PartitioningSatisfaction { + if partition_exprs.is_empty() || required_exprs.is_empty() { + return PartitioningSatisfaction::NotSatisfied; + } + + if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { + return PartitioningSatisfaction::Exact; + } + + let eq_groups = eq_properties.eq_group(); + if !eq_groups.is_empty() { + if allow_subset { + let normalized_partition_exprs = + normalize_exprs(partition_exprs, eq_properties); + let normalized_required_exprs = + normalize_exprs(required_exprs, eq_properties); + if Self::is_subset_partitioning( + &normalized_partition_exprs, + &normalized_required_exprs, + ) { + return PartitioningSatisfaction::Subset; + } + } + } else if allow_subset + && Self::is_subset_partitioning(partition_exprs, required_exprs) + { + return PartitioningSatisfaction::Subset; + } + + PartitioningSatisfaction::NotSatisfied + } + /// Calculate the output partitioning after applying the given projection. pub fn project( &self, mapping: &ProjectionMapping, input_eq_properties: &EquivalenceProperties, ) -> Self { - if let Partitioning::Hash(exprs, part) = self { - let normalized_exprs = input_eq_properties - .project_expressions(exprs, mapping) - .zip(exprs) - .map(|(proj_expr, expr)| { - proj_expr.unwrap_or_else(|| { - Arc::new(UnKnownColumn::new(&expr.to_string())) + match self { + Partitioning::Hash(exprs, part) => { + let normalized_exprs = input_eq_properties + .project_expressions(exprs, mapping) + .zip(exprs) + .map(|(proj_expr, expr)| { + proj_expr.unwrap_or_else(|| { + Arc::new(UnKnownColumn::new(&expr.to_string())) + }) }) - }) - .collect(); - Partitioning::Hash(normalized_exprs, *part) - } else { - self.clone() + .collect(); + Partitioning::Hash(normalized_exprs, *part) + } + Partitioning::Range(range) => { + if let Some(projected) = range.project(mapping, input_eq_properties) { + Partitioning::Range(projected) + } else { + Partitioning::UnknownPartitioning(range.partition_count()) + } + } + Partitioning::RoundRobinBatch(_) | Partitioning::UnknownPartitioning(_) => { + self.clone() + } } } } @@ -306,6 +529,7 @@ impl PartialEq for Partitioning { { true } + (Partitioning::Range(left), Partitioning::Range(right)) => left == right, _ => false, } } @@ -319,12 +543,39 @@ pub enum Distribution { UnspecifiedDistribution, /// A single partition is required SinglePartition, - /// Requires children to be distributed in such a way that the same - /// values of the keys end up in the same partition + /// Deprecated historical name for [`Distribution::KeyPartitioned`]. + /// See for details. + #[deprecated(since = "55.0.0", note = "Use Distribution::KeyPartitioned")] HashPartitioned(Vec>), + /// Requires rows with equal values for the given keys to be colocated in + /// the same partition, without requiring a specific partitioning algorithm. + /// + /// Unlike [`Self::HashPartitioned`], this can be satisfied by non-hash + /// partitioning such as range partitioning. A partitioning on a subset of + /// these keys can also satisfy this requirement because rows equal on all + /// required keys are also equal on any subset. + /// + /// For multi-input operators, satisfaction alone is not enough: each input + /// may satisfy its own key requirement while using incompatible partition + /// boundaries. Such operators must separately require compatible + /// co-partitioning before pairing partitions by index. + KeyPartitioned(Vec>), } +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] impl Distribution { + /// Returns key expressions for distribution variants that require + /// co-locating equal key values. + pub fn key_exprs(&self) -> Option<&[Arc]> { + match self { + Self::HashPartitioned(exprs) | Self::KeyPartitioned(exprs) => Some(exprs), + Self::UnspecifiedDistribution | Self::SinglePartition => None, + } + } + /// Creates a `Partitioning` that satisfies this `Distribution` pub fn create_partitioning(self, partition_count: usize) -> Partitioning { match self { @@ -332,13 +583,17 @@ impl Distribution { Partitioning::UnknownPartitioning(partition_count) } Distribution::SinglePartition => Partitioning::UnknownPartitioning(1), - Distribution::HashPartitioned(expr) => { + Distribution::HashPartitioned(expr) | Distribution::KeyPartitioned(expr) => { Partitioning::Hash(expr, partition_count) } } } } +#[expect( + deprecated, + reason = "HashPartitioned display is preserved during the KeyPartitioned migration" +)] impl Display for Distribution { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { @@ -347,6 +602,9 @@ impl Display for Distribution { Distribution::HashPartitioned(exprs) => { write!(f, "HashPartitioned[{}])", format_physical_expr_list(exprs)) } + Distribution::KeyPartitioned(exprs) => { + write!(f, "KeyPartitioned[{}])", format_physical_expr_list(exprs)) + } } } } @@ -356,56 +614,188 @@ mod tests { use super::*; use crate::expressions::Column; + use crate::projection::ProjectionTargets; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_common::{Result, ScalarValue}; + + struct PartitioningTestFixture { + schema: SchemaRef, + cols: Vec>, + eq_properties: EquivalenceProperties, + } + + impl PartitioningTestFixture { + fn new(fields: Vec<(&str, DataType)>) -> Result { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, data_type)| Field::new(*name, data_type.clone(), false)) + .collect::>(), + )); + let cols = fields + .iter() + .map(|(name, _)| { + Ok(Arc::new(Column::new_with_schema(name, &schema)?) + as Arc) + }) + .collect::>()?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + Ok(Self { + schema, + cols, + eq_properties, + }) + } + + fn int64(names: &[&str]) -> Result { + Self::new(names.iter().map(|name| (*name, DataType::Int64)).collect()) + } - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::Result; + fn col(&self, index: usize) -> Arc { + Arc::clone(&self.cols[index]) + } + + fn cols( + &self, + indices: impl IntoIterator, + ) -> Vec> { + indices.into_iter().map(|index| self.col(index)).collect() + } + + fn hash_partitioning( + &self, + indices: impl IntoIterator, + partition_count: usize, + ) -> Partitioning { + Partitioning::Hash(self.cols(indices), partition_count) + } + + fn key_distribution( + &self, + indices: impl IntoIterator, + ) -> Distribution { + Distribution::KeyPartitioned(self.cols(indices)) + } + + fn key_partitioned_distribution( + &self, + indices: impl IntoIterator, + ) -> Distribution { + Distribution::KeyPartitioned(self.cols(indices)) + } + + fn range_sort_expr( + &self, + index: usize, + options: SortOptions, + ) -> PhysicalSortExpr { + PhysicalSortExpr::new(self.col(index), options) + } + + fn range_ordering( + &self, + indices: impl IntoIterator, + ) -> LexOrdering { + LexOrdering::new( + indices + .into_iter() + .map(|index| PhysicalSortExpr::new_default(self.col(index))), + ) + .expect("ordering must not be empty") + } + + fn range( + &self, + indices: impl IntoIterator, + split_points: Vec, + ) -> RangePartitioning { + RangePartitioning::try_new(self.range_ordering(indices), split_points) + .expect("test range partitioning should be valid") + } + + fn range_partitioning( + &self, + indices: impl IntoIterator, + split_points: Vec, + ) -> Partitioning { + Partitioning::Range(self.range(indices, split_points)) + } + + fn range_partitioning_with_ordering( + &self, + ordering: LexOrdering, + split_points: Vec, + ) -> Partitioning { + Partitioning::Range( + RangePartitioning::try_new(ordering, split_points) + .expect("test range partitioning should be valid"), + ) + } + } + + fn assert_satisfaction( + desc: &str, + partitioning: &Partitioning, + required: &Distribution, + eq_properties: &EquivalenceProperties, + expected_with_subset: PartitioningSatisfaction, + expected_without_subset: PartitioningSatisfaction, + ) { + assert_eq!( + partitioning.satisfaction(required, eq_properties, true), + expected_with_subset, + "Failed for {desc} with subset enabled" + ); + assert_eq!( + partitioning.satisfaction(required, eq_properties, false), + expected_without_subset, + "Failed for {desc} with subset disabled" + ); + } #[test] + #[expect( + deprecated, + reason = "test intentionally covers deprecated HashPartitioned compatibility" + )] fn partitioning_satisfy_distribution() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("column_1", DataType::Int64, false), - Field::new("column_2", DataType::Utf8, false), - ])); - - let partition_exprs1: Vec> = vec![ - Arc::new(Column::new_with_schema("column_1", &schema).unwrap()), - Arc::new(Column::new_with_schema("column_2", &schema).unwrap()), - ]; - - let partition_exprs2: Vec> = vec![ - Arc::new(Column::new_with_schema("column_2", &schema).unwrap()), - Arc::new(Column::new_with_schema("column_1", &schema).unwrap()), - ]; + let fixture = PartitioningTestFixture::new(vec![ + ("column_1", DataType::Int64), + ("column_2", DataType::Utf8), + ])?; let distribution_types = vec![ Distribution::UnspecifiedDistribution, Distribution::SinglePartition, - Distribution::HashPartitioned(partition_exprs1.clone()), + Distribution::HashPartitioned(fixture.cols([0, 1])), + Distribution::KeyPartitioned(fixture.cols([0, 1])), ]; let single_partition = Partitioning::UnknownPartitioning(1); let unspecified_partition = Partitioning::UnknownPartitioning(10); let round_robin_partition = Partitioning::RoundRobinBatch(10); - let hash_partition1 = Partitioning::Hash(partition_exprs1, 10); - let hash_partition2 = Partitioning::Hash(partition_exprs2, 10); - let eq_properties = EquivalenceProperties::new(schema); + let hash_partition1 = fixture.hash_partitioning([0, 1], 10); + let hash_partition2 = fixture.hash_partitioning([1, 0], 10); for distribution in distribution_types { let result = ( single_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), unspecified_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), round_robin_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), hash_partition1 - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), hash_partition2 - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), ); @@ -416,7 +806,7 @@ mod tests { Distribution::SinglePartition => { assert_eq!(result, (true, false, false, false, false)) } - Distribution::HashPartitioned(_) => { + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) => { assert_eq!(result, (true, false, false, true, false)) } } @@ -426,239 +816,129 @@ mod tests { } #[test] - fn test_partitioning_satisfy_by_subset() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + #[expect( + deprecated, + reason = "test intentionally covers deprecated HashPartitioned compatibility" + )] + fn deprecated_hash_partitioned_matches_key_partitioned() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let partitioning = fixture.hash_partitioning([0, 1], 4); + let hash_distribution = Distribution::HashPartitioned(fixture.cols([0, 1])); + let key_distribution = fixture.key_distribution([0, 1]); + + assert_eq!( + partitioning.satisfaction(&hash_distribution, &fixture.eq_properties, false), + partitioning.satisfaction(&key_distribution, &fixture.eq_properties, false) + ); + assert_eq!( + hash_distribution.create_partitioning(4), + key_distribution.create_partitioning(4) + ); + + Ok(()) + } + + #[test] + fn hash_partitioning_key_distribution_satisfaction() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); let test_cases = vec![ ( - "Hash([a]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + "exact: KeyPartitioned([a, b]) satisfied by Hash([a, b])", + fixture.hash_partitioning([0, 1], 4), + fixture.key_distribution([0, 1]), + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ), + ( + "subset: KeyPartitioned([a, b]) satisfied by Hash([a])", + fixture.hash_partitioning([0], 4), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + "subset: KeyPartitioned([a, b, c]) satisfied by Hash([b])", + fixture.hash_partitioning([1], 4), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + "subset reordered: KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", + fixture.hash_partitioning([1, 0], 4), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([b]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), - PartitioningSatisfaction::Subset, + "superset: KeyPartitioned([a]) not satisfied by Hash([a, b])", + fixture.hash_partitioning([0, 1], 4), + fixture.key_distribution([0]), + PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([b, a]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), - PartitioningSatisfaction::Subset, + "superset: KeyPartitioned([a, b]) not satisfied by Hash([a, b, c])", + fixture.hash_partitioning([0, 1, 2], 4), + fixture.key_distribution([0, 1]), + PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_current_superset() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let test_cases = vec![ ( - "Hash([a, b]) vs Hash([a])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + "partial overlap: KeyPartitioned([a, b]) not satisfied by Hash([a, c])", + fixture.hash_partitioning([0, 2], 4), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b, c]) vs Hash([a])", - Partitioning::Hash( - vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_c)], - 4, - ), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + "no overlap: KeyPartitioned([b, c]) not satisfied by Hash([a])", + fixture.hash_partitioning([0], 4), + fixture.key_distribution([1, 2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b, c]) vs Hash([a, b])", - Partitioning::Hash( - vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_c)], - 4, - ), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + "unknown partition expr", + Partitioning::Hash(vec![Arc::clone(&unknown)], 4), + fixture.key_distribution([0, 1]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "unknown required expr", + fixture.hash_partitioning([0, 1], 4), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_partial_overlap() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let test_cases = vec![( - "Partial overlap: Hash([a, c]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_c)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a), Arc::clone(&col_b)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - )]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_no_overlap() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let test_cases = vec![ ( - "Hash([a]) vs Hash([b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + "same unknown expr", + Partitioning::Hash(vec![Arc::clone(&unknown)], 4), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b]) vs Hash([c])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_c)]), + "unknown partition expr is not a valid subset", + Partitioning::Hash(vec![Arc::clone(&unknown)], 4), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown), fixture.col(0)]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "empty hash partitioning", + Partitioning::Hash(vec![], 4), + fixture.key_distribution([0]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "empty key distribution", + fixture.hash_partitioning([0], 4), + Distribution::KeyPartitioned(vec![]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -667,16 +947,13 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" + assert_satisfaction( + desc, + &partition, + &required, + &fixture.eq_properties, + expected_with_subset, + expected_without_subset, ); } @@ -684,164 +961,302 @@ mod tests { } #[test] - fn test_partitioning_exact_match() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let test_cases = vec![ + fn key_partitioned_satisfaction_is_exact() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let hash_a = fixture.hash_partitioning([0], 2); + let hash_ab = fixture.hash_partitioning([0, 1], 2); + let range_a = fixture + .range_partitioning([0], vec![int_split_point([10]), int_split_point([20])]); + let range_ab = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 20])]); + + let exact_cases = [ ( - "Hash([a, b]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, + "Hash([a]) vs KeyPartitioned([a])", + &hash_a, + fixture.key_partitioned_distribution([0]), ), ( - "Hash([a]) vs Hash([a])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, + "Range([a]) vs KeyPartitioned([a])", + &range_a, + fixture.key_partitioned_distribution([0]), ), ]; + for (desc, partitioning, requirement) in exact_cases { + for allow_subset in [true, false] { + assert_eq!( + partitioning.satisfaction( + &requirement, + &fixture.eq_properties, + allow_subset, + ), + PartitioningSatisfaction::Exact, + "Failed for {desc} with allow_subset={allow_subset}" + ); + } + } - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); + let subset_cases = [ + ( + "Hash([a]) vs KeyPartitioned([a, b])", + &hash_a, + fixture.key_partitioned_distribution([0, 1]), + ), + ( + "Range([a]) vs KeyPartitioned([a, b])", + &range_a, + fixture.key_partitioned_distribution([0, 1]), + ), + ]; + for (desc, partitioning, requirement) in subset_cases { assert_eq!( - result, expected_with_subset, + partitioning.satisfaction(&requirement, &fixture.eq_properties, true), + PartitioningSatisfaction::Subset, "Failed for {desc} with subset enabled" ); - - let result = partition.satisfaction(&required, &eq_properties, false); assert_eq!( - result, expected_without_subset, + partitioning.satisfaction(&requirement, &fixture.eq_properties, false), + PartitioningSatisfaction::NotSatisfied, "Failed for {desc} with subset disabled" ); } - Ok(()) - } - - #[test] - fn test_partitioning_unknown() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let test_cases = vec![ + let not_satisfied_cases = [ ( - "Hash([unknown]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, + "Range([a]) vs KeyPartitioned([b])", + &range_a, + fixture.key_partitioned_distribution([1]), ), ( - "Hash([a, b]) vs Hash([unknown])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&unknown)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, + "Hash([a, b]) vs KeyPartitioned([a])", + &hash_ab, + fixture.key_partitioned_distribution([0]), ), ( - "Hash([unknown]) vs Hash([unknown])", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&unknown)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, + "Range([a, b]) vs KeyPartitioned([a])", + &range_ab, + fixture.key_partitioned_distribution([0]), ), ]; + for (desc, partitioning, requirement) in not_satisfied_cases { + for allow_subset in [true, false] { + assert_eq!( + partitioning.satisfaction( + &requirement, + &fixture.eq_properties, + allow_subset, + ), + PartitioningSatisfaction::NotSatisfied, + "Failed for {desc} with allow_subset={allow_subset}" + ); + } + } - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); + Ok(()) + } - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } + fn int_split_point(values: impl IntoIterator) -> SplitPoint { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::Int64(Some(value))) + .collect(), + ) + } + + fn assert_range_try_new_error( + ordering: LexOrdering, + split_points: Vec, + expected: &str, + ) { + let error = RangePartitioning::try_new(ordering, split_points) + .unwrap_err() + .to_string(); + assert!(error.contains(expected), "{error}"); + } + + #[test] + fn test_range_partitioning_metadata() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + + let range_partitioning = + fixture.range([0], vec![int_split_point([10]), int_split_point([20])]); + assert_eq!(range_partitioning.ordering()[0].to_string(), "a@0 ASC"); + assert_eq!( + range_partitioning.split_points(), + &[int_split_point([10]), int_split_point([20])] + ); + let partitioning = Partitioning::Range(range_partitioning); + + assert_eq!(partitioning.partition_count(), 3); + assert_eq!( + partitioning.to_string(), + "Range([a@0 ASC], [(10), (20)], 3)" + ); Ok(()) } #[test] - fn test_partitioning_empty_hash() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + fn test_range_partitioning_try_new_validates_split_points() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let asc_a = fixture.range_ordering([0]); + let ordering_ab = fixture.range_ordering([0, 1]); + + assert_range_try_new_error( + ordering_ab.clone(), + vec![int_split_point([10])], + "split point 0 has width 1, but ordering has width 2", + ); + + RangePartitioning::try_new( + [fixture.range_sort_expr(0, SortOptions::new(true, false))].into(), + vec![int_split_point([20]), int_split_point([10])], + )?; + + assert_range_try_new_error( + asc_a, + vec![int_split_point([20]), int_split_point([10])], + "split points must be strictly ordered", + ); + + assert_range_try_new_error( + [fixture.range_sort_expr(0, SortOptions::new(false, false))].into(), + vec![ + SplitPoint::new(vec![ScalarValue::Int64(None)]), + int_split_point([10]), + ], + "split points must be strictly ordered", + ); + + RangePartitioning::try_new( + ordering_ab.clone(), + vec![int_split_point([10, 20]), int_split_point([10, 30])], + )?; + + assert_range_try_new_error( + ordering_ab, + vec![int_split_point([10, 30]), int_split_point([10, 20])], + "split points must be strictly ordered", + ); - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + Ok(()) + } - let test_cases = vec![ - ( - "Hash([]) vs Hash([a])", - Partitioning::Hash(vec![], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), + #[test] + fn test_range_partitioning_project_preserves_or_degrades() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let range_partitioning = fixture.range_partitioning_with_ordering( + [fixture.range_sort_expr(1, SortOptions::new(true, false))].into(), + vec![int_split_point([10])], + ); + + let keep_b_mapping = ProjectionMapping::from_indices(&[1], &fixture.schema)?; + let projected = + range_partitioning.project(&keep_b_mapping, &fixture.eq_properties); + assert_eq!( + projected.to_string(), + "Range([b@0 DESC NULLS LAST], [(10)], 2)" + ); + + let drop_b_mapping = ProjectionMapping::from_indices(&[0], &fixture.schema)?; + let projected = + range_partitioning.project(&drop_b_mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); + + Ok(()) + } + + #[test] + fn test_range_partitioning_project_degrades_if_ordering_collapses() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let target: Arc = Arc::new(Column::new("x", 0)); + let range_partitioning = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); + let mapping = ProjectionMapping::from_iter([ ( - "Hash([a]) vs Hash([])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, + fixture.col(0), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), ), ( - "Hash([]) vs Hash([])", - Partitioning::Hash(vec![], 4), - Distribution::HashPartitioned(vec![]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, + fixture.col(1), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), ), - ]; + ]); - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); + let projected = range_partitioning.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } + Ok(()) + } + + #[test] + fn range_partitioning_key_distribution_satisfaction() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let range_a = fixture.range_partitioning([0], vec![int_split_point([10])]); + let range_ab = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); + + assert_satisfaction( + "exact single key", + &range_a, + &fixture.key_distribution([0]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "exact compound key", + &range_ab, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "subset key", + &range_a, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + assert_satisfaction( + "incompatible key", + &range_a, + &fixture.key_distribution([1]), + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(2))?; + assert_satisfaction( + "equivalent subset key", + &range_a, + &fixture.key_distribution([1, 2]), + &eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?; + assert_satisfaction( + "equivalent exact key", + &range_a, + &fixture.key_distribution([1]), + &eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); Ok(()) } diff --git a/datafusion/physical-expr/src/physical_expr.rs b/datafusion/physical-expr/src/physical_expr.rs index 77ede76e1daa8..6ff5be4e38229 100644 --- a/datafusion/physical-expr/src/physical_expr.rs +++ b/datafusion/physical-expr/src/physical_expr.rs @@ -21,15 +21,17 @@ use crate::expressions::{self, Column}; use crate::{LexOrdering, PhysicalSortExpr, create_physical_expr}; use arrow::compute::SortOptions; -use arrow::datatypes::{Schema, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{DFSchema, HashMap}; +use datafusion_common::{DFSchema, HashMap, ScalarValue, SplitPoint}; use datafusion_common::{Result, plan_err}; use datafusion_expr::execution_props::ExecutionProps; -use datafusion_expr::{Expr, SortExpr}; +use datafusion_expr::{Expr, Partitioning as LogicalPartitioning, SortExpr}; +use datafusion_expr_common::casts::try_cast_literal_to_type; use itertools::izip; // Exports: +use crate::{Partitioning, RangePartitioning}; pub(crate) use datafusion_physical_expr_common::physical_expr::PhysicalExpr; /// Adds the `offset` value to `Column` indices inside `expr`. This function is @@ -216,6 +218,99 @@ pub fn create_physical_sort_exprs( .collect() } +/// Create physical partitioning from logical partitioning. +pub fn create_physical_partitioning( + partitioning: &LogicalPartitioning, + input_dfschema: &DFSchema, + execution_props: &ExecutionProps, +) -> Result { + match partitioning { + LogicalPartitioning::RoundRobinBatch(n) => Ok(Partitioning::RoundRobinBatch(*n)), + LogicalPartitioning::Hash(exprs, partition_count) => { + let exprs = exprs + .iter() + .map(|expr| create_physical_expr(expr, input_dfschema, execution_props)) + .collect::>>()?; + Ok(Partitioning::Hash(exprs, *partition_count)) + } + LogicalPartitioning::Range(range) => { + let ordering = create_physical_sort_exprs( + range.ordering(), + input_dfschema, + execution_props, + )?; + let Some(ordering) = LexOrdering::new(ordering) else { + return plan_err!("Range partitioning requires non-empty ordering"); + }; + let split_points = normalize_range_split_points( + &ordering, + range.split_points(), + input_dfschema.as_arrow(), + )?; + let range = RangePartitioning::try_new(ordering, split_points)?; + Ok(Partitioning::Range(range)) + } + LogicalPartitioning::DistributeBy(_) => { + datafusion_common::not_impl_err!( + "Physical plan does not support DistributeBy partitioning" + ) + } + } +} + +fn normalize_range_split_points( + ordering: &LexOrdering, + split_points: &[SplitPoint], + schema: &Schema, +) -> Result> { + split_points + .iter() + .enumerate() + .map(|(split_idx, split_point)| { + let values = split_point + .values() + .iter() + .zip(ordering.iter()) + .enumerate() + .map(|(value_idx, (value, sort_expr))| { + let target_type = sort_expr.expr.data_type(schema)?; + normalize_range_split_point_value( + value, + &target_type, + split_idx, + value_idx, + ) + }) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect() +} + +fn normalize_range_split_point_value( + value: &ScalarValue, + target_type: &DataType, + split_idx: usize, + value_idx: usize, +) -> Result { + let value_type = value.data_type(); + if &value_type == target_type { + return Ok(value.clone()); + } + + if let Some(casted) = try_cast_literal_to_type(value, target_type) { + // Split points define physical partition boundaries, so normalization + // must reject casts that would change the advertised boundary. + if try_cast_literal_to_type(&casted, &value_type).as_ref() == Some(value) { + return Ok(casted); + } + } + + plan_err!( + "Range output partitioning split point {split_idx} value {value_idx} with type {value_type} cannot be represented exactly as ordering expression type {target_type}" + ) +} + pub fn add_offset_to_physical_sort_exprs( sort_exprs: impl IntoIterator, offset: isize, diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/enforce_distribution.rs index c522867c05196..3b520fd3de60c 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -41,7 +41,8 @@ use datafusion_expr::logical_plan::{Aggregate, JoinType}; use datafusion_physical_expr::expressions::{Column, NoOp}; use datafusion_physical_expr::utils::map_columns_before_projection; use datafusion_physical_expr::{ - EquivalenceProperties, PhysicalExpr, PhysicalExprRef, physical_exprs_equal, + EquivalenceProperties, OrderingRequirements, PhysicalExpr, PhysicalExprRef, + physical_exprs_equal, }; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::aggregates::{ @@ -59,7 +60,10 @@ use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::union::{InterleaveExec, UnionExec, can_interleave}; use datafusion_physical_plan::windows::WindowAggExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, get_best_fitting_window}; -use datafusion_physical_plan::{Distribution, ExecutionPlan, Partitioning}; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, +}; use itertools::izip; @@ -853,72 +857,34 @@ fn add_roundrobin_on_top( } } -/// Adds a hash repartition operator: -/// - to increase parallelism, and/or -/// - to satisfy requirements of the subsequent operators. -/// -/// Repartition(Hash) is added on top of operator `input`. -/// -/// # Arguments -/// -/// * `input`: Current node. -/// * `hash_exprs`: Stores Physical Exprs that are used during hashing. -/// * `n_target`: desired target partition number, if partition number of the -/// current executor is less than this value. Partition number will be increased. -/// * `allow_subset_satisfy_partitioning`: Whether to allow subset partitioning logic in satisfaction checks. -/// Set to `false` for partitioned hash joins to ensure exact hash matching. -/// -/// # Returns -/// -/// A [`Result`] object that contains new execution plan where the desired -/// distribution is satisfied by adding a Hash repartition. -fn add_hash_on_top( - input: DistributionContext, - hash_exprs: Vec>, - n_target: usize, +// Partial aggregates require unspecified input distribution, but their output +// may already satisfy the final aggregate's key distribution because partial +// aggregation preserves/projects input partitioning. Keep that reusable output +// partitioning intact when preserve_file_partitions would otherwise insert +// RoundRobin below the partial aggregate. +fn partial_aggregate_output_satisfies_final_partitioning( + plan: &Arc, allow_subset_satisfy_partitioning: bool, -) -> Result { - // Early return if hash repartition is unnecessary - // `RepartitionExec: partitioning=Hash([...], 1), input_partitions=1` is unnecessary. - if n_target == 1 && input.plan.output_partitioning().partition_count() == 1 { - return Ok(input); - } - - let dist = Distribution::HashPartitioned(hash_exprs); - let satisfaction = input.plan.output_partitioning().satisfaction( - &dist, - input.plan.equivalence_properties(), - allow_subset_satisfy_partitioning, - ); - - // Add hash repartitioning when: - // - When subset satisfaction is enabled (current >= threshold): only repartition if not satisfied - // - When below threshold (current < threshold): repartition if expressions don't match OR to increase parallelism - let needs_repartition = if allow_subset_satisfy_partitioning { - !satisfaction.is_satisfied() - } else { - !satisfaction.is_satisfied() - || n_target > input.plan.output_partitioning().partition_count() +) -> bool { + let Some(aggregate) = plan.downcast_ref::() else { + return false; }; - - if needs_repartition { - // When there is an existing ordering, we preserve ordering during - // repartition. This will be rolled back in the future if any of the - // following conditions is true: - // - Preserving ordering is not helpful in terms of satisfying ordering - // requirements. - // - Usage of order preserving variants is not desirable (per the flag - // `config.optimizer.prefer_existing_sort`). - let partitioning = dist.create_partitioning(n_target); - let repartition = - RepartitionExec::try_new(Arc::clone(&input.plan), partitioning)? - .with_preserve_order(); - let plan = Arc::new(repartition) as _; - - return Ok(DistributionContext::new(plan, true, vec![input])); + if aggregate.mode() != &AggregateMode::Partial + || aggregate.group_expr().is_empty() + || aggregate.group_expr().has_grouping_set() + { + return false; } - Ok(input) + let key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr()); + + plan.output_partitioning() + .satisfaction( + &key_distribution, + plan.equivalence_properties(), + allow_subset_satisfy_partitioning, + ) + .is_satisfied() } /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator @@ -1094,6 +1060,27 @@ struct RepartitionRequirementStatus { hash_necessary: bool, } +fn requirement_includes_grouping_id(requirement: &Distribution) -> bool { + // Grouping set aggregates (ROLLUP, CUBE, GROUPING SETS) require exact hash + // partitioning on all group columns including __grouping_id to ensure + // partial aggregates from different partitions are correctly combined. + requirement.key_exprs().is_some_and(|exprs| { + exprs.iter().any(|expr| { + (expr.as_ref() as &dyn Any) + .downcast_ref::() + .is_some_and(|col| col.name() == Aggregate::INTERNAL_GROUPING_ID) + }) + }) +} + +/// Per-child state while enforcing a parent's distribution requirements. +struct DistributionChildState { + context: DistributionContext, + required_input_ordering: Option, + maintains_input_order: bool, + requirement: Distribution, +} + /// Calculates the `RepartitionRequirementStatus` for each children to generate /// consistent and sensible (in terms of performance) distribution requirements. /// As an example, a hash join's left (build) child might produce @@ -1133,7 +1120,7 @@ fn get_repartition_requirement_status( let mut needs_alignment = false; let children = plan.children(); let rr_beneficial = plan.benefits_from_input_partitioning(); - let requirements = plan.required_input_distribution(); + let requirements = plan.input_distribution_requirements().into_per_child(); let mut repartition_status_flags = vec![]; for (child, requirement, roundrobin_beneficial) in izip!(children.into_iter(), requirements, rr_beneficial) @@ -1146,30 +1133,31 @@ fn get_repartition_requirement_status( Precision::Inexact(n_rows) => !should_use_estimates || (n_rows > batch_size), Precision::Absent => true, }; - let is_hash = matches!(requirement, Distribution::HashPartitioned(_)); - // Hash re-partitioning is necessary when the input has more than one - // partitions: + let is_partitioned_requirement = requirement.key_exprs().is_some(); + // Hash repartitioning is necessary when the input has more than one + // partition. let multi_partitions = child.output_partitioning().partition_count() > 1; let roundrobin_sensible = roundrobin_beneficial && roundrobin_beneficial_stats; - needs_alignment |= is_hash && (multi_partitions || roundrobin_sensible); + needs_alignment |= + is_partitioned_requirement && (multi_partitions || roundrobin_sensible); repartition_status_flags.push(( - is_hash, + is_partitioned_requirement, RepartitionRequirementStatus { requirement, roundrobin_beneficial, roundrobin_beneficial_stats, - hash_necessary: is_hash && multi_partitions, + hash_necessary: is_partitioned_requirement && multi_partitions, }, )); } - // Align hash necessary flags for hash partitions to generate consistent + // Align hash necessary flags for key partitions to generate consistent // hash partitions at each children: if needs_alignment { - // When there is at least one hash requirement that is necessary or - // beneficial according to statistics, make all children require hash - // repartitioning: - for (is_hash, status) in &mut repartition_status_flags { - if *is_hash { + // When there is at least one key-partitioned requirement that is necessary + // or beneficial according to statistics, make all key-partitioned + // children require hash repartitioning: + for (is_partitioned_requirement, status) in &mut repartition_status_flags { + if *is_partitioned_requirement { status.hash_necessary = true; } } @@ -1180,6 +1168,90 @@ fn get_repartition_requirement_status( .collect()) } +/// Enforce cross-child distribution relationships after each child has already +/// satisfied its own distribution requirement. +/// +/// See [`InputDistributionRequirements`] for the distinction between +/// independent per-child requirements and co-partitioned child relationships. +/// +/// Currently, unsatisfied co-partitioning is repaired by hash repartitioning +/// key-partitioned children and other relationship kinds are rejected. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] +fn enforce_distribution_relationships( + plan_name: &str, + input_distributions: &InputDistributionRequirements, + children: &mut [DistributionChildState], + target_partitions: usize, +) -> Result<()> { + let mut repartitioned_for_relationship = vec![false; children.len()]; + + loop { + let child_plan_refs = children + .iter() + .map(|child| child.context.plan.as_ref()) + .collect::>(); + let unsatisfied_children = input_distributions + .unsatisfied_co_partitioned_children(plan_name, &child_plan_refs)?; + + if unsatisfied_children.is_empty() { + return Ok(()); + } + + let mut changed = false; + for child_idx in unsatisfied_children { + if repartitioned_for_relationship[child_idx] { + continue; + } + + let (Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs)) = &children[child_idx].requirement + else { + continue; + }; + + let already_target_hash = matches!( + children[child_idx].context.plan.output_partitioning(), + Partitioning::Hash(_, partition_count) if *partition_count == target_partitions + ) && input_distributions + .child_satisfaction( + child_idx, + children[child_idx].context.plan.as_ref(), + ChildSatisfactionOptions::new(), + )? + .is_satisfied(); + + if already_target_hash { + continue; + } + + let partitioning = Distribution::KeyPartitioned(exprs.to_vec()) + .create_partitioning(target_partitions); + let repartition = RepartitionExec::try_new( + Arc::clone(&children[child_idx].context.plan), + partitioning, + )? + .with_preserve_order(); + let plan = Arc::new(repartition) as _; + let original_child = std::mem::replace( + &mut children[child_idx].context, + DistributionContext::new(plan, true, vec![]), + ); + children[child_idx].context.children = vec![original_child]; + repartitioned_for_relationship[child_idx] = true; + changed = true; + } + + if !changed { + return datafusion_common::internal_err!( + "{plan_name} has distribution relationships that could not be enforced" + ); + } + } +} + /// This function checks whether we need to add additional data exchange /// operators to satisfy distribution requirements. Since this function /// takes care of such requirements, we should avoid manually adding data @@ -1188,6 +1260,10 @@ fn get_repartition_requirement_status( /// This function is intended to be used in a bottom up traversal, as it /// can first repartition (or newly partition) at the datasources -- these /// source partitions may be later repartitioned with additional data exchange operators. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] pub fn ensure_distribution( dist_context: DistributionContext, config: &ConfigOptions, @@ -1278,6 +1354,7 @@ pub fn ensure_distribution( .is_some_and(|join| join.mode == PartitionMode::Partitioned) || plan.is::(); + let input_distributions = plan.input_distribution_requirements(); let repartition_status_flags = get_repartition_requirement_status(&plan, batch_size, should_use_estimates)?; // This loop iterates over all the children to: @@ -1285,7 +1362,8 @@ pub fn ensure_distribution( // - Satisfy the distribution requirements of every child, if it is not // already satisfied. // We store the updated children in `new_children`. - let children = izip!( + let mut children = izip!( + 0..children.len(), children.into_iter(), plan.required_input_ordering(), plan.maintains_input_order(), @@ -1293,6 +1371,7 @@ pub fn ensure_distribution( ) .map( |( + child_idx, mut child, required_input_ordering, maintains, @@ -1303,33 +1382,18 @@ pub fn ensure_distribution( hash_necessary, }, )| { - let increases_partition_count = - child.plan.output_partitioning().partition_count() < target_partitions; - - let add_roundrobin = enable_round_robin - // Operator benefits from partitioning (e.g. filter): - && roundrobin_beneficial - && roundrobin_beneficial_stats - // Unless partitioning increases the partition count, it is not beneficial: - && increases_partition_count; - // Allow subset satisfaction when: // 1. Current partition count >= threshold // 2. Not a partitioned join since must use exact hash matching for joins // 3. Not a grouping set aggregate (requires exact hash including __grouping_id) + // + // Partitioned joins still require exact satisfaction. If that + // exact check already passes, preserve_file_partitions can skip + // repartitioning whose only purpose is increasing partition count. let current_partitions = child.plan.output_partitioning().partition_count(); - - // Check if the hash partitioning requirement includes __grouping_id column. - // Grouping set aggregates (ROLLUP, CUBE, GROUPING SETS) require exact hash - // partitioning on all group columns including __grouping_id to ensure partial - // aggregates from different partitions are correctly combined. - let requires_grouping_id = matches!(&requirement, Distribution::HashPartitioned(exprs) - if exprs.iter().any(|expr| { - (expr.as_ref() as &dyn Any) - .downcast_ref::() - .is_some_and(|col| col.name() == Aggregate::INTERNAL_GROUPING_ID) - }) - ); + let preserve_file_partition_threshold_met = + config.optimizer.preserve_file_partitions > 0 + && current_partitions >= config.optimizer.preserve_file_partitions; let allow_subset_satisfy_partitioning = (current_partitions >= subset_satisfaction_threshold @@ -1337,10 +1401,27 @@ pub fn ensure_distribution( // partitioning to the optimizer. Respect it when the only // reason to repartition would be to increase partition count // beyond the preserved file-group count. - || (config.optimizer.preserve_file_partitions > 0 + || (preserve_file_partition_threshold_met && current_partitions < target_partitions)) && !is_partitioned_join - && !requires_grouping_id; + && !requirement_includes_grouping_id(&requirement); + + let increases_partition_count = current_partitions < target_partitions; + + let preserve_partial_aggregate_partitioning = + preserve_file_partition_threshold_met + && partial_aggregate_output_satisfies_final_partitioning( + &plan, + allow_subset_satisfy_partitioning, + ); + + let add_roundrobin = enable_round_robin + // Operator benefits from partitioning (e.g. filter): + && roundrobin_beneficial + && roundrobin_beneficial_stats + // Unless partitioning increases the partition count, it is not beneficial: + && increases_partition_count + && !preserve_partial_aggregate_partitioning; // When `repartition_file_scans` is set, attempt to increase // parallelism at the source. @@ -1361,16 +1442,49 @@ pub fn ensure_distribution( Distribution::SinglePartition => { child = add_merge_on_top(child); } - Distribution::HashPartitioned(exprs) => { + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) => { + let child_partitions = + child.plan.output_partitioning().partition_count(); + let partitioning_satisfied = input_distributions + .child_satisfaction( + child_idx, + child.plan.as_ref(), + ChildSatisfactionOptions::new() + .with_allow_subset(allow_subset_satisfy_partitioning), + )? + .is_satisfied(); + let preserve_satisfying_file_partitioning = + preserve_file_partition_threshold_met + && !requirement_includes_grouping_id(&requirement) + && partitioning_satisfied + && target_partitions > child_partitions; + + // When subset satisfaction is enabled, preserve an + // already-satisfying partitioning. Otherwise, hash + // repartition may also increase parallelism. + let needs_hash_repartition = if allow_subset_satisfy_partitioning { + !partitioning_satisfied + } else { + !partitioning_satisfied + || (target_partitions > child_partitions + && !preserve_satisfying_file_partitioning) + }; + let should_add_hash_repartition = + hash_necessary && needs_hash_repartition; + // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background // When inserting hash is necessary to satisfy hash requirement, insert hash repartition. - if hash_necessary { - child = add_hash_on_top( - child, - exprs.to_vec(), - target_partitions, - allow_subset_satisfy_partitioning, - )?; + if should_add_hash_repartition { + let partitioning = Distribution::KeyPartitioned(exprs.to_vec()) + .create_partitioning(target_partitions); + let repartition = RepartitionExec::try_new( + Arc::clone(&child.plan), + partitioning, + )? + .with_preserve_order(); + let plan = Arc::new(repartition) as _; + child = DistributionContext::new(plan, true, vec![child]); } } Distribution::UnspecifiedDistribution => { @@ -1382,73 +1496,101 @@ pub fn ensure_distribution( } }; - let streaming_benefit = if child.data { - preserving_order_enables_streaming(&plan, &child.plan)? - } else { - false - }; + Ok(DistributionChildState { + context: child, + required_input_ordering, + maintains_input_order: maintains, + requirement, + }) + }, + ) + .collect::>>()?; - // There is an ordering requirement of the operator: - if let Some(required_input_ordering) = required_input_ordering { - // Either: - // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or - // - using order preserving variant is not desirable. - let sort_req = required_input_ordering.into_single(); - let ordering_satisfied = child - .plan - .equivalence_properties() - .ordering_satisfy_requirement(sort_req.clone())?; - - if (!ordering_satisfied || !order_preserving_variants_desirable) - && !streaming_benefit - && child.data - { - child = replace_order_preserving_variants(child)?; - // If ordering requirements were satisfied before repartitioning, - // make sure ordering requirements are still satisfied after. - if ordering_satisfied { - // Make sure to satisfy ordering requirement: - child = add_sort_above_with_check( - child, - sort_req, - plan.downcast_ref::() - .map(|output| output.fetch()) - .unwrap_or(None), - )?; - } - } - // Stop tracking distribution changing operators - child.data = false; - } else { - let streaming_benefit = if child.data { - preserving_order_enables_streaming(&plan, &child.plan)? + // This is called after each child satisfies its own distribution requirement. + // It enforces relationships between child partition layouts for multi-child + // operators that process matching partition indexes together. + enforce_distribution_relationships( + plan.name(), + &input_distributions, + &mut children, + target_partitions, + )?; + + let children = children + .into_iter() + .map( + |DistributionChildState { + mut context, + required_input_ordering, + maintains_input_order, + requirement, + }| { + let streaming_benefit = if context.data { + preserving_order_enables_streaming(&plan, &context.plan)? } else { false }; - // no ordering requirement - match requirement { - // Operator requires specific distribution. - Distribution::SinglePartition | Distribution::HashPartitioned(_) => { - // If the parent doesn't maintain input order, preserving - // ordering is pointless. However, if it does maintain - // input order, we keep order-preserving variants so - // ordering can flow through to ancestors that need it. - if !maintains && !streaming_benefit { - child = replace_order_preserving_variants(child)?; + + // There is an ordering requirement of the operator: + if let Some(required_input_ordering) = required_input_ordering { + // Either: + // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or + // - using order preserving variant is not desirable. + let sort_req = required_input_ordering.into_single(); + let ordering_satisfied = context + .plan + .equivalence_properties() + .ordering_satisfy_requirement(sort_req.clone())?; + + if (!ordering_satisfied || !order_preserving_variants_desirable) + && !streaming_benefit + && context.data + { + context = replace_order_preserving_variants(context)?; + // If ordering requirements were satisfied before repartitioning, + // make sure ordering requirements are still satisfied after. + if ordering_satisfied { + // Make sure to satisfy ordering requirement: + context = add_sort_above_with_check( + context, + sort_req, + plan.downcast_ref::() + .map(|output| output.fetch()) + .unwrap_or(None), + )?; } } - Distribution::UnspecifiedDistribution => { - // Since ordering is lost, trying to preserve ordering is pointless - if !maintains || plan.is::() { - child = replace_order_preserving_variants(child)?; + // Stop tracking distribution changing operators + context.data = false; + } else { + // no ordering requirement + match requirement { + // Operator requires specific distribution. + Distribution::SinglePartition + | Distribution::HashPartitioned(_) + | Distribution::KeyPartitioned(_) => { + // If the parent doesn't maintain input order, preserving + // ordering is pointless. However, if it does maintain + // input order, we keep order-preserving variants so + // ordering can flow through to ancestors that need it. + if !maintains_input_order && !streaming_benefit { + context = replace_order_preserving_variants(context)?; + } + } + Distribution::UnspecifiedDistribution => { + // Since ordering is lost, trying to preserve ordering is pointless + if !maintains_input_order + || plan.is::() + { + context = replace_order_preserving_variants(context)?; + } } } } - } - Ok(child) - }, - ) - .collect::>>()?; + Ok(context) + }, + ) + .collect::>>()?; let children_plans = children .iter() @@ -1514,7 +1656,8 @@ fn update_children(mut dist_context: DistributionContext) -> Result Result { if node.data { let requires_single_partition = matches!( - parent.required_input_distribution()[child_idx], - Distribution::SinglePartition + parent + .input_distribution_requirements() + .child_distribution(child_idx), + Some(Distribution::SinglePartition) ); node = remove_corresponding_sort_from_sub_plan(node, requires_single_partition)?; } @@ -749,7 +757,7 @@ fn remove_corresponding_sort_from_sub_plan( } } else { let mut any_connection = false; - let required_dist = node.plan.required_input_distribution(); + let required_dist = node.plan.input_distribution_requirements().into_per_child(); node.children = node .children .into_iter() diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 24eb3af5f564c..48cc41b8a21c4 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -28,7 +28,7 @@ use crate::PhysicalOptimizerRule; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{Result, Statistics}; +use datafusion_common::{Result, Statistics, internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::Distribution; use datafusion_physical_expr_common::sort_expr::OrderingRequirements; @@ -205,7 +205,15 @@ impl ExecutionPlan for OutputRequirementExec { } fn required_input_distribution(&self) -> Vec { - vec![self.dist_requirement.clone()] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + self.dist_requirement.clone(), + ]) } fn maintains_input_order(&self) -> Vec { @@ -244,6 +252,10 @@ impl ExecutionPlan for OutputRequirementExec { self.input.partition_statistics(partition) } + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] fn try_swapping_with_projection( &self, projection: &ProjectionExec, @@ -268,8 +280,12 @@ impl ExecutionPlan for OutputRequirementExec { requirements = OrderingRequirements::new_alternatives(updated_reqs, soft); } - let dist_req = match &self.required_input_distribution()[0] { - Distribution::HashPartitioned(exprs) => { + let input_distributions = self.input_distribution_requirements(); + let dist_req = match input_distributions.child_distribution(0) { + Some( + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs), + ) => { let mut updated_exprs = vec![]; for expr in exprs { let Some(new_expr) = update_expr(expr, projection.expr(), false)? @@ -278,9 +294,14 @@ impl ExecutionPlan for OutputRequirementExec { }; updated_exprs.push(new_expr); } - Distribution::HashPartitioned(updated_exprs) + Distribution::KeyPartitioned(updated_exprs) + } + Some(dist) => dist.clone(), + None => { + return internal_err!( + "OutputRequirementExec missing input distribution requirement" + ); } - dist => dist.clone(), }; make_with_child(projection, &self.input()).map(|input| { @@ -355,7 +376,10 @@ fn require_top_ordering_helper( // In case of constant columns, output ordering of the `SortExec` would // be an empty set. Therefore; we check the sort expression field to // assign the requirements. - let req_dist = sort_exec.required_input_distribution().swap_remove(0); + let req_dist = sort_exec + .input_distribution_requirements() + .into_per_child() + .swap_remove(0); let req_ordering = sort_exec.expr(); let reqs = OrderingRequirements::from(req_ordering.clone()); let fetch = sort_exec.fetch(); diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 40c6245d894d4..713213b70612d 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -30,9 +30,13 @@ use datafusion_common::config::{ConfigOptions, OptimizerOptions}; use datafusion_common::plan_err; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported}; -use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, InvariantLevel, +}; use datafusion_physical_plan::joins::SymmetricHashJoinExec; -use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string}; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, ExecutionPlanProperties, get_plan_string, +}; use crate::PhysicalOptimizerRule; use datafusion_physical_expr_common::sort_expr::format_physical_sort_requirement_list; @@ -141,11 +145,12 @@ pub fn check_plan_sanity( optimizer_options: &OptimizerOptions, ) -> Result<()> { check_finiteness_requirements(plan.as_ref(), optimizer_options)?; + let input_distributions = plan.input_distribution_requirements(); for ((idx, child), sort_req, dist_req) in izip!( plan.children().into_iter().enumerate(), plan.required_input_ordering(), - plan.required_input_distribution(), + input_distributions.per_child_distributions(), ) { let child_eq_props = child.equivalence_properties(); if let Some(sort_req) = sort_req { @@ -162,9 +167,12 @@ pub fn check_plan_sanity( } } - if !child - .output_partitioning() - .satisfaction(&dist_req, child_eq_props, true) + if !input_distributions + .child_satisfaction( + idx, + child.as_ref(), + ChildSatisfactionOptions::new().with_allow_subset(true), + )? .is_satisfied() { let plan_str = get_plan_string(plan); @@ -178,6 +186,8 @@ pub fn check_plan_sanity( } } + plan.check_invariants(InvariantLevel::Executable)?; + Ok(()) } diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 40dbddfbdf9fb..9459376054fda 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -46,6 +46,7 @@ use datafusion_physical_expr::window::StandardWindowExpr; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; @@ -119,17 +120,19 @@ impl WindowTopN { // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.) let (col_idx, limit_n) = extract_window_limit(filter.predicate())?; - // Step 3: Walk through optional ProjectionExec to find BoundedWindowAggExec + // Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec let child = filter.input(); - let (window_exec, proj_between) = find_window_below(child)?; + let (window_exec, intermediates) = find_window_below(child)?; // Step 4: Verify col_idx references a ROW_NUMBER window output column - let input_field_count = window_exec.input().schema().fields().len(); + let window_exec_typed = window_exec.downcast_ref::()?; + let sort_exec = window_exec_typed.input().downcast_ref::()?; + let input_field_count = window_exec_typed.input().schema().fields().len(); if col_idx < input_field_count { return None; // Filter is on an input column, not a window column } let window_expr_idx = col_idx - input_field_count; - let window_exprs = window_exec.window_expr(); + let window_exprs = window_exec_typed.window_expr(); if window_expr_idx >= window_exprs.len() { return None; } @@ -137,8 +140,7 @@ impl WindowTopN { return None; } - // Step 5: Verify child of window is SortExec - let sort_exec = window_exec.input().downcast_ref::()?; + // Step 5: child of window is SortExec (verified above) let sort_child = sort_exec.input(); // Step 6: Determine partition_prefix_len from the window expression @@ -161,28 +163,19 @@ impl WindowTopN { .ok()?; // Step 8: Rebuild window with new child - let new_window = Arc::clone(&child_as_arc(window_exec)) + let mut result = window_exec .with_new_children(vec![Arc::new(partitioned_topk)]) .ok()?; - // Step 9: If ProjectionExec was between Filter and Window, rebuild it - let result = match proj_between { - Some(proj) => Arc::clone(&child_as_arc(proj)) - .with_new_children(vec![new_window]) - .ok()?, - None => new_window, - }; + // Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec) + for node in intermediates.into_iter().rev() { + result = node.with_new_children(vec![result]).ok()?; + } Some(result) } } -/// Helper to get an `Arc` from a reference. -/// We need this because `with_new_children` takes `Arc`. -fn child_as_arc(plan: &T) -> Arc { - Arc::new(plan.clone()) -} - impl PhysicalOptimizerRule for WindowTopN { fn optimize( &self, @@ -303,29 +296,32 @@ fn is_row_number(expr: &Arc) - udf.fun().name() == "row_number" } +type PlanAndIntermediates = (Arc, Vec>); + /// Walk below a plan node looking for a [`BoundedWindowAggExec`]. /// -/// Handles two cases: -/// - Direct child: `FilterExec → BoundedWindowAggExec` -/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec` +/// Handles sequences of `ProjectionExec` and `RepartitionExec`. +/// This is safe because `PartitionedTopKExec` can be pushed below them: +/// projections only provide aliases, and pushing the limit below repartitions +/// is safe because the limit is computed per-partition. /// -/// Returns the window exec and an optional `ProjectionExec` in between, -/// or `None` if no `BoundedWindowAggExec` is found within one or two levels. -fn find_window_below( - plan: &Arc, -) -> Option<(&BoundedWindowAggExec, Option<&ProjectionExec>)> { - // Direct child is BoundedWindowAggExec - if let Some(window) = plan.downcast_ref::() { - return Some((window, None)); - } +/// Returns the window exec and a list of intermediate nodes to rebuild, +/// or `None` if no `BoundedWindowAggExec` is found. +fn find_window_below(plan: &Arc) -> Option { + let mut current = Arc::clone(plan); + let mut intermediates = Vec::new(); - // Child is ProjectionExec with BoundedWindowAggExec below - if let Some(proj) = plan.downcast_ref::() { - let proj_child = proj.input(); - if let Some(window) = proj_child.downcast_ref::() { - return Some((window, Some(proj))); + loop { + if current.downcast_ref::().is_some() { + return Some((current, intermediates)); + } else if current.downcast_ref::().is_some() + || current.downcast_ref::().is_some() + { + let next = Arc::clone(current.children().first()?); + intermediates.push(current); + current = next; + } else { + return None; } } - - None } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 5e6b8505764a2..61e88f0f75db3 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -32,8 +32,8 @@ use crate::filter_pushdown::{ }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode, - SendableRecordBatchStream, Statistics, check_if_same_properties, + DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + InputOrderMode, SendableRecordBatchStream, Statistics, check_if_same_properties, }; use datafusion_common::config::ConfigOptions; use datafusion_physical_expr::utils::collect_columns; @@ -1544,17 +1544,21 @@ impl ExecutionPlan for AggregateExec { } fn required_input_distribution(&self) -> Vec { - match &self.mode { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(match &self.mode { AggregateMode::Partial | AggregateMode::PartialReduce => { vec![Distribution::UnspecifiedDistribution] } AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned => { - vec![Distribution::HashPartitioned(self.group_by.input_exprs())] + vec![Distribution::KeyPartitioned(self.group_by.input_exprs())] } AggregateMode::Final | AggregateMode::Single => { vec![Distribution::SinglePartition] } - } + }) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 582af8f1e3dae..708f6b4c95ac3 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -144,7 +144,13 @@ impl ExecutionPlan for AnalyzeExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn with_new_children( diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs new file mode 100644 index 0000000000000..6405b1f121ef7 --- /dev/null +++ b/datafusion/physical-plan/src/distribution_requirements.rs @@ -0,0 +1,359 @@ +// 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. + +//! Input distribution requirements for physical execution plans. + +use datafusion_common::{Result, internal_err}; +use datafusion_physical_expr::{Distribution, Partitioning, PartitioningSatisfaction}; + +use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; + +/// Distribution requirements for an [`ExecutionPlan`]'s inputs. +/// +/// [`InputDistributionRequirements`] describes what distribution an operator +/// requires from each child. +/// +/// - [`Self::new`] describes independent per-child requirements. +/// - [`Self::co_partitioned`] additionally requires child partitions with the +/// same index to cover compatible key ranges. +/// +/// For a single-input aggregate: +/// +/// ```text +/// AggregateExec +/// child 0 requirement: KeyPartitioned(group_exprs) +/// ``` +/// +/// each input partition can aggregate its own key domain independently. +/// +/// For a partitioned join: +/// +/// ```text +/// HashJoinExec +/// child 0 requirement: KeyPartitioned(left_keys) +/// child 1 requirement: KeyPartitioned(right_keys) +/// +/// partition 0: join(left partition 0, right partition 0) +/// partition 1: join(left partition 1, right partition 1) +/// partition 2: join(left partition 2, right partition 2) +/// ``` +/// +/// each child must satisfy its own key requirement. In addition, matching +/// partition indexes must be safe to process together. +#[non_exhaustive] +#[derive(Debug, Clone)] +pub struct InputDistributionRequirements { + /// Per-child distribution requirements, indexed by child position. + children: Vec, + /// Child indexes that must also have compatible partition layouts. + co_partitioned: Option>, +} + +/// Options for checking child distribution satisfaction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ChildSatisfactionOptions { + allow_subset: bool, +} + +impl ChildSatisfactionOptions { + /// Create default satisfaction options. + pub fn new() -> Self { + Self::default() + } + + /// Allow a child partitioning whose key expressions are a subset of the + /// required key expressions to satisfy the requirement. + pub fn with_allow_subset(mut self, allow_subset: bool) -> Self { + self.allow_subset = allow_subset; + self + } + + /// Whether subset satisfaction is enabled. + pub fn allow_subset(&self) -> bool { + self.allow_subset + } +} + +impl InputDistributionRequirements { + /// Create independent per-child requirements. + pub fn new(per_child: Vec) -> Self { + let children = per_child + .into_iter() + .map(|distribution| ChildDistributionRequirement { distribution }) + .collect(); + + Self { + children, + co_partitioned: None, + } + } + + /// Create a requirement that all children are co-partitioned. + /// + /// Each child must satisfy its own [`Distribution`]. Matching partition + /// indexes are processed together: + /// + /// ```text + /// left: Range(left.a ASC, split_points=[10, 20]) + /// right: Range(right.x ASC, split_points=[10, 20]) + /// + /// partition 0 from both sides contains keys before 10 + /// partition 1 from both sides contains keys in [10, 20) + /// partition 2 from both sides contains keys at/after 20 + /// ``` + /// + /// If the split points differ, partition `i` from one side no longer covers + /// the same key range as partition `i` from the other side. + pub fn co_partitioned(per_child: Vec) -> Self { + debug_assert!( + per_child.len() >= 2, + "co-partitioned distribution requirements need at least two children" + ); + let co_partitioned = (0..per_child.len()).collect(); + let mut result = Self::new(per_child); + result.co_partitioned = Some(co_partitioned); + result + } + + /// Return the per-child distribution requirements. + pub fn per_child_distributions( + &self, + ) -> impl ExactSizeIterator + '_ { + self.children.iter().map(|child| &child.distribution) + } + + /// Return the distribution requirement for a child. + pub fn child_distribution(&self, child_idx: usize) -> Option<&Distribution> { + self.children + .get(child_idx) + .map(|child| &child.distribution) + } + + /// Return the per-child distribution requirements. + /// + /// WARNING: This intentionally drops any grouped relationship. + pub fn into_per_child(self) -> Vec { + self.children + .into_iter() + .map(|child| child.distribution) + .collect() + } + + /// Returns how a child satisfies its distribution requirement. + /// + /// This preserves the requirement set's satisfaction policy. + pub fn child_satisfaction( + &self, + child_idx: usize, + child: &dyn ExecutionPlan, + options: ChildSatisfactionOptions, + ) -> Result { + let Some(requirement) = self.children.get(child_idx) else { + return internal_err!( + "missing distribution requirement for child {child_idx}" + ); + }; + + Ok(child.output_partitioning().satisfaction( + &requirement.distribution, + child.equivalence_properties(), + options.allow_subset(), + )) + } + + /// Return child indexes whose co-partitioning requirements are + /// unsatisfied by the provided candidate children. + /// + /// Independent per-child requirements are intentionally ignored here, use + /// [`Self::child_satisfaction`] for those checks. An empty result means all + /// co-partitioning requirements are satisfied. + #[doc(hidden)] + pub fn unsatisfied_co_partitioned_children( + &self, + plan_name: &str, + children: &[&dyn ExecutionPlan], + ) -> Result> { + self.validate_shape(plan_name, children.len())?; + + let Some(co_partitioned) = &self.co_partitioned else { + return Ok(vec![]); + }; + if self.co_partitioning_satisfied(co_partitioned, children) { + return Ok(vec![]); + } + + Ok(co_partitioned.clone()) + } + + /// Validate the requirements against a plan's children. + pub(crate) fn check_invariants( + &self, + plan: &P, + check: InvariantLevel, + ) -> Result<()> { + let children = plan.children(); + self.validate_shape(plan.name(), children.len())?; + + let children = children + .into_iter() + .map(|child| child.as_ref()) + .collect::>(); + if matches!(check, InvariantLevel::Executable) + && let Some(co_partitioned) = &self.co_partitioned + && !self.co_partitioning_satisfied(co_partitioned, &children) + { + return internal_err!( + "{} requires children {:?} to be co-partitioned", + plan.name(), + co_partitioned + ); + } + + Ok(()) + } + + fn validate_shape(&self, plan_name: &str, children_len: usize) -> Result<()> { + if self.children.len() != children_len { + return internal_err!( + "{plan_name}::input_distribution_requirements returned incorrect child count: {} != {}", + self.children.len(), + children_len + ); + } + + if let Some(co_partitioned) = &self.co_partitioned { + if co_partitioned.len() < 2 { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: at least two children are required" + ); + } + let mut seen = vec![false; self.children.len()]; + for &child in co_partitioned { + validate_child_index(plan_name, child, self.children.len(), &mut seen)?; + if matches!( + self.children[child].distribution, + Distribution::UnspecifiedDistribution + ) { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: child {child} has unspecified distribution" + ); + } + } + } + + Ok(()) + } + + fn co_partitioning_satisfied( + &self, + co_partitioned: &[usize], + children: &[&dyn ExecutionPlan], + ) -> bool { + let first_idx = co_partitioned[0]; + let first_requirement = &self.children[first_idx]; + let first = children[first_idx]; + let first_partitioning = first.output_partitioning(); + + if !first_partitioning + .satisfaction( + &first_requirement.distribution, + first.equivalence_properties(), + false, + ) + .is_satisfied() + { + return false; + } + + for &child_idx in co_partitioned.iter().skip(1) { + let requirement = &self.children[child_idx]; + let child = children[child_idx]; + if !child + .output_partitioning() + .satisfaction( + &requirement.distribution, + child.equivalence_properties(), + false, + ) + .is_satisfied() + || !compatible_co_partitioning_layout( + first_partitioning, + child.output_partitioning(), + ) + { + return false; + } + } + + true + } +} + +/// A distribution requirement for a single child. +#[derive(Debug, Clone)] +struct ChildDistributionRequirement { + distribution: Distribution, +} + +fn validate_child_index( + plan_name: &str, + child_idx: usize, + child_count: usize, + seen: &mut [bool], +) -> Result<()> { + if child_idx >= child_count { + return internal_err!( + "{plan_name} has invalid distribution requirement: child index {child_idx} out of bounds" + ); + } + if seen[child_idx] { + return internal_err!( + "{plan_name} has invalid distribution requirement: child {child_idx} appears more than once" + ); + } + seen[child_idx] = true; + Ok(()) +} + +fn compatible_co_partitioning_layout( + first_partitioning: &Partitioning, + other_partitioning: &Partitioning, +) -> bool { + if first_partitioning.partition_count() == 1 + && other_partitioning.partition_count() == 1 + { + return true; + } + + if first_partitioning.partition_count() != other_partitioning.partition_count() { + return false; + } + + match (first_partitioning, other_partitioning) { + (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, + (Partitioning::Range(left), Partitioning::Range(right)) => { + left.split_points() == right.split_points() + && left.ordering().len() == right.ordering().len() + && left + .ordering() + .iter() + .zip(right.ordering()) + .all(|(left, right)| left.options == right.options) + } + _ => false, + } +} diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 50eac566d90ef..5c20f85656e90 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -16,6 +16,7 @@ // under the License. pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +use crate::distribution_requirements::InputDistributionRequirements; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -161,12 +162,30 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { check_default_invariants(self, check) } - /// Specifies the data distribution requirements for all the - /// children for this `ExecutionPlan`, By default it's [[Distribution::UnspecifiedDistribution]] for each child, + /// Specifies simple per-child input distribution requirements. + /// + /// Deprecated: override [`Self::input_distribution_requirements`] instead. + /// + /// By default, each child has [`Distribution::UnspecifiedDistribution`]. + #[deprecated(since = "55.0.0", note = "Use input_distribution_requirements")] fn required_input_distribution(&self) -> Vec { vec![Distribution::UnspecifiedDistribution; self.children().len()] } + /// Specifies the input distribution requirements for this plan. + /// + /// The default implementation wraps [`Self::required_input_distribution`]. + /// Override this method for richer requirements, such as allowing alternate + /// satisfaction policies or requiring multiple children to be co-partitioned. + /// See [`InputDistributionRequirements`] for details. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + #[expect( + deprecated, + reason = "compatibility shim for external ExecutionPlan implementations" + )] + InputDistributionRequirements::new(self.required_input_distribution()) + } + /// Specifies the ordering required for all of the children of this /// `ExecutionPlan`. /// @@ -213,8 +232,8 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { fn benefits_from_input_partitioning(&self) -> Vec { // By default try to maximize parallelism with more CPUs if // possible - self.required_input_distribution() - .into_iter() + self.input_distribution_requirements() + .per_child_distributions() .map(|dist| !matches!(dist, Distribution::SinglePartition)) .collect() } @@ -1197,14 +1216,15 @@ macro_rules! check_len { /// Returns an error if the given node does not conform. pub fn check_default_invariants( plan: &P, - _check: InvariantLevel, + check: InvariantLevel, ) -> Result<(), DataFusionError> { let children_len = plan.children().len(); check_len!(plan, maintains_input_order, children_len); check_len!(plan, required_input_ordering, children_len); - check_len!(plan, required_input_distribution, children_len); check_len!(plan, benefits_from_input_partitioning, children_len); + plan.input_distribution_requirements() + .check_invariants(plan, check)?; Ok(()) } diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 6661d2782b212..fde1718001252 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -306,10 +306,14 @@ impl ExecutionPlan for CrossJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn execute( diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index f7391feb29cc3..0ee918bb66e70 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -54,8 +54,9 @@ use crate::projection::{ use crate::repartition::REPARTITION_RANDOM_STATE; use crate::spill::get_record_batch_memory_size; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, - PlanProperties, SendableRecordBatchStream, Statistics, + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + InputDistributionRequirements, Partitioning, PlanProperties, + SendableRecordBatchStream, Statistics, common::can_project, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, @@ -845,21 +846,65 @@ impl HashJoinExec { return false; } - // `preserve_file_partitions` can report Hash partitioning for Hive-style - // file groups, but those partitions are not actually hash-distributed. - // Partitioned dynamic filters rely on hash routing, so disable them in - // this mode to avoid incorrect results. Follow-up work: enable dynamic - // filtering for preserve_file_partitioned scans (issue #20195). - // https://github.com/apache/datafusion/issues/20195 - if config.optimizer.preserve_file_partitions > 0 - && self.mode == PartitionMode::Partitioned + // Bounds and membership filters derived from the build side do not + // account for null-equal matching: a probe-side NULL key evaluates + // such predicates to NULL and would be pruned, even though it can + // match a build-side NULL when nulls compare equal. + if self.null_equality == NullEquality::NullEqualsNull { + return false; + } + + if self.mode == PartitionMode::Partitioned { + // `preserve_file_partitions` can report Hash partitioning for + // Hive-style file groups, but those partitions are not actually + // hash-distributed. Partitioned dynamic filters rely on hash + // routing, so disable them in this mode to avoid incorrect + // results. Follow-up work: enable dynamic filtering for + // preserve_file_partitioned scans (issue #20195). + // https://github.com/apache/datafusion/issues/20195 + if config.optimizer.preserve_file_partitions > 0 { + return false; + } + + // Partitioned dynamic filters route probe rows with + // `hash(join_key) % partition_count`. That is only valid when + // partition ids are hash buckets with the same bucket count, or + // when there is only one partition and no routing choice exists. + // This also rejects non-hash partitioning such as + // `Partitioning::Range`. + if !self.has_partitioned_dynamic_filter_routing() { + return false; + } + } + + if self.mode == PartitionMode::Partitioned + && !self.has_partitioned_dynamic_filter_routing() { + // TODO: support partition-routed dynamic filters for compatible + // range co-partitioned joins. + // . return false; } true } + fn has_partitioned_dynamic_filter_routing(&self) -> bool { + match ( + self.left.output_partitioning(), + self.right.output_partitioning(), + ) { + ( + Partitioning::Hash(_, left_partition_count), + Partitioning::Hash(_, right_partition_count), + ) => left_partition_count == right_partition_count, + (left_partitioning, right_partitioning) => { + left_partitioning.partition_count() == 1 + && right_partitioning.partition_count() == 1 + } + } + } + /// left (build) side which gets hashed pub fn left(&self) -> &Arc { &self.left @@ -1203,26 +1248,30 @@ impl ExecutionPlan for HashJoinExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { match self.mode { - PartitionMode::CollectLeft => vec![ - Distribution::SinglePartition, - Distribution::UnspecifiedDistribution, - ], PartitionMode::Partitioned => { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) } - PartitionMode::Auto => vec![ + PartitionMode::CollectLeft => InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, Distribution::UnspecifiedDistribution, + ]), + PartitionMode::Auto => InputDistributionRequirements::new(vec![ Distribution::UnspecifiedDistribution, - ], + Distribution::UnspecifiedDistribution, + ]), } } @@ -2119,7 +2168,90 @@ mod tests { Ok((left_schema, right_schema, on)) } + #[derive(Debug)] + struct PartitionedTestInput { + input: Arc, + cache: Arc, + } + + impl PartitionedTestInput { + fn new(input: Arc, partitioning: Partitioning) -> Self { + let cache = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + partitioning, + input.pipeline_behavior(), + input.boundedness(), + )); + Self { input, cache } + } + } + + impl DisplayAs for PartitionedTestInput { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "PartitionedTestInput") + } + DisplayFormatType::TreeRender => write!(f, ""), + } + } + } + + impl ExecutionPlan for PartitionedTestInput { + fn name(&self) -> &'static str { + "PartitionedTestInput" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!( + "PartitionedTestInput expected one child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new( + Arc::clone(&children[0]), + self.cache.output_partitioning().clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(partition, context) + } + } + + fn declared_hash_test_input( + input: Arc, + key: &str, + key_index: usize, + partition_count: usize, + ) -> Result> { + Ok(Arc::new(PartitionedTestInput::new( + input, + Partitioning::Hash( + vec![Arc::new(Column::new(key, key_index))], + partition_count, + ), + ))) + } + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::execution_plan::Boundedness; use crate::joins::hash_join::stream::lookup_join_hashmap; use crate::test::{TestMemoryExec, assert_join_metrics}; use crate::{ @@ -2142,11 +2274,67 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion_physical_expr::{ + EquivalenceProperties, PhysicalSortExpr, RangePartitioning, SplitPoint, + }; use hashbrown::HashTable; use insta::{allow_duplicates, assert_snapshot}; use rstest::*; use rstest_reuse::*; + #[derive(Debug)] + struct PartitionedTestExec { + cache: Arc, + } + + impl PartitionedTestExec { + fn try_new(schema: SchemaRef, partitioning: Partitioning) -> Result { + Ok(Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )), + }) + } + } + + impl DisplayAs for PartitionedTestExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PartitionedTestExec") + } + } + + impl ExecutionPlan for PartitionedTestExec { + fn name(&self) -> &'static str { + "PartitionedTestExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } + } + fn div_ceil(a: usize, b: usize) -> usize { a.div_ceil(b) } @@ -5375,6 +5563,7 @@ mod tests { None, ) .unwrap(); + let left = declared_hash_test_input(left, "b1", 1, 2)?; let right_batch = build_table_i32( ("a2", &vec![10, 11]), ("b2", &vec![12, 13]), @@ -5386,6 +5575,7 @@ mod tests { None, ) .unwrap(); + let right = declared_hash_test_input(right, "b2", 1, 2)?; let on = vec![( Arc::new(Column::new_with_schema("b1", &left_batch.schema())?) as _, Arc::new(Column::new_with_schema("b2", &right_batch.schema())?) as _, @@ -5413,8 +5603,8 @@ mod tests { let task_ctx = Arc::new(task_ctx); let join = HashJoinExec::try_new( - Arc::clone(&left) as Arc, - Arc::clone(&right) as Arc, + Arc::clone(&left), + Arc::clone(&right), on.clone(), None, &join_type, @@ -5717,6 +5907,8 @@ mod tests { Arc::clone(&child_right_schema), None, )?; + let child_left = declared_hash_test_input(child_left, "child_key", 1, 4)?; + let child_right = declared_hash_test_input(child_right, "child_right_key", 1, 4)?; let parent_left: Arc = TestMemoryExec::try_new_exec( &[ vec![build_table_i32( @@ -5735,6 +5927,7 @@ mod tests { Arc::clone(&parent_left_schema), None, )?; + let parent_left = declared_hash_test_input(parent_left, "parent_key", 1, 4)?; let child_on = vec![( Arc::new(Column::new_with_schema("child_key", &child_left_schema)?) as _, @@ -6341,6 +6534,193 @@ mod tests { Ok(()) } + #[test] + fn test_dynamic_filter_pushdown_rejects_null_equal_join() -> Result<()> { + let (_, _, on) = build_schema_and_on()?; + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![1]), ("b1", &vec![1]), ("c2", &vec![1])); + + let session_config = join_dynamic_filter_session_config(0); + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::RightSemi, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + fn range_partitioned_test_input( + schema: SchemaRef, + range_key: &str, + ) -> Result> { + let input = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let range_expr = Arc::new(Column::new_with_schema(range_key, &schema)?); + let range_partitioning = Partitioning::Range(RangePartitioning::new( + [PhysicalSortExpr::new_default(range_expr)].into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )); + RepartitionExec::try_new(input, range_partitioning) + .map(|exec| Arc::new(exec) as _) + } + + fn hash_partitioned_test_input( + schema: SchemaRef, + hash_key: &str, + partition_count: usize, + ) -> Result> { + let input = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let hash_expr = Arc::new(Column::new_with_schema(hash_key, &schema)?); + RepartitionExec::try_new( + input, + Partitioning::Hash(vec![hash_expr], partition_count), + ) + .map(|exec| Arc::new(exec) as _) + } + + fn join_dynamic_filter_session_config( + preserve_file_partitions: usize, + ) -> SessionConfig { + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + session_config + .options_mut() + .optimizer + .preserve_file_partitions = preserve_file_partitions; + session_config + } + + fn partitioned_inner_hash_join( + left: Arc, + right: Arc, + on: JoinOn, + ) -> Result { + HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + } + + #[test] + fn dynamic_filter_rejects_range_partitioning() -> Result<()> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left = range_partitioned_test_input(left_schema, "b1")?; + let right = range_partitioned_test_input(right_schema, "b1")?; + + let session_config = join_dynamic_filter_session_config(0); + let join = partitioned_inner_hash_join(left, right, on)?; + + let requirements = join.input_distribution_requirements().into_per_child(); + assert!(matches!( + requirements.as_slice(), + [ + Distribution::KeyPartitioned(_), + Distribution::KeyPartitioned(_) + ] + )); + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn dynamic_filter_rejects_preserve_file_partitions() -> Result<()> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left = hash_partitioned_test_input(left_schema, "b1", 2)?; + let right = hash_partitioned_test_input(right_schema, "b1", 2)?; + + let session_config = join_dynamic_filter_session_config(1); + let join = partitioned_inner_hash_join(left, right, on)?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn dynamic_filter_rejects_mismatched_hash_counts() -> Result<()> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left = hash_partitioned_test_input(left_schema, "b1", 2)?; + let right = hash_partitioned_test_input(right_schema, "b1", 3)?; + + let session_config = join_dynamic_filter_session_config(0); + let join = partitioned_inner_hash_join(left, right, on)?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_partitioned_dynamic_filter_pushdown_rejects_range_partitioning() -> Result<()> + { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].0), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?); + let right_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].1), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?); + let left = Arc::new(PartitionedTestExec::try_new( + left_schema, + left_partitioning, + )?); + let right = Arc::new(PartitionedTestExec::try_new( + right_schema, + right_partitioning, + )?); + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + #[test] fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> { let (_, _, on) = build_schema_and_on()?; diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 15af23b447836..20d05983cf4ed 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -565,10 +565,14 @@ impl ExecutionPlan for NestedLoopJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 50e9252a21131..a2a1d6c787bb7 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -508,10 +508,14 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 9e87b52696a57..6a7064bf0417b 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -40,7 +40,8 @@ use crate::projection::{ use crate::spill::spill_manager::SpillManager; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties, + InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use arrow::compute::SortOptions; @@ -423,15 +424,19 @@ impl ExecutionPlan for SortMergeJoinExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) } fn required_input_ordering(&self) -> Vec> { diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index ef92964fadf84..b798a166e9cc2 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -53,7 +53,8 @@ use crate::projection::{ use crate::stream::EmptyRecordBatchStream; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, + InputDistributionRequirements, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, joins::StreamJoinPartitionMode, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; @@ -426,6 +427,10 @@ impl ExecutionPlan for SymmetricHashJoinExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { match self.mode { StreamJoinPartitionMode::Partitioned => { let (left_expr, right_expr) = self @@ -433,13 +438,16 @@ impl ExecutionPlan for SymmetricHashJoinExec { .iter() .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) .unzip(); - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) } StreamJoinPartitionMode::SinglePartition => { - vec![Distribution::SinglePartition, Distribution::SinglePartition] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) } } } diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index b4aa295562b67..134d650a9cd6b 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -33,7 +33,8 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, ProjectionExpr}; use crate::{ - ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, Statistics, + ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, + RangePartitioning, Statistics, }; // compatibility pub use super::join_filter::JoinFilter; @@ -67,7 +68,7 @@ use datafusion_common::hash_utils::create_hashes; use datafusion_common::stats::Precision; use datafusion_common::{ DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, - not_impl_err, plan_err, + internal_datafusion_err, not_impl_err, plan_err, }; use datafusion_expr::Operator; use datafusion_expr::interval_arithmetic::Interval; @@ -144,6 +145,21 @@ pub fn adjust_right_output_partitioning( .collect::>()?; Partitioning::Hash(new_exprs, *size) } + Partitioning::Range(range) => { + let ordering = add_offset_to_physical_sort_exprs( + range.ordering().iter().cloned(), + left_columns_len as _, + )?; + let ordering = LexOrdering::new(ordering).ok_or_else(|| { + internal_datafusion_err!( + "Offsetting range partitioning produced an empty ordering" + ) + })?; + Partitioning::Range(RangePartitioning::new( + ordering, + range.split_points().to_vec(), + )) + } result => result.clone(), }; Ok(result) @@ -2144,7 +2160,7 @@ mod tests { use arrow::datatypes::{DataType, Fields}; use arrow::error::{ArrowError, Result as ArrowResult}; use datafusion_common::stats::Precision::{Absent, Exact, Inexact}; - use datafusion_common::{ScalarValue, arrow_datafusion_err, arrow_err}; + use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err}; use datafusion_physical_expr::PhysicalSortExpr; use rstest::rstest; @@ -3239,6 +3255,53 @@ mod tests { Ok(()) } + #[test] + fn test_adjust_right_output_partitioning_preserves_range() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int32(Some(10)), + ScalarValue::Int32(Some(100)), + ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(20)), + ScalarValue::Int32(Some(50)), + ]), + ]; + let range = RangePartitioning::try_new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 2)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points.clone(), + )?; + + let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?; + let expected = Partitioning::Range(RangePartitioning::new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 3)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 5)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points, + )); + + assert_eq!(adjusted, expected); + Ok(()) + } + #[test] fn test_calculate_join_output_ordering() -> Result<()> { let left_ordering = LexOrdering::new(vec![ diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 3005e975424b4..8190663b2b23d 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -37,10 +37,13 @@ pub use datafusion_expr::{Accumulator, ColumnarValue}; use datafusion_physical_expr::PhysicalSortExpr; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ - Distribution, Partitioning, PhysicalExpr, expressions, + Distribution, Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, expressions, }; pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +pub use crate::distribution_requirements::{ + ChildSatisfactionOptions, InputDistributionRequirements, +}; pub use crate::execution_plan::{ ExecutionPlan, ExecutionPlanProperties, PlanProperties, collect, collect_partitioned, displayable, execute_input_stream, execute_stream, execute_stream_partitioned, @@ -71,6 +74,7 @@ pub mod column_rewriter; pub mod common; pub mod coop; pub mod display; +pub mod distribution_requirements; pub mod empty; pub mod execution_plan; pub mod explain; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 7f42c33a79ca0..94ef00d8567aa 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -162,7 +162,11 @@ impl ExecutionPlan for GlobalLimitExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index f34aac3744557..99767df73a734 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -163,10 +163,14 @@ impl ExecutionPlan for RecursiveQueryExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ crate::Distribution::SinglePartition, crate::Distribution::SinglePartition, - ] + ]) } fn with_new_children( diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 465ca4a99e961..a20714d1e8647 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -19,10 +19,11 @@ //! partitions to M output partitions based on a partitioning scheme, optionally //! maintaining the order of the input rows in the output. +use std::cmp::Ordering; use std::fmt::{Debug, Formatter}; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::task::{Context, Poll}; use std::vec; @@ -45,20 +46,22 @@ use crate::{ check_if_same_properties, }; -use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions}; +use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions}; use arrow::compute::take_arrays; use arrow::datatypes::{SchemaRef, UInt32Type}; +use arrow_schema::SortOptions; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; -use datafusion_common::utils::transpose; +use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, assert_or_internal_err, internal_err, + ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, + assert_or_internal_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryConsumer; -use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr, RangePartitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use crate::filter_pushdown::{ @@ -255,7 +258,7 @@ impl SharedCoalescer { /// sender, finalize the coalescer and return its residual batches; if /// other senders are still active, return `Ok(None)`. fn finalize(&self) -> Result> { - let was_last = self.active_senders.fetch_sub(1, Ordering::AcqRel) == 1; + let was_last = self.active_senders.fetch_sub(1, AtomicOrdering::AcqRel) == 1; if !was_last { return Ok(vec![]); } @@ -569,6 +572,18 @@ enum BatchPartitionerState { num_partitions: usize, next_idx: usize, }, + Range { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Sort options from the `LexOrdering` + sort_options: Vec, + /// Boundaries between adjacent partitions. + split_points: Vec, + /// Row indices grouped by output partition + indices: Vec>, + /// Buffer of `ScalarValue` used to represent the values for a row - based on the `LexOrdering` ordering - to compare against split points + partition_buffer: Vec, + }, } /// Fixed RandomState used for hash repartitioning to ensure consistent behavior across @@ -705,13 +720,40 @@ impl BatchPartitioner { timer, } } + + /// Create a new [`BatchPartitioner`] for range-based repartitioning. + /// + /// # Parameters + /// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions + /// - `timer`: Metric used to record time spent during repartitioning. + pub fn new_range_partitioner( + range_partitioning: &RangePartitioning, + timer: metrics::Time, + ) -> Self { + let ordering = range_partitioning.ordering().clone(); + let split_points = range_partitioning.split_points().to_vec(); + let num_partitions = range_partitioning.partition_count(); + let sort_options: Vec = ordering.iter().map(|e| e.options).collect(); + + Self { + state: BatchPartitionerState::Range { + partition_buffer: Vec::with_capacity(ordering.len()), + ordering, + sort_options, + split_points, + indices: vec![vec![]; num_partitions], + }, + timer, + } + } + /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme. /// /// This is a convenience constructor that delegates to the specialized - /// hash or round-robin constructors depending on the partitioning variant. + /// hash, round-robin, or range constructors depending on the partitioning variant. /// /// # Parameters - /// - `partitioning`: Partitioning scheme to apply (hash or round-robin). + /// - `partitioning`: Partitioning scheme to apply (hash, round-robin, or range). /// - `timer`: Metric used to record time spent during repartitioning. /// - `input_partition`: Index of the current input partition. /// - `num_input_partitions`: Total number of input partitions. @@ -737,6 +779,9 @@ impl BatchPartitioner { num_input_partitions, )) } + Partitioning::Range(range_repartitioning) => { + Ok(Self::new_range_partitioner(&range_repartitioning, timer)) + } other => { not_impl_err!("Unsupported repartitioning scheme {other:?}") } @@ -823,22 +868,95 @@ impl BatchPartitioner { Box::new(partitioned_batches.into_iter()) } + BatchPartitionerState::Range { + ordering, + sort_options, + split_points, + indices, + partition_buffer, + } => { + // Tracking time required for distributing indexes across output partitions + let timer = self.timer.timer(); + if split_points.is_empty() { + timer.done(); + Box::new(std::iter::once(Ok((0, batch)))) + } else { + let arrays = evaluate_expressions_to_arrays( + ordering.iter().map(|e| &e.expr), + &batch, + )?; + + indices.iter_mut().for_each(|v| v.clear()); + + Self::partition_range_indices( + &arrays, + split_points, + sort_options, + partition_buffer, + indices, + )?; + + // Finished building index-arrays for output partitions + timer.done(); + + let partitioned_batches = + Self::partition_grouped_take(&batch, indices, &self.timer)?; + + Box::new(partitioned_batches.into_iter()) + } + } }; Ok(it) } + /// Groups input row indices by range partition. This populates `indices[p]` with the + /// row indices from `arrays` that belong in output partition `p` according to `split_points` and `sort_options`. + fn partition_range_indices( + arrays: &[Arc], + split_points: &[SplitPoint], + sort_options: &[SortOptions], + row_key_buffer: &mut Vec, + indices: &mut [Vec], + ) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + // Note that `extract_row_at_idx_to_buf` clears the `row_key_buffer` on each invocation, creating a new row key for comparison for each row + extract_row_at_idx_to_buf(arrays, row_idx, row_key_buffer)?; + + let mut low = 0; + let mut high = split_points.len(); + while low < high { + let mid = low + (high - low) / 2; + let comparison = compare_rows( + row_key_buffer, + split_points[mid].values(), + sort_options, + )?; + match comparison { + Ordering::Less => high = mid, + Ordering::Equal | Ordering::Greater => low = mid + 1, + } + } + + indices[low].push(row_idx as u32) + } + + Ok(()) + } + // return the number of output partitions fn num_partitions(&self) -> usize { match &self.state { BatchPartitionerState::RoundRobin { num_partitions, .. } => *num_partitions, - BatchPartitionerState::Hash { indices, .. } => indices.len(), + BatchPartitionerState::Hash { indices, .. } + | BatchPartitionerState::Range { indices, .. } => indices.len(), } } - /// Build repartitioned hash output batches using one `take` per input batch. + /// Build repartitioned hash/range output batches using one `take` per input batch. /// - /// The hash router first fills one index vector per output partition. This method + /// The routers first fills one index vector per output partition. This method /// concatenates those index vectors, performs one grouped `take_arrays`, and /// then returns each output partition as a slice of the reordered batch. /// @@ -1430,6 +1548,30 @@ impl ExecutionPlan for RepartitionExec { } Partitioning::Hash(new_partitions, *size) } + Partitioning::Range(range_partitioning) => { + // Rewrite range key expressions through the projection. + let mut sort_exprs = + Vec::with_capacity(range_partitioning.ordering().len()); + for sort_expr in range_partitioning.ordering() { + let Some(new_expr) = + update_expr(&sort_expr.expr, projection.expr(), false)? + else { + return Ok(None); + }; + sort_exprs.push(PhysicalSortExpr::new(new_expr, sort_expr.options)); + } + + let Some(ordering) = LexOrdering::new(sort_exprs) else { + return internal_err!( + "failed to create LexOrdering for range partitioning" + ); + }; + + Partitioning::Range(RangePartitioning::try_new( + ordering, + range_partitioning.split_points().to_vec(), + )?) + } others => others.clone(), }; @@ -1488,6 +1630,10 @@ impl ExecutionPlan for RepartitionExec { new_properties.partitioning = match new_properties.partitioning { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), + Range(_) => { + // Number of partitions is constrained by the split points and cannot be changed + return Ok(None); + } UnknownPartitioning(_) => UnknownPartitioning(target_partitions), }; Ok(Some(Arc::new(Self { @@ -1601,26 +1747,12 @@ impl RepartitionExec { input_partition: usize, num_input_partitions: usize, ) -> Result<()> { - let mut partitioner = match &partitioning { - Partitioning::Hash(exprs, num_partitions) => { - BatchPartitioner::new_hash_partitioner( - exprs.clone(), - *num_partitions, - metrics.repartition_time.clone(), - )? - } - Partitioning::RoundRobinBatch(num_partitions) => { - BatchPartitioner::new_round_robin_partitioner( - *num_partitions, - metrics.repartition_time.clone(), - input_partition, - num_input_partitions, - ) - } - other => { - return not_impl_err!("Unsupported repartitioning scheme {other:?}"); - } - }; + let mut partitioner = BatchPartitioner::try_new( + partitioning, + metrics.repartition_time.clone(), + input_partition, + num_input_partitions, + )?; // While there are still outputs to send to, keep pulling inputs let mut batches_until_yield = partitioner.num_partitions(); @@ -1954,6 +2086,8 @@ mod tests { use std::collections::HashSet; use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExpr; use crate::test::TestMemoryExec; use crate::{ test::{ @@ -1968,12 +2102,14 @@ mod tests { use arrow::array::{ArrayRef, StringArray, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::cast::as_string_array; + use datafusion_common::ScalarValue; + use datafusion_common::cast::{as_string_array, as_uint32_array}; use datafusion_common::exec_err; use datafusion_common::test_util::batches_to_sort_string; use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint}; use insta::assert_snapshot; #[test] @@ -2071,7 +2207,7 @@ mod tests { #[tokio::test] async fn one_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition]; @@ -2094,7 +2230,7 @@ mod tests { #[tokio::test] async fn many_to_one_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2111,7 +2247,7 @@ mod tests { #[tokio::test] async fn many_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2132,7 +2268,7 @@ mod tests { #[tokio::test] async fn many_to_many_hash_partition() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2154,9 +2290,542 @@ mod tests { Ok(()) } + #[tokio::test] + async fn many_to_many_range_partition() -> Result<()> { + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; + + // create_batch values are [1, 2, 3, 4, 5, 6, 7, 8]; split at 3 and 6 yields + // 2, 3, and 3 rows per batch respectively + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![3, 6])?; + + let output_partitions = repartition(&schema, partitions, partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!(300, partition_row_count(&output_partitions[0])); + assert_eq!(450, partition_row_count(&output_partitions[1])); + assert_eq!(450, partition_row_count(&output_partitions[2])); + assert_eq!( + collect_partition_u32_values(&output_partitions[0]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([1, 2]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[1]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([3, 4, 5]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[2]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([6, 7, 8]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_compound_keys() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![5, 10, 10, 10, 10, 15])), + Arc::new(UInt32Array::from(vec![1, 1, 3, 5, 7, 0])), + ], + )?; + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()), + ] + .into(), + vec![ + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(1)), + ]), + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(5)), + ]), + ], + )?); + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![(5, 1)], + collect_partition_u32_pairs(&output_partitions[0]) + ); + assert_eq!( + vec![(10, 1), (10, 3)], + collect_partition_u32_pairs(&output_partitions[1]) + ); + assert_eq!( + vec![(10, 5), (10, 7), (15, 0)], + collect_partition_u32_pairs(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_last() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, false), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![None, Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_first() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, true), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![None, Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_asc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![10, 20])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_desc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 20, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(true, false), vec![20, 10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(15), Some(20)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(5), Some(10)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_string_rows() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["bar", "baz", "foo", "qux"])) as ArrayRef, + )])?; + + let schema = batch.schema(); + let expr = col("my_awesome_field", &schema)?; + let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema)); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new_default(expr)].into(), + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "foo".to_string(), + ))])], + )?); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?; + + let mut partition_0 = Vec::new(); + let mut stream = exec.execute(0, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + partition_0.push(result?); + } + + let mut partition_1 = Vec::new(); + let mut stream = exec.execute(1, task_ctx)?; + while let Some(result) = stream.next().await { + partition_1.push(result?); + } + + assert_eq!( + vec!["bar", "baz"], + collect_partition_string_values(&partition_0) + ); + assert_eq!( + vec!["foo", "qux"], + collect_partition_string_values(&partition_1) + ); + + Ok(()) + } + + #[test] + fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> { + // Three columns so the projection both narrows the schema (required for + // swap) and moves the range key from @0 to @1. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; + + let swapped = repartition + .try_swapping_with_projection(&projection)? + .expect("swap should succeed when projection keeps the range key"); + let swapped_repartition = swapped + .downcast_ref::() + .expect("top node should be RepartitionExec"); + + assert!(swapped_repartition.input().is::()); + let range = expect_range_partitioning(swapped_repartition.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@1 ASC"); + assert_eq!( + range.split_points(), + &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])] + ); + + Ok(()) + } + + #[test] + fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> { + // Drop a simple range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + // Drop part of a compound range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_when_maintains_order() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source with preserve_order: Range maintains input order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new( + RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )? + .with_preserve_order(), + ); + assert!(repartition.maintains_input_order()[0]); + + match repartition.try_pushdown_sort(ordering.as_ref())? { + SortOrderPushdownResult::Exact { inner } => { + let pushed = inner + .downcast_ref::() + .expect("pushdown should keep RepartitionExec"); + + assert!(pushed.preserve_order()); + assert!(pushed.maintains_input_order()[0]); + + let range = expect_range_partitioning(pushed.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); + assert_eq!( + inner.properties().output_ordering().map(|o| o.to_string()), + Some(ordering.to_string()), + "pushed repartition output ordering should match the requested sort" + ); + } + other => panic!("expected Exact sort pushdown, got {other:?}"), + } + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_unsupported_without_order_maintenance() + -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source without preserve_order: Range does not maintain order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new(RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + assert!(!repartition.maintains_input_order()[0]); + + assert!(matches!( + repartition.try_pushdown_sort(ordering.as_ref())?, + SortOrderPushdownResult::Unsupported + )); + + Ok(()) + } + + fn range_partitioning_on_columns( + schema: &SchemaRef, + key_columns: &[&str], + split_points: Vec>, + ) -> Result { + let Some(ordering) = LexOrdering::new( + key_columns + .iter() + .map(|name| { + Ok(PhysicalSortExpr::new( + col(name, schema)?, + SortOptions::default(), + )) + }) + .collect::>>()?, + ) else { + return exec_err!("range ordering must not be empty"); + }; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points + .into_iter() + .map(|values| { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::UInt32(Some(value))) + .collect(), + ) + }) + .collect(), + )?)) + } + + fn projection_on_columns( + input: &Arc, + names: &[&str], + ) -> Result { + let exprs = names + .iter() + .map(|name| { + Ok(ProjectionExpr { + expr: col(name, &input.schema())?, + alias: (*name).to_string(), + }) + }) + .collect::>>()?; + ProjectionExec::try_new(exprs, Arc::clone(input)) + } + + fn expect_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning { + match partitioning { + Partitioning::Range(range) => range, + other => panic!("expected Range partitioning, got {other:?}"), + } + } + + /// Test source that claims Exact support for any sort pushdown request. + #[derive(Debug, Clone)] + struct ExactSortPushdownExec { + cache: Arc, + } + + impl ExactSortPushdownExec { + fn new(schema: SchemaRef, num_partitions: usize, ordering: LexOrdering) -> Self { + use crate::execution_plan::{Boundedness, EmissionType}; + Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings(schema, [ordering]), + Partitioning::UnknownPartitioning(num_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } + } + + impl DisplayAs for ExactSortPushdownExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ExactSortPushdownExec") + } + } + + impl ExecutionPlan for ExactSortPushdownExec { + fn name(&self) -> &str { + "ExactSortPushdownExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(EmptyRecordBatchStream::new(self.schema()))) + } + + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result>> { + Ok(SortOrderPushdownResult::Exact { + inner: Arc::new(self.clone()), + }) + } + } + #[tokio::test] async fn test_repartition_with_coalescing() -> Result<()> { - let schema = test_schema(); + let schema = test_schema(false); // create 50 batches, each having 8 rows let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone()]; @@ -2180,8 +2849,76 @@ mod tests { Ok(()) } - fn test_schema() -> Arc { - Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) + fn test_schema(nullable: bool) -> Arc { + Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::UInt32, + nullable, + )])) + } + + fn u32_range_partitioning( + schema: &SchemaRef, + sort_options: SortOptions, + split_values: Vec, + ) -> Result { + let expr = col("c0", schema)?; + Ok(Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(expr, sort_options)].into(), + split_values + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))])) + .collect(), + )?)) + } + + fn partition_row_count(batches: &[RecordBatch]) -> usize { + batches.iter().map(|batch| batch.num_rows()).sum() + } + + fn collect_partition_u32_values(batches: &[RecordBatch]) -> Vec> { + batches + .iter() + .flat_map(|batch| { + let array = + as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + (0..array.len()) + .map(|idx| { + if array.is_null(idx) { + None + } else { + Some(array.value(idx)) + } + }) + .collect::>() + }) + .collect() + } + + fn collect_partition_u32_pairs(batches: &[RecordBatch]) -> Vec<(u32, u32)> { + batches + .iter() + .flat_map(|batch| { + let a = as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + let b = as_uint32_array(batch.column(1)).expect("expected UInt32 column"); + (0..a.len()) + .map(|idx| (a.value(idx), b.value(idx))) + .collect::>() + }) + .collect() + } + + fn collect_partition_string_values(batches: &[RecordBatch]) -> Vec<&str> { + batches + .iter() + .flat_map(|batch| { + let array = + as_string_array(batch.column(0)).expect("expected Utf8 column"); + (0..array.len()) + .map(|idx| array.value(idx)) + .collect::>() + }) + .collect() } async fn repartition( @@ -2214,7 +2951,7 @@ mod tests { let handle: SpawnedTask>>> = SpawnedTask::spawn(async move { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2574,7 +3311,7 @@ mod tests { #[tokio::test] async fn repartition_with_spilling() -> Result<()> { // Test that repartition successfully spills to disk when memory is constrained - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2636,7 +3373,7 @@ mod tests { #[tokio::test] async fn repartition_with_partial_spilling() -> Result<()> { // Test that repartition can handle partial spilling (some batches in memory, some spilled) - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2706,7 +3443,7 @@ mod tests { #[tokio::test] async fn repartition_without_spilling() -> Result<()> { // Test that repartition does not spill when there's ample memory - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2768,7 +3505,7 @@ mod tests { use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; // Test that repartition fails with OOM when disk manager is disabled - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2811,7 +3548,7 @@ mod tests { /// Create batch fn create_batch() -> RecordBatch { - let schema = test_schema(); + let schema = test_schema(false); RecordBatch::try_new( schema, vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))], @@ -2821,7 +3558,7 @@ mod tests { /// Create batches with sequential values for ordering tests fn create_ordered_batches(num_batches: usize) -> Vec { - let schema = test_schema(); + let schema = test_schema(false); (0..num_batches) .map(|i| { let start = (i * 8) as u32; @@ -2842,7 +3579,7 @@ mod tests { // This tests the state machine fix where we must block on spill_stream // when a Spilled marker is received, rather than continuing to poll the channel - let schema = test_schema(); + let schema = test_schema(false); // Create batches with sequential values: batch 0 has [0,1,2,3,4,5,6,7], // batch 1 has [8,9,10,11,12,13,14,15], etc. let partition = create_ordered_batches(20); @@ -3214,6 +3951,33 @@ mod test { Ok(()) } + #[test] + fn test_range_repartitioned_returns_none() -> Result<()> { + let schema = test_schema(); + let source = memory_exec(&schema); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new( + col("c0", &schema)?, + SortOptions::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ], + )?); + let exec = RepartitionExec::try_new(source, partitioning)?; + + // Range partition count is fixed by split points, so repartitioned() + // cannot change it to an arbitrary target. + let result = exec.repartitioned(10, &Default::default())?; + assert!( + result.is_none(), + "range repartitioning should not support changing partition count" + ); + Ok(()) + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) } diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 3bf16af36c62b..3eeefa3acd7c7 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -268,11 +268,15 @@ impl ExecutionPlan for PartialSortExec { } fn required_input_distribution(&self) -> Vec { - if self.preserve_partitioning { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { vec![Distribution::UnspecifiedDistribution] } else { vec![Distribution::SinglePartition] - } + }) } fn benefits_from_input_partitioning(&self) -> Vec { diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index fe876eeddf7f2..17eb70ef12131 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -302,12 +302,18 @@ impl ExecutionPlan for PartitionedTopKExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { let partition_exprs: Vec> = self.expr [..self.partition_prefix_len] .iter() .map(|e| Arc::clone(&e.expr)) .collect(); - vec![Distribution::HashPartitioned(partition_exprs)] + crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + partition_exprs, + )]) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index f715de0b5964b..551f7e0ffa6dd 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1137,13 +1137,18 @@ impl ExecutionPlan for SortExec { } fn required_input_distribution(&self) -> Vec { - if self.preserve_partitioning { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { vec![Distribution::UnspecifiedDistribution] } else { // global sort - // TODO support RangePartition and OrderedDistribution + // TODO support range partitioning and OrderedDistribution. + // See https://github.com/apache/datafusion/issues/22395 vec![Distribution::SinglePartition] - } + }) } fn children(&self) -> Vec<&Arc> { diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 09570f14ba734..053c57a5a6a62 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -266,7 +266,13 @@ impl ExecutionPlan for SortPreservingMergeExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn benefits_from_input_partitioning(&self) -> Vec { @@ -1486,11 +1492,7 @@ mod tests { let task_ctx = Arc::new(TaskContext::default()); let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]); let properties = CongestedExec::compute_properties(Arc::new(schema.clone())); - let &partition_count = match properties.output_partitioning() { - Partitioning::RoundRobinBatch(partitions) => partitions, - Partitioning::Hash(_, partitions) => partitions, - Partitioning::UnknownPartitioning(partitions) => partitions, - }; + let partition_count = properties.output_partitioning().partition_count(); let source = CongestedExec { schema: schema.clone(), cache: Arc::new(properties), diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index cdf4b08f718c6..61a9b9cc6d0de 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -35,6 +35,7 @@ use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::{Result, internal_err, plan_err}; use datafusion_execution::TaskContext; +use datafusion_physical_expr::projection::ProjectionMapping; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; use async_trait::async_trait; @@ -100,7 +101,7 @@ impl StreamingTableExec { let cache = Self::compute_properties( Arc::clone(&projected_schema), projected_output_ordering.clone(), - &partitions, + Partitioning::UnknownPartitioning(partitions.len()), infinite, ); Ok(Self { @@ -115,6 +116,25 @@ impl StreamingTableExec { }) } + /// Declares the output partitioning of this stream. + /// + /// `output_partitioning` must describe this plan's current output and have + /// the same number of partitions as the stream. + pub fn with_output_partitioning( + mut self, + output_partitioning: Partitioning, + ) -> Result { + if output_partitioning.partition_count() != self.partitions.len() { + return plan_err!( + "Output partitioning has {} partitions but stream has {} partitions", + output_partitioning.partition_count(), + self.partitions.len() + ); + } + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; + Ok(self) + } + pub fn partitions(&self) -> &Vec> { &self.partitions } @@ -147,14 +167,12 @@ impl StreamingTableExec { fn compute_properties( schema: SchemaRef, orderings: Vec, - partitions: &[Arc], + output_partitioning: Partitioning, infinite: bool, ) -> PlanProperties { // Calculate equivalence properties: let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings); - // Get output partitioning: - let output_partitioning = Partitioning::UnknownPartitioning(partitions.len()); let boundedness = if infinite { Boundedness::Unbounded { requires_infinite_memory: false, @@ -204,6 +222,16 @@ impl DisplayAs for StreamingTableExec { if let Some(fetch) = self.limit { write!(f, ", fetch={fetch}")?; } + if !matches!( + self.cache.output_partitioning(), + Partitioning::UnknownPartitioning(_) + ) { + write!( + f, + ", output_partitioning={}", + self.cache.output_partitioning() + )?; + } display_orderings(f, &self.projected_output_ordering)?; @@ -306,6 +334,17 @@ impl ExecutionPlan for StreamingTableExec { }; lex_orderings.push(ordering); } + let projection_mapping = ProjectionMapping::try_new( + projection + .expr() + .iter() + .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())), + &self.schema(), + )?; + let output_partitioning = self + .cache + .output_partitioning() + .project(&projection_mapping, self.cache.equivalence_properties()); StreamingTableExec::try_new( Arc::clone(self.partition_schema()), @@ -315,6 +354,7 @@ impl ExecutionPlan for StreamingTableExec { self.is_infinite(), self.limit(), ) + .and_then(|exec| exec.with_output_partitioning(output_partitioning)) .map(|e| Some(Arc::new(e) as _)) } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 3ea2eb5402fe5..08c47ab3aee93 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -462,7 +462,10 @@ impl ExecutionPlan for UnionExec { /// Combines multiple input streams by interleaving them. /// -/// This only works if all inputs have the same hash-partitioning. +/// All inputs must share an identical [`Partitioning::Hash`] or [`Partitioning::Range`] so that +/// partition `k` covers the same data across every input. Each output partition is the +/// interleaving of the same-indexed partition from all inputs: +/// `output[k] = input[0][k] + input[1][k] + ... + input[n-1][k]` /// /// # Data Flow /// ```text @@ -507,7 +510,7 @@ impl InterleaveExec { pub fn try_new(inputs: Vec>) -> Result { assert_or_internal_err!( can_interleave(inputs.iter()), - "Not all InterleaveExec children have a consistent hash partitioning" + "Not all InterleaveExec children have a consistent hash or range partitioning" ); let cache = Self::compute_properties(&inputs)?; Ok(InterleaveExec { @@ -655,8 +658,12 @@ impl ExecutionPlan for InterleaveExec { } } -/// If all the input partitions have the same Hash partition spec with the first_input_partition -/// The InterleaveExec is partition aware. +/// Returns true if all inputs have the same [`Partitioning::Hash`] or [`Partitioning::Range`] +/// spec, making them safe to interleave. Two inputs are interleave-compatible when partition +/// `k` covers the identical key range or hash bucket across every input. +/// +/// Note: compatibility is checked sequentially against the first input, so +/// `InputDistributionRequirements::co_partitioned` is not needed here. /// /// It might be too strict here in the case that the input partition specs are compatible but not exactly the same. /// For example one input partition has the partition spec Hash('a','b','c') and @@ -669,7 +676,7 @@ pub fn can_interleave>>( }; let reference = first.borrow().output_partitioning(); - matches!(reference, Partitioning::Hash(_, _)) + matches!(reference, Partitioning::Hash(_, _) | Partitioning::Range(_)) && inputs .map(|plan| plan.borrow().output_partitioning().clone()) .all(|partition| partition == *reference) @@ -833,10 +840,13 @@ mod tests { use arrow::compute::SortOptions; use arrow::datatypes::DataType; + use datafusion_common::SplitPoint; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_physical_expr::RangePartitioning; use datafusion_physical_expr::equivalence::convert_to_orderings; use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; // Generate a schema which consists of 7 columns (a, b, c, d, e, f, g) fn create_test_schema() -> Result { @@ -1273,6 +1283,124 @@ mod tests { ); } + fn make_hash_exec( + schema: &SchemaRef, + hash_cols: Vec<&str>, + buckets: usize, + ) -> Result> { + let exprs = hash_cols + .iter() + .map(|c| col(c, schema)) + .collect::>>()?; + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Hash(exprs, buckets), + )?)) + } + + fn make_range_exec( + schema: &SchemaRef, + split_values: Vec, + sort_options: SortOptions, + ) -> Result> { + let sort_expr = + PhysicalSortExpr::new(col(schema.field(0).name(), schema)?, sort_options); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let split_points = split_values + .into_iter() + .map(|v| SplitPoint::new(vec![ScalarValue::Int32(Some(v))])) + .collect(); + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?), + )?)) + } + + #[test] + fn test_can_interleave_matrix() -> Result<()> { + let name_column = "name"; + let age_column = "age"; + let schema = Arc::new(Schema::new(vec![ + Field::new(name_column, DataType::Int32, true), + Field::new(age_column, DataType::Int32, true), + ])); + + let ascending = SortOptions { + descending: false, + nulls_first: false, + }; + struct Case { + inputs: Vec>, + expected: bool, + label: &'static str, + } + + let cases = vec![ + // compatible + Case { + label: "matching hash on single column", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column], 3)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + Case { + label: "matching hash on multiple columns", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + ], + }, + Case { + label: "matching range same splits and order", + expected: true, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 20], ascending)?, + ], + }, + // incompatible + Case { + label: "subset range partition", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 15], ascending)?, + ], + }, + Case { + label: "range different split points", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 30], ascending)?, + ], + }, + Case { + label: "mixed range and hash", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + ]; + + for case in cases { + assert_eq!( + can_interleave(case.inputs.iter()), + case.expected, + "{}", + case.label + ); + } + Ok(()) + } + #[test] fn test_union_cardinality_effect() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index c31d0dd23fa68..4f345a6c2dc54 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -253,7 +253,13 @@ impl ExecutionPlan for UnnestExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn execute( diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 6c6b26c9cf49f..1acda21196008 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -35,8 +35,9 @@ use crate::windows::{ }; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputOrderMode, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, + ExecutionPlanProperties, InputDistributionRequirements, InputOrderMode, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, + check_if_same_properties, }; use arrow::compute::take_record_batch; @@ -331,11 +332,17 @@ impl ExecutionPlan for BoundedWindowAggExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { if self.partition_keys().is_empty() { debug!("No partition defined for BoundedWindowAggExec!!!"); - vec![Distribution::SinglePartition] + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::HashPartitioned(self.partition_keys().clone())] + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) } } diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index ee3b071fc9167..e909dd2017827 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -31,8 +31,9 @@ use crate::windows::{ }; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PhysicalExpr, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, + ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, + check_if_same_properties, }; use arrow::array::ArrayRef; @@ -240,10 +241,16 @@ impl ExecutionPlan for WindowAggExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { if self.partition_keys().is_empty() { - vec![Distribution::SinglePartition] + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::HashPartitioned(self.partition_keys())] + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) } } diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 15a744f5c3500..bfa40e61a8983 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -148,9 +148,19 @@ message RepartitionNode { oneof partition_method { uint64 round_robin = 2; HashRepartition hash = 3; + RangeRepartition range = 4; } } +message RangeSplitPoint { + repeated datafusion_common.ScalarValue value = 1; +} + +message RangeRepartition { + repeated SortExprNode sort_expr = 1; + repeated RangeSplitPoint split_point = 2; +} + message HashRepartition { repeated LogicalExprNode hash_expr = 1; uint64 partition_count = 2; @@ -1160,6 +1170,8 @@ message FileScanExecConf { optional uint64 batch_size = 12; optional ProjectionExprs projection_exprs = 13; + optional bool partitioned_by_file_group = 14; + optional Partitioning output_partitioning = 15; } message ParquetScanExecNode { @@ -1420,13 +1432,22 @@ message PhysicalHashRepartition { uint64 partition_count = 2; } +message PhysicalRangePartitioning { + repeated PhysicalSortExprNode sort_expr = 1; + repeated PhysicalRangeSplitPoint split_point = 2; +} + +message PhysicalRangeSplitPoint { + repeated datafusion_common.ScalarValue value = 1; +} + message RepartitionExecNode{ PhysicalPlanNode input = 1; - // oneof partition_method { + // Legacy direct partitioning fields: // uint64 round_robin = 2; // PhysicalHashRepartition hash = 3; // uint64 unknown = 4; - // } + // New partitioning variants are stored in `partitioning`. Partitioning partitioning = 5; bool preserve_order = 6; } @@ -1436,6 +1457,7 @@ message Partitioning { uint64 round_robin = 1; PhysicalHashRepartition hash = 2; uint64 unknown = 3; + PhysicalRangePartitioning range = 4; } } diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index bd859f2c080b7..9f5785f4cc64a 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -6848,6 +6848,12 @@ impl serde::Serialize for FileScanExecConf { if self.projection_exprs.is_some() { len += 1; } + if self.partitioned_by_file_group.is_some() { + len += 1; + } + if self.output_partitioning.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.FileScanExecConf", len)?; if !self.file_groups.is_empty() { struct_ser.serialize_field("fileGroups", &self.file_groups)?; @@ -6884,6 +6890,12 @@ impl serde::Serialize for FileScanExecConf { if let Some(v) = self.projection_exprs.as_ref() { struct_ser.serialize_field("projectionExprs", v)?; } + if let Some(v) = self.partitioned_by_file_group.as_ref() { + struct_ser.serialize_field("partitionedByFileGroup", v)?; + } + if let Some(v) = self.output_partitioning.as_ref() { + struct_ser.serialize_field("outputPartitioning", v)?; + } struct_ser.end() } } @@ -6911,6 +6923,10 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "batchSize", "projection_exprs", "projectionExprs", + "partitioned_by_file_group", + "partitionedByFileGroup", + "output_partitioning", + "outputPartitioning", ]; #[allow(clippy::enum_variant_names)] @@ -6926,6 +6942,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { Constraints, BatchSize, ProjectionExprs, + PartitionedByFileGroup, + OutputPartitioning, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -6958,6 +6976,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "constraints" => Ok(GeneratedField::Constraints), "batchSize" | "batch_size" => Ok(GeneratedField::BatchSize), "projectionExprs" | "projection_exprs" => Ok(GeneratedField::ProjectionExprs), + "partitionedByFileGroup" | "partitioned_by_file_group" => Ok(GeneratedField::PartitionedByFileGroup), + "outputPartitioning" | "output_partitioning" => Ok(GeneratedField::OutputPartitioning), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -6988,6 +7008,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { let mut constraints__ = None; let mut batch_size__ = None; let mut projection_exprs__ = None; + let mut partitioned_by_file_group__ = None; + let mut output_partitioning__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::FileGroups => { @@ -7061,6 +7083,18 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { } projection_exprs__ = map_.next_value()?; } + GeneratedField::PartitionedByFileGroup => { + if partitioned_by_file_group__.is_some() { + return Err(serde::de::Error::duplicate_field("partitionedByFileGroup")); + } + partitioned_by_file_group__ = map_.next_value()?; + } + GeneratedField::OutputPartitioning => { + if output_partitioning__.is_some() { + return Err(serde::de::Error::duplicate_field("outputPartitioning")); + } + output_partitioning__ = map_.next_value()?; + } } } Ok(FileScanExecConf { @@ -7075,6 +7109,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { constraints: constraints__, batch_size: batch_size__, projection_exprs: projection_exprs__, + partitioned_by_file_group: partitioned_by_file_group__, + output_partitioning: output_partitioning__, }) } } @@ -16138,6 +16174,9 @@ impl serde::Serialize for Partitioning { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("unknown", ToString::to_string(&v).as_str())?; } + partitioning::PartitionMethod::Range(v) => { + struct_ser.serialize_field("range", v)?; + } } } struct_ser.end() @@ -16154,6 +16193,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin", "hash", "unknown", + "range", ]; #[allow(clippy::enum_variant_names)] @@ -16161,6 +16201,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { RoundRobin, Hash, Unknown, + Range, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -16185,6 +16226,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin" | "round_robin" => Ok(GeneratedField::RoundRobin), "hash" => Ok(GeneratedField::Hash), "unknown" => Ok(GeneratedField::Unknown), + "range" => Ok(GeneratedField::Range), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -16226,6 +16268,13 @@ impl<'de> serde::Deserialize<'de> for Partitioning { } partition_method__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| partitioning::PartitionMethod::Unknown(x.0)); } + GeneratedField::Range => { + if partition_method__.is_some() { + return Err(serde::de::Error::duplicate_field("range")); + } + partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(partitioning::PartitionMethod::Range) +; + } } } Ok(Partitioning { @@ -19810,6 +19859,207 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { deserializer.deserialize_struct("datafusion.PhysicalPlanNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalRangePartitioning { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.sort_expr.is_empty() { + len += 1; + } + if !self.split_point.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalRangePartitioning", len)?; + if !self.sort_expr.is_empty() { + struct_ser.serialize_field("sortExpr", &self.sort_expr)?; + } + if !self.split_point.is_empty() { + struct_ser.serialize_field("splitPoint", &self.split_point)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalRangePartitioning { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "sort_expr", + "sortExpr", + "split_point", + "splitPoint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + SortExpr, + SplitPoint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "sortExpr" | "sort_expr" => Ok(GeneratedField::SortExpr), + "splitPoint" | "split_point" => Ok(GeneratedField::SplitPoint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalRangePartitioning; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalRangePartitioning") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut sort_expr__ = None; + let mut split_point__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::SortExpr => { + if sort_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("sortExpr")); + } + sort_expr__ = Some(map_.next_value()?); + } + GeneratedField::SplitPoint => { + if split_point__.is_some() { + return Err(serde::de::Error::duplicate_field("splitPoint")); + } + split_point__ = Some(map_.next_value()?); + } + } + } + Ok(PhysicalRangePartitioning { + sort_expr: sort_expr__.unwrap_or_default(), + split_point: split_point__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalRangePartitioning", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalRangeSplitPoint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.value.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalRangeSplitPoint", len)?; + if !self.value.is_empty() { + struct_ser.serialize_field("value", &self.value)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalRangeSplitPoint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "value", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Value, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "value" => Ok(GeneratedField::Value), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalRangeSplitPoint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalRangeSplitPoint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut value__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Value => { + if value__.is_some() { + return Err(serde::de::Error::duplicate_field("value")); + } + value__ = Some(map_.next_value()?); + } + } + } + Ok(PhysicalRangeSplitPoint { + value: value__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalRangeSplitPoint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalScalarSubqueryExprNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -21988,6 +22238,207 @@ impl<'de> serde::Deserialize<'de> for ProjectionNode { deserializer.deserialize_struct("datafusion.ProjectionNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for RangeRepartition { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.sort_expr.is_empty() { + len += 1; + } + if !self.split_point.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.RangeRepartition", len)?; + if !self.sort_expr.is_empty() { + struct_ser.serialize_field("sortExpr", &self.sort_expr)?; + } + if !self.split_point.is_empty() { + struct_ser.serialize_field("splitPoint", &self.split_point)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for RangeRepartition { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "sort_expr", + "sortExpr", + "split_point", + "splitPoint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + SortExpr, + SplitPoint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "sortExpr" | "sort_expr" => Ok(GeneratedField::SortExpr), + "splitPoint" | "split_point" => Ok(GeneratedField::SplitPoint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = RangeRepartition; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.RangeRepartition") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut sort_expr__ = None; + let mut split_point__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::SortExpr => { + if sort_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("sortExpr")); + } + sort_expr__ = Some(map_.next_value()?); + } + GeneratedField::SplitPoint => { + if split_point__.is_some() { + return Err(serde::de::Error::duplicate_field("splitPoint")); + } + split_point__ = Some(map_.next_value()?); + } + } + } + Ok(RangeRepartition { + sort_expr: sort_expr__.unwrap_or_default(), + split_point: split_point__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.RangeRepartition", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for RangeSplitPoint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.value.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.RangeSplitPoint", len)?; + if !self.value.is_empty() { + struct_ser.serialize_field("value", &self.value)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for RangeSplitPoint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "value", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Value, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "value" => Ok(GeneratedField::Value), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = RangeSplitPoint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.RangeSplitPoint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut value__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Value => { + if value__.is_some() { + return Err(serde::de::Error::duplicate_field("value")); + } + value__ = Some(map_.next_value()?); + } + } + } + Ok(RangeSplitPoint { + value: value__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.RangeSplitPoint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for RecursionUnnestOption { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -22416,6 +22867,9 @@ impl serde::Serialize for RepartitionNode { repartition_node::PartitionMethod::Hash(v) => { struct_ser.serialize_field("hash", v)?; } + repartition_node::PartitionMethod::Range(v) => { + struct_ser.serialize_field("range", v)?; + } } } struct_ser.end() @@ -22432,6 +22886,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { "round_robin", "roundRobin", "hash", + "range", ]; #[allow(clippy::enum_variant_names)] @@ -22439,6 +22894,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { Input, RoundRobin, Hash, + Range, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -22463,6 +22919,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { "input" => Ok(GeneratedField::Input), "roundRobin" | "round_robin" => Ok(GeneratedField::RoundRobin), "hash" => Ok(GeneratedField::Hash), + "range" => Ok(GeneratedField::Range), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -22503,6 +22960,13 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { return Err(serde::de::Error::duplicate_field("hash")); } partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(repartition_node::PartitionMethod::Hash) +; + } + GeneratedField::Range => { + if partition_method__.is_some() { + return Err(serde::de::Error::duplicate_field("range")); + } + partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(repartition_node::PartitionMethod::Range) ; } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index ea25dba4abd8c..06db2b35bbf72 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -214,7 +214,7 @@ pub struct SortNode { pub struct RepartitionNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - #[prost(oneof = "repartition_node::PartitionMethod", tags = "2, 3")] + #[prost(oneof = "repartition_node::PartitionMethod", tags = "2, 3, 4")] pub partition_method: ::core::option::Option, } /// Nested message and enum types in `RepartitionNode`. @@ -225,9 +225,23 @@ pub mod repartition_node { RoundRobin(u64), #[prost(message, tag = "3")] Hash(super::HashRepartition), + #[prost(message, tag = "4")] + Range(super::RangeRepartition), } } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeSplitPoint { + #[prost(message, repeated, tag = "1")] + pub value: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeRepartition { + #[prost(message, repeated, tag = "1")] + pub sort_expr: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct HashRepartition { #[prost(message, repeated, tag = "1")] pub hash_expr: ::prost::alloc::vec::Vec, @@ -1735,6 +1749,10 @@ pub struct FileScanExecConf { pub batch_size: ::core::option::Option, #[prost(message, optional, tag = "13")] pub projection_exprs: ::core::option::Option, + #[prost(bool, optional, tag = "14")] + pub partitioned_by_file_group: ::core::option::Option, + #[prost(message, optional, tag = "15")] + pub output_partitioning: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ParquetScanExecNode { @@ -2100,14 +2118,26 @@ pub struct PhysicalHashRepartition { pub partition_count: u64, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangePartitioning { + #[prost(message, repeated, tag = "1")] + pub sort_expr: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangeSplitPoint { + #[prost(message, repeated, tag = "1")] + pub value: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct RepartitionExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - /// oneof partition_method { + /// Legacy direct partitioning fields: /// uint64 round_robin = 2; /// PhysicalHashRepartition hash = 3; /// uint64 unknown = 4; - /// } + /// New partitioning variants are stored in `partitioning`. #[prost(message, optional, tag = "5")] pub partitioning: ::core::option::Option, #[prost(bool, tag = "6")] @@ -2115,7 +2145,7 @@ pub struct RepartitionExecNode { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct Partitioning { - #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3")] + #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3, 4")] pub partition_method: ::core::option::Option, } /// Nested message and enum types in `Partitioning`. @@ -2128,6 +2158,8 @@ pub mod partitioning { Hash(super::PhysicalHashRepartition), #[prost(uint64, tag = "3")] Unknown(u64), + #[prost(message, tag = "4")] + Range(super::PhysicalRangePartitioning), } } #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index cff90ba7e31b7..bed64879f5a92 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field}; use datafusion_common::datatype::DataTypeExt; use datafusion_common::{ - NullEquality, RecursionUnnestOption, Result, ScalarValue, TableReference, + NullEquality, RecursionUnnestOption, Result, ScalarValue, SplitPoint, TableReference, UnnestOptions, exec_datafusion_err, internal_err, plan_datafusion_err, }; use datafusion_execution::TaskContext; @@ -834,3 +834,14 @@ fn parse_required_expr( fn proto_error>(message: S) -> Error { Error::General(message.into()) } + +pub(super) fn parse_protobuf_range_split_point( + split_point: &protobuf::RangeSplitPoint, +) -> Result { + let values = split_point + .value + .iter() + .map(ScalarValue::try_from) + .collect::>()?; + Ok(SplitPoint::new(values)) +} diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 4e01ebabb5f69..5b2acfcac0e60 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -56,8 +56,8 @@ use datafusion_datasource_json::file_format::{ #[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RecursiveQuery, SkipType, - TableSource, Unnest, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, + RecursiveQuery, SkipType, TableSource, Unnest, }; use datafusion_expr::{ DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, @@ -70,6 +70,7 @@ use datafusion_expr::{ }; use self::to_proto::{serialize_expr, serialize_exprs}; +use crate::logical_plan::to_proto::serialize_range_split_point; use crate::logical_plan::to_proto::serialize_sorts; use datafusion_catalog::TableProvider; use datafusion_catalog::default_table_source::{provider_as_source, source_as_provider}; @@ -675,6 +676,16 @@ impl AsLogicalPlan for LogicalPlanNode { PartitionMethod::RoundRobin(partition_count) => { Partitioning::RoundRobinBatch(*partition_count as usize) } + PartitionMethod::Range(protobuf::RangeRepartition { + sort_expr: pb_sort_expr, + split_point, + }) => Partitioning::Range(RangePartitioning::try_new( + from_proto::parse_sorts(pb_sort_expr, ctx, extension_codec)?, + split_point + .iter() + .map(from_proto::parse_protobuf_range_split_point) + .collect::, _>>()?, + )?), }; LogicalPlanBuilder::from(input) @@ -1631,6 +1642,19 @@ impl AsLogicalPlan for LogicalPlanNode { Partitioning::RoundRobinBatch(partition_count) => { PartitionMethod::RoundRobin(*partition_count as u64) } + Partitioning::Range(range_partitioning) => { + let ordering = range_partitioning.ordering(); + let split_point = range_partitioning + .split_points() + .iter() + .map(serialize_range_split_point) + .collect::, _>>()?; + + PartitionMethod::Range(protobuf::RangeRepartition { + sort_expr: serialize_sorts(ordering, extension_codec)?, + split_point, + }) + } Partitioning::DistributeBy(_) => { return not_impl_err!("DistributeBy"); } diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 0599e48d2795c..86b980a109b90 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -21,7 +21,7 @@ use std::collections::HashMap; -use datafusion_common::{NullEquality, TableReference, UnnestOptions}; +use datafusion_common::{NullEquality, SplitPoint, TableReference, UnnestOptions}; use datafusion_expr::WriteOp; use datafusion_expr::dml::InsertOp; use datafusion_expr::expr::{ @@ -706,6 +706,18 @@ where .collect::, Error>>() } +pub(super) fn serialize_range_split_point( + split_point: &SplitPoint, +) -> Result { + Ok(protobuf::RangeSplitPoint { + value: split_point + .values() + .iter() + .map(TryInto::::try_into) + .collect::>()?, + }) +} + impl From for protobuf::TableReference { fn from(t: TableReference) -> Self { use protobuf::table_reference::TableReferenceEnum; diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 62f989a111ed2..37e7d6b5d84c2 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -24,7 +24,9 @@ use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; use chrono::{TimeZone, Utc}; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, not_impl_err}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, internal_datafusion_err, not_impl_err, +}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; @@ -51,7 +53,9 @@ use datafusion_physical_plan::expressions::{ }; use datafusion_physical_plan::joins::{HashExpr, SeededRandomState}; use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field}; -use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; +use datafusion_physical_plan::{ + Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, +}; use datafusion_proto_common::common::proto_error; use object_store::ObjectMeta; use object_store::path::Path; @@ -674,6 +678,14 @@ pub fn parse_protobuf_partitioning( proto_converter, ) } + Some(protobuf::partitioning::PartitionMethod::Range(range_partitioning)) => { + Ok(Some(parse_protobuf_range_partitioning( + range_partitioning, + ctx, + input_schema, + proto_converter, + )?)) + } Some(protobuf::partitioning::PartitionMethod::Unknown(partition_count)) => { Ok(Some(Partitioning::UnknownPartitioning( *partition_count as usize, @@ -685,6 +697,49 @@ pub fn parse_protobuf_partitioning( } } +fn parse_protobuf_range_partitioning( + range_partitioning: &protobuf::PhysicalRangePartitioning, + ctx: &PhysicalPlanDecodeContext<'_>, + input_schema: &Schema, + proto_converter: &dyn PhysicalProtoConverterExtension, +) -> Result { + let sort_exprs = parse_physical_sort_exprs( + &range_partitioning.sort_expr, + ctx, + input_schema, + proto_converter, + )?; + let sort_expr_count = sort_exprs.len(); + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!("Range partitioning requires non-empty ordering") + })?; + if ordering.len() != sort_expr_count { + return Err(internal_datafusion_err!( + "Range partitioning ordering must not contain duplicate expressions" + )); + } + let split_points = range_partitioning + .split_point + .iter() + .map(parse_protobuf_range_split_point) + .collect::>()?; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) +} + +fn parse_protobuf_range_split_point( + split_point: &protobuf::PhysicalRangeSplitPoint, +) -> Result { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::>()?; + Ok(SplitPoint::new(values)) +} + pub fn parse_protobuf_file_scan_schema( proto: &protobuf::FileScanExecConf, ) -> Result> { @@ -754,6 +809,12 @@ pub fn parse_protobuf_file_scan_config( )?; output_ordering.extend(LexOrdering::new(sort_exprs)); } + let output_partitioning = parse_protobuf_partitioning( + proto.output_partitioning.as_ref(), + ctx, + &schema, + proto_converter, + )?; // Parse projection expressions if present and apply to file source let file_source = if let Some(proto_projection_exprs) = &proto.projection_exprs { @@ -782,14 +843,18 @@ pub fn parse_protobuf_file_scan_config( file_source }; - let config = FileScanConfigBuilder::new(object_store_url, file_source) + let mut config_builder = FileScanConfigBuilder::new(object_store_url, file_source) .with_file_groups(file_groups) .with_constraints(constraints) .with_statistics(statistics) .with_limit(proto.limit.as_ref().map(|sl| sl.limit as usize)) .with_output_ordering(output_ordering) - .with_batch_size(proto.batch_size.map(|s| s as usize)) - .build(); + .with_output_partitioning(output_partitioning) + .with_batch_size(proto.batch_size.map(|s| s as usize)); + if proto.partitioned_by_file_group.unwrap_or(false) { + config_builder = config_builder.with_partitioned_by_file_group(true); + } + let config = config_builder.build(); Ok(config) } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 73fca1bbe6070..5d3fed8737ecf 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -44,7 +44,9 @@ use datafusion_physical_plan::expressions::{ use datafusion_physical_plan::joins::{HashExpr, HashTableLookupExpr}; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; -use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; +use datafusion_physical_plan::{ + Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, +}; use super::{ DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -658,6 +660,11 @@ pub fn serialize_partitioning( )), } } + Partitioning::Range(range) => protobuf::Partitioning { + partition_method: Some(protobuf::partitioning::PartitionMethod::Range( + serialize_range_partitioning(range, codec, proto_converter)?, + )), + }, Partitioning::UnknownPartitioning(partition_count) => protobuf::Partitioning { partition_method: Some(protobuf::partitioning::PartitionMethod::Unknown( *partition_count as u64, @@ -667,6 +674,40 @@ pub fn serialize_partitioning( Ok(serialized_partitioning) } +fn serialize_range_partitioning( + range: &RangePartitioning, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, +) -> Result { + Ok(protobuf::PhysicalRangePartitioning { + sort_expr: serialize_physical_sort_exprs( + range.ordering().iter().cloned(), + codec, + proto_converter, + )?, + split_point: range + .split_points() + .iter() + .map(serialize_range_split_point) + .collect::>()?, + }) +} + +fn serialize_range_split_point( + split_point: &SplitPoint, +) -> Result { + Ok(protobuf::PhysicalRangeSplitPoint { + value: split_point + .values() + .iter() + .map(|value| { + TryInto::::try_into(value) + .map_err(Into::into) + }) + .collect::>()?, + }) +} + fn serialize_when_then_expr( when_expr: &Arc, then_expr: &Arc, @@ -745,6 +786,11 @@ pub fn serialize_file_scan_config( serialize_physical_sort_exprs(order.to_vec(), codec, proto_converter)?; output_orderings.push(ordering) } + let output_partitioning = conf + .output_partitioning + .as_ref() + .map(|partitioning| serialize_partitioning(partitioning, codec, proto_converter)) + .transpose()?; // Fields must be added to the schema so that they can persist in the protobuf, // and then they are to be removed from the schema in `parse_protobuf_file_scan_config` @@ -804,6 +850,8 @@ pub fn serialize_file_scan_config( constraints: Some(conf.constraints.clone().into()), batch_size: conf.batch_size.map(|s| s as u64), projection_exprs, + partitioned_by_file_group: Some(conf.partitioned_by_file_group), + output_partitioning, }) } diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index fbe1af5617e42..555690d4db419 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -45,6 +45,7 @@ use std::vec; use datafusion::catalog::{TableProvider, TableProviderFactory}; use datafusion::datasource::DefaultTableSource; +use datafusion::datasource::empty::EmptyTable; use datafusion::datasource::file_format::arrow::ArrowFormatFactory; use datafusion::datasource::file_format::csv::CsvFormatFactory; use datafusion::datasource::file_format::parquet::ParquetFormatFactory; @@ -71,8 +72,8 @@ use datafusion_common::config::TableOptions; use datafusion_common::format::ExplainFormat; use datafusion_common::scalar::ScalarStructBuilder; use datafusion_common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference, - internal_datafusion_err, internal_err, not_impl_err, plan_err, + DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SplitPoint, + TableReference, internal_datafusion_err, internal_err, not_impl_err, plan_err, }; use datafusion_execution::TaskContext; use datafusion_expr::dml::CopyTo; @@ -84,9 +85,9 @@ use datafusion_expr::logical_plan::{Extension, UserDefinedLogicalNodeCore}; use datafusion_expr::{ Accumulator, AggregateUDF, ColumnarValue, ExprFunctionExt, ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, - PartitionEvaluator, ScalarUDF, Signature, TryCast, Volatility, WindowFrame, - WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, WindowUDF, - WindowUDFImpl, + PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, + Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, + WindowFunctionDefinition, WindowUDF, WindowUDFImpl, }; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::expr_fn::{ @@ -3276,9 +3277,7 @@ async fn roundtrip_empty_table_scan() -> Result<()> { Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( - &schema, - ))); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3300,9 +3299,7 @@ async fn roundtrip_empty_table_scan_with_projection() -> Result<()> { Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( - &schema, - ))); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3378,14 +3375,8 @@ async fn roundtrip_join_null_equality() -> Result<()> { let left_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); let right_schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); - ctx.register_table( - "t1", - Arc::new(datafusion::datasource::empty::EmptyTable::new(left_schema)), - )?; - ctx.register_table( - "t2", - Arc::new(datafusion::datasource::empty::EmptyTable::new(right_schema)), - )?; + ctx.register_table("t1", Arc::new(EmptyTable::new(left_schema)))?; + ctx.register_table("t2", Arc::new(EmptyTable::new(right_schema)))?; let left = ctx.table("t1").await?.into_optimized_plan()?; let right = ctx.table("t2").await?.into_optimized_plan()?; @@ -3406,3 +3397,91 @@ async fn roundtrip_join_null_equality() -> Result<()> { Ok(()) } + +// Single column, single split point range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_single_col() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} + +// Multi-column compound key with multiple split points for range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_multi_col() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("region", DataType::Utf8, false), + ])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("ts").sort(true, true), col("region").sort(true, true)], + vec![ + SplitPoint::new(vec![ + ScalarValue::Int64(Some(1000)), + ScalarValue::Utf8(Some("east".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(2000)), + ScalarValue::Utf8(Some("west".to_string())), + ]), + ], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} + +// Non-default sort options: descending with nulls last for range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_desc_nulls_last() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "score", + DataType::Float64, + true, + )])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("score").sort(false, false)], + vec![SplitPoint::new(vec![ScalarValue::Float64(Some(50.0))])], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index a0b7180ba90bc..674bf5a2d60ed 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -89,7 +89,8 @@ use datafusion::physical_plan::windows::{ }; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, - PhysicalExpr, PlanProperties, SendableRecordBatchStream, Statistics, displayable, + PhysicalExpr, PlanProperties, RangePartitioning, SendableRecordBatchStream, + SplitPoint, Statistics, displayable, }; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion::scalar::ScalarValue; @@ -1995,6 +1996,21 @@ fn roundtrip_repartition_preserve_order() -> Result<()> { roundtrip_test(Arc::new(repartition)) } +#[test] +fn roundtrip_range_partitioning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let range_partitioning = Partitioning::Range(RangePartitioning::new( + [PhysicalSortExpr::new_default(col("a", &schema)?)].into(), + vec![SplitPoint::new(vec![ScalarValue::Int64(Some(10))])], + )); + // RepartitionExec is used only to carry the partitioning through proto. + // Executing range repartitioning is intentionally unsupported. + let repartition = RepartitionExec::try_new(input, range_partitioning)?; + + roundtrip_test(Arc::new(repartition)) +} + #[test] fn roundtrip_interleave() -> Result<()> { let field_a = Field::new("col", DataType::Int64, false); @@ -4065,3 +4081,98 @@ fn test_sort_topk_with_dynamic_filter_roundtrip() -> Result<()> { Ok(()) } + +fn roundtrip_file_scan_config(scan_config: FileScanConfig) -> Result { + let exec_plan: Arc = DataSourceExec::from_data_source(scan_config); + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result_plan = + roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + let data_source_exec = result_plan + .downcast_ref::() + .expect("Expected DataSourceExec"); + let file_scan_config = data_source_exec + .data_source() + .downcast_ref::() + .expect("Expected FileScanConfig"); + Ok(file_scan_config.clone()) +} + +#[test] +fn roundtrip_parquet_exec_partitioned_by_file_group() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_partitioned_by_file_group(true) + .build(); + + assert!(roundtrip_file_scan_config(scan_config)?.partitioned_by_file_group); + Ok(()) +} + +#[test] +fn roundtrip_parquet_exec_output_partitioning() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let output_partitioning = + Partitioning::Hash(vec![Arc::new(Column::new("col", 0))], 1); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_output_partitioning(Some(output_partitioning.clone())) + .build(); + + assert_eq!( + roundtrip_file_scan_config(scan_config)?.output_partitioning, + Some(output_partitioning) + ); + + Ok(()) +} + +#[test] +fn roundtrip_parquet_exec_range_output_partitioning() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let output_partitioning = Partitioning::Range(RangePartitioning::new( + LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( + "col", 0, + )))]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![ + FileGroup::new(vec![PartitionedFile::new( + "/path/to/file-1.parquet".to_string(), + 1024, + )]), + FileGroup::new(vec![PartitionedFile::new( + "/path/to/file-2.parquet".to_string(), + 1024, + )]), + ]) + .with_output_partitioning(Some(output_partitioning.clone())) + .build(); + + assert_eq!( + roundtrip_file_scan_config(scan_config)?.output_partitioning, + Some(output_partitioning) + ); + + Ok(()) +} diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 0edde71b939f4..a83db2bfb947f 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -53,6 +53,8 @@ use datafusion::{ use datafusion_spark::SessionStateBuilderSpark; use crate::is_spark_path; +use range_partitioning::register_range_partitioned_table; + use async_trait::async_trait; use datafusion::common::cast::as_float64_array; use datafusion::execution::SessionStateBuilder; @@ -61,6 +63,8 @@ use log::info; use sqlparser::ast; use tempfile::TempDir; +mod range_partitioning; + /// Context for running tests pub struct TestContext { /// Context for running queries @@ -167,6 +171,10 @@ impl TestContext { info!("Registering table with many types"); register_table_with_many_types(test_ctx.session_ctx()).await; } + "range_partitioning.slt" => { + info!("Registering range partitioned table"); + register_range_partitioned_table(test_ctx.session_ctx()); + } "metadata.slt" | "arrow_field.slt" => { info!("Registering metadata table tables"); register_metadata_tables(test_ctx.session_ctx()).await; diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs new file mode 100644 index 0000000000000..4141e000145a8 --- /dev/null +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -0,0 +1,288 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fs::{create_dir_all, remove_dir_all, write}; +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion::catalog::streaming::StreamingTable; +use datafusion::common::{ScalarValue, SplitPoint}; +use datafusion::datasource::file_format::csv::CsvFormat; +use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, +}; +use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; +use datafusion::physical_expr::{ + Partitioning as PhysicalPartitioning, PhysicalSortExpr, + RangePartitioning as PhysicalRangePartitioning, expressions::col as physical_col, +}; +use datafusion::physical_plan::streaming::PartitionStream; +use datafusion::physical_plan::test::TestPartitionStream; +use datafusion::prelude::SessionContext; + +// ============================================================================== +// Range Partitioned Table (sqllogictest-only) +// ============================================================================== + +/// Registers a simple range-partitioned listing table for testing before +/// declaring such tables is supported via SQL. +pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { + let schema = Arc::new(Schema::new(vec![ + Field::new("range_key", DataType::Int32, false), + Field::new("non_range_key", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + let range_table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned"); + + register_csv_listing_table( + ctx, + "range_partitioned", + &range_table_dir, + Arc::clone(&schema), + [ + "1,1,10\n5,2,50\n", + "10,1,100\n15,2,150\n", + "20,1,200\n25,2,250\n", + "30,1,300\n35,2,350\n", + ], + Some(output_partitioning), + ); + + register_unbounded_range_stream_table( + ctx, + "unbounded_range_like", + Arc::clone(&schema), + [10, 20, 30], + [ + vec![(1, 1, 10), (5, 2, 50)], + vec![(10, 1, 100), (15, 2, 150)], + vec![(20, 1, 200), (25, 2, 250)], + vec![(30, 1, 300), (35, 2, 350)], + ], + ); + register_unbounded_range_stream_table( + ctx, + "unbounded_range_like_shifted", + Arc::clone(&schema), + [15, 20, 30], + [ + vec![(1, 1, 10), (5, 2, 50), (10, 1, 100)], + vec![(15, 2, 150)], + vec![(20, 1, 200), (25, 2, 250)], + vec![(30, 1, 300), (35, 2, 350)], + ], + ); + + let shifted_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(15))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_shifted", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_shifted"), + Arc::clone(&schema), + [ + "1,1,10\n5,2,50\n10,1,100\n", + "15,2,150\n", + "20,1,200\n25,2,250\n", + "30,1,300\n35,2,350\n", + ], + Some(shifted_output_partitioning), + ); + + // Same rows as `range_partitioned` but split into only three range + // partitions on `range_key`. Used to exercise the co-partition check when + // two Range inputs disagree on partition count. + let narrow_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_narrow", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_narrow"), + Arc::clone(&schema), + [ + "1,1,10\n5,2,50\n", + "10,1,100\n15,2,150\n", + "20,1,200\n25,2,250\n30,1,300\n35,2,350\n", + ], + Some(narrow_output_partitioning), + ); + + let sparse_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_csv_listing_table( + ctx, + "range_partitioned_sparse", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_sparse"), + schema, + [ + "5,2,50\n8,3,80\n", + "10,1,100\n", + "20,1,200\n", + "30,1,300\n40,4,400\n", + ], + Some(sparse_output_partitioning), + ); +} + +fn register_csv_listing_table( + ctx: &SessionContext, + name: &str, + table_dir: impl AsRef, + schema: Arc, + partitions: impl IntoIterator, + output_partitioning: Option, +) { + let table_dir = table_dir.as_ref(); + if table_dir.exists() { + remove_dir_all(table_dir).expect("test table dir should be removable"); + } + create_dir_all(table_dir).expect("test table dir should be created"); + for (idx, rows) in partitions.into_iter().enumerate() { + write(table_dir.join(format!("part-{idx}.csv")), rows) + .expect("test table csv partition should be written"); + } + + let table_path = format!( + "{}/", + table_dir + .to_str() + .expect("test table path should be valid utf8") + ); + let table_url = + ListingTableUrl::parse(&table_path).expect("test table url should parse"); + let options = + ListingOptions::new(Arc::new(CsvFormat::default().with_has_header(false))) + .with_output_partitioning(output_partitioning); + let config = ListingTableConfig::new(table_url) + .with_listing_options(options) + .with_schema(schema); + let table = + ListingTable::try_new(config).expect("test listing table should be valid"); + + ctx.register_table(name, Arc::new(table)) + .expect("test listing table registration should succeed"); +} + +fn register_unbounded_range_stream_table( + ctx: &SessionContext, + name: &str, + schema: Arc, + split_points: [i32; 3], + partition_rows: [Vec<(i32, i32, i32)>; 4], +) { + let output_partitioning = PhysicalPartitioning::Range( + PhysicalRangePartitioning::try_new( + [PhysicalSortExpr { + expr: physical_col("range_key", &schema) + .expect("range key should exist in stream schema"), + options: SortOptions::default(), + }] + .into(), + split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(), + ) + .expect("range partitioning should be valid"), + ); + let partitions = partition_rows + .into_iter() + .map(|rows| range_stream_partition(Arc::clone(&schema), &rows)) + .collect(); + + ctx.register_table( + name, + Arc::new( + StreamingTable::try_new(schema, partitions) + .expect("range stream table should be valid") + .with_infinite_table(true) + .with_output_partitioning(output_partitioning), + ), + ) + .expect("test stream table registration should succeed"); +} + +fn range_stream_partition( + schema: SchemaRef, + rows: &[(i32, i32, i32)], +) -> Arc { + let range_key: Vec = rows.iter().map(|(range_key, _, _)| *range_key).collect(); + let non_range_key: Vec = rows + .iter() + .map(|(_, non_range_key, _)| *non_range_key) + .collect(); + let value: Vec = rows.iter().map(|(_, _, value)| *value).collect(); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(range_key)) as ArrayRef, + Arc::new(Int32Array::from(non_range_key)) as ArrayRef, + Arc::new(Int32Array::from(value)) as ArrayRef, + ], + ) + .expect("range stream batch should be valid"); + Arc::new(TestPartitionStream::new_with_batches(vec![batch])) +} diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt new file mode 100644 index 0000000000000..ec51888c24b1e --- /dev/null +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -0,0 +1,1897 @@ +# 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. + +# The sqllogictest harness registers range_partitioned(range_key, non_range_key, value) +# as a CSV ListingTable with four declared range-partitioned file groups: +# +# partition 0: range_key in [..., 10), rows (1, 1, 10), (5, 2, 50) +# partition 1: range_key in [10, 20), rows (10, 1, 100), (15, 2, 150) +# partition 2: range_key in [20, 30), rows (20, 1, 200), (25, 2, 250) +# partition 3: range_key in [30, ...), rows (30, 1, 300), (35, 2, 350) + +statement ok +set datafusion.explain.physical_plan_only = true; + +########## +# TEST 1: Aggregate on Range Partition Column +# Scanning range_key preserves source Range partitioning metadata. +# Planning still inserts Hash repartitioning today; later optimizer PRs can +# use this baseline to show when the repartition is removed. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 2: Aggregate on Non-Range Column +# Projecting away range_key means the scan output no longer contains the +# expression needed to describe range partitioning, so it reports +# UnknownPartitioning with the same partition count. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query II +SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key ORDER BY non_range_key; +---- +1 610 +2 800 + + +########## +# TEST 3: Aggregate Reuses Range Subset Partitioning +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies grouping by (range_key, non_range_key). +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key ORDER BY range_key, non_range_key; +---- +1 1 10 +5 2 50 +10 1 100 +15 2 150 +20 1 200 +25 2 250 +30 1 300 +35 2 350 + + +########## +# TEST 4: Exact Range Aggregate Below Subset Threshold +# Even when subset satisfaction is disabled, exact Range([range_key]) +# satisfies GROUP BY range_key when repartitioning would not increase +# partition count. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + + +########## +# TEST 5: Range Subset Aggregate Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of GROUP BY (range_key, non_range_key), +# so it should not satisfy the aggregate key when subset satisfaction is +# disabled. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + + +########## +# TEST 6: Aggregate Rehashes Below Subset Threshold +# With subset threshold 5 and only 4 input partitions, planning repartitions +# to increase parallelism instead of reusing Range partitioning. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + + +########## +# TEST 7: Aggregate Preserves Range When Preserve File Threshold Met +# With preserve-file threshold 1 and 4 input partitions, Range is preserved +# even though target_partitions is 5. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 8: Aggregate Rehashes When Preserve File Threshold Not Met +# With preserve-file threshold 5 and only 4 input partitions, planning can +# repartition to increase parallelism. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 5; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + + +########## +# TEST 9: Join on Range Partition Column +# A partitioned inner hash join requires co-partitioned KeyPartitioned inputs. +# Compatible Range layouts satisfy both the per-child key requirements and the +# cross-child layout requirement, so no Hash repartitioning is inserted. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 10: Incompatible Range Join Repartitions +# Both inputs are independently range partitioned on range_key, but their split +# points differ. The per-child key requirements can be satisfied by Range, but +# the co-partitioned layout requirement cannot, so Hash repartitioning repairs +# both sides. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 11: Non-Range Join Repartitions +# Range([range_key]) does not satisfy KeyPartitioned([non_range_key]), so +# planning inserts Hash repartitioning on the actual join key. +########## + +query TT +EXPLAIN SELECT l.non_range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(non_range_key@0, non_range_key@0)], projection=[non_range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT l.non_range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.non_range_key = r.non_range_key +ORDER BY l.non_range_key, l.value, r.value; +---- +1 10 10 +1 10 100 +1 10 200 +1 10 300 +1 100 10 +1 100 100 +1 100 200 +1 100 300 +1 200 10 +1 200 100 +1 200 200 +1 200 300 +1 300 10 +1 300 100 +1 300 200 +1 300 300 +2 50 50 +2 50 150 +2 50 250 +2 50 350 +2 150 50 +2 150 150 +2 150 250 +2 150 350 +2 250 50 +2 250 150 +2 250 250 +2 250 350 +2 350 50 +2 350 150 +2 350 250 +2 350 350 + +########## +# TEST 12: Left-Side Range Hash Joins +# Compatible Range layouts satisfy left-side partitioned hash join +# requirements without Hash repartitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--FilterExec: value@1 <= 150 +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 NULL +25 250 NULL +30 300 NULL +35 350 NULL + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(range_key@0, range_key@0)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--FilterExec: value@1 <= 150, projection=[range_key@0] +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(range_key@0, range_key@0)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--FilterExec: value@1 <= 150, projection=[range_key@0] +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +20 200 +25 250 +30 300 +35 350 + +########## +# TEST 13: Left-Side Range Hash Joins With Incomplete Range Keys +# Range partitioning covers only range_key, so joins requiring additional +# or different keys are repaired with Hash repartitioning. +########## + +# Range([range_key]) is only a subset of the composite join key, so the +# co-partitioned hash join requirement is repaired with Hash repartitioning. +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, non_range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0), (non_range_key@1, non_range_key@1)], projection=[range_key@0, non_range_key@1, value@2, value@5] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)----FilterExec: value@2 <= 150 +06)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +# Range([range_key]) does not satisfy a join keyed on non_range_key. +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(non_range_key@1, non_range_key@0)], projection=[range_key@0, non_range_key@1, value@2, value@4] +02)--RepartitionExec: partitioning=Hash([non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +########## +# TEST 14: Left-Side Range Hash Joins With Incompatible Range Layouts +# Different split points or partition counts do not satisfy the +# co-partitioned layout requirement. +########## + +# Different split points do not satisfy the co-partitioned layout requirement. +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +# Different partition counts do not satisfy the co-partitioned layout +# requirement. +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_narrow r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=3 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20)], 3), file_type=csv, has_header=false + +########## +# TEST 15: LeftMark Subqueries Over Range Hash Joins +# SQL IN subqueries decorrelate to LeftMark joins. These queries pin matched, +# unmatched, and NULL marker behavior over compatible Range inputs. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150); +---- +physical_plan +01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2] +02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----FilterExec: value@1 <= 150, projection=[range_key@0] +05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150) +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT CASE WHEN value <= 150 THEN range_key ELSE NULL END + FROM range_partitioned) +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +########## +# TEST 16: Compatible Range Join Repartitions to Increase Parallelism +# Co-partitioning satisfaction does not prevent a repartition that increases +# parallelism. With target_partitions larger than the Range partition count, +# both sides are hash repartitioned. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 17: Preserve File Partitions Preserves Range Join Inputs +# preserve_file_partitions preserves compatible Range inputs for partitioned +# joins even when target_partitions is higher than the input partition count. +########## + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +########## +# TEST 18: Nested Range Joins +# Compatible Range partitioning is preserved through the lower join, allowing +# the upper join to consume it without Hash repartitioning either input. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value, s.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +JOIN range_partitioned s ON r.range_key = s.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@2, range_key@0)], projection=[range_key@0, value@1, value@3, value@5] +02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT l.range_key, l.value, r.value, s.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +JOIN range_partitioned s ON r.range_key = s.range_key +ORDER BY l.range_key; +---- +1 10 10 10 +5 50 50 50 +10 100 100 100 +15 150 150 150 +20 200 200 200 +25 250 250 250 +30 300 300 300 +35 350 350 350 + +########## +# TEST 19: Range Aggregates Feed Range Join +# Aggregates on range_key preserve reusable partitioning for the downstream +# partitioned join. +########## + +query TT +EXPLAIN WITH + l AS ( + SELECT range_key, SUM(value) AS l_sum + FROM range_partitioned + GROUP BY range_key + ), + r AS ( + SELECT range_key, SUM(value) AS r_sum + FROM range_partitioned + GROUP BY range_key + ) +SELECT l.range_key, l.l_sum, r.r_sum +FROM l JOIN r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, l_sum@1, r_sum@3] +02)--ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value)@1 as l_sum] +03)----AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value)@1 as r_sum] +06)----AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +07)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +WITH + l AS ( + SELECT range_key, SUM(value) AS l_sum + FROM range_partitioned + GROUP BY range_key + ), + r AS ( + SELECT range_key, SUM(value) AS r_sum + FROM range_partitioned + GROUP BY range_key + ) +SELECT l.range_key, l.l_sum, r.r_sum +FROM l JOIN r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 20: Range Join Feeds Aggregate +# The join preserves compatible Range partitioning on range_key, allowing the +# aggregate above it to avoid Hash repartitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, SUM(l.value + r.value) +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +GROUP BY l.range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)] +02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT l.range_key, SUM(l.value + r.value) +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +GROUP BY l.range_key +ORDER BY l.range_key; +---- +1 20 +5 100 +10 200 +15 300 +20 400 +25 500 +30 600 +35 700 + +########## +# TEST 21: Right Join on Range Partition Column +# Compatible Range inputs satisfy the join's partitioning requirements, so no +# Hash repartitioning is inserted. The left filter keeps its Range partitioning +# and the unmatched right rows above 150 are preserved. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--FilterExec: value@1 <= 150 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +NULL 20 200 +NULL 25 250 +NULL 30 300 +NULL 35 350 + +########## +# TEST 22: Right Semi Join on Range Partition Column +# Compatible Range inputs avoid Hash repartitioning for RightSemi joins. +# Only right rows with a match on the filtered left side are returned. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT SEMI JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(range_key@0, range_key@0)] +02)--FilterExec: value@1 <= 150, projection=[range_key@0] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT SEMI JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 + +########## +# TEST 23: Right Anti Join on Range Partition Column +# Compatible Range inputs avoid Hash repartitioning for RightAnti joins. +# Only right rows without a match on the filtered left side are returned. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT ANTI JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=RightAnti, on=[(range_key@0, range_key@0)] +02)--FilterExec: value@1 <= 150, projection=[range_key@0] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT ANTI JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +20 200 +25 250 +30 300 +35 350 + +########## +# TEST 24: Incompatible Range Right Join Repartitions +# The split points of the two inputs differ, so the co-partitioned layout +# requirement cannot be satisfied and Hash repartitioning repairs both sides +# of the right join. Results stay correct on the repartitioned path. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----FilterExec: value@1 <= 150 +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +06)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +NULL 20 200 +NULL 25 250 +NULL 30 300 +NULL 35 350 + +########## +# TEST 25: Composite-Key Right Join Repartitions +# Range([range_key]) does not satisfy a partitioned join on +# (range_key, non_range_key), so both sides repartition on the full key. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0), (non_range_key@1, non_range_key@1)], projection=[range_key@0, non_range_key@1, value@2, value@5] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key +ORDER BY l.range_key; +---- +1 1 10 10 +5 2 50 50 +10 1 100 100 +15 2 150 150 +20 1 200 200 +25 2 250 250 +30 1 300 300 +35 2 350 350 + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +########## +# TEST 26: Right Join with Mismatched Range Partition Counts Repartitions +# Both inputs are range partitioned on range_key, but declare a different number +# of partitions (four vs three). The per-child key requirements can be satisfied +# by Range, but the co-partitioned layout requirement cannot, so Hash +# repartitioning repairs both sides of the right join. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned_narrow r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=3 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20)], 3), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned_narrow r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +200 20 200 +250 25 250 +300 30 300 +350 35 350 + +########## +# TEST 27: Right Join on Non-Range Key Repartitions +# Both inputs expose Range([range_key]), but the join key is non_range_key. +# Range([range_key]) does not satisfy KeyPartitioned([non_range_key]), so +# planning inserts Hash repartitioning on the actual join key for the right join. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT non_range_key, value FROM range_partitioned WHERE range_key < 10) l +RIGHT JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(non_range_key@0, non_range_key@1)], projection=[value@1, range_key@2, value@4] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----FilterExec: range_key@0 < 10, projection=[non_range_key@1, value@2] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)--RepartitionExec: partitioning=Hash([non_range_key@1], 4), input_partitions=4 +06)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT non_range_key, value FROM range_partitioned WHERE range_key < 10) l +RIGHT JOIN range_partitioned r ON l.non_range_key = r.non_range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +10 10 100 +50 15 150 +10 20 200 +50 25 250 +10 30 300 +50 35 350 + +########## +# TEST 28: Mark Join Marker Semantics +# Mark joins preserve matched, unmatched, and NULL-key marker behavior over +# range-partitioned inputs. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150); +---- +physical_plan +01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2] +02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----FilterExec: value@1 <= 150, projection=[range_key@0] +05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +# Matched rows have mark=true and are returned; unmatched rows have +# mark=false and are only returned when non_range_key = 2. +query II +SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150) +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +# NULL join keys on the build side never match: rows whose keys only "match" +# the NULL entries keep a non-true marker and are filtered out unless the +# non_range_key = 2 disjunct covers them. +query II +SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT CASE WHEN value <= 150 THEN range_key ELSE NULL END FROM range_partitioned) +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +########## +# TEST 29: Sort Merge Join Avoids Repartition for Compatible Range Inputs +# Compatible Range inputs satisfy SortMergeJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] +03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +05)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +06)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 30: Sort Merge Join Repartitions Incompatible Range Inputs +# Different Range split points do not satisfy SortMergeJoinExec's +# co-partitioned KeyPartitioned requirements. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] +03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +06)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +07)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +08)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +########## +# TEST 31: Symmetric Hash Join Avoids Repartition for Compatible Range Inputs +# Compatible Range streams satisfy SymmetricHashJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +03)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +04)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) + +query III rowsort +SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +1 10 10 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 +5 50 50 + +########## +# TEST 32: Symmetric Hash Join Repartitions Incompatible Range Inputs +# Different Range split points do not satisfy SymmetricHashJoinExec's +# co-partitioned KeyPartitioned requirements. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +04)------StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +05)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +06)------StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4) + +query III rowsort +SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; +---- +1 10 10 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 +5 50 50 + +########## +# TEST 33: Full Outer Join on Range Partition Column +# Full partitioned hash joins also opt in to Range satisfying KeyPartitioned +# requirements, so compatible Range layouts avoid Hash repartitioning here too. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 34: Full Outer Join Incompatible Range Repartitions +# Same as TEST 10, but for Full: differing split points between the two +# Range-partitioned inputs still require Hash repartitioning to co-partition. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 35: Full Outer Join Produces Matched and Unmatched Rows +# `range_partitioned` and `range_partitioned_sparse` share the same Range +# split points/partition count but only partially overlapping range_key +# values, so this exercises matched rows, left-only unmatched rows (NULLs on +# the right), and right-only unmatched rows (NULLs on the left) while still +# avoiding Hash repartitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, r.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, range_key@2, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT l.range_key, r.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +ORDER BY l.range_key, r.range_key; +---- +1 NULL 10 NULL +5 5 50 50 +10 10 100 100 +15 NULL 150 NULL +20 20 200 200 +25 NULL 250 NULL +30 30 300 300 +35 NULL 350 NULL +NULL 8 NULL 80 +NULL 40 NULL 400 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +statement ok +reset datafusion.optimizer.repartition_joins; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +########## +# TEST 36: Union of Range Partitioned Inputs +# Each input exposes the same Range partitioning on range_key, so the optimizer +# converts UnionExec to InterleaveExec to avoid redundant repartitioning. +########## + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned; +---- +physical_plan +01)InterleaveExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +ORDER BY range_key, value; +---- +1 10 +1 10 +5 50 +5 50 +10 100 +10 100 +15 150 +15 150 +20 200 +20 200 +25 250 +25 250 +30 300 +30 300 +35 350 +35 350 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + + +########## +# TEST 37: Window on Range Partition Column +# Range([range_key]) colocates equal range_key values, so +# PARTITION BY range_key is satisfied without a hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 38: Unbounded-Frame Window on Range Partition Column +# The unbounded frame makes DataFusion use WindowAggExec instead of +# BoundedWindowAggExec, which likewise reuses Range partitioning without a +# hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] +02)--WindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 39: Window on Non-Range Column Rehashes +# Range([range_key]) does not colocate non_range_key values, so +# PARTITION BY non_range_key still requires a hash repartition. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[non_range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value) FROM range_partitioned ORDER BY non_range_key, value; +---- +1 10 10 +1 100 110 +1 200 310 +1 300 610 +2 50 50 +2 150 200 +2 250 450 +2 350 800 + + +########## +# TEST 40: Unbounded-Frame Window on Non-Range Column Rehashes +# The unbounded frame makes DataFusion use WindowAggExec; Range([range_key]) +# does not colocate non_range_key values, so PARTITION BY non_range_key +# still requires a hash repartition. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] +02)--WindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +03)----SortExec: expr=[non_range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned ORDER BY non_range_key, value; +---- +1 10 610 +1 100 610 +1 200 610 +1 300 610 +2 50 800 +2 150 800 +2 250 800 +2 350 800 + + +########## +# TEST 41: Window Subset Satisfaction on Range Partition Column +# With the subset threshold met, Range([range_key]) satisfies +# PARTITION BY (range_key, non_range_key): equal composite keys share the +# same range_key, so they are already colocated. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, non_range_key@1 ASC NULLS LAST, value@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 42: Window Subset Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of PARTITION BY +# (range_key, non_range_key), so it should not satisfy the window key when +# subset satisfaction is disabled. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, non_range_key@1 ASC NULLS LAST, value@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 43: Window Without Partition Keys Uses a Single Partition +# A window with no PARTITION BY requires a single partition; range +# partitioning is not applicable. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortPreservingMergeExec: [value@1 ASC NULLS LAST] +04)------SortExec: expr=[value@1 ASC NULLS LAST], preserve_partitioning=[true] +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 60 +10 160 +15 310 +20 510 +25 760 +30 1060 +35 1410 + + + +########## +# TEST 44: PartitionedTopK on Range Partition Column +# Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning. +########## + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0], order=[value@1 DESC] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key; +---- +1 10 1 +5 50 1 +10 100 1 +15 150 1 +20 200 1 +25 250 1 +30 300 1 +35 350 1 + + +########## +# TEST 45: PartitionedTopK on Non-Range Column +# Partitioning on a non-range key cannot reuse Range([range_key]) and +# requires hash repartitioning. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fetch=1, partition=[non_range_key@0], order=[value@1 DESC] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY non_range_key; +---- +1 300 1 +2 350 1 + + +########## +# TEST 46: PartitionedTopK Reuses Range Subset Partitioning +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies partitioning by (range_key, non_range_key). +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key, non_range_key; +---- +1 1 10 1 +5 2 50 1 +10 1 100 1 +15 2 150 1 +20 1 200 1 +25 2 250 1 +30 1 300 1 +35 2 350 1 + + +########## +# TEST 47: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), +# so it should not satisfy the TopK partition key when subset satisfaction is +# disabled. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +statement ok +reset datafusion.explain.physical_plan_only; + +statement ok +reset datafusion.optimizer.enable_window_topn; + +########## +# TEST 48: Subset of Inputs Compatible Does Not Trigger InterleaveExec +# In a three-way union, two inputs share the same Range split points [10,20,30] +# while the third has a partially-overlapping but different set [15,20,30]. +# can_interleave requires ALL inputs to match, so UnionExec is kept. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +1 10 +5 50 +5 50 +5 50 +10 100 +10 100 +10 100 +15 150 +15 150 +15 150 +20 200 +20 200 +20 200 +25 250 +25 250 +25 250 +30 300 +30 300 +30 300 +35 350 +35 350 +35 350 + +########## +# TEST 49: Incompatible Range Split Points Falls Back to UnionExec +# Two range-partitioned inputs with different split points cannot be interleaved, +# so the optimizer keeps UnionExec instead of converting to InterleaveExec. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +5 50 +5 50 +10 100 +10 100 +15 150 +15 150 +20 200 +20 200 +25 250 +25 250 +30 300 +30 300 +35 350 +35 350 + +########## +# TEST 50: InterleaveExec Propagates Range Partitioning to Aggregate +# InterleaveExec outputs the same Range partitioning as its compatible inputs, +# allowing a downstream aggregate on range_key to run SinglePartitioned without +# a Hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(value)] +02)--InterleaveExec +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query II +SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key ORDER BY range_key; +---- +1 20 +5 100 +10 200 +15 300 +20 400 +25 500 +30 600 +35 700 + +statement ok +reset datafusion.explain.physical_plan_only; diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index bf9ce26b35537..2f885991e3497 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -624,3 +624,48 @@ DROP TABLE window_topn_nulls; # Reset config to default (false) statement ok SET datafusion.optimizer.enable_window_topn = false; + +statement ok +create table t(c1 int, c2 int) as values (1, 2), (3, 4); + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.repartition_windows = false; + +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +query TT +EXPLAIN SELECT * FROM ( + SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2 DESC) as rn + FROM t +) WHERE rn <= 1; +---- +logical_plan +01)Projection: t.c1, t.c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: t projection=[c1, c2] +physical_plan +01)ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true +03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------PartitionedTopKExec: fetch=1, partition=[c1@0], order=[c2@1 DESC] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_windows = true; + +statement ok +set datafusion.execution.batch_size = 8192; + +statement ok +set datafusion.optimizer.enable_window_topn = false; 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"