feat: kernel snapshot abstractions

Signed-off-by: Robert Pack <[email protected]>

Author: roeap

crates/core/src/kernel/mod.rs

@@ -10,6 +10,8 @@ pub mod error;
 pub mod models;
 pub mod scalars;
 mod snapshot;
+#[allow(unused)]
+mod snapshot_next;
 pub mod transaction;
 
 pub use error::*;

crates/core/src/kernel/snapshot/mod.rs

@@ -336,7 +336,7 @@ impl Snapshot {
 /// A snapshot of a Delta table that has been eagerly loaded into memory.
 #[derive(Debug, Clone, PartialEq)]
 pub struct EagerSnapshot {
-    snapshot: Snapshot,
+    pub(crate) snapshot: Snapshot,
     // additional actions that should be tracked during log replay.
     tracked_actions: HashSet<ActionType>,
 

crates/core/src/kernel/snapshot_next/arrow_ext.rs

@@ -0,0 +1,134 @@
+use delta_kernel::arrow::array::BooleanArray;
+use delta_kernel::arrow::compute::filter_record_batch;
+use delta_kernel::arrow::record_batch::RecordBatch;
+use delta_kernel::engine::arrow_data::ArrowEngineData;
+use delta_kernel::scan::{Scan, ScanMetadata};
+use delta_kernel::{
+    DeltaResult, Engine, EngineData, ExpressionEvaluator, ExpressionRef, PredicateRef, Version,
+};
+use itertools::Itertools;
+
+/// [`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>>,
+}
+
+pub trait ScanExt {
+    fn scan_metadata_arrow(
+        &self,
+        engine: &dyn Engine,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<ScanMetadataArrow>>>;
+
+    fn scan_metadata_from_arrow(
+        &self,
+        engine: &dyn Engine,
+        existing_version: Version,
+        existing_data: Box<dyn Iterator<Item = RecordBatch>>,
+        existing_predicate: Option<PredicateRef>,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<ScanMetadataArrow>>>;
+}
+
+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(kernel_to_arrow)
+            .flatten())
+    }
+
+    fn scan_metadata_from_arrow(
+        &self,
+        engine: &dyn Engine,
+        existing_version: Version,
+        existing_data: Box<dyn Iterator<Item = RecordBatch>>,
+        existing_predicate: Option<PredicateRef>,
+    ) -> DeltaResult<impl Iterator<Item = DeltaResult<ScanMetadataArrow>>> {
+        let engine_iter =
+            existing_data.map(|batch| Box::new(ArrowEngineData::new(batch)) as Box<dyn EngineData>);
+        Ok(self
+            .scan_metadata_from(engine, existing_version, engine_iter, existing_predicate)?
+            .map_ok(kernel_to_arrow)
+            .flatten())
+    }
+}
+
+fn kernel_to_arrow(metadata: ScanMetadata) -> DeltaResult<ScanMetadataArrow> {
+    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,
+    })
+}
+
+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 delta_kernel::arrow::array::Int32Array;
+    use delta_kernel::arrow::datatypes::{DataType, Field, Schema};
+    use delta_kernel::arrow::record_batch::RecordBatch;
+    use delta_kernel::engine::arrow_conversion::TryIntoKernel;
+    use delta_kernel::engine::arrow_expression::ArrowEvaluationHandler;
+    use delta_kernel::expressions::*;
+    use delta_kernel::EvaluationHandler;
+
+    #[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_kernel().unwrap()),
+            expression,
+            delta_kernel::schema::DataType::INTEGER,
+        );
+
+        let result = expr.evaluate_arrow(batch);
+        assert!(result.is_ok());
+    }
+}

crates/core/src/kernel/snapshot_next/eager.rs

