Skip to main content

liquid_cache/liquid_array/byte_view_array/
operator.rs

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