Repository: apache/datafusion (rust)

Issue:
`count` fails for FFI Table Providers
### Describe the bug


When using FFI Table Providers, we generate an error because the input schemas do not match for cases like `count` where the input schema is irrelevant. See the minimal reproducible example below.


### To Reproduce

```
use datafusion::arrow::{datatypes::SchemaRef, record_batch::RecordBatch};
use datafusion::{
    catalog::streaming::StreamingTable,
    error::Result,
    physical_plan::{
        stream::RecordBatchStreamAdapter, streaming::PartitionStream,
        SendableRecordBatchStream,
    },
    prelude::*,
};
use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
use std::sync::Arc;

fn create_partition() -> RecordBatch {
    datafusion::common::record_batch!(
        ("a", Int32, vec![1, 2, 3]),
        ("b", Float64, vec![Some(4.0), None, Some(5.0)]),
        ("c", Utf8, vec!["alpha", "beta", "gamma"])
    )
    .unwrap()
}

#[derive(Debug)]
struct CustomStream {
    schema: SchemaRef,
}

impl Default for CustomStream {
    fn default() -> Self {
        Self {
            schema: create_partition().schema(),
        }
    }
}

impl PartitionStream for CustomStream {
    fn schema(&self) -> &SchemaRef {
        &self.schema
    }

    fn execute(
        &self,
        _ctx: Arc<datafusion::execution::TaskContext>,
    ) -> SendableRecordBatchStream {
        let stream_iter = [0..3].into_iter().map(|_| Ok(create_partition()));
        let stream = futures::stream::iter(stream_iter);
        let adapter = RecordBatchStreamAdapter::new(Arc::clone(&self.schema), stream);
        Box::pin(adapter)
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let ctx = SessionContext::new();

    let partition_stream = Arc::new(CustomStream::default());

    let table = Arc::new(StreamingTable::try_new(
        create_partition().schema(),
        vec![partition_stream],
    )?);

    let ffi_table = FFI_TableProvider::new(table, false, None);
    let foreign_table: ForeignTableProvider = (&ffi_table).into();
    ctx.register_table("my-table", Arc::new(foreign_table))?;

    ctx.table("my-table").await?.show().await?;

    let num_rows = ctx.table("my-table").await?.count().await?;
    println!("Found {num_rows} rows.");

    Ok(())
}
```

The table does work, but count fails.

```
+---+-----+-------+
| a | b   | c     |
+---+-----+-------+
| 1 | 4.0 | alpha |
| 2 |     | beta  |
| 3 | 5.0 | gamma |
+---+-----+-------+
Error: Internal("Physical input schema should be the same as the one converted from logical input schema. Differences: \n\t- Different number of fields: (physical) 3 vs (logical) 0")
```

### Expected behavior

These should work the same as when not using FFI.

### Additional context

_No response_

Pull request:
fix: ffi aggregation
## Which issue does this PR close?

- Closes #15569.

## Rationale for this change

In ffi physical plan, if projection is empty, full schema will be returned

## What changes are included in this PR?

remove this logic. But I don't know are there any other scenarios that depend on this trikcy logic.

## Are these changes tested?

UT

## Are there any user-facing changes?

No


Patch:
diff --git a/benchmarks/README.md b/benchmarks/README.md
index 8acaa298b..86b2e1b3b 100644
--- a/benchmarks/README.md
+++ b/benchmarks/README.md
@@ -200,6 +200,16 @@ cargo run --release --bin tpch -- convert --input ./data --output /mnt/tpch-parq
 
 Or if you want to verify and run all the queries in the benchmark, you can just run `cargo test`.
 
+#### Sorted Conversion
+
+The TPCH tables generated by the dbgen utility are sorted by their first column (their primary key for most tables, the `l_orderkey` column for the `lineitem` table.)
+
+To preserve this sorted order information during conversion (useful for benchmarking execution on pre-sorted data) include the `--sort` flag:
+
+```bash
+cargo run --release --bin tpch -- convert --input ./data --output /mnt/tpch-sorted-parquet --format parquet --sort
+```
+
 ### Comparing results between runs
 
 Any `dfbench` execution with `-o <dir>` argument will produce a
