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 polars_utils::float16::pf16;
4
use polars_utils::floor_divmod::FloorDivMod;
5
use polars_utils::total_ord::ToTotalOrd;
6
use simplify_functions::optimize_functions;
7
mod arity;
8

9
use crate::plans::*;
10

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

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

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

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

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

101
macro_rules! eval_binary_cmp_same_type {
102
    ($lhs:expr, $operand: tt, $rhs:expr) => {{
103
    if let (AExpr::Literal(lit_left), AExpr::Literal(lit_right)) = ($lhs, $rhs) {
104
        match (lit_left, lit_right) {
105
            (LiteralValue::Scalar(l), LiteralValue::Scalar(r)) => match (l.as_any_value(), r.as_any_value()) {
106
                (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())),
107
                (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())),
108
                (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())),
109

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

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

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

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

140
    }}
141
}
142

143
pub struct SimplifyBooleanRule {}
144

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

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

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

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

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

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

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

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

355
pub struct SimplifyExprRule {}
356

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

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

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

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

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

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

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

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

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

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

658
                                (
659
                                    LiteralValue::Dyn(DynLiteralValue::Float(x)),
660
                                    LiteralValue::Dyn(DynLiteralValue::Float(y)),
661
                                ) => Some(AExpr::Literal(Scalar::from(*x / *y).into())),
662
                                (
663
                                    LiteralValue::Dyn(DynLiteralValue::Int(x)),
664
                                    LiteralValue::Dyn(DynLiteralValue::Int(y)),
665
                                ) => {
666
                                    Some(AExpr::Literal(Scalar::from(*x as f64 / *y as f64).into()))
667
                                },
668
                                _ => None,
669
                            }
670
                        } else {
671
                            None
672
                        }
673
                    },
674
                    Modulus => eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| l
675
                        .wrapping_floor_div_mod(r)
676
                        .1),
677
                    Lt => eval_binary_cmp_same_type!(left_aexpr, <, right_aexpr),
678
                    Gt => eval_binary_cmp_same_type!(left_aexpr, >, right_aexpr),
679
                    Eq | EqValidity => eval_binary_cmp_same_type!(left_aexpr, ==, right_aexpr),
680
                    NotEq | NotEqValidity => {
681
                        eval_binary_cmp_same_type!(left_aexpr, !=, right_aexpr)
682
                    },
683
                    GtEq => eval_binary_cmp_same_type!(left_aexpr, >=, right_aexpr),
684
                    LtEq => eval_binary_cmp_same_type!(left_aexpr, <=, right_aexpr),
685
                    And | LogicalAnd => eval_bitwise(left_aexpr, right_aexpr, |l, r| l & r),
686
                    Or | LogicalOr => eval_bitwise(left_aexpr, right_aexpr, |l, r| l | r),
687
                    Xor => eval_bitwise(left_aexpr, right_aexpr, |l, r| l ^ r),
688
                    FloorDivide => eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| l
689
                        .wrapping_floor_div_mod(r)
690
                        .0),
691
                };
692
                if out.is_some() {
693
                    return Ok(out);
694
                }
695

696
                None
697
            },
698
            AExpr::Function {
699
                input,
700
                function,
701
                options,
702
                ..
703
            } => {
704
                return optimize_functions(input.clone(), function.clone(), *options, expr_arena);
705
            },
706
            _ => None,
707
        };
708
        Ok(out)
709
    }
710
}
711

712
Product

Resources

Company
