partition/
collider.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Provides a Collider tool to convert [`PartitionExpr`] into a form that is easier to operate by program.
16//!
17//! This mod provides the following major structs:
18//!
19//! - [`Collider`]: The main struct that converts [`PartitionExpr`].
20//! - [`AtomicExpr`]: An "atomic" Expression, which isn't composed (OR-ed) of other expressions.
21//! - [`NucleonExpr`]: A simplified expression representation.
22//! - [`GluonOp`]: Further restricted operation set.
23//!
24//! On the naming aspect, "collider" is a high-energy machine that cracks particles, "atomic" is a typical
25//! non-divisible particle before ~100 years ago, "nucleon" is what composes an atom and "gluon" is the
26//! force inside nucleons.
27
28use std::collections::HashMap;
29use std::fmt::{Debug, Display, Formatter};
30use std::sync::Arc;
31
32use datafusion_expr::Operator;
33use datafusion_physical_expr::PhysicalExpr;
34use datafusion_physical_expr::expressions::{BinaryExpr, col, lit};
35use datatypes::arrow::datatypes::Schema;
36use datatypes::value::{OrderedF64, OrderedFloat, Value};
37
38use crate::error;
39use crate::error::Result;
40use crate::expr::{Operand, PartitionExpr, RestrictedOp};
41use crate::function::IcebergHash;
42
43const ZERO: OrderedF64 = OrderedFloat(0.0f64);
44pub(crate) const NORMALIZE_STEP: OrderedF64 = OrderedFloat(1.0f64);
45pub(crate) const CHECK_STEP: OrderedF64 = OrderedFloat(0.5f64);
46
47/// A dimension of the partition space: a column, or the hash of a column.
48#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
49pub enum PartitionDim {
50    Column(String),
51    IcebergHash(IcebergHash),
52}
53
54impl PartitionDim {
55    fn from_operand(operand: &Operand) -> Option<Self> {
56        match operand {
57            Operand::Column(column) => Some(Self::Column(column.clone())),
58            Operand::IcebergHash(hash) => Some(Self::IcebergHash(hash.clone())),
59            Operand::Value(_) | Operand::Expr(_) => None,
60        }
61    }
62
63    pub fn to_operand(&self) -> Operand {
64        match self {
65            Self::Column(column) => Operand::Column(column.clone()),
66            Self::IcebergHash(hash) => Operand::IcebergHash(hash.clone()),
67        }
68    }
69}
70
71impl Display for PartitionDim {
72    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
73        match self {
74            Self::Column(column) => write!(f, "{column}"),
75            Self::IcebergHash(hash) => write!(f, "{hash}"),
76        }
77    }
78}
79
80/// Represents an "atomic" Expression, which isn't composed (OR-ed) of other expressions.
81#[allow(unused)]
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub struct AtomicExpr {
84    /// A (ordered) list of simplified expressions. They are [`RestrictedOp::And`]'ed together.
85    pub nucleons: Vec<NucleonExpr>,
86    /// Index to reference the [`PartitionExpr`] that this [`AtomicExpr`] is derived from.
87    /// This index is used with `exprs` field in [`MultiDimPartitionRule`](crate::multi_dim::MultiDimPartitionRule).
88    pub source_expr_index: usize,
89}
90
91impl AtomicExpr {
92    /// Builds the conjunction over `schema`, where `field_name` maps each dimension to its field.
93    pub fn to_physical_expr(
94        &self,
95        schema: &Schema,
96        field_name: impl Fn(&PartitionDim) -> String,
97    ) -> Arc<dyn PhysicalExpr> {
98        let mut exprs = Vec::with_capacity(self.nucleons.len());
99        for nucleon in &self.nucleons {
100            exprs.push(nucleon.to_physical_expr(schema, &field_name(&nucleon.column)));
101        }
102        let result: Arc<dyn PhysicalExpr> = exprs
103            .into_iter()
104            .reduce(|l, r| Arc::new(BinaryExpr::new(l, Operator::And, r)))
105            .unwrap();
106        result
107    }
108}
109
110impl PartialOrd for AtomicExpr {
111    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
112        Some(self.nucleons.cmp(&other.nucleons))
113    }
114}
115
116/// A simplified expression representation.
117///
118/// This struct is used to compose [`AtomicExpr`], hence "nucleon".
119#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
120pub struct NucleonExpr {
121    column: PartitionDim,
122    op: GluonOp,
123    /// Normalized [`Value`].
124    value: OrderedF64,
125}
126
127impl NucleonExpr {
128    pub fn to_physical_expr(&self, schema: &Schema, field_name: &str) -> Arc<dyn PhysicalExpr> {
129        Arc::new(BinaryExpr::new(
130            col(field_name, schema).unwrap(),
131            self.op.to_operator(),
132            lit(*self.value.as_ref()),
133        ))
134    }
135
136    /// Get the dimension
137    pub fn column(&self) -> &PartitionDim {
138        &self.column
139    }
140
141    /// Get the normalized value
142    pub fn value(&self) -> OrderedF64 {
143        self.value
144    }
145
146    /// Get the operation
147    pub fn op(&self) -> &GluonOp {
148        &self.op
149    }
150
151    pub fn new(column: PartitionDim, op: GluonOp, value: OrderedF64) -> Self {
152        Self { column, op, value }
153    }
154}
155
156/// Further restricted operation set.
157///
158/// Conjunction operations are removed from [`RestrictedOp`].
159/// This enumeration is used to bind elements in [`NucleonExpr`], hence "gluon".
160#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
161pub enum GluonOp {
162    Eq,
163    NotEq,
164    Lt,
165    LtEq,
166    Gt,
167    GtEq,
168}
169
170impl GluonOp {
171    pub fn to_operator(&self) -> Operator {
172        match self {
173            GluonOp::Eq => Operator::Eq,
174            GluonOp::NotEq => Operator::NotEq,
175            GluonOp::Lt => Operator::Lt,
176            GluonOp::LtEq => Operator::LtEq,
177            GluonOp::Gt => Operator::Gt,
178            GluonOp::GtEq => Operator::GtEq,
179        }
180    }
181}
182
183/// Collider is used to collide a list of [`PartitionExpr`] into a list of [`AtomicExpr`]
184///
185/// It also normalizes the values of the columns in the expressions.
186#[allow(unused)]
187pub struct Collider<'a> {
188    source_exprs: &'a [PartitionExpr],
189
190    pub atomic_exprs: Vec<AtomicExpr>,
191    /// A map of dimension to a list of `(value, normalized value)` pairs.
192    ///
193    /// The normalized value is used for comparison. The normalization process keeps the order of the values.
194    pub normalized_values: HashMap<PartitionDim, Vec<(Value, OrderedF64)>>,
195}
196
197impl<'a> Collider<'a> {
198    pub fn new(source_exprs: &'a [PartitionExpr]) -> Result<Self> {
199        // first walk to collect all values
200        let mut values: HashMap<PartitionDim, Vec<Value>> = HashMap::new();
201        for expr in source_exprs {
202            Self::collect_column_values_from_expr(expr, &mut values)?;
203        }
204
205        // normalize values, assumes all values on a column are the same type
206        let mut normalized_values: HashMap<PartitionDim, HashMap<Value, OrderedF64>> =
207            HashMap::with_capacity(values.len());
208        for (column, mut column_values) in values {
209            column_values.sort_unstable();
210            column_values.dedup(); // Remove duplicates
211
212            // allowed because we have carefully implemented `Hash` to eliminate the mutable
213            #[allow(clippy::mutable_key_type)]
214            let mut value_map = HashMap::with_capacity(column_values.len());
215            let mut start_value = ZERO;
216            for value in column_values {
217                value_map.insert(value, start_value);
218                start_value += NORMALIZE_STEP;
219            }
220            normalized_values.insert(column, value_map);
221        }
222
223        // second walk to get atomic exprs
224        let mut atomic_exprs = Vec::with_capacity(source_exprs.len());
225        for (index, expr) in source_exprs.iter().enumerate() {
226            Self::collide_expr(expr, index, &normalized_values, &mut atomic_exprs)?;
227        }
228        // sort nucleon exprs
229        for expr in &mut atomic_exprs {
230            expr.nucleons.sort_unstable();
231        }
232
233        // convert normalized values to a map
234        let normalized_values = normalized_values
235            .into_iter()
236            .map(|(col, values)| {
237                let mut values = values.into_iter().collect::<Vec<_>>();
238                values.sort_unstable_by_key(|(_, v)| *v);
239                (col, values)
240            })
241            .collect();
242
243        Ok(Self {
244            source_exprs,
245            atomic_exprs,
246            normalized_values,
247        })
248    }
249
250    /// Helper to collect values with their associated columns from an expression
251    fn collect_column_values_from_expr(
252        expr: &PartitionExpr,
253        values: &mut HashMap<PartitionDim, Vec<Value>>,
254    ) -> Result<()> {
255        // Handle binary operations between column and value
256        match (&*expr.lhs, &*expr.rhs) {
257            (dim, Operand::Value(val)) | (Operand::Value(val), dim)
258                if let Some(dim) = PartitionDim::from_operand(dim) =>
259            {
260                values.entry(dim).or_default().push(val.clone());
261                Ok(())
262            }
263            (Operand::Expr(left_expr), Operand::Expr(right_expr)) => {
264                Self::collect_column_values_from_expr(left_expr, values)?;
265                Self::collect_column_values_from_expr(right_expr, values)
266            }
267            // Other combinations don't directly contribute column-value pairs
268            _ => error::InvalidExprSnafu { expr: expr.clone() }.fail(),
269        }
270    }
271
272    /// Collide a [`PartitionExpr`] into multiple [`AtomicExpr`]s.
273    ///
274    /// Split the [`PartitionExpr`] on every [`RestrictedOp::Or`] (disjunction), each branch is an [`AtomicExpr`].
275    /// Since [`PartitionExpr`] doesn't allow parentheses, Expression like `(a = 1 OR b = 2) AND c = 3` won't occur.
276    /// We can safely split on every [`RestrictedOp::Or`].
277    fn collide_expr(
278        expr: &PartitionExpr,
279        index: usize,
280        normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
281        result: &mut Vec<AtomicExpr>,
282    ) -> Result<()> {
283        match expr.op {
284            RestrictedOp::Or => {
285                // Split on OR operation - each side becomes a separate atomic expression
286
287                // Process left side
288                match &*expr.lhs {
289                    Operand::Expr(left_expr) => {
290                        Self::collide_expr(left_expr, index, normalized_values, result)?;
291                    }
292                    _ => {
293                        // Single operand - this shouldn't happen with OR
294                        // OR should always connect two sub-expressions
295                        return error::InvalidExprSnafu { expr: expr.clone() }.fail();
296                    }
297                }
298
299                // Process right side
300                match &*expr.rhs {
301                    Operand::Expr(right_expr) => {
302                        Self::collide_expr(right_expr, index, normalized_values, result)?;
303                    }
304                    _ => {
305                        // Single operand - this shouldn't happen with OR
306                        // OR should always connect two sub-expressions
307                        return error::InvalidExprSnafu { expr: expr.clone() }.fail();
308                    }
309                }
310            }
311            RestrictedOp::And => {
312                // For AND operations, we need to combine nucleons
313                let mut nucleons = Vec::new();
314                Self::collect_nucleons_from_expr(expr, &mut nucleons, normalized_values)?;
315
316                result.push(AtomicExpr {
317                    nucleons,
318                    source_expr_index: index,
319                });
320            }
321            _ => {
322                // For other operations, create a single atomic expression
323                let mut nucleons = Vec::new();
324                Self::collect_nucleons_from_expr(expr, &mut nucleons, normalized_values)?;
325
326                result.push(AtomicExpr {
327                    nucleons,
328                    source_expr_index: index,
329                });
330            }
331        }
332        Ok(())
333    }
334
335    /// Collect nucleons from an expression (handles AND operations recursively)
336    fn collect_nucleons_from_expr(
337        expr: &PartitionExpr,
338        nucleons: &mut Vec<NucleonExpr>,
339        normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
340    ) -> Result<()> {
341        match expr.op {
342            RestrictedOp::And => {
343                // For AND operations, collect nucleons from both sides
344                Self::collect_nucleons_from_operand(&expr.lhs, nucleons, normalized_values)?;
345                Self::collect_nucleons_from_operand(&expr.rhs, nucleons, normalized_values)?;
346            }
347            _ => {
348                // For non-AND operations, try to create a nucleon directly
349                nucleons.push(Self::try_create_nucleon(
350                    &expr.lhs,
351                    &expr.op,
352                    &expr.rhs,
353                    normalized_values,
354                )?);
355            }
356        }
357        Ok(())
358    }
359
360    /// Collect nucleons from an operand
361    fn collect_nucleons_from_operand(
362        operand: &Operand,
363        nucleons: &mut Vec<NucleonExpr>,
364        normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
365    ) -> Result<()> {
366        match operand {
367            Operand::Expr(expr) => {
368                Self::collect_nucleons_from_expr(expr, nucleons, normalized_values)
369            }
370            _ => {
371                // Only `Operand::Expr` can be conjuncted by AND.
372                error::NoExprOperandSnafu {
373                    operand: operand.clone(),
374                }
375                .fail()
376            }
377        }
378    }
379
380    /// Try to create a nucleon from operands
381    fn try_create_nucleon(
382        lhs: &Operand,
383        op: &RestrictedOp,
384        rhs: &Operand,
385        normalized_values: &HashMap<PartitionDim, HashMap<Value, OrderedF64>>,
386    ) -> Result<NucleonExpr> {
387        let gluon_op = match op {
388            RestrictedOp::Eq => GluonOp::Eq,
389            RestrictedOp::NotEq => GluonOp::NotEq,
390            RestrictedOp::Lt => GluonOp::Lt,
391            RestrictedOp::LtEq => GluonOp::LtEq,
392            RestrictedOp::Gt => GluonOp::Gt,
393            RestrictedOp::GtEq => GluonOp::GtEq,
394            RestrictedOp::And | RestrictedOp::Or => {
395                // These should be handled elsewhere
396                return error::UnexpectedSnafu {
397                    err_msg: format!("Conjunction operation {:?} should be handled elsewhere", op),
398                }
399                .fail();
400            }
401        };
402
403        match (lhs, rhs) {
404            (dim, Operand::Value(val)) if let Some(dim) = PartitionDim::from_operand(dim) => {
405                if let Some(column_values) = normalized_values.get(&dim)
406                    && let Some(&normalized_val) = column_values.get(val)
407                {
408                    return Ok(NucleonExpr {
409                        column: dim,
410                        op: gluon_op,
411                        value: normalized_val,
412                    });
413                }
414            }
415            (Operand::Value(val), dim) if let Some(dim) = PartitionDim::from_operand(dim) => {
416                if let Some(column_values) = normalized_values.get(&dim)
417                    && let Some(&normalized_val) = column_values.get(val)
418                {
419                    // Flip the operation for value op column
420                    let flipped_op = match gluon_op {
421                        GluonOp::Lt => GluonOp::Gt,
422                        GluonOp::LtEq => GluonOp::GtEq,
423                        GluonOp::Gt => GluonOp::Lt,
424                        GluonOp::GtEq => GluonOp::LtEq,
425                        op => op, // Eq and NotEq remain the same
426                    };
427                    return Ok(NucleonExpr {
428                        column: dim,
429                        op: flipped_op,
430                        value: normalized_val,
431                    });
432                }
433            }
434            _ => {}
435        }
436
437        // Other combinations not supported for nucleons
438        error::InvalidExprSnafu {
439            expr: PartitionExpr::new(lhs.clone(), op.clone(), rhs.clone()),
440        }
441        .fail()
442    }
443}
444
445#[cfg(test)]
446mod test {
447    use super::*;
448    use crate::expr::col;
449
450    #[test]
451    fn test_collider_basic_value_normalization() {
452        // Test with different value types in different columns
453        let exprs = vec![
454            // Integer values
455            col("age").eq(Value::UInt32(25)),
456            col("age").eq(Value::UInt32(30)),
457            col("age").eq(Value::UInt32(25)), // Duplicate should be handled
458            // String values
459            col("name").eq(Value::String("alice".into())),
460            col("name").eq(Value::String("bob".into())),
461            // Boolean values
462            col("active").eq(Value::Boolean(true)),
463            col("active").eq(Value::Boolean(false)),
464            // Float values
465            col("score").eq(Value::Float64(OrderedFloat(95.5))),
466            col("score").eq(Value::Float64(OrderedFloat(87.2))),
467        ];
468
469        let collider = Collider::new(&exprs).expect("Failed to create collider");
470
471        // Check that we have the right number of columns
472        assert_eq!(collider.normalized_values.len(), 4);
473
474        // Check age column - should have 2 unique values (25, 30)
475        let age_values = &collider.normalized_values[&PartitionDim::Column("age".to_string())];
476        assert_eq!(age_values.len(), 2);
477        assert_eq!(
478            age_values,
479            &[
480                (Value::UInt32(25), OrderedFloat(0.0f64)),
481                (Value::UInt32(30), OrderedFloat(1.0f64))
482            ]
483        );
484
485        // Check name column - should have 2 values
486        let name_values = &collider.normalized_values[&PartitionDim::Column("name".to_string())];
487        assert_eq!(name_values.len(), 2);
488        assert_eq!(
489            name_values,
490            &[
491                (Value::String("alice".into()), OrderedFloat(0.0f64)),
492                (Value::String("bob".into()), OrderedFloat(1.0f64))
493            ]
494        );
495
496        // Check active column - should have 2 values
497        let active_values =
498            &collider.normalized_values[&PartitionDim::Column("active".to_string())];
499        assert_eq!(active_values.len(), 2);
500        assert_eq!(
501            active_values,
502            &[
503                (Value::Boolean(false), OrderedFloat(0.0f64)),
504                (Value::Boolean(true), OrderedFloat(1.0f64))
505            ]
506        );
507
508        // Check score column - should have 2 values
509        let score_values = &collider.normalized_values[&PartitionDim::Column("score".to_string())];
510        assert_eq!(score_values.len(), 2);
511        assert_eq!(
512            score_values,
513            &[
514                (Value::Float64(OrderedFloat(87.2)), OrderedFloat(0.0f64)),
515                (Value::Float64(OrderedFloat(95.5)), OrderedFloat(1.0f64))
516            ]
517        );
518    }
519
520    #[test]
521    fn test_collider_simple_expressions() {
522        // Test simple equality
523        let exprs = vec![col("id").eq(Value::UInt32(1))];
524
525        let collider = Collider::new(&exprs).unwrap();
526        assert_eq!(collider.atomic_exprs.len(), 1);
527        assert_eq!(collider.atomic_exprs[0].nucleons.len(), 1);
528        assert_eq!(collider.atomic_exprs[0].source_expr_index, 0);
529
530        // Test simple AND
531        let exprs = vec![
532            col("id")
533                .eq(Value::UInt32(1))
534                .and(col("status").eq(Value::String("active".into()))),
535        ];
536
537        let collider = Collider::new(&exprs).unwrap();
538        assert_eq!(collider.atomic_exprs.len(), 1);
539        assert_eq!(collider.atomic_exprs[0].nucleons.len(), 2);
540
541        // Test simple OR - should create 2 atomic expressions
542        let expr = PartitionExpr::new(
543            Operand::Expr(col("id").eq(Value::UInt32(1))),
544            RestrictedOp::Or,
545            Operand::Expr(col("id").eq(Value::UInt32(2))),
546        );
547        let exprs = vec![expr];
548
549        let collider = Collider::new(&exprs).unwrap();
550        assert_eq!(collider.atomic_exprs.len(), 2);
551        assert_eq!(collider.atomic_exprs[0].nucleons.len(), 1);
552        assert_eq!(collider.atomic_exprs[1].nucleons.len(), 1);
553    }
554
555    #[test]
556    fn test_collider_complex_nested_expressions() {
557        // Test: (id = 1 AND status = 'active') OR (id = 2 AND status = 'inactive') OR (id = 3)
558        let branch1 = col("id")
559            .eq(Value::UInt32(1))
560            .and(col("status").eq(Value::String("active".into())));
561        let branch2 = col("id")
562            .eq(Value::UInt32(2))
563            .and(col("status").eq(Value::String("inactive".into())));
564        let branch3 = col("id").eq(Value::UInt32(3));
565
566        let expr = PartitionExpr::new(
567            Operand::Expr(PartitionExpr::new(
568                Operand::Expr(branch1),
569                RestrictedOp::Or,
570                Operand::Expr(branch2),
571            )),
572            RestrictedOp::Or,
573            Operand::Expr(branch3),
574        );
575
576        let exprs = vec![expr];
577        let collider = Collider::new(&exprs).unwrap();
578
579        assert_eq!(collider.atomic_exprs.len(), 3);
580
581        let total_nucleons: usize = collider
582            .atomic_exprs
583            .iter()
584            .map(|ae| ae.nucleons.len())
585            .sum();
586        assert_eq!(total_nucleons, 5);
587    }
588
589    #[test]
590    fn test_collider_deep_nesting() {
591        // Test deeply nested AND operations: a = 1 AND b = 2 AND c = 3 AND d = 4
592        let expr = col("a")
593            .eq(Value::UInt32(1))
594            .and(col("b").eq(Value::UInt32(2)))
595            .and(col("c").eq(Value::UInt32(3)))
596            .and(col("d").eq(Value::UInt32(4)));
597
598        let exprs = vec![expr];
599        let collider = Collider::new(&exprs).unwrap();
600
601        assert_eq!(collider.atomic_exprs.len(), 1);
602        assert_eq!(collider.atomic_exprs[0].nucleons.len(), 4);
603
604        // All nucleons should have Eq operation
605        for nucleon in &collider.atomic_exprs[0].nucleons {
606            assert_eq!(nucleon.op, GluonOp::Eq);
607        }
608    }
609
610    #[test]
611    fn test_collider_multiple_expressions() {
612        // Test multiple separate expressions
613        let exprs = vec![
614            col("id").eq(Value::UInt32(1)),
615            col("name").eq(Value::String("alice".into())),
616            col("score").gt_eq(Value::Float64(OrderedFloat(90.0))),
617        ];
618
619        let collider = Collider::new(&exprs).unwrap();
620
621        // Should create 3 atomic expressions (one for each input expression)
622        assert_eq!(collider.atomic_exprs.len(), 3);
623
624        // Each should have exactly 1 nucleon
625        for atomic_expr in &collider.atomic_exprs {
626            assert_eq!(atomic_expr.nucleons.len(), 1);
627        }
628
629        // Check that source indices are correct
630        let indices: Vec<usize> = collider
631            .atomic_exprs
632            .iter()
633            .map(|ae| ae.source_expr_index)
634            .collect();
635        assert!(indices.contains(&0));
636        assert!(indices.contains(&1));
637        assert!(indices.contains(&2));
638    }
639
640    #[test]
641    fn test_collider_value_column_order() {
642        // Test expressions where value comes before column (should flip operation)
643        let expr1 = PartitionExpr::new(
644            Operand::Value(Value::UInt32(10)),
645            RestrictedOp::Lt,
646            Operand::Column("age".to_string()),
647        ); // 10 < age should become age > 10
648
649        let expr2 = PartitionExpr::new(
650            Operand::Value(Value::UInt32(20)),
651            RestrictedOp::GtEq,
652            Operand::Column("score".to_string()),
653        ); // 20 >= score should become score <= 20
654
655        let exprs = vec![expr1, expr2];
656        let collider = Collider::new(&exprs).unwrap();
657
658        assert_eq!(collider.atomic_exprs.len(), 2);
659
660        // Check that operations were flipped correctly
661        let operations: Vec<GluonOp> = collider
662            .atomic_exprs
663            .iter()
664            .map(|ae| ae.nucleons[0].op.clone())
665            .collect();
666
667        assert!(operations.contains(&GluonOp::Gt)); // 10 < age -> age > 10
668        assert!(operations.contains(&GluonOp::LtEq)); // 20 >= score -> score <= 20
669    }
670
671    #[test]
672    fn test_collider_complex_or_with_different_columns() {
673        // Test: (name = 'alice' AND age = 25) OR (status = 'active' AND score > 90)
674        let branch1 = col("name")
675            .eq(Value::String("alice".into()))
676            .and(col("age").eq(Value::UInt32(25)));
677
678        let branch2 = col("status")
679            .eq(Value::String("active".into()))
680            .and(PartitionExpr::new(
681                Operand::Column("score".to_string()),
682                RestrictedOp::Gt,
683                Operand::Value(Value::Float64(OrderedFloat(90.0))),
684            ));
685
686        let expr = PartitionExpr::new(
687            Operand::Expr(branch1),
688            RestrictedOp::Or,
689            Operand::Expr(branch2),
690        );
691
692        let exprs = vec![expr];
693        let collider = Collider::new(&exprs).expect("Failed to create collider");
694
695        // Should create 2 atomic expressions
696        assert_eq!(collider.atomic_exprs.len(), 2);
697
698        // Each atomic expression should have 2 nucleons
699        for atomic_expr in &collider.atomic_exprs {
700            assert_eq!(atomic_expr.nucleons.len(), 2);
701        }
702
703        // Should have normalized values for all 4 columns
704        assert_eq!(collider.normalized_values.len(), 4);
705        assert!(
706            collider
707                .normalized_values
708                .contains_key(&PartitionDim::Column("name".to_string()))
709        );
710        assert!(
711            collider
712                .normalized_values
713                .contains_key(&PartitionDim::Column("age".to_string()))
714        );
715        assert!(
716            collider
717                .normalized_values
718                .contains_key(&PartitionDim::Column("status".to_string()))
719        );
720        assert!(
721            collider
722                .normalized_values
723                .contains_key(&PartitionDim::Column("score".to_string()))
724        );
725    }
726
727    #[test]
728    fn test_try_create_nucleon_edge_cases() {
729        let normalized_values = HashMap::new();
730
731        // Test with AND operation
732        let result = Collider::try_create_nucleon(
733            &col("a"),
734            &RestrictedOp::And,
735            &Operand::Value(Value::UInt32(1)),
736            &normalized_values,
737        );
738        assert!(result.is_err());
739
740        // Test with OR operation
741        let result = Collider::try_create_nucleon(
742            &col("a"),
743            &RestrictedOp::Or,
744            &Operand::Value(Value::UInt32(1)),
745            &normalized_values,
746        );
747        assert!(result.is_err());
748
749        // Test with Column-Column
750        let result = Collider::try_create_nucleon(
751            &col("a"),
752            &RestrictedOp::Eq,
753            &col("b"),
754            &normalized_values,
755        );
756        assert!(result.is_err());
757
758        // Test with Value-Value
759        let result = Collider::try_create_nucleon(
760            &Operand::Value(Value::UInt32(1)),
761            &RestrictedOp::Eq,
762            &Operand::Value(Value::UInt32(2)),
763            &normalized_values,
764        );
765        assert!(result.is_err());
766
767        // Test empty expression list
768        let exprs = vec![];
769        let collider = Collider::new(&exprs).unwrap();
770        assert_eq!(collider.atomic_exprs.len(), 0);
771        assert_eq!(collider.normalized_values.len(), 0);
772    }
773}
partition/collider.rs

partition/
collider.rs
