Split by layout chunks of fields used in query (#2022)

Co-authored-by: Nicholas Gates <[email protected]>

Author: joseph-isaacs

vortex-datafusion/src/persistent/opener.rs

@@ -15,7 +15,7 @@ use vortex_error::VortexResult;
 use vortex_expr::datafusion::convert_expr_to_vortex;
 use vortex_expr::transform::simplify_typed::simplify_typed;
 use vortex_expr::{and, ident, lit, select, ExprRef};
-use vortex_file::{ExecutionMode, Scan, SplitBy, VortexOpenOptions};
+use vortex_file::{ExecutionMode, Scan, VortexOpenOptions};
 use vortex_io::ObjectStoreReadAt;
 
 use super::cache::FileLayoutCache;
@@ -52,7 +52,7 @@ impl VortexFileOpener {
                     .reduce(and)
                     .unwrap_or_else(|| lit(true));
 
-                simplify_typed(expr, dtype)
+                simplify_typed(expr, &dtype)
             })
             .transpose()?;
 
@@ -88,8 +88,6 @@ impl FileOpener for VortexFileOpener {
                         .try_get(&file_meta.object_meta, this.object_store.clone())
                         .await?,
                 )
-                // Create larger splits in so that each chunk has more rows
-                .with_split_by(SplitBy::RowCount(2 << 15))
                 .with_execution_mode(ExecutionMode::TokioRuntime(Handle::current()))
                 .open(read_at)
                 .await?;

vortex-dtype/src/field.rs

@@ -127,17 +127,27 @@ impl FieldPath {
     }
 
     /// Pushes a new field selector to the end of this path
