fix: align ASOF logical contracts

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)