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        // Row routing compares `Value`s, which order NULL before every other value.
282        // Batch routing and region filters treat a NULL result as false, so only the
283        // comparisons that hold for NULL need an explicit `IS NULL`.
284        let comparison = match (self.lhs.as_ref(), self.rhs.as_ref()) {
285            (column @ Operand::Column(_), Operand::Value(value)) => {
286                Some((column, value, self.op.clone()))
287            }
288            (Operand::Value(value), column @ Operand::Column(_)) => {
289                Some((column, value, self.op.invert_for_swap()))
290            }
291            _ => None,
292        };
293        if let Some((column, value, op)) = comparison {
294            let column = column.try_as_logical_expr()?;
295            if matches!(value, Value::Null) {
296                return match op {
297                    RestrictedOp::Eq | RestrictedOp::LtEq => Ok(column.is_null()),
298                    RestrictedOp::NotEq | RestrictedOp::Gt => Ok(column.is_not_null()),
299                    RestrictedOp::Lt => Ok(datafusion_expr::lit(false)),
300                    RestrictedOp::GtEq => Ok(datafusion_expr::lit(true)),
301                    RestrictedOp::And | RestrictedOp::Or => {
302                        error::InvalidExprSnafu { expr: self.clone() }.fail()
303                    }
304                };
305            }
306            let bound = Operand::Value(value.clone()).try_as_logical_expr()?;
307            return match op {
308                RestrictedOp::Eq => Ok(column.eq(bound)),
309                RestrictedOp::NotEq => Ok(column.clone().not_eq(bound).or(column.is_null())),
310                RestrictedOp::Lt => Ok(column.clone().lt(bound).or(column.is_null())),
311                RestrictedOp::LtEq => Ok(column.clone().lt_eq(bound).or(column.is_null())),
312                RestrictedOp::Gt => Ok(column.gt(bound)),
313                RestrictedOp::GtEq => Ok(column.gt_eq(bound)),
314                RestrictedOp::And | RestrictedOp::Or => {
315                    error::InvalidExprSnafu { expr: self.clone() }.fail()
316                }
317            };
318        }
319
320        // Normal cases handling, without NULL
321        let lhs = self.lhs.try_as_logical_expr()?;
322        let rhs = self.rhs.try_as_logical_expr()?;
323
324        let expr = match &self.op {
325            RestrictedOp::And => datafusion_expr::and(lhs, rhs),
326            RestrictedOp::Or => datafusion_expr::or(lhs, rhs),
327            RestrictedOp::Gt => lhs.gt(rhs),
328            RestrictedOp::GtEq => lhs.gt_eq(rhs),
329            RestrictedOp::Lt => lhs.lt(rhs),
330            RestrictedOp::LtEq => lhs.lt_eq(rhs),
331            RestrictedOp::Eq => lhs.eq(rhs),
332            RestrictedOp::NotEq => lhs.not_eq(rhs),
333        };
334        Ok(expr)
335    }
336
337    /// Get the left-hand side operand
338    pub fn lhs(&self) -> &Operand {
339        &self.lhs
340    }
341
342    /// Get the right-hand side operand
343    pub fn rhs(&self) -> &Operand {
344        &self.rhs
345    }
346
347    /// Get the operation
348    pub fn op(&self) -> &RestrictedOp {
349        &self.op
350    }
351
352    pub fn try_as_physical_expr(
353        &self,
354        schema: &arrow::datatypes::SchemaRef,
355    ) -> error::Result<Arc<dyn PhysicalExpr>> {
356        let df_schema = schema
357            .clone()
358            .to_dfschema_ref()
359            .context(error::ToDFSchemaSnafu)?;
360        let execution_props = &ExecutionProps::default();
361        let expr = self.try_as_logical_expr()?;
362        create_physical_expr(
363            &expr,
364            &df_schema,
365            execution_props,
366            &PhysicalPlanningContext::default(),
367        )
368        .context(error::CreatePhysicalExprSnafu)
369    }
370
371    pub fn and(self, rhs: PartitionExpr) -> PartitionExpr {
372        PartitionExpr::new(Operand::Expr(self), RestrictedOp::And, Operand::Expr(rhs))
373    }
374
375    /// Serializes `PartitionExpr` to json string.
376    ///
377    /// Wraps `PartitionBound::Expr` for compatibility.
378    pub fn as_json_str(&self) -> error::Result<String> {
379        serde_json::to_string(&PartitionBound::Expr(self.clone()))
380            .context(error::SerializeJsonSnafu)
381    }
382
383    /// Deserializes `PartitionExpr` from json string.
384    ///
385    /// Deserializes to `PartitionBound` for compatibility.
386    pub fn from_json_str(s: &str) -> error::Result<Option<Self>> {
387        if s.is_empty() {
388            return Ok(None);
389        }
390
391        let bound: PartitionBound = serde_json::from_str(s).context(error::DeserializeJsonSnafu)?;
392        match bound {
393            PartitionBound::Expr(expr) => Ok(Some(expr.canonicalize())),
394            _ => Ok(None),
395        }
396    }
397
398    /// Converts [Self] to [Partition].
399    pub fn as_pb_partition(&self) -> error::Result<Partition> {
400        Ok(Partition {
401            expression: self.as_json_str()?,
402            ..Default::default()
403        })
404    }
405
406    /// Collects all column names referenced by this expression.
407    pub fn collect_column_names(&self, columns: &mut HashSet<String>) {
408        Self::collect_operand_columns(&self.lhs, columns);
409        Self::collect_operand_columns(&self.rhs, columns);
410    }
411
412    fn collect_operand_columns(operand: &Operand, columns: &mut HashSet<String>) {
413        match operand {
414            Operand::Column(c) => {
415                columns.insert(c.clone());
416            }
417            Operand::Expr(e) => {
418                e.collect_column_names(columns);
419            }
420            Operand::Value(_) => {}
421        }
422    }
423}
424
425impl Display for PartitionExpr {
426    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
427        write!(f, "{} {} {}", self.lhs, self.op, self.rhs)
428    }
429}
430
431#[cfg(test)]
432mod tests {
433    use std::collections::HashMap;
434
435    use common_meta::ddl::alter_table::validate_region_partition_bounds;
436    use common_meta::error::Error;
437    use common_meta::rpc::router::{Region, RegionRoute};
438    use datatypes::data_type::ConcreteDataType;
439    use store_api::storage::RegionId;
440
441    use super::*;
442
443    #[test]
444    fn test_partition_expr() {
445        let cases = [
446            (
447                Operand::Column("a".to_string()),
448                RestrictedOp::Eq,
449                Operand::Value(Value::UInt32(10)),
450                "a = 10",
451            ),
452            (
453                Operand::Column("a".to_string()),
454                RestrictedOp::NotEq,
455                Operand::Value(Value::UInt32(10)),
456                "a <> 10",
457            ),
458            (
459                Operand::Column("a".to_string()),
460                RestrictedOp::Lt,
461                Operand::Value(Value::UInt32(10)),
462                "a < 10",
463            ),
464            (
465                Operand::Column("a".to_string()),
466                RestrictedOp::LtEq,
467                Operand::Value(Value::UInt32(10)),
468                "a <= 10",
469            ),
470            (
471                Operand::Column("a".to_string()),
472                RestrictedOp::Gt,
473                Operand::Value(Value::UInt32(10)),
474                "a > 10",
475            ),
476            (
477                Operand::Column("a".to_string()),
478                RestrictedOp::GtEq,
479                Operand::Value(Value::UInt32(10)),
480                "a >= 10",
481            ),
482            (
483                Operand::Column("a".to_string()),
484                RestrictedOp::And,
485                Operand::Column("b".to_string()),
486                "a AND b",
487            ),
488            (
489                Operand::Column("a".to_string()),
490                RestrictedOp::Or,
491                Operand::Column("b".to_string()),
492                "a OR b",
493            ),
494            (
495                Operand::Column("a".to_string()),
496                RestrictedOp::Or,
497                Operand::Expr(PartitionExpr::new(
498                    Operand::Column("c".to_string()),
499                    RestrictedOp::And,
500                    Operand::Column("d".to_string()),
501                )),
502                "a OR c AND d",
503            ),
504        ];
505
506        for case in cases {
507            let expr = PartitionExpr::new(case.0, case.1.clone(), case.2);
508            assert_eq!(case.3, expr.to_string());
509        }
510    }
511
512    #[test]
513    fn test_try_as_logical_expr_null_equality() {
514        let eq_expr = PartitionExpr::new(
515            Operand::Column("a".to_string()),
516            RestrictedOp::Eq,
517            Operand::Value(Value::Null),
518        );
519        assert_eq!(
520            eq_expr.try_as_logical_expr().unwrap().to_string(),
521            "a IS NULL"
522        );
523
524        let neq_expr = PartitionExpr::new(
525            Operand::Column("a".to_string()),
526            RestrictedOp::NotEq,
527            Operand::Value(Value::Null),
528        );
529        assert_eq!(
530            neq_expr.try_as_logical_expr().unwrap().to_string(),
531            "a IS NOT NULL"
532        );
533    }
534
535    #[test]
536    fn test_try_as_logical_expr_null_range_comparison() {
537        // Test Lt with column on LHS
538        let lt_expr = PartitionExpr::new(
539            Operand::Column("a".to_string()),
540            RestrictedOp::Lt,
541            Operand::Value(Value::Int64(10)),
542        );
543        assert_eq!(
544            lt_expr.try_as_logical_expr().unwrap().to_string(),
545            "a < Int64(10) OR a IS NULL"
546        );
547
548        // Test Lt with column on RHS
549        let lt_expr_rhs_column = PartitionExpr::new(
550            Operand::Value(Value::Int64(10)),
551            RestrictedOp::Lt,
552            Operand::Column("a".to_string()),
553        );
554        assert_eq!(
555            lt_expr_rhs_column
556                .try_as_logical_expr()
557                .unwrap()
558                .to_string(),
559            "a > Int64(10)"
560        );
561
562        // Test Gt with column on LHS
563        let gt_expr = PartitionExpr::new(
564            Operand::Column("a".to_string()),
565            RestrictedOp::Gt,
566            Operand::Value(Value::Int64(10)),
567        );
568        assert_eq!(
569            gt_expr.try_as_logical_expr().unwrap().to_string(),
570            "a > Int64(10)"
571        );
572
573        // Test Gt with column on RHS
574        let gt_expr_rhs_column = PartitionExpr::new(
575            Operand::Value(Value::Int64(10)),
576            RestrictedOp::Gt,
577            Operand::Column("a".to_string()),
578        );
579        assert_eq!(
580            gt_expr_rhs_column
581                .try_as_logical_expr()
582                .unwrap()
583                .to_string(),
584            "a < Int64(10) OR a IS NULL"
585        );
586
587        // Test GtEq with column on LHS
588        let gteq_expr = PartitionExpr::new(
589            Operand::Column("a".to_string()),
590            RestrictedOp::GtEq,
591            Operand::Value(Value::Int64(10)),
592        );
593        assert_eq!(
594            gteq_expr.try_as_logical_expr().unwrap().to_string(),
595            "a >= Int64(10)"
596        );
597
598        // Test LtEq with column on LHS
599        let lteq_expr = PartitionExpr::new(
600            Operand::Column("a".to_string()),
601            RestrictedOp::LtEq,
602            Operand::Value(Value::Int64(10)),
603        );
604        assert_eq!(
605            lteq_expr.try_as_logical_expr().unwrap().to_string(),
606            "a <= Int64(10) OR a IS NULL"
607        );
608
609        let gteq_expr_rhs_column = PartitionExpr::new(
610            Operand::Value(Value::Int64(10)),
611            RestrictedOp::GtEq,
612            Operand::Column("a".to_string()),
613        );
614        assert_eq!(
615            gteq_expr_rhs_column
616                .try_as_logical_expr()
617                .unwrap()
618                .to_string(),
619            "a <= Int64(10) OR a IS NULL"
620        );
621
622        let lteq_expr_rhs_column = PartitionExpr::new(
623            Operand::Value(Value::Int64(10)),
624            RestrictedOp::LtEq,
625            Operand::Column("a".to_string()),
626        );
627        assert_eq!(
628            lteq_expr_rhs_column
629                .try_as_logical_expr()
630                .unwrap()
631                .to_string(),
632            "a >= Int64(10)"
633        );
634
635        let and_expr = PartitionExpr::new(
636            Operand::Expr(PartitionExpr::new(
637                Operand::Column("a".to_string()),
638                RestrictedOp::LtEq,
639                Operand::Value(Value::Int64(10)),
640            )),
641            RestrictedOp::And,
642            Operand::Expr(PartitionExpr::new(
643                Operand::Column("b".to_string()),
644                RestrictedOp::Gt,
645                Operand::Value(Value::Int64(5)),
646            )),
647        );
648        assert_eq!(
649            and_expr.try_as_logical_expr().unwrap().to_string(),
650            "(a <= Int64(10) OR a IS NULL) AND b > Int64(5)"
651        );
652
653        let and_expr = PartitionExpr::new(
654            Operand::Expr(PartitionExpr::new(
655                Operand::Column("a".to_string()),
656                RestrictedOp::LtEq,
657                Operand::Value(Value::Int64(10)),
658            )),
659            RestrictedOp::And,
660            Operand::Expr(PartitionExpr::new(
661                Operand::Column("a".to_string()),
662                RestrictedOp::Gt,
663                Operand::Value(Value::Int64(5)),
664            )),
665        );
666        assert_eq!(
667            and_expr.try_as_logical_expr().unwrap().to_string(),
668            "(a <= Int64(10) OR a IS NULL) AND a > Int64(5)"
669        );
670
671        let and_expr_strict_lower = PartitionExpr::new(
672            Operand::Expr(PartitionExpr::new(
673                Operand::Column("a".to_string()),
674                RestrictedOp::Lt,
675                Operand::Value(Value::Int64(10)),
676            )),
677            RestrictedOp::And,
678            Operand::Expr(PartitionExpr::new(
679                Operand::Column("a".to_string()),
680                RestrictedOp::GtEq,
681                Operand::Value(Value::Int64(5)),
682            )),
683        );
684        assert_eq!(
685            and_expr_strict_lower
686                .try_as_logical_expr()
687                .unwrap()
688                .to_string(),
689            "(a < Int64(10) OR a IS NULL) AND a >= Int64(5)"
690        );
691
692        let and_expr_rhs_column = PartitionExpr::new(
693            Operand::Expr(PartitionExpr::new(
694                Operand::Value(Value::Int64(10)),
695                RestrictedOp::GtEq,
696                Operand::Column("a".to_string()),
697            )),
698            RestrictedOp::And,
699            Operand::Expr(PartitionExpr::new(
700                Operand::Value(Value::Int64(5)),
701                RestrictedOp::Lt,
702                Operand::Column("a".to_string()),
703            )),
704        );
705        assert_eq!(
706            and_expr_rhs_column
707                .try_as_logical_expr()
708                .unwrap()
709                .to_string(),
710            "(a <= Int64(10) OR a IS NULL) AND a > Int64(5)"
711        );
712
713        let or_expr_same_column = PartitionExpr::new(
714            Operand::Expr(PartitionExpr::new(
715                Operand::Column("a".to_string()),
716                RestrictedOp::LtEq,
717                Operand::Value(Value::Int64(10)),
718            )),
719            RestrictedOp::Or,
720            Operand::Expr(PartitionExpr::new(
721                Operand::Column("a".to_string()),
722                RestrictedOp::Gt,
723                Operand::Value(Value::Int64(5)),
724            )),
725        );
726        assert_eq!(
727            or_expr_same_column
728                .try_as_logical_expr()
729                .unwrap()
730                .to_string(),
731            "a <= Int64(10) OR a IS NULL OR a > Int64(5)"
732        );
733    }
734
735    #[test]
736    fn test_try_as_logical_expr_rhs_column_without_canonicalize() {
737        let gt_expr_rhs_column = PartitionExpr {
738            lhs: Box::new(Operand::Value(Value::Int64(10))),
739            op: RestrictedOp::Gt,
740            rhs: Box::new(Operand::Column("a".to_string())),
741        };
742        assert_eq!(
743            gt_expr_rhs_column
744                .try_as_logical_expr()
745                .unwrap()
746                .to_string(),
747            "a < Int64(10) OR a IS NULL"
748        );
749
750        let gteq_expr_rhs_column = PartitionExpr {
751            lhs: Box::new(Operand::Value(Value::Int64(10))),
752            op: RestrictedOp::GtEq,
753            rhs: Box::new(Operand::Column("a".to_string())),
754        };
755        assert_eq!(
756            gteq_expr_rhs_column
757                .try_as_logical_expr()
758                .unwrap()
759                .to_string(),
760            "a <= Int64(10) OR a IS NULL"
761        );
762    }
763
764    #[test]
765    fn test_serde_partition_expr() {
766        let expr = PartitionExpr::new(
767            Operand::Column("a".to_string()),
768            RestrictedOp::Eq,
769            Operand::Value(Value::UInt32(10)),
770        );
771        let json = expr.as_json_str().unwrap();
772        assert_eq!(
773            json,
774            "{\"Expr\":{\"lhs\":{\"Column\":\"a\"},\"op\":\"Eq\",\"rhs\":{\"Value\":{\"UInt32\":10}}}}"
775        );
776
777        let json = r#"{"Expr":{"lhs":{"Column":"a"},"op":"GtEq","rhs":{"Value":{"UInt32":10}}}}"#;
778        let expr2 = PartitionExpr::from_json_str(json).unwrap().unwrap();
779        let expected = PartitionExpr::new(
780            Operand::Column("a".to_string()),
781            RestrictedOp::GtEq,
782            Operand::Value(Value::UInt32(10)),
783        );
784        assert_eq!(expr2, expected);
785
786        // empty string
787        let json = "";
788        let expr3 = PartitionExpr::from_json_str(json).unwrap();
789        assert!(expr3.is_none());
790
791        // variants other than Expr
792        let json = r#""MaxValue""#;
793        let expr4 = PartitionExpr::from_json_str(json).unwrap();
794        assert!(expr4.is_none());
795
796        let json = r#"{"Value":{"UInt32":10}}"#;
797        let expr5 = PartitionExpr::from_json_str(json).unwrap();
798        assert!(expr5.is_none());
799    }
800
801    #[test]
802    fn test_collect_column_names() {
803        // Simple expression: col_a = 1 should give {col_a}
804        let expr = col("a").eq(Value::Int64(1));
805        let mut columns = HashSet::new();
806        expr.collect_column_names(&mut columns);
807        assert_eq!(columns.len(), 1);
808        assert!(columns.contains("a"));
809
810        // Compound AND with same column: col_a >= 0 AND col_a < 10 should give {col_a}
811        let expr = col("a")
812            .gt_eq(Value::Int64(0))
813            .and(col("a").lt(Value::Int64(10)));
814        let mut columns = HashSet::new();
815        expr.collect_column_names(&mut columns);
816        assert_eq!(columns.len(), 1);
817        assert!(columns.contains("a"));
818
819        // Multiple columns: col_a >= 0 AND col_b < 10 should give {col_a, col_b}
820        let expr = col("a")
821            .gt_eq(Value::Int64(0))
822            .and(col("b").lt(Value::Int64(10)));
823        let mut columns = HashSet::new();
824        expr.collect_column_names(&mut columns);
825        assert_eq!(columns.len(), 2);
826        assert!(columns.contains("a"));
827        assert!(columns.contains("b"));
828
829        // Nested expression: (col_a >= 0 AND col_b < 10) AND col_c = 5
830        let expr = col("a")
831            .gt_eq(Value::Int64(0))
832            .and(col("b").lt(Value::Int64(10)))
833            .and(col("c").eq(Value::Int64(5)));
834        let mut columns = HashSet::new();
835        expr.collect_column_names(&mut columns);
836        assert_eq!(columns.len(), 3);
837        assert!(columns.contains("a"));
838        assert!(columns.contains("b"));
839        assert!(columns.contains("c"));
840    }
841
842    /// Real serialized expressions must survive the round trip through
843    /// `common-meta`'s persisted-expression mirror: whenever that mirror drifts
844    /// from these types, deserialization fails, validation is skipped and the
845    /// overflowing cases below stop producing `PartitionExprIncompatible`.
846    #[test]
847    fn test_serialized_expressions_validate_against_type_changes() {
848        let overflow_ms = serde_json::from_str::<Value>(
849            r#"{"Timestamp":{"value":32503680000000,"unit":"Millisecond"}}"#,
850        )
851        .unwrap();
852        let start_ms =
853            serde_json::from_str::<Value>(r#"{"Timestamp":{"value":0,"unit":"Millisecond"}}"#)
854                .unwrap();
855        let ts_ms = ConcreteDataType::timestamp_millisecond_datatype();
856        let ts_ns = ConcreteDataType::timestamp_nanosecond_datatype();
857
858        let routes = |expr: &PartitionExpr| {
859            vec![RegionRoute {
860                region: Region {
861                    id: RegionId::new(1024, 1),
862                    partition_expr: expr.as_json_str().unwrap(),
863                    ..Default::default()
864                },
865                ..Default::default()
866            }]
867        };
868        let changed = |column: &str, old_type: ConcreteDataType, new_type: ConcreteDataType| {
869            HashMap::from([(column.to_string(), (old_type, new_type))])
870        };
871
872        // Compile-time tripwire: exhaustively matching the operand and operator
873        // variants makes a new variant break this test at compile time, so the
874        // `partition_compat::Persisted*` mirror and these cases must be extended
875        // together instead of silently skipping validation.
876        fn walk(expr: &PartitionExpr) {
877            walk_operand(&expr.lhs);
878            walk_operand(&expr.rhs);
879            match expr.op {
880                RestrictedOp::Eq
881                | RestrictedOp::NotEq
882                | RestrictedOp::Lt
883                | RestrictedOp::LtEq
884                | RestrictedOp::Gt
885                | RestrictedOp::GtEq
886                | RestrictedOp::And
887                | RestrictedOp::Or => {}
888            }
889        }
890        fn walk_operand(operand: &Operand) {
891            match operand {
892                Operand::Column(_) | Operand::Value(_) => {}
893                Operand::Expr(expr) => walk(expr),
894            }
895        }
896
897        // Every expression shape must be rejected while the bound overflows
898        // nanoseconds, with the error naming the offending column.
899        let simple = col("ts").lt(overflow_ms.clone());
900        let nested_and = col("ts")
901            .gt_eq(start_ms.clone())
902            .and(col("ts").lt(overflow_ms.clone()));
903        let nested_or = PartitionExpr::new(
904            Operand::Expr(col("host").eq(Value::String("a".into()))),
905            RestrictedOp::Or,
906            Operand::Expr(col("ts").lt(overflow_ms.clone())),
907        );
908        for expr in [&simple, &nested_and, &nested_or] {
909            walk(expr);
910            let err = validate_region_partition_bounds(
911                &changed("ts", ts_ms.clone(), ts_ns.clone()),
912                &routes(expr),
913            )
914            .unwrap_err();
915            assert!(
916                matches!(
917                    err,
918                    Error::PartitionExprIncompatible { ref column, .. } if column == "ts"
919                ),
920                "unexpected error: {err}"
921            );
922        }
923
924        // A representable bound passes.
925        let compatible = col("ts").lt(start_ms);
926        validate_region_partition_bounds(&changed("ts", ts_ms, ts_ns), &routes(&compatible))
927            .unwrap();
928
929        // Narrowing a field column must be rejected as well: the region's write
930        // path compares the persisted Int64 bound against Int32 rows directly
931        // and panics, even though the scan would coerce both sides.
932        let narrowing = col("v").lt(Value::Int64(3_000_000_000));
933        let err = validate_region_partition_bounds(
934            &changed(
935                "v",
936                ConcreteDataType::int64_datatype(),
937                ConcreteDataType::int32_datatype(),
938            ),
939            &routes(&narrowing),
940        )
941        .unwrap_err();
942        assert!(
943            matches!(
944                err,
945                Error::PartitionExprIncompatible { ref column, .. } if column == "v"
946            ),
947            "unexpected error: {err}"
948        );
949    }
950}
partition/expr.rs

partition/
expr.rs
