test: preserve full join schema and trim unused probe dependencies

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

scripts/paper-plan-choice/Cargo.lock

@@ -609,7 +609,6 @@ dependencies = [
  "indexmap",
  "itertools 0.15.0",
  "serde_json",
- "sqlparser",
 ]
 
 [[package]]
@@ -908,9 +907,9 @@ dependencies = [
 
 [[package]]
 name = "futures"
-version = "0.3.31"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876"
+checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3"
 dependencies = [
  "futures-channel",
  "futures-core",
@@ -1503,11 +1502,9 @@ dependencies = [
  "datafusion-common",
  "datafusion-datasource",
  "datafusion-execution",
- "datafusion-expr",
  "datafusion-physical-expr",
  "datafusion-physical-optimizer",
  "datafusion-physical-plan",
- "futures",
  "tokio",
 ]
 
@@ -1850,27 +1847,6 @@ version = "1.16.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9395f0f0eee849a9b707b2f06bb92a6a422090e2123bb2ef8e87a0e61892a8e"
 
-[[package]]
-name = "sqlparser"
-version = "0.62.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec"
-dependencies = [
- "log",
- "sqlparser_derive",
-]
-
-[[package]]
-name = "sqlparser_derive"
-version = "0.5.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.119",
-]
-
 [[package]]
 name = "stable_deref_trait"
 version = "1.2.1"

scripts/paper-plan-choice/Cargo.toml

@@ -12,9 +12,7 @@ arrow-schema = "=59.2.0"
 datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" }
 datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" }
 datafusion-execution = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" }
-datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" }
 datafusion-physical-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" }
 datafusion-physical-optimizer = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" }
 datafusion-physical-plan = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "4bd567a884dcbfd89d0f0a7859ef8367e38fff97" }
-futures = "=0.3.31"
 tokio = { version = "=1.52.3", features = ["macros", "rt-multi-thread"] }

scripts/paper-plan-choice/src/main.rs

@@ -213,19 +213,11 @@ fn fact_build(plan: &Arc<dyn ExecutionPlan>) -> bool {
 
 type Row = (Option<i64>, Option<i64>, Option<i64>, Option<i64>);
 
-async fn output(plan: Arc<dyn ExecutionPlan>) -> Result<Vec<Row>> {
+async fn output(plan: Arc<dyn ExecutionPlan>, expected_schema: &SchemaRef) -> Result<Vec<Row>> {
     let batches = collect(plan, Arc::new(TaskContext::default())).await?;
     let mut rows = Vec::new();
     for batch in batches {
-        assert_eq!(
-            batch
-                .schema()
-                .fields()
-                .iter()
-                .map(|f| f.name().as_str())
-                .collect::<Vec<_>>(),
-            ["fact_key", "fact_value", "dim_key", "dim_weight"]
-        );
+        assert_eq!(batch.schema().as_ref(), expected_schema.as_ref());
         for i in 0..batch.num_rows() {
             let value = |column: usize| -> Option<i64> {
                 let array = batch
@@ -296,7 +288,7 @@ fn source_stats(
         sources.extend(source_stats(child.as_ref(), registry)?);
     }
     if plan.downcast_ref::<DataSourceExec>().is_some() {
-        sources.push((**registry.compute(plan)?.base_arc()).clone());
+        sources.push(registry.compute_base(plan)?);
     }
     Ok(sources)
 }
@@ -308,6 +300,7 @@ async fn check_case(fact_first: bool, facts: usize, dims: usize) -> Result<()> {
     config.optimizer.hash_join_single_partition_threshold = 0;
     config.optimizer.hash_join_single_partition_threshold_rows = 0;
     let original = join(fact_first, facts, dims)?;
+    let expected_schema = original.schema();
     let expected = oracle(facts, dims);
     for stats_kind in [RowStats::Absent, RowStats::Exact] {
         let registry = StatisticsRegistry::with_providers(vec![Arc::new(FixtureStatistics {
@@ -333,6 +326,7 @@ async fn check_case(fact_first: bool, facts: usize, dims: usize) -> Result<()> {
             hash_join(plan.as_ref()).unwrap().mode,
             PartitionMode::Partitioned
         );
+        assert_eq!(plan.schema().as_ref(), expected_schema.as_ref());
         let build_is_fact = fact_build(&plan);
         assert_eq!(
             build_is_fact,
@@ -341,7 +335,7 @@ async fn check_case(fact_first: bool, facts: usize, dims: usize) -> Result<()> {
                 RowStats::Exact => facts <= dims,
             }
         );
-        let result = output(plan).await?;
+        let result = output(plan, &expected_schema).await?;
         assert_eq!(result, expected);
         println!(
             "fact={facts} dim={dims} orientation={} stats={stats_kind:?} build={} rows={}",