test: consolidate ASOF SQL coverage

datafusion/core/tests/sql/joins.rs

@@ -17,10 +17,12 @@
 
 use insta::assert_snapshot;
 
-use datafusion::assert_batches_eq;
 use datafusion::catalog::MemTable;
 use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable};
+use datafusion::physical_plan::joins::AsOfJoinExec;
+use datafusion::physical_plan::{Distribution, ExecutionPlanProperties};
 use datafusion::test_util::register_unbounded_file_with_ordering;
+use datafusion::{assert_batches_eq, assert_batches_sorted_eq};
 use datafusion_sql::unparser::plan_to_sql;
 
 use super::*;
@@ -297,3 +299,312 @@ async fn unparse_cross_join() -> Result<()> {
 
     Ok(())
 }
+
+fn register_asof_test_tables(ctx: &SessionContext) -> Result<()> {
+    let trades_schema = Arc::new(Schema::new(vec![
+        Field::new("symbol", DataType::Utf8, true),
+        Field::new("ts", DataType::Int64, true),
+        Field::new("trade_id", DataType::Int32, false),
+    ]));
+    let trades = vec![
+        RecordBatch::try_new(
+            Arc::clone(&trades_schema),
+            vec![
+                Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])),
+                Arc::new(Int64Array::from(vec![Some(7), Some(2), Some(3)])),
+                Arc::new(Int32Array::from(vec![3, 4, 6])),
+            ],
+        )?,
+        RecordBatch::try_new(
+            Arc::clone(&trades_schema),
+            vec![
+                Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])),
+                Arc::new(Int64Array::from(vec![Some(1), Some(4), Some(8)])),
+                Arc::new(Int32Array::from(vec![1, 2, 5])),
+            ],
+        )?,
+    ];
+    ctx.register_table(
+        "trades",
+        Arc::new(MemTable::try_new(
+            trades_schema,
+            trades.into_iter().map(|batch| vec![batch]).collect(),
+        )?),
+    )?;
+
+    let prices_schema = Arc::new(Schema::new(vec![
+        Field::new("symbol", DataType::Utf8, true),
+        Field::new("ts", DataType::Int64, true),
+        Field::new("price", DataType::Int32, false),
+    ]));
+    let prices = vec![
+        RecordBatch::try_new(
+            Arc::clone(&prices_schema),
+            vec![
+                Arc::new(StringArray::from(vec![Some("A"), Some("B"), None])),
+                Arc::new(Int64Array::from(vec![Some(6), Some(1), Some(2)])),
+                Arc::new(Int32Array::from(vec![60, 101, 999])),
+            ],
+        )?,
+        RecordBatch::try_new(
+            Arc::clone(&prices_schema),
+            vec![
+                Arc::new(StringArray::from(vec![Some("A"), Some("A"), Some("B")])),
+                Arc::new(Int64Array::from(vec![Some(2), Some(4), Some(6)])),
+                Arc::new(Int32Array::from(vec![20, 40, 106])),
+            ],
+        )?,
+    ];
+    ctx.register_table(
+        "prices",
+        Arc::new(MemTable::try_new(
+            prices_schema,
+            prices.into_iter().map(|batch| vec![batch]).collect(),
+        )?),
+    )?;
+    Ok(())
+}
+
+fn find_asof_exec(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
+    if plan.downcast_ref::<AsOfJoinExec>().is_some() {
+        return Some(Arc::clone(plan));
+    }
+    plan.children().into_iter().find_map(find_asof_exec)
+}
+
+#[tokio::test]
+async fn asof_join_all_match_directions_across_batches() -> Result<()> {
+    let config = SessionConfig::new()
+        .with_batch_size(2)
+        .with_target_partitions(2);
+    let ctx = SessionContext::new_with_config(config);
+    register_asof_test_tables(&ctx)?;
+
+    for (op, expected) in [
+        (
+            ">=",
+            [
+                "+----------+-------+",
+                "| trade_id | price |",
+                "+----------+-------+",
+                "| 1        |       |",
+                "| 2        | 40    |",
+                "| 3        | 60    |",
+                "| 4        | 101   |",
+                "| 5        | 106   |",
+                "| 6        |       |",
+                "+----------+-------+",
+            ],
+        ),
+        (
+            ">",
+            [
+                "+----------+-------+",
+                "| trade_id | price |",
+                "+----------+-------+",
+                "| 1        |       |",
+                "| 2        | 20    |",
+                "| 3        | 60    |",
+                "| 4        | 101   |",
+                "| 5        | 106   |",
+                "| 6        |       |",
+                "+----------+-------+",
+            ],
+        ),
+        (
+            "<=",
+            [
+                "+----------+-------+",
+                "| trade_id | price |",
+                "+----------+-------+",
+                "| 1        | 20    |",
+                "| 2        | 40    |",
+                "| 3        |       |",
+                "| 4        | 106   |",
+                "| 5        |       |",
+                "| 6        |       |",
+                "+----------+-------+",
+            ],
+        ),
+        (
+            "<",
+            [
+                "+----------+-------+",
+                "| trade_id | price |",
+                "+----------+-------+",
+                "| 1        | 20    |",
+                "| 2        | 60    |",
+                "| 3        |       |",
+                "| 4        | 106   |",
+                "| 5        |       |",
+                "| 6        |       |",
+                "+----------+-------+",
+            ],
+        ),
+    ] {
+        let batches = ctx
+            .sql(&format!(
+                "SELECT t.trade_id, p.price FROM trades t \
+                 ASOF JOIN prices p MATCH_CONDITION (t.ts {op} p.ts) \
+                 ON t.symbol = p.symbol ORDER BY t.trade_id"
+            ))
+            .await?
+            .collect()
+            .await?;
+        assert_batches_eq!(expected, &batches);
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn asof_join_broadcasts_multi_partition_right_input() -> Result<()> {
+    let config = SessionConfig::new().with_target_partitions(4);
+    let ctx = SessionContext::new_with_config(config);
+    register_asof_test_tables(&ctx)?;
+    let df = ctx
+        .sql(
+            "SELECT t.trade_id, p.price FROM trades t ASOF JOIN \
+             (SELECT ts, price FROM prices WHERE symbol = 'A') p \
+             MATCH_CONDITION (t.ts >= p.ts)",
+        )
+        .await?;
+    let sql = plan_to_sql(df.logical_plan())?.to_string();
+    assert_contains!(sql.as_str(), "ASOF JOIN");
+    assert!(!sql.contains(" ON "), "unexpected equality clause: {sql}");
+    ctx.sql(&sql).await?;
+    let plan = df.create_physical_plan().await?;
+    let asof = find_asof_exec(&plan).expect("physical ASOF join must be present");
+    let output_partitions = asof.output_partitioning().partition_count();
+    assert_eq!(
+        output_partitions,
+        asof.children()[0].output_partitioning().partition_count()
+    );
+    assert!(
+        output_partitions > 1,
+        "ASOF join did not preserve left-side parallelism"
+    );
+    assert_eq!(
+        asof.children()[1].output_partitioning().partition_count(),
+        1
+    );
+    let right_plan = displayable(asof.children()[1].as_ref())
+        .indent(true)
+        .to_string();
+    assert_contains!(right_plan.as_str(), "SortPreservingMergeExec");
+    assert_contains!(right_plan.as_str(), "DataSourceExec: partitions=2");
+    assert!(asof.output_ordering().is_some());
+    assert!(matches!(
+        &asof.input_distribution_requirements().into_per_child()[..],
+        [
+            Distribution::UnspecifiedDistribution,
+            Distribution::SinglePartition
+        ]
+    ));
+    let batches = collect(plan, ctx.task_ctx()).await?;
+    assert_batches_sorted_eq!(
+        [
+            "+----------+-------+",
+            "| trade_id | price |",
+            "+----------+-------+",
+            "| 1        |       |",
+            "| 2        | 40    |",
+            "| 3        | 60    |",
+            "| 4        | 20    |",
+            "| 5        | 60    |",
+            "| 6        | 20    |",
+            "+----------+-------+",
+        ],
+        &batches
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn asof_join_rejects_unbounded_inputs_during_physical_planning() -> Result<()> {
+    let ctx = SessionContext::new();
+    let tmp_dir = TempDir::new()?;
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("symbol", DataType::UInt32, false),
+        Field::new("ts", DataType::UInt32, false),
+    ]));
+    let ordering = vec![vec![
+        col("symbol").sort(true, true),
+        col("ts").sort(true, true),
+    ]];
+    for table in ["left_stream", "right_stream"] {
+        let path = tmp_dir.path().join(format!("{table}.csv"));
+        File::create(&path)?;
+        register_unbounded_file_with_ordering(
+            &ctx,
+            Arc::clone(&schema),
+            &path,
+            table,
+            ordering.clone(),
+        )?;
+    }
+    let error = ctx
+        .sql(
+            "SELECT * FROM left_stream l ASOF JOIN right_stream r \
+             MATCH_CONDITION (l.ts >= r.ts) ON l.symbol = r.symbol",
+        )
+        .await?
+        .create_physical_plan()
+        .await
+        .expect_err("ASOF physical planning must reject unbounded inputs");
+    assert_contains!(error.to_string(), "AsOfJoinExec requires bounded inputs");
+    Ok(())
+}
+
+#[tokio::test]
+async fn asof_join_using_unparser_round_trips() -> Result<()> {
+    let ctx = SessionContext::new();
+    register_asof_test_tables(&ctx)?;
+    let df = ctx
+        .sql(
+            "SELECT * FROM trades t ASOF JOIN prices p \
+             MATCH_CONDITION (t.ts >= p.ts) USING (symbol)",
+        )
+        .await?;
+    let sql = plan_to_sql(df.logical_plan())?.to_string();
+    assert!(sql.contains("ASOF JOIN"));
+    assert!(sql.contains("MATCH_CONDITION"));
+    assert!(sql.contains("USING(symbol)"), "unexpected SQL: {sql}");
+    ctx.sql(&sql).await?;
+    Ok(())
+}
+
+#[tokio::test]
+async fn asof_join_unparser_preserves_right_preselection() -> Result<()> {
+    let ctx = SessionContext::new();
+    register_asof_test_tables(&ctx)?;
+    for query in [
+        "SELECT t.trade_id, p.price FROM trades t \
+         ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \
+         MATCH_CONDITION (t.ts >= p.ts) ON t.symbol = p.symbol \
+         ORDER BY t.trade_id",
+        "SELECT * FROM trades t \
+         ASOF JOIN (SELECT * FROM prices WHERE price < 100) p \
+         MATCH_CONDITION (t.ts >= p.ts) USING (symbol) \
+         ORDER BY t.trade_id",
+        "SELECT t.trade_id, p.price FROM trades t \
+         JOIN prices q ON t.symbol = q.symbol AND t.ts = q.ts \
+         ASOF JOIN prices p MATCH_CONDITION (t.ts >= p.ts) \
+         ON q.symbol = p.symbol ORDER BY t.trade_id",
+        "SELECT t.trade_id, q.trade_id FROM trades t \
+         ASOF JOIN (prices p JOIN trades q \
+         ON p.symbol = q.symbol AND p.ts = q.ts) \
+         MATCH_CONDITION (t.ts >= q.ts) ON t.symbol = p.symbol \
+         ORDER BY t.trade_id",
+    ] {
+        let expected = ctx.sql(query).await?.collect().await?;
+        let plan = ctx.sql(query).await?.into_optimized_plan()?;
+        let sql = plan_to_sql(&plan)?.to_string();
+        let actual = ctx.sql(&sql).await?.collect().await?;
+        assert_eq!(
+            datafusion_common::test_util::batches_to_string(&expected),
+            datafusion_common::test_util::batches_to_string(&actual),
+            "unparsed SQL changed ASOF candidate preselection: {sql}"
+        );
+    }
+    Ok(())
+}