Skip to main content

risingwave_frontend/utils/
iceberg_predicate.rs

1// Copyright 2026 RisingWave Labs
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 chrono::Datelike;
16use iceberg::expr::{Predicate as IcebergPredicate, Reference};
17use iceberg::spec::Datum as IcebergDatum;
18use risingwave_common::catalog::Field;
19use risingwave_common::types::{DataType, ScalarImpl};
20
21use crate::expr::{Expr, ExprImpl, ExprType, InputRef, Literal};
22use crate::utils::Condition;
23
24pub struct ExtractIcebergPredicateResult {
25    pub iceberg_predicate: IcebergPredicate,
26    pub extracted_condition: Condition,
27    pub remaining_condition: Condition,
28}
29
30/// NOTE(kwannoel): We do predicate pushdown to the iceberg-sdk here.
31/// zone-map is used to evaluate predicates on iceberg tables.
32/// Without zone-map, iceberg-sdk will still apply the predicate on its own.
33/// See: <https://github.com/apache/iceberg-rust/blob/5c1a9e68da346819072a15327080a498ad91c488/crates/iceberg/src/arrow/reader.rs#L229-L235>.
34///
35/// `fields` must carry the iceberg-side column types: for engine tables, pass the original
36/// source types rather than the remapped Hummock output types.
37pub fn extract_iceberg_predicate(
38    predicate: Condition,
39    fields: &[Field],
40) -> ExtractIcebergPredicateResult {
41    if predicate.always_true() {
42        return ExtractIcebergPredicateResult {
43            iceberg_predicate: IcebergPredicate::AlwaysTrue,
44            extracted_condition: Condition {
45                conjunctions: vec![],
46            },
47            remaining_condition: Condition {
48                conjunctions: vec![],
49            },
50        };
51    }
52
53    let mut iceberg_predicates = Vec::new();
54    let mut extracted_conjunctions = Vec::new();
55    let mut remaining_conjunctions = Vec::new();
56
57    for conjunction in predicate.conjunctions {
58        match rw_expr_to_iceberg_predicate(&conjunction, fields) {
59            Some(iceberg_predicate) => {
60                iceberg_predicates.push(iceberg_predicate);
61                extracted_conjunctions.push(conjunction);
62            }
63            None => remaining_conjunctions.push(conjunction),
64        }
65    }
66
67    let iceberg_predicate = iceberg_predicates
68        .into_iter()
69        .reduce(IcebergPredicate::and)
70        .unwrap_or(IcebergPredicate::AlwaysTrue);
71
72    ExtractIcebergPredicateResult {
73        iceberg_predicate,
74        extracted_condition: Condition {
75            conjunctions: extracted_conjunctions,
76        },
77        remaining_condition: Condition {
78            conjunctions: remaining_conjunctions,
79        },
80    }
81}
82
83fn rw_literal_to_iceberg_datum(literal: &Literal) -> Option<IcebergDatum> {
84    let Some(scalar) = literal.get_data() else {
85        return None;
86    };
87    match scalar {
88        ScalarImpl::Bool(b) => Some(IcebergDatum::bool(*b)),
89        ScalarImpl::Int32(i) => Some(IcebergDatum::int(*i)),
90        ScalarImpl::Int64(i) => Some(IcebergDatum::long(*i)),
91        ScalarImpl::Float32(f) => Some(IcebergDatum::float(*f)),
92        ScalarImpl::Float64(f) => Some(IcebergDatum::double(*f)),
93        ScalarImpl::Decimal(_) => {
94            // TODO(iceberg): iceberg-rust doesn't support decimal predicate pushdown yet.
95            None
96        }
97        ScalarImpl::Date(d) => {
98            let Ok(datum) = IcebergDatum::date_from_ymd(d.0.year(), d.0.month(), d.0.day()) else {
99                return None;
100            };
101            Some(datum)
102        }
103        ScalarImpl::Timestamp(t) => Some(IcebergDatum::timestamp_micros(
104            t.0.and_utc().timestamp_micros(),
105        )),
106        ScalarImpl::Timestamptz(t) => Some(IcebergDatum::timestamptz_micros(t.timestamp_micros())),
107        ScalarImpl::Utf8(s) => Some(IcebergDatum::string(s)),
108        ScalarImpl::Bytea(b) => Some(IcebergDatum::binary(b.clone())),
109        _ => None,
110    }
111}
112
113/// Whether a column of this iceberg-side type can appear in a pushed-down predicate:
114/// iceberg-rust only builds field accessors for primitive columns, so a pushed-down
115/// predicate referencing others fails at scan time with "Accessor for Field ... not found".
116fn is_iceberg_predicate_pushable_column_type(ty: &DataType) -> bool {
117    matches!(
118        ty,
119        DataType::Boolean
120            | DataType::Int16
121            | DataType::Int32
122            | DataType::Int64
123            | DataType::Float32
124            | DataType::Float64
125            | DataType::Decimal
126            | DataType::Date
127            | DataType::Time
128            | DataType::Timestamp
129            | DataType::Timestamptz
130            | DataType::Varchar
131            | DataType::Bytea
132    )
133}
134
135/// Build the `Reference` for a pushed-down predicate, or `None` for an unpushable column.
136/// Every match arm must obtain its `Reference` here so the column-type check cannot be
137/// bypassed.
138fn input_ref_to_reference(input_ref: &InputRef, fields: &[Field]) -> Option<Reference> {
139    let field = &fields[input_ref.index];
140    if !is_iceberg_predicate_pushable_column_type(&field.data_type) {
141        return None;
142    }
143    Some(Reference::new(&field.name))
144}
145
146/// Convert `<col> <op> <literal>` (either operand order) into an iceberg predicate.
147fn comparison_to_iceberg_predicate(
148    func_type: ExprType,
149    args: &[ExprImpl],
150    fields: &[Field],
151) -> Option<IcebergPredicate> {
152    // `flipped` means the literal is on the left: `lit <op> col` ⇔ `col <mirrored op> lit`.
153    let (input_ref, literal, flipped) = match [&args[0], &args[1]] {
154        [ExprImpl::InputRef(input_ref), ExprImpl::Literal(literal)] => (input_ref, literal, false),
155        [ExprImpl::Literal(literal), ExprImpl::InputRef(input_ref)] => (input_ref, literal, true),
156        _ => return None,
157    };
158    let reference = input_ref_to_reference(input_ref, fields)?;
159    let datum = rw_literal_to_iceberg_datum(literal)?;
160    let predicate = match (func_type, flipped) {
161        (ExprType::Equal, _) => reference.equal_to(datum),
162        (ExprType::NotEqual, _) => reference.not_equal_to(datum),
163        (ExprType::GreaterThan, false) | (ExprType::LessThan, true) => {
164            reference.greater_than(datum)
165        }
166        (ExprType::GreaterThanOrEqual, false) | (ExprType::LessThanOrEqual, true) => {
167            reference.greater_than_or_equal_to(datum)
168        }
169        (ExprType::LessThan, false) | (ExprType::GreaterThan, true) => reference.less_than(datum),
170        (ExprType::LessThanOrEqual, false) | (ExprType::GreaterThanOrEqual, true) => {
171            reference.less_than_or_equal_to(datum)
172        }
173        _ => return None,
174    };
175    Some(predicate)
176}
177
178fn rw_expr_to_iceberg_predicate(expr: &ExprImpl, fields: &[Field]) -> Option<IcebergPredicate> {
179    match expr {
180        ExprImpl::Literal(l) => match l.get_data() {
181            Some(ScalarImpl::Bool(b)) => {
182                if *b {
183                    Some(IcebergPredicate::AlwaysTrue)
184                } else {
185                    Some(IcebergPredicate::AlwaysFalse)
186                }
187            }
188            _ => None,
189        },
190        ExprImpl::FunctionCall(f) => {
191            let args = f.inputs();
192            match f.func_type() {
193                ExprType::Not => {
194                    let arg = rw_expr_to_iceberg_predicate(&args[0], fields)?;
195                    Some(IcebergPredicate::negate(arg))
196                }
197                ExprType::And => {
198                    let arg0 = rw_expr_to_iceberg_predicate(&args[0], fields)?;
199                    let arg1 = rw_expr_to_iceberg_predicate(&args[1], fields)?;
200                    Some(IcebergPredicate::and(arg0, arg1))
201                }
202                ExprType::Or => {
203                    let arg0 = rw_expr_to_iceberg_predicate(&args[0], fields)?;
204                    let arg1 = rw_expr_to_iceberg_predicate(&args[1], fields)?;
205                    Some(IcebergPredicate::or(arg0, arg1))
206                }
207                ExprType::Equal
208                | ExprType::NotEqual
209                | ExprType::GreaterThan
210                | ExprType::GreaterThanOrEqual
211                | ExprType::LessThan
212                | ExprType::LessThanOrEqual
213                    if args[0].return_type() == args[1].return_type() =>
214                {
215                    comparison_to_iceberg_predicate(f.func_type(), args, fields)
216                }
217                ExprType::IsNull => match &args[0] {
218                    ExprImpl::InputRef(lhs) => {
219                        let reference = input_ref_to_reference(lhs, fields)?;
220                        Some(reference.is_null())
221                    }
222                    _ => None,
223                },
224                ExprType::IsNotNull => match &args[0] {
225                    ExprImpl::InputRef(lhs) => {
226                        let reference = input_ref_to_reference(lhs, fields)?;
227                        Some(reference.is_not_null())
228                    }
229                    _ => None,
230                },
231                ExprType::In => match &args[0] {
232                    ExprImpl::InputRef(lhs) => {
233                        let reference = input_ref_to_reference(lhs, fields)?;
234                        let mut datums = Vec::with_capacity(args.len() - 1);
235                        for arg in &args[1..] {
236                            if args[0].return_type() != arg.return_type() {
237                                return None;
238                            }
239                            if let ExprImpl::Literal(l) = arg {
240                                let datum = rw_literal_to_iceberg_datum(l)?;
241                                datums.push(datum);
242                            } else {
243                                return None;
244                            }
245                        }
246                        Some(reference.is_in(datums))
247                    }
248                    _ => None,
249                },
250                _ => None,
251            }
252        }
253        _ => None,
254    }
255}
256
257#[cfg(test)]
258mod tests {
259    use risingwave_common::types::{MapType, ScalarImpl, StructType};
260
261    use super::*;
262    use crate::expr::{FunctionCall, InputRef};
263
264    #[test]
265    fn comparison_pushdown_mirrors_flipped_operators() {
266        let fields = vec![Field::new("x", DataType::Int64)];
267        let col = || InputRef::new(0, DataType::Int64).into();
268        let lit = || Literal::new(Some(ScalarImpl::Int64(5)), DataType::Int64).into();
269
270        let cases: [(ExprType, ExprImpl, ExprImpl, &str); 12] = [
271            (ExprType::Equal, col(), lit(), "x = 5"),
272            (ExprType::Equal, lit(), col(), "x = 5"),
273            (ExprType::NotEqual, col(), lit(), "x != 5"),
274            (ExprType::NotEqual, lit(), col(), "x != 5"),
275            (ExprType::GreaterThan, col(), lit(), "x > 5"),
276            (ExprType::GreaterThan, lit(), col(), "x < 5"),
277            (ExprType::GreaterThanOrEqual, col(), lit(), "x >= 5"),
278            (ExprType::GreaterThanOrEqual, lit(), col(), "x <= 5"),
279            (ExprType::LessThan, col(), lit(), "x < 5"),
280            (ExprType::LessThan, lit(), col(), "x > 5"),
281            (ExprType::LessThanOrEqual, col(), lit(), "x <= 5"),
282            (ExprType::LessThanOrEqual, lit(), col(), "x >= 5"),
283        ];
284        for (op, arg0, arg1, expected) in cases {
285            let expr = FunctionCall::new(op, vec![arg0, arg1]).unwrap().into();
286            let predicate = rw_expr_to_iceberg_predicate(&expr, &fields)
287                .unwrap_or_else(|| panic!("{op:?} should be pushable"));
288            assert_eq!(predicate.to_string(), expected);
289        }
290    }
291
292    fn null_check(op: ExprType, input_type: DataType) -> ExprImpl {
293        let col: ExprImpl = InputRef::new(0, input_type).into();
294        FunctionCall::new_unchecked(op, vec![col], DataType::Boolean).into()
295    }
296
297    #[test]
298    fn non_primitive_columns_are_not_pushed_down() {
299        // iceberg-rust only builds field accessors for primitive columns, so a pushed-down
300        // predicate on any other type fails the scan with "Accessor for Field ... not found".
301        let unpushable = [
302            DataType::Struct(StructType::new(vec![("a", DataType::Int32)])),
303            DataType::list(DataType::Int32),
304            DataType::Map(MapType::from_kv(DataType::Varchar, DataType::Int32)),
305            DataType::Variant,
306            DataType::Jsonb,
307        ];
308        for data_type in unpushable {
309            let fields = vec![Field::new("c", data_type.clone())];
310            for op in [ExprType::IsNull, ExprType::IsNotNull] {
311                let expr = null_check(op, data_type.clone());
312                assert!(
313                    rw_expr_to_iceberg_predicate(&expr, &fields).is_none(),
314                    "{op:?} on {data_type} must not be pushed down"
315                );
316            }
317        }
318
319        let fields = vec![Field::new("c", DataType::Int64)];
320        for op in [ExprType::IsNull, ExprType::IsNotNull] {
321            let expr = null_check(op, DataType::Int64);
322            assert!(
323                rw_expr_to_iceberg_predicate(&expr, &fields).is_some(),
324                "{op:?} on a primitive column should be pushable"
325            );
326        }
327    }
328
329    #[test]
330    fn pushability_follows_the_iceberg_side_column_type() {
331        // Engine tables remap the scan schema to Hummock types (here jsonb backed by an iceberg
332        // string), so the `InputRef` carries jsonb while the iceberg-side field carries varchar.
333        // Judging by the remapped type would silently drop the pushdown.
334        let expr = null_check(ExprType::IsNotNull, DataType::Jsonb);
335
336        let iceberg_side = vec![Field::new("c", DataType::Varchar)];
337        assert!(rw_expr_to_iceberg_predicate(&expr, &iceberg_side).is_some());
338
339        let remapped = vec![Field::new("c", DataType::Jsonb)];
340        assert!(rw_expr_to_iceberg_predicate(&expr, &remapped).is_none());
341    }
342}