feat: add hidden datafusion bench bridge

Signed-off-by: Ethan Urbanski <ethan@urbanskitech.com>

Author: ethan-tyler

crates/core/src/delta_datafusion/bench_support.rs

@@ -0,0 +1,91 @@
+#![doc(hidden)]
+
+//! Diagnostic only, unstable bridge for file selection internals.
+//! This is not a supported public API.
+
+use datafusion::catalog::Session;
+use datafusion::common::{HashSet, Result as DataFusionResult};
+use datafusion::datasource::MemTable;
+use datafusion::logical_expr::{Expr, LogicalPlan};
+
+use crate::errors::DeltaResult;
+use crate::kernel::{Add, EagerSnapshot};
+use crate::logstore::LogStoreRef;
+
+use super::DeltaSessionExt as _;
+
+/// Wrapper result for the internal file-selection entrypoint.
+pub struct FindFilesResult {
+    pub candidates: Vec<Add>,
+    pub partition_scan: bool,
+}
+
+impl From<super::FindFiles> for FindFilesResult {
+    fn from(result: super::FindFiles) -> Self {
+        Self {
+            candidates: result.candidates,
+            partition_scan: result.partition_scan,
+        }
+    }
+}
+
+/// Minimal wrapper around the internal matched-files scan plan.
+pub struct MatchedFilesScan(super::MatchedFilesScan);
+
+impl MatchedFilesScan {
+    pub fn scan(&self) -> &LogicalPlan {
+        self.0.scan()
+    }
+
+    pub fn files_set(&self) -> HashSet<String> {
+        self.0.files_set()
+    }
+
+    pub fn predicate(&self) -> &Expr {
+        &self.0.predicate
+    }
+}
+
+fn prepare_session(session: &dyn Session, log_store: &LogStoreRef) -> DeltaResult<()> {
+    super::update_datafusion_session(session, log_store.as_ref(), None)?;
+    session.ensure_log_store_registered(log_store.as_ref())?;
+    Ok(())
+}
+
+pub async fn find_files(
+    snapshot: &EagerSnapshot,
+    log_store: LogStoreRef,
+    session: &dyn Session,
+    predicate: Option<Expr>,
+) -> DeltaResult<FindFilesResult> {
+    prepare_session(session, &log_store)?;
+    super::find_files(snapshot, log_store, session, predicate)
+        .await
+        .map(Into::into)
+}
+
+pub async fn find_files_scan(
+    snapshot: &EagerSnapshot,
+    log_store: LogStoreRef,
+    session: &dyn Session,
+    predicate: Expr,
+) -> DeltaResult<Vec<Add>> {
+    prepare_session(session, &log_store)?;
+    super::find_files_scan(snapshot, log_store, session, predicate).await
+}
+
+pub async fn scan_files_where_matches(
+    session: &dyn Session,
+    snapshot: &EagerSnapshot,
+    log_store: LogStoreRef,
+    predicate: Expr,
+) -> DataFusionResult<Option<MatchedFilesScan>> {
+    prepare_session(session, &log_store)?;
+    super::scan_files_where_matches(session, snapshot, predicate)
+        .await
+        .map(|scan| scan.map(MatchedFilesScan))
+}
+
+pub fn add_actions_partition_mem_table(snapshot: &EagerSnapshot) -> DeltaResult<Option<MemTable>> {
+    super::add_actions_partition_mem_table(snapshot)
+}

crates/core/src/delta_datafusion/find_files.rs

@@ -262,7 +262,7 @@ fn join_batches_with_add_actions(
         matching_files = field::Empty
     )
 )]
