Shared identity expression (#1916)

spiraldb · Jan 13, 2025 · b281020 · b281020
1 parent 98e2968
commit b281020
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 15 deletions.
diff --git a/vortex-expr/src/identity.rs b/vortex-expr/src/identity.rs
@@ -1,18 +1,20 @@
 use std::any::Any;
 use std::fmt::Display;
-use std::sync::Arc;
+use std::sync::{Arc, LazyLock};
 
 use vortex_array::ArrayData;
 use vortex_error::VortexResult;
 
 use crate::{ExprRef, VortexExpr};
 
+static IDENTITY: LazyLock<ExprRef> = LazyLock::new(|| Arc::new(Identity));
+
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct Identity;
 
 impl Identity {
     pub fn new_expr() -> ExprRef {
-        Arc::new(Identity)
+        IDENTITY.clone()
     }
 }
 

diff --git a/vortex-expr/src/transform/partition.rs b/vortex-expr/src/transform/partition.rs
@@ -295,7 +295,7 @@ mod tests {
     use vortex_dtype::{DType, StructDType};
 
     use super::*;
-    use crate::transform::simplify::Simplify;
+    use crate::transform::simplify::simplify;
     use crate::{and, get_item, ident, lit, pack, select, Pack};
 
     fn struct_dtype() -> StructDType {
@@ -348,7 +348,7 @@ mod tests {
             &get_item("0", get_item(split_a.name.clone(), ident()))
         );
         assert_eq!(
-            &Simplify::simplify(split_a.expr.clone()).unwrap(),
+            &simplify(split_a.expr.clone()).unwrap(),
             &pack(vec!["0".into()], vec![get_item("b", ident())])
         );
     }
@@ -369,15 +369,15 @@ mod tests {
 
         let split_a = partitioned.find_partition(&"a".into()).unwrap();
         assert_eq!(
-            &Simplify::simplify(split_a.expr.clone()).unwrap(),
+            &simplify(split_a.expr.clone()).unwrap(),
             &pack(
                 vec!["0".into(), "1".into()],
                 vec![get_item("a", ident()), get_item("b", ident())]
             )
         );
         let split_c = partitioned.find_partition(&"c".into()).unwrap();
         assert_eq!(
-            &Simplify::simplify(split_c.expr.clone()).unwrap(),
+            &simplify(split_c.expr.clone()).unwrap(),
             &pack(vec!["0".into()], vec![ident()])
         )
     }

diff --git a/vortex-expr/src/transform/simplify.rs b/vortex-expr/src/transform/simplify.rs
@@ -1,18 +1,16 @@
 use vortex_error::VortexResult;
 
 use crate::traversal::{FoldChildren, FoldUp, FolderMut, Node};
-use crate::{ExprRef, GetItem, Pack};
+use crate::{get_item, ident, Column, ExprRef, GetItem, Pack};
 
-pub struct Simplify;
-
-impl Simplify {
-    pub fn simplify(e: ExprRef) -> VortexResult<ExprRef> {
-        let mut folder = Simplify;
-        e.transform_with_context(&mut folder, ())
-            .map(|e| e.result())
-    }
+pub fn simplify(e: ExprRef) -> VortexResult<ExprRef> {
+    let mut folder = Simplify;
+    e.transform_with_context(&mut folder, ())
+        .map(|e| e.result())
 }
 
+struct Simplify;
+
 impl FolderMut for Simplify {
     type NodeTy = ExprRef;
     type Out = ExprRef;
@@ -24,6 +22,11 @@ impl FolderMut for Simplify {
         _context: Self::Context,
         children: FoldChildren<Self::Out>,
     ) -> VortexResult<FoldUp<Self::Out>> {
+        if let Some(column) = node.as_any().downcast_ref::<Column>() {
+            // TODO(ngates): deprecate Column, or keep at and simplify GetItem(Ident).
+            return Ok(FoldUp::Continue(get_item(column.field().clone(), ident())));
+        }
+
         if let Some(get_item) = node.as_any().downcast_ref::<GetItem>() {
             if let Some(pack) = get_item.child().as_any().downcast_ref::<Pack>() {
                 let expr = pack.field(get_item.field())?;