Merge branch 'xuanwo/asof-logical' into xuanwo/asof-proto

Author: Xuanwo

datafusion/core/src/physical_planner.rs

@@ -94,8 +94,8 @@ use datafusion_expr::physical_planning_context::{
 use datafusion_expr::utils::{expr_to_columns, split_conjunction};
 use datafusion_expr::{
     Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension,
-    FetchType, Filter, JoinConstraint, JoinType, Operator, RecursiveQuery, SkipType,
-    StringifiedPlan, WindowFrame, WindowFrameBound, WriteOp,
+    FetchType, Filter, JoinType, Operator, RecursiveQuery, SkipType, StringifiedPlan,
+    WindowFrame, WindowFrameBound, WriteOp,
 };
 use datafusion_physical_expr::aggregate::{
     AggregateFunctionExpr, LoweredAggregate, LoweredAggregateBuilder,
@@ -1903,22 +1903,8 @@ impl DefaultPhysicalPlanner {
                         planning_ctx,
                     )?,
                 );
-                let omitted_right = if join.join_constraint == JoinConstraint::Using {
-                    join.on
-                        .iter()
-                        .map(|(_, right)| {
-                            let column = right.get_as_join_column().ok_or_else(|| {
-                                internal_datafusion_err!("ASOF USING key is not a column")
-                            })?;
-                            join.right.schema().index_of_column(column)
-                        })
-                        .collect::<Result<HashSet<_>>>()?
-                } else {
-                    HashSet::new()
-                };
-                let right_output_indices = (0..join.right.schema().fields().len())
-                    .filter(|index| !omitted_right.contains(index))
-                    .collect();
+                let right_output_indices =
+                    (0..join.right.schema().fields().len()).collect();
                 Arc::new(AsOfJoinExec::try_new(
                     physical_left,
                     physical_right,

datafusion/expr/src/logical_plan/builder.rs

@@ -1840,38 +1840,10 @@ pub fn build_join_schema(
 
 /// Creates the schema for a left-preserving ASOF join.
 ///
-/// `ON` emits all left fields followed by nullable right fields. `USING` emits
-/// each equality key once by omitting the corresponding right field.
-pub fn build_asof_join_schema(
-    left: &DFSchema,
-    right: &DFSchema,
-    on: &[(Expr, Expr)],
-    join_constraint: JoinConstraint,
-) -> Result<DFSchema> {
-    let omitted_right_indices = if join_constraint == JoinConstraint::Using {
-        on.iter()
-            .map(|(_, right_expr)| {
-                let column = right_expr.get_as_join_column().ok_or_else(|| {
-                    plan_datafusion_err!("ASOF USING keys must be columns")
-                })?;
-                right.index_of_column(column)
-            })
-            .collect::<Result<HashSet<_>>>()?
-    } else {
-        HashSet::new()
-    };
-
-    let full_schema = build_join_schema(left, right, &JoinType::Left)?;
-    let left_len = left.fields().len();
-    let fields = full_schema
-        .iter()
-        .enumerate()
-        .filter(|(index, _)| {
-            *index < left_len || !omitted_right_indices.contains(&(*index - left_len))
-        })
-        .map(|(_, (qualifier, field))| (qualifier.cloned(), Arc::clone(field)))
-        .collect();
-    DFSchema::new_with_metadata(fields, full_schema.metadata().clone())?
+/// Both `ON` and `USING` preserve all qualified input fields. SQL wildcard
+/// expansion handles the unqualified `USING` key as a single column.
+pub fn build_asof_join_schema(left: &DFSchema, right: &DFSchema) -> Result<DFSchema> {
+    build_join_schema(left, right, &JoinType::Left)?
         .with_functional_dependencies(left.functional_dependencies().clone())
 }
 

datafusion/expr/src/logical_plan/plan.rs

@@ -240,9 +240,6 @@ pub enum LogicalPlan {
     /// Join two logical plans on one or more join columns.
     /// This is used to implement SQL `JOIN`
     Join(Join),
-    /// Match each left row with at most one ordered row from the right input.
-    /// This is used to implement SQL `ASOF JOIN`.
-    AsOfJoin(AsOfJoin),
     /// Repartitions the input based on a partitioning scheme. This is
     /// used to add parallelism and is sometimes referred to as an
     /// "exchange" operator in other systems
@@ -299,6 +296,9 @@ pub enum LogicalPlan {
     Unnest(Unnest),
     /// A variadic query (e.g. "Recursive CTEs")
     RecursiveQuery(RecursiveQuery),
+    /// Match each left row with at most one ordered row from the right input.
+    /// This is used to implement SQL `ASOF JOIN`.
+    AsOfJoin(AsOfJoin),
 }
 
 impl Default for LogicalPlan {
@@ -4455,8 +4455,7 @@ impl AsOfJoin {
             return plan_err!("ASOF USING keys must be columns");
         }
 
-        let schema =
-            build_asof_join_schema(left.schema(), right.schema(), &on, join_constraint)?;
+        let schema = build_asof_join_schema(left.schema(), right.schema())?;
         Ok(Self {
             left,
             right,
@@ -6867,6 +6866,32 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_asof_using_preserves_qualified_keys() -> Result<()> {
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("ts", DataType::Int64, false),
+        ]);
+        let left = Arc::new(table_scan(Some("t1"), &schema, None)?.build()?);
+        let right = Arc::new(table_scan(Some("t2"), &schema, None)?.build()?);
+        let join = AsOfJoin::try_new(
+            left,
+            right,
+            vec![(col("t1.id"), col("t2.id"))],
+            AsOfMatch::new(col("t1.ts"), Operator::GtEq, col("t2.ts")),
+            JoinConstraint::Using,
+        )?;
+
+        assert_eq!(join.schema.fields().len(), 4);
+        assert_eq!(
+            join.schema
+                .index_of_column(&Column::from_qualified_name("t2.id"))?,
+            2
+        );
+        assert!(join.schema.field(2).is_nullable());
+        Ok(())
+    }
+
     #[test]
     fn test_join_try_new_schema_validation() -> Result<()> {
         let left_schema = Schema::new(vec![

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -24,9 +24,8 @@ use crate::{OptimizerConfig, OptimizerRule};
 use std::sync::Arc;
 
 use datafusion_common::{
-    Column, DFSchema, HashMap, JoinConstraint, JoinType, Result,
-    assert_eq_or_internal_err, get_required_group_by_exprs_indices,
-    internal_datafusion_err, internal_err,
+    Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err,
+    get_required_group_by_exprs_indices, internal_datafusion_err, internal_err,
 };
 use datafusion_expr::expr::Alias;
 use datafusion_expr::{
@@ -410,31 +409,13 @@ fn optimize_projections(
         }
         LogicalPlan::AsOfJoin(join) => {
             let left_len = join.left.schema().fields().len();
-            let omitted_right = if join.join_constraint == JoinConstraint::Using {
-                join.on
-                    .iter()
-                    .map(|(_, right)| {
-                        let column = right.get_as_join_column().ok_or_else(|| {
-                            internal_datafusion_err!("ASOF USING key is not a column")
-                        })?;
-                        join.right.schema().index_of_column(column)
-                    })
-                    .collect::<Result<std::collections::HashSet<_>>>()?
-            } else {
-                std::collections::HashSet::new()
-            };
-            let right_output_indices = (0..join.right.schema().fields().len())
-                .filter(|index| !omitted_right.contains(index))
-                .collect::<Vec<_>>();
             let mut left_required = Vec::new();
             let mut right_required = Vec::new();
             for index in indices.indices() {
                 if *index < left_len {
                     left_required.push(*index);
-                } else if let Some(right_index) =
-                    right_output_indices.get(*index - left_len)
-                {
-                    right_required.push(*right_index);
+                } else {
+                    right_required.push(*index - left_len);
                 }
             }
             let left_indices = RequiredIndices::new_from_indices(left_required)

datafusion/physical-plan/src/joins/asof_join.rs

@@ -62,6 +62,7 @@ use arrow::compute::{SortOptions, interleave};
 use arrow::datatypes::{Schema, SchemaRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::stats::Precision;
+use datafusion_common::utils::memory::RecordBatchMemoryCounter;
 use datafusion_common::utils::{
     compare_rows, get_row_at_idx, normalize_float_zero_scalar,
 };
@@ -93,7 +94,6 @@ use crate::metrics::{
     BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder,
     MetricCategory, MetricsSet, RecordOutput, Time,
 };
-use crate::spill::get_record_batch_memory_size;
 use crate::statistics::{ChildStats, StatisticsArgs};
 use crate::stream::RecordBatchStreamAdapter;
 use crate::{
@@ -752,9 +752,10 @@ async fn collect_right_input(
     metrics: AsOfJoinMetrics,
 ) -> Result<BroadcastRightInput> {
     let schema = input.schema();
+    let mut memory_counter = RecordBatchMemoryCounter::new();
     let batches = input
         .try_fold(Vec::new(), |mut batches, batch| {
-            let batch_size = get_record_batch_memory_size(&batch);
+            let batch_size = memory_counter.count_batch(&batch);
             futures::future::ready(reservation.try_grow(batch_size).map(|_| {
                 metrics.build_mem_used.add(batch_size);
                 metrics.build_input_batches.add(1);
@@ -1213,6 +1214,7 @@ mod tests {
     };
     use arrow::datatypes::{DataType, Field, Int8Type};
     use datafusion_execution::config::SessionConfig;
+    use datafusion_execution::runtime_env::RuntimeEnvBuilder;
     use datafusion_expr::ColumnarValue;
     use datafusion_physical_expr_common::metrics::MetricValue;
     use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
@@ -1544,6 +1546,90 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn shared_right_buffers_are_reserved_once() -> Result<()> {
+        let left_schema = Arc::new(Schema::new(vec![
+            Field::new("key", DataType::Utf8, false),
+            Field::new("ts", DataType::Int64, false),
+            Field::new("id", DataType::Int32, false),
+        ]));
+        let left = TestMemoryExec::try_new_exec(
+            &[vec![make_batch(
+                &left_schema,
+                vec![Some("A")],
+                vec![Some(4095)],
+                vec![0],
+            )?]],
+            Arc::clone(&left_schema),
+            None,
+        )?;
+
+        let right_schema = Arc::new(Schema::new(vec![
+            Field::new("key", DataType::Utf8, false),
+            Field::new("ts", DataType::Int64, false),
+            Field::new("price", DataType::Int32, false),
+        ]));
+        let row_count = 4096;
+        let parent = make_batch(
+            &right_schema,
+            vec![Some("A"); row_count],
+            (0..row_count).map(|value| Some(value as i64)).collect(),
+            (0..row_count as i32).collect(),
+        )?;
+        let mut memory_counter = RecordBatchMemoryCounter::new();
+        let retained_size = memory_counter.count_batch(&parent);
+        let right_batches = (0..16)
+            .map(|index| parent.slice(index * 256, 256))
+            .collect();
+        let right = TestMemoryExec::try_new_exec(
+            &[right_batches],
+            Arc::clone(&right_schema),
+            None,
+        )?;
+
+        let exec = Arc::new(AsOfJoinExec::try_new(
+            left,
+            right,
+            vec![(
+                Arc::new(PhysicalColumn::new("key", 0)),
+                Arc::new(PhysicalColumn::new("key", 0)),
+            )],
+            AsOfMatchExpr::new(
+                Arc::new(PhysicalColumn::new("ts", 1)),
+                Operator::GtEq,
+                Arc::new(PhysicalColumn::new("ts", 1)),
+            ),
+            vec![2],
+        )?);
+        let runtime = RuntimeEnvBuilder::new()
+            .with_memory_limit(retained_size, 1.0)
+            .build_arc()?;
+        let context = Arc::new(TaskContext::default().with_runtime(runtime));
+
+        let batches = collect(Arc::clone(&exec) as _, context).await?;
+        let prices = batches
+            .iter()
+            .flat_map(|batch| {
+                batch
+                    .column(3)
+                    .as_any()
+                    .downcast_ref::<Int32Array>()
+                    .unwrap()
+                    .iter()
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(prices, vec![Some(4095)]);
+
+        let metrics = exec.metrics().expect("ASOF metrics must be present");
+        assert_eq!(
+            metrics
+                .sum_by_name("build_mem_used")
+                .map(|value| value.as_usize()),
+            Some(retained_size)
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn preserves_dictionary_outputs_across_large_flush() -> Result<()> {
         let dictionary_type =