Merge branch 'main' into dynamic-disk-limit

Author: Bukhtawar

datafusion/functions-aggregate-common/src/min_max.rs

@@ -339,11 +339,28 @@ where
                 payload
             } else {
                 // no existing value, make a new one
-                let value: &[u8] = values.value(i).as_ref();
-                let payload = make_payload_fn(Some(value));
+                let (new_view, payload) = if len <= 12 {
+                    // Inline path: bytes are already packed in view_u128.
+                    // The inline ByteView format is [len:u32 LE][data:12 bytes zero-padded],
+                    // so extracting bytes from the u128 avoids a round-trip through
+                    // values.value(i) (which reads the views buffer and returns the same slice).
+                    let view_bytes = view_u128.to_le_bytes();
+                    let value = &view_bytes[4..4 + len as usize];
+                    let payload = make_payload_fn(Some(value));
+                    // For inline strings, the stored view is identical to the input view:
+                    // make_view(value, 0, 0) produces the same u128 as view_u128.
+                    //
+                    // SAFETY: view_u128 was a valid view, and the enclosing `len <= 12`
+                    // ensures it is inline
+                    let new_view = unsafe { self.append_inline_view(view_u128) };
+                    (new_view, payload)
+                } else {
+                    let value: &[u8] = values.value(i).as_ref();
+                    let payload = make_payload_fn(Some(value));
+                    let new_view = self.append_value(value);
+                    (new_view, payload)
+                };
 
-                // Create view pointing to our buffers
-                let new_view = self.append_value(value);
                 let new_header = Entry {
                     view: new_view,
                     hash,
@@ -389,6 +406,26 @@ where
         }
     }
 
+    /// Append an already-computed inline view (len <= 12) directly, bypassing
+    /// buffer allocation.
+    ///
+    /// Returns the view that was stored (identical to the argument).
+    ///
+    /// # Safety
+    ///
+    /// `view` must be a valid inline `ByteView`: the length field in the low
+    /// 32 bits must be <= 12, and the remaining 12 bytes must hold the
+    /// value's bytes (zero-padded if shorter). Calling with a non-inline view
+    /// would store a value that downstream `views` consumers interpret as
+    /// `[buffer_index, offset]` into the `completed`/`in_progress` buffers,
+    /// which is unsound for any view that didn't originate from a real
+    /// allocation in those buffers.
+    unsafe fn append_inline_view(&mut self, view: u128) -> u128 {
+        self.views.push(view);
+        self.nulls.append_non_null();
+        view
+    }
+
     /// Append a value to our buffers and return the view pointing to it
     fn append_value(&mut self, value: &[u8]) -> u128 {
         let len = value.len();

datafusion/physical-expr-common/src/binary_view_map.rs

@@ -52,8 +52,7 @@ use datafusion_common::stats::Precision;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::utils::transpose;
 use datafusion_common::{
-    ColumnStatistics, DataFusionError, HashMap, assert_or_internal_err,
-    internal_datafusion_err, internal_err,
+    ColumnStatistics, DataFusionError, HashMap, assert_or_internal_err, internal_err,
 };
 use datafusion_common::{Result, not_impl_err};
 use datafusion_common_runtime::SpawnedTask;
@@ -681,46 +680,8 @@ impl BatchPartitioner {
                     // Finished building index-arrays for output partitions
                     timer.done();
 
-                    // Borrowing partitioner timer to prevent moving `self` to closure
-                    let partitioner_timer = &self.timer;
-
-                    let mut partitioned_batches = vec![];
-                    for (partition, p_indices) in indices.iter_mut().enumerate() {
-                        if !p_indices.is_empty() {
-                            let taken_indices = std::mem::take(p_indices);
-                            let indices_array: PrimitiveArray<UInt32Type> =
-                                taken_indices.into();
-
-                            // Tracking time required for repartitioned batches construction
-                            let _timer = partitioner_timer.timer();
-
-                            // Produce batches based on indices
-                            let columns =
-                                take_arrays(batch.columns(), &indices_array, None)?;
-
-                            let mut options = RecordBatchOptions::new();
-                            options = options.with_row_count(Some(indices_array.len()));
-                            let batch = RecordBatch::try_new_with_options(
-                                batch.schema(),
-                                columns,
-                                &options,
-                            )
-                            .unwrap();
-
-                            partitioned_batches.push(Ok((partition, batch)));
-
-                            // Return the taken vec
-                            let (_, buffer, _) = indices_array.into_parts();
-                            let mut vec =
-                                buffer.into_inner().into_vec::<u32>().map_err(|e| {
-                                    internal_datafusion_err!(
-                                        "Could not convert buffer to vec: {e:?}"
-                                    )
-                                })?;
-                            vec.clear();
-                            *p_indices = vec;
-                        }
-                    }
+                    let partitioned_batches =
+                        Self::partition_grouped_take(&batch, indices, &self.timer)?;
 
                     Box::new(partitioned_batches.into_iter())
                 }
@@ -736,6 +697,66 @@ impl BatchPartitioner {
             BatchPartitionerState::Hash { indices, .. } => indices.len(),
         }
     }