-async fn find_files_scan(
+pub(in crate::delta_datafusion) async fn find_files_scan(
     snapshot: &EagerSnapshot,
     log_store: LogStoreRef,
     session: &dyn Session,

crates/core/src/delta_datafusion/mod.rs

@@ -83,6 +83,8 @@ pub(crate) use table_provider::{
 
 pub(crate) const PATH_COLUMN: &str = "__delta_rs_path";
 
+#[doc(hidden)]
+pub mod bench_support;
 pub mod cdf;
 mod data_validation;
 pub mod engine;

crates/core/tests/file_selection_bench_bridge.rs

@@ -0,0 +1,124 @@
+#![cfg(feature = "datafusion")]
+
+use std::sync::Arc;
+
+use arrow_array::{Int32Array, RecordBatch, StringArray};
+use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema};
+use datafusion::datasource::TableProvider;
+use datafusion::physical_plan::collect;
+use datafusion::prelude::{SessionContext, col, lit};
+use deltalake_core::DeltaTable;
+use deltalake_core::delta_datafusion::bench_support;
+use deltalake_core::kernel::{DataType, PrimitiveType, StructField};
+use deltalake_core::protocol::SaveMode;
+use deltalake_test::TestResult;
+
+fn batch(ids: Vec<i32>, parts: Vec<&str>) -> TestResult<RecordBatch> {
+    Ok(RecordBatch::try_new(
+        Arc::new(ArrowSchema::new(vec![
+            ArrowField::new("id", ArrowDataType::Int32, false),
+            ArrowField::new("part", ArrowDataType::Utf8, false),
+        ])),
+        vec![
+            Arc::new(Int32Array::from(ids)),
+            Arc::new(StringArray::from(parts)),
+        ],
+    )?)
+}
+
+#[tokio::test]
+async fn test_out_of_crate_bridge_exposes_file_selection_paths() -> TestResult {
+    let table = DeltaTable::new_in_memory()
+        .create()
+        .with_columns(vec![
+            StructField::new(
+                "id".to_string(),
+                DataType::Primitive(PrimitiveType::Integer),
+                false,
+            ),
+            StructField::new(
+                "part".to_string(),
+                DataType::Primitive(PrimitiveType::String),
+                false,
+            ),
+        ])
+        .with_partition_columns(["part"])
+        .await?;
+    let table = table
+        .write(vec![batch(vec![1, 2], vec!["a", "a"])?])
+        .with_save_mode(SaveMode::Append)
+        .await?;
+    let table = table
+        .write(vec![batch(vec![100, 101], vec!["b", "b"])?])
+        .with_save_mode(SaveMode::Append)
+        .await?;
+
+    let snapshot = table.snapshot()?.snapshot().clone();
+    let log_store = table.log_store();
+    let session = SessionContext::new().state();
+
+    let partition_result = bench_support::find_files(
+        &snapshot,
+        log_store.clone(),
+        &session,
+        Some(col("part").eq(lit("a"))),
+    )
+    .await?;
+    assert!(partition_result.partition_scan);
+    assert_eq!(partition_result.candidates.len(), 1);
+    assert!(
+        partition_result.candidates[0].path.contains("part=a/"),
+        "expected partition-only path to match partition file, got {}",
+        partition_result.candidates[0].path
+    );
+
+    let mem_table = bench_support::add_actions_partition_mem_table(&snapshot)?
+        .expect("partition mem table should exist");
+    let field_names = mem_table
+        .schema()
+        .fields()
+        .iter()
+        .map(|field| field.name().clone())
+        .collect::<Vec<_>>();
+    assert_eq!(field_names, vec!["__delta_rs_path", "part"]);
+
+    let data_predicate = col("id").gt(lit(50i32));
+    let data_result = bench_support::find_files(
+        &snapshot,
+        log_store.clone(),
+        &session,
+        Some(data_predicate.clone()),
+    )
+    .await?;
+    assert!(!data_result.partition_scan);
+    assert_eq!(data_result.candidates.len(), 1);
+    assert!(
+        data_result.candidates[0].path.contains("part=b/"),
+        "expected data predicate to match partition b file, got {}",
+        data_result.candidates[0].path
+    );
+
+    let direct_scan = bench_support::find_files_scan(
+        &snapshot,
+        log_store.clone(),
+        &session,
+        data_predicate.clone(),
+    )
+    .await?;
+    assert_eq!(direct_scan.len(), 1);
+    assert_eq!(direct_scan[0].path, data_result.candidates[0].path);
+
+    let matched_scan =
+        bench_support::scan_files_where_matches(&session, &snapshot, log_store, data_predicate)
+            .await?
+            .expect("matched file scan should exist");
+    assert_eq!(matched_scan.files_set().len(), 1);
+    assert_eq!(matched_scan.predicate(), &col("id").gt(lit(50i32)));
+
+    let plan = session.create_physical_plan(matched_scan.scan()).await?;
+    let batches = collect(plan, session.task_ctx()).await?;
+    let row_count = batches.iter().map(|batch| batch.num_rows()).sum::<usize>();
+    assert_eq!(row_count, 2);
+
+    Ok(())
+}