Merge branch 'main' into fix-offset-overflow

Author: maxburke

datafusion-cli/src/exec.rs

@@ -423,16 +423,17 @@ async fn create_plan(
 
         // Expose stdin (e.g. `cat data.csv | datafusion-cli`) as a `stdin://`
         // object store, registered like any other scheme in `get_object_store`.
-        cmd.location = StdinUtils::rewrite_location(&cmd.location, format.as_ref());
-
-        register_object_store_and_config_extensions(
-            ctx,
-            &cmd.location,
-            &cmd.options,
-            format,
-            resolve_region,
-        )
-        .await?;
+        for location in &mut cmd.locations {
+            *location = StdinUtils::rewrite_location(location, format.as_ref());
+            register_object_store_and_config_extensions(
+                ctx,
+                location,
+                &cmd.options,
+                format.clone(),
+                resolve_region,
+            )
+            .await?;
+        }
     }
 
     if let LogicalPlan::Copy(copy_to) = &mut plan {
@@ -535,14 +536,16 @@ mod tests {
 
         if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan {
             let format = config_file_type_from_str(&cmd.file_type);
-            register_object_store_and_config_extensions(
-                &ctx,
-                &cmd.location,
-                &cmd.options,
-                format,
-                false,
-            )
-            .await?;
+            for location in &cmd.locations {
+                register_object_store_and_config_extensions(
+                    &ctx,
+                    location,
+                    &cmd.options,
+                    format.clone(),
+                    false,
+                )
+                .await?;
+            }
         } else {
             return plan_err!("LogicalPlan is not a CreateExternalTable");
         }

datafusion-examples/examples/dataframe/cache_factory.rs

@@ -29,6 +29,7 @@ use datafusion::error::Result;
 use datafusion::execution::context::QueryPlanner;
 use datafusion::execution::session_state::CacheFactory;
 use datafusion::execution::{SessionState, SessionStateBuilder};
+use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion::logical_expr::{
     Extension, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
 };
@@ -146,6 +147,7 @@ impl ExtensionPlanner for CacheNodePlanner {
         logical_inputs: &[&LogicalPlan],
         physical_inputs: &[Arc<dyn ExecutionPlan>],
         session_state: &SessionState,
+        _planning_ctx: &PhysicalPlanningContext,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         if let Some(cache_node) = node.as_any().downcast_ref::<CacheNode>() {
             assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs");

datafusion-examples/examples/query_planning/expr_api.rs

@@ -33,6 +33,7 @@ use datafusion::functions_aggregate::first_last::first_value_udaf;
 use datafusion::logical_expr::execution_props::ExecutionProps;
 use datafusion::logical_expr::expr::BinaryExpr;
 use datafusion::logical_expr::interval_arithmetic::Interval;
+use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion::logical_expr::simplify::SimplifyContext;
 use datafusion::logical_expr::{ColumnarValue, ExprFunctionExt, ExprSchemable, Operator};
 use datafusion::optimizer::analyzer::type_coercion::TypeCoercionRewriter;
@@ -541,8 +542,12 @@ fn type_coercion_demo() -> Result<()> {
 
     // Evaluation with an expression that has not been type coerced cannot succeed.
     let props = ExecutionProps::default();
-    let physical_expr =
-        datafusion::physical_expr::create_physical_expr(&expr, &df_schema, &props)?;
+    let physical_expr = datafusion::physical_expr::create_physical_expr(
+        &expr,
+        &df_schema,
+        &props,
+        &PhysicalPlanningContext::default(),
+    )?;
     let e = physical_expr.evaluate(&batch).unwrap_err();
     assert!(
         e.find_root()
@@ -566,6 +571,7 @@ fn type_coercion_demo() -> Result<()> {
         &coerced_expr,
         &df_schema,
         &props,
+        &PhysicalPlanningContext::default(),
     )?;
     assert!(physical_expr.evaluate(&batch).is_ok());
 
@@ -578,6 +584,7 @@ fn type_coercion_demo() -> Result<()> {
         &coerced_expr,
         &df_schema,
         &props,
+        &PhysicalPlanningContext::default(),
     )?;
     assert!(physical_expr.evaluate(&batch).is_ok());
 
@@ -606,6 +613,7 @@ fn type_coercion_demo() -> Result<()> {
         &coerced_expr,
         &df_schema,
         &props,
+        &PhysicalPlanningContext::default(),
     )?;
     assert!(physical_expr.evaluate(&batch).is_ok());
 

datafusion-examples/examples/query_planning/pruning.rs

@@ -26,6 +26,7 @@ use datafusion::common::pruning::PruningStatistics;
 use datafusion::common::{DFSchema, ScalarValue};
 use datafusion::error::Result;
 use datafusion::execution::context::ExecutionProps;
+use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion::physical_expr::create_physical_expr;
 use datafusion::physical_optimizer::pruning::PruningPredicate;
 use datafusion::prelude::*;
@@ -194,7 +195,13 @@ impl PruningStatistics for MyCatalog {
 fn create_pruning_predicate(expr: Expr, schema: &SchemaRef) -> PruningPredicate {
     let df_schema = DFSchema::try_from(Arc::clone(schema)).unwrap();
     let props = ExecutionProps::new();
-    let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap();
+    let physical_expr = create_physical_expr(
+        &expr,
+        &df_schema,
+        &props,
+        &PhysicalPlanningContext::default(),
+    )
+    .unwrap();
     PruningPredicate::try_new(physical_expr, Arc::clone(schema)).unwrap()
 }
 

datafusion-examples/examples/relation_planner/table_sample.rs

@@ -119,6 +119,7 @@ use datafusion_common::{
     DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err,
     plan_datafusion_err, plan_err,
 };
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{
     UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
     logical_plan::{Extension, LogicalPlan, LogicalPlanBuilder},
@@ -587,6 +588,7 @@ impl ExtensionPlanner for TableSampleExtensionPlanner {
         _logical_inputs: &[&LogicalPlan],
         physical_inputs: &[Arc<dyn ExecutionPlan>],
         _session_state: &SessionState,
+        _planning_ctx: &PhysicalPlanningContext,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         let Some(sample_node) = node.as_any().downcast_ref::<TableSamplePlanNode>()
         else {

datafusion/catalog-listing/src/helpers.rs

@@ -34,6 +34,7 @@ use arrow::{
     record_batch::RecordBatch,
 };
 use datafusion_expr::execution_props::ExecutionProps;
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use futures::stream::FuturesUnordered;
 use futures::{StreamExt, TryStreamExt, stream::BoxStream};
 use log::{debug, trace};
@@ -328,7 +329,12 @@ pub fn filter_partitioned_file(
 
     let filter = utils::conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true));
     let props = ExecutionProps::new();
-    let expr = create_physical_expr(&filter, df_schema, &props)?;
+    let expr = create_physical_expr(
+        &filter,
+        df_schema,
+        &props,
+        &PhysicalPlanningContext::default(),
+    )?;
 
     // Since we're only operating on a single file, our batch and resulting "array" holds only one
     // value indicating if the input file matches the provided filters

datafusion/catalog-listing/src/table.rs

@@ -44,6 +44,7 @@ use datafusion_execution::cache::cache_manager::{
 };
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::execution_props::ExecutionProps;
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{
     Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType,
 };
@@ -54,7 +55,7 @@ use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::empty::EmptyExec;
 use futures::{Stream, StreamExt, TryStreamExt, future, stream};
 use object_store::ObjectStore;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 /// Result of a file listing operation from [`ListingTable::list_files_for_scan`].
@@ -614,6 +615,7 @@ impl TableProvider for ListingTable {
                         output_partitioning,
                         &df_schema,
                         state.execution_props(),
+                        &PhysicalPlanningContext::default(),
                     )?
                 }
             };
@@ -816,7 +818,15 @@ impl ListingTable {
         .await?;
         let meta_fetch_concurrency =
             ctx.config_options().execution.meta_fetch_concurrency.get();
-        let file_list = stream::iter(file_list).flatten_unordered(meta_fetch_concurrency);
+        // Table paths can overlap, for example when one path is a directory and
+        // another names a file inside it. A ListingTable uses one object store,
+        // so the object path uniquely identifies a file within this scan.
+        let mut seen_files = HashSet::new();
+        let file_list = stream::iter(file_list)
+            .flatten_unordered(meta_fetch_concurrency)
+            .try_filter(move |file| {
+                future::ready(seen_files.insert(file.object_meta.location.clone()))
+            });
         // collect the statistics and ordering if required by the config
         let files = file_list
             .map(|part_file| async {

datafusion/catalog/src/memory/table.rs

@@ -36,6 +36,7 @@ use datafusion_datasource::memory::{MemSink, MemorySourceConfig};
 use datafusion_datasource::sink::DataSinkExec;
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_expr::dml::InsertOp;
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{Expr, SortExpr, TableType};
 use datafusion_physical_expr::{
     LexOrdering, PhysicalExpr, create_physical_expr, create_physical_sort_exprs,
@@ -209,8 +210,12 @@ impl TableProvider for MemTable {
             let eqp = state.execution_props();
             let mut file_sort_order = vec![];
             for sort_exprs in sort_order.iter() {
-                let physical_exprs =
-                    create_physical_sort_exprs(sort_exprs, &df_schema, eqp)?;
+                let physical_exprs = create_physical_sort_exprs(
+                    sort_exprs,
+                    &df_schema,
+                    eqp,
+                    &PhysicalPlanningContext::default(),
+                )?;
                 file_sort_order.extend(LexOrdering::new(physical_exprs));
             }
             source = source.try_with_sort_information(file_sort_order)?;
@@ -356,8 +361,12 @@ impl TableProvider for MemTable {
         let physical_assignments: HashMap<String, Arc<dyn PhysicalExpr>> = assignments
             .iter()
             .map(|(name, expr)| {
-                let physical_expr =
-                    create_physical_expr(expr, &df_schema, state.execution_props())?;
+                let physical_expr = create_physical_expr(
+                    expr,
+                    &df_schema,
+                    state.execution_props(),
+                    &PhysicalPlanningContext::default(),
+                )?;
                 Ok((name.clone(), physical_expr))
             })
             .collect::<Result<_>>()?;
@@ -470,8 +479,12 @@ fn evaluate_filters_to_mask(
     let mut combined_mask: Option<BooleanArray> = None;
 
     for filter_expr in filters {
-        let physical_expr =
-            create_physical_expr(filter_expr, df_schema, execution_props)?;
+        let physical_expr = create_physical_expr(
+            filter_expr,
+            df_schema,
+            execution_props,
+            &PhysicalPlanningContext::default(),
+        )?;
 
         let result = physical_expr.evaluate(batch)?;
         let array = result.into_array(batch.num_rows())?;

datafusion/catalog/src/stream.rs

@@ -53,7 +53,15 @@ impl TableProviderFactory for StreamTableFactory {
         cmd: &CreateExternalTable,
     ) -> Result<Arc<dyn TableProvider>> {
         let schema: SchemaRef = Arc::clone(cmd.schema.inner());
-        let location = cmd.location.clone();
+        let location = match cmd.locations.as_slice() {
+            [single] => single.clone(),
+            _ => {
+                return config_err!(
+                    "Stream tables support exactly one location; \
+                     use a listing table to read multiple files"
+                );
+            }
+        };
         let encoding = cmd.file_type.parse()?;
         let header = if let Ok(opt) = cmd
             .options

datafusion/catalog/src/streaming.rs

@@ -22,6 +22,7 @@ use std::sync::Arc;
 use arrow::datatypes::SchemaRef;
 use async_trait::async_trait;
 use datafusion_common::{DFSchema, Result, plan_err};
+use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
 use datafusion_expr::{Expr, SortExpr, TableType};
 use datafusion_physical_expr::equivalence::project_ordering;
 use datafusion_physical_expr::projection::ProjectionMapping;
@@ -131,8 +132,12 @@ impl TableProvider for StreamingTable {
             let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?;
             let eqp = state.execution_props();
 
-            let original_sort_exprs =
-                create_physical_sort_exprs(&self.sort_order, &df_schema, eqp)?;
+            let original_sort_exprs = create_physical_sort_exprs(
+                &self.sort_order,
+                &df_schema,
+                eqp,
+                &PhysicalPlanningContext::default(),
+            )?;
 
             if let Some(p) = projection {
                 // When performing a projection, the output columns will not match