wip

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

encodings/alp/src/alp/compute/expr_pushdown.rs

@@ -6,7 +6,7 @@ use vortex_array::compute::Operator;
 use vortex_array::expr::{Binary, Literal, Root, VTableExt, lit, root};
 use vortex_array::transform::{ArrayParentReduceRule, ArrayRuleContext};
 use vortex_array::{ArrayRef, IntoArray};
-use vortex_error::VortexResult;
+use vortex_error::{VortexExpect, VortexResult};
 use vortex_scalar::{PrimitiveScalar, Scalar};
 
 use super::compare_common::{EncodedComparison, encode_for_comparison};
@@ -57,29 +57,10 @@ impl ArrayParentReduceRule<ALPVTable, ExprVTable> for ALPExprPushdownRule {
             return Ok(None);
         };
 
-        // Get the comparison operator - only handle comparison operators
         let operator = binary_view.operator();
-        if !matches!(
-            operator,
-            vortex_array::expr::Operator::Eq
-                | vortex_array::expr::Operator::NotEq
-                | vortex_array::expr::Operator::Lt
-                | vortex_array::expr::Operator::Lte
-                | vortex_array::expr::Operator::Gt
-                | vortex_array::expr::Operator::Gte
-        ) {
-            return Ok(None);
-        }
 
-        // Convert to compute operator
-        let compute_op = match operator {
-            vortex_array::expr::Operator::Eq => Operator::Eq,
-            vortex_array::expr::Operator::NotEq => Operator::NotEq,
-            vortex_array::expr::Operator::Lt => Operator::Lt,
-            vortex_array::expr::Operator::Lte => Operator::Lte,
-            vortex_array::expr::Operator::Gt => Operator::Gt,
-            vortex_array::expr::Operator::Gte => Operator::Gte,
-            _ => return Ok(None),
+        let Some(compute_op) = operator.maybe_cmp_operator() else {
+            return Ok(None);
         };
 
         // Check if this is a comparison of root() with a literal
@@ -95,16 +76,13 @@ impl ArrayParentReduceRule<ALPVTable, ExprVTable> for ALPExprPushdownRule {
                 return Ok(None);
             };
 
-        // Get the literal scalar - literals evaluate to a constant array with one element
         let literal_value = literal_expr.as_::<Literal>().data().clone();
 
-        // Don't optimize nullable comparisons
         if literal_value.dtype().is_nullable() {
             return Ok(None);
         }
 
-        // Convert to primitive scalar
-        let Ok(pscalar) = PrimitiveScalar::try_from(&literal_value) else {
+        let Some(pscalar) =  literal_value.as_primitive_opt() else {
             return Ok(None);
         };
 
@@ -183,7 +161,7 @@ mod tests {
 
         // Apply the optimization
         let session = ArraySession::default();
-        crate::register_alp_rules(&session);
+        crate::initialize(&session);
         let expr_session = ExprSession::default();
         let optimizer = session.optimizer(ExprOptimizer::new(&expr_session));
         let optimized = optimizer.optimize_array(expr_array.into_array()).unwrap();
@@ -242,7 +220,7 @@ mod tests {
         let expr_array = ExprArray::new_infer_dtype(alp.clone().into_array(), expr).unwrap();
 
         let session = ArraySession::default();
-        crate::register_alp_rules(&session);
+        crate::initialize(&session);
         let expr_session = ExprSession::default();
         let optimizer = session.optimizer(ExprOptimizer::new(&expr_session));
         let optimized = optimizer.optimize_array(expr_array.into_array()).unwrap();
@@ -297,7 +275,7 @@ mod tests {
         assert!(expr_array.child().is::<ALPVTable>());
 
         let session = ArraySession::default();
-        crate::register_alp_rules(&session);
+        crate::initialize(&session);
         let expr_session = ExprSession::default();
         let optimizer = session.optimizer(ExprOptimizer::new(&expr_session));
         let optimized = optimizer.optimize_array(expr_array.into_array()).unwrap();
@@ -343,7 +321,7 @@ mod tests {
             ExprArray::new_infer_dtype(alp.clone().into_array(), expr.clone()).unwrap();
 
         let session = ArraySession::default();
-        crate::register_alp_rules(&session);
+        crate::initialize(&session);
         let expr_session = ExprSession::default();
         let optimizer = session.optimizer(ExprOptimizer::new(&expr_session));
         let optimized = optimizer.optimize_array(expr_array.into_array()).unwrap();
@@ -398,7 +376,7 @@ mod tests {
         let test_value = 0.06051f32;
 
         let session = ArraySession::default();
-        crate::register_alp_rules(&session);
+        crate::initialize(&session);
         let expr_session = ExprSession::default();
         let expr_optimizer = ExprOptimizer::new(&expr_session);
 

encodings/alp/src/lib.rs

@@ -18,14 +18,17 @@
 
 pub use alp::*;
 pub use alp_rd::*;
+use vortex_array::EncodingRef;
 
 mod alp;
 mod alp_rd;
 
 /// Register ALP optimization rules with an ArraySession.
-pub fn register_alp_rules(session: &vortex_array::ArraySession) {
+pub fn initialize(session: &vortex_array::ArraySession) {
     use vortex_array::arrays::{ExprEncoding, ExprVTable};
 
+    session.register(EncodingRef::new_ref(ALPEncoding.as_ref()));
+
     // Register the comparison pushdown rule for ALP arrays wrapped in ExprArray
     session.register_parent_rule::<ALPVTable, ExprVTable, _>(
         &ALPEncoding,

vortex-file/src/lib.rs

@@ -106,7 +106,7 @@ pub use footer::*;
 pub use forever_constant::*;
 pub use open::*;
 pub use strategy::*;
-use vortex_alp::{ALPEncoding, ALPRDEncoding, register_alp_rules};
+use vortex_alp::{ALPEncoding, ALPRDEncoding};
 use vortex_array::arrays::DictEncoding;
 use vortex_array::{ArraySessionExt, EncodingRef};
 use vortex_bytebool::ByteBoolEncoding;
@@ -160,7 +160,6 @@ mod forever_constant {
 /// Vortex "Editions" that may support different sets of encodings.
 pub fn register_default_encodings(session: &VortexSession) {
     session.arrays().register_many([
-        EncodingRef::new_ref(ALPEncoding.as_ref()),
         EncodingRef::new_ref(ALPRDEncoding.as_ref()),
         EncodingRef::new_ref(BitPackedEncoding.as_ref()),
         EncodingRef::new_ref(ByteBoolEncoding.as_ref()),
@@ -180,5 +179,5 @@ pub fn register_default_encodings(session: &VortexSession) {
         EncodingRef::new_ref(vortex_zstd::ZstdEncoding.as_ref()),
     ]);
 
-    register_alp_rules(&session.arrays())
+    vortex_alp::initialize(&session.arrays())
 }