Skip to main content

liquid_cache/liquid_array/byte_view_array/
operator.rs

1use std::sync::Arc;
2
3use datafusion_common::ScalarValue;
4use datafusion_expr_common::operator::Operator;
5use datafusion_physical_expr::PhysicalExpr;
6use datafusion_physical_expr::expressions::{
7    BinaryExpr, DynamicFilterPhysicalExpr, LikeExpr, Literal,
8};
9
10use crate::utils::get_bytes_needle;
11
12/// Supported ordering comparisons for byte views.
13#[derive(Debug)]
14pub enum Comparison {
15    /// Less-than.
16    Lt,
17    /// Greater-than.
18    Gt,
19    /// Less-than or equal.
20    LtEq,
21    /// Greater-than or equal.
22    GtEq,
23}
24
25/// Supported equality comparisons for byte views.
26#[derive(Debug)]
27pub enum Equality {
28    /// Equal.
29    Eq,
30    /// Not equal.
31    NotEq,
32}
33
34#[derive(Debug, PartialEq, Eq, Copy, Clone)]
35/// Supported substring predicate kinds.
36pub enum SubString {
37    /// Contains a substring.
38    Contains,
39    /// Does not contain a substring.
40    NotContains,
41}
42
43/// Supported operators for byte view predicates.
44#[derive(Debug)]
45pub enum ByteViewOperator {
46    /// Ordering comparison.
47    Comparison(Comparison),
48    /// Equality comparison.
49    Equality(Equality),
50    /// Substring predicate.
51    SubString(SubString),
52}
53
54impl ByteViewOperator {
55    fn from_like_expr(like: &LikeExpr) -> Result<Self, UnsupportedOperator> {
56        match (like.negated(), like.case_insensitive()) {
57            (false, false) => Ok(ByteViewOperator::SubString(SubString::Contains)),
58            (true, false) => Ok(ByteViewOperator::SubString(SubString::NotContains)),
59            _ => Err(UnsupportedOperator),
60        }
61    }
62}
63
64pub struct UnsupportedOperator;
65
66impl TryFrom<&Operator> for ByteViewOperator {
67    type Error = UnsupportedOperator;
68
69    fn try_from(operator: &Operator) -> Result<Self, UnsupportedOperator> {
70        match operator {
71            Operator::Eq => Ok(ByteViewOperator::Equality(Equality::Eq)),
72            Operator::NotEq => Ok(ByteViewOperator::Equality(Equality::NotEq)),
73            Operator::Lt => Ok(ByteViewOperator::Comparison(Comparison::Lt)),
74            Operator::Gt => Ok(ByteViewOperator::Comparison(Comparison::Gt)),
75            Operator::LtEq => Ok(ByteViewOperator::Comparison(Comparison::LtEq)),
76            Operator::GtEq => Ok(ByteViewOperator::Comparison(Comparison::GtEq)),
77            Operator::LikeMatch => Ok(ByteViewOperator::SubString(SubString::Contains)),
78            Operator::NotLikeMatch => Ok(ByteViewOperator::SubString(SubString::NotContains)),
79            _ => Err(UnsupportedOperator),
80        }
81    }
82}
83
84impl From<&ByteViewOperator> for Operator {
85    fn from(byte_view_operator: &ByteViewOperator) -> Self {
86        match byte_view_operator {
87            ByteViewOperator::Comparison(comparison) => match comparison {
88                Comparison::Lt => Operator::Lt,
89                Comparison::Gt => Operator::Gt,
90                Comparison::LtEq => Operator::LtEq,
91                Comparison::GtEq => Operator::GtEq,
92            },
93            ByteViewOperator::Equality(equality) => match equality {
94                Equality::Eq => Operator::Eq,
95                Equality::NotEq => Operator::NotEq,
96            },
97            ByteViewOperator::SubString(substring) => match substring {
98                SubString::Contains => Operator::LikeMatch,
99                SubString::NotContains => Operator::NotLikeMatch,
100            },
101        }
102    }
103}
104
105#[derive(Debug)]
106pub(super) struct ByteViewExpression {
107    op: ByteViewOperator,
108    literal: Vec<u8>,
109}
110
111pub(super) enum UnsupportedExpression {
112    Op,
113    Expr,
114    // This is frequently the case with dynamic filters
115    Constant(bool),
116}
117
118impl From<UnsupportedOperator> for UnsupportedExpression {
119    fn from(_op: UnsupportedOperator) -> Self {
120        UnsupportedExpression::Op
121    }
122}
123
124impl ByteViewExpression {
125    pub(super) fn op(&self) -> &ByteViewOperator {
126        &self.op
127    }
128
129    pub(super) fn literal(&self) -> &[u8] {
130        &self.literal
131    }
132}
133
134impl TryFrom<&Arc<dyn PhysicalExpr>> for ByteViewExpression {
135    type Error = UnsupportedExpression;
136    fn try_from(expr: &Arc<dyn PhysicalExpr>) -> Result<Self, UnsupportedExpression> {
137        let expr = if let Some(dynamic_filter) =
138            expr.as_any().downcast_ref::<DynamicFilterPhysicalExpr>()
139        {
140            dynamic_filter.current().unwrap()
141        } else {
142            expr.clone()
143        };
144
145        if let Some(literal) = expr.as_any().downcast_ref::<Literal>()
146            && let ScalarValue::Boolean(Some(v)) = literal.value()
147        {
148            return Err(UnsupportedExpression::Constant(*v));
149        }
150
151        if let Some(binary_expr) = expr.as_any().downcast_ref::<BinaryExpr>() {
152            if let Some(literal) = binary_expr.right().as_any().downcast_ref::<Literal>() {
153                let op = binary_expr.op();
154                let byte_view_operator = ByteViewOperator::try_from(op)?;
155                let literal =
156                    get_bytes_needle(literal.value()).ok_or(UnsupportedExpression::Expr)?;
157                return Ok(ByteViewExpression {
158                    op: byte_view_operator,
159                    literal,
160                });
161            }
162        }
163        // Handle like expressions
164        else if let Some(like_expr) = expr.as_any().downcast_ref::<LikeExpr>()
165            && let Some(literal) = like_expr.pattern().as_any().downcast_ref::<Literal>()
166        {
167            let byte_view_operator = ByteViewOperator::from_like_expr(like_expr)?;
168            let literal = get_bytes_needle(literal.value()).ok_or(UnsupportedExpression::Expr)?;
169            return Ok(ByteViewExpression {
170                op: byte_view_operator,
171                literal,
172            });
173        }
174        Err(UnsupportedExpression::Expr)
175    }
176}