+
+    /// Build repartitioned hash output batches using one `take` per input batch.
+    ///
+    /// The hash router first fills one index vector per output partition. This method
+    /// concatenates those index vectors, performs one grouped `take_arrays`, and
+    /// then returns each output partition as a slice of the reordered batch.
+    ///
+    /// For example, given partition indices:
+    ///
+    /// ```text
+    /// partition 0: [2, 5]
+    /// partition 1: []
+    /// partition 2: [0, 3, 4]
+    /// ```
+    ///
+    /// this method takes rows in `[2, 5, 0, 3, 4]` order once, then returns
+    /// `partition 0 = slice(0, 2)` and `partition 2 = slice(2, 3)`.
+    fn partition_grouped_take(
+        batch: &RecordBatch,
+        indices: &mut [Vec<u32>],
+        timer: &metrics::Time,
+    ) -> Result<Vec<Result<(usize, RecordBatch)>>> {
+        let mut partition_ranges = Vec::with_capacity(indices.len());
+        let mut reordered_indices = Vec::with_capacity(batch.num_rows());
+
+        for (partition, p_indices) in indices.iter_mut().enumerate() {
+            if p_indices.is_empty() {
+                continue;
+            }
+
+            let start = reordered_indices.len();
+            reordered_indices.extend_from_slice(p_indices);
+            partition_ranges.push((partition, start, p_indices.len()));
+            p_indices.clear();
+        }
+
+        if reordered_indices.is_empty() {
+            return Ok(vec![]);
+        }
+
+        let batches = {
+            let _timer = timer.timer();
+            let indices_array: PrimitiveArray<UInt32Type> = reordered_indices.into();
+            let columns = take_arrays(batch.columns(), &indices_array, None)?;
+
+            let mut options = RecordBatchOptions::new();
+            options = options.with_row_count(Some(indices_array.len()));
+            let reordered_batch =
+                RecordBatch::try_new_with_options(batch.schema(), columns, &options)?;
+
+            partition_ranges
+                .into_iter()
+                .map(|(partition, start, len)| {
+                    Ok((partition, reordered_batch.slice(start, len)))
+                })
+                .collect()
+        };
+
+        Ok(batches)
+    }
 }
 
 /// Maps `N` input partitions to `M` output partitions based on a

datafusion/physical-plan/src/repartition/mod.rs

@@ -886,14 +886,18 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                         "Execute statement with DEFAULT is not supported"
                     );
                 }
+                let name = name.ok_or_else(|| {
+                    plan_datafusion_err!("EXECUTE statement requires a name")
+                })?;
+
                 let empty_schema = DFSchema::empty();
                 let parameters = parameters
                     .into_iter()
                     .map(|expr| self.sql_to_expr(expr, &empty_schema, planner_context))
                     .collect::<Result<Vec<Expr>>>()?;
 
                 Ok(LogicalPlan::Statement(PlanStatement::Execute(Execute {
-                    name: object_name_to_string(&name.unwrap()),
+                    name: object_name_to_string(&name),
                     parameters,
                 })))
             }

datafusion/sql/src/statement.rs

@@ -48,9 +48,9 @@ use datafusion_common::{
 };
 use datafusion_expr::expr::{OUTER_REFERENCE_COLUMN_PREFIX, UNNEST_COLUMN_PREFIX};
 use datafusion_expr::{
-    BinaryExpr, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
+    Aggregate, BinaryExpr, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan,
     LogicalPlanBuilder, Operator, Projection, SortExpr, TableScan, Unnest,
-    UserDefinedLogicalNode, expr::Alias,
+    UserDefinedLogicalNode, Window, expr::Alias,
 };
 use sqlparser::ast::{self, Ident, OrderByKind, SetExpr, TableAliasColumnDef};
 use std::{sync::Arc, vec};
@@ -478,6 +478,80 @@ impl Unparser<'_> {
         Ok(false)
     }
 
