CoCalc -- mod.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-plan/src/plans/optimizer/simplify_expr/mod.rs
⁸⁴¹⁶ views
1
mod simplify_functions;
2

3
use num_traits::Zero;
4
use polars_utils::float16::pf16;
5
use polars_utils::floor_divmod::FloorDivMod;
6
use polars_utils::total_ord::ToTotalOrd;
7
use simplify_functions::optimize_functions;
8
mod arity;
9

10
use crate::plans::*;
11

12
fn new_null_count(input: &[ExprIR]) -> AExpr {
13
    let function = IRFunctionExpr::NullCount;
14
    let options = function.function_options();
15
    AExpr::Function {
16
        input: input.to_vec(),
17
        function,
18
        options,
19
    }
20
}
21

22
macro_rules! eval_binary_same_type {
23
    ($lhs:expr, $rhs:expr, |$l: ident, $r: ident| $ret: expr) => {{
24
        if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) = ($lhs, $rhs) {
25
            match (lit_left, lit_right) {
26
                (LiteralValue::Scalar(l), LiteralValue::Scalar(r)) => {
27
                    match (l.as_any_value(), r.as_any_value()) {
28
                        (AnyValue::Float16($l), AnyValue::Float16($r)) => {
29
                            Some(AExpr::Literal(Scalar::from($ret).into()))
30
                        },
31
                        (AnyValue::Float32($l), AnyValue::Float32($r)) => {
32
                            Some(AExpr::Literal(<Scalar as From<f32>>::from($ret).into()))
33
                        },
34
                        (AnyValue::Float64($l), AnyValue::Float64($r)) => {
35
                            Some(AExpr::Literal(<Scalar as From<f64>>::from($ret).into()))
36
                        },
37

38
                        (AnyValue::Int8($l), AnyValue::Int8($r)) => {
39
                            Some(AExpr::Literal(<Scalar as From<i8>>::from($ret).into()))
40
                        },
41
                        (AnyValue::Int16($l), AnyValue::Int16($r)) => {
42
                            Some(AExpr::Literal(<Scalar as From<i16>>::from($ret).into()))
43
                        },
44
                        (AnyValue::Int32($l), AnyValue::Int32($r)) => {
45
                            Some(AExpr::Literal(<Scalar as From<i32>>::from($ret).into()))
46
                        },
47
                        (AnyValue::Int64($l), AnyValue::Int64($r)) => {
48
                            Some(AExpr::Literal(<Scalar as From<i64>>::from($ret).into()))
49
                        },
50
                        (AnyValue::Int128($l), AnyValue::Int128($r)) => {
51
                            Some(AExpr::Literal(<Scalar as From<i128>>::from($ret).into()))
52
                        },
53

54
                        (AnyValue::UInt8($l), AnyValue::UInt8($r)) => {
55
                            Some(AExpr::Literal(<Scalar as From<u8>>::from($ret).into()))
56
                        },
57
                        (AnyValue::UInt16($l), AnyValue::UInt16($r)) => {
58
                            Some(AExpr::Literal(<Scalar as From<u16>>::from($ret).into()))
59
                        },
60
                        (AnyValue::UInt32($l), AnyValue::UInt32($r)) => {
61
                            Some(AExpr::Literal(<Scalar as From<u32>>::from($ret).into()))
62
                        },
63
                        (AnyValue::UInt64($l), AnyValue::UInt64($r)) => {
64
                            Some(AExpr::Literal(<Scalar as From<u64>>::from($ret).into()))
65
                        },
66
                        (AnyValue::UInt128($l), AnyValue::UInt128($r)) => {
67
                            Some(AExpr::Literal(<Scalar as From<u128>>::from($ret).into()))
68
                        },
69

70
                        _ => None,
71
                    }
72
                    .into()
73
                },
74
                (
75
                    LiteralValue::Dyn(DynLiteralValue::Float($l)),
76
                    LiteralValue::Dyn(DynLiteralValue::Float($r)),
77
                ) => {
78
                    let $l = *$l;
79
                    let $r = *$r;
80
                    Some(AExpr::Literal(LiteralValue::Dyn(DynLiteralValue::Float(
81
                        $ret,
82
                    ))))
83
                },
84
                (
85
                    LiteralValue::Dyn(DynLiteralValue::Int($l)),
86
                    LiteralValue::Dyn(DynLiteralValue::Int($r)),
87
                ) => {
88
                    let $l = *$l;
89
                    let $r = *$r;
90
                    Some(AExpr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(
91
                        $ret,
92
                    ))))
93
                },
94
                _ => None,
95
            }
96
        } else {
97
            None
98
        }
