diff --git a/vortex-expr/src/identity.rs b/vortex-expr/src/identity.rs index 6c684de85d..08df539576 100644 --- 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 = 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 index afb745fcf1..08bb57a3e5 100644 --- 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,7 +369,7 @@ 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())] @@ -377,7 +377,7 @@ mod tests { ); 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 index 6e021122ba..49df3532cf 100644 --- 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 { - let mut folder = Simplify; - e.transform_with_context(&mut folder, ()) - .map(|e| e.result()) - } +pub fn simplify(e: ExprRef) -> VortexResult { + 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, ) -> VortexResult> { + if let Some(column) = node.as_any().downcast_ref::() { + // 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::() { if let Some(pack) = get_item.child().as_any().downcast_ref::() { let expr = pack.field(get_item.field())?;