feat: kernel snapshot abstractions

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

Author: roeap

.pre-commit-config.yaml

@@ -1,5 +1,29 @@
+ci:
+  skip: [cargo-fmt]
+
 repos:
   - repo: https://github.com/crate-ci/typos
     rev: v1.32.0
     hooks:
       - id: typos
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-merge-conflict
+      - id: trailing-whitespace
+        args: [--markdown-linebreak-ext=md]
+
+  - repo: https://github.com/gitleaks/gitleaks
+    rev: v8.26.0
+    hooks:
+      - id: gitleaks
+
+  - repo: local
+    hooks:
+      - id: cargo-fmt
+        name: cargo fmt
+        entry: cargo fmt --
+        language: system
+        types: [rust]
+        pass_filenames: false # This makes it a lot faster

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

@@ -21,6 +21,7 @@ use std::io::{BufRead, BufReader, Cursor};
 use ::serde::{Deserialize, Serialize};
 use arrow_array::RecordBatch;
 use delta_kernel::path::{LogPathFileType, ParsedLogPath};
+use delta_kernel::Version;
 use futures::stream::BoxStream;
 use futures::{StreamExt, TryStreamExt};
 use object_store::path::Path;
@@ -241,7 +242,7 @@ impl Snapshot {
         &self,
         log_store: &dyn LogStore,
         limit: Option<usize>,
-    ) -> DeltaResult<BoxStream<'_, DeltaResult<Option<CommitInfo>>>> {
+    ) -> DeltaResult<BoxStream<'static, DeltaResult<Option<(Version, CommitInfo)>>>> {
         let store = log_store.root_object_store(None);
 
         let log_root = self.table_root_path()?.child("_delta_log");
@@ -266,21 +267,24 @@ impl Snapshot {
             let dummy_path = dummy_url.join(meta.location.as_ref()).unwrap();
             if let Some(parsed_path) = ParsedLogPath::try_from(dummy_path)? {
                 if matches!(parsed_path.file_type, LogPathFileType::Commit) {
-                    commit_files.push(meta);
+                    commit_files.push((meta, parsed_path));
                 }
             }
         }
-        commit_files.sort_unstable_by(|a, b| b.location.cmp(&a.location));
+        commit_files.sort_unstable_by(|a, b| b.0.location.cmp(&a.0.location));
         Ok(futures::stream::iter(commit_files)
-            .map(move |meta| {
+            .map(move |(meta, parsed_path)| {
                 let store = store.clone();
                 async move {
                     let commit_log_bytes = store.get(&meta.location).await?.bytes().await?;
                     let reader = BufReader::new(Cursor::new(commit_log_bytes));
                     for line in reader.lines() {
                         let action: Action = serde_json::from_str(line?.as_str())?;
                         if let Action::CommitInfo(commit_info) = action {
-                            return Ok::<_, DeltaTableError>(Some(commit_info));
+                            return Ok::<_, DeltaTableError>(Some((
+                                parsed_path.version,
+                                commit_info,
+                            )));
                         }
                     }
                     Ok(None)
@@ -336,7 +340,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,141 @@
+//! Utilities for interacting with Kernel APIs using Arrow data structures.
+//!
+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 {
+    /// Get the metadata for a table scan.
+    ///
+    /// This method handles translation between `EngineData` and `RecordBatch`
+    /// and will already apply any selection vectors to the data.
+    /// See [`Scan::scan_metadata`] for details.
+    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,152 @@
+use std::sync::Arc;
+
+use arrow::compute::concat_batches;
+use arrow_array::RecordBatch;
+use arrow_schema::SchemaRef as ArrowSchemaRef;
+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 futures::TryStreamExt;
+use itertools::Itertools;
+use object_store::ObjectStore;
+use url::Url;
+use uuid::Uuid;
+
+use super::arrow_ext::ScanExt;
+use super::lazy::LazySnapshot;
+use super::stream::SendableRBStream;
+use super::{LogStoreHandler, 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 version(&self) -> Version {
+        self.snapshot.version()
+    }
+
+    fn schema(&self) -> ArrowSchemaRef {
+        self.snapshot.schema()
+    }
+
+    fn table_properties(&self) -> &TableProperties {
+        self.snapshot.table_properties()
+    }
+
+    fn logical_files(
+        &self,
+        log_store: &LogStoreHandler,
+        predicate: Option<PredicateRef>,
+    ) -> SendableRBStream {
+        if let Some(predicate) = predicate {
+            self.snapshot.logical_files(log_store, Some(predicate))
+        } else {
+            let batch = self.files.clone();
+            return Box::pin(futures::stream::once(async move { Ok(batch) }));
+        }
+    }
+
+    async fn application_transaction_version(
+        &self,
+        log_store: &LogStoreHandler,
+        app_id: String,
+    ) -> DeltaResult<Option<i64>> {
+        self.snapshot
+            .application_transaction_version(log_store, app_id)
+            .await
+    }
+
+    async fn commit_infos(
+        &self,
+        log_store: &LogStoreHandler,
+        start_version: Option<Version>,
+        limit: Option<usize>,
+    ) -> DeltaResult<Vec<(Version, CommitInfo)>> {
+        self.snapshot
+            .commit_infos(log_store, start_version, limit)
+            .await
+    }
+
+    async fn update(
+        &mut self,
+        log_store: &LogStoreHandler,
+        target_version: Option<Version>,
+    ) -> DeltaResult<bool> {
+        self.update_impl(log_store, target_version).await
+    }
+}
+
+impl EagerSnapshot {
+    /// Create a new [`EagerSnapshot`] instance
+    pub(crate) async fn try_new(
+        log_store: &LogStoreHandler,
+        config: DeltaTableConfig,
+        version: impl Into<Option<Version>>,
+        predicate: Option<PredicateRef>,
+    ) -> DeltaResult<Self> {
+        let snapshot = LazySnapshot::try_new(log_store, version).await?;
+        let all: Vec<_> = snapshot
+            .logical_files(log_store, predicate.clone())
+            .try_collect()
+            .await?;
+        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,
+        log_store: &LogStoreHandler,
+        target_version: Option<Version>,
+    ) -> DeltaResult<bool> {
+        let mut snapshot = self.snapshot.clone();
+        let current = snapshot.version();
+        if !snapshot.update(log_store, target_version.clone()).await? {
+            return Ok(false);
+        }
+
+        let scan = snapshot.inner.clone().scan_builder().build()?;
+        let engine = log_store.engine();
+        // 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)
+    }
+}