Merge branch 'main' into fix/sort-merge-reservation-starvation

Author: xudong963

datafusion-cli/src/exec.rs

@@ -196,6 +196,7 @@ pub async fn exec_from_repl(
             }
             Err(ReadlineError::Interrupted) => {
                 println!("^C");
+                rl.helper().unwrap().reset_hint();
                 continue;
             }
             Err(ReadlineError::Eof) => {

datafusion-cli/src/helper.rs

@@ -19,6 +19,7 @@
 //! and auto-completion for file name during creating external table.
 
 use std::borrow::Cow;
+use std::cell::Cell;
 
 use crate::highlighter::{Color, NoSyntaxHighlighter, SyntaxHighlighter};
 
@@ -40,6 +41,10 @@ pub struct CliHelper {
     completer: FilenameCompleter,
     dialect: Dialect,
     highlighter: Box<dyn Highlighter>,
+    /// Tracks whether to show the default hint. Set to `false` once the user
+    /// types anything, so the hint doesn't reappear after deleting back to
+    /// an empty line. Reset to `true` when the line is submitted.
+    show_hint: Cell<bool>,
 }
 
 impl CliHelper {
@@ -53,6 +58,7 @@ impl CliHelper {
             completer: FilenameCompleter::new(),
             dialect: *dialect,
             highlighter,
+            show_hint: Cell::new(true),
         }
     }
 
@@ -62,6 +68,11 @@ impl CliHelper {
         }
     }
 
+    /// Re-enable the default hint for the next prompt.
+    pub fn reset_hint(&self) {
+        self.show_hint.set(true);
+    }
+
     fn validate_input(&self, input: &str) -> Result<ValidationResult> {
         if let Some(sql) = input.strip_suffix(';') {
             let dialect = match dialect_from_str(self.dialect) {
@@ -119,12 +130,11 @@ impl Hinter for CliHelper {
     type Hint = String;
 
     fn hint(&self, line: &str, _pos: usize, _ctx: &Context<'_>) -> Option<String> {
-        if line.trim().is_empty() {
-            let suggestion = Color::gray(DEFAULT_HINT_SUGGESTION);
-            Some(suggestion)
-        } else {
-            None
+        if !line.is_empty() {
+            self.show_hint.set(false);
         }
+        (self.show_hint.get() && line.trim().is_empty())
+            .then(|| Color::gray(DEFAULT_HINT_SUGGESTION))
     }
 }
 
@@ -133,12 +143,9 @@ impl Hinter for CliHelper {
 fn is_open_quote_for_location(line: &str, pos: usize) -> bool {
     let mut sql = line[..pos].to_string();
     sql.push('\'');
-    if let Ok(stmts) = DFParser::parse_sql(&sql)
-        && let Some(Statement::CreateExternalTable(_)) = stmts.back()
-    {
-        return true;
-    }
-    false
+    DFParser::parse_sql(&sql).is_ok_and(|stmts| {
+        matches!(stmts.back(), Some(Statement::CreateExternalTable(_)))
+    })
 }
 
 impl Completer for CliHelper {
@@ -161,7 +168,9 @@ impl Completer for CliHelper {
 impl Validator for CliHelper {
     fn validate(&self, ctx: &mut ValidationContext<'_>) -> Result<ValidationResult> {
         let input = ctx.input().trim_end();
-        self.validate_input(input)
+        let result = self.validate_input(input);
+        self.reset_hint();
+        result
     }
 }
 

datafusion-examples/examples/custom_data_source/custom_datasource.rs

@@ -27,6 +27,7 @@ use async_trait::async_trait;
 use datafusion::arrow::array::{UInt8Builder, UInt64Builder};
 use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::datasource::{TableProvider, TableType, provider_as_source};
 use datafusion::error::Result;
 use datafusion::execution::context::TaskContext;
@@ -283,4 +284,20 @@ impl ExecutionPlan for CustomExec {
             None,
         )?))
     }
+
+    fn apply_expressions(
+        &self,
+        f: &mut dyn FnMut(
+            &dyn datafusion::physical_plan::PhysicalExpr,
+        ) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        // Visit expressions in the output ordering from equivalence properties
+        let mut tnr = TreeNodeRecursion::Continue;
+        if let Some(ordering) = self.cache.output_ordering() {
+            for sort_expr in ordering {
+                tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?;
+            }
+        }
+        Ok(tnr)
+    }
 }

datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs

@@ -29,6 +29,7 @@
 use arrow::record_batch::RecordBatch;
 use arrow_schema::SchemaRef;
 use datafusion::common::record_batch;
+use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::common::{exec_datafusion_err, internal_err};
 use datafusion::datasource::{DefaultTableSource, memory::MemTable};
 use datafusion::error::Result;
@@ -296,4 +297,20 @@ impl ExecutionPlan for BufferingExecutionPlan {
             }),
         )))
     }
+
+    fn apply_expressions(
+        &self,
+        f: &mut dyn FnMut(
+            &dyn datafusion::physical_plan::PhysicalExpr,
+        ) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        // Visit expressions in the output ordering from equivalence properties
+        let mut tnr = TreeNodeRecursion::Continue;
+        if let Some(ordering) = self.properties.output_ordering() {
+            for sort_expr in ordering {
+                tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?;
+            }
+        }
+        Ok(tnr)
+    }
 }

