feat: arrow convenience extensions

Author: roeap

kernel/examples/read-table-single-threaded/src/main.rs

@@ -2,10 +2,8 @@ use std::collections::HashMap;
 use std::process::ExitCode;
 use std::sync::Arc;
 
-use arrow::compute::filter_record_batch;
-use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::print_batches;
-use delta_kernel::engine::arrow_data::ArrowEngineData;
+use delta_kernel::engine::arrow_extensions::ScanExt;
 use delta_kernel::engine::default::executor::tokio::TokioBackgroundExecutor;
 use delta_kernel::engine::default::DefaultEngine;
 use delta_kernel::engine::sync::SyncEngine;
@@ -119,24 +117,7 @@ fn try_main() -> DeltaResult<()> {
         .with_schema_opt(read_schema_opt)
         .build()?;
 
-    let batches: Vec<RecordBatch> = scan
-        .execute(engine)?
-        .map(|scan_result| -> DeltaResult<_> {
-            let scan_result = scan_result?;
-            let mask = scan_result.full_mask();
-            let data = scan_result.raw_data?;
-            let record_batch: RecordBatch = data
-                .into_any()
-                .downcast::<ArrowEngineData>()
-                .map_err(|_| delta_kernel::Error::EngineDataType("ArrowEngineData".to_string()))?
-                .into();
-            if let Some(mask) = mask {
-                Ok(filter_record_batch(&record_batch, &mask.into())?)
-            } else {
-                Ok(record_batch)
-            }
-        })
-        .try_collect()?;
+    let batches: Vec<_> = scan.execute_arrow(engine)?.try_collect()?;
     print_batches(&batches)?;
     Ok(())
 }

kernel/src/engine/arrow_extensions/evaluator.rs

@@ -0,0 +1,49 @@
+use crate::arrow::record_batch::RecordBatch;
+
+use crate::{DeltaResult, ExpressionEvaluator};
+
+use super::super::arrow_data::ArrowEngineData;
+
+pub trait ExpressionEvaluatorExt {
+    fn evaluate_arrow(&self, batch: RecordBatch) -> DeltaResult<RecordBatch>;
+}
+
+impl<T: ExpressionEvaluator + ?Sized> ExpressionEvaluatorExt for T {
+    fn evaluate_arrow(&self, batch: RecordBatch) -> DeltaResult<RecordBatch> {
+        let engine_data = ArrowEngineData::new(batch);
+        Ok(ArrowEngineData::try_from_engine_data(T::evaluate(&self, &engine_data)?)?.into())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use super::ExpressionEvaluatorExt;
+
+    use crate::arrow::array::Int32Array;
+    use crate::arrow::datatypes::{DataType, Field, Schema};
+    use crate::arrow::record_batch::RecordBatch;
+    use crate::engine::arrow_expression::ArrowEvaluationHandler;
+    use crate::expressions::*;
+    use crate::EvaluationHandler;
+
+    #[test_log::test]
+    fn test_evaluate_arrow() {
+        let handler = ArrowEvaluationHandler;
+
+        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+        let values = Int32Array::from(vec![1, 2, 3]);
+        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(values)]).unwrap();
+
+        let expression = column_expr!("a");
+        let expr = handler.new_expression_evaluator(
+            Arc::new((&schema).try_into().unwrap()),
+            expression,
+            crate::schema::DataType::INTEGER,
+        );
+
+        let result = expr.evaluate_arrow(batch);
+        assert!(result.is_ok());
+    }
+}

kernel/src/engine/arrow_extensions/mod.rs

@@ -0,0 +1,5 @@
+mod evaluator;
+mod scan;
+
+pub use evaluator::*;
+pub use scan::*;

kernel/src/engine/arrow_extensions/scan.rs