-    pub fn push<F: Into<Field>>(&mut self, field: F) {
+    pub fn push<F: Into<Field>>(mut self, field: F) -> Self {
         self.0.push(field.into());
+        self
+    }
+
+    /// Whether the path starts with the given field name
+    /// TODO(joe): handle asserts better.
+    pub fn starts_with_field(&self, field: &Field) -> bool {
+        assert!(matches!(field, Field::Name(_)));
+        let first = self.0.first();
+        assert!(matches!(first, Some(Field::Name(_))));
+        first.is_some_and(|f| f == field)
     }
 
     /// Steps into the next field in the path
-    pub fn step_into(mut self) -> VortexResult<Self> {
+    pub fn step_into(mut self) -> Option<Self> {
         if self.0.is_empty() {
-            return Err(vortex_err!("Cannot step into root path"));
+            return None;
         }
         self.0 = self.0.iter().skip(1).cloned().collect();
-        Ok(self)
+        Some(self)
     }
 }
 
@@ -171,9 +181,7 @@ mod tests {
 
     #[test]
     fn test_field_path() {
-        let mut path = FieldPath::from_name("A");
-        path.push("B");
-        path.push("C");
+        let path = FieldPath::from_name("A").push("B").push("C");
         assert_eq!(path.to_string(), "$A.$B.$C");
 
         let fields = vec!["A", "B", "C"]

vortex-dtype/src/field_mask.rs

@@ -0,0 +1,68 @@
+//! Field mask represents a field projection, which leads to a set of field paths under a given layout.
+
+use vortex_error::{vortex_bail, VortexResult};
+
+use crate::{Field, FieldPath};
+
+/// Represents a field mask, which is a projection of fields under a layout.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum FieldMask {
+    /// Select all fields in the layout
+    All,
+    /// Select all with the `FieldPath` prefix
+    Prefix(FieldPath),
+    /// Select a field matching exactly the `FieldPath`
+    Exact(FieldPath),
+}
+
+impl FieldMask {
+    /// Creates a new field mask stepping one level into the layout structure.
+    pub fn step_into(self) -> VortexResult<Self> {
+        match self {
+            FieldMask::All => Ok(FieldMask::All),
+            FieldMask::Prefix(fp) => {
+                let Some(stepped_fp) = fp.step_into() else {
+                    return Ok(FieldMask::All);
+                };
+                if stepped_fp.is_root() {
+                    Ok(FieldMask::All)
+                } else {
+                    Ok(FieldMask::Prefix(stepped_fp))
+                }
+            }
+            FieldMask::Exact(fp) => {
+                if let Some(stepped_fp) = fp.step_into() {
+                    Ok(FieldMask::Exact(stepped_fp))
+                } else {
+                    vortex_bail!("Cannot step into exact root field path");
+                }
+            }
+        }
+    }
+
+    /// Returns the first field explicit select mask, if there is one, failing if mask = `All`.
+    pub fn starting_field(&self) -> VortexResult<Option<&Field>> {
+        match self {
+            FieldMask::All => vortex_bail!("Cannot get starting field from All mask"),
+            // We know that fp is non-empty
+            FieldMask::Prefix(fp) | FieldMask::Exact(fp) => Ok(fp.path().first()),
+        }
+    }
+
+    /// True iff all fields are selected (including self).
+    pub fn matches_all(&self) -> bool {
+        match self {
+            FieldMask::All => true,
+            FieldMask::Prefix(path) => path.is_root(),
+            FieldMask::Exact(_) => false,
+        }
+    }
+
+    /// True if the mask matches the root field.
+    pub fn matches_root(&self) -> bool {
+        match self {
+            FieldMask::All => true,
+            FieldMask::Prefix(path) | FieldMask::Exact(path) => path.is_root(),
+        }
+    }
+}

vortex-dtype/src/lib.rs

@@ -9,6 +9,7 @@
 pub use dtype::*;
 pub use extension::*;
 pub use field::*;
+pub use field_mask::*;
 pub use half;
 pub use nullability::*;
 pub use ptype::*;
@@ -19,6 +20,7 @@ mod arbitrary;
 mod dtype;
 mod extension;
 mod field;
+mod field_mask;
 mod nullability;
 mod ptype;
 mod serde;

vortex-dtype/src/struct_.rs

@@ -238,6 +238,14 @@ impl StructDType {
         &self.names
     }
 
+    /// Find the index of a field.
+    pub fn find(&self, field: &Field) -> Option<usize> {
+        match field {
+            Field::Name(name) => self.find_name(name),
+            Field::Index(idx) => Some(*idx),
+        }
+    }
+
     /// Find the index of a field by name
     /// Returns `None` if the field is not found
     pub fn find_name(&self, name: &str) -> Option<usize> {

vortex-expr/src/transform/field_mask.rs

@@ -0,0 +1,109 @@
+use vortex_array::aliases::hash_set::HashSet;
+use vortex_dtype::{DType, Field, FieldPath};
+use vortex_error::{vortex_bail, VortexResult};
+
+use crate::traversal::{FoldUp, Folder, Node};
+use crate::{ExprRef, GetItem, Identity, Select};
+
+/// Returns the field mask for the given expression.
+///
+/// This defines a mask over the scope of the fields that are accessed by the expression.
+pub fn field_mask(expr: &ExprRef, scope_dtype: &DType) -> VortexResult<HashSet<FieldPath>> {
+    // I know it's unused now, but we will for sure need the scope DType for future expressions.
+    let DType::Struct(_scope_dtype, _) = scope_dtype else {
+        vortex_bail!("Mismatched dtype {} for struct layout", scope_dtype);
+    };
+
+    Ok(match expr.accept_with_context(&mut FieldMaskFolder, ())? {
+        FoldUp::Abort(out) => out,
+        FoldUp::Continue(out) => out,
+    })
+}
+
+struct FieldMaskFolder;
+
+impl<'a> Folder<'a> for FieldMaskFolder {
+    type NodeTy = ExprRef;
+    type Out = HashSet<FieldPath>;
+    type Context = ();
+
+    fn visit_up(
+        &mut self,
+        node: &'a Self::NodeTy,
+        _context: Self::Context,
+        children: Vec<Self::Out>,
+    ) -> VortexResult<FoldUp<Self::Out>> {
+        // The identity returns a field path covering the root.
+        if node.as_any().is::<Identity>() {
+            return Ok(FoldUp::Continue([FieldPath::root()].into()));
+        }
+
+        // GetItem pushes an element to each field path
+        if let Some(getitem) = node.as_any().downcast_ref::<GetItem>() {
+            let fields = children
+                .into_iter()
+                .flat_map(|field_mask| field_mask.into_iter())
+                .map(|field_path| field_path.push(Field::Name(getitem.field().clone())))
+                .collect();
+            return Ok(FoldUp::Continue(fields));
+        }
+
+        if node.as_any().is::<Select>() {
+            vortex_bail!("Expression must be simplified")
+        }
+
+        // Otherwise, return the field paths from the children
+        Ok(FoldUp::Continue(children.into_iter().flatten().collect()))
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::iter;
+
+    use itertools::Itertools;
+    use vortex_dtype::Nullability::NonNullable;
+    use vortex_dtype::{DType, FieldPath, PType, StructDType};
+
+    use crate::transform::field_mask::field_mask;
+    use crate::{get_item, ident};
+
+    fn dtype() -> DType {
+        DType::Struct(
+            StructDType::new(
+                ["A".into(), "B".into(), "C".into()].into(),
+                iter::repeat(DType::Primitive(PType::I32, NonNullable))
+                    .take(3)
+                    .collect(),
+            ),
+            NonNullable,
+        )
+    }
+
+    #[test]
+    fn field_mask_ident() {
+        let mask = field_mask(&ident(), &dtype())
+            .unwrap()
+            .into_iter()
+            .collect_vec();
+        assert_eq!(mask.as_slice(), &[FieldPath::root()]);
+    }
+
+    #[test]
+    fn field_mask_get_item() {
+        let mask = field_mask(&get_item("A", ident()), &dtype())
+            .unwrap()
+            .into_iter()
+            .collect_vec();
+        assert_eq!(mask.as_slice(), &[FieldPath::from_name("A")]);
+    }
+
+    #[test]
+    fn field_mask_get_item_nested() {
+        let mask = field_mask(&get_item("B", get_item("A", ident())), &dtype())
+            .unwrap()
+            .into_iter()
+            .collect_vec();
+        assert_eq!(mask.as_slice(), &[FieldPath::from_name("A").push("B")]);
+    }
+}

vortex-expr/src/transform/immediate_access.rs

@@ -2,7 +2,7 @@ use itertools::Itertools;
 use vortex_array::aliases::hash_map::HashMap;
 use vortex_array::aliases::hash_set::HashSet;
 use vortex_dtype::{FieldName, StructDType};
-use vortex_error::VortexResult;
+use vortex_error::{vortex_err, VortexResult};
 
 use crate::traversal::{Node, NodeVisitor, TraversalOrder};
 use crate::{ExprRef, GetItem, Identity, Select};
@@ -23,6 +23,19 @@ pub fn immediate_scope_accesses<'a>(
     ImmediateScopeAccessesAnalysis::<'a>::analyze(expr, scope_dtype)
 }
 
+/// This returns the immediate scope_access (as explained `immediate_scope_accesses`) for `expr`.
+pub fn immediate_scope_access<'a>(
+    expr: &'a ExprRef,
+    scope_dtype: &'a StructDType,
+) -> VortexResult<HashSet<FieldName>> {
+    ImmediateScopeAccessesAnalysis::<'a>::analyze(expr, scope_dtype)?
+        .get(expr)
+        .ok_or_else(|| {
+            vortex_err!("Expression missing from scope accesses, this is a internal bug")
+        })
+        .cloned()
+}
+
 struct ImmediateScopeAccessesAnalysis<'a> {
     sub_expressions: FieldAccesses<'a>,
     scope_dtype: &'a StructDType,

vortex-expr/src/transform/mod.rs

@@ -1,5 +1,6 @@
 //! A collection of transformations that can be applied to a [`crate::ExprRef`].
-mod immediate_access;
+pub mod field_mask;
+pub mod immediate_access;
 pub mod partition;
 pub(crate) mod remove_select;
 pub mod simplify;

vortex-expr/src/transform/partition.rs

@@ -114,7 +114,7 @@ impl<'a> StructFieldExpressionSplitter<'a> {
                 };
                 VortexResult::Ok(Partition {
                     name,
-                    expr: simplify_typed(expr, field_dtype)?,
+                    expr: simplify_typed(expr, &field_dtype)?,
                 })
             })
             .try_collect()?;
