partition/
expr.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
15use std::collections::HashSet;
16use std::fmt::{Debug, Display, Formatter};
17use std::sync::Arc;
18
19use api::v1::meta::Partition;
20use datafusion_common::{ScalarValue, ToDFSchema};
21use datafusion_expr::Expr;
22use datafusion_expr::execution_props::ExecutionProps;
23use datafusion_expr::physical_planning_context::PhysicalPlanningContext;
24use datafusion_physical_expr::{PhysicalExpr, create_physical_expr};
25use datatypes::arrow;
26use datatypes::value::{
27    Value, duration_to_scalar_value, time_to_scalar_value, timestamp_to_scalar_value,
28};
29use serde::{Deserialize, Serialize};
30use snafu::ResultExt;
31use sql::statements::value_to_sql_value;
32use sqlparser::ast::{BinaryOperator as ParserBinaryOperator, Expr as ParserExpr, Ident};
33
34use crate::error;
35use crate::partition::PartitionBound;
36
37/// Struct for partition expression. This can be converted back to sqlparser's [Expr].
38/// by [`Self::to_parser_expr`].
39///
40/// [Expr]: sqlparser::ast::Expr
41#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
42pub struct PartitionExpr {
43    pub lhs: Box<Operand>,
44    pub op: RestrictedOp,
45    pub rhs: Box<Operand>,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
49pub enum Operand {
50    Column(String),
51    Value(Value),
52    Expr(PartitionExpr),
53}
54
55pub fn col(column_name: impl Into<String>) -> Operand {
56    Operand::Column(column_name.into())
57}
58
59impl From<Value> for Operand {
60    fn from(value: Value) -> Self {
61        Operand::Value(value)
62    }
63}
64
65impl Operand {
66    pub fn try_as_logical_expr(&self) -> error::Result<Expr> {
67        match self {
68            Self::Column(c) => Ok(datafusion_expr::col(format!(r#""{}""#, c))),
69            Self::Value(v) => {
70                let scalar_value = match v {
71                    Value::Boolean(v) => ScalarValue::Boolean(Some(*v)),
72                    Value::UInt8(v) => ScalarValue::UInt8(Some(*v)),
73                    Value::UInt16(v) => ScalarValue::UInt16(Some(*v)),
74                    Value::UInt32(v) => ScalarValue::UInt32(Some(*v)),
75                    Value::UInt64(v) => ScalarValue::UInt64(Some(*v)),
76                    Value::Int8(v) => ScalarValue::Int8(Some(*v)),
77                    Value::Int16(v) => ScalarValue::Int16(Some(*v)),
78                    Value::Int32(v) => ScalarValue::Int32(Some(*v)),
79                    Value::Int64(v) => ScalarValue::Int64(Some(*v)),
80                    Value::Float32(v) => ScalarValue::Float32(Some(v.0)),
81                    Value::Float64(v) => ScalarValue::Float64(Some(v.0)),
82                    Value::String(v) => ScalarValue::Utf8(Some(v.as_utf8().to_string())),
83                    Value::Binary(v) => ScalarValue::Binary(Some(v.to_vec())),
84                    Value::Date(v) => ScalarValue::Date32(Some(v.val())),
85                    Value::Null => ScalarValue::Null,
86                    Value::Timestamp(t) => timestamp_to_scalar_value(t.unit(), Some(t.value())),
87                    Value::Time(t) => time_to_scalar_value(*t.unit(), Some(t.value()))
88                        .context(error::ConvertPartitionExprValueSnafu { value: v.clone() })?,
89                    Value::IntervalYearMonth(v) => ScalarValue::IntervalYearMonth(Some(v.to_i32())),
90                    Value::IntervalDayTime(v) => ScalarValue::IntervalDayTime(Some((*v).into())),
91                    Value::IntervalMonthDayNano(v) => {
92                        ScalarValue::IntervalMonthDayNano(Some((*v).into()))
93                    }
94                    Value::Duration(d) => duration_to_scalar_value(d.unit(), Some(d.value())),
95                    Value::Decimal128(d) => {
96                        let (v, p, s) = d.to_scalar_value();
97                        ScalarValue::Decimal128(v, p, s)
98                    }
99                    other => {
100                        return error::UnsupportedPartitionExprValueSnafu {
101                            value: other.clone(),
102                        }
103                        .fail();
104                    }
105                };
106                Ok(datafusion_expr::lit(scalar_value))
107            }
108            Self::Expr(e) => e.try_as_logical_expr(),
109        }
110    }
111
112    pub fn lt(self, rhs: impl Into<Self>) -> PartitionExpr {
113        PartitionExpr::new(self, RestrictedOp::Lt, rhs.into())
114    }
115
116    pub fn gt_eq(self, rhs: impl Into<Self>) -> PartitionExpr {
117        PartitionExpr::new(self, RestrictedOp::GtEq, rhs.into())
118    }
119
120    pub fn eq(self, rhs: impl Into<Self>) -> PartitionExpr {
121        PartitionExpr::new(self, RestrictedOp::Eq, rhs.into())
122    }
123
124    pub fn not_eq(self, rhs: impl Into<Self>) -> PartitionExpr {
125        PartitionExpr::new(self, RestrictedOp::NotEq, rhs.into())
126    }
127
128    pub fn gt(self, rhs: impl Into<Self>) -> PartitionExpr {
129        PartitionExpr::new(self, RestrictedOp::Gt, rhs.into())
130    }
131
132    pub fn lt_eq(self, rhs: impl Into<Self>) -> PartitionExpr {
133        PartitionExpr::new(self, RestrictedOp::LtEq, rhs.into())
134    }
135}
136
137impl Display for Operand {
138    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
139        match self {
140            Self::Column(v) => write!(f, "{v}"),
141            Self::Value(v) => write!(f, "{v}"),
142            Self::Expr(v) => write!(f, "{v}"),
143        }
144    }
145}
146
147/// A restricted set of [Operator](datafusion_expr::Operator) that can be used in
148/// partition expressions.
149#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
150pub enum RestrictedOp {
151    // Evaluate to binary
152    Eq,
153    NotEq,
154    Lt,
155    LtEq,
156    Gt,
157    GtEq,
158
159    // Conjunction
160    And,
161    Or,
162}
163
164impl RestrictedOp {
165    pub fn try_from_parser(op: &ParserBinaryOperator) -> Option<Self> {
166        match op {
167            ParserBinaryOperator::Eq => Some(Self::Eq),
168            ParserBinaryOperator::NotEq => Some(Self::NotEq),
169            ParserBinaryOperator::Lt => Some(Self::Lt),
170            ParserBinaryOperator::LtEq => Some(Self::LtEq),
171            ParserBinaryOperator::Gt => Some(Self::Gt),
172            ParserBinaryOperator::GtEq => Some(Self::GtEq),
173            ParserBinaryOperator::And => Some(Self::And),
174            ParserBinaryOperator::Or => Some(Self::Or),
175            _ => None,
176        }
177    }
178
179    pub fn to_parser_op(&self) -> ParserBinaryOperator {
180        match self {
181            Self::Eq => ParserBinaryOperator::Eq,
182            Self::NotEq => ParserBinaryOperator::NotEq,
183            Self::Lt => ParserBinaryOperator::Lt,
184            Self::LtEq => ParserBinaryOperator::LtEq,
185            Self::Gt => ParserBinaryOperator::Gt,
186            Self::GtEq => ParserBinaryOperator::GtEq,
187            Self::And => ParserBinaryOperator::And,
188            Self::Or => ParserBinaryOperator::Or,
189        }
190    }
191
192    fn invert_for_swap(&self) -> Self {
193        match self {
194            Self::Eq => Self::Eq,
195            Self::NotEq => Self::NotEq,
196            Self::Lt => Self::Gt,
197            Self::LtEq => Self::GtEq,
198            Self::Gt => Self::Lt,
199            Self::GtEq => Self::LtEq,
200            Self::And => Self::And,
201            Self::Or => Self::Or,
202        }
203    }
204}
205impl Display for RestrictedOp {
206    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
207        match self {
208            Self::Eq => write!(f, "="),
209            Self::NotEq => write!(f, "<>"),
210            Self::Lt => write!(f, "<"),
211            Self::LtEq => write!(f, "<="),
212            Self::Gt => write!(f, ">"),
213            Self::GtEq => write!(f, ">="),
214            Self::And => write!(f, "AND"),
215            Self::Or => write!(f, "OR"),
216        }
217    }
218}
219
220impl PartitionExpr {
221    pub fn new(lhs: Operand, op: RestrictedOp, rhs: Operand) -> Self {
222        Self {
223            lhs: Box::new(lhs),
224            op,
225            rhs: Box::new(rhs),
226        }
227        .canonicalize()
228    }
229
230    /// Canonicalize to `Column op Value` form when possible for consistent equality checks.
231    pub fn canonicalize(self) -> Self {
232        let lhs = Self::canonicalize_operand(*self.lhs);
233        let rhs = Self::canonicalize_operand(*self.rhs);
234        let mut expr = Self {
235            lhs: Box::new(lhs),
236            op: self.op,
237            rhs: Box::new(rhs),
238        };
239
240        if matches!(&*expr.lhs, Operand::Value(_)) && matches!(&*expr.rhs, Operand::Column(_)) {
241            std::mem::swap(&mut expr.lhs, &mut expr.rhs);
242            expr.op = expr.op.invert_for_swap();
243        }
244
245        expr
246    }
247
248    fn canonicalize_operand(operand: Operand) -> Operand {
249        match operand {
250            Operand::Expr(expr) => Operand::Expr(expr.canonicalize()),
251            other => other,
252        }
253    }
254
255    /// Convert [Self] back to sqlparser's [Expr]
256    ///
257    /// [Expr]: ParserExpr
258    pub fn to_parser_expr(&self) -> ParserExpr {
259        // Safety: Partition rule won't contains unsupported value type.
260        // Otherwise it will be rejected by the parser.
261        let lhs = match &*self.lhs {
262            Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())),
263            Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()),
264            Operand::Expr(e) => e.to_parser_expr(),
265        };
266
267        let rhs = match &*self.rhs {
268            Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())),
269            Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()),
270            Operand::Expr(e) => e.to_parser_expr(),
271        };
272
273        ParserExpr::BinaryOp {
274            left: Box::new(lhs),
275            op: self.op.to_parser_op(),
276            right: Box::new(rhs),
277        }
278    }
279
280    pub fn try_as_logical_expr(&self) -> error::Result<Expr> {
281        // Special handling for null equality.
282        // `col = NULL` -> `col IS NULL` to match SQL (DataFusion) semantics.
283        let lhs_is_null = matches!(self.lhs.as_ref(), Operand::Value(Value::Null));
284        let rhs_is_null = matches!(self.rhs.as_ref(), Operand::Value(Value::Null));
285
286        match (self.op.clone(), lhs_is_null, rhs_is_null) {
287            (RestrictedOp::Eq, _, true) => {
288                return Ok(self.lhs.try_as_logical_expr()?.is_null());
289            }
290            (RestrictedOp::Eq, true, _) => {
291                return Ok(self.rhs.try_as_logical_expr()?.is_null());
292            }
293            (RestrictedOp::NotEq, _, true) => {
294                return Ok(self.lhs.try_as_logical_expr()?.is_not_null());
295            }
296            (RestrictedOp::NotEq, true, _) => {
297                return Ok(self.rhs.try_as_logical_expr()?.is_not_null());
298            }
299            _ => {}
300        }
301
302        if matches!(
303            self.op,
304            RestrictedOp::Lt | RestrictedOp::LtEq | RestrictedOp::Gt | RestrictedOp::GtEq
305        ) {
306            // Keep filtering semantics aligned with direct PartitionExpr evaluation (null-first ordering).
307            // In DataFusion SQL semantics, range comparisons with NULL yield NULL, so we inject
308            // `OR col IS NULL` on the null-first side of the comparison.
309            if matches!(self.lhs.as_ref(), Operand::Column(_)) {
310                let column_expr = self.lhs.try_as_logical_expr()?;
311                let other_expr = self.rhs.try_as_logical_expr()?;
312                let base = match self.op {
313                    RestrictedOp::Lt => {
314                        column_expr.clone().lt(other_expr).or(column_expr.is_null())
315                    }
316                    RestrictedOp::LtEq => column_expr
317                        .clone()
318                        .lt_eq(other_expr)
319                        .or(column_expr.is_null()),
320                    RestrictedOp::Gt => column_expr
321                        .clone()
322                        .gt(other_expr)
323                        .and(column_expr.is_not_null()),
324                    RestrictedOp::GtEq => column_expr
325                        .clone()
326                        .gt_eq(other_expr)
327                        .and(column_expr.is_not_null()),
328                    _ => unreachable!(),
329                };
330                return Ok(base);
331            } else if matches!(self.rhs.as_ref(), Operand::Column(_)) {
332                let other_expr = self.lhs.try_as_logical_expr()?;
333                let column_expr = self.rhs.try_as_logical_expr()?;
334                let base = match self.op {
335                    RestrictedOp::Lt => other_expr
336                        .lt(column_expr.clone())
337                        .and(column_expr.is_not_null()),
338                    RestrictedOp::LtEq => other_expr
339                        .lt_eq(column_expr.clone())
340                        .and(column_expr.is_not_null()),
341                    RestrictedOp::Gt => {
342                        other_expr.gt(column_expr.clone()).or(column_expr.is_null())
343                    }
344                    RestrictedOp::GtEq => other_expr
345                        .gt_eq(column_expr.clone())
346                        .or(column_expr.is_null()),
347                    _ => unreachable!(),
348                };
349                return Ok(base);
350            }
351        }
352
353        // Normal cases handling, without NULL
354        let lhs = self.lhs.try_as_logical_expr()?;
355        let rhs = self.rhs.try_as_logical_expr()?;
356
357        let expr = match &self.op {
358            RestrictedOp::And => datafusion_expr::and(lhs, rhs),
359            RestrictedOp::Or => datafusion_expr::or(lhs, rhs),
360            RestrictedOp::Gt => lhs.gt(rhs),
361            RestrictedOp::GtEq => lhs.gt_eq(rhs),
362            RestrictedOp::Lt => lhs.lt(rhs),
363            RestrictedOp::LtEq => lhs.lt_eq(rhs),
364            RestrictedOp::Eq => lhs.eq(rhs),
365            RestrictedOp::NotEq => lhs.not_eq(rhs),
366        };
367        Ok(expr)
368    }
369
370    /// Get the left-hand side operand
371    pub fn lhs(&self) -> &Operand {
372        &self.lhs
373    }
374
375    /// Get the right-hand side operand
376    pub fn rhs(&self) -> &Operand {
377        &self.rhs
378    }
379
380    /// Get the operation
381    pub fn op(&self) -> &RestrictedOp {
382        &self.op
383    }
384
385    pub fn try_as_physical_expr(
386        &self,
387        schema: &arrow::datatypes::SchemaRef,
388    ) -> error::Result<Arc<dyn PhysicalExpr>> {
389        let df_schema = schema
390            .clone()
391            .to_dfschema_ref()
392            .context(error::ToDFSchemaSnafu)?;
393        let execution_props = &ExecutionProps::default();
394        let expr = self.try_as_logical_expr()?;
395        create_physical_expr(
396            &expr,
397            &df_schema,
398            execution_props,
399            &PhysicalPlanningContext::default(),
400        )
401        .context(error::CreatePhysicalExprSnafu)
402    }
403
404    pub fn and(self, rhs: PartitionExpr) -> PartitionExpr {
405        PartitionExpr::new(Operand::Expr(self), RestrictedOp::And, Operand::Expr(rhs))
406    }
407
408    /// Serializes `PartitionExpr` to json string.
409    ///
410    /// Wraps `PartitionBound::Expr` for compatibility.
411    pub fn as_json_str(&self) -> error::Result<String> {
412        serde_json::to_string(&PartitionBound::Expr(self.clone()))
413            .context(error::SerializeJsonSnafu)
414    }
415
416    /// Deserializes `PartitionExpr` from json string.
417    ///
418    /// Deserializes to `PartitionBound` for compatibility.
419    pub fn from_json_str(s: &str) -> error::Result<Option<Self>> {
420        if s.is_empty() {
421            return Ok(None);
422        }
423
424        let bound: PartitionBound = serde_json::from_str(s).context(error::DeserializeJsonSnafu)?;
425        match bound {
426            PartitionBound::Expr(expr) => Ok(Some(expr.canonicalize())),
427            _ => Ok(None),
428        }
429    }
430
431    /// Converts [Self] to [Partition].
432    pub fn as_pb_partition(&self) -> error::Result<Partition> {
433        Ok(Partition {
434            expression: self.as_json_str()?,
435            ..Default::default()
436        })
437    }
438
439    /// Collects all column names referenced by this expression.
440    pub fn collect_column_names(&self, columns: &mut HashSet<String>) {
441        Self::collect_operand_columns(&self.lhs, columns);
442        Self::collect_operand_columns(&self.rhs, columns);
443    }
444
445    fn collect_operand_columns(operand: &Operand, columns: &mut HashSet<String>) {
446        match operand {
447            Operand::Column(c) => {
448                columns.insert(c.clone());
449            }
450            Operand::Expr(e) => {
451                e.collect_column_names(columns);
452            }
453            Operand::Value(_) => {}
454        }
455    }
456}
457
458impl Display for PartitionExpr {
459    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
460        write!(f, "{} {} {}", self.lhs, self.op, self.rhs)
461    }
462}
463
464#[cfg(test)]
465mod tests {
466    use super::*;
467
468    #[test]
469    fn test_partition_expr() {
470        let cases = [
471            (
472                Operand::Column("a".to_string()),
473                RestrictedOp::Eq,
474                Operand::Value(Value::UInt32(10)),
475                "a = 10",
476            ),
477            (
478                Operand::Column("a".to_string()),
479                RestrictedOp::NotEq,
480                Operand::Value(Value::UInt32(10)),
481                "a <> 10",
482            ),
483            (
484                Operand::Column("a".to_string()),
485                RestrictedOp::Lt,
486                Operand::Value(Value::UInt32(10)),
487                "a < 10",
488            ),
489            (
490                Operand::Column("a".to_string()),
491                RestrictedOp::LtEq,
492                Operand::Value(Value::UInt32(10)),
493                "a <= 10",
494            ),
495            (
496                Operand::Column("a".to_string()),
497                RestrictedOp::Gt,
498                Operand::Value(Value::UInt32(10)),
499                "a > 10",
500            ),
501            (
502                Operand::Column("a".to_string()),
503                RestrictedOp::GtEq,
504                Operand::Value(Value::UInt32(10)),
505                "a >= 10",
506            ),
507            (
508                Operand::Column("a".to_string()),
509                RestrictedOp::And,
510                Operand::Column("b".to_string()),
511                "a AND b",
512            ),
513            (
514                Operand::Column("a".to_string()),
515                RestrictedOp::Or,
516                Operand::Column("b".to_string()),
517                "a OR b",
518            ),
519            (
520                Operand::Column("a".to_string()),
521                RestrictedOp::Or,
522                Operand::Expr(PartitionExpr::new(
523                    Operand::Column("c".to_string()),
524                    RestrictedOp::And,
525                    Operand::Column("d".to_string()),
526                )),
527                "a OR c AND d",
528            ),
529        ];
530
531        for case in cases {
532            let expr = PartitionExpr::new(case.0, case.1.clone(), case.2);
533            assert_eq!(case.3, expr.to_string());
534        }
535    }
536
537    #[test]
538    fn test_try_as_logical_expr_null_equality() {
539        let eq_expr = PartitionExpr::new(
540            Operand::Column("a".to_string()),
541            RestrictedOp::Eq,
542            Operand::Value(Value::Null),
543        );
544        assert_eq!(
545            eq_expr.try_as_logical_expr().unwrap().to_string(),
546            "a IS NULL"
547        );
548
549        let neq_expr = PartitionExpr::new(
550            Operand::Column("a".to_string()),
551            RestrictedOp::NotEq,
552            Operand::Value(Value::Null),
553        );
554        assert_eq!(
555            neq_expr.try_as_logical_expr().unwrap().to_string(),
556            "a IS NOT NULL"
557        );
558    }
559
560    #[test]
561    fn test_try_as_logical_expr_null_range_comparison() {
562        // Test Lt with column on LHS
563        let lt_expr = PartitionExpr::new(
564            Operand::Column("a".to_string()),
565            RestrictedOp::Lt,
566            Operand::Value(Value::Int64(10)),
567        );
568        assert_eq!(
569            lt_expr.try_as_logical_expr().unwrap().to_string(),
570            "a < Int64(10) OR a IS NULL"
571        );
572
573        // Test Lt with column on RHS
574        let lt_expr_rhs_column = PartitionExpr::new(
575            Operand::Value(Value::Int64(10)),
576            RestrictedOp::Lt,
577            Operand::Column("a".to_string()),
578        );
579        assert_eq!(
580            lt_expr_rhs_column
581                .try_as_logical_expr()
582                .unwrap()
583                .to_string(),
584            "a > Int64(10) AND a IS NOT NULL"
585        );
586
587        // Test Gt with column on LHS
588        let gt_expr = PartitionExpr::new(
589            Operand::Column("a".to_string()),
590            RestrictedOp::Gt,
591            Operand::Value(Value::Int64(10)),
592        );
593        assert_eq!(
594            gt_expr.try_as_logical_expr().unwrap().to_string(),
595            "a > Int64(10) AND a IS NOT NULL"
596        );
597
598        // Test Gt with column on RHS
599        let gt_expr_rhs_column = PartitionExpr::new(
600            Operand::Value(Value::Int64(10)),
601            RestrictedOp::Gt,
602            Operand::Column("a".to_string()),
603        );
604        assert_eq!(
605            gt_expr_rhs_column
606                .try_as_logical_expr()
607                .unwrap()
608                .to_string(),
609            "a < Int64(10) OR a IS NULL"
610        );
611
612        // Test GtEq with column on LHS
613        let gteq_expr = PartitionExpr::new(
614            Operand::Column("a".to_string()),
615            RestrictedOp::GtEq,
616            Operand::Value(Value::Int64(10)),
617        );
618        assert_eq!(
619            gteq_expr.try_as_logical_expr().unwrap().to_string(),
620            "a >= Int64(10) AND a IS NOT NULL"
621        );
622
623        // Test LtEq with column on LHS
624        let lteq_expr = PartitionExpr::new(
625            Operand::Column("a".to_string()),
626            RestrictedOp::LtEq,
627            Operand::Value(Value::Int64(10)),
628        );
629        assert_eq!(
630            lteq_expr.try_as_logical_expr().unwrap().to_string(),
631            "a <= Int64(10) OR a IS NULL"
632        );
633
634        let gteq_expr_rhs_column = PartitionExpr::new(
635            Operand::Value(Value::Int64(10)),
636            RestrictedOp::GtEq,
637            Operand::Column("a".to_string()),
638        );
639        assert_eq!(
640            gteq_expr_rhs_column
641                .try_as_logical_expr()
642                .unwrap()
643                .to_string(),
644            "a <= Int64(10) OR a IS NULL"
645        );
646
647        let lteq_expr_rhs_column = PartitionExpr::new(
648            Operand::Value(Value::Int64(10)),
649            RestrictedOp::LtEq,
650            Operand::Column("a".to_string()),
651        );
652        assert_eq!(
653            lteq_expr_rhs_column
654                .try_as_logical_expr()
655                .unwrap()
656                .to_string(),
657            "a >= Int64(10) AND a IS NOT NULL"
658        );
659
660        let and_expr = PartitionExpr::new(
661            Operand::Expr(PartitionExpr::new(
662                Operand::Column("a".to_string()),
663                RestrictedOp::LtEq,
664                Operand::Value(Value::Int64(10)),
665            )),
666            RestrictedOp::And,
667            Operand::Expr(PartitionExpr::new(
668                Operand::Column("b".to_string()),
669                RestrictedOp::Gt,
670                Operand::Value(Value::Int64(5)),
671            )),
672        );
673        assert_eq!(
674            and_expr.try_as_logical_expr().unwrap().to_string(),
675            "(a <= Int64(10) OR a IS NULL) AND b > Int64(5) AND b IS NOT NULL"
676        );
677
678        let and_expr = PartitionExpr::new(
679            Operand::Expr(PartitionExpr::new(
680                Operand::Column("a".to_string()),
681                RestrictedOp::LtEq,
682                Operand::Value(Value::Int64(10)),
683            )),
684            RestrictedOp::And,
685            Operand::Expr(PartitionExpr::new(
686                Operand::Column("a".to_string()),
687                RestrictedOp::Gt,
688                Operand::Value(Value::Int64(5)),
689            )),
690        );
691        assert_eq!(
692            and_expr.try_as_logical_expr().unwrap().to_string(),
693            "(a <= Int64(10) OR a IS NULL) AND a > Int64(5) AND a IS NOT NULL"
694        );
695
696        let and_expr_strict_lower = PartitionExpr::new(
697            Operand::Expr(PartitionExpr::new(
698                Operand::Column("a".to_string()),
699                RestrictedOp::Lt,
700                Operand::Value(Value::Int64(10)),
701            )),
702            RestrictedOp::And,
703            Operand::Expr(PartitionExpr::new(
704                Operand::Column("a".to_string()),
705                RestrictedOp::GtEq,
706                Operand::Value(Value::Int64(5)),
707            )),
708        );
709        assert_eq!(
710            and_expr_strict_lower
711                .try_as_logical_expr()
712                .unwrap()
713                .to_string(),
714            "(a < Int64(10) OR a IS NULL) AND a >= Int64(5) AND a IS NOT NULL"
715        );
716
717        let and_expr_rhs_column = PartitionExpr::new(
718            Operand::Expr(PartitionExpr::new(
719                Operand::Value(Value::Int64(10)),
720                RestrictedOp::GtEq,
721                Operand::Column("a".to_string()),
722            )),
723            RestrictedOp::And,
724            Operand::Expr(PartitionExpr::new(
725                Operand::Value(Value::Int64(5)),
726                RestrictedOp::Lt,
727                Operand::Column("a".to_string()),
728            )),
729        );
730        assert_eq!(
731            and_expr_rhs_column
732                .try_as_logical_expr()
733                .unwrap()
734                .to_string(),
735            "(a <= Int64(10) OR a IS NULL) AND a > Int64(5) AND a IS NOT NULL"
736        );
737
738        let or_expr_same_column = PartitionExpr::new(
739            Operand::Expr(PartitionExpr::new(
740                Operand::Column("a".to_string()),
741                RestrictedOp::LtEq,
742                Operand::Value(Value::Int64(10)),
743            )),
744            RestrictedOp::Or,
745            Operand::Expr(PartitionExpr::new(
746                Operand::Column("a".to_string()),
747                RestrictedOp::Gt,
748                Operand::Value(Value::Int64(5)),
749            )),
750        );
751        assert_eq!(
752            or_expr_same_column
753                .try_as_logical_expr()
754                .unwrap()
755                .to_string(),
756            "a <= Int64(10) OR a IS NULL OR a > Int64(5) AND a IS NOT NULL"
757        );
758    }
759
760    #[test]
761    fn test_try_as_logical_expr_rhs_column_without_canonicalize() {
762        let gt_expr_rhs_column = PartitionExpr {
763            lhs: Box::new(Operand::Value(Value::Int64(10))),
764            op: RestrictedOp::Gt,
765            rhs: Box::new(Operand::Column("a".to_string())),
766        };
767        assert_eq!(
768            gt_expr_rhs_column
769                .try_as_logical_expr()
770                .unwrap()
771                .to_string(),
772            "Int64(10) > a OR a IS NULL"
773        );
774
775        let gteq_expr_rhs_column = PartitionExpr {
776            lhs: Box::new(Operand::Value(Value::Int64(10))),
777            op: RestrictedOp::GtEq,
778            rhs: Box::new(Operand::Column("a".to_string())),
779        };
780        assert_eq!(
781            gteq_expr_rhs_column
782                .try_as_logical_expr()
783                .unwrap()
784                .to_string(),
785            "Int64(10) >= a OR a IS NULL"
786        );
787    }
788
789    #[test]
790    fn test_serde_partition_expr() {
791        let expr = PartitionExpr::new(
792            Operand::Column("a".to_string()),
793            RestrictedOp::Eq,
794            Operand::Value(Value::UInt32(10)),
795        );
796        let json = expr.as_json_str().unwrap();
797        assert_eq!(
798            json,
799            "{\"Expr\":{\"lhs\":{\"Column\":\"a\"},\"op\":\"Eq\",\"rhs\":{\"Value\":{\"UInt32\":10}}}}"
800        );
801
802        let json = r#"{"Expr":{"lhs":{"Column":"a"},"op":"GtEq","rhs":{"Value":{"UInt32":10}}}}"#;
803        let expr2 = PartitionExpr::from_json_str(json).unwrap().unwrap();
804        let expected = PartitionExpr::new(
805            Operand::Column("a".to_string()),
806            RestrictedOp::GtEq,
807            Operand::Value(Value::UInt32(10)),
808        );
809        assert_eq!(expr2, expected);
810
811        // empty string
812        let json = "";
813        let expr3 = PartitionExpr::from_json_str(json).unwrap();
814        assert!(expr3.is_none());
815
816        // variants other than Expr
817        let json = r#""MaxValue""#;
818        let expr4 = PartitionExpr::from_json_str(json).unwrap();
819        assert!(expr4.is_none());
820
821        let json = r#"{"Value":{"UInt32":10}}"#;
822        let expr5 = PartitionExpr::from_json_str(json).unwrap();
823        assert!(expr5.is_none());
824    }
825
826    #[test]
827    fn test_collect_column_names() {
828        // Simple expression: col_a = 1 should give {col_a}
829        let expr = col("a").eq(Value::Int64(1));
830        let mut columns = HashSet::new();
831        expr.collect_column_names(&mut columns);
832        assert_eq!(columns.len(), 1);
833        assert!(columns.contains("a"));
834
835        // Compound AND with same column: col_a >= 0 AND col_a < 10 should give {col_a}
836        let expr = col("a")
837            .gt_eq(Value::Int64(0))
838            .and(col("a").lt(Value::Int64(10)));
839        let mut columns = HashSet::new();
840        expr.collect_column_names(&mut columns);
841        assert_eq!(columns.len(), 1);
842        assert!(columns.contains("a"));
843
844        // Multiple columns: col_a >= 0 AND col_b < 10 should give {col_a, col_b}
845        let expr = col("a")
846            .gt_eq(Value::Int64(0))
847            .and(col("b").lt(Value::Int64(10)));
848        let mut columns = HashSet::new();
849        expr.collect_column_names(&mut columns);
850        assert_eq!(columns.len(), 2);
851        assert!(columns.contains("a"));
852        assert!(columns.contains("b"));
853
854        // Nested expression: (col_a >= 0 AND col_b < 10) AND col_c = 5
855        let expr = col("a")
856            .gt_eq(Value::Int64(0))
857            .and(col("b").lt(Value::Int64(10)))
858            .and(col("c").eq(Value::Int64(5)));
859        let mut columns = HashSet::new();
860        expr.collect_column_names(&mut columns);
861        assert_eq!(columns.len(), 3);
862        assert!(columns.contains("a"));
863        assert!(columns.contains("b"));
864        assert!(columns.contains("c"));
865    }
866}
partition/expr.rs

partition/
expr.rs