datafusion-examples/examples/proto/composed_extension_codec.rs

@@ -38,6 +38,7 @@ use std::sync::Arc;
 
 use datafusion::common::Result;
 use datafusion::common::internal_err;
+use datafusion::common::tree_node::TreeNodeRecursion;
 use datafusion::execution::TaskContext;
 use datafusion::physical_plan::{DisplayAs, ExecutionPlan};
 use datafusion::prelude::SessionContext;
@@ -128,6 +129,15 @@ impl ExecutionPlan for ParentExec {
     ) -> Result<datafusion::physical_plan::SendableRecordBatchStream> {
         unreachable!()
     }
+
+    fn apply_expressions(
+        &self,
+        _f: &mut dyn FnMut(
+            &dyn datafusion::physical_plan::PhysicalExpr,
+        ) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
+    }
 }
 
 /// A PhysicalExtensionCodec that can serialize and deserialize ParentExec
@@ -204,6 +214,15 @@ impl ExecutionPlan for ChildExec {
     ) -> Result<datafusion::physical_plan::SendableRecordBatchStream> {
         unreachable!()
     }
+
+    fn apply_expressions(
+        &self,
+        _f: &mut dyn FnMut(
+            &dyn datafusion::physical_plan::PhysicalExpr,
+        ) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
+    }
 }
 
 /// A PhysicalExtensionCodec that can serialize and deserialize ChildExec

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -117,7 +117,7 @@ use datafusion::{
 };
 use datafusion_common::{
     DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err,
-    plan_datafusion_err, plan_err,
+    plan_datafusion_err, plan_err, tree_node::TreeNodeRecursion,
 };
 use datafusion_expr::{
     UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
@@ -743,6 +743,22 @@ impl ExecutionPlan for SampleExec {
 
         Ok(stats)
     }
+
+    fn apply_expressions(
+        &self,
+        f: &mut dyn FnMut(
+            &dyn datafusion::physical_plan::PhysicalExpr,
+        ) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        // Visit expressions in the output ordering from equivalence properties
+        let mut tnr = TreeNodeRecursion::Continue;
+        if let Some(ordering) = self.cache.output_ordering() {
+            for sort_expr in ordering {
+                tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?;
+            }
+        }
+        Ok(tnr)
+    }
 }
 
 /// Bernoulli sampler: includes each row with probability `(upper - lower)`.

datafusion/catalog/src/memory/table.rs

@@ -32,6 +32,7 @@ use arrow::compute::{and, filter_record_batch};
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::error::Result;
+use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::{Constraints, DFSchema, SchemaExt, not_impl_err, plan_err};
 use datafusion_common_runtime::JoinSet;
 use datafusion_datasource::memory::{MemSink, MemorySourceConfig};
@@ -46,7 +47,7 @@ use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning,
-    PlanProperties, common,
+    PhysicalExpr, PlanProperties, common,
 };
 use datafusion_session::Session;
 
@@ -400,10 +401,7 @@ impl TableProvider for MemTable {
         let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
 
         // Create physical expressions for assignments upfront (outside batch loop)
-        let physical_assignments: HashMap<
-            String,
-            Arc<dyn datafusion_physical_plan::PhysicalExpr>,
-        > = assignments
+        let physical_assignments: HashMap<String, Arc<dyn PhysicalExpr>> = assignments
             .iter()
             .map(|(name, expr)| {
                 let physical_expr =
@@ -638,4 +636,11 @@ impl ExecutionPlan for DmlResultExec {
             stream,
         )))
     }
+
+    fn apply_expressions(
+        &self,
+        _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+    ) -> Result<TreeNodeRecursion> {
+        Ok(TreeNodeRecursion::Continue)
+    }
 }

datafusion/common/src/join_type.rs

@@ -97,6 +97,35 @@ impl JoinType {
         }
     }
 
+    /// Whether each side of the join is preserved for ON-clause filter pushdown.
+    ///
+    /// It is only correct to push ON-clause filters below a join for preserved
+    /// inputs.
+    ///
+    /// # "Preserved" input definition
+    ///
+    /// A join side is preserved if the join returns all or a subset of the rows
+    /// from that side, such that each output row directly maps to an input row.
+    /// If a side is not preserved, the join can produce extra null rows that
+    /// don't map to any input row.
+    ///
+    /// # Return Value
+    ///
+    /// A tuple of booleans - (left_preserved, right_preserved).
+    pub fn on_lr_is_preserved(&self) -> (bool, bool) {
+        match self {
+            JoinType::Inner => (true, true),
+            JoinType::Left => (false, true),
+            JoinType::Right => (true, false),
+            JoinType::Full => (false, false),
+            JoinType::LeftSemi | JoinType::RightSemi => (true, true),
+            JoinType::LeftAnti => (false, true),
+            JoinType::RightAnti => (true, false),
+            JoinType::LeftMark => (false, true),
+            JoinType::RightMark => (true, false),
+        }
+    }
+
     /// Does the join type support swapping inputs?
     pub fn supports_swap(&self) -> bool {
         matches!(