Rescope expression then the identity is a struct dtype (#1887)

Defines a transform on expression with an scope type of struct(..). The
transform splits the expr in to a series of expressions that can be
evaluated only using a single field of the struct, and a combination
expr which combines the series of smaller expr into the original expr.

This is used #1893.

---------

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

Author: joseph-isaacs

vortex-expr/src/forms/nnf.rs

@@ -1,6 +1,6 @@
 use vortex_error::VortexResult;
 
-use crate::traversal::{FoldChildren, FoldDown, FoldUp, Folder, Node as _};
+use crate::traversal::{FoldChildren, FoldDown, FoldUp, FolderMut, Node as _};
 use crate::{not, BinaryExpr, ExprRef, Not, Operator};
 
 /// Return an equivalent expression in Negative Normal Form (NNF).
@@ -63,7 +63,7 @@ pub fn nnf(expr: ExprRef) -> VortexResult<ExprRef> {
 #[derive(Default)]
 struct NNFVisitor {}
 
-impl Folder for NNFVisitor {
+impl FolderMut for NNFVisitor {
     type NodeTy = ExprRef;
     type Out = ExprRef;
     type Context = bool;

vortex-expr/src/lib.rs

@@ -19,6 +19,7 @@ mod project;
 pub mod pruning;
 mod row_filter;
 mod select;
+pub mod transform;
 #[allow(dead_code)]
 mod traversal;
 

vortex-expr/src/pack.rs

@@ -7,8 +7,8 @@ use itertools::Itertools as _;
 use vortex_array::array::StructArray;
 use vortex_array::validity::Validity;
 use vortex_array::{ArrayData, IntoArrayData};
-use vortex_dtype::FieldNames;
-use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
+use vortex_dtype::{Field, FieldNames};
+use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult};
 
 use crate::{ExprRef, VortexExpr};
 
@@ -51,6 +51,33 @@ impl Pack {
         }
         Ok(Arc::new(Pack { names, values }))
     }
+
+    pub fn names(&self) -> &FieldNames {
+        &self.names
+    }
+
+    pub fn field(&self, f: &Field) -> VortexResult<ExprRef> {
+        let idx = match f {
+            Field::Name(n) => self
+                .names
+                .iter()
+                .position(|name| name == n)
+                .ok_or_else(|| {
+                    vortex_err!("Cannot find field {} in pack fields {:?}", n, self.names)
+                })?,
+            Field::Index(idx) => *idx,
+        };
+
+        self.values
+            .get(idx)
+            .cloned()
+            .ok_or_else(|| vortex_err!("field index out of bounds: {}", idx))
+    }
+}
+
+pub fn pack(names: impl Into<FieldNames>, values: Vec<ExprRef>) -> ExprRef {
+    Pack::try_new_expr(names.into(), values)
+        .vortex_expect("pack names and values have the same length")
 }
 
 impl PartialEq<dyn Any> for Pack {

vortex-expr/src/select.rs

@@ -23,6 +23,14 @@ pub struct Select {
     child: ExprRef,
 }
 
+pub fn select(fields: Vec<Field>, child: ExprRef) -> ExprRef {
+    Select::include_expr(fields, child)
+}
+
+pub fn select_exclude(columns: Vec<Field>, child: ExprRef) -> ExprRef {
+    Select::exclude_expr(columns, child)
+}
+
 impl Select {
     pub fn new_expr(fields: SelectField, child: ExprRef) -> ExprRef {
         Arc::new(Self { fields, child })

vortex-expr/src/transform/mod.rs

@@ -0,0 +1,2 @@
+pub mod partition;
+pub mod simplify;