1use 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
30pub 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 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
113fn 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
135fn 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
146fn comparison_to_iceberg_predicate(
148 func_type: ExprType,
149 args: &[ExprImpl],
150 fields: &[Field],
151) -> Option<IcebergPredicate> {
152 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 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 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}