From 27de32b7186bfb7af5c571f1780f34f0f23e9f82 Mon Sep 17 00:00:00 2001 From: Adam Reeve Date: Tue, 4 Aug 2026 16:39:17 +1200 Subject: [PATCH 1/5] Add ProgressiveEvalExec operator --- datafusion/common/src/config.rs | 11 + datafusion/physical-plan/src/sorts/mod.rs | 1 + .../src/sorts/progressive_eval.rs | 2018 +++++++++++++++++ .../test_files/information_schema.slt | 2 + docs/source/user-guide/configs.md | 1 + 5 files changed, 2033 insertions(+) create mode 100644 datafusion/physical-plan/src/sorts/progressive_eval.rs diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index f5742f09f9b08..0e2055d5d9403 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1076,6 +1076,17 @@ config_namespace! { /// /// Disabled by default, set to a number greater than 0 for enabling it. pub hash_join_buffering_capacity: usize, default = 0 + + /// Number of input streams to prefetch ahead-of-time for `ProgressiveEvalExec`. + /// Since `ProgressiveEvalExec` only polls one stream at a time in order, + /// we do not need to prefetch all streams at once, saving resources. However, if the + /// streams' IO time is much greater than their CPU/processing time, prefetching them will + /// help improve performance. + /// Default is 1 which means we will prefetch one extra stream before it is polled. + /// 0 means streams are only fetched immediately before they are required. + /// Increase this value if IO time to read a stream is often much more than CPU time to + /// process the previous one. + pub progressive_eval_num_prefetch_input_streams: usize, default = 1 } } diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs index ca8d4a4400c49..6f73ea758c91c 100644 --- a/datafusion/physical-plan/src/sorts/mod.rs +++ b/datafusion/physical-plan/src/sorts/mod.rs @@ -23,6 +23,7 @@ mod merge; mod multi_level_merge; pub mod partial_sort; pub mod partitioned_topk; +pub mod progressive_eval; pub mod sort; pub mod sort_preserving_merge; mod stream; diff --git a/datafusion/physical-plan/src/sorts/progressive_eval.rs b/datafusion/physical-plan/src/sorts/progressive_eval.rs new file mode 100644 index 0000000000000..969b02113b2a6 --- /dev/null +++ b/datafusion/physical-plan/src/sorts/progressive_eval.rs @@ -0,0 +1,2018 @@ +// 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. + +//! Defines the progressive eval plan + +use std::borrow::Cow::Borrowed; +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use crate::common::spawn_buffered; +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, Metric, MetricBuilder, MetricValue, + MetricsSet, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, +}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, ScalarValue, Statistics, internal_err}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use futures::{Stream, StreamExt, ready}; +use log::{debug, trace}; + +/// ProgressiveEval returns a stream of record batches in the order of its inputs. +/// It will stop when the number of output rows reaches the given limit. +/// +/// This takes an input execution plan and an optional limit N, and provided each partition of +/// the input plan is in the expected order, this operator will return the top N rows +/// in the order of the input plan (truncating the record batch that crosses the limit). +/// +/// ```text +/// ┌─────────────────────────┐ +/// │ ┌───┬───┬───┬───┐ │ +/// │ │ A │ B │ C │ D │ │──┐ +/// │ └───┴───┴───┴───┘ │ │ +/// └─────────────────────────┘ │ ┌───────────────────┐ ┌───────────────────────────────┐ +/// Stream 1 │ │ │ │ ┌───┬───╦═══╦───┬───╦═══╗ │ +/// ├─▶│ ProgressiveEval │───▶│ │ A │ B ║ C ║ D │ M ║ N ║ ... │ +/// │ │ │ │ └───┴─▲─╩═══╩───┴───╩═══╝ │ +/// ┌─────────────────────────┐ │ └───────────────────┘ └─┬─────┴───────────────────────┘ +/// │ ╔═══╦═══╗ │ │ +/// │ ║ M ║ N ║ │──┘ │ +/// │ ╚═══╩═══╝ │ Output only includes the top record batches that cover top N rows +/// └─────────────────────────┘ +/// Stream 2 +/// +/// +/// Input Streams Output stream +/// (in some order) (in same order) +/// ``` +#[derive(Debug, Clone)] +pub struct ProgressiveEvalExec { + /// Input plan + input: Arc, + + /// Corresponding value ranges of the input plan. + /// None if the value ranges are not available. + value_ranges: Option>, + + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + + /// Optional number of rows to fetch. Stops producing rows after this fetch + fetch: Option, + + /// Cache holding plan properties like equivalences, output partitioning, output ordering etc. + cache: Arc, +} + +impl ProgressiveEvalExec { + /// Create a new progressive-evaluation execution plan. + /// + // Requires that the input partitions are in order with respect to the input ordering, + // and non-overlapping. + pub fn new( + input: Arc, + value_ranges: Option>, + fetch: Option, + ) -> Self { + let cache = Arc::new(Self::compute_properties(&input, fetch)); + Self { + input, + value_ranges, + metrics: ExecutionPlanMetricsSet::new(), + fetch, + cache, + } + } + + /// Input plan + pub fn input(&self) -> &Arc { + &self.input + } + + /// Creates the cache object that stores the plan properties such as equivalence properties, partitioning, ordering, etc. + fn compute_properties( + input: &Arc, + fetch: Option, + ) -> PlanProperties { + // Progressive eval does not change the equivalence properties of its input. + // This assumes that if the input is ordered, then the input partitions are non-overlapping + // with respect to the ordering and in-order. + let eq_properties = input.equivalence_properties().clone(); + + // This node serializes all the data to a single partition + let output_partitioning = Partitioning::UnknownPartitioning(1); + + // A fetch limit makes the output finite even if the input is unbounded + let boundedness = if fetch.is_some() { + Boundedness::Bounded + } else { + input.boundedness() + }; + + PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + boundedness, + ) + } +} + +impl DisplayAs for ProgressiveEvalExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter<'_>, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ProgressiveEvalExec: ")?; + if let Some(fetch) = self.fetch { + write!(f, "fetch={fetch}, ")?; + }; + if let Some(value_ranges) = &self.value_ranges { + write!(f, "input_ranges={value_ranges:?}")?; + }; + } + DisplayFormatType::TreeRender => { + writeln!(f, "ProgressiveEvalExec")?; + if let Some(fetch) = self.fetch { + writeln!(f, "fetch={fetch}")?; + }; + } + } + Ok(()) + } +} + +impl ExecutionPlan for ProgressiveEvalExec { + fn name(&self) -> &'static str { + "ProgressiveEvalExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::UnspecifiedDistribution] + } + + fn required_input_ordering(&self) -> Vec> { + let input_ordering = self + .input() + .properties() + .output_ordering() + .cloned() + .map(OrderingRequirements::from); + + vec![input_ordering] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!( + "ProgressiveEvalExec expected 1 child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new( + Arc::::clone(&children[0]), + self.value_ranges.clone(), + self.fetch, + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + trace!("Start ProgressiveEvalExec::execute for partition: {partition}"); + if 0 != partition { + return internal_err!("ProgressiveEvalExec invalid partition {partition}"); + } + + let input_partitions = self + .input + .properties() + .output_partitioning() + .partition_count(); + trace!( + "Number of input partitions of ProgressiveEvalExec::execute: {input_partitions}" + ); + let schema = self.schema(); + + // Add a metric to record the number of inputs + let num_inputs = Count::new(); + num_inputs.add(input_partitions); + self.metrics.register(Arc::new(Metric::new( + MetricValue::Count { + name: Borrowed("num_inputs"), + count: num_inputs, + }, + None, + ))); + // Add a metric to record the number of inputs that are actually read which is <= num_inputs + let num_read_inputs_counter = + MetricBuilder::new(&self.metrics).global_counter("num_read_inputs"); + // Add other baseline metrics + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + + let result = ProgressiveEvalStream::new( + Arc::clone(&self.input), + Arc::clone(&context), + schema, + baseline_metrics, + num_read_inputs_counter, + self.fetch, + )?; + + debug!("Got stream result from ProgressiveEvalStream::new_from_receivers"); + + Ok(Box::pin(result)) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + // The single output partition carries the input's combined statistics, + // capped by the fetch limit if one is set. + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) + } + + fn child_stats_requests(&self, _partition: Option) -> Vec { + vec![ChildStats::At(None)] + } + + fn with_fetch(&self, limit: Option) -> Option> { + // Rebuild rather than clone so the cached plan properties reflect the new fetch + Some(Arc::new(Self::new( + Arc::::clone(&self.input), + self.value_ranges.clone(), + limit, + ))) + } + + fn fetch(&self) -> Option { + self.fetch + } +} + +/// Handle when to prefetch input streams and how to poll next record batch +struct InputStreams { + /// Input plan of the progressive eval exec + input_plan: Arc, + + /// Context of the progressive eval exec + context: Arc, + + /// Total input streams + input_stream_count: usize, + + /// Number of input streams to prefetch ahead of time + num_input_streams_to_prefetch: usize, + + /// Index of current stream + current_stream_idx: usize, + + /// Input stream to poll data + current_input_stream: Option, + + /// Prefetched Input streams + prefetched_input_streams: VecDeque, + + /// Used to record number of actually read input streams + num_read_inputs_counter: Count, +} + +impl InputStreams { + fn new( + input_plan: Arc, + context: Arc, + num_input_streams_to_prefetch: usize, + num_read_inputs_counter: Count, + ) -> Result { + let input_stream_count = input_plan + .properties() + .output_partitioning() + .partition_count(); + + let current_stream_idx = 0; + let mut current_input_stream = None; + // The capacity required for prefetched streams is 1 more than the number of streams to + // prefetch, because we push a new stream before popping the new current stream. It is + // also bounded by the total number of inputs, excluding the current stream. + let prefetch_capacity = num_input_streams_to_prefetch + .saturating_add(1) + .min(input_stream_count.saturating_sub(1)); + let mut prefetched_input_streams = VecDeque::with_capacity(prefetch_capacity); + + // Always start fetching the first input stream, and also start + // fetching an additional `num_input_streams_to_prefetch` inputs. + for i in 0..=num_input_streams_to_prefetch { + if i >= input_stream_count { + break; + } + + let input_stream = spawn_buffered( + input_plan.execute(i, Arc::::clone(&context))?, + 1, + ); + num_read_inputs_counter.add(1); + + if i == 0 { + current_input_stream = Some(input_stream); + } else { + prefetched_input_streams.push_back(input_stream); + } + } + + Ok(Self { + input_plan, + context, + input_stream_count, + num_input_streams_to_prefetch, + current_stream_idx, + current_input_stream, + prefetched_input_streams, + num_read_inputs_counter, + }) + } + + /// Set next available stream to current_input_stream + /// Also prefetch one more input stream if not all of them are prefetched yet + fn next_stream(&mut self) -> Result<()> { + // No more input stream + if self.current_stream_idx >= self.input_stream_count { + // all input streams must have been consumed already + if !self.prefetched_input_streams.is_empty() { + return internal_err!( + "Internal error in ProgressiveEvalStream: Expected no input streams left to read" + ); + } + + self.current_input_stream = None; + } else { + // prefetch one more input stream before setting next stream to the current input stream + let next_prefetch_idx = self + .current_stream_idx + .saturating_add(self.num_input_streams_to_prefetch) + .saturating_add(1); + if next_prefetch_idx < self.input_stream_count { + self.num_read_inputs_counter.add(1); + self.prefetched_input_streams.push_back(spawn_buffered( + self.input_plan.execute( + next_prefetch_idx, + Arc::::clone(&self.context), + )?, + 1, + )); + } + + self.current_stream_idx += 1; + self.current_input_stream = self.prefetched_input_streams.pop_front(); + } + Ok(()) + } + + fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll>> { + // All input streams have been read + if self.current_input_stream.is_none() { + return Poll::Ready(None); + } + + // Get next record batch + let mut poll; + loop { + poll = self + .current_input_stream + .as_mut() + .unwrap() + .poll_next_unpin(cx); + match poll { + // This input stream no longer has data, move to next stream + Poll::Ready(None) => { + if let Err(e) = self.next_stream() { + return Poll::Ready(Some(Err(e))); + } + if self.current_input_stream.is_none() { + // Have reached the end of all input streams + return Poll::Ready(None); + } + } + _ => break, + } + } + + poll + } +} + +/// Concat input streams until reaching the fetch limit +struct ProgressiveEvalStream { + /// Input streams + input_streams: InputStreams, + + /// The schema of the input and output. + schema: SchemaRef, + + /// used to record execution baseline metrics + baseline_metrics: BaselineMetrics, + + /// If the stream has encountered an error + aborted: bool, + + /// Optional number of rows to fetch + fetch: Option, + + /// number of rows produced + produced: usize, +} + +impl ProgressiveEvalStream { + fn new( + input_plan: Arc, + context: Arc, + schema: SchemaRef, + baseline_metrics: BaselineMetrics, + num_read_inputs_counter: Count, + fetch: Option, + ) -> Result { + let num_input_streams_to_prefetch = context + .session_config() + .options() + .execution + .progressive_eval_num_prefetch_input_streams; + let input_streams = InputStreams::new( + input_plan, + context, + num_input_streams_to_prefetch, + num_read_inputs_counter, + )?; + + Ok(Self { + input_streams, + schema, + baseline_metrics, + aborted: false, + fetch, + produced: 0, + }) + } +} + +impl Stream for ProgressiveEvalStream { + type Item = Result; + + // Return the next record batch until reaching the fetch limit or the end of all input streams + // Return pending if the next record batch is not ready + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + // Error in previous poll + if self.aborted { + return Poll::Ready(None); + } + + // Have reached the fetch limit + if self.produced >= self.fetch.unwrap_or(usize::MAX) { + return Poll::Ready(None); + } + + let poll = self.input_streams.poll_next(cx); + + let poll = match ready!(poll) { + // This input stream has data, return its next record batch, + // truncated to the remaining fetch budget + Some(Ok(batch)) => { + let remaining = self.fetch.unwrap_or(usize::MAX) - self.produced; + let batch = if batch.num_rows() > remaining { + batch.slice(0, remaining) + } else { + batch + }; + self.produced += batch.num_rows(); + Poll::Ready(Some(Ok(batch))) + } + // This input stream has an error, return the error and set aborted to true to stop polling next round + Some(Err(e)) => { + self.aborted = true; + Poll::Ready(Some(Err(e))) + } + // This input stream has no more data, return None (aka finished) + None => { + // Reaching here means data of all streams have read + Poll::Ready(None) + } + }; + + self.baseline_metrics.record_poll(poll) + } +} + +impl RecordBatchStream for ProgressiveEvalStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::collect; + use crate::metrics::Timestamp; + use crate::statistics::StatisticsContext; + use crate::stream::RecordBatchStreamAdapter; + use crate::streaming::{PartitionStream, StreamingTableExec}; + use crate::test::exec::{BlockingExec, assert_strong_count_converges_to_zero}; + use crate::test::{TestMemoryExec, TestPartitionStream}; + use arrow::array::ArrayRef; + use arrow::array::{Int32Array, StringArray, TimestampNanosecondArray}; + use arrow::datatypes::Schema; + use arrow::datatypes::{DataType, Field}; + use arrow::record_batch::RecordBatch; + use datafusion_common::DataFusionError; + use datafusion_common::assert_batches_eq; + use datafusion_common::stats::Precision; + use datafusion_execution::config::SessionConfig; + use futures::FutureExt; + use std::iter::FromIterator; + + #[tokio::test] + async fn test_no_input_stream() { + let task_ctx = Arc::new(TaskContext::default()); + + let empty_table_result = ["++", "++"]; + + // no fetch limit --> return all rows + run_progressive_eval_test( + &[], + None, + None, + &empty_table_result, + 0, // 0 input streams + 0, // 0 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // limit = 0 means select nothing + run_progressive_eval_test( + &[], + None, + Some(0), + &empty_table_result, + 0, // 0 input streams + 0, // 0 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // limit = 1 on no data + run_progressive_eval_test( + &[], + None, + Some(1), + &empty_table_result, + 0, // 0 input streams + 0, // 0 input streams are fetched and polled + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_one_input_stream() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("c"), + Some("e"), + Some("g"), + Some("j"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let all_rows = [ + "+---+---+-------------------------------+", + "| a | b | c |", + "+---+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | c | 1970-01-01T00:00:00.000000007 |", + "| 7 | e | 1970-01-01T00:00:00.000000006 |", + "| 9 | g | 1970-01-01T00:00:00.000000005 |", + "| 3 | j | 1970-01-01T00:00:00.000000008 |", + "+---+---+-------------------------------+", + ]; + + // return all + run_progressive_eval_test( + &[vec![b1.clone()]], + None, + None, // no fetch limit --> return all rows + &all_rows, + 1, // 1 input stream + 1, // 1 input stream is fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // fetch no rows + run_progressive_eval_test( + &[vec![b1.clone()]], + None, + Some(0), + &["++", "++"], + 1, + 1, + Arc::clone(&task_ctx), + ) + .await; + + // return exactly 3 rows: the first record batch is truncated at the limit + run_progressive_eval_test( + &[vec![b1.clone()]], + None, + Some(3), + &[ + "+---+---+-------------------------------+", + "| a | b | c |", + "+---+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | c | 1970-01-01T00:00:00.000000007 |", + "| 7 | e | 1970-01-01T00:00:00.000000006 |", + "+---+---+-------------------------------+", + ], + 1, // 1 input stream + 1, // 1 input stream is fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // return all because fetch limit is larger + run_progressive_eval_test( + &[vec![b1.clone()]], + None, + Some(7), + &all_rows, + 1, // 1 input stream + 1, // 1 input stream is fetched and polled + Arc::clone(&task_ctx), + ) + .await; + } + + #[tokio::test] + async fn test_return_all() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("c"), + Some("e"), + Some("g"), + Some("j"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("b"), + Some("d"), + Some("f"), + Some("h"), + Some("j"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2, 2, 6])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let b1_b2 = [ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | c | 1970-01-01T00:00:00.000000007 |", + "| 7 | e | 1970-01-01T00:00:00.000000006 |", + "| 9 | g | 1970-01-01T00:00:00.000000005 |", + "| 3 | j | 1970-01-01T00:00:00.000000008 |", + "| 10 | b | 1970-01-01T00:00:00.000000004 |", + "| 20 | d | 1970-01-01T00:00:00.000000006 |", + "| 70 | f | 1970-01-01T00:00:00.000000002 |", + "| 90 | h | 1970-01-01T00:00:00.000000002 |", + "| 30 | j | 1970-01-01T00:00:00.000000006 |", + "+----+---+-------------------------------+", + ]; + + let b2_b1 = [ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 10 | b | 1970-01-01T00:00:00.000000004 |", + "| 20 | d | 1970-01-01T00:00:00.000000006 |", + "| 70 | f | 1970-01-01T00:00:00.000000002 |", + "| 90 | h | 1970-01-01T00:00:00.000000002 |", + "| 30 | j | 1970-01-01T00:00:00.000000006 |", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | c | 1970-01-01T00:00:00.000000007 |", + "| 7 | e | 1970-01-01T00:00:00.000000006 |", + "| 9 | g | 1970-01-01T00:00:00.000000005 |", + "| 3 | j | 1970-01-01T00:00:00.000000008 |", + "+----+---+-------------------------------+", + ]; + + // [b1, b2] + // return all by not specifying fetch limit + run_progressive_eval_test( + &[vec![b1.clone()], vec![b2.clone()]], + None, + None, // no fetch limit --> return all rows + &b1_b2, + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2] + // return all by specifying large limit + run_progressive_eval_test( + &[vec![b1.clone()], vec![b2.clone()]], + None, + Some(10), // limit = max num rows --> return all rows + &b1_b2, + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // [b2, b1] + // return all by not specifying fetch limit + run_progressive_eval_test( + &[vec![b2.clone()], vec![b1.clone()]], + None, + None, + &b2_b1, + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // [b2, b1] + // return all by specifying large limit + run_progressive_eval_test( + &[vec![b2], vec![b1]], + None, + Some(20), + &b2_b1, + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_return_all_on_different_length_batches() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // [b1, b2] + run_progressive_eval_test( + &[vec![b1.clone()], vec![b2.clone()]], + None, + None, + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | d | 1970-01-01T00:00:00.000000005 |", + "| 3 | e | 1970-01-01T00:00:00.000000008 |", + "| 70 | c | 1970-01-01T00:00:00.000000004 |", + "| 90 | d | 1970-01-01T00:00:00.000000006 |", + "| 30 | e | 1970-01-01T00:00:00.000000002 |", + "+----+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // [b2, b1] + run_progressive_eval_test( + &[vec![b2], vec![b1]], + None, + None, + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 70 | c | 1970-01-01T00:00:00.000000004 |", + "| 90 | d | 1970-01-01T00:00:00.000000006 |", + "| 30 | e | 1970-01-01T00:00:00.000000002 |", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | d | 1970-01-01T00:00:00.000000005 |", + "| 3 | e | 1970-01-01T00:00:00.000000008 |", + "+----+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_multiple_batches_per_partition() { + let task_ctx = Arc::new(TaskContext::default()); + let make_batch = |values: Vec| { + let a: ArrayRef = Arc::new(Int32Array::from(values)); + RecordBatch::try_from_iter(vec![("a", a)]).unwrap() + }; + let partitions = [ + vec![make_batch(vec![1, 2]), make_batch(vec![3, 4])], + vec![make_batch(vec![5, 6]), make_batch(vec![7, 8])], + ]; + + // No fetch limit: all batches of all partitions are returned in + // partition order + run_progressive_eval_test( + &partitions, + None, + None, + &[ + "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "| 4 |", "| 5 |", + "| 6 |", "| 7 |", "| 8 |", "+---+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // Fetch limit in the middle of the first partition's second batch: + // that batch is truncated + run_progressive_eval_test( + &partitions, + None, + Some(3), + &[ + "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", + ], + 2, // 2 input streams + 2, // the second stream is prefetched even though it is never polled + Arc::clone(&task_ctx), + ) + .await; + + // Fetch limit exactly at the end of the first partition: both of its + // batches are returned untruncated and nothing from the second + // partition is emitted + run_progressive_eval_test( + &partitions, + None, + Some(4), + &[ + "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "| 4 |", "+---+", + ], + 2, // 2 input streams + 2, // the second stream is prefetched even though it is never polled + Arc::clone(&task_ctx), + ) + .await; + + // Fetch limit in the middle of the second partition's first batch: + // all of the first partition plus a truncated batch from the second + run_progressive_eval_test( + &partitions, + None, + Some(5), + &[ + "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "| 4 |", "| 5 |", + "+---+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // With prefetch disabled, a fetch limit satisfied part-way through + // the first partition's batches never starts the second stream + run_progressive_eval_test( + &partitions, + None, + Some(3), + &[ + "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", + ], + 2, // 2 input streams + 1, // only the first stream is started + task_ctx_with_prefetch_depth(0), + ) + .await; + } + + #[tokio::test] + async fn test_fetch_limit_1() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // [b2, b1] + // b2 has 3 rows. b1 has 5 rows + // Fetch limit is 1 --> return the first row of the first batch (b2) + run_progressive_eval_test( + &[vec![b2.clone()], vec![b1.clone()]], + None, + Some(1), + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 70 | c | 1970-01-01T00:00:00.000000004 |", + "+----+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched by default even though only the first one is actually polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2] + // b1 has 5 rows. b2 has 3 rows + // Fetch limit is 1 --> return the first row of the first batch (b1) + run_progressive_eval_test( + &[vec![b1], vec![b2]], + None, + Some(1), + &[ + "+---+---+-------------------------------+", + "| a | b | c |", + "+---+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "+---+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched by default even though only the first one is actually polled + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_fetch_limit_equal_first_batch_size() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // [b2, b1] + // b2 has 3 rows. b1 has 5 rows + // Fetch limit is 3 --> return all 3 rows of the first batch (b2) that covers that limit + run_progressive_eval_test( + &[vec![b2.clone()], vec![b1.clone()]], + None, + Some(3), + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 70 | c | 1970-01-01T00:00:00.000000004 |", + "| 90 | d | 1970-01-01T00:00:00.000000006 |", + "| 30 | e | 1970-01-01T00:00:00.000000002 |", + "+----+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched by default even though only the first one is actually polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2] + // b1 has 5 rows. b2 has 3 rows + // Fetch limit is 5 --> return all 5 rows of first batch (b1) that covers that limit + run_progressive_eval_test( + &[vec![b1], vec![b2]], + None, + Some(5), + &[ + "+---+---+-------------------------------+", + "| a | b | c |", + "+---+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | d | 1970-01-01T00:00:00.000000005 |", + "| 3 | e | 1970-01-01T00:00:00.000000008 |", + "+---+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched by default even though only the first one is actually polled + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_fetch_limit_over_first_batch_size() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // [b2, b1] + // b2 has 3 rows. b1 has 5 rows + // Fetch limit is 4 --> return all of b2 plus the first row of b1 + run_progressive_eval_test( + &[vec![b2.clone()], vec![b1.clone()]], + None, + Some(4), + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 70 | c | 1970-01-01T00:00:00.000000004 |", + "| 90 | d | 1970-01-01T00:00:00.000000006 |", + "| 30 | e | 1970-01-01T00:00:00.000000002 |", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "+----+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2] + // b1 has 5 rows. b2 has 3 rows + // Fetch limit is 6 --> return all of b1 plus the first row of b2 + run_progressive_eval_test( + &[vec![b1], vec![b2]], + None, + Some(6), + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | d | 1970-01-01T00:00:00.000000005 |", + "| 3 | e | 1970-01-01T00:00:00.000000008 |", + "| 70 | c | 1970-01-01T00:00:00.000000004 |", + "+----+---+-------------------------------+", + ], + 2, // 2 input streams + 2, // all 2 input streams are fetched and polled + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_three_partitions_with_nulls() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + None, + Some("f"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 70])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("e"), + Some("g"), + Some("h"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![40, 60, 20])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 700, 900])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + None, + Some("g"), + Some("h"), + Some("i"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2, 2])); + let b3 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // [b1, b2, b3] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows + // Fetch limit is 1 --> return the first row of b1 + run_progressive_eval_test( + &[vec![b1.clone()], vec![b2.clone()], vec![b3.clone()]], + None, + Some(1), + &[ + "+---+---+-------------------------------+", + "| a | b | c |", + "+---+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "+---+---+-------------------------------+", + ], + 3, // 3 input streams + 2, // 2 input streams are fetched by default even though only the first one is polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows + // Fetch limit is 7 --> return all rows of b1 plus the first 2 rows of b2 + run_progressive_eval_test( + &[vec![b1.clone()], vec![b2.clone()], vec![b3.clone()]], + None, + Some(7), + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | | 1970-01-01T00:00:00.000000005 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "| 10 | e | 1970-01-01T00:00:00.000000040 |", + "| 20 | g | 1970-01-01T00:00:00.000000060 |", + "+----+---+-------------------------------+", + ], + 3, // 3 input streams + 3, // since we need to poll 2 input streams, 1 extra stream is prefetched + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows + // Fetch limit is 50 --> return all rows of all batches in the order of b1, b2, b3 + run_progressive_eval_test( + &[vec![b1], vec![b2], vec![b3]], + None, + Some(50), + &[ + "+-----+---+-------------------------------+", + "| a | b | c |", + "+-----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | | 1970-01-01T00:00:00.000000005 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "| 10 | e | 1970-01-01T00:00:00.000000040 |", + "| 20 | g | 1970-01-01T00:00:00.000000060 |", + "| 70 | h | 1970-01-01T00:00:00.000000020 |", + "| 100 | | 1970-01-01T00:00:00.000000004 |", + "| 200 | g | 1970-01-01T00:00:00.000000006 |", + "| 700 | h | 1970-01-01T00:00:00.000000002 |", + "| 900 | i | 1970-01-01T00:00:00.000000002 |", + "+-----+---+-------------------------------+", + ], + 3, // 3 input streams + 3, // 3 input streams are fetched and polled + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_four_partitions_with_nulls() { + let task_ctx = Arc::new(TaskContext::default()); + + // partition 1 + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + None, + Some("f"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // partition 2 + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 70])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("e"), + Some("g"), + Some("h"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![40, 60, 20])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // partition 3 + let a: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 700, 900])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + None, + Some("g"), + Some("h"), + Some("i"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2, 2])); + let b3 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // partition 4 + let a: ArrayRef = Arc::new(Int32Array::from(vec![1000, 2000])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![None, Some("x")])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![40, 60])); + let b4 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + // [b1, b2, b3, b4] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows. b4 has 2 rows + // Fetch limit is 0 --> return nothing. + run_progressive_eval_test( + &[ + vec![b1.clone()], + vec![b2.clone()], + vec![b3.clone()], + vec![b4.clone()], + ], + None, + Some(0), + &["++", "++"], + 4, // 4 input streams + 2, // 2 input streams are fetched by default even though nothing is polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3, b4] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows. b4 has 2 rows + // Fetch limit is 3 --> return the first 3 rows of b1 + run_progressive_eval_test( + &[ + vec![b1.clone()], + vec![b2.clone()], + vec![b3.clone()], + vec![b4.clone()], + ], + None, + Some(3), + &[ + "+---+---+-------------------------------+", + "| a | b | c |", + "+---+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "+---+---+-------------------------------+", + ], + 4, // 4 input streams + 2, // 2 input streams are fetched and one stream is polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3, b4] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows. b4 has 2 rows + // Fetch limit is 5 --> return all 5 rows of b1 + run_progressive_eval_test( + &[ + vec![b1.clone()], + vec![b2.clone()], + vec![b3.clone()], + vec![b4.clone()], + ], + None, + Some(5), + &[ + "+---+---+-------------------------------+", + "| a | b | c |", + "+---+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | | 1970-01-01T00:00:00.000000005 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "+---+---+-------------------------------+", + ], + 4, // 4 input streams + 2, // 2 input streams are fetched and one stream is polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3, b4] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows. b4 has 2 rows + // Fetch limit is 8 --> return all 8 rows of b1 and b2 + // Fetched 3 input streams since we will always prefetch one extra one + run_progressive_eval_test( + &[ + vec![b1.clone()], + vec![b2.clone()], + vec![b3.clone()], + vec![b4.clone()], + ], + None, + Some(8), + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | | 1970-01-01T00:00:00.000000005 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "| 10 | e | 1970-01-01T00:00:00.000000040 |", + "| 20 | g | 1970-01-01T00:00:00.000000060 |", + "| 70 | h | 1970-01-01T00:00:00.000000020 |", + "+----+---+-------------------------------+", + ], + 4, // 4 input streams + 3, // 3 input streams are fetched and 2 streams are polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3, b4] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows. b4 has 2 rows + // Fetch limit is 12 --> return all 12 rows of b1, b2 and b3 + // Fetches 4 input streams since we will always prefetch one extra one + run_progressive_eval_test( + &[ + vec![b1.clone()], + vec![b2.clone()], + vec![b3.clone()], + vec![b4.clone()], + ], + None, + Some(12), + &[ + "+-----+---+-------------------------------+", + "| a | b | c |", + "+-----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | | 1970-01-01T00:00:00.000000005 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "| 10 | e | 1970-01-01T00:00:00.000000040 |", + "| 20 | g | 1970-01-01T00:00:00.000000060 |", + "| 70 | h | 1970-01-01T00:00:00.000000020 |", + "| 100 | | 1970-01-01T00:00:00.000000004 |", + "| 200 | g | 1970-01-01T00:00:00.000000006 |", + "| 700 | h | 1970-01-01T00:00:00.000000002 |", + "| 900 | i | 1970-01-01T00:00:00.000000002 |", + "+-----+---+-------------------------------+", + ], + 4, // 4 input streams + 4, // 4 input streams are fetched and 3 streams are polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3, b4] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows. b4 has 2 rows + // Fetch limit is 15 --> return all 15 rows of b1, b2, b3 and b4 + // Fetches all 4 input streams + run_progressive_eval_test( + &[ + vec![b1.clone()], + vec![b2.clone()], + vec![b3.clone()], + vec![b4.clone()], + ], + None, + Some(15), + &[ + "+------+---+-------------------------------+", + "| a | b | c |", + "+------+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | | 1970-01-01T00:00:00.000000005 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "| 10 | e | 1970-01-01T00:00:00.000000040 |", + "| 20 | g | 1970-01-01T00:00:00.000000060 |", + "| 70 | h | 1970-01-01T00:00:00.000000020 |", + "| 100 | | 1970-01-01T00:00:00.000000004 |", + "| 200 | g | 1970-01-01T00:00:00.000000006 |", + "| 700 | h | 1970-01-01T00:00:00.000000002 |", + "| 900 | i | 1970-01-01T00:00:00.000000002 |", + "| 1000 | | 1970-01-01T00:00:00.000000040 |", + "| 2000 | x | 1970-01-01T00:00:00.000000060 |", + "+------+---+-------------------------------+", + ], + 4, // 4 input streams + 4, // 4 input streams are fetched and polled + Arc::clone(&task_ctx), + ) + .await; + + // [b1, b2, b3, b4] + // b1 has 5 rows. b2 has 3 rows. b3 has 4 rows. b4 has 2 rows + // No fetch limit--> return all 15 rows of b1, b2, b3 and b4 + run_progressive_eval_test( + &[ + vec![b1.clone()], + vec![b2.clone()], + vec![b3.clone()], + vec![b4.clone()], + ], + None, + None, // No fetch limit + &[ + "+------+---+-------------------------------+", + "| a | b | c |", + "+------+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | | 1970-01-01T00:00:00.000000005 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "| 10 | e | 1970-01-01T00:00:00.000000040 |", + "| 20 | g | 1970-01-01T00:00:00.000000060 |", + "| 70 | h | 1970-01-01T00:00:00.000000020 |", + "| 100 | | 1970-01-01T00:00:00.000000004 |", + "| 200 | g | 1970-01-01T00:00:00.000000006 |", + "| 700 | h | 1970-01-01T00:00:00.000000002 |", + "| 900 | i | 1970-01-01T00:00:00.000000002 |", + "| 1000 | | 1970-01-01T00:00:00.000000040 |", + "| 2000 | x | 1970-01-01T00:00:00.000000060 |", + "+------+---+-------------------------------+", + ], + 4, // 4 input streams + 4, // all input streams end up read (lazily, 2 at a time) because no fetch limit stops early + Arc::clone(&task_ctx), + ) + .await; + } + + #[tokio::test] + async fn test_prefetch_depth_config() { + let make_partition = |values: Vec| { + let a: ArrayRef = Arc::new(Int32Array::from(values)); + vec![RecordBatch::try_from_iter(vec![("a", a)]).unwrap()] + }; + let partitions = [ + make_partition(vec![1, 2]), + make_partition(vec![3, 4]), + make_partition(vec![5, 6]), + make_partition(vec![7, 8]), + ]; + + let first_row = ["+---+", "| a |", "+---+", "| 1 |", "+---+"]; + let first_batch = ["+---+", "| a |", "+---+", "| 1 |", "| 2 |", "+---+"]; + let all_rows = [ + "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "| 4 |", "| 5 |", + "| 6 |", "| 7 |", "| 8 |", "+---+", + ]; + + // Prefetch depth 0: only the stream being polled is started, so a + // fetch limit satisfied by the first stream reads nothing else + run_progressive_eval_test( + &partitions, + None, + Some(1), + &first_row, + 4, // 4 input streams + 1, // only the first stream is started + task_ctx_with_prefetch_depth(0), + ) + .await; + + // Prefetch depth 0 without a fetch limit: streams are started one at + // a time until all of them have been read + run_progressive_eval_test( + &partitions, + None, + None, + &all_rows, + 4, // 4 input streams + 4, // all streams are eventually read + task_ctx_with_prefetch_depth(0), + ) + .await; + + // Prefetch depth 2: the current stream plus two more are started up + // front. The fetch limit is satisfied by the first stream, so no + // further streams are started + run_progressive_eval_test( + &partitions, + None, + Some(2), + &first_batch, + 4, // 4 input streams + 3, // the current stream plus 2 prefetched streams are started + task_ctx_with_prefetch_depth(2), + ) + .await; + + // Prefetch depth 3: all four streams are started up front even + // though only the first one is polled + run_progressive_eval_test( + &partitions, + None, + Some(1), + &first_row, + 4, // 4 input streams + 4, // all streams are started up front + task_ctx_with_prefetch_depth(3), + ) + .await; + + // A prefetch depth larger than the number of streams is capped + run_progressive_eval_test( + &partitions, + None, + Some(1), + &first_row, + 4, // 4 input streams + 4, // all streams are started up front + task_ctx_with_prefetch_depth(10), + ) + .await; + } + + #[test] + fn test_partition_statistics_account_for_fetch() { + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); + let batch = RecordBatch::try_from_iter(vec![("a", a)]).unwrap(); + let schema = batch.schema(); + let input = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); + + // Without a fetch limit the input statistics pass through unchanged + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None, None); + let stats = StatisticsContext::new() + .compute(&progressive, &StatisticsArgs::new().with_partition(Some(0))) + .unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(5)); + + // A fetch limit below the input row count caps the reported row count + let progressive = + ProgressiveEvalExec::new(Arc::clone(&input) as _, None, Some(3)); + let stats = StatisticsContext::new() + .compute(&progressive, &StatisticsArgs::new()) + .unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(3)); + + // A fetch limit above the input row count has no effect + let progressive = + ProgressiveEvalExec::new(Arc::clone(&input) as _, None, Some(10)); + let stats = StatisticsContext::new() + .compute(&progressive, &StatisticsArgs::new()) + .unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(5)); + + // Setting a fetch limit on an existing plan is reflected in its statistics + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None, None); + let limited = progressive.with_fetch(Some(2)).unwrap(); + let stats = StatisticsContext::new() + .compute(limited.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(2)); + } + + #[test] + fn test_boundedness_accounts_for_fetch() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + // An infinite streaming table reports `Boundedness::Unbounded`. The + // stream is never polled, so it doesn't need to produce any batches. + let input = Arc::new( + StreamingTableExec::try_new( + Arc::clone(&schema), + vec![Arc::new(TestPartitionStream { + schema, + batches: vec![], + }) as _], + None, + None, + true, + None, + ) + .unwrap(), + ); + + // Without a fetch limit an unbounded input makes the output unbounded + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None, None); + assert!(matches!( + progressive.properties().boundedness, + Boundedness::Unbounded { .. } + )); + + // A fetch limit makes the output finite regardless of the input + let progressive = + ProgressiveEvalExec::new(Arc::clone(&input) as _, None, Some(10)); + assert!(matches!( + progressive.properties().boundedness, + Boundedness::Bounded + )); + + // Removing the fetch limit from an existing plan updates its boundedness + let unlimited = progressive.with_fetch(None).unwrap(); + assert!(matches!( + unlimited.properties().boundedness, + Boundedness::Unbounded { .. } + )); + } + + /// Create a task context whose session config sets + /// `execution.progressive_eval_num_prefetch_input_streams` to `depth` + fn task_ctx_with_prefetch_depth(depth: usize) -> Arc { + let mut config = SessionConfig::new(); + config + .options_mut() + .execution + .progressive_eval_num_prefetch_input_streams = depth; + Arc::new(TaskContext::default().with_session_config(config)) + } + + async fn run_progressive_eval_test( + partitions: &[Vec], + value_ranges: Option>, + fetch: Option, + expected_result: &[&str], + expected_num_input_streams: usize, + expected_num_read_input_streams: usize, + context: Arc, + ) { + let schema = if partitions.is_empty() { + // Schema is arbitrary + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let batch = RecordBatch::try_from_iter(vec![("a", a)]).unwrap(); + batch.schema() + } else { + partitions[0][0].schema() + }; + + let exec = TestMemoryExec::try_new_exec(partitions, schema, None).unwrap(); + let progressive = Arc::new(ProgressiveEvalExec::new(exec, value_ranges, fetch)); + + let progressive_clone = Arc::clone(&progressive); + + let collected = collect(progressive, context).await.unwrap(); + assert_batches_eq!(expected_result, collected.as_slice()); + + // verify metrics + let metrics = progressive_clone.metrics().unwrap(); + let num_input_streams = Count::new(); + num_input_streams.add(expected_num_input_streams); + let num_read_input_streams = Count::new(); + num_read_input_streams.add(expected_num_read_input_streams); + + assert_eq!( + metrics.sum_by_name("num_inputs"), + Some(MetricValue::Count { + name: Borrowed("num_inputs"), + count: num_input_streams + }) + ); + assert_eq!( + metrics.sum_by_name("num_read_inputs"), + Some(MetricValue::Count { + name: Borrowed("num_read_inputs"), + count: num_read_input_streams + }) + ); + } + + #[tokio::test] + async fn test_merge_metrics() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("a"), Some("c")])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("b"), Some("d")])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + + let schema = b1.schema(); + let exec = + TestMemoryExec::try_new_exec(&[vec![b1], vec![b2]], schema, None).unwrap(); + let progressive = Arc::new(ProgressiveEvalExec::new(exec, None, None)); + + let collected = + collect(Arc::::clone(&progressive), task_ctx) + .await + .unwrap(); + let expected = [ + "+----+---+", + "| a | b |", + "+----+---+", + "| 1 | a |", + "| 2 | c |", + "| 10 | b |", + "| 20 | d |", + "+----+---+", + ]; + assert_batches_eq!(expected, collected.as_slice()); + + // Now, validate metrics + let metrics = progressive.metrics().unwrap(); + + assert_eq!(metrics.output_rows().unwrap(), 4); + assert!(metrics.elapsed_compute().unwrap() > 0); + + let num_input_streams = Count::new(); + num_input_streams.add(2); + assert_eq!( + metrics.sum_by_name("num_inputs"), + Some(MetricValue::Count { + name: Borrowed("num_inputs"), + count: num_input_streams + }) + ); + + let num_read_input_streams = Count::new(); + num_read_input_streams.add(2); + assert_eq!( + metrics.sum_by_name("num_read_inputs"), + Some(MetricValue::Count { + name: Borrowed("num_read_inputs"), + count: num_read_input_streams + }) + ); + + let mut saw_start = false; + let mut saw_end = false; + metrics.iter().for_each(|m| match m.value() { + MetricValue::StartTimestamp(ts) => { + saw_start = true; + assert!(nanos_from_timestamp(ts) > 0); + } + MetricValue::EndTimestamp(ts) => { + saw_end = true; + assert!(nanos_from_timestamp(ts) > 0); + } + _ => {} + }); + + assert!(saw_start); + assert!(saw_end); + } + + fn nanos_from_timestamp(ts: &Timestamp) -> i64 { + ts.value().unwrap().timestamp_nanos_opt().unwrap() + } + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 2)); + let refs = blocking_exec.refs(); + let progressive_exec = + Arc::new(ProgressiveEvalExec::new(blocking_exec, None, None)); + + let fut = collect(progressive_exec, task_ctx); + let mut fut = fut.boxed(); + + assert!( + fut.as_mut().now_or_never().is_none(), + "future should be pending" + ); + drop(fut); + + // The plan and its streams should be dropped along with the future; + // wait for the spawn_buffered tasks to notice and release their + // references. + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + async fn test_error_in_first_stream_aborts_output() { + let task_ctx = Arc::new(TaskContext::default()); + let exec = error_exec(2, 0); + let progressive = ProgressiveEvalExec::new(exec, None, None); + + let mut stream = progressive.execute(0, task_ctx).unwrap(); + + let batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 2); + + let err = stream.next().await.unwrap().unwrap_err(); + assert!(err.to_string().contains("error in partition 0"), "{err}"); + + // The error aborts the output stream: the second input stream still + // holds valid data, but it must not be emitted + assert!(stream.next().await.is_none()); + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn test_error_in_later_stream_propagates() { + let task_ctx = Arc::new(TaskContext::default()); + let exec = error_exec(3, 1); + let progressive = ProgressiveEvalExec::new(exec, None, None); + + let mut stream = progressive.execute(0, task_ctx).unwrap(); + + // Data before the error arrives intact: all of partition 0 + // and the first batch of partition 1 + let batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 2); + let batch = stream.next().await.unwrap().unwrap(); + assert_eq!(batch.num_rows(), 2); + + let err = stream.next().await.unwrap().unwrap_err(); + assert!(err.to_string().contains("error in partition 1"), "{err}"); + + // Partition 2 is not emitted after the error + assert!(stream.next().await.is_none()); + } + + /// A [`PartitionStream`] that yields one two-row batch, followed by an + /// error if `error` is set. Used to verify error propagation. + #[derive(Debug)] + struct ErrorPartitionStream { + schema: SchemaRef, + partition: usize, + error: bool, + } + + impl PartitionStream for ErrorPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { + let a: ArrayRef = Arc::new(Int32Array::from(vec![ + (self.partition * 2 + 1) as i32, + (self.partition * 2 + 2) as i32, + ])); + let batch = RecordBatch::try_new(Arc::clone(&self.schema), vec![a]).unwrap(); + let mut items = vec![Ok(batch)]; + if self.error { + items.push(Err(DataFusionError::Execution(format!( + "error in partition {}", + self.partition + )))); + } + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + futures::stream::iter(items), + )) + } + } + + /// Create an execution plan whose partitions each yield one two-row + /// batch, with the `err_partition` stream yielding an error after its + /// batch. + fn error_exec(n_partitions: usize, err_partition: usize) -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let partitions = (0..n_partitions) + .map(|partition| { + Arc::new(ErrorPartitionStream { + schema: Arc::clone(&schema), + partition, + error: partition == err_partition, + }) as _ + }) + .collect(); + Arc::new( + StreamingTableExec::try_new(schema, partitions, None, None, false, None) + .unwrap(), + ) + } +} diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 573fb04b3451b..493aa3ccba317 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -273,6 +273,7 @@ datafusion.execution.parquet.writer_version 1.0 datafusion.execution.perfect_hash_join_min_key_density 0.15 datafusion.execution.perfect_hash_join_small_build_threshold 1024 datafusion.execution.planning_concurrency 13 +datafusion.execution.progressive_eval_num_prefetch_input_streams 1 datafusion.execution.skip_partial_aggregation_probe_ratio_threshold 0.8 datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 datafusion.execution.skip_physical_aggregate_schema_check false @@ -433,6 +434,7 @@ datafusion.execution.parquet.writer_version 1.0 (writing) Sets parquet writer ve datafusion.execution.perfect_hash_join_min_key_density 0.15 The minimum required density of join keys on the build side to consider a perfect hash join (see `HashJoinExec` for more details). Density is calculated as: `(number of rows) / (max_key - min_key + 1)`. A perfect hash join may be used if the actual key density > this value. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. datafusion.execution.perfect_hash_join_small_build_threshold 1024 A perfect hash join (see `HashJoinExec` for more details) will be considered if the range of keys (max - min) on the build side is < this threshold. This provides a fast path for joins with very small key ranges, bypassing the density check. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. datafusion.execution.planning_concurrency 13 Fan-out during initial physical planning. This is mostly use to plan `UNION` children in parallel. Defaults to the number of CPU cores on the system +datafusion.execution.progressive_eval_num_prefetch_input_streams 1 Number of input streams to prefetch ahead-of-time for `ProgressiveEvalExec`. Since `ProgressiveEvalExec` only polls one stream at a time in order, we do not need to prefetch all streams at once, saving resources. However, if the streams' IO time is much greater than their CPU/processing time, prefetching them will help improve performance. Default is 1 which means we will prefetch one extra stream before it is polled. 0 means streams are only fetched immediately before they are required. Increase this value if IO time to read a stream is often much more than CPU time to process the previous one. datafusion.execution.skip_partial_aggregation_probe_ratio_threshold 0.8 Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode datafusion.execution.skip_physical_aggregate_schema_check false When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index e02ada03fc413..47203d97ff1d5 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -144,6 +144,7 @@ The following configuration settings are available: | datafusion.execution.objectstore_writer_buffer_size | 10485760 | Size (bytes) of data buffer DataFusion uses when writing output files. This affects the size of the data chunks that are uploaded to remote object stores (e.g. AWS S3). If very large (>= 100 GiB) output files are being written, it may be necessary to increase this size to avoid errors from the remote end point. | | datafusion.execution.enable_ansi_mode | false | Whether to enable ANSI SQL mode. The flag is experimental and relevant only for DataFusion Spark built-in functions When `enable_ansi_mode` is set to `true`, the query engine follows ANSI SQL semantics for expressions, casting, and error handling. This means: - **Strict type coercion rules:** implicit casts between incompatible types are disallowed. - **Standard SQL arithmetic behavior:** operations such as division by zero, numeric overflow, or invalid casts raise runtime errors rather than returning `NULL` or adjusted values. - **Consistent ANSI behavior** for string concatenation, comparisons, and `NULL` handling. When `enable_ansi_mode` is `false` (the default), the engine uses a more permissive, non-ANSI mode designed for user convenience and backward compatibility. In this mode: - Implicit casts between types are allowed (e.g., string to integer when possible). - Arithmetic operations are more lenient — for example, `abs()` on the minimum representable integer value returns the input value instead of raising overflow. - Division by zero or invalid casts may return `NULL` instead of failing. # Default `false` — ANSI SQL mode is disabled by default. | | datafusion.execution.hash_join_buffering_capacity | 0 | How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. | +| datafusion.execution.progressive_eval_num_prefetch_input_streams | 1 | Number of input streams to prefetch ahead-of-time for `ProgressiveEvalExec`. Since `ProgressiveEvalExec` only polls one stream at a time in order, we do not need to prefetch all streams at once, saving resources. However, if the streams' IO time is much greater than their CPU/processing time, prefetching them will help improve performance. Default is 1 which means we will prefetch one extra stream before it is polled. 0 means streams are only fetched immediately before they are required. Increase this value if IO time to read a stream is often much more than CPU time to process the previous one. | | datafusion.optimizer.enable_distinct_aggregation_soft_limit | true | When set to true, the optimizer will push a limit operation into grouped aggregations which have no aggregate expressions, as a soft limit, emitting groups once the limit is reached, before all rows in the group are read. | | datafusion.optimizer.enable_round_robin_repartition | true | When set to true, the physical plan optimizer will try to add round robin repartitioning to increase parallelism to leverage more CPU cores | | datafusion.optimizer.enable_topk_aggregation | true | When set to true, the optimizer will attempt to perform limit operations during aggregations, if possible | From d3837925349cd6d235f3e6a48b257881622cd86d Mon Sep 17 00:00:00 2001 From: Adam Reeve Date: Tue, 18 Aug 2026 15:51:39 +1200 Subject: [PATCH 2/5] Add ReorderPartitionsExec --- .../src/joins/sort_merge_join/tests.rs | 13 +- datafusion/physical-plan/src/sorts/mod.rs | 1 + .../src/sorts/reorder_partitions.rs | 255 ++++++++++++++++++ datafusion/physical-plan/src/test.rs | 12 +- 4 files changed, 271 insertions(+), 10 deletions(-) create mode 100644 datafusion/physical-plan/src/sorts/reorder_partitions.rs diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index 91d1b893f1b29..c03206b853035 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -3443,9 +3443,8 @@ fn test_partition_statistics() -> Result<()> { // Test partition-specific statistics (partition = Some(0)) // The implementation correctly passes `partition` to children. - // Since the child TestMemoryExec returns unknown stats for specific partitions, - // the join output will also have Absent num_rows. This is expected behavior - // as the statistics depend on what the children can provide. + // The inputs have a single partition, so the statistics for partition 0 + // match the aggregate statistics. let partition_stats = StatisticsContext::new() .compute(&join_exec, &StatisticsArgs::new().with_partition(Some(0)))?; assert_eq!( @@ -3453,11 +3452,9 @@ fn test_partition_statistics() -> Result<()> { expected_cols, "Partition stats column count failed for {join_type:?}" ); - // When children return unknown stats, the join's partition stats will be Absent - assert!( - partition_stats.num_rows == Precision::Absent, - "Partition stats should have Absent num_rows when children return unknown for {join_type:?}, got {:?}", - partition_stats.num_rows + assert_eq!( + partition_stats.num_rows, stats.num_rows, + "Partition stats num_rows should match aggregate stats for {join_type:?}" ); } diff --git a/datafusion/physical-plan/src/sorts/mod.rs b/datafusion/physical-plan/src/sorts/mod.rs index 6f73ea758c91c..39d4c3b0a03e2 100644 --- a/datafusion/physical-plan/src/sorts/mod.rs +++ b/datafusion/physical-plan/src/sorts/mod.rs @@ -24,6 +24,7 @@ mod multi_level_merge; pub mod partial_sort; pub mod partitioned_topk; pub mod progressive_eval; +pub mod reorder_partitions; pub mod sort; pub mod sort_preserving_merge; mod stream; diff --git a/datafusion/physical-plan/src/sorts/reorder_partitions.rs b/datafusion/physical-plan/src/sorts/reorder_partitions.rs new file mode 100644 index 0000000000000..640c324eae062 --- /dev/null +++ b/datafusion/physical-plan/src/sorts/reorder_partitions.rs @@ -0,0 +1,255 @@ +// 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. + +//! Defines the reorder-partitions plan. + +use std::sync::Arc; + +use crate::{ + ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + StatisticsArgs, +}; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, Statistics, internal_err}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + +/// Pass-through operator that reorders partitions based on a provided permutation. +#[derive(Debug)] +pub struct ReorderPartitionsExec { + input: Arc, + /// For each output partition, the corresponding input partition + permutation: Vec, + properties: Arc, +} + +impl ReorderPartitionsExec { + pub fn new(input: Arc, permutation: Vec) -> Self { + let properties = Arc::clone(input.properties()); + Self { + input, + permutation, + properties, + } + } + + fn map_partition(&self, partition: usize) -> Result { + self.permutation.get(partition).copied().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ReorderPartitionsExec invalid partition {partition} for permutation of length {}", + self.permutation.len() + ) + }) + } +} + +impl DisplayAs for ReorderPartitionsExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter<'_>, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ReorderPartitionsExec: order={:?}", self.permutation) + } + DisplayFormatType::TreeRender => writeln!(f, "ReorderPartitionsExec"), + } + } +} + +impl ExecutionPlan for ReorderPartitionsExec { + fn name(&self) -> &'static str { + "ReorderPartitionsExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!( + "ReorderPartitionsExec expected 1 child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new( + Arc::::clone(&children[0]), + self.permutation.clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.input.execute(self.map_partition(partition)?, context) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + if input_stats.is_empty() { + return internal_err!( + "Could not get required input stats for ReorderPartitionsExec" + ); + } + Ok(Arc::clone(&input_stats[0])) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + match partition { + None => vec![ChildStats::At(None)], + Some(partition) => match self.map_partition(partition) { + Ok(input_partition) => vec![ChildStats::At(Some(input_partition))], + Err(_) => vec![], + }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::collect; + use crate::statistics::StatisticsContext; + use crate::test::TestMemoryExec; + use arrow::array::{ArrayRef, Int32Array}; + use arrow::record_batch::RecordBatch; + use datafusion_common::assert_batches_eq; + use datafusion_common::stats::Precision; + + /// Returns an input plan with `num_partitions` partitions, where + /// partition `i` contains a single batch with the value `i` + fn partitioned_input(num_partitions: usize) -> Arc { + let partitions: Vec> = (0..num_partitions) + .map(|i| { + let arr: ArrayRef = Arc::new(Int32Array::from(vec![i as i32])); + vec![RecordBatch::try_from_iter(vec![("i", arr)]).unwrap()] + }) + .collect(); + let schema = partitions[0][0].schema(); + TestMemoryExec::try_new_exec(&partitions, schema, None).unwrap() + } + + #[tokio::test] + async fn test_reorders_partitions() { + let task_ctx = Arc::new(TaskContext::default()); + let exec = ReorderPartitionsExec::new(partitioned_input(3), vec![2, 0, 1]); + + for (output_partition, input_partition) in [(0, 2), (1, 0), (2, 1)] { + let stream = exec + .execute(output_partition, Arc::clone(&task_ctx)) + .unwrap(); + let batches = collect(stream).await.unwrap(); + let expected_row = format!("| {input_partition} |"); + assert_batches_eq!( + ["+---+", "| i |", "+---+", &expected_row, "+---+"], + &batches + ); + } + } + + #[tokio::test] + async fn test_execute_invalid_partition() { + let task_ctx = Arc::new(TaskContext::default()); + let exec = ReorderPartitionsExec::new(partitioned_input(3), vec![2, 0, 1]); + + let err = match exec.execute(3, task_ctx) { + Ok(_) => panic!("Expected an Err result"), + Err(e) => e, + }; + assert!( + err.to_string() + .contains("invalid partition 3 for permutation of length 3"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_child_stats_requests_map_partitions() { + let exec = ReorderPartitionsExec::new(partitioned_input(3), vec![2, 0, 1]); + + // Overall stats come from overall stats of the input + assert_eq!(exec.child_stats_requests(None), vec![ChildStats::At(None)]); + + // Partition-specific stats require stats from remapped partition index + for (output_partition, input_partition) in [(0, 2), (1, 0), (2, 1)] { + assert_eq!( + exec.child_stats_requests(Some(output_partition)), + vec![ChildStats::At(Some(input_partition))] + ); + } + + // An out-of-range partition shouldn't panic + assert_eq!(exec.child_stats_requests(Some(3)), vec![]); + } + + #[test] + fn test_statistics_passed_through() { + // Input where partition i contains i + 1 rows, so partition statistics can be + // matched to their index. + let partitions: Vec> = (0..3) + .map(|i| { + let arr: ArrayRef = Arc::new(Int32Array::from(vec![i as i32; i + 1])); + vec![RecordBatch::try_from_iter(vec![("i", arr)]).unwrap()] + }) + .collect(); + let schema = partitions[0][0].schema(); + let input = TestMemoryExec::try_new_exec(&partitions, schema, None).unwrap(); + let exec = ReorderPartitionsExec::new(input, vec![2, 0, 1]); + + let ctx = StatisticsContext::new(); + + // Overall stats come from overall stats of the input + let stats = ctx.compute(&exec, &StatisticsArgs::new()).unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(6)); + + // Partition-specific stats map to stats from the remapped partition index + for (output_partition, input_partition) in [(0, 2), (1, 0), (2, 1)] { + let stats = ctx + .compute( + &exec, + &StatisticsArgs::new().with_partition(Some(output_partition)), + ) + .unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(input_partition + 1)); + } + } +} diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index b38a46d160755..8f111f90244bb 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -191,8 +191,16 @@ impl ExecutionPlan for TestMemoryExec { _input_stats: &[Arc], args: &StatisticsArgs, ) -> Result> { - if args.partition().is_some() { - Ok(Arc::new(Statistics::new_unknown(&self.schema))) + if let Some(partition) = args.partition() { + if let Some(batches) = self.partitions.get(partition) { + Ok(Arc::new(common::compute_record_batch_statistics( + std::slice::from_ref(batches), + &self.schema, + self.projection.clone(), + ))) + } else { + Ok(Arc::new(Statistics::new_unknown(&self.projected_schema))) + } } else { Ok(Arc::new(self.statistics_inner()?)) } From 4ef5e61746e3e8d1a399573403e70e6bbd32ce33 Mon Sep 17 00:00:00 2001 From: Adam Reeve Date: Wed, 19 Aug 2026 13:01:47 +1200 Subject: [PATCH 3/5] Remove value_ranges parameter from ProgressiveEvalExec This was unused except for display purposes. --- .../src/sorts/progressive_eval.rs | 84 +++---------------- 1 file changed, 13 insertions(+), 71 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/progressive_eval.rs b/datafusion/physical-plan/src/sorts/progressive_eval.rs index 969b02113b2a6..08b460cf81d2a 100644 --- a/datafusion/physical-plan/src/sorts/progressive_eval.rs +++ b/datafusion/physical-plan/src/sorts/progressive_eval.rs @@ -36,7 +36,7 @@ use crate::{ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{Result, ScalarValue, Statistics, internal_err}; +use datafusion_common::{Result, Statistics, internal_err}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{Distribution, OrderingRequirements, Partitioning}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -75,10 +75,6 @@ pub struct ProgressiveEvalExec { /// Input plan input: Arc, - /// Corresponding value ranges of the input plan. - /// None if the value ranges are not available. - value_ranges: Option>, - /// Execution metrics metrics: ExecutionPlanMetricsSet, @@ -94,15 +90,10 @@ impl ProgressiveEvalExec { /// // Requires that the input partitions are in order with respect to the input ordering, // and non-overlapping. - pub fn new( - input: Arc, - value_ranges: Option>, - fetch: Option, - ) -> Self { + pub fn new(input: Arc, fetch: Option) -> Self { let cache = Arc::new(Self::compute_properties(&input, fetch)); Self { input, - value_ranges, metrics: ExecutionPlanMetricsSet::new(), fetch, cache, @@ -155,9 +146,6 @@ impl DisplayAs for ProgressiveEvalExec { if let Some(fetch) = self.fetch { write!(f, "fetch={fetch}, ")?; }; - if let Some(value_ranges) = &self.value_ranges { - write!(f, "input_ranges={value_ranges:?}")?; - }; } DisplayFormatType::TreeRender => { writeln!(f, "ProgressiveEvalExec")?; @@ -229,7 +217,6 @@ impl ExecutionPlan for ProgressiveEvalExec { } Ok(Arc::new(Self::new( Arc::::clone(&children[0]), - self.value_ranges.clone(), self.fetch, ))) } @@ -307,7 +294,6 @@ impl ExecutionPlan for ProgressiveEvalExec { // Rebuild rather than clone so the cached plan properties reflect the new fetch Some(Arc::new(Self::new( Arc::::clone(&self.input), - self.value_ranges.clone(), limit, ))) } @@ -608,7 +594,6 @@ mod tests { run_progressive_eval_test( &[], None, - None, &empty_table_result, 0, // 0 input streams 0, // 0 input streams are fetched and polled @@ -619,7 +604,6 @@ mod tests { // limit = 0 means select nothing run_progressive_eval_test( &[], - None, Some(0), &empty_table_result, 0, // 0 input streams @@ -631,7 +615,6 @@ mod tests { // limit = 1 on no data run_progressive_eval_test( &[], - None, Some(1), &empty_table_result, 0, // 0 input streams @@ -670,7 +653,6 @@ mod tests { // return all run_progressive_eval_test( &[vec![b1.clone()]], - None, None, // no fetch limit --> return all rows &all_rows, 1, // 1 input stream @@ -682,7 +664,6 @@ mod tests { // fetch no rows run_progressive_eval_test( &[vec![b1.clone()]], - None, Some(0), &["++", "++"], 1, @@ -694,7 +675,6 @@ mod tests { // return exactly 3 rows: the first record batch is truncated at the limit run_progressive_eval_test( &[vec![b1.clone()]], - None, Some(3), &[ "+---+---+-------------------------------+", @@ -714,7 +694,6 @@ mod tests { // return all because fetch limit is larger run_progressive_eval_test( &[vec![b1.clone()]], - None, Some(7), &all_rows, 1, // 1 input stream @@ -787,7 +766,6 @@ mod tests { // return all by not specifying fetch limit run_progressive_eval_test( &[vec![b1.clone()], vec![b2.clone()]], - None, None, // no fetch limit --> return all rows &b1_b2, 2, // 2 input streams @@ -800,7 +778,6 @@ mod tests { // return all by specifying large limit run_progressive_eval_test( &[vec![b1.clone()], vec![b2.clone()]], - None, Some(10), // limit = max num rows --> return all rows &b1_b2, 2, // 2 input streams @@ -814,7 +791,6 @@ mod tests { run_progressive_eval_test( &[vec![b2.clone()], vec![b1.clone()]], None, - None, &b2_b1, 2, // 2 input streams 2, // all 2 input streams are fetched and polled @@ -826,7 +802,6 @@ mod tests { // return all by specifying large limit run_progressive_eval_test( &[vec![b2], vec![b1]], - None, Some(20), &b2_b1, 2, // 2 input streams @@ -863,7 +838,6 @@ mod tests { run_progressive_eval_test( &[vec![b1.clone()], vec![b2.clone()]], None, - None, &[ "+----+---+-------------------------------+", "| a | b | c |", @@ -888,7 +862,6 @@ mod tests { run_progressive_eval_test( &[vec![b2], vec![b1]], None, - None, &[ "+----+---+-------------------------------+", "| a | b | c |", @@ -927,7 +900,6 @@ mod tests { run_progressive_eval_test( &partitions, None, - None, &[ "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "| 4 |", "| 5 |", "| 6 |", "| 7 |", "| 8 |", "+---+", @@ -942,7 +914,6 @@ mod tests { // that batch is truncated run_progressive_eval_test( &partitions, - None, Some(3), &[ "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", @@ -958,7 +929,6 @@ mod tests { // partition is emitted run_progressive_eval_test( &partitions, - None, Some(4), &[ "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "| 4 |", "+---+", @@ -973,7 +943,6 @@ mod tests { // all of the first partition plus a truncated batch from the second run_progressive_eval_test( &partitions, - None, Some(5), &[ "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "| 4 |", "| 5 |", @@ -989,7 +958,6 @@ mod tests { // the first partition's batches never starts the second stream run_progressive_eval_test( &partitions, - None, Some(3), &[ "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", @@ -1029,7 +997,6 @@ mod tests { // Fetch limit is 1 --> return the first row of the first batch (b2) run_progressive_eval_test( &[vec![b2.clone()], vec![b1.clone()]], - None, Some(1), &[ "+----+---+-------------------------------+", @@ -1049,7 +1016,6 @@ mod tests { // Fetch limit is 1 --> return the first row of the first batch (b1) run_progressive_eval_test( &[vec![b1], vec![b2]], - None, Some(1), &[ "+---+---+-------------------------------+", @@ -1093,7 +1059,6 @@ mod tests { // Fetch limit is 3 --> return all 3 rows of the first batch (b2) that covers that limit run_progressive_eval_test( &[vec![b2.clone()], vec![b1.clone()]], - None, Some(3), &[ "+----+---+-------------------------------+", @@ -1115,7 +1080,6 @@ mod tests { // Fetch limit is 5 --> return all 5 rows of first batch (b1) that covers that limit run_progressive_eval_test( &[vec![b1], vec![b2]], - None, Some(5), &[ "+---+---+-------------------------------+", @@ -1163,7 +1127,6 @@ mod tests { // Fetch limit is 4 --> return all of b2 plus the first row of b1 run_progressive_eval_test( &[vec![b2.clone()], vec![b1.clone()]], - None, Some(4), &[ "+----+---+-------------------------------+", @@ -1186,7 +1149,6 @@ mod tests { // Fetch limit is 6 --> return all of b1 plus the first row of b2 run_progressive_eval_test( &[vec![b1], vec![b2]], - None, Some(6), &[ "+----+---+-------------------------------+", @@ -1245,7 +1207,6 @@ mod tests { // Fetch limit is 1 --> return the first row of b1 run_progressive_eval_test( &[vec![b1.clone()], vec![b2.clone()], vec![b3.clone()]], - None, Some(1), &[ "+---+---+-------------------------------+", @@ -1265,7 +1226,6 @@ mod tests { // Fetch limit is 7 --> return all rows of b1 plus the first 2 rows of b2 run_progressive_eval_test( &[vec![b1.clone()], vec![b2.clone()], vec![b3.clone()]], - None, Some(7), &[ "+----+---+-------------------------------+", @@ -1291,7 +1251,6 @@ mod tests { // Fetch limit is 50 --> return all rows of all batches in the order of b1, b2, b3 run_progressive_eval_test( &[vec![b1], vec![b2], vec![b3]], - None, Some(50), &[ "+-----+---+-------------------------------+", @@ -1371,7 +1330,6 @@ mod tests { vec![b3.clone()], vec![b4.clone()], ], - None, Some(0), &["++", "++"], 4, // 4 input streams @@ -1390,7 +1348,6 @@ mod tests { vec![b3.clone()], vec![b4.clone()], ], - None, Some(3), &[ "+---+---+-------------------------------+", @@ -1417,7 +1374,6 @@ mod tests { vec![b3.clone()], vec![b4.clone()], ], - None, Some(5), &[ "+---+---+-------------------------------+", @@ -1447,7 +1403,6 @@ mod tests { vec![b3.clone()], vec![b4.clone()], ], - None, Some(8), &[ "+----+---+-------------------------------+", @@ -1480,7 +1435,6 @@ mod tests { vec![b3.clone()], vec![b4.clone()], ], - None, Some(12), &[ "+-----+---+-------------------------------+", @@ -1517,7 +1471,6 @@ mod tests { vec![b3.clone()], vec![b4.clone()], ], - None, Some(15), &[ "+------+---+-------------------------------+", @@ -1555,7 +1508,6 @@ mod tests { vec![b3.clone()], vec![b4.clone()], ], - None, None, // No fetch limit &[ "+------+---+-------------------------------+", @@ -1608,7 +1560,6 @@ mod tests { // fetch limit satisfied by the first stream reads nothing else run_progressive_eval_test( &partitions, - None, Some(1), &first_row, 4, // 4 input streams @@ -1622,7 +1573,6 @@ mod tests { run_progressive_eval_test( &partitions, None, - None, &all_rows, 4, // 4 input streams 4, // all streams are eventually read @@ -1635,7 +1585,6 @@ mod tests { // further streams are started run_progressive_eval_test( &partitions, - None, Some(2), &first_batch, 4, // 4 input streams @@ -1648,7 +1597,6 @@ mod tests { // though only the first one is polled run_progressive_eval_test( &partitions, - None, Some(1), &first_row, 4, // 4 input streams @@ -1660,7 +1608,6 @@ mod tests { // A prefetch depth larger than the number of streams is capped run_progressive_eval_test( &partitions, - None, Some(1), &first_row, 4, // 4 input streams @@ -1678,30 +1625,28 @@ mod tests { let input = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); // Without a fetch limit the input statistics pass through unchanged - let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None, None); + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None); let stats = StatisticsContext::new() .compute(&progressive, &StatisticsArgs::new().with_partition(Some(0))) .unwrap(); assert_eq!(stats.num_rows, Precision::Exact(5)); // A fetch limit below the input row count caps the reported row count - let progressive = - ProgressiveEvalExec::new(Arc::clone(&input) as _, None, Some(3)); + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, Some(3)); let stats = StatisticsContext::new() .compute(&progressive, &StatisticsArgs::new()) .unwrap(); assert_eq!(stats.num_rows, Precision::Exact(3)); // A fetch limit above the input row count has no effect - let progressive = - ProgressiveEvalExec::new(Arc::clone(&input) as _, None, Some(10)); + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, Some(10)); let stats = StatisticsContext::new() .compute(&progressive, &StatisticsArgs::new()) .unwrap(); assert_eq!(stats.num_rows, Precision::Exact(5)); // Setting a fetch limit on an existing plan is reflected in its statistics - let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None, None); + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None); let limited = progressive.with_fetch(Some(2)).unwrap(); let stats = StatisticsContext::new() .compute(limited.as_ref(), &StatisticsArgs::new()) @@ -1730,15 +1675,14 @@ mod tests { ); // Without a fetch limit an unbounded input makes the output unbounded - let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None, None); + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, None); assert!(matches!( progressive.properties().boundedness, Boundedness::Unbounded { .. } )); // A fetch limit makes the output finite regardless of the input - let progressive = - ProgressiveEvalExec::new(Arc::clone(&input) as _, None, Some(10)); + let progressive = ProgressiveEvalExec::new(Arc::clone(&input) as _, Some(10)); assert!(matches!( progressive.properties().boundedness, Boundedness::Bounded @@ -1765,7 +1709,6 @@ mod tests { async fn run_progressive_eval_test( partitions: &[Vec], - value_ranges: Option>, fetch: Option, expected_result: &[&str], expected_num_input_streams: usize, @@ -1782,7 +1725,7 @@ mod tests { }; let exec = TestMemoryExec::try_new_exec(partitions, schema, None).unwrap(); - let progressive = Arc::new(ProgressiveEvalExec::new(exec, value_ranges, fetch)); + let progressive = Arc::new(ProgressiveEvalExec::new(exec, fetch)); let progressive_clone = Arc::clone(&progressive); @@ -1826,7 +1769,7 @@ mod tests { let schema = b1.schema(); let exec = TestMemoryExec::try_new_exec(&[vec![b1], vec![b2]], schema, None).unwrap(); - let progressive = Arc::new(ProgressiveEvalExec::new(exec, None, None)); + let progressive = Arc::new(ProgressiveEvalExec::new(exec, None)); let collected = collect(Arc::::clone(&progressive), task_ctx) @@ -1900,8 +1843,7 @@ mod tests { let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 2)); let refs = blocking_exec.refs(); - let progressive_exec = - Arc::new(ProgressiveEvalExec::new(blocking_exec, None, None)); + let progressive_exec = Arc::new(ProgressiveEvalExec::new(blocking_exec, None)); let fut = collect(progressive_exec, task_ctx); let mut fut = fut.boxed(); @@ -1924,7 +1866,7 @@ mod tests { async fn test_error_in_first_stream_aborts_output() { let task_ctx = Arc::new(TaskContext::default()); let exec = error_exec(2, 0); - let progressive = ProgressiveEvalExec::new(exec, None, None); + let progressive = ProgressiveEvalExec::new(exec, None); let mut stream = progressive.execute(0, task_ctx).unwrap(); @@ -1944,7 +1886,7 @@ mod tests { async fn test_error_in_later_stream_propagates() { let task_ctx = Arc::new(TaskContext::default()); let exec = error_exec(3, 1); - let progressive = ProgressiveEvalExec::new(exec, None, None); + let progressive = ProgressiveEvalExec::new(exec, None); let mut stream = progressive.execute(0, task_ctx).unwrap(); From 06343010750b98ecf597c7c9f415f1539cf40d9e Mon Sep 17 00:00:00 2001 From: Adam Reeve Date: Wed, 19 Aug 2026 11:01:03 +1200 Subject: [PATCH 4/5] Add "sequence sorted inputs" optimization rule --- datafusion/common/src/config.rs | 6 + datafusion/physical-optimizer/src/lib.rs | 1 + .../physical-optimizer/src/optimizer.rs | 5 + .../src/sequence_sorted_inputs.rs | 668 ++++++++++++++++++ .../test_files/information_schema.slt | 2 + docs/source/user-guide/configs.md | 1 + 6 files changed, 683 insertions(+) create mode 100644 datafusion/physical-optimizer/src/sequence_sorted_inputs.rs diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 0e2055d5d9403..854e238bee3bc 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1764,6 +1764,12 @@ config_namespace! { /// Default: true pub enable_sort_pushdown: bool, default = true + /// When set to true, the physical plan optimizer will replace + /// `SortPreservingMergeExec` with `ProgressiveEvalExec` when the input + /// partitions are non-overlapping ranges of the merge ordering, + /// avoiding a merge by emitting the partitions sequentially. + pub sequence_sorted_inputs: bool, default = false + /// When set to true, the optimizer will extract leaf expressions /// (such as `get_field`) from filter/sort/join nodes into projections /// closer to the leaf table scans, and push those projections down diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs index b9eb248f6e843..58384d7c84c8a 100644 --- a/datafusion/physical-optimizer/src/lib.rs +++ b/datafusion/physical-optimizer/src/lib.rs @@ -45,6 +45,7 @@ pub use datafusion_pruning as pruning; pub mod hash_join_buffering; pub mod pushdown_sort; pub mod sanity_checker; +pub mod sequence_sorted_inputs; pub mod topk_aggregation; pub mod topk_repartition; pub mod update_aggr_exprs; diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index aed25546cd09b..a43ed5ff4096e 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -38,6 +38,7 @@ use crate::update_aggr_exprs::OptimizeAggregateOrder; use crate::hash_join_buffering::HashJoinBuffering; use crate::limit_pushdown_past_window::LimitPushPastWindows; use crate::pushdown_sort::PushdownSort; +use crate::sequence_sorted_inputs::SequenceSortedInputs; use crate::window_topn::WindowTopN; use datafusion_common::config::ConfigOptions; @@ -174,6 +175,10 @@ impl PhysicalOptimizer { Arc::new(ProjectionPushdown::new()), // PushdownSort: Detect sorts that can be pushed down to data sources. Arc::new(PushdownSort::new()), + // SequenceSortedInputs: Replace SortPreservingMergeExec with ProgressiveEvalExec + // for partitions that don't overlap in the sort columns. + // Runs after PushdownSort which might introduce a SortPreservingMergeExec. + Arc::new(SequenceSortedInputs::new()), Arc::new(EnsureCooperative::new()), // This FilterPushdown handles dynamic filters that may have references to the source ExecutionPlan. // Therefore, it should be run at the end of the optimization process since any changes to the plan may break the dynamic filter's references. diff --git a/datafusion/physical-optimizer/src/sequence_sorted_inputs.rs b/datafusion/physical-optimizer/src/sequence_sorted_inputs.rs new file mode 100644 index 0000000000000..f1d20982f5902 --- /dev/null +++ b/datafusion/physical-optimizer/src/sequence_sorted_inputs.rs @@ -0,0 +1,668 @@ +// 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::PhysicalOptimizerRule; +use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::Statistics; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_physical_expr::LexOrdering; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_plan::sorts::{ + progressive_eval::ProgressiveEvalExec, reorder_partitions::ReorderPartitionsExec, + sort_preserving_merge::SortPreservingMergeExec, +}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties as _}; + +use std::cmp::Ordering; +use std::sync::Arc; + +/// Optimization that replaces [`SortPreservingMergeExec`] with [`ProgressiveEvalExec`] +/// when its input partitions are non-overlapping with respect to the merge +/// ordering. The partitions are arranged into the required order (via +/// [`ReorderPartitionsExec`] when they are not already laid out that way) so +/// that concatenating them yields globally ordered output. +#[derive(Debug, Default)] +pub struct SequenceSortedInputs; + +impl SequenceSortedInputs { + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for SequenceSortedInputs { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + if !config.optimizer.sequence_sorted_inputs { + return Ok(plan); + } + plan.transform_down(|plan| { + let Some(merge) = (plan.as_ref() as &dyn ExecutionPlan) + .downcast_ref::() + else { + return Ok(Transformed::no(plan)); + }; + let input = merge.input(); + let Some(permutation) = ordered_partition_permutation(input, merge.expr()) + else { + return Ok(Transformed::no(plan)); + }; + let ordered_input = + if permutation.iter().enumerate().all(|(idx, &src)| idx == src) { + Arc::clone(input) + } else { + Arc::new(ReorderPartitionsExec::new(Arc::clone(input), permutation)) + as Arc + }; + let replacement = ProgressiveEvalExec::new(ordered_input, merge.fetch()); + Ok(Transformed::yes(Arc::new(replacement) as _)) + }) + .data() + } + + fn name(&self) -> &str { + "SequenceSortedInputs" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Per-partition ordering statistics: the (start, end) value of each sort +/// column (oriented to the sort direction) plus each sort column's null count. +struct PartitionOrderStats { + starts: Vec, + ends: Vec, + null_counts: Vec, +} + +/// Find an arrangement of `plan`'s partitions whose concatenation is globally +/// ordered by `ordering`. Returns None if no such arrangement exists. +/// +/// Each partition is assumed to be internally ordered by `ordering` (the caller +/// only applies this to a merge whose input already satisfies the ordering). +fn ordered_partition_permutation( + plan: &Arc, + ordering: &LexOrdering, +) -> Option> { + let partition_count = plan.output_partitioning().partition_count(); + let mut stats: Vec = Vec::with_capacity(partition_count); + let stats_ctx = StatisticsContext::new(); + for partition_idx in 0..partition_count { + let partition_stats = stats_ctx + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(partition_idx)), + ) + .ok()?; + stats.push(get_ordering_stats(&partition_stats, ordering)?); + } + + // Order the partitions by the ordering + let mut perm: Vec = (0..partition_count).collect(); + perm.sort_by(|&a, &b| compare_partitions(&stats[a], &stats[b], ordering)); + + // Check for overlap in the ordering columns + for pair in perm.windows(2) { + if !boundary_ordered(&stats[pair[0]], &stats[pair[1]], ordering) { + return None; + } + } + + Some(perm) +} + +/// Compare two partitions by their boundary values corresponding to `ordering`. +fn compare_partitions( + a: &PartitionOrderStats, + b: &PartitionOrderStats, + ordering: &LexOrdering, +) -> Ordering { + compare_boundary_values(&a.starts, &b.starts, ordering) + .then_with(|| compare_boundary_values(&a.ends, &b.ends, ordering)) +} + +/// Compare two partition boundaries based on the ordering columns and sort directions. +/// Incomparable values are treated as equal. +fn compare_boundary_values( + a: &[ScalarValue], + b: &[ScalarValue], + ordering: &LexOrdering, +) -> Ordering { + for (i, sort_expr) in ordering.iter().enumerate() { + let Some(cmp) = a[i].partial_cmp(&b[i]) else { + return Ordering::Equal; + }; + let cmp = if sort_expr.options.descending { + cmp.reverse() + } else { + cmp + }; + if cmp != Ordering::Equal { + return cmp; + } + } + Ordering::Equal +} + +/// Whether we can prove that `cur` is ordered after `prev` under `ordering`. +fn boundary_ordered( + prev: &PartitionOrderStats, + cur: &PartitionOrderStats, + ordering: &LexOrdering, +) -> bool { + for (i, sort_expr) in ordering.iter().enumerate() { + // Reject nulls that could sort onto the wrong side of this boundary. + let boundary_null_count = if sort_expr.options.nulls_first { + cur.null_counts[i] + } else { + prev.null_counts[i] + }; + if boundary_null_count != 0 { + return false; + } + // Incomparable values (partial_cmp is None) are rejected. + let Some(cmp) = cur.starts[i].partial_cmp(&prev.ends[i]) else { + return false; + }; + let cmp = if sort_expr.options.descending { + cmp.reverse() + } else { + cmp + }; + match cmp { + Ordering::Greater => return true, + Ordering::Less => return false, + Ordering::Equal => continue, // Continue to the next sort column + } + } + // If we reach here, all sort columns are equal + true +} + +fn get_ordering_stats( + stats: &Arc, + ordering: &LexOrdering, +) -> Option { + let mut starts = Vec::with_capacity(ordering.len()); + let mut ends = Vec::with_capacity(ordering.len()); + let mut null_counts = Vec::with_capacity(ordering.len()); + + for sort_expr in ordering.iter() { + let column = sort_expr.expr.downcast_ref::()?; + let col_stats = stats.column_statistics.get(column.index())?; + // We require exact stats to guarantee no overlap in partition ranges. + if !(col_stats.null_count.is_exact()? + && col_stats.min_value.is_exact()? + && col_stats.max_value.is_exact()?) + { + return None; + } + // Note that for secondary sort columns, the start and end values are bounds + // on the actual start and end, as we only have access to the min/max stats, + // and not the sort-column values from the first and last rows. + // This means we may reject some orderings that are actually valid because we + // can't prove they're non-overlapping. + let (start, end) = if sort_expr.options.descending { + ( + col_stats.max_value.get_value()?, + col_stats.min_value.get_value()?, + ) + } else { + ( + col_stats.min_value.get_value()?, + col_stats.max_value.get_value()?, + ) + }; + // Stats may be null for all-null or empty partitions. + // For now, don't try to optimize this case: + if start.is_null() || end.is_null() { + return None; + } + starts.push(start.clone()); + ends.push(end.clone()); + null_counts.push(*col_stats.null_count.get_value()?); + } + + Some(PartitionOrderStats { + starts, + ends, + null_counts, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_common::stats::{ColumnStatistics, Precision}; + use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + use datafusion_physical_expr::{ + EquivalenceProperties, Partitioning, PhysicalSortExpr, + }; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PlanProperties}; + + /// Test plan with fixed per-partition statistics + #[derive(Debug)] + struct StatsTestExec { + stats: Vec, + cache: Arc, + } + + impl StatsTestExec { + fn new(stats: Vec) -> Self { + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(test_schema()), + Partitioning::UnknownPartitioning(stats.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { stats, cache } + } + } + + impl DisplayAs for StatsTestExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter<'_>, + ) -> std::fmt::Result { + write!(f, "StatsTestExec") + } + } + + impl ExecutionPlan for StatsTestExec { + fn name(&self) -> &'static str { + "StatsTestExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!("StatsTestExec is only used for planning") + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + match args.partition() { + Some(idx) => Ok(Arc::new(self.stats[idx].clone())), + None => Ok(Arc::new(Statistics::new_unknown(&self.schema()))), + } + } + } + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("t", DataType::Int64, true), + Field::new("id", DataType::Int64, true), + ])) + } + + fn sort_expr(index: usize, name: &str, options: SortOptions) -> PhysicalSortExpr { + PhysicalSortExpr::new(Arc::new(Column::new(name, index)), options) + } + + /// Ascending order on the specified column, with nulls last + fn asc(index: usize, name: &str) -> PhysicalSortExpr { + sort_expr( + index, + name, + SortOptions { + descending: false, + nulls_first: false, + }, + ) + } + + /// Descending order on the specified column, with nulls first + fn desc(index: usize, name: &str) -> PhysicalSortExpr { + sort_expr( + index, + name, + SortOptions { + descending: true, + nulls_first: true, + }, + ) + } + + fn exact_i64(min: i64, max: i64, null_count: usize) -> ColumnStatistics { + ColumnStatistics { + null_count: Precision::Exact(null_count), + min_value: Precision::Exact(ScalarValue::Int64(Some(min))), + max_value: Precision::Exact(ScalarValue::Int64(Some(max))), + ..Default::default() + } + } + + fn partition(column_statistics: Vec) -> Statistics { + Statistics { + num_rows: Precision::Exact(10), + total_byte_size: Precision::Exact(100), + column_statistics, + } + } + + #[test] + fn equal_first_column_boundary_ordered_by_second_column() { + // The partitions share t = 100 on the boundary; the disjoint id + // ranges disambiguate, so the check falls through to the second + // sort column and accepts. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 100, 0), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(100, 200, 0), exact_i64(11, 20, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + let perm = ordered_partition_permutation(&plan, &ordering) + .expect("expected an ordered permutation"); + assert_eq!(perm, vec![0, 1]); + } + + #[test] + fn equal_first_column_boundary_overlapping_second_column() { + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 100, 0), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(100, 200, 0), exact_i64(5, 20, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn boundary_equal_on_all_sort_columns_is_accepted() { + // min == prev max on every sort column is considered ordered. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 100, 0), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(100, 200, 0), exact_i64(10, 20, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_some()); + } + + #[test] + fn nulls_only_allowed_in_last_partition_for_nulls_last() { + let ordering = LexOrdering::new(vec![asc(0, "t")]).unwrap(); + + // Nulls in the last partition sort after all values: accepted. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 99, 0), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(100, 200, 2), exact_i64(0, 10, 0)]), + ])); + assert!(ordered_partition_permutation(&plan, &ordering).is_some()); + + // Nulls in the first partition would surface mid-stream: rejected. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 99, 2), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(100, 200, 0), exact_i64(0, 10, 0)]), + ])); + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn nulls_only_allowed_in_first_partition_for_nulls_first() { + let ordering = LexOrdering::new(vec![desc(0, "t")]).unwrap(); + + // Descending partitions with nulls leading in the first: accepted. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(100, 200, 2), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(0, 99, 0), exact_i64(0, 10, 0)]), + ])); + assert!(ordered_partition_permutation(&plan, &ordering).is_some()); + + // Nulls in the last partition sort before its values: rejected. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(100, 200, 0), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(0, 99, 2), exact_i64(0, 10, 0)]), + ])); + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn nulls_in_deeper_sort_column_hidden_by_first_column_break() { + // The middle partition carries a null in the second sort column. Its + // boundary with the first partition is strict on the first column, so + // the second column is never inspected for it. The boundary with the + // last partition is equal on the first column and falls through to + // the second: the middle partition's non-null max (4) < the last + // partition's min (5) looks ordered, but the middle partition's null + // rows sort after every non-null value (nulls last), so a row like + // (20, NULL) would precede (20, 5) in the concatenation. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 9, 0), exact_i64(0, 9, 0)]), + partition(vec![exact_i64(10, 20, 0), exact_i64(0, 4, 1)]), + partition(vec![exact_i64(20, 30, 0), exact_i64(5, 8, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn nulls_in_deeper_sort_column_harmless_when_boundaries_strict_on_first() { + // Both of the middle partition's boundaries are strict on the first + // column, so the second sort column is never relied on and its nulls + // cannot surface out of order. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 9, 0), exact_i64(0, 9, 0)]), + partition(vec![exact_i64(10, 19, 0), exact_i64(0, 4, 1)]), + partition(vec![exact_i64(20, 30, 0), exact_i64(5, 8, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_some()); + } + + #[test] + fn nulls_in_first_partition_with_all_equal_values_rejected() { + // Every partition shares the same value on both sort columns, so no + // boundary breaks early and every column is inspected. The first + // partition's nulls (sorting last) would surface before the later + // partitions' rows; the first boundary must reject them. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 1)]), + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 0)]), + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn nulls_in_middle_partition_with_all_equal_values_rejected() { + // As above, but the nulls sit in the middle partition: its boundary + // with the *next* partition is the one that must reject them. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 0)]), + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 1)]), + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn nulls_in_last_partition_with_all_equal_values_accepted() { + // Nulls sorting last in the last partition stream at the very end of + // the concatenation: correct, and there is no later boundary to + // invalidate. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 0)]), + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 0)]), + partition(vec![exact_i64(5, 5, 0), exact_i64(7, 7, 1)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t"), asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_some()); + } + + #[test] + fn equal_starts_ordered_by_end_values() { + // Both partitions start at 0, but the [0, 0] partition must come + // before [0, 1]: placing the partition with the smaller end first is + // the only arrangement whose concatenation stays ordered. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 1, 0), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(0, 0, 0), exact_i64(0, 10, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t")]).unwrap(); + + let perm = ordered_partition_permutation(&plan, &ordering) + .expect("expected a permutation"); + assert_eq!(perm, vec![1, 0]); + } + + #[test] + fn descending_ordering_reorders_ascending_layout() { + // Ascending partition layout under a descending ordering: the + // partitions are non-overlapping, so a permutation reverses them into + // descending order rather than bailing. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 99, 0), exact_i64(0, 10, 0)]), + partition(vec![exact_i64(100, 200, 0), exact_i64(0, 10, 0)]), + ])); + let ordering = LexOrdering::new(vec![desc(0, "t")]).unwrap(); + + let perm = ordered_partition_permutation(&plan, &ordering) + .expect("expected a permutation"); + assert_eq!(perm, vec![1, 0]); + } + + #[test] + fn incomparable_statistics_types_bail_out() { + // Mismatched stat types across partitions are incomparable; they + // must not be treated as an equal boundary. + let utf8 = |value: &str| ScalarValue::Utf8(Some(value.to_string())); + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 100, 0), exact_i64(0, 10, 0)]), + partition(vec![ + ColumnStatistics { + null_count: Precision::Exact(0), + min_value: Precision::Exact(utf8("a")), + max_value: Precision::Exact(utf8("b")), + ..Default::default() + }, + exact_i64(0, 10, 0), + ]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn inexact_statistics_bail_out() { + let inexact = ColumnStatistics { + null_count: Precision::Exact(0), + min_value: Precision::Inexact(ScalarValue::Int64(Some(100))), + max_value: Precision::Exact(ScalarValue::Int64(Some(200))), + ..Default::default() + }; + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 99, 0), exact_i64(0, 10, 0)]), + partition(vec![inexact, exact_i64(0, 10, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(0, "t")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn null_statistics_values_bail_out() { + // All-null or empty partitions report exact but null min/max values; + // they prove nothing about the partition's range. An all-null first + // partition under an ascending nulls-first ordering is the dangerous + // layout: the null guard passes (the later partition has no nulls) + // and a null scalar compares before any value, so without the + // explicit bail-out the boundary would look ordered. + let all_null = ColumnStatistics { + null_count: Precision::Exact(10), + min_value: Precision::Exact(ScalarValue::Int64(None)), + max_value: Precision::Exact(ScalarValue::Int64(None)), + ..Default::default() + }; + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![all_null, exact_i64(0, 10, 0)]), + partition(vec![exact_i64(0, 99, 0), exact_i64(0, 10, 0)]), + ])); + let ordering = LexOrdering::new(vec![sort_expr( + 0, + "t", + SortOptions { + descending: false, + nulls_first: true, + }, + )]) + .unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } + + #[test] + fn missing_column_statistics_bail_out() { + // The ordering references a column index beyond the available + // statistics; the lookup must bail out rather than panic. + let plan: Arc = Arc::new(StatsTestExec::new(vec![ + partition(vec![exact_i64(0, 99, 0)]), + partition(vec![exact_i64(100, 200, 0)]), + ])); + let ordering = LexOrdering::new(vec![asc(1, "id")]).unwrap(); + + assert!(ordered_partition_permutation(&plan, &ordering).is_none()); + } +} diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 493aa3ccba317..b702022a97dee 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -339,6 +339,7 @@ datafusion.optimizer.repartition_file_scans true datafusion.optimizer.repartition_joins true datafusion.optimizer.repartition_sorts true datafusion.optimizer.repartition_windows true +datafusion.optimizer.sequence_sorted_inputs false datafusion.optimizer.skip_failed_rules false datafusion.optimizer.subset_repartition_threshold 4 datafusion.optimizer.top_down_join_key_reordering true @@ -500,6 +501,7 @@ datafusion.optimizer.repartition_file_scans true When set to `true`, datasource datafusion.optimizer.repartition_joins true Should DataFusion repartition data using the join keys to execute joins in parallel using the provided `target_partitions` level datafusion.optimizer.repartition_sorts true Should DataFusion execute sorts in a per-partition fashion and merge afterwards instead of coalescing first and sorting globally. With this flag is enabled, plans in the form below ```text "SortExec: [a@0 ASC]", " CoalescePartitionsExec", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` would turn into the plan below which performs better in multithreaded environments ```text "SortPreservingMergeExec: [a@0 ASC]", " SortExec: [a@0 ASC]", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` datafusion.optimizer.repartition_windows true Should DataFusion repartition data using the partitions keys to execute window functions in parallel using the provided `target_partitions` level +datafusion.optimizer.sequence_sorted_inputs false When set to true, the physical plan optimizer will replace `SortPreservingMergeExec` with `ProgressiveEvalExec` when the input partitions are non-overlapping ranges of the merge ordering, avoiding a merge by emitting the partitions sequentially. datafusion.optimizer.skip_failed_rules false When set to true, the logical plan optimizer will produce warning messages if any optimization rules produce errors and then proceed to the next rule. When set to false, any rules that produce errors will cause the query to fail datafusion.optimizer.subset_repartition_threshold 4 Partition count threshold for subset satisfaction optimization. When the current partition count is >= this threshold, DataFusion will skip repartitioning if the required partitioning expression is a subset of the current partition expression such as Hash(a) satisfies Hash(a, b). When the current partition count is < this threshold, DataFusion will repartition to increase parallelism even when subset satisfaction applies. Set to 0 to always repartition (disable subset satisfaction optimization). Set to a high value to always use subset satisfaction. Example (subset_repartition_threshold = 4): ```text Hash([a]) satisfies Hash([a, b]) because (Hash([a, b]) is subset of Hash([a]) If current partitions (3) < threshold (4), repartition: AggregateExec: mode=FinalPartitioned, gby=[a, b], aggr=[SUM(x)] RepartitionExec: partitioning=Hash([a, b], 8), input_partitions=3 AggregateExec: mode=Partial, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 3) If current partitions (8) >= threshold (4), use subset satisfaction: AggregateExec: mode=SinglePartitioned, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 8) ``` datafusion.optimizer.top_down_join_key_reordering true When set to true, the physical plan optimizer will run a top down process to reorder the join keys diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 47203d97ff1d5..b9611d71ccd16 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -182,6 +182,7 @@ The following configuration settings are available: | datafusion.optimizer.prefer_existing_union | false | When set to true, the optimizer will not attempt to convert Union to Interleave | | datafusion.optimizer.expand_views_at_output | false | When set to true, if the returned type is a view type then the output will be coerced to a non-view. Coerces `Utf8View` to `LargeUtf8`, and `BinaryView` to `LargeBinary`. | | datafusion.optimizer.enable_sort_pushdown | true | Enable sort pushdown optimization. When enabled, attempts to push sort requirements down to data sources that can natively handle them (e.g., by reversing file/row group read order). Returns **inexact ordering**: Sort operator is kept for correctness, but optimized input enables early termination for TopK queries (ORDER BY ... LIMIT N), providing significant speedup. Memory: No additional overhead (only changes read order). Future: Will add option to detect perfectly sorted data and eliminate Sort completely. Default: true | +| datafusion.optimizer.sequence_sorted_inputs | false | When set to true, the physical plan optimizer will replace `SortPreservingMergeExec` with `ProgressiveEvalExec` when the input partitions are non-overlapping ranges of the merge ordering, avoiding a merge by emitting the partitions sequentially. | | datafusion.optimizer.enable_leaf_expression_pushdown | true | When set to true, the optimizer will extract leaf expressions (such as `get_field`) from filter/sort/join nodes into projections closer to the leaf table scans, and push those projections down towards the leaf nodes. | | datafusion.optimizer.enable_unions_to_filter | false | When set to true, the logical optimizer will rewrite `UNION DISTINCT` branches that read from the same source and differ only by filter predicates into a single branch with a combined filter. This optimization is conservative and only applies when the branches share the same source and compatible wrapper nodes such as identical projections or aliases. | | datafusion.explain.logical_plan_only | false | When set to true, the explain statement will only print logical plans | From a261615d618e926d01fddad5e801a5c7a02e1762 Mon Sep 17 00:00:00 2001 From: Adam Reeve Date: Fri, 21 Aug 2026 14:38:48 +1200 Subject: [PATCH 5/5] Ensure inputs preserve their order when using ProgressiveEval --- .../src/sequence_sorted_inputs.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-optimizer/src/sequence_sorted_inputs.rs b/datafusion/physical-optimizer/src/sequence_sorted_inputs.rs index f1d20982f5902..adefe725ec7f0 100644 --- a/datafusion/physical-optimizer/src/sequence_sorted_inputs.rs +++ b/datafusion/physical-optimizer/src/sequence_sorted_inputs.rs @@ -66,11 +66,12 @@ impl PhysicalOptimizerRule for SequenceSortedInputs { else { return Ok(Transformed::no(plan)); }; + let input = preserve_input_order(Arc::clone(input))?; let ordered_input = if permutation.iter().enumerate().all(|(idx, &src)| idx == src) { - Arc::clone(input) + input } else { - Arc::new(ReorderPartitionsExec::new(Arc::clone(input), permutation)) + Arc::new(ReorderPartitionsExec::new(input, permutation)) as Arc }; let replacement = ProgressiveEvalExec::new(ordered_input, merge.fetch()); @@ -88,6 +89,17 @@ impl PhysicalOptimizerRule for SequenceSortedInputs { } } +/// Mark every node in `plan` as order-sensitive, so data sources keep their +/// partition-to-data mapping. +/// Otherwise, work stealing can modify which files end up being read in a given partition. +fn preserve_input_order(plan: Arc) -> Result> { + plan.transform_down(|plan| match plan.with_preserve_order(true) { + Some(pinned) => Ok(Transformed::yes(pinned)), + None => Ok(Transformed::no(plan)), + }) + .data() +} + /// Per-partition ordering statistics: the (start, end) value of each sort /// column (oriented to the sort direction) plus each sort column's null count. struct PartitionOrderStats {