+    fn project_window_output(
+        &self,
+        window_expr: &[Expr],
+        select: &mut SelectBuilder,
+        agg: Option<&Aggregate>,
+    ) -> Result<()> {
+        let mut items = if select.already_projected() {
+            select.pop_projections()
+        } else {
+            vec![ast::SelectItem::Wildcard(
+                ast::WildcardAdditionalOptions::default(),
+            )]
+        };
+
+        items.extend(
+            window_expr
+                .iter()
+                .map(|expr| {
+                    let expr = if let Some(agg) = agg {
+                        unproject_agg_exprs(expr.clone(), agg, None)?
+                    } else {
+                        expr.clone()
+                    };
+                    self.select_item_to_sql(&expr)
+                })
+                .collect::<Result<Vec<_>>>()?,
+        );
+        select.projection(items);
+
+        Ok(())
+    }
+
+    fn window_input_requires_derived_subquery(plan: &LogicalPlan) -> bool {
+        // These operators either produce a SELECT list or apply SQL clauses
+        // that are evaluated after window functions in a single SELECT block.
+        // Keep them below the Window node by emitting a derived table.
+        matches!(
+            plan,
+            LogicalPlan::Projection(_)
+                | LogicalPlan::Distinct(_)
+                | LogicalPlan::Limit(_)
+                | LogicalPlan::Sort(_)
+                | LogicalPlan::Union(_)
+        )
+    }
+
+    fn window_to_sql_with_derived_input(
+        &self,
+        window: &Window,
+        select: &mut SelectBuilder,
+        relation: &mut RelationBuilder,
+    ) -> Result<()> {
+        let input_alias = "derived_window_input";
+        self.derive(
+            window.input.as_ref(),
+            relation,
+            Some(self.new_table_alias(input_alias.to_string(), vec![])),
+            false,
+        )?;
+
+        let input_schema = window.input.schema();
+        let mut alias_rewriter = TableAliasRewriter {
+            table_schema: input_schema.as_arrow(),
+            alias_name: TableReference::bare(input_alias),
+        };
+        let window_expr = window
+            .window_expr
+            .iter()
+            .map(|expr| expr.clone().rewrite(&mut alias_rewriter).data())
+            .collect::<Result<Vec<_>>>()?;
+
+        self.project_window_output(&window_expr, select, None)
+    }
+
     #[cfg_attr(feature = "recursive_protection", recursive::recursive)]
     fn select_to_sql_recursively(
         &self,
@@ -610,25 +684,28 @@ impl Unparser<'_> {
                 self.select_to_sql_recursively(p.input.as_ref(), query, select, relation)
             }
             LogicalPlan::Filter(filter) => {
-                if let Some(agg) =
-                    find_agg_node_within_select(plan, select.already_projected())
+                let window = find_window_nodes_within_select(
+                    plan,
+                    None,
+                    select.already_projected(),
+                );
+                let agg = find_agg_node_within_select(plan, select.already_projected());
+
+                if let (Some(window), true) =
+                    (window.as_deref(), self.dialect.supports_qualify())
                 {
+                    let mut unprojected =
+                        unproject_window_exprs(filter.predicate.clone(), window)?;
+                    if let Some(agg) = agg {
+                        unprojected = unproject_agg_exprs(unprojected, agg, None)?;
+                    }
+                    let filter_expr = self.expr_to_sql(&unprojected)?;
+                    select.qualify(Some(filter_expr));
+                } else if let Some(agg) = agg {
                     let unprojected =
                         unproject_agg_exprs(filter.predicate.clone(), agg, None)?;
                     let filter_expr = self.expr_to_sql(&unprojected)?;
                     select.having(Some(filter_expr));
-                } else if let (Some(window), true) = (
-                    find_window_nodes_within_select(
-                        plan,
-                        None,
-                        select.already_projected(),
-                    ),
-                    self.dialect.supports_qualify(),
-                ) {
-                    let unprojected =
-                        unproject_window_exprs(filter.predicate.clone(), &window)?;
-                    let filter_expr = self.expr_to_sql(&unprojected)?;
-                    select.qualify(Some(filter_expr));
                 } else {
                     let filter_expr = self.expr_to_sql(&filter.predicate)?;
                     select.selection(Some(filter_expr));
@@ -1143,13 +1220,37 @@ impl Unparser<'_> {
                 Ok(())
             }
             LogicalPlan::Window(window) => {
-                // Window nodes are handled simultaneously with Projection nodes
+                // Window nodes are usually handled simultaneously with Projection
+                // nodes, where projected columns are unprojected back into their
+                // corresponding window expressions. Manually built plans can have
+                // Window nodes without an enclosing Projection, so in that case
+                // the Window node itself must contribute its output expressions.
+                let project_window_output = !select.already_projected();
+                if project_window_output
+                    && Self::window_input_requires_derived_subquery(window.input.as_ref())
+                {
+                    return self
+                        .window_to_sql_with_derived_input(window, select, relation);
+                }
+
+                let agg = if project_window_output {
+                    find_agg_node_within_select(plan, false)
+                } else {
+                    None
+                };
+
                 self.select_to_sql_recursively(
                     window.input.as_ref(),
                     query,
                     select,
                     relation,
-                )
+                )?;
+
+                if project_window_output {
+                    self.project_window_output(&window.window_expr, select, agg)?;
+                }
+
+                Ok(())
             }
             LogicalPlan::EmptyRelation(_) => {
                 // An EmptyRelation could be behind an UNNEST node. If the dialect supports UNNEST as a table factor,

datafusion/sql/src/unparser/plan.rs

@@ -54,7 +54,12 @@ pub(crate) fn find_agg_node_within_select(
     // Agg nodes explicitly return immediately with a single node
     if let LogicalPlan::Aggregate(agg) = input {
         Some(agg)
-    } else if let LogicalPlan::TableScan(_) = input {
+    } else if matches!(
+        input,
+        LogicalPlan::TableScan(_)
+            | LogicalPlan::Subquery(_)
+            | LogicalPlan::SubqueryAlias(_)
+    ) {
         None
     } else if let LogicalPlan::Projection(_) = input {
         if already_projected {