@@ -445,20 +455,29 @@ Test performance of end-to-end sort SQL queries. (While the `Sort` benchmark foc
 
 Sort integration benchmark runs whole table sort queries on TPCH `lineitem` table, with different characteristics. For example, different number of sort keys, different sort key cardinality, different number of payload columns, etc.
 
+If the TPCH tables have been converted as sorted on their first column (see [Sorted Conversion](#sorted-conversion)), you can use the `--sorted` flag to indicate that the input data is pre-sorted, allowing DataFusion to leverage that order during query execution.
+
+Additionally, an optional `--limit` flag is available for the sort benchmark. When specified, this flag appends a `LIMIT n` clause to the SQL query, effectively converting the query into a TopK query. Combining the `--sorted` and `--limit` options enables benchmarking of TopK queries on pre-sorted inputs.
+
 See [`sort_tpch.rs`](src/sort_tpch.rs) for more details.
 
 ### Sort TPCH Benchmark Example Runs
 1. Run all queries with default setting:
 ```bash
- cargo run --release --bin dfbench -- sort-tpch -p '....../datafusion/benchmarks/data/tpch_sf1' -o '/tmp/sort_tpch.json'
+ cargo run --release --bin dfbench -- sort-tpch -p './datafusion/benchmarks/data/tpch_sf1' -o '/tmp/sort_tpch.json'
 ```
 
 2. Run a specific query:
 ```bash
- cargo run --release --bin dfbench -- sort-tpch -p '....../datafusion/benchmarks/data/tpch_sf1' -o '/tmp/sort_tpch.json' --query 2
+ cargo run --release --bin dfbench -- sort-tpch -p './datafusion/benchmarks/data/tpch_sf1' -o '/tmp/sort_tpch.json' --query 2
 ```
 
-3. Run all queries with `bench.sh` script:
+3. Run all queries as TopK queries on presorted data:
+```bash
+ cargo run --release --bin dfbench -- sort-tpch --sorted --limit 10 -p './datafusion/benchmarks/data/tpch_sf1' -o '/tmp/sort_tpch.json'
+```
+
+4. Run all queries with `bench.sh` script:
 ```bash
 ./bench.sh run sort_tpch
 ```
diff --git a/benchmarks/src/sort_tpch.rs b/benchmarks/src/sort_tpch.rs
index 956bb92b6..176234eca 100644
--- a/benchmarks/src/sort_tpch.rs
+++ b/benchmarks/src/sort_tpch.rs
@@ -63,6 +63,15 @@ pub struct RunOpt {
     /// Load the data into a MemTable before executing the query
     #[structopt(short = "m", long = "mem-table")]
     mem_table: bool,
+
+    /// Mark the first column of each table as sorted in ascending order.
+    /// The tables should have been created with the `--sort` option for this to have any effect.
+    #[structopt(short = "t", long = "sorted")]
+    sorted: bool,
+
+    /// Append a `LIMIT n` clause to the query
+    #[structopt(short = "l", long = "limit")]
+    limit: Option<usize>,
 }
 
 struct QueryResult {
@@ -163,7 +172,7 @@ impl RunOpt {
         r#"
         SELECT l_shipmode, l_comment, l_partkey
         FROM lineitem
-        ORDER BY l_shipmode;
+        ORDER BY l_shipmode
         "#,
     ];
 
@@ -212,9 +221,14 @@ impl RunOpt {
             let start = Instant::now();
 
             let query_idx = query_id - 1; // 1-indexed -> 0-indexed
-            let sql = Self::SORT_QUERIES[query_idx];
+            let base_sql = Self::SORT_QUERIES[query_idx].to_string();
+            let sql = if let Some(limit) = self.limit {
+                format!("{base_sql} LIMIT {limit}")
+            } else {
+                base_sql
+            };
 
-            let row_count = self.execute_query(&ctx, sql).await?;
+            let row_count = self.execute_query(&ctx, sql.as_str()).await?;
 
             let elapsed = start.elapsed(); //.as_secs_f64() * 1000.0;
             let ms = elapsed.as_secs_f64() * 1000.0;
@@ -315,8 +329,18 @@ impl RunOpt {
             .with_collect_stat(state.config().collect_statistics());
 
         let table_path = ListingTableUrl::parse(path)?;
-        let config = ListingTableConfig::new(table_path).with_listing_options(options);
-        let config = config.infer_schema(&state).await?;
+        let schema = options.infer_schema(&state, &table_path).await?;
+        let options = if self.sorted {
+            let key_column_name = schema.fields()[0].name();
+            options
+                .with_file_sort_order(vec![vec![col(key_column_name).sort(true, false)]])
+        } else {
+            options
+        };
+
+        let config = ListingTableConfig::new(table_path)
+            .with_listing_options(options)
+            .with_schema(schema);
 
         Ok(Arc::new(ListingTable::try_new(config)?))
     }
diff --git a/benchmarks/src/tpch/convert.rs b/benchmarks/src/tpch/convert.rs
index 7f391d930..5219e09cd 100644
--- a/benchmarks/src/tpch/convert.rs
+++ b/benchmarks/src/tpch/convert.rs
@@ -22,15 +22,14 @@ use std::path::{Path, PathBuf};
 
 use datafusion::common::not_impl_err;
 
+use super::get_tbl_tpch_table_schema;
+use super::TPCH_TABLES;
 use datafusion::error::Result;
 use datafusion::prelude::*;
 use parquet::basic::Compression;
 use parquet::file::properties::WriterProperties;
 use structopt::StructOpt;
 
-use super::get_tbl_tpch_table_schema;
-use super::TPCH_TABLES;
-
 /// Convert tpch .slt files to .parquet or .csv files
 #[derive(Debug, StructOpt)]
 pub struct ConvertOpt {
@@ -57,6 +56,10 @@ pub struct ConvertOpt {
     /// Batch size when reading CSV or Parquet files
     #[structopt(short = "s", long = "batch-size", default_value = "8192")]
     batch_size: usize,
+
+    /// Sort each table by its first column in ascending order.
+    #[structopt(short = "t", long = "sort")]
+    sort: bool,
 }
 
 impl ConvertOpt {
@@ -70,6 +73,7 @@ impl ConvertOpt {
         for table in TPCH_TABLES {
             let start = Instant::now();
             let schema = get_tbl_tpch_table_schema(table);
+            let key_column_name = schema.fields()[0].name();
 
             let input_path = format!("{input_path}/{table}.tbl");
             let options = CsvReadOptions::new()
@@ -77,6 +81,13 @@ impl ConvertOpt {
                 .has_header(false)
                 .delimiter(b'|')
                 .file_extension(".tbl");
+            let options = if self.sort {
+                // indicated that the file is already sorted by its first column to speed up the conversion
+                options
+                    .file_sort_order(vec![vec![col(key_column_name).sort(true, false)]])
+            } else {
+                options
+            };
 
             let config = SessionConfig::new().with_batch_size(self.batch_size);
             let ctx = SessionContext::new_with_config(config);
@@ -99,6 +110,11 @@ impl ConvertOpt {
             if partitions > 1 {
                 csv = csv.repartition(Partitioning::RoundRobinBatch(partitions))?
             }
+            let csv = if self.sort {
+                csv.sort_by(vec![col(key_column_name)])?
+            } else {
+                csv
+            };
 
             // create the physical plan
             let csv = csv.create_physical_plan().await?;
diff --git a/benchmarks/src/tpch/run.rs b/benchmarks/src/tpch/run.rs
index eb9db821d..752a5a1a6 100644
--- a/benchmarks/src/tpch/run.rs
+++ b/benchmarks/src/tpch/run.rs
@@ -90,6 +90,11 @@ pub struct RunOpt {
     /// True by default.
     #[structopt(short = "j", long = "prefer_hash_join", default_value = "true")]
     prefer_hash_join: BoolDefaultTrue,
+
+    /// Mark the first column of each table as sorted in ascending order.
+    /// The tables should have been created with the `--sort` option for this to have any effect.
+    #[structopt(short = "t", long = "sorted")]
+    sorted: bool,
 }
 
 const TPCH_QUERY_START_ID: usize = 1;
@@ -275,20 +280,28 @@ impl RunOpt {
                 }
             };
 
+        let table_path = ListingTableUrl::parse(path)?;
         let options = ListingOptions::new(format)
             .with_file_extension(extension)
             .with_target_partitions(target_partitions)
             .with_collect_stat(state.config().collect_statistics());
-
-        let table_path = ListingTableUrl::parse(path)?;
-        let config = ListingTableConfig::new(table_path).with_listing_options(options);
-
-        let config = match table_format {
-            "parquet" => config.infer_schema(&state).await?,
-            "tbl" => config.with_schema(Arc::new(get_tbl_tpch_table_schema(table))),
-            "csv" => config.with_schema(Arc::new(get_tpch_table_schema(table))),
+        let schema = match table_format {
+            "parquet" => options.infer_schema(&state, &table_path).await?,
+            "tbl" => Arc::new(get_tbl_tpch_table_schema(table)),
+            "csv" => Arc::new(get_tpch_table_schema(table)),
             _ => unreachable!(),
         };
+        let options = if self.sorted {
+            let key_column_name = schema.fields()[0].name();
+            options
+                .with_file_sort_order(vec![vec![col(key_column_name).sort(true, false)]])
+        } else {
+            options
+        };
+
+        let config = ListingTableConfig::new(table_path)
+            .with_listing_options(options)
+            .with_schema(schema);
 
         Ok(Arc::new(ListingTable::try_new(config)?))
     }
@@ -357,6 +370,7 @@ mod tests {
             output_path: None,
             disable_statistics: false,
             prefer_hash_join: true,
+            sorted: false,
         };
         opt.register_tables(&ctx).await?;
         let queries = get_query_sql(query)?;
@@ -393,6 +407,7 @@ mod tests {
             output_path: None,
             disable_statistics: false,
             prefer_hash_join: true,
+            sorted: false,
         };
         opt.register_tables(&ctx).await?;
         let queries = get_query_sql(query)?;
diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs
index 6049614f3..cfb1f49d4 100644
--- a/datafusion/core/src/datasource/listing/table.rs
+++ b/datafusion/core/src/datasource/listing/table.rs
@@ -875,15 +875,13 @@ impl TableProvider for ListingTable {
             filters.iter().cloned().partition(|filter| {
                 can_be_evaluted_for_partition_pruning(&table_partition_col_names, filter)
             });
-        // TODO (https://github.com/apache/datafusion/issues/11600) remove downcast_ref from here?
-        let session_state = state.as_any().downcast_ref::<SessionState>().unwrap();
 
         // We should not limit the number of partitioned files to scan if there are filters and limit
         // at the same time. This is because the limit should be applied after the filters are applied.
         let statistic_file_limit = if filters.is_empty() { limit } else { None };
 
         let (mut partitioned_file_lists, statistics) = self
-            .list_files_for_scan(session_state, &partition_filters, statistic_file_limit)
+            .list_files_for_scan(state, &partition_filters, statistic_file_limit)
             .await?;
 
         // if no files need to be read, return an `EmptyExec`
@@ -942,7 +940,7 @@ impl TableProvider for ListingTable {
         self.options
             .format
             .create_physical_plan(
-                session_state,
+                state,
                 FileScanConfigBuilder::new(
                     object_store_url,
                     Arc::clone(&self.file_schema),
@@ -1022,10 +1020,8 @@ impl TableProvider for ListingTable {
         // Get the object store for the table path.
         let store = state.runtime_env().object_store(table_path)?;
 
-        // TODO (https://github.com/apache/datafusion/issues/11600) remove downcast_ref from here?
-        let session_state = state.as_any().downcast_ref::<SessionState>().unwrap();
         let file_list_stream = pruned_partition_list(
-            session_state,
+            state,
             store.as_ref(),
             table_path,
             &[],
@@ -1073,7 +1069,7 @@ impl TableProvider for ListingTable {
 
         self.options()
             .format
-            .create_writer_physical_plan(input, session_state, config, order_requirements)
+            .create_writer_physical_plan(input, state, config, order_requirements)
             .await
     }
 
diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs
index a7391a850..c86c07eba 100644
--- a/datafusion/ffi/src/table_provider.rs
+++ b/datafusion/ffi/src/table_provider.rs
@@ -259,14 +259,10 @@ unsafe extern "C" fn scan_fn_wrapper(
         };
 
         let projections: Vec<_> = projections.into_iter().collect();
-        let maybe_projections = match projections.is_empty() {
-            true => None,
-            false => Some(&projections),
-        };
 
         let plan = rresult_return!(
             internal_provider
-                .scan(&ctx.state(), maybe_projections, &filters, limit.into())
+                .scan(&ctx.state(), Some(&projections), &filters, limit.into())
                 .await
         );
 
@@ -600,4 +596,49 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_aggregation() -> Result<()> {
+        use arrow::datatypes::Field;
+        use datafusion::arrow::{
+            array::Float32Array, datatypes::DataType, record_batch::RecordBatch,
+        };
+        use datafusion::common::assert_batches_eq;
+        use datafusion::datasource::MemTable;
+
+        let schema =
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)]));
+
+        // define data in two partitions
+        let batch1 = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))],
+        )?;
+
+        let ctx = SessionContext::new();
+
+        let provider = Arc::new(MemTable::try_new(schema, vec![vec![batch1]])?);
+
+        let ffi_provider = FFI_TableProvider::new(provider, true, None);
+
+        let foreign_table_provider: ForeignTableProvider = (&ffi_provider).into();
+
+        ctx.register_table("t", Arc::new(foreign_table_provider))?;
+
+        let result = ctx
+            .sql("SELECT COUNT(*) as cnt FROM t")
+            .await?
+            .collect()
+            .await?;
+        #[rustfmt::skip]
+        let expected = [
+            "+-----+",
+            "| cnt |",
+            "+-----+",
+            "| 3   |",
+            "+-----+"
+        ];
+        assert_batches_eq!(expected, &result);
+        Ok(())
+    }
 }
diff --git a/datafusion/physical-plan/src/joins/sort_merge_join.rs b/datafusion/physical-plan/src/joins/sort_merge_join.rs
index 716cff939..89f2e3c91 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join.rs
@@ -823,42 +823,65 @@ impl BufferedBatch {
 /// Sort-Merge join stream that consumes streamed and buffered data streams
 /// and produces joined output stream.
 struct SortMergeJoinStream {
-    /// Current state of the stream
-    pub state: SortMergeJoinState,
+    // ========================================================================
+    // PROPERTIES:
+    // These fields are initialized at the start and remain constant throughout
+    // the execution.
+    // ========================================================================
     /// Output schema
     pub schema: SchemaRef,
-    /// Sort options of join columns used to sort streamed and buffered data stream
-    pub sort_options: Vec<SortOptions>,
     /// null == null?
     pub null_equals_null: bool,
+    /// Sort options of join columns used to sort streamed and buffered data stream
+    pub sort_options: Vec<SortOptions>,
+    /// optional join filter
+    pub filter: Option<JoinFilter>,
+    /// How the join is performed
+    pub join_type: JoinType,
+    /// Target output batch size
+    pub batch_size: usize,
+
+    // ========================================================================
+    // STREAMED FIELDS:
+    // These fields manage the properties and state of the streamed input.
+    // ========================================================================
     /// Input schema of streamed
     pub streamed_schema: SchemaRef,
-    /// Input schema of buffered
-    pub buffered_schema: SchemaRef,
     /// Streamed data stream
     pub streamed: SendableRecordBatchStream,
-    /// Buffered data stream
-    pub buffered: SendableRecordBatchStream,
     /// Current processing record batch of streamed
     pub streamed_batch: StreamedBatch,
-    /// Current buffered data
-    pub buffered_data: BufferedData,
     /// (used in outer join) Is current streamed row joined at least once?
     pub streamed_joined: bool,
-    /// (used in outer join) Is current buffered batches joined at least once?
-    pub buffered_joined: bool,
     /// State of streamed
     pub streamed_state: StreamedState,
-    /// State of buffered
-    pub buffered_state: BufferedState,
-    /// The comparison result of current streamed row and buffered batches
-    pub current_ordering: Ordering,
     /// Join key columns of streamed
     pub on_streamed: Vec<PhysicalExprRef>,
+
+    // ========================================================================
+    // BUFFERED FIELDS:
+    // These fields manage the properties and state of the buffered input.
+    // ========================================================================
+    /// Input schema of buffered
+    pub buffered_schema: SchemaRef,
+    /// Buffered data stream
+    pub buffered: SendableRecordBatchStream,
+    /// Current buffered data
+    pub buffered_data: BufferedData,
+    /// (used in outer join) Is current buffered batches joined at least once?
+    pub buffered_joined: bool,
+    /// State of buffered
+    pub buffered_state: BufferedState,
     /// Join key columns of buffered
     pub on_buffered: Vec<PhysicalExprRef>,
-    /// optional join filter
-    pub filter: Option<JoinFilter>,
+
+    // ========================================================================
+    // MERGE JOIN STATES:
+    // These fields track the execution state of merge join and are updated
+    // during the execution.
+    // ========================================================================
+    /// Current state of the stream
+    pub state: SortMergeJoinState,
     /// Staging output array builders
     pub staging_output_record_batches: JoinedRecordBatches,
     /// Output buffer. Currently used by filtering as it requires double buffering
@@ -868,18 +891,21 @@ struct SortMergeJoinStream {
     /// Increased when we put rows into buffer and decreased after we actually output batches.
     /// Used to trigger output when sufficient rows are ready
     pub output_size: usize,
-    /// Target output batch size
-    pub batch_size: usize,
-    /// How the join is performed
-    pub join_type: JoinType,
+    /// The comparison result of current streamed row and buffered batches
+    pub current_ordering: Ordering,
+    /// Manages the process of spilling and reading back intermediate data
+    pub spill_manager: SpillManager,
+
+    // ========================================================================
+    // EXECUTION RESOURCES:
+    // Fields related to managing execution resources and monitoring performance.
+    // ========================================================================
     /// Metrics
     pub join_metrics: SortMergeJoinMetrics,
     /// Memory reservation
     pub reservation: MemoryReservation,
     /// Runtime env
     pub runtime_env: Arc<RuntimeEnv>,
-    /// Manages the process of spilling and reading back intermediate data
-    pub spill_manager: SpillManager,
     /// A unique number for each batch
     pub streamed_batch_counter: AtomicUsize,
 }


Test patch:
diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs
index edc1fe357..2804a1de0 100644
--- a/datafusion/sql/tests/sql_integration.rs
+++ b/datafusion/sql/tests/sql_integration.rs
@@ -3388,26 +3388,15 @@ fn ident_normalization_parser_options_ident_normalization() -> ParserOptions {
     }
 }
 
-fn prepare_stmt_quick_test(
-    sql: &str,
-    expected_plan: &str,
-    expected_data_types: &str,
-) -> LogicalPlan {
+fn generate_prepare_stmt_and_data_types(sql: &str) -> (LogicalPlan, String) {
     let plan = logical_plan(sql).unwrap();
-
-    let assert_plan = plan.clone();
-    // verify plan
-    assert_eq!(format!("{assert_plan}"), expected_plan);
-
-    // verify data types
-    if let LogicalPlan::Statement(Statement::Prepare(Prepare { data_types, .. })) =
-        assert_plan
-    {
-        let dt = format!("{data_types:?}");
-        assert_eq!(dt, expected_data_types);
-    }
-
-    plan
+    let data_types = match &plan {
+        LogicalPlan::Statement(Statement::Prepare(Prepare { data_types, .. })) => {
+            format!("{data_types:?}")
+        }
+        _ => panic!("Expected a Prepare statement"),
+    };
+    (plan, data_types)
 }
 
 #[test]
@@ -4383,8 +4372,12 @@ fn test_prepare_statement_to_plan_panic_prepare_wrong_syntax() {
 #[test]
 fn test_prepare_statement_to_plan_panic_no_relation_and_constant_param() {
     let sql = "PREPARE my_plan(INT) AS SELECT id + $1";
-    let expected = "Schema error: No field named id.";
-    assert_eq!(logical_plan(sql).unwrap_err().strip_backtrace(), expected);
+
+    let plan = logical_plan(sql).unwrap_err().strip_backtrace();
+    assert_snapshot!(
+        plan,
+        @r"Schema error: No field named id."
+    );
 }
 
 #[test]
@@ -4426,15 +4419,17 @@ fn test_prepare_statement_to_plan_panic_is_param() {
 fn test_prepare_statement_to_plan_no_param() {
     // no embedded parameter but still declare it
     let sql = "PREPARE my_plan(INT) AS SELECT id, age  FROM person WHERE age = 10";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int32] \
-        \n  Projection: person.id, person.age\
-        \n    Filter: person.age = Int64(10)\
-        \n      TableScan: person";
-
-    let expected_dt = "[Int32]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int32] 
+      Projection: person.id, person.age
+        Filter: person.age = Int64(10)
+          TableScan: person
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int32]"#);
 
     ///////////////////
     // replace params with values
@@ -4452,15 +4447,17 @@ fn test_prepare_statement_to_plan_no_param() {
     //////////////////////////////////////////
     // no embedded parameter and no declare it
     let sql = "PREPARE my_plan AS SELECT id, age  FROM person WHERE age = 10";
-
-    let expected_plan = "Prepare: \"my_plan\" [] \
-        \n  Projection: person.id, person.age\
-        \n    Filter: person.age = Int64(10)\
-        \n      TableScan: person";
-
-    let expected_dt = "[]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [] 
+      Projection: person.id, person.age
+        Filter: person.age = Int64(10)
+          TableScan: person
+    "#
+    );
+    assert_snapshot!(dt, @r#"[]"#);
 
     ///////////////////
     // replace params with values
@@ -4532,12 +4529,16 @@ fn test_prepare_statement_to_plan_no_param_on_value_panic() {
 #[test]
 fn test_prepare_statement_to_plan_params_as_constants() {
     let sql = "PREPARE my_plan(INT) AS SELECT $1";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int32] \
-        \n  Projection: $1\n    EmptyRelation";
-    let expected_dt = "[Int32]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int32] 
+      Projection: $1
+        EmptyRelation
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int32]"#);
 
     ///////////////////
     // replace params with values
@@ -4553,12 +4554,16 @@ fn test_prepare_statement_to_plan_params_as_constants() {
 
     ///////////////////////////////////////
     let sql = "PREPARE my_plan(INT) AS SELECT 1 + $1";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int32] \
-        \n  Projection: Int64(1) + $1\n    EmptyRelation";
-    let expected_dt = "[Int32]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int32] 
+      Projection: Int64(1) + $1
+        EmptyRelation
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int32]"#);
 
     ///////////////////
     // replace params with values
@@ -4574,12 +4579,16 @@ fn test_prepare_statement_to_plan_params_as_constants() {
 
     ///////////////////////////////////////
     let sql = "PREPARE my_plan(INT, DOUBLE) AS SELECT 1 + $1 + $2";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int32, Float64] \
-        \n  Projection: Int64(1) + $1 + $2\n    EmptyRelation";
-    let expected_dt = "[Int32, Float64]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int32, Float64] 
+      Projection: Int64(1) + $1 + $2
+        EmptyRelation
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int32, Float64]"#);
 
     ///////////////////
     // replace params with values
@@ -4598,20 +4607,20 @@ fn test_prepare_statement_to_plan_params_as_constants() {
 }
 
 #[test]
-fn test_prepare_statement_infer_types_from_join() {
+fn test_infer_types_from_join() {
     let sql =
         "SELECT id, order_id FROM person JOIN orders ON id = customer_id and age = $1";
 
-    let expected_plan = r#"
-Projection: person.id, orders.order_id
-  Inner Join:  Filter: person.id = orders.customer_id AND person.age = $1
-    TableScan: person
-    TableScan: orders
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: person.id, orders.order_id
+      Inner Join:  Filter: person.id = orders.customer_id AND person.age = $1
+        TableScan: person
+        TableScan: orders
     "#
-    .trim();
-
-    let expected_dt = "[Int32]";
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    );
 
     let actual_types = plan.get_parameter_types().unwrap();
     let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
@@ -4633,18 +4642,17 @@ Projection: person.id, orders.order_id
 }
 
 #[test]
-fn test_prepare_statement_infer_types_from_predicate() {
+fn test_infer_types_from_predicate() {
     let sql = "SELECT id, age FROM person WHERE age = $1";
-
-    let expected_plan = r#"
-Projection: person.id, person.age
-  Filter: person.age = $1
-    TableScan: person
-        "#
-    .trim();
-
-    let expected_dt = "[Int32]";
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: person.id, person.age
+      Filter: person.age = $1
+        TableScan: person
+    "#
+    );
 
     let actual_types = plan.get_parameter_types().unwrap();
     let expected_types = HashMap::from([("$1".to_string(), Some(DataType::Int32))]);
@@ -4665,18 +4673,18 @@ Projection: person.id, person.age
 }
 
 #[test]
-fn test_prepare_statement_infer_types_from_between_predicate() {
+fn test_infer_types_from_between_predicate() {
     let sql = "SELECT id, age FROM person WHERE age BETWEEN $1 AND $2";
 
-    let expected_plan = r#"
-Projection: person.id, person.age
-  Filter: person.age BETWEEN $1 AND $2
-    TableScan: person
-        "#
-    .trim();
-
-    let expected_dt = "[Int32]";
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: person.id, person.age
+      Filter: person.age BETWEEN $1 AND $2
+        TableScan: person
+    "#
+    );
 
     let actual_types = plan.get_parameter_types().unwrap();
     let expected_types = HashMap::from([
@@ -4700,23 +4708,23 @@ Projection: person.id, person.age
 }
 
 #[test]
-fn test_prepare_statement_infer_types_subquery() {
+fn test_infer_types_subquery() {
     let sql = "SELECT id, age FROM person WHERE age = (select max(age) from person where id = $1)";
 
-    let expected_plan = r#"
-Projection: person.id, person.age
-  Filter: person.age = (<subquery>)
-    Subquery:
-      Projection: max(person.age)
-        Aggregate: groupBy=[[]], aggr=[[max(person.age)]]
-          Filter: person.id = $1
-            TableScan: person
-    TableScan: person
-        "#
-    .trim();
-
-    let expected_dt = "[Int32]";
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Projection: person.id, person.age
+      Filter: person.age = (<subquery>)
+        Subquery:
+          Projection: max(person.age)
+            Aggregate: groupBy=[[]], aggr=[[max(person.age)]]
+              Filter: person.id = $1
+                TableScan: person
+        TableScan: person
+    "#
+    );
 
     let actual_types = plan.get_parameter_types().unwrap();
     let expected_types = HashMap::from([("$1".to_string(), Some(DataType::UInt32))]);
@@ -4742,19 +4750,19 @@ Projection: person.id, person.age
 }
 
 #[test]
-fn test_prepare_statement_update_infer() {
+fn test_update_infer() {
     let sql = "update person set age=$1 where id=$2";
 
-    let expected_plan = r#"
-Dml: op=[Update] table=[person]
-  Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, $1 AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
-    Filter: person.id = $2
-      TableScan: person
-        "#
-        .trim();
-
-    let expected_dt = "[Int32]";
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Dml: op=[Update] table=[person]
+      Projection: person.id AS id, person.first_name AS first_name, person.last_name AS last_name, $1 AS age, person.state AS state, person.salary AS salary, person.birth_date AS birth_date, person.😀 AS 😀
+        Filter: person.id = $2
+          TableScan: person
+    "#
+    );
 
     let actual_types = plan.get_parameter_types().unwrap();
     let expected_types = HashMap::from([
@@ -4779,17 +4787,17 @@ Dml: op=[Update] table=[person]
 }
 
 #[test]
-fn test_prepare_statement_insert_infer() {
+fn test_insert_infer() {
     let sql = "insert into person (id, first_name, last_name) values ($1, $2, $3)";
-
-    let expected_plan = "Dml: op=[Insert Into] table=[person]\
-                        \n  Projection: column1 AS id, column2 AS first_name, column3 AS last_name, \
-                                    CAST(NULL AS Int32) AS age, CAST(NULL AS Utf8) AS state, CAST(NULL AS Float64) AS salary, \
-                                    CAST(NULL AS Timestamp(Nanosecond, None)) AS birth_date, CAST(NULL AS Int32) AS 😀\
-                        \n    Values: ($1, $2, $3)";
-
-    let expected_dt = "[Int32]";
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @r#"
+    Dml: op=[Insert Into] table=[person]
+      Projection: column1 AS id, column2 AS first_name, column3 AS last_name, CAST(NULL AS Int32) AS age, CAST(NULL AS Utf8) AS state, CAST(NULL AS Float64) AS salary, CAST(NULL AS Timestamp(Nanosecond, None)) AS birth_date, CAST(NULL AS Int32) AS 😀
+        Values: ($1, $2, $3)
+    "#
+    );
 
     let actual_types = plan.get_parameter_types().unwrap();
     let expected_types = HashMap::from([
@@ -4819,15 +4827,17 @@ fn test_prepare_statement_insert_infer() {
 #[test]
 fn test_prepare_statement_to_plan_one_param() {
     let sql = "PREPARE my_plan(INT) AS SELECT id, age  FROM person WHERE age = $1";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int32] \
-        \n  Projection: person.id, person.age\
-        \n    Filter: person.age = $1\
-        \n      TableScan: person";
-
-    let expected_dt = "[Int32]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int32] 
+      Projection: person.id, person.age
+        Filter: person.age = $1
+          TableScan: person
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int32]"#);
 
     ///////////////////
     // replace params with values
@@ -4848,16 +4858,19 @@ fn test_prepare_statement_to_plan_one_param() {
 fn test_prepare_statement_to_plan_data_type() {
     let sql = "PREPARE my_plan(DOUBLE) AS SELECT id, age  FROM person WHERE age = $1";
 
-    // age is defined as Int32 but prepare statement declares it as DOUBLE/Float64
-    // Prepare statement and its logical plan should be created successfully
-    let expected_plan = "Prepare: \"my_plan\" [Float64] \
-        \n  Projection: person.id, person.age\
-        \n    Filter: person.age = $1\
-        \n      TableScan: person";
-
-    let expected_dt = "[Float64]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        // age is defined as Int32 but prepare statement declares it as DOUBLE/Float64
+        // Prepare statement and its logical plan should be created successfully
+        @r#"
+    Prepare: "my_plan" [Float64] 
+      Projection: person.id, person.age
+        Filter: person.age = $1
+          TableScan: person
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Float64]"#);
 
     ///////////////////
     // replace params with values still succeed and use Float64
@@ -4880,15 +4893,17 @@ fn test_prepare_statement_to_plan_multi_params() {
         SELECT id, age, $6
         FROM person
         WHERE age IN ($1, $4) AND salary > $3 and salary < $5 OR first_name < $2";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int32, Utf8, Float64, Int32, Float64, Utf8] \
-        \n  Projection: person.id, person.age, $6\
-        \n    Filter: person.age IN ([$1, $4]) AND person.salary > $3 AND person.salary < $5 OR person.first_name < $2\
-        \n      TableScan: person";
-
-    let expected_dt = "[Int32, Utf8, Float64, Int32, Float64, Utf8]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int32, Utf8, Float64, Int32, Float64, Utf8] 
+      Projection: person.id, person.age, $6
+        Filter: person.age IN ([$1, $4]) AND person.salary > $3 AND person.salary < $5 OR person.first_name < $2
+          TableScan: person
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int32, Utf8, Float64, Int32, Float64, Utf8]"#);
 
     ///////////////////
     // replace params with values
@@ -4921,17 +4936,19 @@ fn test_prepare_statement_to_plan_having() {
         GROUP BY id
         HAVING sum(age) < $1 AND sum(age) > 10 OR sum(age) in ($3, $4)\
         ";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int32, Float64, Float64, Float64] \
-        \n  Projection: person.id, sum(person.age)\
-        \n    Filter: sum(person.age) < $1 AND sum(person.age) > Int64(10) OR sum(person.age) IN ([$3, $4])\
-        \n      Aggregate: groupBy=[[person.id]], aggr=[[sum(person.age)]]\
-        \n        Filter: person.salary > $2\
-        \n          TableScan: person";
-
-    let expected_dt = "[Int32, Float64, Float64, Float64]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int32, Float64, Float64, Float64] 
+      Projection: person.id, sum(person.age)
+        Filter: sum(person.age) < $1 AND sum(person.age) > Int64(10) OR sum(person.age) IN ([$3, $4])
+          Aggregate: groupBy=[[person.id]], aggr=[[sum(person.age)]]
+            Filter: person.salary > $2
+              TableScan: person
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int32, Float64, Float64, Float64]"#);
 
     ///////////////////
     // replace params with values
@@ -4960,15 +4977,17 @@ fn test_prepare_statement_to_plan_limit() {
     let sql = "PREPARE my_plan(BIGINT, BIGINT) AS
         SELECT id FROM person \
         OFFSET $1 LIMIT $2";
-
-    let expected_plan = "Prepare: \"my_plan\" [Int64, Int64] \
-        \n  Limit: skip=$1, fetch=$2\
-        \n    Projection: person.id\
-        \n      TableScan: person";
-
-    let expected_dt = "[Int64, Int64]";
-
-    let plan = prepare_stmt_quick_test(sql, expected_plan, expected_dt);
+    let (plan, dt) = generate_prepare_stmt_and_data_types(sql);
+    assert_snapshot!(
+        plan,
+        @r#"
+    Prepare: "my_plan" [Int64, Int64] 
+      Limit: skip=$1, fetch=$2
+        Projection: person.id
+          TableScan: person
+    "#
+    );
+    assert_snapshot!(dt, @r#"[Int64, Int64]"#);
 
     // replace params with values
     let param_values = vec![ScalarValue::Int64(Some(10)), ScalarValue::Int64(Some(200))];
diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt
index ca86dbfcc..ddf701ba0 100644
--- a/datafusion/sqllogictest/test_files/joins.slt
+++ b/datafusion/sqllogictest/test_files/joins.slt
@@ -4742,3 +4742,51 @@ drop table person;
 
 statement count 0
 drop table orders;
+
+# Create tables for testing compound field access in JOIN conditions
+statement ok
+CREATE TABLE compound_field_table_t
+AS VALUES
+({r: 'a', c: 1}),
+({r: 'b', c: 2.3});
+
+statement ok
+CREATE TABLE compound_field_table_u
+AS VALUES
+({r: 'a', c: 1}),
+({r: 'b', c: 2.3});
+
+# Test compound field access in JOIN condition with table aliases
+query ??
+SELECT * FROM compound_field_table_t tee JOIN compound_field_table_u you ON tee.column1['r'] = you.column1['r']
+----
+{r: a, c: 1.0} {r: a, c: 1.0}
+{r: b, c: 2.3} {r: b, c: 2.3}
+
+# Test compound field access in JOIN condition without table aliases
+query ??
+SELECT * FROM compound_field_table_t JOIN compound_field_table_u ON compound_field_table_t.column1['r'] = compound_field_table_u.column1['r']
+----
+{r: a, c: 1.0} {r: a, c: 1.0}
+{r: b, c: 2.3} {r: b, c: 2.3}
+
+# Test compound field access with numeric field access
+query ??
+SELECT * FROM compound_field_table_t tee JOIN compound_field_table_u you ON tee.column1['c'] = you.column1['c']
+----
+{r: a, c: 1.0} {r: a, c: 1.0}
+{r: b, c: 2.3} {r: b, c: 2.3}
+
+# Test compound field access with mixed field types
+query ??
+SELECT * FROM compound_field_table_t tee JOIN compound_field_table_u you ON tee.column1['r'] = you.column1['r'] AND tee.column1['c'] = you.column1['c']
+----
+{r: a, c: 1.0} {r: a, c: 1.0}
+{r: b, c: 2.3} {r: b, c: 2.3}
+
+# Clean up compound field tables
+statement ok
+DROP TABLE compound_field_table_t;
+
+statement ok
+DROP TABLE compound_field_table_u;