@@ -0,0 +1,95 @@
+use std::sync::Arc;
+
+use crate::arrow::array::BooleanArray;
+use crate::arrow::compute::filter_record_batch;
+use crate::arrow::record_batch::RecordBatch;
+use itertools::Itertools;
+
+use crate::scan::{Scan, ScanMetadata, ScanResult};
+use crate::{DeltaResult, Engine, Error, ExpressionRef};
+
+use super::super::arrow_data::ArrowEngineData;
+
+/// [`ScanMetadata`] contains (1) a [`RecordBatch`] specifying data files to be scanned
+/// and (2) a vector of transforms (one transform per scan file) that must be applied to the data read
+/// from those files.
+pub struct ScanMetadataArrow {
+    /// Record batch with one row per file to scan
+    pub scan_files: RecordBatch,
+
+    /// Row-level transformations to apply to data read from files.
+    ///
+    /// Each entry in this vector corresponds to a row in the `scan_files` data. The entry is an
+    /// expression that must be applied to convert the file's data into the logical schema
+    /// expected by the scan:
+    ///
+    /// - `Some(expr)`: Apply this expression to transform the data to match [`Scan::schema()`].
+    /// - `None`: No transformation is needed; the data is already in the correct logical form.
+    ///
+    /// Note: This vector can be indexed by row number.
+    pub scan_file_transforms: Vec<Option<ExpressionRef>>,
+}
+
+impl TryFrom<ScanMetadata> for ScanMetadataArrow {
+    type Error = Error;
+
+    fn try_from(metadata: ScanMetadata) -> Result<Self, Self::Error> {
+        let scan_file_transforms = metadata
+            .scan_file_transforms
+            .into_iter()
+            .enumerate()
+            .filter_map(|(i, v)| metadata.scan_files.selection_vector[i].then_some(v))
+            .collect();
+        let batch = ArrowEngineData::try_from_engine_data(metadata.scan_files.data)?.into();
+        let scan_files = filter_record_batch(
+            &batch,
+            &BooleanArray::from(metadata.scan_files.selection_vector),
+        )?;
+        Ok(ScanMetadataArrow {
+            scan_files,
+            scan_file_transforms,
+        })
+    }
+}
+
+impl TryFrom<ScanResult> for RecordBatch {
+    type Error = Error;
+
+    fn try_from(result: ScanResult) -> Result<Self, Self::Error> {
+        let (mask, data) = (result.full_mask(), result.raw_data?);
+        let record_batch = ArrowEngineData::try_from_engine_data(data)?.into();
+        mask.map(|m| Ok(filter_record_batch(&record_batch, &m.into())?))
+            .unwrap_or(Ok(record_batch))
+    }
+}
+
+pub trait ScanExt {
+    fn scan_metadata_arrow(
+        &self,
+        engine: &dyn Engine,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<ScanMetadataArrow>>>;
+
+    fn execute_arrow(
+        &self,
+        engine: Arc<dyn Engine>,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<RecordBatch>>>;
+}
+
+impl ScanExt for Scan {
+    fn scan_metadata_arrow(
+        &self,
+        engine: &dyn Engine,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<ScanMetadataArrow>>> {
+        Ok(self
+            .scan_metadata(engine)?
+            .map_ok(TryFrom::try_from)
+            .flatten())
+    }
+
+    fn execute_arrow(
+        &self,
+        engine: Arc<dyn Engine>,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<RecordBatch>>> {
+        Ok(self.execute(engine)?.map_ok(TryFrom::try_from).flatten())
+    }
+}

kernel/src/engine/mod.rs

@@ -22,6 +22,8 @@ pub mod sync;
 #[cfg(any(feature = "default-engine-base", feature = "sync-engine"))]
 pub mod arrow_data;
 #[cfg(any(feature = "default-engine-base", feature = "sync-engine"))]
+pub mod arrow_extensions;
+#[cfg(any(feature = "default-engine-base", feature = "sync-engine"))]
 pub(crate) mod arrow_get_data;
 #[cfg(any(feature = "default-engine-base", feature = "sync-engine"))]
 pub(crate) mod ensure_data_types;

kernel/tests/arrow_extensions.rs

@@ -0,0 +1,62 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use delta_kernel::engine::arrow_extensions::ScanExt;
+use delta_kernel::engine::sync::SyncEngine;
+use delta_kernel::Table;
+use itertools::Itertools;
+
+mod common;
+
+#[test_log::test]
+fn test_scan_metadata_arrow() {
+    let path =
+        std::fs::canonicalize(PathBuf::from("./tests/data/table-without-dv-small/")).unwrap();
+    let url = url::Url::from_directory_path(path).unwrap();
+    let engine = Arc::new(SyncEngine::new());
+
+    let table = Table::new(url);
+    let snapshot = table.snapshot(engine.as_ref(), None).unwrap();
+    let scan = snapshot.into_scan_builder().build().unwrap();
+    let files: Vec<_> = scan
+        .scan_metadata_arrow(engine.as_ref())
+        .unwrap()
+        .try_collect()
+        .unwrap();
+
+    assert_eq!(files.len(), 1);
+    let num_rows = files[0].scan_files.num_rows();
+    assert_eq!(num_rows, 1)
+}
+
+#[test_log::test]
+fn test_execute_arrow() {
+    let path =
+        std::fs::canonicalize(PathBuf::from("./tests/data/table-without-dv-small/")).unwrap();
+    let url = url::Url::from_directory_path(path).unwrap();
+    let engine = Arc::new(SyncEngine::new());
+
+    let table = Table::new(url);
+    let snapshot = table.snapshot(engine.as_ref(), None).unwrap();
+    let scan = snapshot.into_scan_builder().build().unwrap();
+    let files: Vec<_> = scan.execute_arrow(engine).unwrap().try_collect().unwrap();
+
+    let expected = vec![
+        "+-------+",
+        "| value |",
+        "+-------+",
+        "| 0     |",
+        "| 1     |",
+        "| 2     |",
+        "| 3     |",
+        "| 4     |",
+        "| 5     |",
+        "| 6     |",
+        "| 7     |",
+        "| 8     |",
+        "| 9     |",
+        "+-------+",
+    ];
+
+    assert_batches_sorted_eq!(expected, &files);
+}

kernel/tests/common/mod.rs

@@ -3,7 +3,7 @@ use delta_kernel::arrow::record_batch::RecordBatch;
 use delta_kernel::arrow::util::pretty::pretty_format_batches;
 use itertools::Itertools;
 
-use crate::ArrowEngineData;
+use delta_kernel::engine::arrow_data::ArrowEngineData;
 use delta_kernel::scan::Scan;
 use delta_kernel::{DeltaResult, Engine, EngineData, Table};