99
    }};
100
}
101

102
macro_rules! eval_binary_cmp_same_type {
103
    ($lhs:expr, $operand: tt, $rhs:expr) => {{
104
    if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) = ($lhs, $rhs) {
105
        match (lit_left, lit_right) {
106
            (LiteralValue::Scalar(l), LiteralValue::Scalar(r)) => match (l.as_any_value(), r.as_any_value()) {
107
                (AnyValue::Float16(l), AnyValue::Float16(r)) => Some(AExpr::Literal({ let x: bool = l.to_total_ord() $operand r.to_total_ord(); Scalar::from(x) }.into())),
108
                (AnyValue::Float32(l), AnyValue::Float32(r)) => Some(AExpr::Literal({ let x: bool = l.to_total_ord() $operand r.to_total_ord(); Scalar::from(x) }.into())),
109
                (AnyValue::Float64(l), AnyValue::Float64(r)) => Some(AExpr::Literal({ let x: bool = l.to_total_ord() $operand r.to_total_ord(); Scalar::from(x) }.into())),
110

111
                (AnyValue::Boolean(l), AnyValue::Boolean(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
112

113
                (AnyValue::Int8(l), AnyValue::Int8(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
114
                (AnyValue::Int16(l), AnyValue::Int16(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
115
                (AnyValue::Int32(l), AnyValue::Int32(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
116
                (AnyValue::Int64(l), AnyValue::Int64(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
117
                (AnyValue::Int128(l), AnyValue::Int128(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
118

119
                (AnyValue::UInt8(l), AnyValue::UInt8(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
120
                (AnyValue::UInt16(l), AnyValue::UInt16(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
121
                (AnyValue::UInt32(l), AnyValue::UInt32(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
122
                (AnyValue::UInt64(l), AnyValue::UInt64(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
123
                (AnyValue::UInt128(l), AnyValue::UInt128(r)) => Some(AExpr::Literal({ let x: bool = l $operand r; Scalar::from(x) }.into())),
124

125
                _ => None,
126
            }.into(),
127
            (LiteralValue::Dyn(DynLiteralValue::Float(l)), LiteralValue::Dyn(DynLiteralValue::Float(r))) => {
128
                let x: bool = l.to_total_ord() $operand r.to_total_ord();
129
                Some(AExpr::Literal(Scalar::from(x).into()))
130
            },
131
            (LiteralValue::Dyn(DynLiteralValue::Int(l)), LiteralValue::Dyn(DynLiteralValue::Int(r))) => {
132
                let x: bool = l $operand r;
133
                Some(AExpr::Literal(Scalar::from(x).into()))
134
            },
135
            _ => None,
136
        }
137
    } else {
138
        None
139
    }
140

141
    }}
142
}
143

144
pub struct SimplifyBooleanRule {}
145

146
impl OptimizationRule for SimplifyBooleanRule {
147
    fn optimize_expr(
148
        &mut self,
149
        expr_arena: &mut Arena<AExpr>,
150
        expr_node: Node,
151
        _schema: &Schema,
152
        ctx: OptimizeExprContext,
153
    ) -> PolarsResult<Option<AExpr>> {
154
        let expr = expr_arena.get(expr_node);
155

156
        let out = match expr {
157
            // true AND x => x
158
            AExpr::BinaryExpr { left, op, right } => {
159
                return Ok(arity::simplify_binary(*left, *op, *right, ctx, expr_arena));
160
            },
161
            AExpr::Ternary {
162
                predicate,
163
                truthy,
164
                falsy,
165
            } => {
166
                return Ok(arity::simplify_ternary(
167
                    *predicate, *truthy, *falsy, expr_arena,
168
                ));
169
            },
170
            AExpr::Function {
171
                input,
172
                function: IRFunctionExpr::Negate,
173
                ..
174
            } if input.len() == 1 => {
175
                let input = &input[0];
176
                let ae = expr_arena.get(input.node());
177
                eval_negate(ae)
178
            },
179
            _ => None,
180
        };
181
        Ok(out)
182
    }
183
}
184

185
fn eval_negate(ae: &AExpr) -> Option<AExpr> {
186
    use std::ops::Neg;
187
    let out = match ae {
188
        AExpr::Literal(lv) => match lv {
189
            LiteralValue::Scalar(sc) => match sc.as_any_value() {
190
                AnyValue::Int8(v) => Scalar::from(v.checked_neg()?),
191
                AnyValue::Int16(v) => Scalar::from(v.checked_neg()?),
192
                AnyValue::Int32(v) => Scalar::from(v.checked_neg()?),
193
                AnyValue::Int64(v) => Scalar::from(v.checked_neg()?),
194
                AnyValue::Int128(v) => Scalar::from(v.checked_neg()?),
195
                AnyValue::Float16(v) => Scalar::from(v.neg()),
196
                AnyValue::Float32(v) => Scalar::from(v.neg()),
197
                AnyValue::Float64(v) => Scalar::from(v.neg()),
198
                _ => return None,
199
            }
200
            .into(),
201
            LiteralValue::Dyn(d) => LiteralValue::Dyn(match d {
202
                DynLiteralValue::Int(v) => DynLiteralValue::Int(v.checked_neg()?),
203
                DynLiteralValue::Float(v) => DynLiteralValue::Float(v.neg()),
204
                _ => return None,
205
            }),
206
            _ => return None,
207
        },
208
        _ => return None,
209
    };
210
    Some(AExpr::Literal(out))
211
}
212

213
fn eval_bitwise<F>(left: &AExpr, right: &AExpr, operation: F) -> Option<AExpr>
214
where
215
    F: Fn(bool, bool) -> bool,
216
{
217
    if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) = (left, right) {
218
        return match (lit_left.bool(), lit_right.bool()) {
219
            (Some(x), Some(y)) => Some(AExpr::Literal(Scalar::from(operation(x, y)).into())),
220
            _ => None,
221
        };
222
    }
223
    None
224
}
225

226
#[cfg(all(feature = "strings", feature = "concat_str"))]
227
fn string_addition_to_linear_concat(
228
    expr_arena: &Arena<AExpr>,
229
    left_node: Node,
230
    right_node: Node,
231
    left_aexpr: &AExpr,
232
    right_aexpr: &AExpr,
233
    input_schema: &Schema,
234
) -> Option<AExpr> {
235
    {
236
        let left_e = ExprIR::from_node(left_node, expr_arena);
237
        let right_e = ExprIR::from_node(right_node, expr_arena);
238

239
        let get_type = |ae: &AExpr| {
240
            ae.to_dtype(&ToFieldContext::new(expr_arena, input_schema))
241
                .ok()
242
        };
243
        let type_a = get_type(left_aexpr).or_else(|| get_type(right_aexpr))?;
244
        let type_b = get_type(right_aexpr).or_else(|| get_type(right_aexpr))?;
245

246
        if type_a != type_b {
247
            return None;
248
        }
249

250
        if type_a.is_string() {
251
            match (left_aexpr, right_aexpr) {
252
                // concat + concat
253
                (
254
                    AExpr::Function {
255
                        input: input_left,
256
                        function:
257
                            fun_l @ IRFunctionExpr::StringExpr(IRStringFunction::ConcatHorizontal {
258
                                delimiter: sep_l,
259
                                ignore_nulls: ignore_nulls_l,
260
                            }),
261
                        options,
262
                    },
263
                    AExpr::Function {
264
                        input: input_right,
265
                        function:
266
                            IRFunctionExpr::StringExpr(IRStringFunction::ConcatHorizontal {
267
                                delimiter: sep_r,
268
                                ignore_nulls: ignore_nulls_r,
269
                            }),
270
                        ..
271
                    },
272
                ) => {
273
                    if sep_l.is_empty() && sep_r.is_empty() && ignore_nulls_l == ignore_nulls_r {
274
                        let mut input = Vec::with_capacity(input_left.len() + input_right.len());
275
                        input.extend_from_slice(input_left);
276
                        input.extend_from_slice(input_right);
277
                        Some(AExpr::Function {
278
                            input,
279
                            function: fun_l.clone(),
280
                            options: *options,
281
                        })
282
                    } else {
283
                        None
284
                    }
285
                },
286
                // concat + str
287
                (
288
                    AExpr::Function {
289
                        input,
290
                        function:
291
                            fun @ IRFunctionExpr::StringExpr(IRStringFunction::ConcatHorizontal {
292
                                delimiter: sep,
293
                                ignore_nulls,
294
                            }),
295
                        options,
296
                    },
297
                    _,
298
                ) => {
299
                    if sep.is_empty() && !ignore_nulls {
300
                        let mut input = input.clone();
301
                        input.push(right_e);
302
                        Some(AExpr::Function {
303
                            input,
304
                            function: fun.clone(),
305
                            options: *options,
306
                        })
307
                    } else {
308
                        None
309
                    }
310
                },
311
                // str + concat
312
                (
313
                    _,
314
                    AExpr::Function {
315
                        input: input_right,
316
                        function:
317
                            fun @ IRFunctionExpr::StringExpr(IRStringFunction::ConcatHorizontal {
318
                                delimiter: sep,
319
                                ignore_nulls,
320
                            }),
321
                        options,
322
                    },
323
                ) => {
324
                    if sep.is_empty() && !ignore_nulls {
325
                        let mut input = Vec::with_capacity(1 + input_right.len());
326
                        input.push(left_e);
327
                        input.extend_from_slice(input_right);
328
                        Some(AExpr::Function {
329
                            input,
330
                            function: fun.clone(),
331
                            options: *options,
332
                        })
333
                    } else {
334
                        None
335
                    }
336
                },
337
                _ => {
338
                    let function = IRStringFunction::ConcatHorizontal {
339
                        delimiter: "".into(),
340
                        ignore_nulls: false,
341
                    };
342
                    let options = function.function_options();
343
                    Some(AExpr::Function {
344
                        input: vec![left_e, right_e],
345
                        function: function.into(),
346
                        options,
347
                    })
348
                },
349
            }
350
        } else {
351
            None
352
        }
353
    }
354
}
355

356
pub struct SimplifyExprRule {}
357

358
impl OptimizationRule for SimplifyExprRule {
359
    #[allow(clippy::float_cmp)]
360
    fn optimize_expr(
361
        &mut self,
362
        expr_arena: &mut Arena<AExpr>,
363
        expr_node: Node,
364
        schema: &Schema,
365
        _ctx: OptimizeExprContext,
366
    ) -> PolarsResult<Option<AExpr>> {
367
        let expr = expr_arena.get(expr_node);
368

369
        let out = match &expr {
370
            AExpr::SortBy { expr, by, .. } if by.is_empty() => Some(expr_arena.get(*expr).clone()),
371
            // drop_nulls().len() -> len() - null_count()
372
            // drop_nulls().count() -> len() - null_count()
373
            AExpr::Agg(IRAggExpr::Count {
374
                input,
375
                include_nulls: _,
376
            }) => {
377
                let input_expr = expr_arena.get(*input);
378
                match input_expr {
379
                    AExpr::Function {
380
                        input,
381
                        function: IRFunctionExpr::DropNulls,
382
                        options: _,
383
                    } => {
384
                        // we should perform optimization only if the original expression is a column
385
                        // so in case of disabled CSE, we will not suffer from performance regression
386
                        if input.len() == 1 {
387
                            let drop_nulls_input_node = input[0].node();
388
                            match expr_arena.get(drop_nulls_input_node) {
389
                                AExpr::Column(_) => Some(AExpr::BinaryExpr {
390
                                    op: Operator::Minus,
391
                                    right: expr_arena.add(new_null_count(input)),
392
                                    left: expr_arena.add(AExpr::Agg(IRAggExpr::Count {
393
                                        input: drop_nulls_input_node,
394
                                        include_nulls: true,
395
                                    })),
396
                                }),
397
                                _ => None,
398
                            }
399
                        } else {
400
                            None
401
                        }
402
                    },
403
                    _ => None,
404
                }
405
            },
406
            // is_null().sum() -> null_count()
407
            // is_not_null().sum() -> len() - null_count()
408
            AExpr::Agg(IRAggExpr::Sum(input)) => {
409
                let input_expr = expr_arena.get(*input);
410
                match input_expr {
411
                    AExpr::Function {
412
                        input,
413
                        function: IRFunctionExpr::Boolean(IRBooleanFunction::IsNull),
414
                        options: _,
415
                    } => Some(new_null_count(input)),
416
                    AExpr::Function {
417
                        input,
418
                        function: IRFunctionExpr::Boolean(IRBooleanFunction::IsNotNull),
419
                        options: _,
420
                    } => {
421
                        // we should perform optimization only if the original expression is a column
422
                        // so in case of disabled CSE, we will not suffer from performance regression
423
                        if input.len() == 1 {
424
                            let is_not_null_input_node = input[0].node();
425
                            match expr_arena.get(is_not_null_input_node) {
426
                                AExpr::Column(_) => Some(AExpr::BinaryExpr {
427
                                    op: Operator::Minus,
428
                                    right: expr_arena.add(new_null_count(input)),
429
                                    left: expr_arena.add(AExpr::Agg(IRAggExpr::Count {
430
                                        input: is_not_null_input_node,
431
                                        include_nulls: true,
432
                                    })),
433
                                }),
434
                                _ => None,
435
                            }
436
                        } else {
437
                            None
438
                        }
439
                    },
440
                    _ => None,
441
                }
442
            },
443
            // lit(left) + lit(right) => lit(left + right)
444
            // and null propagation
445
            AExpr::BinaryExpr { left, op, right } => {
446
                let left_aexpr = expr_arena.get(*left);
447
                let right_aexpr = expr_arena.get(*right);
448

449
                // lit(left) + lit(right) => lit(left + right)
450
                use Operator::*;
451
                #[allow(clippy::manual_map)]
452
                let out = match op {
453
                    Plus => {
454
                        match eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| l + r) {
455
                            Some(new) => Some(new),
456
                            None => {
457
                                // try to replace addition of string columns with `concat_str`
458
                                #[cfg(all(feature = "strings", feature = "concat_str"))]
459
                                {
460
                                    string_addition_to_linear_concat(
461
                                        expr_arena,
462
                                        *left,
463
                                        *right,
464
                                        left_aexpr,
465
                                        right_aexpr,
466
                                        schema,
467
                                    )
468
                                }
469
                                #[cfg(not(all(feature = "strings", feature = "concat_str")))]
470
                                {
471
                                    None
472
                                }
473
                            },
474
                        }
475
                    },
476
                    Minus => eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| l - r),
477
                    Multiply => eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| l * r),
478
                    RustDivide => {
479
                        if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) =
480
                            (left_aexpr, right_aexpr)
481
                        {
482
                            match (lit_left, lit_right) {
483
                                (LiteralValue::Scalar(l), LiteralValue::Scalar(r)) => {
484
                                    match (l.as_any_value(), r.as_any_value()) {
485
                                        (AnyValue::Float16(x), AnyValue::Float16(y)) => {
486
                                            Some(AExpr::Literal(
487
                                                <Scalar as From<pf16>>::from(x / y).into(),
488
                                            ))
489
                                        },
490
                                        (AnyValue::Float32(x), AnyValue::Float32(y)) => {
491
                                            Some(AExpr::Literal(
492
                                                <Scalar as From<f32>>::from(x / y).into(),
493
                                            ))
494
                                        },
495
                                        (AnyValue::Float64(x), AnyValue::Float64(y)) => {
496
                                            Some(AExpr::Literal(
497
                                                <Scalar as From<f64>>::from(x / y).into(),
498
                                            ))
499
                                        },
500

501
                                        (AnyValue::Int8(x), AnyValue::Int8(y)) => {
502
                                            Some(AExpr::Literal(
503
                                                <Scalar as From<i8>>::from(
504
                                                    x.wrapping_floor_div_mod(y).0,
505
                                                )
506
                                                .into(),
507
                                            ))
508
                                        },
509
                                        (AnyValue::Int16(x), AnyValue::Int16(y)) => {
510
                                            Some(AExpr::Literal(
511
                                                <Scalar as From<i16>>::from(
512
                                                    x.wrapping_floor_div_mod(y).0,
513
                                                )
514
                                                .into(),
515
                                            ))
516
                                        },
517
                                        (AnyValue::Int32(x), AnyValue::Int32(y)) => {
518
                                            Some(AExpr::Literal(
519
                                                <Scalar as From<i32>>::from(
520
                                                    x.wrapping_floor_div_mod(y).0,
521
                                                )
522
                                                .into(),
523
                                            ))
524
                                        },
525
                                        (AnyValue::Int64(x), AnyValue::Int64(y)) => {
526
                                            Some(AExpr::Literal(
527
                                                <Scalar as From<i64>>::from(
528
                                                    x.wrapping_floor_div_mod(y).0,
529
                                                )
530
                                                .into(),
531
                                            ))
532
                                        },
533
                                        (AnyValue::Int128(x), AnyValue::Int128(y)) => {
534
                                            Some(AExpr::Literal(
535
                                                <Scalar as From<i128>>::from(
536
                                                    x.wrapping_floor_div_mod(y).0,
537
                                                )
538
                                                .into(),
539
                                            ))
540
                                        },
541

542
                                        (AnyValue::UInt8(x), AnyValue::UInt8(y)) => {
543
                                            Some(AExpr::Literal(
544
                                                <Scalar as From<u8>>::from(x / y).into(),
545
                                            ))
546
                                        },
547
                                        (AnyValue::UInt16(x), AnyValue::UInt16(y)) => {
548
                                            Some(AExpr::Literal(
549
                                                <Scalar as From<u16>>::from(x / y).into(),
550
                                            ))
551
                                        },
552
                                        (AnyValue::UInt32(x), AnyValue::UInt32(y)) => {
553
                                            Some(AExpr::Literal(
554
                                                <Scalar as From<u32>>::from(x / y).into(),
555
                                            ))
556
                                        },
557
                                        (AnyValue::UInt64(x), AnyValue::UInt64(y)) => {
558
                                            Some(AExpr::Literal(
559
                                                <Scalar as From<u64>>::from(x / y).into(),
560
                                            ))
561
                                        },
562
                                        (AnyValue::UInt128(x), AnyValue::UInt128(y)) => {
563
                                            Some(AExpr::Literal(
564
                                                <Scalar as From<u128>>::from(x / y).into(),
565
                                            ))
566
                                        },
567

568
                                        _ => None,
569
                                    }
570
                                },
571

572
                                (
573
                                    LiteralValue::Dyn(DynLiteralValue::Float(x)),
574
                                    LiteralValue::Dyn(DynLiteralValue::Float(y)),
575
                                ) => {
576
                                    Some(AExpr::Literal(<Scalar as From<f64>>::from(x / y).into()))
577
                                },
578
                                (
579
                                    LiteralValue::Dyn(DynLiteralValue::Int(x)),
580
                                    LiteralValue::Dyn(DynLiteralValue::Int(y)),
581
                                ) => Some(AExpr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(
582
                                    x.wrapping_floor_div_mod(*y).0,
583
                                )))),
584
                                _ => None,
585
                            }
586
                        } else {
587
                            None
588
                        }
589
                    },
590
                    TrueDivide => {
591
                        if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) =
592
                            (left_aexpr, right_aexpr)
593
                        {
594
                            match (lit_left, lit_right) {
595
                                (LiteralValue::Scalar(l), LiteralValue::Scalar(r)) => {
596
                                    match (l.as_any_value(), r.as_any_value()) {
597
                                        #[cfg(feature = "dtype-f16")]
598
                                        (AnyValue::Float16(x), AnyValue::Float16(y)) => {
599
                                            Some(AExpr::Literal(Scalar::from(x / y).into()))
600
                                        },
601
                                        (AnyValue::Float32(x), AnyValue::Float32(y)) => {
602
                                            Some(AExpr::Literal(Scalar::from(x / y).into()))
603
                                        },
604
                                        (AnyValue::Float64(x), AnyValue::Float64(y)) => {
605
                                            Some(AExpr::Literal(Scalar::from(x / y).into()))
606
                                        },
607

608
                                        (AnyValue::Int8(x), AnyValue::Int8(y)) => {
609
                                            Some(AExpr::Literal(
610
                                                Scalar::from(x as f64 / y as f64).into(),
611
                                            ))
612
                                        },
613
                                        (AnyValue::Int16(x), AnyValue::Int16(y)) => {
614
                                            Some(AExpr::Literal(
615
                                                Scalar::from(x as f64 / y as f64).into(),
616
                                            ))
617
                                        },
618
                                        (AnyValue::Int32(x), AnyValue::Int32(y)) => {
619
                                            Some(AExpr::Literal(
620
                                                Scalar::from(x as f64 / y as f64).into(),
621
                                            ))
622
                                        },
623
                                        (AnyValue::Int64(x), AnyValue::Int64(y)) => {
624
                                            Some(AExpr::Literal(
625
                                                Scalar::from(x as f64 / y as f64).into(),
626
                                            ))
627
                                        },
628
                                        (AnyValue::Int128(x), AnyValue::Int128(y)) => {
629
                                            Some(AExpr::Literal(
630
                                                Scalar::from(x as f64 / y as f64).into(),
631
                                            ))
632
                                        },
633

634
                                        (AnyValue::UInt8(x), AnyValue::UInt8(y)) => {
635
                                            Some(AExpr::Literal(
636
                                                Scalar::from(x as f64 / y as f64).into(),
637
                                            ))
638
                                        },
639
                                        (AnyValue::UInt16(x), AnyValue::UInt16(y)) => {
640
                                            Some(AExpr::Literal(
641
                                                Scalar::from(x as f64 / y as f64).into(),
642
                                            ))
643
                                        },
644
                                        (AnyValue::UInt32(x), AnyValue::UInt32(y)) => {
645
                                            Some(AExpr::Literal(
646
                                                Scalar::from(x as f64 / y as f64).into(),
647
                                            ))
648
                                        },
649
                                        (AnyValue::UInt64(x), AnyValue::UInt64(y)) => {
650
                                            Some(AExpr::Literal(
651
                                                Scalar::from(x as f64 / y as f64).into(),
652
                                            ))
653
                                        },
654

655
                                        _ => None,
656
                                    }
657
                                },
658

659
                                (
660
                                    LiteralValue::Dyn(DynLiteralValue::Float(x)),
661
                                    LiteralValue::Dyn(DynLiteralValue::Float(y)),
662
                                ) => Some(AExpr::Literal(Scalar::from(*x / *y).into())),
663
                                (
664
                                    LiteralValue::Dyn(DynLiteralValue::Int(x)),
665
                                    LiteralValue::Dyn(DynLiteralValue::Int(y)),
666
                                ) => {
667
                                    Some(AExpr::Literal(Scalar::from(*x as f64 / *y as f64).into()))
668
                                },
669
                                _ => None,
670
                            }
671
                        } else {
672
                            None
673
                        }
674
                    },
675
                    Modulus => eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| {
676
                        if r.is_zero() {
677
                            // TODO: this should optimize to `null` once we can express "is dynamic int but actually contains null"
678
                            return Ok(None);
679
                        }
680

681
                        l.wrapping_floor_div_mod(r).1
682
                    }),
683
                    Lt => eval_binary_cmp_same_type!(left_aexpr, <, right_aexpr),
684
                    Gt => eval_binary_cmp_same_type!(left_aexpr, >, right_aexpr),
685
                    Eq | EqValidity => eval_binary_cmp_same_type!(left_aexpr, ==, right_aexpr),
686
                    NotEq | NotEqValidity => {
687
                        eval_binary_cmp_same_type!(left_aexpr, !=, right_aexpr)
688
                    },
689
                    GtEq => eval_binary_cmp_same_type!(left_aexpr, >=, right_aexpr),
690
                    LtEq => eval_binary_cmp_same_type!(left_aexpr, <=, right_aexpr),
691
                    And | LogicalAnd => eval_bitwise(left_aexpr, right_aexpr, |l, r| l & r),
692
                    Or | LogicalOr => eval_bitwise(left_aexpr, right_aexpr, |l, r| l | r),
693
                    Xor => eval_bitwise(left_aexpr, right_aexpr, |l, r| l ^ r),
694
                    FloorDivide => eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| {
695
                        if r.is_zero() {
696
                            // TODO: this should optimize to `null` once we can express "is dynamic int but actually contains null"
697
                            return Ok(None);
698
                        }
699

700
                        l.wrapping_floor_div_mod(r).0
701
                    }),
702
                };
703
                if out.is_some() {
704
                    return Ok(out);
705
                }
706

707
                None
708
            },
709
            AExpr::Function {
710
                input,
711
                function,
712
                options,
713
                ..
714
            } => {
715
                return optimize_functions(input.clone(), function.clone(), *options, expr_arena);
716
            },
717
            _ => None,
718
        };
719
        Ok(out)
720
    }
721
}
722

723
Product

Resources

Company
