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_expr::simplify::SimplifyContext;
25use datafusion_optimizer::simplify_expressions::ExprSimplifier;
26use datafusion_physical_expr::{PhysicalExpr, create_physical_expr};
27use datatypes::arrow;
28use datatypes::value::{
29    Value, duration_to_scalar_value, time_to_scalar_value, timestamp_to_scalar_value,
30};
31use serde::{Deserialize, Serialize};
32use snafu::ResultExt;
33use sql::statements::value_to_sql_value;
34use sqlparser::ast::{BinaryOperator as ParserBinaryOperator, Expr as ParserExpr, Ident};
35
36use crate::error;
37use crate::partition::PartitionBound;
38
39/// Struct for partition expression. This can be converted back to sqlparser's [Expr].
40/// by [`Self::to_parser_expr`].
41///
42/// [Expr]: sqlparser::ast::Expr
43#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
44pub struct PartitionExpr {
45    pub lhs: Box<Operand>,
46    pub op: RestrictedOp,
47    pub rhs: Box<Operand>,
48}
49
50#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
51pub enum Operand {
52    Column(String),
53    Value(Value),
54    Expr(PartitionExpr),
55}
56
57pub fn col(column_name: impl Into<String>) -> Operand {
58    Operand::Column(column_name.into())
59}
60
61impl From<Value> for Operand {
62    fn from(value: Value) -> Self {
63        Operand::Value(value)
64    }
65}
66
67impl Operand {
68    pub fn try_as_logical_expr(&self) -> error::Result<Expr> {
69        match self {
70            Self::Column(c) => Ok(datafusion_expr::col(format!(r#""{}""#, c))),
71            Self::Value(v) => {
72                let scalar_value = match v {
73                    Value::Boolean(v) => ScalarValue::Boolean(Some(*v)),
74                    Value::UInt8(v) => ScalarValue::UInt8(Some(*v)),
75                    Value::UInt16(v) => ScalarValue::UInt16(Some(*v)),
76                    Value::UInt32(v) => ScalarValue::UInt32(Some(*v)),
77                    Value::UInt64(v) => ScalarValue::UInt64(Some(*v)),
78                    Value::Int8(v) => ScalarValue::Int8(Some(*v)),
79                    Value::Int16(v) => ScalarValue::Int16(Some(*v)),
80                    Value::Int32(v) => ScalarValue::Int32(Some(*v)),
81                    Value::Int64(v) => ScalarValue::Int64(Some(*v)),
82                    Value::Float32(v) => ScalarValue::Float32(Some(v.0)),
83                    Value::Float64(v) => ScalarValue::Float64(Some(v.0)),
84                    Value::String(v) => ScalarValue::Utf8(Some(v.as_utf8().to_string())),
85                    Value::Binary(v) => ScalarValue::Binary(Some(v.to_vec())),
86                    Value::Date(v) => ScalarValue::Date32(Some(v.val())),
87                    Value::Null => ScalarValue::Null,
88                    Value::Timestamp(t) => timestamp_to_scalar_value(t.unit(), Some(t.value())),
89                    Value::Time(t) => time_to_scalar_value(*t.unit(), Some(t.value()))
90                        .context(error::ConvertPartitionExprValueSnafu { value: v.clone() })?,
91                    Value::IntervalYearMonth(v) => ScalarValue::IntervalYearMonth(Some(v.to_i32())),
92                    Value::IntervalDayTime(v) => ScalarValue::IntervalDayTime(Some((*v).into())),
93                    Value::IntervalMonthDayNano(v) => {
94                        ScalarValue::IntervalMonthDayNano(Some((*v).into()))
95                    }
96                    Value::Duration(d) => duration_to_scalar_value(d.unit(), Some(d.value())),
97                    Value::Decimal128(d) => {
98                        let (v, p, s) = d.to_scalar_value();
99                        ScalarValue::Decimal128(v, p, s)
100                    }
101                    other => {
102                        return error::UnsupportedPartitionExprValueSnafu {
103                            value: other.clone(),
104                        }
105                        .fail();
106                    }
107                };
108                Ok(datafusion_expr::lit(scalar_value))
109            }
110            Self::Expr(e) => e.try_as_logical_expr(),
111        }
112    }
113
114    pub fn lt(self, rhs: impl Into<Self>) -> PartitionExpr {
115        PartitionExpr::new(self, RestrictedOp::Lt, rhs.into())
116    }
117
118    pub fn gt_eq(self, rhs: impl Into<Self>) -> PartitionExpr {
119        PartitionExpr::new(self, RestrictedOp::GtEq, rhs.into())
120    }
121
122    pub fn eq(self, rhs: impl Into<Self>) -> PartitionExpr {
123        PartitionExpr::new(self, RestrictedOp::Eq, rhs.into())
124    }
125
126    pub fn not_eq(self, rhs: impl Into<Self>) -> PartitionExpr {
127        PartitionExpr::new(self, RestrictedOp::NotEq, rhs.into())
128    }
129
130    pub fn gt(self, rhs: impl Into<Self>) -> PartitionExpr {
131        PartitionExpr::new(self, RestrictedOp::Gt, rhs.into())
132    }
133
134    pub fn lt_eq(self, rhs: impl Into<Self>) -> PartitionExpr {
135        PartitionExpr::new(self, RestrictedOp::LtEq, rhs.into())
136    }
137}
138
139impl Display for Operand {
140    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
141        match self {
142            Self::Column(v) => write!(f, "{v}"),
143            Self::Value(v) => write!(f, "{v}"),
144            Self::Expr(v) => write!(f, "{v}"),
145        }
146    }
147}
148
149/// A restricted set of [Operator](datafusion_expr::Operator) that can be used in
150/// partition expressions.
151#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
152pub enum RestrictedOp {
153    // Evaluate to binary
154    Eq,
155    NotEq,
156    Lt,
157    LtEq,
158    Gt,
159    GtEq,
160
161    // Conjunction
162    And,
163    Or,
164}
165
166impl RestrictedOp {
167    pub fn try_from_parser(op: &ParserBinaryOperator) -> Option<Self> {
168        match op {
169            ParserBinaryOperator::Eq => Some(Self::Eq),
170            ParserBinaryOperator::NotEq => Some(Self::NotEq),
171            ParserBinaryOperator::Lt => Some(Self::Lt),
172            ParserBinaryOperator::LtEq => Some(Self::LtEq),
173            ParserBinaryOperator::Gt => Some(Self::Gt),
174            ParserBinaryOperator::GtEq => Some(Self::GtEq),
175            ParserBinaryOperator::And => Some(Self::And),
176            ParserBinaryOperator::Or => Some(Self::Or),
177            _ => None,
178        }
179    }
180
181    pub fn to_parser_op(&self) -> ParserBinaryOperator {
182        match self {
183            Self::Eq => ParserBinaryOperator::Eq,
184            Self::NotEq => ParserBinaryOperator::NotEq,
185            Self::Lt => ParserBinaryOperator::Lt,
186            Self::LtEq => ParserBinaryOperator::LtEq,
187            Self::Gt => ParserBinaryOperator::Gt,
188            Self::GtEq => ParserBinaryOperator::GtEq,
189            Self::And => ParserBinaryOperator::And,
190            Self::Or => ParserBinaryOperator::Or,
191        }
192    }
193
194    fn invert_for_swap(&self) -> Self {
195        match self {
196            Self::Eq => Self::Eq,
197            Self::NotEq => Self::NotEq,
198            Self::Lt => Self::Gt,
199            Self::LtEq => Self::GtEq,
200            Self::Gt => Self::Lt,
201            Self::GtEq => Self::LtEq,
202            Self::And => Self::And,
203            Self::Or => Self::Or,
204        }
205    }
206}
207impl Display for RestrictedOp {
208    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
209        match self {
210            Self::Eq => write!(f, "="),
211            Self::NotEq => write!(f, "<>"),
212            Self::Lt => write!(f, "<"),
213            Self::LtEq => write!(f, "<="),
214            Self::Gt => write!(f, ">"),
215            Self::GtEq => write!(f, ">="),
216            Self::And => write!(f, "AND"),
217            Self::Or => write!(f, "OR"),
218        }
219    }
220}
221
222impl PartitionExpr {
223    pub fn new(lhs: Operand, op: RestrictedOp, rhs: Operand) -> Self {
224        Self {
225            lhs: Box::new(lhs),
226            op,
227            rhs: Box::new(rhs),
228        }
229        .canonicalize()
230    }
231
232    /// Canonicalize to `Column op Value` form when possible for consistent equality checks.
233    pub fn canonicalize(self) -> Self {
234        let lhs = Self::canonicalize_operand(*self.lhs);
235        let rhs = Self::canonicalize_operand(*self.rhs);
236        let mut expr = Self {
237            lhs: Box::new(lhs),
238            op: self.op,
239            rhs: Box::new(rhs),
240        };
241
242        if matches!(&*expr.lhs, Operand::Value(_)) && matches!(&*expr.rhs, Operand::Column(_)) {
243            std::mem::swap(&mut expr.lhs, &mut expr.rhs);
244            expr.op = expr.op.invert_for_swap();
245        }
246
247        expr
248    }
249
250    fn canonicalize_operand(operand: Operand) -> Operand {
251        match operand {
252            Operand::Expr(expr) => Operand::Expr(expr.canonicalize()),
253            other => other,
254        }
255    }
256
257    /// Convert [Self] back to sqlparser's [Expr]
258    ///
259    /// [Expr]: ParserExpr
260    pub fn to_parser_expr(&self) -> ParserExpr {
261        // Safety: Partition rule won't contains unsupported value type.
262        // Otherwise it will be rejected by the parser.
263        let lhs = match &*self.lhs {
264            Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())),
265            Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()),
266            Operand::Expr(e) => e.to_parser_expr(),
267        };
268
269        let rhs = match &*self.rhs {
270            Operand::Column(c) => ParserExpr::Identifier(Ident::new(c.clone())),
271            Operand::Value(v) => ParserExpr::Value(value_to_sql_value(v).unwrap().into()),
272            Operand::Expr(e) => e.to_parser_expr(),
273        };
274
275        ParserExpr::BinaryOp {
276            left: Box::new(lhs),
277            op: self.op.to_parser_op(),
278            right: Box::new(rhs),
279        }
280    }
281
282    pub fn try_as_logical_expr(&self) -> error::Result<Expr> {
283        // Row routing compares `Value`s, which order NULL before every other value.
284        // Batch routing and region filters treat a NULL result as false, so only the
285        // comparisons that hold for NULL need an explicit `IS NULL`.
286        let comparison = match (self.lhs.as_ref(), self.rhs.as_ref()) {
287            (column @ Operand::Column(_), Operand::Value(value)) => {
288                Some((column, value, self.op.clone()))
289            }
290            (Operand::Value(value), column @ Operand::Column(_)) => {
291                Some((column, value, self.op.invert_for_swap()))
292            }
293            _ => None,
294        };
295        if let Some((column, value, op)) = comparison {
296            let column = column.try_as_logical_expr()?;
297            if matches!(value, Value::Null) {
298                return match op {
299                    RestrictedOp::Eq | RestrictedOp::LtEq => Ok(column.is_null()),
300                    RestrictedOp::NotEq | RestrictedOp::Gt => Ok(column.is_not_null()),
301                    RestrictedOp::Lt => Ok(datafusion_expr::lit(false)),
302                    RestrictedOp::GtEq => Ok(datafusion_expr::lit(true)),
303                    RestrictedOp::And | RestrictedOp::Or => {
304                        error::InvalidExprSnafu { expr: self.clone() }.fail()
305                    }
306                };
307            }
308            let bound = Operand::Value(value.clone()).try_as_logical_expr()?;
309            return match op {
310                RestrictedOp::Eq => Ok(column.eq(bound)),
311                RestrictedOp::NotEq => Ok(column.clone().not_eq(bound).or(column.is_null())),
312                RestrictedOp::Lt => Ok(column.clone().lt(bound).or(column.is_null())),
313                RestrictedOp::LtEq => Ok(column.clone().lt_eq(bound).or(column.is_null())),
314                RestrictedOp::Gt => Ok(column.gt(bound)),
315                RestrictedOp::GtEq => Ok(column.gt_eq(bound)),
316                RestrictedOp::And | RestrictedOp::Or => {
317                    error::InvalidExprSnafu { expr: self.clone() }.fail()
318                }
319            };
320        }
321
322        // Normal cases handling, without NULL
323        let lhs = self.lhs.try_as_logical_expr()?;
324        let rhs = self.rhs.try_as_logical_expr()?;
325
326        let expr = match &self.op {
327            RestrictedOp::And => datafusion_expr::and(lhs, rhs),
328            RestrictedOp::Or => datafusion_expr::or(lhs, rhs),
329            RestrictedOp::Gt => lhs.gt(rhs),
330            RestrictedOp::GtEq => lhs.gt_eq(rhs),
331            RestrictedOp::Lt => lhs.lt(rhs),
332            RestrictedOp::LtEq => lhs.lt_eq(rhs),
333            RestrictedOp::Eq => lhs.eq(rhs),
334            RestrictedOp::NotEq => lhs.not_eq(rhs),
335        };
336        Ok(expr)
337    }
338
339    /// Get the left-hand side operand
340    pub fn lhs(&self) -> &Operand {
341        &self.lhs
342    }
343
344    /// Get the right-hand side operand
345    pub fn rhs(&self) -> &Operand {
346        &self.rhs
347    }
348
349    /// Get the operation
350    pub fn op(&self) -> &RestrictedOp {
351        &self.op
352    }
353
354    pub fn try_as_physical_expr(
355        &self,
356        schema: &arrow::datatypes::SchemaRef,
357    ) -> error::Result<Arc<dyn PhysicalExpr>> {
358        let df_schema = schema
359            .clone()
360            .to_dfschema_ref()
361            .context(error::ToDFSchemaSnafu)?;
362        let expr = self.try_as_logical_expr()?;
363        // Persisted partition bounds keep the unit (and type) they were created
364        // with, so after a schema change like widening the time index unit the
365        // expression may mix types with the columns it is evaluated against.
366        // Coerce it against the schema so comparisons share a common type, the
367        // same way `PredicateGroup` does for region partition filters.
368        let simplifier = ExprSimplifier::new(
369            SimplifyContext::builder()
370                .with_schema(df_schema.clone())
371                .build(),
372        );
373        let expr = simplifier
374            .coerce(expr, &df_schema)
375            .context(error::CoercePartitionExprSnafu)?;
376        let execution_props = &ExecutionProps::default();
377        create_physical_expr(
378            &expr,
379            &df_schema,
380            execution_props,
381            &PhysicalPlanningContext::default(),
382        )
383        .context(error::CreatePhysicalExprSnafu)
384    }
385
386    pub fn and(self, rhs: PartitionExpr) -> PartitionExpr {
387        PartitionExpr::new(Operand::Expr(self), RestrictedOp::And, Operand::Expr(rhs))
388    }
389
390    /// Serializes `PartitionExpr` to json string.
391    ///
392    /// Wraps `PartitionBound::Expr` for compatibility.
393    pub fn as_json_str(&self) -> error::Result<String> {
394        serde_json::to_string(&PartitionBound::Expr(self.clone()))
395            .context(error::SerializeJsonSnafu)
396    }
397
398    /// Deserializes `PartitionExpr` from json string.
399    ///
400    /// Deserializes to `PartitionBound` for compatibility.
401    pub fn from_json_str(s: &str) -> error::Result<Option<Self>> {
402        if s.is_empty() {
403            return Ok(None);
404        }
405
406        let bound: PartitionBound = serde_json::from_str(s).context(error::DeserializeJsonSnafu)?;
407        match bound {
408            PartitionBound::Expr(expr) => Ok(Some(expr.canonicalize())),
409            _ => Ok(None),
410        }
411    }
412
413    /// Converts [Self] to [Partition].
414    pub fn as_pb_partition(&self) -> error::Result<Partition> {
415        Ok(Partition {
416            expression: self.as_json_str()?,
417            ..Default::default()
418        })
419    }
420
421    /// Collects all column names referenced by this expression.
422    pub fn collect_column_names(&self, columns: &mut HashSet<String>) {
423        Self::collect_operand_columns(&self.lhs, columns);
424        Self::collect_operand_columns(&self.rhs, columns);
425    }
426
427    fn collect_operand_columns(operand: &Operand, columns: &mut HashSet<String>) {
428        match operand {
429            Operand::Column(c) => {
430                columns.insert(c.clone());
431            }
432            Operand::Expr(e) => {
433                e.collect_column_names(columns);
434            }
435            Operand::Value(_) => {}
436        }
437    }
438}
439
440impl Display for PartitionExpr {
441    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
442        write!(f, "{} {} {}", self.lhs, self.op, self.rhs)
443    }
444}
445
446#[cfg(test)]
447mod tests {
448    use std::collections::HashMap;
449
450    use common_meta::ddl::alter_table::validate_region_partition_bounds;
451    use common_meta::error::Error;
452    use common_meta::rpc::router::{Region, RegionRoute};
453    use datatypes::data_type::ConcreteDataType;
454    use store_api::storage::RegionId;
455
456    use super::*;
457
458    #[test]
459    fn test_partition_expr() {
460        let cases = [
461            (
462                Operand::Column("a".to_string()),
463                RestrictedOp::Eq,
464                Operand::Value(Value::UInt32(10)),
465                "a = 10",
466            ),
467            (
468                Operand::Column("a".to_string()),
469                RestrictedOp::NotEq,
470                Operand::Value(Value::UInt32(10)),
471                "a <> 10",
472            ),
473            (
474                Operand::Column("a".to_string()),
475                RestrictedOp::Lt,
476                Operand::Value(Value::UInt32(10)),
477                "a < 10",
478            ),
479            (
480                Operand::Column("a".to_string()),
481                RestrictedOp::LtEq,
482                Operand::Value(Value::UInt32(10)),
483                "a <= 10",
484            ),
485            (
486                Operand::Column("a".to_string()),
487                RestrictedOp::Gt,
488                Operand::Value(Value::UInt32(10)),
489                "a > 10",
490            ),
491            (
492                Operand::Column("a".to_string()),
493                RestrictedOp::GtEq,
494                Operand::Value(Value::UInt32(10)),
495                "a >= 10",
496            ),
497            (
498                Operand::Column("a".to_string()),
499                RestrictedOp::And,
500                Operand::Column("b".to_string()),
501                "a AND b",
502            ),
503            (
504                Operand::Column("a".to_string()),
505                RestrictedOp::Or,
506                Operand::Column("b".to_string()),
507                "a OR b",
508            ),
509            (
510                Operand::Column("a".to_string()),
511                RestrictedOp::Or,
512                Operand::Expr(PartitionExpr::new(
513                    Operand::Column("c".to_string()),
514                    RestrictedOp::And,
515                    Operand::Column("d".to_string()),
516                )),
517                "a OR c AND d",
518            ),
519        ];
520
521        for case in cases {
522            let expr = PartitionExpr::new(case.0, case.1.clone(), case.2);
523            assert_eq!(case.3, expr.to_string());
524        }
525    }
526
527    #[test]
528    fn test_try_as_logical_expr_null_equality() {
529        let eq_expr = PartitionExpr::new(
530            Operand::Column("a".to_string()),
531            RestrictedOp::Eq,
532            Operand::Value(Value::Null),
533        );
534        assert_eq!(
535            eq_expr.try_as_logical_expr().unwrap().to_string(),
536            "a IS NULL"
537        );
538
539        let neq_expr = PartitionExpr::new(
540            Operand::Column("a".to_string()),
541            RestrictedOp::NotEq,
542            Operand::Value(Value::Null),
543        );
544        assert_eq!(
545            neq_expr.try_as_logical_expr().unwrap().to_string(),
546            "a IS NOT NULL"
547        );
548    }
549
550    #[test]
551    fn test_try_as_logical_expr_null_range_comparison() {
552        // Test Lt with column on LHS
553        let lt_expr = PartitionExpr::new(
554            Operand::Column("a".to_string()),
555            RestrictedOp::Lt,
556            Operand::Value(Value::Int64(10)),
557        );
558        assert_eq!(
559            lt_expr.try_as_logical_expr().unwrap().to_string(),
560            "a < Int64(10) OR a IS NULL"
561        );
562
563        // Test Lt with column on RHS
564        let lt_expr_rhs_column = PartitionExpr::new(
565            Operand::Value(Value::Int64(10)),
566            RestrictedOp::Lt,
567            Operand::Column("a".to_string()),
568        );
569        assert_eq!(
570            lt_expr_rhs_column
571                .try_as_logical_expr()
572                .unwrap()
573                .to_string(),
574            "a > Int64(10)"
575        );
576
577        // Test Gt with column on LHS
578        let gt_expr = PartitionExpr::new(
579            Operand::Column("a".to_string()),
580            RestrictedOp::Gt,
581            Operand::Value(Value::Int64(10)),
582        );
583        assert_eq!(
584            gt_expr.try_as_logical_expr().unwrap().to_string(),
585            "a > Int64(10)"
586        );
587
588        // Test Gt with column on RHS
589        let gt_expr_rhs_column = PartitionExpr::new(
590            Operand::Value(Value::Int64(10)),
591            RestrictedOp::Gt,
592            Operand::Column("a".to_string()),
593        );
594        assert_eq!(
595            gt_expr_rhs_column
596                .try_as_logical_expr()
597                .unwrap()
598                .to_string(),
599            "a < Int64(10) OR a IS NULL"
600        );
601
602        // Test GtEq with column on LHS
603        let gteq_expr = PartitionExpr::new(
604            Operand::Column("a".to_string()),
605            RestrictedOp::GtEq,
606            Operand::Value(Value::Int64(10)),
607        );
608        assert_eq!(
609            gteq_expr.try_as_logical_expr().unwrap().to_string(),
610            "a >= Int64(10)"
611        );
612
613        // Test LtEq with column on LHS
614        let lteq_expr = PartitionExpr::new(
615            Operand::Column("a".to_string()),
616            RestrictedOp::LtEq,
617            Operand::Value(Value::Int64(10)),
618        );
619        assert_eq!(
620            lteq_expr.try_as_logical_expr().unwrap().to_string(),
621            "a <= Int64(10) OR a IS NULL"
622        );
623
624        let gteq_expr_rhs_column = PartitionExpr::new(
625            Operand::Value(Value::Int64(10)),
626            RestrictedOp::GtEq,
627            Operand::Column("a".to_string()),
628        );
629        assert_eq!(
630            gteq_expr_rhs_column
631                .try_as_logical_expr()
632                .unwrap()
633                .to_string(),
634            "a <= Int64(10) OR a IS NULL"
635        );
636
637        let lteq_expr_rhs_column = PartitionExpr::new(
638            Operand::Value(Value::Int64(10)),
639            RestrictedOp::LtEq,
640            Operand::Column("a".to_string()),
641        );
642        assert_eq!(
643            lteq_expr_rhs_column
644                .try_as_logical_expr()
645                .unwrap()
646                .to_string(),
647            "a >= Int64(10)"
648        );
649
650        let and_expr = PartitionExpr::new(
651            Operand::Expr(PartitionExpr::new(
652                Operand::Column("a".to_string()),
653                RestrictedOp::LtEq,
654                Operand::Value(Value::Int64(10)),
655            )),
656            RestrictedOp::And,
657            Operand::Expr(PartitionExpr::new(
658                Operand::Column("b".to_string()),
659                RestrictedOp::Gt,
660                Operand::Value(Value::Int64(5)),
661            )),
662        );
663        assert_eq!(
664            and_expr.try_as_logical_expr().unwrap().to_string(),
665            "(a <= Int64(10) OR a IS NULL) AND b > Int64(5)"
666        );
667
668        let and_expr = PartitionExpr::new(
669            Operand::Expr(PartitionExpr::new(
670                Operand::Column("a".to_string()),
671                RestrictedOp::LtEq,
672                Operand::Value(Value::Int64(10)),
673            )),
674            RestrictedOp::And,
675            Operand::Expr(PartitionExpr::new(
676                Operand::Column("a".to_string()),
677                RestrictedOp::Gt,
678                Operand::Value(Value::Int64(5)),
679            )),
680        );
681        assert_eq!(
682            and_expr.try_as_logical_expr().unwrap().to_string(),
683            "(a <= Int64(10) OR a IS NULL) AND a > Int64(5)"
684        );
685
686        let and_expr_strict_lower = PartitionExpr::new(
687            Operand::Expr(PartitionExpr::new(
688                Operand::Column("a".to_string()),
689                RestrictedOp::Lt,
690                Operand::Value(Value::Int64(10)),
691            )),
692            RestrictedOp::And,
693            Operand::Expr(PartitionExpr::new(
694                Operand::Column("a".to_string()),
695                RestrictedOp::GtEq,
696                Operand::Value(Value::Int64(5)),
697            )),
698        );
699        assert_eq!(
700            and_expr_strict_lower
701                .try_as_logical_expr()
702                .unwrap()
703                .to_string(),
704            "(a < Int64(10) OR a IS NULL) AND a >= Int64(5)"
705        );
706
707        let and_expr_rhs_column = PartitionExpr::new(
708            Operand::Expr(PartitionExpr::new(
709                Operand::Value(Value::Int64(10)),
710                RestrictedOp::GtEq,
711                Operand::Column("a".to_string()),
712            )),
713            RestrictedOp::And,
714            Operand::Expr(PartitionExpr::new(
715                Operand::Value(Value::Int64(5)),
716                RestrictedOp::Lt,
717                Operand::Column("a".to_string()),
718            )),
719        );
720        assert_eq!(
721            and_expr_rhs_column
722                .try_as_logical_expr()
723                .unwrap()
724                .to_string(),
725            "(a <= Int64(10) OR a IS NULL) AND a > Int64(5)"
726        );
727
728        let or_expr_same_column = PartitionExpr::new(
729            Operand::Expr(PartitionExpr::new(
730                Operand::Column("a".to_string()),
731                RestrictedOp::LtEq,
732                Operand::Value(Value::Int64(10)),
733            )),
734            RestrictedOp::Or,
735            Operand::Expr(PartitionExpr::new(
736                Operand::Column("a".to_string()),
737                RestrictedOp::Gt,
738                Operand::Value(Value::Int64(5)),
739            )),
740        );
741        assert_eq!(
742            or_expr_same_column
743                .try_as_logical_expr()
744                .unwrap()
745                .to_string(),
746            "a <= Int64(10) OR a IS NULL OR a > Int64(5)"
747        );
748    }
749
750    #[test]
751    fn test_try_as_logical_expr_rhs_column_without_canonicalize() {
752        let gt_expr_rhs_column = PartitionExpr {
753            lhs: Box::new(Operand::Value(Value::Int64(10))),
754            op: RestrictedOp::Gt,
755            rhs: Box::new(Operand::Column("a".to_string())),
756        };
757        assert_eq!(
758            gt_expr_rhs_column
759                .try_as_logical_expr()
760                .unwrap()
761                .to_string(),
762            "a < Int64(10) OR a IS NULL"
763        );
764
765        let gteq_expr_rhs_column = PartitionExpr {
766            lhs: Box::new(Operand::Value(Value::Int64(10))),
767            op: RestrictedOp::GtEq,
768            rhs: Box::new(Operand::Column("a".to_string())),
769        };
770        assert_eq!(
771            gteq_expr_rhs_column
772                .try_as_logical_expr()
773                .unwrap()
774                .to_string(),
775            "a <= Int64(10) OR a IS NULL"
776        );
777    }
778
779    #[test]
780    fn test_serde_partition_expr() {
781        let expr = PartitionExpr::new(
782            Operand::Column("a".to_string()),
783            RestrictedOp::Eq,
784            Operand::Value(Value::UInt32(10)),
785        );
786        let json = expr.as_json_str().unwrap();
787        assert_eq!(
788            json,
789            "{\"Expr\":{\"lhs\":{\"Column\":\"a\"},\"op\":\"Eq\",\"rhs\":{\"Value\":{\"UInt32\":10}}}}"
790        );
791
792        let json = r#"{"Expr":{"lhs":{"Column":"a"},"op":"GtEq","rhs":{"Value":{"UInt32":10}}}}"#;
793        let expr2 = PartitionExpr::from_json_str(json).unwrap().unwrap();
794        let expected = PartitionExpr::new(
795            Operand::Column("a".to_string()),
796            RestrictedOp::GtEq,
797            Operand::Value(Value::UInt32(10)),
798        );
799        assert_eq!(expr2, expected);
800
801        // empty string
802        let json = "";
803        let expr3 = PartitionExpr::from_json_str(json).unwrap();
804        assert!(expr3.is_none());
805
806        // variants other than Expr
807        let json = r#""MaxValue""#;
808        let expr4 = PartitionExpr::from_json_str(json).unwrap();
809        assert!(expr4.is_none());
810
811        let json = r#"{"Value":{"UInt32":10}}"#;
812        let expr5 = PartitionExpr::from_json_str(json).unwrap();
813        assert!(expr5.is_none());
814    }
815
816    #[test]
817    fn test_collect_column_names() {
818        // Simple expression: col_a = 1 should give {col_a}
819        let expr = col("a").eq(Value::Int64(1));
820        let mut columns = HashSet::new();
821        expr.collect_column_names(&mut columns);
822        assert_eq!(columns.len(), 1);
823        assert!(columns.contains("a"));
824
825        // Compound AND with same column: col_a >= 0 AND col_a < 10 should give {col_a}
826        let expr = col("a")
827            .gt_eq(Value::Int64(0))
828            .and(col("a").lt(Value::Int64(10)));
829        let mut columns = HashSet::new();
830        expr.collect_column_names(&mut columns);
831        assert_eq!(columns.len(), 1);
832        assert!(columns.contains("a"));
833
834        // Multiple columns: col_a >= 0 AND col_b < 10 should give {col_a, col_b}
835        let expr = col("a")
836            .gt_eq(Value::Int64(0))
837            .and(col("b").lt(Value::Int64(10)));
838        let mut columns = HashSet::new();
839        expr.collect_column_names(&mut columns);
840        assert_eq!(columns.len(), 2);
841        assert!(columns.contains("a"));
842        assert!(columns.contains("b"));
843
844        // Nested expression: (col_a >= 0 AND col_b < 10) AND col_c = 5
845        let expr = col("a")
846            .gt_eq(Value::Int64(0))
847            .and(col("b").lt(Value::Int64(10)))
848            .and(col("c").eq(Value::Int64(5)));
849        let mut columns = HashSet::new();
850        expr.collect_column_names(&mut columns);
851        assert_eq!(columns.len(), 3);
852        assert!(columns.contains("a"));
853        assert!(columns.contains("b"));
854        assert!(columns.contains("c"));
855    }
856
857    /// Real serialized expressions must survive the round trip through
858    /// `common-meta`'s persisted-expression mirror: whenever that mirror drifts
859    /// from these types, deserialization fails, validation is skipped and the
860    /// overflowing cases below stop producing `PartitionExprIncompatible`.
861    #[test]
862    fn test_serialized_expressions_validate_against_type_changes() {
863        let overflow_ms = serde_json::from_str::<Value>(
864            r#"{"Timestamp":{"value":32503680000000,"unit":"Millisecond"}}"#,
865        )
866        .unwrap();
867        let start_ms =
868            serde_json::from_str::<Value>(r#"{"Timestamp":{"value":0,"unit":"Millisecond"}}"#)
869                .unwrap();
870        let ts_ms = ConcreteDataType::timestamp_millisecond_datatype();
871        let ts_ns = ConcreteDataType::timestamp_nanosecond_datatype();
872
873        let routes = |expr: &PartitionExpr| {
874            vec![RegionRoute {
875                region: Region {
876                    id: RegionId::new(1024, 1),
877                    partition_expr: expr.as_json_str().unwrap(),
878                    ..Default::default()
879                },
880                ..Default::default()
881            }]
882        };
883        let changed = |column: &str, old_type: ConcreteDataType, new_type: ConcreteDataType| {
884            HashMap::from([(column.to_string(), (old_type, new_type))])
885        };
886
887        // Compile-time tripwire: exhaustively matching the operand and operator
888        // variants makes a new variant break this test at compile time, so the
889        // `partition_compat::Persisted*` mirror and these cases must be extended
890        // together instead of silently skipping validation.
891        fn walk(expr: &PartitionExpr) {
892            walk_operand(&expr.lhs);
893            walk_operand(&expr.rhs);
894            match expr.op {
895                RestrictedOp::Eq
896                | RestrictedOp::NotEq
897                | RestrictedOp::Lt
898                | RestrictedOp::LtEq
899                | RestrictedOp::Gt
900                | RestrictedOp::GtEq
901                | RestrictedOp::And
902                | RestrictedOp::Or => {}
903            }
904        }
905        fn walk_operand(operand: &Operand) {
906            match operand {
907                Operand::Column(_) | Operand::Value(_) => {}
908                Operand::Expr(expr) => walk(expr),
909            }
910        }
911
912        // Every expression shape must be rejected while the bound overflows
913        // nanoseconds, with the error naming the offending column.
914        let simple = col("ts").lt(overflow_ms.clone());
915        let nested_and = col("ts")
916            .gt_eq(start_ms.clone())
917            .and(col("ts").lt(overflow_ms.clone()));
918        let nested_or = PartitionExpr::new(
919            Operand::Expr(col("host").eq(Value::String("a".into()))),
920            RestrictedOp::Or,
921            Operand::Expr(col("ts").lt(overflow_ms.clone())),
922        );
923        for expr in [&simple, &nested_and, &nested_or] {
924            walk(expr);
925            let err = validate_region_partition_bounds(
926                &changed("ts", ts_ms.clone(), ts_ns.clone()),
927                &routes(expr),
928            )
929            .unwrap_err();
930            assert!(
931                matches!(
932                    err,
933                    Error::PartitionExprIncompatible { ref column, .. } if column == "ts"
934                ),
935                "unexpected error: {err}"
936            );
937        }
938
939        // A representable bound passes.
940        let compatible = col("ts").lt(start_ms);
941        validate_region_partition_bounds(&changed("ts", ts_ms, ts_ns), &routes(&compatible))
942            .unwrap();
943
944        // Narrowing a field column must be rejected as well: the region's write
945        // path compares the persisted Int64 bound against Int32 rows directly
946        // and panics, even though the scan would coerce both sides.
947        let narrowing = col("v").lt(Value::Int64(3_000_000_000));
948        let err = validate_region_partition_bounds(
949            &changed(
950                "v",
951                ConcreteDataType::int64_datatype(),
952                ConcreteDataType::int32_datatype(),
953            ),
954            &routes(&narrowing),
955        )
956        .unwrap_err();
957        assert!(
958            matches!(
959                err,
960                Error::PartitionExprIncompatible { ref column, .. } if column == "v"
961            ),
962            "unexpected error: {err}"
963        );
964    }
965}
partition/expr.rs

partition/
expr.rs