@@ -0,0 +1,141 @@
+use std::sync::Arc;
+
+use arrow::compute::concat_batches;
+use arrow_array::RecordBatch;
+use delta_kernel::actions::{Metadata, Protocol};
+use delta_kernel::engine::arrow_conversion::TryIntoArrow;
+use delta_kernel::scan::scan_row_schema;
+use delta_kernel::schema::Schema;
+use delta_kernel::table_properties::TableProperties;
+use delta_kernel::{PredicateRef, Table, Version};
+use itertools::Itertools;
+use object_store::ObjectStore;
+use url::Url;
+use uuid::Uuid;
+
+use super::arrow_ext::ScanExt;
+use super::lazy::LazySnapshot;
+use super::Snapshot;
+use crate::kernel::CommitInfo;
+use crate::logstore::LogStore;
+use crate::{DeltaResult, DeltaTableConfig};
+
+/// An eager snapshot of a Delta Table at a specific version.
+///
+/// This snapshot loads some log data eagerly and keeps it in memory.
+#[derive(Clone)]
+pub struct EagerSnapshot {
+    snapshot: LazySnapshot,
+    files: RecordBatch,
+    predicate: Option<PredicateRef>,
+}
+
+#[async_trait::async_trait]
+impl Snapshot for EagerSnapshot {
+    fn table_root(&self) -> &Url {
+        self.snapshot.table_root()
+    }
+
+    fn version(&self) -> Version {
+        self.snapshot.version()
+    }
+
+    fn schema(&self) -> Arc<Schema> {
+        self.snapshot.schema()
+    }
+
+    fn table_properties(&self) -> &TableProperties {
+        self.snapshot.table_properties()
+    }
+
+    async fn logical_files(
+        &self,
+        predicate: Option<PredicateRef>,
+    ) -> DeltaResult<Box<dyn Iterator<Item = DeltaResult<RecordBatch>> + '_>> {
+        let snapshot = self.snapshot.inner.clone();
+        let scan = snapshot.scan_builder().with_predicate(predicate).build()?;
+        let iter: Vec<_> = scan
+            .scan_metadata_from_arrow(
+                self.snapshot.engine_ref().as_ref(),
+                self.version(),
+                Box::new(std::iter::once(self.files.clone())),
+                self.predicate.clone(),
+            )?
+            .collect();
+        Ok(Box::new(iter.into_iter().map(|sc| Ok(sc?.scan_files))))
+    }
+
+    async fn application_transaction_version(&self, app_id: String) -> DeltaResult<Option<i64>> {
+        self.snapshot.application_transaction_version(app_id).await
+    }
+
+    fn commit_infos(
+        &self,
+        start_version: Option<Version>,
+        limit: Option<usize>,
+    ) -> DeltaResult<Box<dyn Iterator<Item = (Version, CommitInfo)>>> {
+        self.snapshot.commit_infos(start_version, limit)
+    }
+
+    async fn update(&mut self, target_version: Option<Version>) -> DeltaResult<bool> {
+        self.update_impl(target_version).await
+    }
+}
+
+impl EagerSnapshot {
+    /// Create a new [`EagerSnapshot`] instance
+    pub(crate) async fn try_new(
+        log_store: Arc<dyn LogStore>,
+        config: DeltaTableConfig,
+        version: impl Into<Option<Version>>,
+        predicate: Option<PredicateRef>,
+        operation_id: impl Into<Option<Uuid>>,
+    ) -> DeltaResult<Self> {
+        let snapshot = LazySnapshot::try_new(log_store, version, operation_id).await?;
+        let all: Vec<_> = snapshot
+            .logical_files(predicate.clone())
+            .await?
+            .try_collect()?;
+        let files = if all.is_empty() {
+            RecordBatch::new_empty(Arc::new((scan_row_schema().as_ref()).try_into_arrow()?))
+        } else {
+            concat_batches(&all[0].schema(), &all)?
+        };
+        Ok(Self {
+            snapshot,
+            files,
+            predicate,
+        })
+    }
+
+    /// Get the number of files in the current snapshot
+    pub fn files_count(&self) -> usize {
+        self.files.num_rows()
+    }
+
+    async fn update_impl(&mut self, target_version: Option<Version>) -> DeltaResult<bool> {
+        let mut snapshot = self.snapshot.clone();
+        let current = snapshot.version();
+        if !snapshot.update(target_version.clone()).await? {
+            return Ok(false);
+        }
+
+        let scan = snapshot.inner.clone().scan_builder().build()?;
+        let engine = snapshot.engine_ref().clone();
+        // TODO: process blocking iterator
+        let files: Vec<_> = scan
+            .scan_metadata_from_arrow(
+                engine.as_ref(),
+                current,
+                Box::new(std::iter::once(self.files.clone())),
+                self.predicate.clone(),
+            )?
+            .map_ok(|s| s.scan_files)
+            .try_collect()?;
+
+        self.files = concat_batches(&files[0].schema(), &files)?;
+        self.snapshot = snapshot;
+
+        Ok(true)
+    }
+}

crates/core/src/kernel/snapshot_next/iterators.rs

@@ -0,0 +1,67 @@
+use arrow_array::cast::AsArray;
+use arrow_array::types::Int64Type;
+use arrow_array::{Array, RecordBatch};
+use chrono::{DateTime, Utc};
+use delta_kernel::expressions::{Scalar, StructData};
+
+use crate::kernel::scalars::ScalarExt;
+use crate::{DeltaResult, DeltaTableError};
+
+#[derive(Clone)]
+pub struct LogicalFileView {
+    pub(super) files: RecordBatch,
+    pub(super) index: usize,
+}
+
+impl LogicalFileView {
+    /// Path of the file.
+    pub fn path(&self) -> &str {
+        self.files.column(0).as_string::<i32>().value(self.index)
+    }
+
+    /// Size of the file in bytes.
+    pub fn size(&self) -> i64 {
+        self.files
+            .column(1)
+            .as_primitive::<Int64Type>()
+            .value(self.index)
+    }
+
+    /// Modification time of the file in milliseconds since epoch.
+    pub fn modification_time(&self) -> i64 {
+        self.files
+            .column(2)
+            .as_primitive::<Int64Type>()
+            .value(self.index)
+    }
+
+    /// Datetime of the last modification time of the file.
+    pub fn modification_datetime(&self) -> DeltaResult<chrono::DateTime<Utc>> {
+        DateTime::from_timestamp_millis(self.modification_time()).ok_or(
+            DeltaTableError::MetadataError(format!(
+                "invalid modification_time: {:?}",
+                self.modification_time()
+            )),
+        )
+    }
+
+    pub fn stats(&self) -> Option<&str> {
+        let col = self.files.column(3).as_string::<i32>();
+        col.is_valid(self.index).then(|| col.value(self.index))
+    }
+
+    pub fn partition_values(&self) -> Option<StructData> {
+        self.files
+            .column_by_name("fileConstantValues")
+            .and_then(|col| col.as_struct_opt())
+            .and_then(|s| s.column_by_name("partitionValues"))
+            .and_then(|arr| {
+                arr.is_valid(self.index)
+                    .then(|| match Scalar::from_array(arr, self.index) {
+                        Some(Scalar::Struct(s)) => Some(s),
+                        _ => None,
+                    })
+                    .flatten()
+            })
+    }
+}