@@ -131,7 +131,7 @@ impl<'a> StructFieldExpressionSplitter<'a> {
             .transform(&mut ReplaceAccessesWithChild(remove_accesses))?;
 
         Ok(PartitionedExpr {
-            root: simplify_typed(split.result, dtype.clone())?,
+            root: simplify_typed(split.result, dtype)?,
             partitions: partitions.into_boxed_slice(),
         })
     }
@@ -376,7 +376,7 @@ mod tests {
             get_item("b", get_item("a", ident())),
             select(vec!["a".into(), "b".into()], ident()),
         );
-        let expr = simplify_typed(expr, dtype.clone()).unwrap();
+        let expr = simplify_typed(expr, &dtype).unwrap();
         let partitioned = StructFieldExpressionSplitter::split(expr, &dtype).unwrap();
 
         // One for id.a and id.b

vortex-expr/src/transform/remove_select.rs

@@ -6,28 +6,24 @@ use crate::{get_item, pack, ExprRef, Select};
 
 /// Select is a useful expression, however it can be defined in terms of get_item & pack,
 /// once the expression type is known, this simplifications pass removes the select expression.
-pub fn remove_select(e: ExprRef, scope_dt: DType) -> VortexResult<ExprRef> {
-    let mut transform = RemoveSelectTransform::new(scope_dt);
+pub fn remove_select(e: ExprRef, scope_dt: &DType) -> VortexResult<ExprRef> {
+    let mut transform = RemoveSelectTransform {
+        scope_dtype: scope_dt,
+    };
     e.transform(&mut transform).map(|e| e.result)
 }
 
-struct RemoveSelectTransform {
-    ident_dtype: DType,
+struct RemoveSelectTransform<'a> {
+    scope_dtype: &'a DType,
 }
 
-impl RemoveSelectTransform {
-    fn new(ident_dtype: DType) -> Self {
-        Self { ident_dtype }
-    }
-}
-
-impl MutNodeVisitor for RemoveSelectTransform {
+impl MutNodeVisitor for RemoveSelectTransform<'_> {
     type NodeTy = ExprRef;
 
     fn visit_up(&mut self, node: ExprRef) -> VortexResult<TransformResult<Self::NodeTy>> {
         if let Some(select) = node.as_any().downcast_ref::<Select>() {
             let child = select.child();
-            let child_dtype = child.return_dtype(&self.ident_dtype)?;
+            let child_dtype = child.return_dtype(self.scope_dtype)?;
             let child_dtype = child_dtype.as_struct().ok_or_else(|| {
                 vortex_err!(
                     "Select child must return a struct dtype, however it was a {}",
@@ -76,7 +72,7 @@ mod tests {
             NonNullable,
         );
         let e = select(["a".into(), "b".into()], ident());
-        let e = remove_select(e, dtype).unwrap();
+        let e = remove_select(e, &dtype).unwrap();
 
         assert!(e.as_any().is::<Pack>());
     }