From 2fd1e8d40b2cc0c7fd69360cbf7b19164a6c3892 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=ADa=20Adriana?= Date: Thu, 13 Aug 2026 16:44:04 +0200 Subject: [PATCH] Revert "Branch 54 range partitioning cherry pick" --- datafusion/catalog-listing/src/helpers.rs | 4 +- datafusion/catalog-listing/src/options.rs | 54 +- datafusion/catalog-listing/src/table.rs | 290 +-- datafusion/catalog/src/streaming.rs | 41 +- datafusion/common/src/lib.rs | 2 - datafusion/common/src/partitioning.rs | 104 - .../core/src/datasource/listing/table.rs | 258 +-- datafusion/core/src/physical_planner.rs | 87 +- .../enforce_distribution.rs | 668 +----- .../physical_optimizer/projection_pushdown.rs | 14 +- .../physical_optimizer/sanity_checker.rs | 220 +- .../tests/physical_optimizer/test_utils.rs | 18 +- .../tests/user_defined/user_defined_plan.rs | 8 +- .../datasource/src/file_scan_config/mod.rs | 224 +- datafusion/datasource/src/sink.rs | 14 +- datafusion/expr/src/logical_plan/display.rs | 17 - datafusion/expr/src/logical_plan/mod.rs | 4 +- datafusion/expr/src/logical_plan/plan.rs | 283 +-- datafusion/expr/src/logical_plan/tree_node.rs | 15 - .../ffi/src/physical_expr/partitioning.rs | 156 +- datafusion/ffi/src/plan_properties.rs | 46 +- datafusion/physical-expr/src/lib.rs | 10 +- datafusion/physical-expr/src/partitioning.rs | 1249 ++++------- datafusion/physical-expr/src/physical_expr.rs | 101 +- .../src/enforce_distribution.rs | 481 ++--- .../src/enforce_sorting/mod.rs | 22 +- .../src/output_requirements.rs | 38 +- .../physical-optimizer/src/sanity_checker.rs | 22 +- .../physical-optimizer/src/window_topn.rs | 74 +- .../physical-plan/src/aggregates/mod.rs | 14 +- datafusion/physical-plan/src/analyze.rs | 8 +- .../src/distribution_requirements.rs | 359 ---- .../physical-plan/src/execution_plan.rs | 32 +- .../physical-plan/src/joins/cross_join.rs | 8 +- .../physical-plan/src/joins/hash_join/exec.rs | 424 +--- .../src/joins/nested_loop_join.rs | 8 +- .../src/joins/piecewise_merge_join/exec.rs | 8 +- .../src/joins/sort_merge_join/exec.rs | 15 +- .../src/joins/symmetric_hash_join.rs | 20 +- datafusion/physical-plan/src/joins/utils.rs | 69 +- datafusion/physical-plan/src/lib.rs | 6 +- datafusion/physical-plan/src/limit.rs | 6 +- .../physical-plan/src/recursive_query.rs | 8 +- .../physical-plan/src/repartition/mod.rs | 858 +------- .../physical-plan/src/sorts/partial_sort.rs | 8 +- .../src/sorts/partitioned_topk.rs | 8 +- datafusion/physical-plan/src/sorts/sort.rs | 11 +- .../src/sorts/sort_preserving_merge.rs | 14 +- datafusion/physical-plan/src/streaming.rs | 48 +- datafusion/physical-plan/src/union.rs | 138 +- datafusion/physical-plan/src/unnest.rs | 8 +- .../src/windows/bounded_window_agg_exec.rs | 15 +- .../src/windows/window_agg_exec.rs | 15 +- datafusion/proto/proto/datafusion.proto | 26 +- datafusion/proto/src/generated/pbjson.rs | 464 ---- datafusion/proto/src/generated/prost.rs | 40 +- .../proto/src/logical_plan/from_proto.rs | 13 +- datafusion/proto/src/logical_plan/mod.rs | 28 +- datafusion/proto/src/logical_plan/to_proto.rs | 14 +- .../proto/src/physical_plan/from_proto.rs | 75 +- .../proto/src/physical_plan/to_proto.rs | 50 +- .../tests/cases/roundtrip_logical_plan.rs | 117 +- .../tests/cases/roundtrip_physical_plan.rs | 113 +- datafusion/sqllogictest/src/test_context.rs | 8 - .../src/test_context/range_partitioning.rs | 288 --- .../test_files/range_partitioning.slt | 1897 ----------------- .../sqllogictest/test_files/window_topn.slt | 45 - .../logical_plan/producer/rel/exchange_rel.rs | 10 - 68 files changed, 1008 insertions(+), 8812 deletions(-) delete mode 100644 datafusion/common/src/partitioning.rs delete mode 100644 datafusion/physical-plan/src/distribution_requirements.rs delete mode 100644 datafusion/sqllogictest/src/test_context/range_partitioning.rs delete mode 100644 datafusion/sqllogictest/test_files/range_partitioning.slt diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index edb3294051cde..0389b3cb17fe9 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>( } } -pub fn filter_partitioned_file( +fn filter_partitions( 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_partitioned_file(pf, filters, &df_schema)) + futures::future::ready(filter_partitions(pf, filters, &df_schema)) }) .boxed()) } diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 65444706dd02f..0ab15e05abba1 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::{Partitioning, SortExpr}; +use datafusion_expr::SortExpr; use futures::StreamExt; use futures::TryStreamExt; use itertools::Itertools; @@ -61,46 +61,6 @@ 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 { @@ -118,7 +78,6 @@ impl ListingOptions { collect_stat: false, target_partitions: 1, file_sort_order: vec![], - output_partitioning: None, } } @@ -177,17 +136,6 @@ 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 3627e560a5096..7ee743a6abe71 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -16,17 +16,14 @@ // under the License. use crate::config::SchemaSource; -use crate::helpers::{ - expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list, -}; +use crate::helpers::{expr_applicable_for_cols, 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, DFSchema, SchemaExt, Statistics, internal_datafusion_err, plan_err, - project_schema, + Constraints, SchemaExt, Statistics, internal_datafusion_err, plan_err, project_schema, }; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; @@ -41,10 +38,8 @@ 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, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType, -}; -use datafusion_physical_expr::{create_lex_ordering, create_physical_partitioning}; +use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; +use datafusion_physical_expr::create_lex_ordering; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::ExecutionPlan; @@ -438,20 +433,6 @@ 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( @@ -519,19 +500,9 @@ impl TableProvider for ListingTable { can_be_evaluated_for_partition_pruning(&table_partition_col_names, filter) }); - 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()); + // 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 ListFilesResult { file_groups: mut partitioned_file_lists, @@ -551,19 +522,17 @@ impl TableProvider for ListingTable { state.execution_props(), &partitioned_file_lists, )?; - 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 + match state + .config_options() + .execution + .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, - file_group_count, + self.options.target_partitions, ) }) }) @@ -571,7 +540,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() <= file_group_count { + if new_groups.len() <= self.options.target_partitions { partitioned_file_lists = new_groups; } else { log::debug!( @@ -582,41 +551,6 @@ 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 { @@ -626,23 +560,24 @@ 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, scan_config) + .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(), + ) .await?; Ok(ScanResult::new(plan)) @@ -754,41 +689,28 @@ 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 { - if let Some(output_partitioning) = self.options.output_partitioning.as_ref() { - self.list_files_for_declared_output_partitioning( - ctx, - output_partitioning, - filters, - ) - .await + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? } else { - 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)> { + return Ok(ListFilesResult { + file_groups: vec![], + statistics: Statistics::new_unknown(&self.file_schema), + grouped_by_partition: false, + }); + }; // 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, - listing_time_filters, + filters, &self.options.file_extension, &self.options.table_partition_cols, ) @@ -802,7 +724,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) @@ -814,34 +736,8 @@ impl ListingTable { .boxed() .buffer_unordered(ctx.config_options().execution.meta_fetch_concurrency); - 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?; + let (file_group, inexact_stats) = + get_files_with_limit(files, limit, self.options.collect_stat).await?; // Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N // @@ -850,102 +746,28 @@ 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(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, - ) - } + 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) } else { - (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)? + 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, + ) + } } else { - return Ok(ListFilesResult { - file_groups: vec![], - statistics: Statistics::new_unknown(&self.file_schema), - grouped_by_partition: false, - }); + ( + file_group.split_files(self.options.target_partitions), + 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 5bfecef1fb2ed..e609877c2b778 100644 --- a/datafusion/catalog/src/streaming.rs +++ b/datafusion/catalog/src/streaming.rs @@ -24,10 +24,7 @@ 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::projection::ProjectionMapping; -use datafusion_physical_expr::{ - EquivalenceProperties, LexOrdering, Partitioning, create_physical_sort_exprs, -}; +use datafusion_physical_expr::{LexOrdering, create_physical_sort_exprs}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use log::debug; @@ -41,7 +38,6 @@ pub struct StreamingTable { partitions: Vec>, infinite: bool, sort_order: Vec, - output_partitioning: Option, } impl StreamingTable { @@ -66,7 +62,6 @@ impl StreamingTable { partitions, infinite: false, sort_order: vec![], - output_partitioning: None, }) } @@ -81,33 +76,6 @@ 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] @@ -151,16 +119,13 @@ impl TableProvider for StreamingTable { vec![] }; - let exec = StreamingTableExec::try_new( + Ok(Arc::new(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 2f6d9848b6e55..e865c548bb554 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -30,7 +30,6 @@ mod dfschema; mod functional_dependencies; mod join_type; mod param_value; -mod partitioning; mod schema_reference; mod table_reference; mod unnest; @@ -93,7 +92,6 @@ 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 deleted file mode 100644 index 8a7212c2e3089..0000000000000 --- a/datafusion/common/src/partitioning.rs +++ /dev/null @@ -1,104 +0,0 @@ -// 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 0d6231ce27833..d14ec1f56dce2 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, TimeUnit}; + use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion_catalog::TableProvider; use datafusion_catalog_listing::{ ListingOptions, ListingTable, ListingTableConfig, SchemaSource, @@ -139,17 +139,12 @@ 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, Partitioning as LogicalPartitioning, - RangePartitioning as LogicalRangePartitioning, - }; + use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator}; use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::expressions::{Column, binary}; + use datafusion_physical_expr::expressions::binary; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::empty::EmptyExec; - use datafusion_physical_plan::{ - ExecutionPlanProperties, Partitioning, RangePartitioning, SplitPoint, collect, - }; + use datafusion_physical_plan::{ExecutionPlanProperties, collect}; use std::collections::HashMap; use std::io::Write; use std::sync::Arc; @@ -182,21 +177,6 @@ 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(); @@ -1306,236 +1286,6 @@ 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 2a41d83aac957..43af743caa030 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -32,11 +32,10 @@ use crate::logical_expr::{ Aggregate, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, Values, Window, }; use crate::logical_expr::{ - Expr, LogicalPlan, PlanType, Repartition, UserDefinedLogicalNode, -}; -use crate::physical_expr::{ - create_physical_expr, create_physical_exprs, create_physical_partitioning, + Expr, LogicalPlan, Partitioning as LogicalPartitioning, PlanType, Repartition, + UserDefinedLogicalNode, }; +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; @@ -53,8 +52,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, PhysicalExpr, WindowExpr, - displayable, windows, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, Partitioning, PhysicalExpr, + WindowExpr, displayable, windows, }; use crate::schema_equivalence::schema_satisfied_by; @@ -1248,11 +1247,25 @@ impl DefaultPhysicalPlanner { }) => { let physical_input = children.one()?; let input_dfschema = input.as_ref().schema(); - let physical_partitioning = create_physical_partitioning( - partitioning_scheme, - input_dfschema, - execution_props, - )?; + 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" + ); + } + }; Arc::new(RepartitionExec::try_new( physical_input, physical_partitioning, @@ -3206,8 +3219,8 @@ mod tests { use crate::datasource::MemTable; use crate::datasource::file_format::options::CsvReadOptions; use crate::physical_plan::{ - DisplayAs, DisplayFormatType, Partitioning, PlanProperties, - SendableRecordBatchStream, expressions, + DisplayAs, DisplayFormatType, PlanProperties, SendableRecordBatchStream, + expressions, }; use crate::prelude::{SessionConfig, SessionContext}; use crate::test_util::{scan_empty, scan_empty_with_partitions}; @@ -3218,8 +3231,8 @@ mod tests { use arrow_schema::{FieldRef, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::{ - DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _, - assert_batches_eq, assert_contains, + DFSchemaRef, ScalarValue, TableReference, ToDFSchema as _, assert_batches_eq, + assert_contains, }; use datafusion_execution::TaskContext; use datafusion_execution::runtime_env::RuntimeEnv; @@ -3228,8 +3241,8 @@ mod tests { use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder, - Partitioning as LogicalPartitioning, RangePartitioning, Signature, TableSource, - UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, lit, + Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, + WindowFunctionDefinition, col, lit, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; @@ -3277,46 +3290,6 @@ 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 54b80644a31a3..12abf79041091 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::{ - bounded_window_exec_with_can_repartition, check_integrity, coalesce_partitions_exec, - parquet_exec_with_sort, parquet_exec_with_stats, repartition_exec, schema, sort_exec, + 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,13 +59,9 @@ use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; -use datafusion_physical_expr::{ - Distribution, EquivalenceProperties, RangePartitioning, SplitPoint, -}; +use datafusion_physical_expr::Distribution; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion_physical_plan::execution_plan::{ - Boundedness, EmissionType, ExecutionPlan, -}; +use datafusion_physical_plan::execution_plan::ExecutionPlan; use datafusion_physical_plan::expressions::col; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::JoinOn; @@ -74,8 +70,7 @@ 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, Partitioning, PlanProperties, - displayable, + DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, displayable, }; use insta::Settings; @@ -268,12 +263,8 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { vec![&self.input] } - fn input_distribution_requirements( - &self, - ) -> datafusion_physical_plan::InputDistributionRequirements { - datafusion_physical_plan::InputDistributionRequirements::new(vec![ - Distribution::SinglePartition, - ]) + fn required_input_distribution(&self) -> Vec { + vec![Distribution::SinglePartition] } fn maintains_input_order(&self) -> Vec { @@ -334,104 +325,6 @@ 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![]) } @@ -808,553 +701,6 @@ 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 827a001b59894..6f88e01059fc9 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::KeyPartitioned(vec![ + Distribution::HashPartitioned(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=KeyPartitioned[[a@0, b@1]]) + OutputRequirementExec: order_by=[(b@1, asc), (c@2 + a@0, asc)], dist_by=HashPartitioned[[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=KeyPartitioned[[new_a@1, b@2]]) + OutputRequirementExec: order_by=[(b@2, asc), (c@0 + new_a@1, asc)], dist_by=HashPartitioned[[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,12 +797,10 @@ fn test_output_req_after_projection() -> Result<()> { Arc::new(Column::new("new_a", 1)), Arc::new(Column::new("b", 2)), ]; - if let Distribution::KeyPartitioned(vec) = after_optimize + if let Distribution::HashPartitioned(vec) = after_optimize .downcast_ref::() .unwrap() - .input_distribution_requirements() - .child_distribution(0) - .unwrap() + .required_input_distribution()[0] .clone() { assert!( @@ -811,7 +809,7 @@ fn test_output_req_after_projection() -> Result<()> { .all(|(actual, expected)| actual.eq(&expected)) ); } else { - panic!("Expected KeyPartitioned distribution!"); + panic!("Expected HashPartitioned distribution!"); }; Ok(()) diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index 184125dcbe180..217570846d56e 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -19,9 +19,8 @@ use insta::assert_snapshot; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - 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, + bounded_window_exec, global_limit_exec, local_limit_exec, memory_exec, + projection_exec, repartition_exec, sort_exec, sort_expr, sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; @@ -30,13 +29,12 @@ 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, NullEquality, Result, ScalarValue}; +use datafusion_common::{JoinType, Result, ScalarValue}; +use datafusion_physical_expr::Partitioning; 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}; @@ -402,146 +400,6 @@ 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<()> { @@ -600,76 +458,6 @@ 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 1973813215ed6..09225cb0385a7 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -263,22 +263,6 @@ 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(); @@ -301,7 +285,7 @@ pub fn bounded_window_exec_with_can_repartition( vec![window_expr], Arc::clone(&input), InputOrderMode::Sorted, - can_repartition, + false, ) .unwrap(), ) diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index b837373632f07..e8ff6758ccdd4 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -708,12 +708,8 @@ impl ExecutionPlan for TopKExec { &self.cache } - fn input_distribution_requirements( - &self, - ) -> datafusion_physical_plan::InputDistributionRequirements { - datafusion_physical_plan::InputDistributionRequirements::new(vec![ - Distribution::SinglePartition, - ]) + fn required_input_distribution(&self) -> Vec { + 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 5e4126793cb1f..4bf86e17d387d 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -27,7 +27,8 @@ use crate::{ file_stream::work_source::SharedWorkSource, source::DataSource, statistics::MinMaxStatistics, }; -use arrow::datatypes::{DataType, FieldRef, Schema, SchemaRef}; +use arrow::datatypes::FieldRef; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::{ Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, @@ -39,7 +40,7 @@ use datafusion_expr::Operator; use crate::source::OpenArgs; use datafusion_physical_expr::expressions::{BinaryExpr, Column}; -use datafusion_physical_expr::projection::{ProjectionExprs, ProjectionMapping}; +use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; @@ -204,17 +205,7 @@ 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. @@ -281,7 +272,6 @@ 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>, @@ -305,7 +295,6 @@ impl FileScanConfigBuilder { file_groups: vec![], statistics: None, output_ordering: vec![], - output_partitioning: None, file_compression_type: None, limit: None, preserve_order: false, @@ -472,15 +461,6 @@ 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, @@ -539,7 +519,6 @@ impl FileScanConfigBuilder { file_groups, statistics, output_ordering, - output_partitioning, file_compression_type, batch_size, expr_adapter_factory: expr_adapter, @@ -569,7 +548,6 @@ impl FileScanConfigBuilder { expr_adapter_factory: expr_adapter, statistics, partitioned_by_file_group, - output_partitioning, } } } @@ -582,7 +560,6 @@ 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, @@ -594,52 +571,6 @@ 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, @@ -727,10 +658,6 @@ 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)?; } @@ -754,9 +681,10 @@ impl DataSource for FileScanConfig { repartition_file_min_size: usize, output_ordering: Option, ) -> Result>> { - // 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 { + // 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 { return Ok(None); } @@ -772,18 +700,13 @@ impl DataSource for FileScanConfig { /// Returns the output partitioning for this file scan. /// - /// 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`. + /// 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. /// /// Tradeoffs - /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries whose - /// required distribution is satisfied by the scan's output partitioning. + /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries with + /// `GROUP BY` or `ORDER BY` on partition columns. /// - 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. @@ -792,45 +715,39 @@ 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 { - 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 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()); + } + }; - 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()) + // 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))); + } } - }; - } - output_partitioning + if exprs.len() == partition_cols.len() { + return Partitioning::Hash(exprs, self.file_groups.len()); + } + } + } + Partitioning::UnknownPartitioning(self.file_groups.len()) } /// Computes the effective equivalence properties of this file scan, taking @@ -1124,10 +1041,7 @@ 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.output_partitioning.is_some() - || self.partitioned_by_file_group - { + if self.preserve_order || self.partitioned_by_file_group { return None; } @@ -2531,58 +2445,6 @@ 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 18ebe80773e8a..e3df1ad6381f4 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -31,9 +31,8 @@ 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, - InputDistributionRequirements, Partitioning, PlanProperties, - SendableRecordBatchStream, execute_input_stream, + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, SendableRecordBatchStream, execute_input_stream, }; use async_trait::async_trait; @@ -190,16 +189,9 @@ 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. - InputDistributionRequirements::new(vec![ - Distribution::SinglePartition; - self.children().len() - ]) + 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 27b86a6d8cdd5..58c7feb616179 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -515,23 +515,6 @@ 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 5da6d97d64403..c2b01868c97f3 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, - RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, - Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, + 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 a5530af914164..c572b202f03ce 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -50,7 +50,6 @@ 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; @@ -61,12 +60,10 @@ use datafusion_common::tree_node::{ use datafusion_common::{ Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Dependency, FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues, Result, - 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, + ScalarValue, Spans, TableReference, UnnestOptions, aggregate_functional_dependencies, + assert_eq_or_internal_err, assert_or_internal_err, internal_err, plan_err, }; use indexmap::IndexSet; -use itertools::Itertools as _; // backwards compatibility use crate::display::PgJsonVisitor; @@ -867,32 +864,6 @@ 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 { @@ -2120,9 +2091,6 @@ 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(); @@ -4169,16 +4137,11 @@ impl Debug for Subquery { } } -/// Logical partitioning schemes. +/// Logical partitioning schemes supported by [`LogicalPlan::Repartition`] /// -/// 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. +/// See [`Partitioning`] for more details on partitioning /// -/// For physical execution partitioning, see -/// [`datafusion_physical_expr::Partitioning`]. -/// -/// [`datafusion_physical_expr::Partitioning`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/enum.Partitioning.html# +/// [`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 @@ -4186,118 +4149,10 @@ 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 /// @@ -4635,134 +4490,6 @@ 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 32396d2cd26f2..ef9382a57209a 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -37,7 +37,6 @@ //! * [`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, @@ -415,7 +414,6 @@ 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, .. }) => { @@ -521,19 +519,6 @@ 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 2a9a8528c6c3e..434b6a097e645 100644 --- a/datafusion/ffi/src/physical_expr/partitioning.rs +++ b/datafusion/ffi/src/physical_expr/partitioning.rs @@ -17,35 +17,20 @@ use std::sync::Arc; -use datafusion_common::{DataFusionError, ScalarValue, SplitPoint}; -use datafusion_physical_expr::{ - LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, -}; +use datafusion_physical_expr::Partitioning; 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 { @@ -60,130 +45,44 @@ 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 TryFrom for Partitioning { - type Error = DataFusionError; - - fn try_from(value: FFI_Partitioning) -> Result { - Ok(match value { +impl From<&FFI_Partitioning> for Partitioning { + fn from(value: &FFI_Partitioning) -> Self { + 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) - } - 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)?) + Self::Hash(exprs, *size) } FFI_Partitioning::UnknownPartitioning(size) => { - Self::UnknownPartitioning(size) + Self::UnknownPartitioning(*size) } - }) + } } } #[cfg(test)] mod tests { - 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 datafusion_physical_expr::Partitioning; + use datafusion_physical_expr::expressions::lit; - 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, - )?)) - } + use crate::physical_expr::partitioning::FFI_Partitioning; #[test] - fn round_trip_ffi_partitioning() -> Result<()> { + fn round_trip_ffi_partitioning() { 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.try_into()?; + let returned: Partitioning = (&ffi_partitioning).into(); if let Partitioning::UnknownPartitioning(return_size) = returned { let Partitioning::UnknownPartitioning(original_size) = partitioning @@ -195,32 +94,5 @@ 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 09ef26af32349..b286ee2d7d30c 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, Partitioning}; +use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::PlanProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -172,7 +172,6 @@ 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)) @@ -188,7 +187,7 @@ impl TryFrom for PlanProperties { Ok(PlanProperties::new( eq_properties, - partitioning, + (&partitioning).into(), emission_type, boundedness, )) @@ -261,15 +260,13 @@ 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)])); @@ -285,25 +282,6 @@ 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()?; @@ -336,22 +314,4 @@ 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 ec7bf648e22ea..848bf81d15979 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -55,18 +55,14 @@ 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, PartitioningSatisfaction, RangePartitioning, -}; +pub use partitioning::{Distribution, Partitioning}; pub use physical_expr::{ add_offset_to_expr, add_offset_to_physical_sort_exprs, create_lex_ordering, - create_ordering, create_physical_partitioning, create_physical_sort_expr, - create_physical_sort_exprs, physical_exprs_bag_equal, physical_exprs_contains, - physical_exprs_equal, + create_ordering, 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 14065ab8853ad..d24c60b63e6bd 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -21,10 +21,7 @@ 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; @@ -120,8 +117,6 @@ 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), } @@ -138,7 +133,6 @@ 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})") } @@ -146,191 +140,6 @@ 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 { @@ -358,7 +167,6 @@ impl Partitioning { use Partitioning::*; match self { RoundRobinBatch(n) | Hash(_, n) | UnknownPartitioning(n) => *n, - Range(range) => range.partition_count(), } } @@ -393,10 +201,6 @@ 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, @@ -408,111 +212,84 @@ impl Partitioning { Distribution::SinglePartition if self.partition_count() == 1 => { PartitioningSatisfaction::Exact } - // When partition count is 1, partitioned requirements are satisfied. - Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) - if self.partition_count() == 1 => - { + // When partition count is 1, hash requirement is satisfied. + Distribution::HashPartitioned(_) if self.partition_count() == 1 => { PartitioningSatisfaction::Exact } - Distribution::HashPartitioned(required_exprs) - | Distribution::KeyPartitioned(required_exprs) => match self { + Distribution::HashPartitioned(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, _) => 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(_) => { + 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; + } + PartitioningSatisfaction::NotSatisfied } + _ => PartitioningSatisfaction::NotSatisfied, }, - Distribution::SinglePartition => PartitioningSatisfaction::NotSatisfied, + _ => 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 { - 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())) - }) + 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())) }) - .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() - } + }) + .collect(); + Partitioning::Hash(normalized_exprs, *part) + } else { + self.clone() } } } @@ -529,7 +306,6 @@ impl PartialEq for Partitioning { { true } - (Partitioning::Range(left), Partitioning::Range(right)) => left == right, _ => false, } } @@ -543,39 +319,12 @@ pub enum Distribution { UnspecifiedDistribution, /// A single partition is required SinglePartition, - /// Deprecated historical name for [`Distribution::KeyPartitioned`]. - /// See for details. - #[deprecated(since = "55.0.0", note = "Use Distribution::KeyPartitioned")] + /// Requires children to be distributed in such a way that the same + /// values of the keys end up in the same partition 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 { @@ -583,17 +332,13 @@ impl Distribution { Partitioning::UnknownPartitioning(partition_count) } Distribution::SinglePartition => Partitioning::UnknownPartitioning(1), - Distribution::HashPartitioned(expr) | Distribution::KeyPartitioned(expr) => { + Distribution::HashPartitioned(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 { @@ -602,9 +347,6 @@ 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)) - } } } } @@ -614,188 +356,56 @@ 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()) - } - 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" - ); - } + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::Result; #[test] - #[expect( - deprecated, - reason = "test intentionally covers deprecated HashPartitioned compatibility" - )] fn partitioning_satisfy_distribution() -> Result<()> { - let fixture = PartitioningTestFixture::new(vec![ - ("column_1", DataType::Int64), - ("column_2", DataType::Utf8), - ])?; + 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 distribution_types = vec![ Distribution::UnspecifiedDistribution, Distribution::SinglePartition, - Distribution::HashPartitioned(fixture.cols([0, 1])), - Distribution::KeyPartitioned(fixture.cols([0, 1])), + Distribution::HashPartitioned(partition_exprs1.clone()), ]; let single_partition = Partitioning::UnknownPartitioning(1); let unspecified_partition = Partitioning::UnknownPartitioning(10); let round_robin_partition = Partitioning::RoundRobinBatch(10); - let hash_partition1 = fixture.hash_partitioning([0, 1], 10); - let hash_partition2 = fixture.hash_partitioning([1, 0], 10); + let hash_partition1 = Partitioning::Hash(partition_exprs1, 10); + let hash_partition2 = Partitioning::Hash(partition_exprs2, 10); + let eq_properties = EquivalenceProperties::new(schema); for distribution in distribution_types { let result = ( single_partition - .satisfaction(&distribution, &fixture.eq_properties, true) + .satisfaction(&distribution, &eq_properties, true) .is_satisfied(), unspecified_partition - .satisfaction(&distribution, &fixture.eq_properties, true) + .satisfaction(&distribution, &eq_properties, true) .is_satisfied(), round_robin_partition - .satisfaction(&distribution, &fixture.eq_properties, true) + .satisfaction(&distribution, &eq_properties, true) .is_satisfied(), hash_partition1 - .satisfaction(&distribution, &fixture.eq_properties, true) + .satisfaction(&distribution, &eq_properties, true) .is_satisfied(), hash_partition2 - .satisfaction(&distribution, &fixture.eq_properties, true) + .satisfaction(&distribution, &eq_properties, true) .is_satisfied(), ); @@ -806,7 +416,7 @@ mod tests { Distribution::SinglePartition => { assert_eq!(result, (true, false, false, false, false)) } - Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) => { + Distribution::HashPartitioned(_) => { assert_eq!(result, (true, false, false, true, false)) } } @@ -816,129 +426,239 @@ mod tests { } #[test] - #[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")); + 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)); let test_cases = vec![ ( - "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]), + "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), + ]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "subset: KeyPartitioned([a, b, c]) satisfied by Hash([b])", - fixture.hash_partitioning([1], 4), - fixture.key_distribution([0, 1, 2]), + "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), + ]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "subset reordered: KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", - fixture.hash_partitioning([1, 0], 4), - fixture.key_distribution([0, 1, 2]), + "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), + ]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "superset: KeyPartitioned([a]) not satisfied by Hash([a, b])", - fixture.hash_partitioning([0, 1], 4), - fixture.key_distribution([0]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ( - "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, - ), - ( - "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, - ), - ( - "no overlap: KeyPartitioned([b, c]) not satisfied by Hash([a])", - fixture.hash_partitioning([0], 4), - fixture.key_distribution([1, 2]), - 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, PartitioningSatisfaction::NotSatisfied, ), ( - "unknown partition expr", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - fixture.key_distribution([0, 1]), - 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, 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![ ( - "unknown required expr", - fixture.hash_partitioning([0, 1], 4), - Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), + "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)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "same unknown expr", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), + "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)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "unknown partition expr is not a valid subset", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::KeyPartitioned(vec![Arc::clone(&unknown), fixture.col(0)]), + "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), + ]), 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![ ( - "empty hash partitioning", - Partitioning::Hash(vec![], 4), - fixture.key_distribution([0]), + "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), + ]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "empty key distribution", - fixture.hash_partitioning([0], 4), - Distribution::KeyPartitioned(vec![]), + "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)]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -947,13 +667,16 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - assert_satisfaction( - desc, - &partition, - &required, - &fixture.eq_properties, - expected_with_subset, - expected_without_subset, + 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" ); } @@ -961,302 +684,164 @@ mod tests { } #[test] - 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]) vs KeyPartitioned([a])", - &hash_a, - fixture.key_partitioned_distribution([0]), - ), - ( - "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}" - ); - } - } + 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 subset_cases = [ + let test_cases = vec![ ( - "Hash([a]) vs KeyPartitioned([a, b])", - &hash_a, - fixture.key_partitioned_distribution([0, 1]), + "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, ), ( - "Range([a]) vs KeyPartitioned([a, b])", - &range_a, - fixture.key_partitioned_distribution([0, 1]), + "Hash([a]) vs Hash([a])", + Partitioning::Hash(vec![Arc::clone(&col_a)], 4), + Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, ), ]; - for (desc, partitioning, requirement) in subset_cases { + + for (desc, partition, required, expected_with_subset, expected_without_subset) in + test_cases + { + let result = partition.satisfaction(&required, &eq_properties, true); assert_eq!( - partitioning.satisfaction(&requirement, &fixture.eq_properties, true), - PartitioningSatisfaction::Subset, + result, expected_with_subset, "Failed for {desc} with subset enabled" ); + + let result = partition.satisfaction(&required, &eq_properties, false); assert_eq!( - partitioning.satisfaction(&requirement, &fixture.eq_properties, false), - PartitioningSatisfaction::NotSatisfied, + result, expected_without_subset, "Failed for {desc} with subset disabled" ); } - let not_satisfied_cases = [ + 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![ ( - "Range([a]) vs KeyPartitioned([b])", - &range_a, - fixture.key_partitioned_distribution([1]), + "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, ), ( - "Hash([a, b]) vs KeyPartitioned([a])", - &hash_ab, - fixture.key_partitioned_distribution([0]), + "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, ), ( - "Range([a, b]) vs KeyPartitioned([a])", - &range_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, ), ]; - 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}" - ); - } - } - - Ok(()) - } - 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(()) - } + 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" + ); - #[test] - 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 result = partition.satisfaction(&required, &eq_properties, false); + assert_eq!( + result, expected_without_subset, + "Failed for {desc} with subset disabled" + ); + } Ok(()) } #[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); + fn test_partitioning_empty_hash() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - Ok(()) - } + let col_a: Arc = + Arc::new(Column::new_with_schema("a", &schema)?); + let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - #[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([ + let test_cases = vec![ ( - fixture.col(0), - ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), + "Hash([]) vs Hash([a])", + Partitioning::Hash(vec![], 4), + Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, ), ( - fixture.col(1), - ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), + "Hash([a]) vs Hash([])", + Partitioning::Hash(vec![Arc::clone(&col_a)], 4), + Distribution::HashPartitioned(vec![]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, ), - ]); - - 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); + ( + "Hash([]) vs Hash([])", + Partitioning::Hash(vec![], 4), + Distribution::HashPartitioned(vec![]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ]; - Ok(()) - } + 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" + ); - #[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, - ); + let result = partition.satisfaction(&required, &eq_properties, false); + assert_eq!( + result, expected_without_subset, + "Failed for {desc} with subset disabled" + ); + } Ok(()) } diff --git a/datafusion/physical-expr/src/physical_expr.rs b/datafusion/physical-expr/src/physical_expr.rs index 6ff5be4e38229..77ede76e1daa8 100644 --- a/datafusion/physical-expr/src/physical_expr.rs +++ b/datafusion/physical-expr/src/physical_expr.rs @@ -21,17 +21,15 @@ use crate::expressions::{self, Column}; use crate::{LexOrdering, PhysicalSortExpr, create_physical_expr}; use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Schema, SchemaRef}; +use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{DFSchema, HashMap, ScalarValue, SplitPoint}; +use datafusion_common::{DFSchema, HashMap}; use datafusion_common::{Result, plan_err}; use datafusion_expr::execution_props::ExecutionProps; -use datafusion_expr::{Expr, Partitioning as LogicalPartitioning, SortExpr}; -use datafusion_expr_common::casts::try_cast_literal_to_type; +use datafusion_expr::{Expr, SortExpr}; 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 @@ -218,99 +216,6 @@ 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 3b520fd3de60c..c522867c05196 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/enforce_distribution.rs @@ -41,8 +41,7 @@ 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, OrderingRequirements, PhysicalExpr, PhysicalExprRef, - physical_exprs_equal, + EquivalenceProperties, PhysicalExpr, PhysicalExprRef, physical_exprs_equal, }; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::aggregates::{ @@ -60,10 +59,7 @@ 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::{ - ChildSatisfactionOptions, Distribution, ExecutionPlan, InputDistributionRequirements, - Partitioning, -}; +use datafusion_physical_plan::{Distribution, ExecutionPlan, Partitioning}; use itertools::izip; @@ -857,34 +853,72 @@ fn add_roundrobin_on_top( } } -// 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, +/// 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, allow_subset_satisfy_partitioning: bool, -) -> bool { - let Some(aggregate) = plan.downcast_ref::() else { - return false; - }; - if aggregate.mode() != &AggregateMode::Partial - || aggregate.group_expr().is_empty() - || aggregate.group_expr().has_grouping_set() - { - return false; +) -> 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 key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr()); + 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() + }; - plan.output_partitioning() - .satisfaction( - &key_distribution, - plan.equivalence_properties(), - allow_subset_satisfy_partitioning, - ) - .is_satisfied() + 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])); + } + + Ok(input) } /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator @@ -1060,27 +1094,6 @@ 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 @@ -1120,7 +1133,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.input_distribution_requirements().into_per_child(); + let requirements = plan.required_input_distribution(); let mut repartition_status_flags = vec![]; for (child, requirement, roundrobin_beneficial) in izip!(children.into_iter(), requirements, rr_beneficial) @@ -1133,31 +1146,30 @@ fn get_repartition_requirement_status( Precision::Inexact(n_rows) => !should_use_estimates || (n_rows > batch_size), Precision::Absent => true, }; - let is_partitioned_requirement = requirement.key_exprs().is_some(); - // Hash repartitioning is necessary when the input has more than one - // partition. + let is_hash = matches!(requirement, Distribution::HashPartitioned(_)); + // Hash re-partitioning is necessary when the input has more than one + // partitions: let multi_partitions = child.output_partitioning().partition_count() > 1; let roundrobin_sensible = roundrobin_beneficial && roundrobin_beneficial_stats; - needs_alignment |= - is_partitioned_requirement && (multi_partitions || roundrobin_sensible); + needs_alignment |= is_hash && (multi_partitions || roundrobin_sensible); repartition_status_flags.push(( - is_partitioned_requirement, + is_hash, RepartitionRequirementStatus { requirement, roundrobin_beneficial, roundrobin_beneficial_stats, - hash_necessary: is_partitioned_requirement && multi_partitions, + hash_necessary: is_hash && multi_partitions, }, )); } - // Align hash necessary flags for key partitions to generate consistent + // Align hash necessary flags for hash partitions to generate consistent // hash partitions at each children: if needs_alignment { - // 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 { + // 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 { status.hash_necessary = true; } } @@ -1168,90 +1180,6 @@ 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 @@ -1260,10 +1188,6 @@ fn enforce_distribution_relationships( /// 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, @@ -1354,7 +1278,6 @@ 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: @@ -1362,8 +1285,7 @@ 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 mut children = izip!( - 0..children.len(), + let children = izip!( children.into_iter(), plan.required_input_ordering(), plan.maintains_input_order(), @@ -1371,7 +1293,6 @@ pub fn ensure_distribution( ) .map( |( - child_idx, mut child, required_input_ordering, maintains, @@ -1382,18 +1303,33 @@ 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(); - let preserve_file_partition_threshold_met = - config.optimizer.preserve_file_partitions > 0 - && current_partitions >= config.optimizer.preserve_file_partitions; + + // 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 allow_subset_satisfy_partitioning = (current_partitions >= subset_satisfaction_threshold @@ -1401,27 +1337,10 @@ 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. - || (preserve_file_partition_threshold_met + || (config.optimizer.preserve_file_partitions > 0 && current_partitions < target_partitions)) && !is_partitioned_join - && !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; + && !requires_grouping_id; // When `repartition_file_scans` is set, attempt to increase // parallelism at the source. @@ -1442,49 +1361,16 @@ pub fn ensure_distribution( Distribution::SinglePartition => { child = add_merge_on_top(child); } - 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; - + Distribution::HashPartitioned(exprs) => { // 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 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]); + if hash_necessary { + child = add_hash_on_top( + child, + exprs.to_vec(), + target_partitions, + allow_subset_satisfy_partitioning, + )?; } } Distribution::UnspecifiedDistribution => { @@ -1496,101 +1382,73 @@ pub fn ensure_distribution( } }; - Ok(DistributionChildState { - context: child, - required_input_ordering, - maintains_input_order: maintains, - requirement, - }) - }, - ) - .collect::>>()?; + let streaming_benefit = if child.data { + preserving_order_enables_streaming(&plan, &child.plan)? + } else { + false + }; - // 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)? + // 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)? } else { false }; - - // 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), - )?; + // 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)?; } } - // 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)?; - } + Distribution::UnspecifiedDistribution => { + // Since ordering is lost, trying to preserve ordering is pointless + if !maintains || plan.is::() { + child = replace_order_preserving_variants(child)?; } } } - Ok(context) - }, - ) - .collect::>>()?; + } + Ok(child) + }, + ) + .collect::>>()?; let children_plans = children .iter() @@ -1656,8 +1514,7 @@ fn update_children(mut dist_context: DistributionContext) -> Result Result { if node.data { let requires_single_partition = matches!( - parent - .input_distribution_requirements() - .child_distribution(child_idx), - Some(Distribution::SinglePartition) + parent.required_input_distribution()[child_idx], + Distribution::SinglePartition ); node = remove_corresponding_sort_from_sub_plan(node, requires_single_partition)?; } @@ -757,7 +749,7 @@ fn remove_corresponding_sort_from_sub_plan( } } else { let mut any_connection = false; - let required_dist = node.plan.input_distribution_requirements().into_per_child(); + let required_dist = node.plan.required_input_distribution(); 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 48cc41b8a21c4..24eb3af5f564c 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, internal_err}; +use datafusion_common::{Result, Statistics}; use datafusion_execution::TaskContext; use datafusion_physical_expr::Distribution; use datafusion_physical_expr_common::sort_expr::OrderingRequirements; @@ -205,15 +205,7 @@ impl ExecutionPlan for OutputRequirementExec { } fn required_input_distribution(&self) -> Vec { - 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(), - ]) + vec![self.dist_requirement.clone()] } fn maintains_input_order(&self) -> Vec { @@ -252,10 +244,6 @@ 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, @@ -280,12 +268,8 @@ impl ExecutionPlan for OutputRequirementExec { requirements = OrderingRequirements::new_alternatives(updated_reqs, soft); } - let input_distributions = self.input_distribution_requirements(); - let dist_req = match input_distributions.child_distribution(0) { - Some( - Distribution::HashPartitioned(exprs) - | Distribution::KeyPartitioned(exprs), - ) => { + let dist_req = match &self.required_input_distribution()[0] { + Distribution::HashPartitioned(exprs) => { let mut updated_exprs = vec![]; for expr in exprs { let Some(new_expr) = update_expr(expr, projection.expr(), false)? @@ -294,14 +278,9 @@ impl ExecutionPlan for OutputRequirementExec { }; updated_exprs.push(new_expr); } - Distribution::KeyPartitioned(updated_exprs) - } - Some(dist) => dist.clone(), - None => { - return internal_err!( - "OutputRequirementExec missing input distribution requirement" - ); + Distribution::HashPartitioned(updated_exprs) } + dist => dist.clone(), }; make_with_child(projection, &self.input()).map(|input| { @@ -376,10 +355,7 @@ 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 - .input_distribution_requirements() - .into_per_child() - .swap_remove(0); + let req_dist = sort_exec.required_input_distribution().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 713213b70612d..40c6245d894d4 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -30,13 +30,9 @@ 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, InvariantLevel, -}; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::joins::SymmetricHashJoinExec; -use datafusion_physical_plan::{ - ChildSatisfactionOptions, ExecutionPlanProperties, get_plan_string, -}; +use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string}; use crate::PhysicalOptimizerRule; use datafusion_physical_expr_common::sort_expr::format_physical_sort_requirement_list; @@ -145,12 +141,11 @@ 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(), - input_distributions.per_child_distributions(), + plan.required_input_distribution(), ) { let child_eq_props = child.equivalence_properties(); if let Some(sort_req) = sort_req { @@ -167,12 +162,9 @@ pub fn check_plan_sanity( } } - if !input_distributions - .child_satisfaction( - idx, - child.as_ref(), - ChildSatisfactionOptions::new().with_allow_subset(true), - )? + if !child + .output_partitioning() + .satisfaction(&dist_req, child_eq_props, true) .is_satisfied() { let plan_str = get_plan_string(plan); @@ -186,8 +178,6 @@ 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 9459376054fda..40dbddfbdf9fb 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -46,7 +46,6 @@ 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}; @@ -120,19 +119,17 @@ 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 and RepartitionExec to find BoundedWindowAggExec + // Step 3: Walk through optional ProjectionExec to find BoundedWindowAggExec let child = filter.input(); - let (window_exec, intermediates) = find_window_below(child)?; + let (window_exec, proj_between) = find_window_below(child)?; // Step 4: Verify col_idx references a ROW_NUMBER window output column - 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(); + let input_field_count = window_exec.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_typed.window_expr(); + let window_exprs = window_exec.window_expr(); if window_expr_idx >= window_exprs.len() { return None; } @@ -140,7 +137,8 @@ impl WindowTopN { return None; } - // Step 5: child of window is SortExec (verified above) + // Step 5: Verify child of window is SortExec + let sort_exec = window_exec.input().downcast_ref::()?; let sort_child = sort_exec.input(); // Step 6: Determine partition_prefix_len from the window expression @@ -163,19 +161,28 @@ impl WindowTopN { .ok()?; // Step 8: Rebuild window with new child - let mut result = window_exec + let new_window = Arc::clone(&child_as_arc(window_exec)) .with_new_children(vec![Arc::new(partitioned_topk)]) .ok()?; - // Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec) - for node in intermediates.into_iter().rev() { - result = node.with_new_children(vec![result]).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, + }; 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, @@ -296,32 +303,29 @@ 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 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. +/// Handles two cases: +/// - Direct child: `FilterExec → BoundedWindowAggExec` +/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec` /// -/// 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(); +/// 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)); + } - 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; + // 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))); } } + + None } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 61e88f0f75db3..5e6b8505764a2 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, InputDistributionRequirements, - InputOrderMode, SendableRecordBatchStream, Statistics, check_if_same_properties, + DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode, + SendableRecordBatchStream, Statistics, check_if_same_properties, }; use datafusion_common::config::ConfigOptions; use datafusion_physical_expr::utils::collect_columns; @@ -1544,21 +1544,17 @@ impl ExecutionPlan for AggregateExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> InputDistributionRequirements { - InputDistributionRequirements::new(match &self.mode { + match &self.mode { AggregateMode::Partial | AggregateMode::PartialReduce => { vec![Distribution::UnspecifiedDistribution] } AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned => { - vec![Distribution::KeyPartitioned(self.group_by.input_exprs())] + vec![Distribution::HashPartitioned(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 708f6b4c95ac3..582af8f1e3dae 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -144,13 +144,7 @@ impl ExecutionPlan for AnalyzeExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![ - Distribution::UnspecifiedDistribution, - ]) + 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 deleted file mode 100644 index 6405b1f121ef7..0000000000000 --- a/datafusion/physical-plan/src/distribution_requirements.rs +++ /dev/null @@ -1,359 +0,0 @@ -// 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 5c20f85656e90..50eac566d90ef 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -16,7 +16,6 @@ // under the License. pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; -use crate::distribution_requirements::InputDistributionRequirements; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -162,30 +161,12 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { check_default_invariants(self, check) } - /// 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")] + /// Specifies the data distribution requirements for all the + /// children for this `ExecutionPlan`, By default it's [[Distribution::UnspecifiedDistribution]] for each child, 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`. /// @@ -232,8 +213,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.input_distribution_requirements() - .per_child_distributions() + self.required_input_distribution() + .into_iter() .map(|dist| !matches!(dist, Distribution::SinglePartition)) .collect() } @@ -1216,15 +1197,14 @@ 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 fde1718001252..6661d2782b212 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -306,14 +306,10 @@ impl ExecutionPlan for CrossJoinExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![ + 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 0ee918bb66e70..f7391feb29cc3 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -54,9 +54,8 @@ use crate::projection::{ use crate::repartition::REPARTITION_RANDOM_STATE; use crate::spill::get_record_batch_memory_size; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - InputDistributionRequirements, Partitioning, PlanProperties, - SendableRecordBatchStream, Statistics, + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + PlanProperties, SendableRecordBatchStream, Statistics, common::can_project, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, @@ -846,65 +845,21 @@ impl HashJoinExec { return false; } - // 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() + // `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 { - // 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 @@ -1248,30 +1203,26 @@ 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(); - InputDistributionRequirements::co_partitioned(vec![ - Distribution::KeyPartitioned(left_expr), - Distribution::KeyPartitioned(right_expr), - ]) + vec![ + Distribution::HashPartitioned(left_expr), + Distribution::HashPartitioned(right_expr), + ] } - PartitionMode::CollectLeft => InputDistributionRequirements::new(vec![ - Distribution::SinglePartition, - Distribution::UnspecifiedDistribution, - ]), - PartitionMode::Auto => InputDistributionRequirements::new(vec![ + PartitionMode::Auto => vec![ Distribution::UnspecifiedDistribution, Distribution::UnspecifiedDistribution, - ]), + ], } } @@ -2168,90 +2119,7 @@ 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::{ @@ -2274,67 +2142,11 @@ 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) } @@ -5563,7 +5375,6 @@ 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]), @@ -5575,7 +5386,6 @@ 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 _, @@ -5603,8 +5413,8 @@ mod tests { let task_ctx = Arc::new(task_ctx); let join = HashJoinExec::try_new( - Arc::clone(&left), - Arc::clone(&right), + Arc::clone(&left) as Arc, + Arc::clone(&right) as Arc, on.clone(), None, &join_type, @@ -5907,8 +5717,6 @@ 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( @@ -5927,7 +5735,6 @@ 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 _, @@ -6534,193 +6341,6 @@ 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 20d05983cf4ed..15af23b447836 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -565,14 +565,10 @@ impl ExecutionPlan for NestedLoopJoinExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![ + 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 a2a1d6c787bb7..50e9252a21131 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -508,14 +508,10 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![ + 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 6a7064bf0417b..9e87b52696a57 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -40,8 +40,7 @@ use crate::projection::{ use crate::spill::spill_manager::SpillManager; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - InputDistributionRequirements, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties, }; use arrow::compute::SortOptions; @@ -424,19 +423,15 @@ 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(); - InputDistributionRequirements::co_partitioned(vec![ - Distribution::KeyPartitioned(left_expr), - Distribution::KeyPartitioned(right_expr), - ]) + vec![ + Distribution::HashPartitioned(left_expr), + Distribution::HashPartitioned(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 b798a166e9cc2..ef92964fadf84 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -53,8 +53,7 @@ use crate::projection::{ use crate::stream::EmptyRecordBatchStream; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - InputDistributionRequirements, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, joins::StreamJoinPartitionMode, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; @@ -427,10 +426,6 @@ 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 @@ -438,16 +433,13 @@ impl ExecutionPlan for SymmetricHashJoinExec { .iter() .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) .unzip(); - InputDistributionRequirements::co_partitioned(vec![ - Distribution::KeyPartitioned(left_expr), - Distribution::KeyPartitioned(right_expr), - ]) + vec![ + Distribution::HashPartitioned(left_expr), + Distribution::HashPartitioned(right_expr), + ] } StreamJoinPartitionMode::SinglePartition => { - InputDistributionRequirements::new(vec![ - Distribution::SinglePartition, - Distribution::SinglePartition, - ]) + vec![Distribution::SinglePartition, Distribution::SinglePartition] } } } diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index 134d650a9cd6b..b4aa295562b67 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -33,8 +33,7 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, ProjectionExpr}; use crate::{ - ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, - RangePartitioning, Statistics, + ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, Statistics, }; // compatibility pub use super::join_filter::JoinFilter; @@ -68,7 +67,7 @@ use datafusion_common::hash_utils::create_hashes; use datafusion_common::stats::Precision; use datafusion_common::{ DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, - internal_datafusion_err, not_impl_err, plan_err, + not_impl_err, plan_err, }; use datafusion_expr::Operator; use datafusion_expr::interval_arithmetic::Interval; @@ -145,21 +144,6 @@ 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) @@ -2160,7 +2144,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, SplitPoint, arrow_datafusion_err, arrow_err}; + use datafusion_common::{ScalarValue, arrow_datafusion_err, arrow_err}; use datafusion_physical_expr::PhysicalSortExpr; use rstest::rstest; @@ -3255,53 +3239,6 @@ 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 8190663b2b23d..3005e975424b4 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -37,13 +37,10 @@ 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, RangePartitioning, SplitPoint, expressions, + Distribution, Partitioning, PhysicalExpr, 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, @@ -74,7 +71,6 @@ 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 94ef00d8567aa..7f42c33a79ca0 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -162,11 +162,7 @@ impl ExecutionPlan for GlobalLimitExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition]) + 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 99767df73a734..f34aac3744557 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -163,14 +163,10 @@ impl ExecutionPlan for RecursiveQueryExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![ + 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 a20714d1e8647..465ca4a99e961 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -19,11 +19,10 @@ //! 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 as AtomicOrdering}; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; use std::vec; @@ -46,22 +45,20 @@ use crate::{ check_if_same_properties, }; -use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions}; +use arrow::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::{compare_rows, extract_row_at_idx_to_buf, transpose}; +use datafusion_common::utils::transpose; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, - assert_or_internal_err, internal_err, + ColumnStatistics, DataFusionError, HashMap, 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, RangePartitioning}; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use crate::filter_pushdown::{ @@ -258,7 +255,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, AtomicOrdering::AcqRel) == 1; + let was_last = self.active_senders.fetch_sub(1, Ordering::AcqRel) == 1; if !was_last { return Ok(vec![]); } @@ -572,18 +569,6 @@ 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 @@ -720,40 +705,13 @@ 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, round-robin, or range constructors depending on the partitioning variant. + /// hash or round-robin constructors depending on the partitioning variant. /// /// # Parameters - /// - `partitioning`: Partitioning scheme to apply (hash, round-robin, or range). + /// - `partitioning`: Partitioning scheme to apply (hash or round-robin). /// - `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. @@ -779,9 +737,6 @@ 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:?}") } @@ -868,95 +823,22 @@ 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, .. } - | BatchPartitionerState::Range { indices, .. } => indices.len(), + BatchPartitionerState::Hash { indices, .. } => indices.len(), } } - /// Build repartitioned hash/range output batches using one `take` per input batch. + /// Build repartitioned hash output batches using one `take` per input batch. /// - /// The routers first fills one index vector per output partition. This method + /// The hash router 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. /// @@ -1548,30 +1430,6 @@ 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(), }; @@ -1630,10 +1488,6 @@ 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 { @@ -1747,12 +1601,26 @@ impl RepartitionExec { input_partition: usize, num_input_partitions: usize, ) -> Result<()> { - let mut partitioner = BatchPartitioner::try_new( - partitioning, - metrics.repartition_time.clone(), - input_partition, - num_input_partitions, - )?; + 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:?}"); + } + }; // While there are still outputs to send to, keep pulling inputs let mut batches_until_yield = partitioner.num_partitions(); @@ -2086,8 +1954,6 @@ mod tests { use std::collections::HashSet; use super::*; - use crate::empty::EmptyExec; - use crate::projection::ProjectionExpr; use crate::test::TestMemoryExec; use crate::{ test::{ @@ -2102,14 +1968,12 @@ mod tests { use arrow::array::{ArrayRef, StringArray, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::ScalarValue; - use datafusion_common::cast::{as_string_array, as_uint32_array}; + use datafusion_common::cast::as_string_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] @@ -2207,7 +2071,7 @@ mod tests { #[tokio::test] async fn one_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(false); + let schema = test_schema(); let partition = create_vec_batches(50); let partitions = vec![partition]; @@ -2230,7 +2094,7 @@ mod tests { #[tokio::test] async fn many_to_one_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(false); + let schema = test_schema(); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2247,7 +2111,7 @@ mod tests { #[tokio::test] async fn many_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(false); + let schema = test_schema(); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2268,7 +2132,7 @@ mod tests { #[tokio::test] async fn many_to_many_hash_partition() -> Result<()> { // define input partitions - let schema = test_schema(false); + let schema = test_schema(); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2290,542 +2154,9 @@ 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(false); + let schema = test_schema(); // create 50 batches, each having 8 rows let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone()]; @@ -2849,76 +2180,8 @@ mod tests { Ok(()) } - 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() + fn test_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) } async fn repartition( @@ -2951,7 +2214,7 @@ mod tests { let handle: SpawnedTask>>> = SpawnedTask::spawn(async move { // define input partitions - let schema = test_schema(false); + let schema = test_schema(); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -3311,7 +2574,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(false); + let schema = test_schema(); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -3373,7 +2636,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(false); + let schema = test_schema(); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -3443,7 +2706,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(false); + let schema = test_schema(); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -3505,7 +2768,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(false); + let schema = test_schema(); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -3548,7 +2811,7 @@ mod tests { /// Create batch fn create_batch() -> RecordBatch { - let schema = test_schema(false); + let schema = test_schema(); RecordBatch::try_new( schema, vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))], @@ -3558,7 +2821,7 @@ mod tests { /// Create batches with sequential values for ordering tests fn create_ordered_batches(num_batches: usize) -> Vec { - let schema = test_schema(false); + let schema = test_schema(); (0..num_batches) .map(|i| { let start = (i * 8) as u32; @@ -3579,7 +2842,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(false); + let schema = test_schema(); // 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); @@ -3951,33 +3214,6 @@ 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 3eeefa3acd7c7..3bf16af36c62b 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -268,15 +268,11 @@ impl ExecutionPlan for PartialSortExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(if self.preserve_partitioning { + 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 17eb70ef12131..fe876eeddf7f2 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -302,18 +302,12 @@ 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(); - crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( - partition_exprs, - )]) + vec![Distribution::HashPartitioned(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 551f7e0ffa6dd..f715de0b5964b 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1137,18 +1137,13 @@ impl ExecutionPlan for SortExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(if self.preserve_partitioning { + if self.preserve_partitioning { vec![Distribution::UnspecifiedDistribution] } else { // global sort - // TODO support range partitioning and OrderedDistribution. - // See https://github.com/apache/datafusion/issues/22395 + // TODO support RangePartition and OrderedDistribution 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 053c57a5a6a62..09570f14ba734 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -266,13 +266,7 @@ impl ExecutionPlan for SortPreservingMergeExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![ - Distribution::UnspecifiedDistribution, - ]) + vec![Distribution::UnspecifiedDistribution] } fn benefits_from_input_partitioning(&self) -> Vec { @@ -1492,7 +1486,11 @@ 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 = properties.output_partitioning().partition_count(); + let &partition_count = match properties.output_partitioning() { + Partitioning::RoundRobinBatch(partitions) => partitions, + Partitioning::Hash(_, partitions) => partitions, + Partitioning::UnknownPartitioning(partitions) => partitions, + }; 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 61a9b9cc6d0de..cdf4b08f718c6 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -35,7 +35,6 @@ 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; @@ -101,7 +100,7 @@ impl StreamingTableExec { let cache = Self::compute_properties( Arc::clone(&projected_schema), projected_output_ordering.clone(), - Partitioning::UnknownPartitioning(partitions.len()), + &partitions, infinite, ); Ok(Self { @@ -116,25 +115,6 @@ 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 } @@ -167,12 +147,14 @@ impl StreamingTableExec { fn compute_properties( schema: SchemaRef, orderings: Vec, - output_partitioning: Partitioning, + partitions: &[Arc], 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, @@ -222,16 +204,6 @@ 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)?; @@ -334,17 +306,6 @@ 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()), @@ -354,7 +315,6 @@ 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 08c47ab3aee93..3ea2eb5402fe5 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -462,10 +462,7 @@ impl ExecutionPlan for UnionExec { /// Combines multiple input streams by interleaving them. /// -/// 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]` +/// This only works if all inputs have the same hash-partitioning. /// /// # Data Flow /// ```text @@ -510,7 +507,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 or range partitioning" + "Not all InterleaveExec children have a consistent hash partitioning" ); let cache = Self::compute_properties(&inputs)?; Ok(InterleaveExec { @@ -658,12 +655,8 @@ impl ExecutionPlan for InterleaveExec { } } -/// 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. +/// If all the input partitions have the same Hash partition spec with the first_input_partition +/// The InterleaveExec is partition aware. /// /// 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 @@ -676,7 +669,7 @@ pub fn can_interleave>>( }; let reference = first.borrow().output_partitioning(); - matches!(reference, Partitioning::Hash(_, _) | Partitioning::Range(_)) + matches!(reference, Partitioning::Hash(_, _)) && inputs .map(|plan| plan.borrow().output_partitioning().clone()) .all(|partition| partition == *reference) @@ -840,13 +833,10 @@ 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 { @@ -1283,124 +1273,6 @@ 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 4f345a6c2dc54..c31d0dd23fa68 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -253,13 +253,7 @@ impl ExecutionPlan for UnnestExec { } fn required_input_distribution(&self) -> Vec { - self.input_distribution_requirements().into_per_child() - } - - fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { - crate::InputDistributionRequirements::new(vec![ - Distribution::UnspecifiedDistribution, - ]) + 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 1acda21196008..6c6b26c9cf49f 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -35,9 +35,8 @@ use crate::windows::{ }; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputDistributionRequirements, InputOrderMode, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, - check_if_same_properties, + ExecutionPlanProperties, InputOrderMode, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, }; use arrow::compute::take_record_batch; @@ -332,17 +331,11 @@ 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!!!"); - InputDistributionRequirements::new(vec![Distribution::SinglePartition]) + vec![Distribution::SinglePartition] } else { - InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( - self.partition_keys(), - )]) + vec![Distribution::HashPartitioned(self.partition_keys().clone())] } } diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index e909dd2017827..ee3b071fc9167 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -31,9 +31,8 @@ use crate::windows::{ }; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, PlanProperties, - RecordBatchStream, SendableRecordBatchStream, Statistics, WindowExpr, - check_if_same_properties, + ExecutionPlanProperties, PhysicalExpr, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, }; use arrow::array::ArrayRef; @@ -241,16 +240,10 @@ 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() { - InputDistributionRequirements::new(vec![Distribution::SinglePartition]) + vec![Distribution::SinglePartition] } else { - InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( - self.partition_keys(), - )]) + vec![Distribution::HashPartitioned(self.partition_keys())] } } diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index bfa40e61a8983..15a744f5c3500 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -148,19 +148,9 @@ 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; @@ -1170,8 +1160,6 @@ 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 { @@ -1432,22 +1420,13 @@ 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; - // Legacy direct partitioning fields: + // oneof partition_method { // uint64 round_robin = 2; // PhysicalHashRepartition hash = 3; // uint64 unknown = 4; - // New partitioning variants are stored in `partitioning`. + // } Partitioning partitioning = 5; bool preserve_order = 6; } @@ -1457,7 +1436,6 @@ 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 9f5785f4cc64a..bd859f2c080b7 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -6848,12 +6848,6 @@ 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)?; @@ -6890,12 +6884,6 @@ 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() } } @@ -6923,10 +6911,6 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "batchSize", "projection_exprs", "projectionExprs", - "partitioned_by_file_group", - "partitionedByFileGroup", - "output_partitioning", - "outputPartitioning", ]; #[allow(clippy::enum_variant_names)] @@ -6942,8 +6926,6 @@ 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 @@ -6976,8 +6958,6 @@ 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)), } } @@ -7008,8 +6988,6 @@ 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 => { @@ -7083,18 +7061,6 @@ 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 { @@ -7109,8 +7075,6 @@ 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__, }) } } @@ -16174,9 +16138,6 @@ 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() @@ -16193,7 +16154,6 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin", "hash", "unknown", - "range", ]; #[allow(clippy::enum_variant_names)] @@ -16201,7 +16161,6 @@ 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 @@ -16226,7 +16185,6 @@ 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)), } } @@ -16268,13 +16226,6 @@ 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 { @@ -19859,207 +19810,6 @@ 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 @@ -22238,207 +21988,6 @@ 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 @@ -22867,9 +22416,6 @@ 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() @@ -22886,7 +22432,6 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { "round_robin", "roundRobin", "hash", - "range", ]; #[allow(clippy::enum_variant_names)] @@ -22894,7 +22439,6 @@ 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 @@ -22919,7 +22463,6 @@ 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)), } } @@ -22960,13 +22503,6 @@ 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 06db2b35bbf72..ea25dba4abd8c 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, 4")] + #[prost(oneof = "repartition_node::PartitionMethod", tags = "2, 3")] pub partition_method: ::core::option::Option, } /// Nested message and enum types in `RepartitionNode`. @@ -225,23 +225,9 @@ 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, @@ -1749,10 +1735,6 @@ 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 { @@ -2118,26 +2100,14 @@ 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>, - /// Legacy direct partitioning fields: + /// oneof partition_method { /// 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")] @@ -2145,7 +2115,7 @@ pub struct RepartitionExecNode { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct Partitioning { - #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3, 4")] + #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3")] pub partition_method: ::core::option::Option, } /// Nested message and enum types in `Partitioning`. @@ -2158,8 +2128,6 @@ 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 bed64879f5a92..cff90ba7e31b7 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, SplitPoint, TableReference, + NullEquality, RecursionUnnestOption, Result, ScalarValue, TableReference, UnnestOptions, exec_datafusion_err, internal_err, plan_datafusion_err, }; use datafusion_execution::TaskContext; @@ -834,14 +834,3 @@ 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 5b2acfcac0e60..4e01ebabb5f69 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, RangePartitioning, - RecursiveQuery, SkipType, TableSource, Unnest, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RecursiveQuery, SkipType, + TableSource, Unnest, }; use datafusion_expr::{ DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, @@ -70,7 +70,6 @@ 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}; @@ -676,16 +675,6 @@ 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) @@ -1642,19 +1631,6 @@ 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 86b980a109b90..0599e48d2795c 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, SplitPoint, TableReference, UnnestOptions}; +use datafusion_common::{NullEquality, TableReference, UnnestOptions}; use datafusion_expr::WriteOp; use datafusion_expr::dml::InsertOp; use datafusion_expr::expr::{ @@ -706,18 +706,6 @@ 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 37e7d6b5d84c2..62f989a111ed2 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -24,9 +24,7 @@ use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; use chrono::{TimeZone, Utc}; -use datafusion_common::{ - DataFusionError, Result, ScalarValue, internal_datafusion_err, not_impl_err, -}; +use datafusion_common::{DataFusionError, Result, 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}; @@ -53,9 +51,7 @@ 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, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use datafusion_proto_common::common::proto_error; use object_store::ObjectMeta; use object_store::path::Path; @@ -678,14 +674,6 @@ 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, @@ -697,49 +685,6 @@ 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> { @@ -809,12 +754,6 @@ 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 { @@ -843,18 +782,14 @@ pub fn parse_protobuf_file_scan_config( file_source }; - let mut config_builder = FileScanConfigBuilder::new(object_store_url, file_source) + let config = 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_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(); + .with_batch_size(proto.batch_size.map(|s| s as usize)) + .build(); Ok(config) } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 5d3fed8737ecf..73fca1bbe6070 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -44,9 +44,7 @@ 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, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use super::{ DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -660,11 +658,6 @@ 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, @@ -674,40 +667,6 @@ 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, @@ -786,11 +745,6 @@ 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` @@ -850,8 +804,6 @@ 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 555690d4db419..fbe1af5617e42 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -45,7 +45,6 @@ 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; @@ -72,8 +71,8 @@ use datafusion_common::config::TableOptions; use datafusion_common::format::ExplainFormat; use datafusion_common::scalar::ScalarStructBuilder; use datafusion_common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SplitPoint, - TableReference, internal_datafusion_err, internal_err, not_impl_err, plan_err, + DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference, + internal_datafusion_err, internal_err, not_impl_err, plan_err, }; use datafusion_execution::TaskContext; use datafusion_expr::dml::CopyTo; @@ -85,9 +84,9 @@ use datafusion_expr::logical_plan::{Extension, UserDefinedLogicalNodeCore}; use datafusion_expr::{ Accumulator, AggregateUDF, ColumnarValue, ExprFunctionExt, ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, - PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, TryCast, - Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, - WindowFunctionDefinition, WindowUDF, WindowUDFImpl, + PartitionEvaluator, ScalarUDF, Signature, TryCast, Volatility, WindowFrame, + WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, WindowUDF, + WindowUDFImpl, }; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::expr_fn::{ @@ -3277,7 +3276,9 @@ async fn roundtrip_empty_table_scan() -> Result<()> { Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( + &schema, + ))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3299,7 +3300,9 @@ 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(EmptyTable::new(Arc::clone(&schema))); + let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( + &schema, + ))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3375,8 +3378,14 @@ 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(EmptyTable::new(left_schema)))?; - ctx.register_table("t2", Arc::new(EmptyTable::new(right_schema)))?; + 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)), + )?; let left = ctx.table("t1").await?.into_optimized_plan()?; let right = ctx.table("t2").await?.into_optimized_plan()?; @@ -3397,91 +3406,3 @@ 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 674bf5a2d60ed..a0b7180ba90bc 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -89,8 +89,7 @@ use datafusion::physical_plan::windows::{ }; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, - PhysicalExpr, PlanProperties, RangePartitioning, SendableRecordBatchStream, - SplitPoint, Statistics, displayable, + PhysicalExpr, PlanProperties, SendableRecordBatchStream, Statistics, displayable, }; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion::scalar::ScalarValue; @@ -1996,21 +1995,6 @@ 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); @@ -4081,98 +4065,3 @@ 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 a83db2bfb947f..0edde71b939f4 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -53,8 +53,6 @@ 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; @@ -63,8 +61,6 @@ use log::info; use sqlparser::ast; use tempfile::TempDir; -mod range_partitioning; - /// Context for running tests pub struct TestContext { /// Context for running queries @@ -171,10 +167,6 @@ 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 deleted file mode 100644 index 4141e000145a8..0000000000000 --- a/datafusion/sqllogictest/src/test_context/range_partitioning.rs +++ /dev/null @@ -1,288 +0,0 @@ -// 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 deleted file mode 100644 index ec51888c24b1e..0000000000000 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ /dev/null @@ -1,1897 +0,0 @@ -# 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 2f885991e3497..bf9ce26b35537 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -624,48 +624,3 @@ 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 1b9e91c7c475a..50c4b3da86cbe 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs @@ -32,11 +32,6 @@ 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" @@ -55,11 +50,6 @@ 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"