Cache partitioned expressions in StructLayoutReader (#1947)

They're not trivially cheap to build

Author: gatesn

vortex-array/src/stats/statsset.rs

@@ -1,16 +1,24 @@
-use enum_iterator::all;
+use enum_iterator::{all, Sequence};
 use itertools::{EitherOrBoth, Itertools};
 use vortex_dtype::DType;
 use vortex_error::{vortex_panic, VortexError};
 use vortex_scalar::Scalar;
 
 use crate::stats::Stat;
 
-#[derive(Debug, Clone, Default, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct StatsSet {
     values: Vec<(Stat, Scalar)>,
 }
 
+impl Default for StatsSet {
+    fn default() -> Self {
+        Self {
+            values: Vec::with_capacity(Stat::CARDINALITY),
+        }
+    }
+}
+
 impl StatsSet {
     /// Create new StatSet without validating uniqueness of all the entries
     ///

vortex-file/src/v2/file.rs

@@ -62,7 +62,7 @@ impl<I: IoDriver> VortexFile<I> {
             .reader(segment_channel.reader(), self.ctx.clone())?;
 
         // Now we give one end of the channel to the layout reader...
-        log::info!("Starting scan with {} splits", self.splits.len());
+        log::debug!("Starting scan with {} splits", self.splits.len());
         let exec_stream = stream::iter(ArcIter::new(self.splits.clone()))
             .map(move |row_range| scan.clone().range_scan(row_range))
             .map(move |range_scan| match range_scan {

vortex-layout/src/layouts/struct_/eval_expr.rs

@@ -5,7 +5,6 @@ use vortex_array::array::StructArray;
 use vortex_array::validity::Validity;
 use vortex_array::{ArrayData, IntoArrayData};
 use vortex_error::VortexResult;
-use vortex_expr::transform::partition::partition;
 use vortex_expr::ExprRef;
 use vortex_scan::RowMask;
 
@@ -16,7 +15,7 @@ use crate::ExprEvaluator;
 impl ExprEvaluator for StructReader {
     async fn evaluate_expr(&self, row_mask: RowMask, expr: ExprRef) -> VortexResult<ArrayData> {
         // Partition the expression into expressions that can be evaluated over individual fields
-        let partitioned = partition(expr.clone(), self.struct_dtype())?;
+        let partitioned = self.partition_expr(expr.clone())?;
         let field_readers: Vec<_> = partitioned
             .partitions
             .iter()

vortex-layout/src/layouts/struct_/reader.rs

@@ -1,9 +1,12 @@
-use std::sync::{Arc, OnceLock};
+use std::hash::Hash;
+use std::sync::{Arc, OnceLock, RwLock};
 
-use vortex_array::aliases::hash_map::HashMap;
+use vortex_array::aliases::hash_map::{Entry, HashMap};
 use vortex_array::ContextRef;
 use vortex_dtype::{DType, FieldName, StructDType};
 use vortex_error::{vortex_err, vortex_panic, VortexExpect, VortexResult};
+use vortex_expr::transform::partition::{partition, PartitionedExpr};
+use vortex_expr::ExprRef;
 
 use crate::layouts::struct_::StructLayout;
 use crate::segments::AsyncSegmentReader;
@@ -18,6 +21,8 @@ pub struct StructReader {
 
     field_readers: Arc<[OnceLock<Arc<dyn LayoutReader>>]>,
     field_lookup: HashMap<FieldName, usize>,
+
+    expr_cache: Arc<RwLock<HashMap<ExactExpr, Arc<PartitionedExpr>>>>,
 }
 
 impl StructReader {
@@ -52,6 +57,7 @@ impl StructReader {
             segments,
             field_readers,
             field_lookup,
+            expr_cache: Arc::new(Default::default()),
         })
     }
 
@@ -77,10 +83,46 @@ impl StructReader {
             child_layout.reader(self.segments.clone(), self.ctx.clone())
         })
     }
+
+    /// Utility for partitioning an expression over the fields of a struct.
+    pub(crate) fn partition_expr(&self, expr: ExprRef) -> VortexResult<Arc<PartitionedExpr>> {
+        Ok(
+            match self
+                .expr_cache
+                .write()
+                .map_err(|_| vortex_err!("poisoned lock"))?
+                .entry(ExactExpr(expr.clone()))
+            {
+                Entry::Occupied(entry) => entry.get().clone(),
+                Entry::Vacant(entry) => entry
+                    .insert(Arc::new(partition(expr, self.struct_dtype())?))
+                    .clone(),
+            },
+        )
+    }
 }
 
 impl LayoutReader for StructReader {
     fn layout(&self) -> &LayoutData {
         &self.layout
     }
 }
+
+/// An expression wrapper that performs pointer equality.
+/// NOTE(ngates): we should consider if this shoud live in vortex-expr crate?
+#[derive(Clone)]
+struct ExactExpr(ExprRef);
+
+impl PartialEq for ExactExpr {
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.0, &other.0)
+    }
+}
+
+impl Eq for ExactExpr {}
+
+impl Hash for ExactExpr {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        Arc::as_ptr(&self.0).hash(state)
+    }
+}