Paquet and JSON readers use Arc<Expression> to avoid deep copies (#364)

Today, the engine parquet/json file handler APIs take an `Expression`
arg for predicate pushdown. They cannot take a reference, because the
iterator they return will likely depend on (but outlive) that reference.
Worse, they need to do it for every file the query reads. Unfortunately,
data skipping predicates can be arbitrarily large, and thus
annoying/expensive to copy so much. We already use Arc to protect
schemas (some of the time, at least), and we can start using Arc to
protect expressions as well.

Author: scovich

ffi/src/scan.rs

@@ -135,10 +135,9 @@ fn scan_impl(
     if let Some(predicate) = predicate {
         let mut visitor_state = KernelExpressionVisitorState::new();
         let exprid = (predicate.visitor)(predicate.predicate, &mut visitor_state);
-        if let Some(predicate) = unwrap_kernel_expression(&mut visitor_state, exprid) {
-            debug!("Got predicate: {}", predicate);
-            scan_builder = scan_builder.with_predicate(predicate);
-        }
+        let predicate = unwrap_kernel_expression(&mut visitor_state, exprid);
+        debug!("Got predicate: {:#?}", predicate);
+        scan_builder = scan_builder.with_predicate(predicate.map(Arc::new));
     }
     Ok(Arc::new(scan_builder.build()?).into())
 }

kernel/src/actions/set_transaction.rs

@@ -1,9 +1,9 @@
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use crate::actions::visitors::SetTransactionVisitor;
 use crate::actions::{get_log_schema, SetTransaction, SET_TRANSACTION_NAME};
 use crate::snapshot::Snapshot;
-use crate::{DeltaResult, Engine, EngineData, Expression as Expr, SchemaRef};
+use crate::{DeltaResult, Engine, EngineData, Expression, ExpressionRef, SchemaRef};
 
 pub use crate::actions::visitors::SetTransactionMap;
 pub struct SetTransactionScanner {
@@ -52,10 +52,11 @@ impl SetTransactionScanner {
         // checkpoint part when patitioned by `add.path` like the Delta spec requires. There's no
         // point filtering by a particular app id, even if we have one, because app ids are all in
         // the a single checkpoint part having large min/max range (because they're usually uuids).
-        let meta_predicate = Expr::column("txn.appId").is_not_null();
+        static META_PREDICATE: LazyLock<Option<ExpressionRef>> =
+            LazyLock::new(|| Some(Arc::new(Expression::column("txn.appId").is_not_null())));
         self.snapshot
             .log_segment
-            .replay(engine, schema.clone(), schema, Some(meta_predicate))
+            .replay(engine, schema.clone(), schema, META_PREDICATE.clone())
     }
 
     /// Scan the Delta Log for the latest transaction entry of an application

kernel/src/engine/default/json.rs

@@ -17,7 +17,8 @@ use super::file_stream::{FileOpenFuture, FileOpener, FileStream};
 use crate::engine::arrow_utils::parse_json as arrow_parse_json;
 use crate::schema::SchemaRef;
 use crate::{
-    DeltaResult, EngineData, Error, Expression, FileDataReadResultIterator, FileMeta, JsonHandler,
+    DeltaResult, EngineData, Error, ExpressionRef, FileDataReadResultIterator, FileMeta,
+    JsonHandler,
 };
 
 #[derive(Debug)]
@@ -72,7 +73,7 @@ impl<E: TaskExecutor> JsonHandler for DefaultJsonHandler<E> {
         &self,
         files: &[FileMeta],
         physical_schema: SchemaRef,
-        _predicate: Option<Expression>,
+        _predicate: Option<ExpressionRef>,
     ) -> DeltaResult<FileDataReadResultIterator> {
         if files.is_empty() {
             return Ok(Box::new(std::iter::empty()));

kernel/src/engine/default/parquet.rs

@@ -16,7 +16,9 @@ use crate::engine::arrow_utils::{generate_mask, get_requested_indices, reorder_s
 use crate::engine::default::executor::TaskExecutor;
 use crate::engine::parquet_row_group_skipping::ParquetRowGroupSkipping;
 use crate::schema::SchemaRef;
-use crate::{DeltaResult, Error, Expression, FileDataReadResultIterator, FileMeta, ParquetHandler};
+use crate::{
+    DeltaResult, Error, ExpressionRef, FileDataReadResultIterator, FileMeta, ParquetHandler,
+};
 
 #[derive(Debug)]
 pub struct DefaultParquetHandler<E: TaskExecutor> {
@@ -48,7 +50,7 @@ impl<E: TaskExecutor> ParquetHandler for DefaultParquetHandler<E> {
         &self,
         files: &[FileMeta],
         physical_schema: SchemaRef,
-        predicate: Option<Expression>,
+        predicate: Option<ExpressionRef>,
     ) -> DeltaResult<FileDataReadResultIterator> {
         if files.is_empty() {
             return Ok(Box::new(std::iter::empty()));
@@ -90,7 +92,7 @@ struct ParquetOpener {
     // projection: Arc<[usize]>,
     batch_size: usize,
     table_schema: SchemaRef,
-    predicate: Option<Expression>,
+    predicate: Option<ExpressionRef>,
     limit: Option<usize>,
     store: Arc<DynObjectStore>,
 }
@@ -99,7 +101,7 @@ impl ParquetOpener {
     pub(crate) fn new(
         batch_size: usize,
         table_schema: SchemaRef,
-        predicate: Option<Expression>,
+        predicate: Option<ExpressionRef>,
         store: Arc<DynObjectStore>,
     ) -> Self {
         Self {
@@ -166,14 +168,18 @@ impl FileOpener for ParquetOpener {
 /// Implements [`FileOpener`] for a opening a parquet file from a presigned URL
 struct PresignedUrlOpener {
     batch_size: usize,
-    predicate: Option<Expression>,
+    predicate: Option<ExpressionRef>,
     limit: Option<usize>,
     table_schema: SchemaRef,
     client: reqwest::Client,
 }
 
 impl PresignedUrlOpener {
-    pub(crate) fn new(batch_size: usize, schema: SchemaRef, predicate: Option<Expression>) -> Self {
+    pub(crate) fn new(
+        batch_size: usize,
+        schema: SchemaRef,
+        predicate: Option<ExpressionRef>,
+    ) -> Self {
         Self {
             batch_size,
             table_schema: schema,

kernel/src/engine/sync/json.rs

@@ -7,7 +7,7 @@ use crate::engine::arrow_data::ArrowEngineData;
 use crate::engine::arrow_utils::parse_json as arrow_parse_json;
 use crate::schema::SchemaRef;
 use crate::{
-    DeltaResult, EngineData, Expression, FileDataReadResultIterator, FileMeta, JsonHandler,
+    DeltaResult, EngineData, ExpressionRef, FileDataReadResultIterator, FileMeta, JsonHandler,
 };
 
 pub(crate) struct SyncJsonHandler;
@@ -16,7 +16,7 @@ fn try_create_from_json(
     file: File,
     _schema: SchemaRef,
     arrow_schema: ArrowSchemaRef,
-    _predicate: Option<&Expression>,
+    _predicate: Option<ExpressionRef>,
 ) -> DeltaResult<impl Iterator<Item = DeltaResult<ArrowEngineData>>> {
     let json = arrow_json::ReaderBuilder::new(arrow_schema)
         .build(BufReader::new(file))?
@@ -29,7 +29,7 @@ impl JsonHandler for SyncJsonHandler {
         &self,
         files: &[FileMeta],
         schema: SchemaRef,
-        predicate: Option<Expression>,
+        predicate: Option<ExpressionRef>,
     ) -> DeltaResult<FileDataReadResultIterator> {
         read_files(files, schema, predicate, try_create_from_json)
     }

kernel/src/engine/sync/mod.rs

@@ -3,7 +3,7 @@
 use super::arrow_expression::ArrowExpressionHandler;
 use crate::engine::arrow_data::ArrowEngineData;
 use crate::{
-    DeltaResult, Engine, Error, Expression, ExpressionHandler, FileDataReadResultIterator,
+    DeltaResult, Engine, Error, ExpressionHandler, ExpressionRef, FileDataReadResultIterator,
     FileMeta, FileSystemClient, JsonHandler, ParquetHandler, SchemaRef,
 };
 
@@ -60,12 +60,12 @@ impl Engine for SyncEngine {
 fn read_files<F, I>(
     files: &[FileMeta],
     schema: SchemaRef,
-    predicate: Option<Expression>,
+    predicate: Option<ExpressionRef>,
     mut try_create_from_file: F,
 ) -> DeltaResult<FileDataReadResultIterator>
 where
     I: Iterator<Item = DeltaResult<ArrowEngineData>> + Send + 'static,
-    F: FnMut(File, SchemaRef, ArrowSchemaRef, Option<&Expression>) -> DeltaResult<I>
+    F: FnMut(File, SchemaRef, ArrowSchemaRef, Option<ExpressionRef>) -> DeltaResult<I>
         + Send
         + 'static,
 {
@@ -88,7 +88,7 @@ where
                 File::open(path)?,
                 schema.clone(),
                 arrow_schema.clone(),
-                predicate.as_ref(),
+                predicate.clone(),
             )
         })
         // Flatten to Iterator<DeltaResult<DeltaResult<ArrowEngineData>>>

kernel/src/engine/sync/parquet.rs

@@ -8,15 +8,15 @@ use crate::engine::arrow_data::ArrowEngineData;
 use crate::engine::arrow_utils::{generate_mask, get_requested_indices, reorder_struct_array};
 use crate::engine::parquet_row_group_skipping::ParquetRowGroupSkipping;
 use crate::schema::SchemaRef;
-use crate::{DeltaResult, Expression, FileDataReadResultIterator, FileMeta, ParquetHandler};
+use crate::{DeltaResult, ExpressionRef, FileDataReadResultIterator, FileMeta, ParquetHandler};
 
 pub(crate) struct SyncParquetHandler;
 
 fn try_create_from_parquet(
     file: File,
     schema: SchemaRef,
     _arrow_schema: ArrowSchemaRef,
-    predicate: Option<&Expression>,
+    predicate: Option<ExpressionRef>,
 ) -> DeltaResult<impl Iterator<Item = DeltaResult<ArrowEngineData>>> {
     let metadata = ArrowReaderMetadata::load(&file, Default::default())?;
     let parquet_schema = metadata.schema();
@@ -27,7 +27,7 @@ fn try_create_from_parquet(
         builder = builder.with_projection(mask);
     }
     if let Some(predicate) = predicate {
-        builder = builder.with_row_group_filter(predicate);
+        builder = builder.with_row_group_filter(predicate.as_ref());
     }
     Ok(builder.build()?.map(move |data| {
         let reordered = reorder_struct_array(data?.into(), &requested_ordering)?;
@@ -40,7 +40,7 @@ impl ParquetHandler for SyncParquetHandler {
         &self,
         files: &[FileMeta],
         schema: SchemaRef,
-        predicate: Option<Expression>,
+        predicate: Option<ExpressionRef>,
     ) -> DeltaResult<FileDataReadResultIterator> {
         read_files(files, schema, predicate, try_create_from_parquet)
     }

kernel/src/expressions/mod.rs

@@ -121,6 +121,8 @@ pub enum UnaryOperator {
     IsNull,
 }
 
+pub type ExpressionRef = std::sync::Arc<Expression>;
+
 /// A SQL expression.
 ///
 /// These expressions do not track or validate data types, other than the type

kernel/src/lib.rs

@@ -77,7 +77,7 @@ pub(crate) mod utils;
 
 pub use engine_data::{DataVisitor, EngineData};
 pub use error::{DeltaResult, Error};
-pub use expressions::Expression;
+pub use expressions::{Expression, ExpressionRef};
 pub use table::Table;
 
 #[cfg(any(
@@ -192,9 +192,7 @@ pub trait JsonHandler: Send + Sync {
         &self,
         files: &[FileMeta],
         physical_schema: SchemaRef,
-        // TODO: This should really be an Option<Arc<Expression>>, because otherwise we have to
-        // clone the (potentially large) expression every time we call this function.
-        predicate: Option<Expression>,
+        predicate: Option<ExpressionRef>,
     ) -> DeltaResult<FileDataReadResultIterator>;
 }
 
@@ -216,9 +214,7 @@ pub trait ParquetHandler: Send + Sync {
         &self,
         files: &[FileMeta],
         physical_schema: SchemaRef,
-        // TODO: This should really be an Option<Arc<Expression>>, because otherwise we have to
-        // clone the (potentially large) expression every time we call this function.
-        predicate: Option<Expression>,
+        predicate: Option<ExpressionRef>,
     ) -> DeltaResult<FileDataReadResultIterator>;
 }
 

kernel/src/scan/data_skipping.rs

@@ -1,14 +1,15 @@
 use std::borrow::Cow;
 use std::collections::HashSet;
-use std::ops::Not;
 use std::sync::{Arc, LazyLock};
 
 use tracing::debug;
 
 use crate::actions::visitors::SelectionVectorVisitor;
 use crate::actions::{get_log_schema, ADD_NAME};
 use crate::error::DeltaResult;
-use crate::expressions::{BinaryOperator, Expression as Expr, UnaryOperator, VariadicOperator};
+use crate::expressions::{
+    BinaryOperator, Expression as Expr, ExpressionRef, UnaryOperator, VariadicOperator,
+};
 use crate::schema::{DataType, PrimitiveType, SchemaRef, SchemaTransform, StructField, StructType};
 use crate::{Engine, EngineData, ExpressionEvaluator, JsonHandler};
 
@@ -80,7 +81,7 @@ fn as_inverted_data_skipping_predicate(expr: &Expr) -> Option<Expr> {
             as_data_skipping_predicate(&expr)
         }
         VariadicOperation { op, exprs } => {
-            let expr = Expr::variadic(op.invert(), exprs.iter().cloned().map(Expr::not));
+            let expr = Expr::variadic(op.invert(), exprs.iter().cloned().map(|e| !e));
             as_data_skipping_predicate(&expr)
         }
         _ => None,
@@ -179,7 +180,7 @@ impl DataSkippingFilter {
     pub(crate) fn new(
         engine: &dyn Engine,
         table_schema: &SchemaRef,
-        predicate: &Option<Expr>,
+        predicate: Option<ExpressionRef>,
     ) -> Option<Self> {
         static PREDICATE_SCHEMA: LazyLock<DataType> = LazyLock::new(|| {
             DataType::struct_type([StructField::new("predicate", DataType::BOOLEAN, true)])
@@ -188,11 +189,7 @@ impl DataSkippingFilter {
         static FILTER_EXPR: LazyLock<Expr> =
             LazyLock::new(|| Expr::column("predicate").distinct(Expr::literal(false)));
 
-        let predicate = match predicate {
-            Some(predicate) => predicate,
-            None => return None,
-        };
-
+        let predicate = predicate.as_deref()?;
         debug!("Creating a data skipping filter for {}", &predicate);
         let field_names: HashSet<_> = predicate.references();