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

partition/
expr.rs
