CoCalc -- ir_to

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-plan/src/plans/conversion/ir_to_dsl.rs
⁶⁹⁴⁰ views
1
use super::*;
2

3
/// converts a node from the AExpr arena to Expr
4
#[recursive]
5
pub fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
6
    let expr = expr_arena.get(node).clone();
7

8
    match expr {
9
        AExpr::Explode { expr, skip_empty } => Expr::Explode {
10
            input: Arc::new(node_to_expr(expr, expr_arena)),
11
            skip_empty,
12
        },
13
        AExpr::Column(a) => Expr::Column(a),
14
        AExpr::Literal(s) => Expr::Literal(s),
15
        AExpr::BinaryExpr { left, op, right } => {
16
            let l = node_to_expr(left, expr_arena);
17
            let r = node_to_expr(right, expr_arena);
18
            Expr::BinaryExpr {
19
                left: Arc::new(l),
20
                op,
21
                right: Arc::new(r),
22
            }
23
        },
24
        AExpr::Cast {
25
            expr,
26
            dtype,
27
            options: strict,
28
        } => {
29
            let exp = node_to_expr(expr, expr_arena);
30
            Expr::Cast {
31
                expr: Arc::new(exp),
32
                dtype: dtype.into(),
33
                options: strict,
34
            }
35
        },
36
        AExpr::Sort { expr, options } => {
37
            let exp = node_to_expr(expr, expr_arena);
38
            Expr::Sort {
39
                expr: Arc::new(exp),
40
                options,
41
            }
42
        },
43
        AExpr::Gather {
44
            expr,
45
            idx,
46
            returns_scalar,
47
        } => {
48
            let expr = node_to_expr(expr, expr_arena);
49
            let idx = node_to_expr(idx, expr_arena);
50
            Expr::Gather {
51
                expr: Arc::new(expr),
52
                idx: Arc::new(idx),
53
                returns_scalar,
54
            }
55
        },
56
        AExpr::SortBy {
57
            expr,
58
            by,
59
            sort_options,
60
        } => {
61
            let expr = node_to_expr(expr, expr_arena);
62
            let by = by
63
                .iter()
64
                .map(|node| node_to_expr(*node, expr_arena))
65
                .collect();
66
            Expr::SortBy {
67
                expr: Arc::new(expr),
68
                by,
69
                sort_options,
70
            }
71
        },
72
        AExpr::Filter { input, by } => {
73
            let input = node_to_expr(input, expr_arena);
74
            let by = node_to_expr(by, expr_arena);
75
            Expr::Filter {
76
                input: Arc::new(input),
77
                by: Arc::new(by),
78
            }
79
        },
80
        AExpr::Agg(agg) => match agg {
81
            IRAggExpr::Min {
82
                input,
83
                propagate_nans,
84
            } => {
85
                let exp = node_to_expr(input, expr_arena);
86
                AggExpr::Min {
87
                    input: Arc::new(exp),
88
                    propagate_nans,
89
                }
90
                .into()
91
            },
92
            IRAggExpr::Max {
93
                input,
94
                propagate_nans,
95
            } => {
96
                let exp = node_to_expr(input, expr_arena);
97
                AggExpr::Max {
98
                    input: Arc::new(exp),
99
                    propagate_nans,
100
                }
101
                .into()
102
            },
103

104
            IRAggExpr::Mean(expr) => {
105
                let exp = node_to_expr(expr, expr_arena);
106
                AggExpr::Mean(Arc::new(exp)).into()
107
            },
108
            IRAggExpr::Median(expr) => {
109
                let exp = node_to_expr(expr, expr_arena);
110
                AggExpr::Median(Arc::new(exp)).into()
111
            },
112
            IRAggExpr::NUnique(expr) => {
113
                let exp = node_to_expr(expr, expr_arena);
114
                AggExpr::NUnique(Arc::new(exp)).into()
115
            },
116
            IRAggExpr::First(expr) => {
117
                let exp = node_to_expr(expr, expr_arena);
118
                AggExpr::First(Arc::new(exp)).into()
119
            },
120
            IRAggExpr::Last(expr) => {
121
                let exp = node_to_expr(expr, expr_arena);
122
                AggExpr::Last(Arc::new(exp)).into()
123
            },
124
            IRAggExpr::Implode(expr) => {
125
                let exp = node_to_expr(expr, expr_arena);
126
                AggExpr::Implode(Arc::new(exp)).into()
127
            },
128
            IRAggExpr::Quantile {
129
                expr,
130
                quantile,
131
                method,
132
            } => {
133
                let expr = node_to_expr(expr, expr_arena);
134
                let quantile = node_to_expr(quantile, expr_arena);
135
                AggExpr::Quantile {
136
                    expr: Arc::new(expr),
137
                    quantile: Arc::new(quantile),
138
                    method,
139
                }
140
                .into()
141
            },
142
            IRAggExpr::Sum(expr) => {
143
                let exp = node_to_expr(expr, expr_arena);
144
                AggExpr::Sum(Arc::new(exp)).into()
145
            },
146
            IRAggExpr::Std(expr, ddof) => {
147
                let exp = node_to_expr(expr, expr_arena);
148
                AggExpr::Std(Arc::new(exp), ddof).into()
149
            },
150
            IRAggExpr::Var(expr, ddof) => {
151
                let exp = node_to_expr(expr, expr_arena);
152
                AggExpr::Var(Arc::new(exp), ddof).into()
153
            },
154
            IRAggExpr::AggGroups(expr) => {
155
                let exp = node_to_expr(expr, expr_arena);
156
                AggExpr::AggGroups(Arc::new(exp)).into()
157
            },
158
            IRAggExpr::Count {
159
                input,
160
                include_nulls,
161
            } => {
162
                let input = node_to_expr(input, expr_arena);
163
                AggExpr::Count {
164
                    input: Arc::new(input),
165
                    include_nulls,
166
                }
167
                .into()
168
            },
169
        },
170
        AExpr::Ternary {
171
            predicate,
172
            truthy,
173
            falsy,
174
        } => {
175
            let p = node_to_expr(predicate, expr_arena);
176
            let t = node_to_expr(truthy, expr_arena);
177
            let f = node_to_expr(falsy, expr_arena);
178

179
            Expr::Ternary {
180
                predicate: Arc::new(p),
181
                truthy: Arc::new(t),
182
                falsy: Arc::new(f),
183
            }
184
        },
185
        AExpr::AnonymousFunction {
186
            input,
187
            function,
188
            options,
189
            fmt_str,
190
        } => Expr::AnonymousFunction {
191
            input: expr_irs_to_exprs(input, expr_arena),
192
            function,
193
            options,
194
            fmt_str,
195
        },
196
        AExpr::Eval {
197
            expr,
198
            evaluation,
199
            variant,
200
        } => Expr::Eval {
201
            expr: Arc::new(node_to_expr(expr, expr_arena)),
202
            evaluation: Arc::new(node_to_expr(evaluation, expr_arena)),
203
            variant,
204
        },
205
        AExpr::Function {
206
            input,
207
            function,
208
            options: _,
209
        } => {
210
            let input = expr_irs_to_exprs(input, expr_arena);
211
            ir_function_to_dsl(input, function)
212
        },
213
        AExpr::Window {
214
            function,
215
            partition_by,
216
            order_by,
217
            options,
218
        } => {
219
            let function = Arc::new(node_to_expr(function, expr_arena));
220
            let partition_by = nodes_to_exprs(&partition_by, expr_arena);
221
            let order_by =
222
                order_by.map(|(n, options)| (Arc::new(node_to_expr(n, expr_arena)), options));
223
            Expr::Window {
224
                function,
225
                partition_by,
226
                order_by,
227
                options,
228
            }
229
        },
230
        AExpr::Slice {
231
            input,
232
            offset,
233
            length,
234
        } => Expr::Slice {
235
            input: Arc::new(node_to_expr(input, expr_arena)),
236
            offset: Arc::new(node_to_expr(offset, expr_arena)),
237
            length: Arc::new(node_to_expr(length, expr_arena)),
238
        },
239
        AExpr::Len => Expr::Len,
240
    }
241
}
242

243
fn nodes_to_exprs(nodes: &[Node], expr_arena: &Arena<AExpr>) -> Vec<Expr> {
244
    nodes.iter().map(|n| node_to_expr(*n, expr_arena)).collect()
245
}
246

247
pub fn ir_function_to_dsl(input: Vec<Expr>, function: IRFunctionExpr) -> Expr {
248
    use {FunctionExpr as F, IRFunctionExpr as IF};
249

250
    let function = match function {
251
        #[cfg(feature = "dtype-array")]
252
        IF::ArrayExpr(f) => {
253
            use {ArrayFunction as A, IRArrayFunction as IA};
254
            F::ArrayExpr(match f {
255
                IA::Concat => A::Concat,
256
                IA::Length => A::Length,
257
                IA::Min => A::Min,
258
                IA::Max => A::Max,
259
                IA::Sum => A::Sum,
260
                IA::ToList => A::ToList,
261
                IA::Unique(v) => A::Unique(v),
262
                IA::NUnique => A::NUnique,
263
                IA::Std(v) => A::Std(v),
264
                IA::Var(v) => A::Var(v),
265
                IA::Mean => A::Mean,
266
                IA::Median => A::Median,
267
                #[cfg(feature = "array_any_all")]
268
                IA::Any => A::Any,
269
                #[cfg(feature = "array_any_all")]
270
                IA::All => A::All,
271
                IA::Sort(v) => A::Sort(v),
272
                IA::Reverse => A::Reverse,
273
                IA::ArgMin => A::ArgMin,
274
                IA::ArgMax => A::ArgMax,
275
                IA::Get(v) => A::Get(v),
276
                IA::Join(v) => A::Join(v),
277
                #[cfg(feature = "is_in")]
278
                IA::Contains { nulls_equal } => A::Contains { nulls_equal },
279
                #[cfg(feature = "array_count")]
280
                IA::CountMatches => A::CountMatches,
281
                IA::Shift => A::Shift,
282
                IA::Slice(offset, length) => A::Slice(offset, length),
283
                IA::Explode { skip_empty } => A::Explode { skip_empty },
284
                #[cfg(feature = "array_to_struct")]
285
                IA::ToStruct(ng) => A::ToStruct(ng),
286
            })
287
        },
288
        IF::BinaryExpr(f) => {
289
            use {BinaryFunction as B, IRBinaryFunction as IB};
290
            F::BinaryExpr(match f {
291
                IB::Contains => B::Contains,
292
                IB::StartsWith => B::StartsWith,
293
                IB::EndsWith => B::EndsWith,
294
                #[cfg(feature = "binary_encoding")]
295
                IB::HexDecode(v) => B::HexDecode(v),
296
                #[cfg(feature = "binary_encoding")]
297
                IB::HexEncode => B::HexEncode,
298
                #[cfg(feature = "binary_encoding")]
299
                IB::Base64Decode(v) => B::Base64Decode(v),
300
                #[cfg(feature = "binary_encoding")]
301
                IB::Base64Encode => B::Base64Encode,
302
                IB::Size => B::Size,
303
                #[cfg(feature = "binary_encoding")]
304
                IB::Reinterpret(data_type, v) => B::Reinterpret(data_type.into(), v),
305
            })
306
        },
307
        #[cfg(feature = "dtype-categorical")]
308
        IF::Categorical(f) => {
309
            use {CategoricalFunction as C, IRCategoricalFunction as IC};
310
            F::Categorical(match f {
311
                IC::GetCategories => C::GetCategories,
312
                #[cfg(feature = "strings")]
313
                IC::LenBytes => C::LenBytes,
314
                #[cfg(feature = "strings")]
315
                IC::LenChars => C::LenChars,
316
                #[cfg(feature = "strings")]
317
                IC::StartsWith(v) => C::StartsWith(v),
318
                #[cfg(feature = "strings")]
319
                IC::EndsWith(v) => C::EndsWith(v),
320
                #[cfg(feature = "strings")]
321
                IC::Slice(s, l) => C::Slice(s, l),
322
            })
323
        },
324
        IF::ListExpr(f) => {
325
            use {IRListFunction as IL, ListFunction as L};
326
            F::ListExpr(match f {
327
                IL::Concat => L::Concat,
328
                #[cfg(feature = "is_in")]
329
                IL::Contains { nulls_equal } => L::Contains { nulls_equal },
330
                #[cfg(feature = "list_drop_nulls")]
331
                IL::DropNulls => L::DropNulls,
332
                #[cfg(feature = "list_sample")]
333
                IL::Sample {
334
                    is_fraction,
335
                    with_replacement,
336
                    shuffle,
337
                    seed,
338
                } => L::Sample {
339
                    is_fraction,
340
                    with_replacement,
341
                    shuffle,
342
                    seed,
343
                },
344
                IL::Slice => L::Slice,
345
                IL::Shift => L::Shift,
346
                IL::Get(v) => L::Get(v),
347
                #[cfg(feature = "list_gather")]
348
                IL::Gather(v) => L::Gather(v),
349
                #[cfg(feature = "list_gather")]
350
                IL::GatherEvery => L::GatherEvery,
351
                #[cfg(feature = "list_count")]
352
                IL::CountMatches => L::CountMatches,
353
                IL::Sum => L::Sum,
354
                IL::Length => L::Length,
355
                IL::Max => L::Max,
356
                IL::Min => L::Min,
357
                IL::Mean => L::Mean,
358
                IL::Median => L::Median,
359
                IL::Std(v) => L::Std(v),
360
                IL::Var(v) => L::Var(v),
361
                IL::ArgMin => L::ArgMin,
362
                IL::ArgMax => L::ArgMax,
363
                #[cfg(feature = "diff")]
364
                IL::Diff { n, null_behavior } => L::Diff { n, null_behavior },
365
                IL::Sort(sort_options) => L::Sort(sort_options),
366
                IL::Reverse => L::Reverse,
367
                IL::Unique(v) => L::Unique(v),
368
                IL::NUnique => L::NUnique,
369
                #[cfg(feature = "list_sets")]
370
                IL::SetOperation(set_operation) => L::SetOperation(set_operation),
371
                #[cfg(feature = "list_any_all")]
372
                IL::Any => L::Any,
373
                #[cfg(feature = "list_any_all")]
374
                IL::All => L::All,
375
                IL::Join(v) => L::Join(v),
376
                #[cfg(feature = "dtype-array")]
377
                IL::ToArray(v) => L::ToArray(v),
378
                #[cfg(feature = "list_to_struct")]
379
                IL::ToStruct(list_to_struct_args) => L::ToStruct(list_to_struct_args),
380
            })
381
        },
382
        #[cfg(feature = "strings")]
383
        IF::StringExpr(f) => {
384
            use {IRStringFunction as IB, StringFunction as B};
385
            F::StringExpr(match f {
386
                #[cfg(feature = "concat_str")]
387
                IB::ConcatHorizontal {
388
                    delimiter,
389
                    ignore_nulls,
390
                } => B::ConcatHorizontal {
391
                    delimiter,
392
                    ignore_nulls,
393
                },
394
                #[cfg(feature = "concat_str")]
395
                IB::ConcatVertical {
396
                    delimiter,
397
                    ignore_nulls,
398
                } => B::ConcatVertical {
399
                    delimiter,
400
                    ignore_nulls,
401
                },
402
                #[cfg(feature = "regex")]
403
                IB::Contains { literal, strict } => B::Contains { literal, strict },
404
                IB::CountMatches(v) => B::CountMatches(v),
405
                IB::EndsWith => B::EndsWith,
406
                IB::Extract(v) => B::Extract(v),
407
                IB::ExtractAll => B::ExtractAll,
408
                #[cfg(feature = "extract_groups")]
409
                IB::ExtractGroups { dtype, pat } => B::ExtractGroups { dtype, pat },
410
                #[cfg(feature = "regex")]
411
                IB::Find { literal, strict } => B::Find { literal, strict },
412
                #[cfg(feature = "string_to_integer")]
413
                IB::ToInteger { dtype, strict } => B::ToInteger { dtype, strict },
414
                IB::LenBytes => B::LenBytes,
415
                IB::LenChars => B::LenChars,
416
                IB::Lowercase => B::Lowercase,
417
                #[cfg(feature = "extract_jsonpath")]
418
                IB::JsonDecode(dtype) => B::JsonDecode(dtype.into()),
419
                #[cfg(feature = "extract_jsonpath")]
420
                IB::JsonPathMatch => B::JsonPathMatch,
421
                #[cfg(feature = "regex")]
422
                IB::Replace { n, literal } => B::Replace { n, literal },
423
                #[cfg(feature = "string_normalize")]
424
                IB::Normalize { form } => B::Normalize { form },
425
                #[cfg(feature = "string_reverse")]
426
                IB::Reverse => B::Reverse,
427
                #[cfg(feature = "string_pad")]
428
                IB::PadStart { fill_char } => B::PadStart { fill_char },
429
                #[cfg(feature = "string_pad")]
430
                IB::PadEnd { fill_char } => B::PadEnd { fill_char },
431
                IB::Slice => B::Slice,
432
                IB::Head => B::Head,
433
                IB::Tail => B::Tail,
434
                #[cfg(feature = "string_encoding")]
435
                IB::HexEncode => B::HexEncode,
436
                #[cfg(feature = "binary_encoding")]
437
                IB::HexDecode(v) => B::HexDecode(v),
438
                #[cfg(feature = "string_encoding")]
439
                IB::Base64Encode => B::Base64Encode,
440
                #[cfg(feature = "binary_encoding")]
441
                IB::Base64Decode(v) => B::Base64Decode(v),
442
                IB::StartsWith => B::StartsWith,
443
                IB::StripChars => B::StripChars,
444
                IB::StripCharsStart => B::StripCharsStart,
445
                IB::StripCharsEnd => B::StripCharsEnd,
446
                IB::StripPrefix => B::StripPrefix,
447
                IB::StripSuffix => B::StripSuffix,
448
                #[cfg(feature = "dtype-struct")]
449
                IB::SplitExact { n, inclusive } => B::SplitExact { n, inclusive },
450
                #[cfg(feature = "dtype-struct")]
451
                IB::SplitN(n) => B::SplitN(n),
452
                #[cfg(feature = "temporal")]
453
                IB::Strptime(dtype, strptime_options) => {
454
                    B::Strptime(dtype.into(), strptime_options)
455
                },
456
                IB::Split(v) => B::Split(v),
457
                #[cfg(feature = "dtype-decimal")]
458
                IB::ToDecimal { scale } => B::ToDecimal { scale },
459
                #[cfg(feature = "nightly")]
460
                IB::Titlecase => B::Titlecase,
461
                IB::Uppercase => B::Uppercase,
462
                #[cfg(feature = "string_pad")]
463
                IB::ZFill => B::ZFill,
464
                #[cfg(feature = "find_many")]
465
                IB::ContainsAny {
466
                    ascii_case_insensitive,
467
                } => B::ContainsAny {
468
                    ascii_case_insensitive,
469
                },
470
                #[cfg(feature = "find_many")]
471
                IB::ReplaceMany {
472
                    ascii_case_insensitive,
473
                } => B::ReplaceMany {
474
                    ascii_case_insensitive,
475
                },
476
                #[cfg(feature = "find_many")]
477
                IB::ExtractMany {
478
                    ascii_case_insensitive,
479
                    overlapping,
480
                } => B::ExtractMany {
481
                    ascii_case_insensitive,
482
                    overlapping,
483
                },
484
                #[cfg(feature = "find_many")]
485
                IB::FindMany {
486
                    ascii_case_insensitive,
487
                    overlapping,
488
                } => B::FindMany {
489
                    ascii_case_insensitive,
490
                    overlapping,
491
                },
492
                #[cfg(feature = "regex")]
493
                IB::EscapeRegex => B::EscapeRegex,
494
            })
495
        },
496
        #[cfg(feature = "dtype-struct")]
497
        IF::StructExpr(f) => {
498
            use {IRStructFunction as IB, StructFunction as B};
499
            F::StructExpr(match f {
500
                IB::FieldByName(pl_small_str) => B::FieldByName(pl_small_str),
501
                IB::RenameFields(pl_small_strs) => B::RenameFields(pl_small_strs),
502
                IB::PrefixFields(pl_small_str) => B::PrefixFields(pl_small_str),
503
                IB::SuffixFields(pl_small_str) => B::SuffixFields(pl_small_str),
504
                #[cfg(feature = "json")]
505
                IB::JsonEncode => B::JsonEncode,
506
                IB::WithFields => B::WithFields,
507
                #[cfg(feature = "python")]
508
                IB::MapFieldNames(special_eq) => B::MapFieldNames(special_eq),
509
            })
510
        },
511
        #[cfg(feature = "temporal")]
512
        IF::TemporalExpr(f) => {
513
            use {IRTemporalFunction as IB, TemporalFunction as B};
514
            F::TemporalExpr(match f {
515
                IB::Millennium => B::Millennium,
516
                IB::Century => B::Century,
517
                IB::Year => B::Year,
518
                IB::IsLeapYear => B::IsLeapYear,
519
                IB::IsoYear => B::IsoYear,
520
                IB::Quarter => B::Quarter,
521
                IB::Month => B::Month,
522
                IB::DaysInMonth => B::DaysInMonth,
523
                IB::Week => B::Week,
524
                IB::WeekDay => B::WeekDay,
525
                IB::Day => B::Day,
526
                IB::OrdinalDay => B::OrdinalDay,
527
                IB::Time => B::Time,
528
                IB::Date => B::Date,
529
                IB::Datetime => B::Datetime,
530
                #[cfg(feature = "dtype-duration")]
531
                IB::Duration(time_unit) => B::Duration(time_unit),
532
                IB::Hour => B::Hour,
533
                IB::Minute => B::Minute,
534
                IB::Second => B::Second,
535
                IB::Millisecond => B::Millisecond,
536
                IB::Microsecond => B::Microsecond,
537
                IB::Nanosecond => B::Nanosecond,
538
                #[cfg(feature = "dtype-duration")]
539
                IB::TotalDays => B::TotalDays,
540
                #[cfg(feature = "dtype-duration")]
541
                IB::TotalHours => B::TotalHours,
542
                #[cfg(feature = "dtype-duration")]
543
                IB::TotalMinutes => B::TotalMinutes,
544
                #[cfg(feature = "dtype-duration")]
545
                IB::TotalSeconds => B::TotalSeconds,
546
                #[cfg(feature = "dtype-duration")]
547
                IB::TotalMilliseconds => B::TotalMilliseconds,
548
                #[cfg(feature = "dtype-duration")]
549
                IB::TotalMicroseconds => B::TotalMicroseconds,
550
                #[cfg(feature = "dtype-duration")]
551
                IB::TotalNanoseconds => B::TotalNanoseconds,
552
                IB::ToString(v) => B::ToString(v),
553
                IB::CastTimeUnit(time_unit) => B::CastTimeUnit(time_unit),
554
                IB::WithTimeUnit(time_unit) => B::WithTimeUnit(time_unit),
555
                #[cfg(feature = "timezones")]
556
                IB::ConvertTimeZone(time_zone) => B::ConvertTimeZone(time_zone),
557
                IB::TimeStamp(time_unit) => B::TimeStamp(time_unit),
558
                IB::Truncate => B::Truncate,
559
                #[cfg(feature = "offset_by")]
560
                IB::OffsetBy => B::OffsetBy,
561
                #[cfg(feature = "month_start")]
562
                IB::MonthStart => B::MonthStart,
563
                #[cfg(feature = "month_end")]
564
                IB::MonthEnd => B::MonthEnd,
565
                #[cfg(feature = "timezones")]
566
                IB::BaseUtcOffset => B::BaseUtcOffset,
567
                #[cfg(feature = "timezones")]
568
                IB::DSTOffset => B::DSTOffset,
569
                IB::Round => B::Round,
570
                IB::Replace => B::Replace,
571
                #[cfg(feature = "timezones")]
572
                IB::ReplaceTimeZone(time_zone, non_existent) => {
573
                    B::ReplaceTimeZone(time_zone, non_existent)
574
                },
575
                IB::Combine(time_unit) => B::Combine(time_unit),
576
                IB::DatetimeFunction {
577
                    time_unit,
578
                    time_zone,
579
                } => B::DatetimeFunction {
580
                    time_unit,
581
                    time_zone,
582
                },
583
            })
584
        },
585
        #[cfg(feature = "bitwise")]
586
        IF::Bitwise(f) => {
587
            use {BitwiseFunction as B, IRBitwiseFunction as IB};
588
            F::Bitwise(match f {
589
                IB::CountOnes => B::CountOnes,
590
                IB::CountZeros => B::CountZeros,
591
                IB::LeadingOnes => B::LeadingOnes,
592
                IB::LeadingZeros => B::LeadingZeros,
593
                IB::TrailingOnes => B::TrailingOnes,
594
                IB::TrailingZeros => B::TrailingZeros,
595
                IB::And => B::And,
596
                IB::Or => B::Or,
597
                IB::Xor => B::Xor,
598
            })
599
        },
600
        IF::Boolean(f) => {
601
            use {BooleanFunction as B, IRBooleanFunction as IB};
602
            F::Boolean(match f {
603
                IB::Any { ignore_nulls } => B::Any { ignore_nulls },
604
                IB::All { ignore_nulls } => B::All { ignore_nulls },
605
                IB::IsNull => B::IsNull,
606
                IB::IsNotNull => B::IsNotNull,
607
                IB::IsFinite => B::IsFinite,
608
                IB::IsInfinite => B::IsInfinite,
609
                IB::IsNan => B::IsNan,
610
                IB::IsNotNan => B::IsNotNan,
611
                #[cfg(feature = "is_first_distinct")]
612
                IB::IsFirstDistinct => B::IsFirstDistinct,
613
                #[cfg(feature = "is_last_distinct")]
614
                IB::IsLastDistinct => B::IsLastDistinct,
615
                #[cfg(feature = "is_unique")]
616
                IB::IsUnique => B::IsUnique,
617
                #[cfg(feature = "is_unique")]
618
                IB::IsDuplicated => B::IsDuplicated,
619
                #[cfg(feature = "is_between")]
620
                IB::IsBetween { closed } => B::IsBetween { closed },
621
                #[cfg(feature = "is_in")]
622
                IB::IsIn { nulls_equal } => B::IsIn { nulls_equal },
623
                #[cfg(feature = "is_close")]
624
                IB::IsClose {
625
                    abs_tol,
626
                    rel_tol,
627
                    nans_equal,
628
                } => B::IsClose {
629
                    abs_tol,
630
                    rel_tol,
631
                    nans_equal,
632
                },
633
                IB::AllHorizontal => B::AllHorizontal,
634
                IB::AnyHorizontal => B::AnyHorizontal,
635
                IB::Not => B::Not,
636
            })
637
        },
638
        #[cfg(feature = "business")]
639
        IF::Business(f) => {
640
            use {BusinessFunction as B, IRBusinessFunction as IB};
641
            F::Business(match f {
642
                IB::BusinessDayCount {
643
                    week_mask,
644
                    holidays,
645
                } => B::BusinessDayCount {
646
                    week_mask,
647
                    holidays,
648
                },
649
                IB::AddBusinessDay {
650
                    week_mask,
651
                    holidays,
652
                    roll,
653
                } => B::AddBusinessDay {
654
                    week_mask,
655
                    holidays,
656
                    roll,
657
                },
658
                IB::IsBusinessDay {
659
                    week_mask,
660
                    holidays,
661
                } => B::IsBusinessDay {
662
                    week_mask,
663
                    holidays,
664
                },
665
            })
666
        },
667
        #[cfg(feature = "abs")]
668
        IF::Abs => F::Abs,
669
        IF::Negate => F::Negate,
670
        #[cfg(feature = "hist")]
671
        IF::Hist {
672
            bin_count,
673
            include_category,
674
            include_breakpoint,
675
        } => F::Hist {
676
            bin_count,
677
            include_category,
678
            include_breakpoint,
679
        },
680
        IF::NullCount => F::NullCount,
681
        IF::Pow(f) => {
682
            use {IRPowFunction as IP, PowFunction as P};
683
            F::Pow(match f {
684
                IP::Generic => P::Generic,
685
                IP::Sqrt => P::Sqrt,
686
                IP::Cbrt => P::Cbrt,
687
            })
688
        },
689
        #[cfg(feature = "row_hash")]
690
        IF::Hash(s0, s1, s2, s3) => F::Hash(s0, s1, s2, s3),
691
        #[cfg(feature = "arg_where")]
692
        IF::ArgWhere => F::ArgWhere,
693
        #[cfg(feature = "index_of")]
694
        IF::IndexOf => F::IndexOf,
695
        #[cfg(feature = "search_sorted")]
696
        IF::SearchSorted { side, descending } => F::SearchSorted { side, descending },
697
        #[cfg(feature = "range")]
698
        IF::Range(f) => {
699
            use {IRRangeFunction as IR, RangeFunction as R};
700
            F::Range(match f {
701
                IR::IntRange { step, dtype } => R::IntRange {
702
                    step,
703
                    dtype: dtype.into(),
704
                },
705
                IR::IntRanges { dtype } => R::IntRanges {
706
                    dtype: dtype.into(),
707
                },
708
                IR::LinearSpace { closed } => R::LinearSpace { closed },
709
                IR::LinearSpaces {
710
                    closed,
711
                    array_width,
712
                } => R::LinearSpaces {
713
                    closed,
714
                    array_width,
715
                },
716
                #[cfg(feature = "dtype-date")]
717
                IR::DateRange { interval, closed } => R::DateRange { interval, closed },
718
                #[cfg(feature = "dtype-date")]
719
                IR::DateRanges { interval, closed } => R::DateRanges { interval, closed },
720
                #[cfg(feature = "dtype-datetime")]
721
                IR::DatetimeRange {
722
                    interval,
723
                    closed,
724
                    time_unit,
725
                    time_zone,
726
                } => R::DatetimeRange {
727
                    interval,
728
                    closed,
729
                    time_unit,
730
                    time_zone,
731
                },
732
                #[cfg(feature = "dtype-datetime")]
733
                IR::DatetimeRanges {
734
                    interval,
735
                    closed,
736
                    time_unit,
737
                    time_zone,
738
                } => R::DatetimeRanges {
739
                    interval,
740
                    closed,
741
                    time_unit,
742
                    time_zone,
743
                },
744
                #[cfg(feature = "dtype-time")]
745
                IR::TimeRange { interval, closed } => R::TimeRange { interval, closed },
746
                #[cfg(feature = "dtype-time")]
747
                IR::TimeRanges { interval, closed } => R::TimeRanges { interval, closed },
748
            })
749
        },
750
        #[cfg(feature = "trigonometry")]
751
        IF::Trigonometry(f) => {
752
            use {IRTrigonometricFunction as IT, TrigonometricFunction as T};
753
            F::Trigonometry(match f {
754
                IT::Cos => T::Cos,
755
                IT::Cot => T::Cot,
756
                IT::Sin => T::Sin,
757
                IT::Tan => T::Tan,
758
                IT::ArcCos => T::ArcCos,
759
                IT::ArcSin => T::ArcSin,
760
                IT::ArcTan => T::ArcTan,
761
                IT::Cosh => T::Cosh,
762
                IT::Sinh => T::Sinh,
763
                IT::Tanh => T::Tanh,
764
                IT::ArcCosh => T::ArcCosh,
765
                IT::ArcSinh => T::ArcSinh,
766
                IT::ArcTanh => T::ArcTanh,
767
                IT::Degrees => T::Degrees,
768
                IT::Radians => T::Radians,
769
            })
770
        },
771
        #[cfg(feature = "trigonometry")]
772
        IF::Atan2 => F::Atan2,
773
        #[cfg(feature = "sign")]
774
        IF::Sign => F::Sign,
775
        IF::FillNull => F::FillNull,
776
        IF::FillNullWithStrategy(strategy) => F::FillNullWithStrategy(strategy),
777
        #[cfg(feature = "rolling_window")]
778
        IF::RollingExpr { function, options } => {
779
            use {IRRollingFunction as IR, RollingFunction as R};
780
            FunctionExpr::RollingExpr {
781
                function: match function {
782
                    IR::Min => R::Min,
783
                    IR::Max => R::Max,
784
                    IR::Mean => R::Mean,
785
                    IR::Sum => R::Sum,
786
                    IR::Quantile => R::Quantile,
787
                    IR::Var => R::Var,
788
                    IR::Std => R::Std,
789
                    #[cfg(feature = "moment")]
790
                    IR::Skew => R::Skew,
791
                    #[cfg(feature = "moment")]
792
                    IR::Kurtosis => R::Kurtosis,
793
                    #[cfg(feature = "cov")]
794
                    IR::CorrCov {
795
                        corr_cov_options,
796
                        is_corr,
797
                    } => R::CorrCov {
798
                        corr_cov_options,
799
                        is_corr,
800
                    },
801
                    IR::Map(f) => R::Map(f),
802
                },
803
                options,
804
            }
805
        },
806
        #[cfg(feature = "rolling_window_by")]
807
        IF::RollingExprBy {
808
            function_by,
809
            options,
810
        } => {
811
            use {IRRollingFunctionBy as IR, RollingFunctionBy as R};
812
            FunctionExpr::RollingExprBy {
813
                function_by: match function_by {
814
                    IR::MinBy => R::MinBy,
815
                    IR::MaxBy => R::MaxBy,
816
                    IR::MeanBy => R::MeanBy,
817
                    IR::SumBy => R::SumBy,
818
                    IR::QuantileBy => R::QuantileBy,
819
                    IR::VarBy => R::VarBy,
820
                    IR::StdBy => R::StdBy,
821
                },
822
                options,
823
            }
824
        },
825
        IF::Rechunk => F::Rechunk,
826
        IF::Append { upcast } => F::Append { upcast },
827
        IF::ShiftAndFill => F::ShiftAndFill,
828
        IF::Shift => F::Shift,
829
        IF::DropNans => F::DropNans,
830
        IF::DropNulls => F::DropNulls,
831
        #[cfg(feature = "mode")]
832
        IF::Mode => F::Mode,
833
        #[cfg(feature = "moment")]
834
        IF::Skew(v) => F::Skew(v),
835
        #[cfg(feature = "moment")]
836
        IF::Kurtosis(fisher, bias) => F::Kurtosis(fisher, bias),
837
        #[cfg(feature = "dtype-array")]
838
        IF::Reshape(dims) => F::Reshape(dims),
839
        #[cfg(feature = "repeat_by")]
840
        IF::RepeatBy => F::RepeatBy,
841
        IF::ArgUnique => F::ArgUnique,
842
        IF::ArgMin => F::ArgMin,
843
        IF::ArgMax => F::ArgMax,
844
        IF::ArgSort {
845
            descending,
846
            nulls_last,
847
        } => F::ArgSort {
848
            descending,
849
            nulls_last,
850
        },
851
        IF::Product => F::Product,
852
        #[cfg(feature = "rank")]
853
        IF::Rank { options, seed } => F::Rank { options, seed },
854
        IF::Repeat => F::Repeat,
855
        #[cfg(feature = "round_series")]
856
        IF::Clip { has_min, has_max } => F::Clip { has_min, has_max },
857
        #[cfg(feature = "dtype-struct")]
858
        IF::AsStruct => F::AsStruct,
859
        #[cfg(feature = "top_k")]
860
        IF::TopK { descending } => F::TopK { descending },
861
        #[cfg(feature = "top_k")]
862
        IF::TopKBy { descending } => F::TopKBy { descending },
863
        #[cfg(feature = "cum_agg")]
864
        IF::CumCount { reverse } => F::CumCount { reverse },
865
        #[cfg(feature = "cum_agg")]
866
        IF::CumSum { reverse } => F::CumSum { reverse },
867
        #[cfg(feature = "cum_agg")]
868
        IF::CumProd { reverse } => F::CumProd { reverse },
869
        #[cfg(feature = "cum_agg")]
870
        IF::CumMin { reverse } => F::CumMin { reverse },
871
        #[cfg(feature = "cum_agg")]
872
        IF::CumMax { reverse } => F::CumMax { reverse },
873
        IF::Reverse => F::Reverse,
874
        #[cfg(feature = "dtype-struct")]
875
        IF::ValueCounts {
876
            sort,
877
            parallel,
878
            name,
879
            normalize,
880
        } => F::ValueCounts {
881
            sort,
882
            parallel,
883
            name,
884
            normalize,
885
        },
886
        #[cfg(feature = "unique_counts")]
887
        IF::UniqueCounts => F::UniqueCounts,
888
        #[cfg(feature = "approx_unique")]
889
        IF::ApproxNUnique => F::ApproxNUnique,
890
        IF::Coalesce => F::Coalesce,
891
        #[cfg(feature = "diff")]
892
        IF::Diff(nb) => F::Diff(nb),
893
        #[cfg(feature = "pct_change")]
894
        IF::PctChange => F::PctChange,
895
        #[cfg(feature = "interpolate")]
896
        IF::Interpolate(m) => F::Interpolate(m),
897
        #[cfg(feature = "interpolate_by")]
898
        IF::InterpolateBy => F::InterpolateBy,
899
        #[cfg(feature = "log")]
900
        IF::Entropy { base, normalize } => F::Entropy { base, normalize },
901
        #[cfg(feature = "log")]
902
        IF::Log => F::Log,
903
        #[cfg(feature = "log")]
904
        IF::Log1p => F::Log1p,
905
        #[cfg(feature = "log")]
906
        IF::Exp => F::Exp,
907
        IF::Unique(v) => F::Unique(v),
908
        #[cfg(feature = "round_series")]
909
        IF::Round { decimals, mode } => F::Round { decimals, mode },
910
        #[cfg(feature = "round_series")]
911
        IF::RoundSF { digits } => F::RoundSF { digits },
912
        #[cfg(feature = "round_series")]
913
        IF::Floor => F::Floor,
914
        #[cfg(feature = "round_series")]
915
        IF::Ceil => F::Ceil,
916
        IF::UpperBound => F::UpperBound,
917
        IF::LowerBound => F::LowerBound,
918
        #[cfg(feature = "fused")]
919
        IF::Fused(f) => {
920
            assert_eq!(input.len(), 3);
921
            let mut input = input.into_iter();
922
            let fst = input.next().unwrap();
923
            let snd = input.next().unwrap();
924
            let trd = input.next().unwrap();
925
            return match f {
926
                FusedOperator::MultiplyAdd => (fst * snd) + trd,
927
                FusedOperator::SubMultiply => fst - (snd * trd),
928
                FusedOperator::MultiplySub => (fst * snd) - trd,
929
            };
930
        },
931
        IF::ConcatExpr(v) => F::ConcatExpr(v),
932
        #[cfg(feature = "cov")]
933
        IF::Correlation { method } => {
934
            use {CorrelationMethod as C, IRCorrelationMethod as IC};
935
            F::Correlation {
936
                method: match method {
937
                    IC::Pearson => C::Pearson,
938
                    #[cfg(all(feature = "rank", feature = "propagate_nans"))]
939
                    IC::SpearmanRank(v) => C::SpearmanRank(v),
940
                    IC::Covariance(v) => C::Covariance(v),
941
                },
942
            }
943
        },
944
        #[cfg(feature = "peaks")]
945
        IF::PeakMin => F::PeakMin,
946
        #[cfg(feature = "peaks")]
947
        IF::PeakMax => F::PeakMax,
948
        #[cfg(feature = "cutqcut")]
949
        IF::Cut {
950
            breaks,
951
            labels,
952
            left_closed,
953
            include_breaks,
954
        } => F::Cut {
955
            breaks,
956
            labels,
957
            left_closed,
958
            include_breaks,
959
        },
960
        #[cfg(feature = "cutqcut")]
961
        IF::QCut {
962
            probs,
963
            labels,
964
            left_closed,
965
            allow_duplicates,
966
            include_breaks,
967
        } => F::QCut {
968
            probs,
969
            labels,
970
            left_closed,
971
            allow_duplicates,
972
            include_breaks,
973
        },
974
        #[cfg(feature = "rle")]
975
        IF::RLE => F::RLE,
976
        #[cfg(feature = "rle")]
977
        IF::RLEID => F::RLEID,
978
        IF::ToPhysical => F::ToPhysical,
979
        #[cfg(feature = "random")]
980
        IF::Random { method, seed } => {
981
            use {IRRandomMethod as IR, RandomMethod as R};
982
            F::Random {
983
                method: match method {
984
                    IR::Shuffle => R::Shuffle,
985
                    IR::Sample {
986
                        is_fraction,
987
                        with_replacement,
988
                        shuffle,
989
                    } => R::Sample {
990
                        is_fraction,
991
                        with_replacement,
992
                        shuffle,
993
                    },
994
                },
995
                seed,
996
            }
997
        },
998
        IF::SetSortedFlag(s) => F::SetSortedFlag(s),
999
        #[cfg(feature = "ffi_plugin")]
1000
        IF::FfiPlugin {
1001
            flags,
1002
            lib,
1003
            symbol,
1004
            kwargs,
1005
        } => F::FfiPlugin {
1006
            flags,
1007
            lib,
1008
            symbol,
1009
            kwargs,
1010
        },
1011

1012
        IF::FoldHorizontal {
1013
            callback,
1014
            returns_scalar,
1015
            return_dtype,
1016
        } => F::FoldHorizontal {
1017
            callback,
1018
            returns_scalar,
1019
            return_dtype: return_dtype.map(DataTypeExpr::Literal),
1020
        },
1021
        IF::ReduceHorizontal {
1022
            callback,
1023
            returns_scalar,
1024
            return_dtype,
1025
        } => F::ReduceHorizontal {
1026
            callback,
1027
            returns_scalar,
1028
            return_dtype: return_dtype.map(DataTypeExpr::Literal),
1029
        },
1030
        #[cfg(feature = "dtype-struct")]
1031
        IF::CumReduceHorizontal {
1032
            callback,
1033
            returns_scalar,
1034
            return_dtype,
1035
        } => F::CumReduceHorizontal {
1036
            callback,
1037
            returns_scalar,
1038
            return_dtype: return_dtype.map(DataTypeExpr::Literal),
1039
        },
1040
        #[cfg(feature = "dtype-struct")]
1041
        IF::CumFoldHorizontal {
1042
            callback,
1043
            returns_scalar,
1044
            return_dtype,
1045
            include_init,
1046
        } => F::CumFoldHorizontal {
1047
            callback,
1048
            returns_scalar,
1049
            return_dtype: return_dtype.map(DataTypeExpr::Literal),
1050
            include_init,
1051
        },
1052

1053
        IF::MaxHorizontal => F::MaxHorizontal,
1054
        IF::MinHorizontal => F::MinHorizontal,
1055
        IF::SumHorizontal { ignore_nulls } => F::SumHorizontal { ignore_nulls },
1056
        IF::MeanHorizontal { ignore_nulls } => F::MeanHorizontal { ignore_nulls },
1057
        #[cfg(feature = "ewma")]
1058
        IF::EwmMean { options } => F::EwmMean { options },
1059
        #[cfg(feature = "ewma_by")]
1060
        IF::EwmMeanBy { half_life } => F::EwmMeanBy { half_life },
1061
        #[cfg(feature = "ewma")]
1062
        IF::EwmStd { options } => F::EwmStd { options },
1063
        #[cfg(feature = "ewma")]
1064
        IF::EwmVar { options } => F::EwmVar { options },
1065
        #[cfg(feature = "replace")]
1066
        IF::Replace => F::Replace,
1067
        #[cfg(feature = "replace")]
1068
        IF::ReplaceStrict { return_dtype } => F::ReplaceStrict {
1069
            return_dtype: return_dtype.map(Into::into),
1070
        },
1071
        IF::GatherEvery { n, offset } => F::GatherEvery { n, offset },
1072
        #[cfg(feature = "reinterpret")]
1073
        IF::Reinterpret(v) => F::Reinterpret(v),
1074
        IF::ExtendConstant => F::ExtendConstant,
1075

1076
        IF::RowEncode(_, v) => F::RowEncode(v),
1077
        #[cfg(feature = "dtype-struct")]
1078
        IF::RowDecode(fs, v) => F::RowDecode(
1079
            fs.into_iter().map(|f| (f.name, f.dtype.into())).collect(),
1080
            v,
1081
        ),
1082
    };
1083

1084
    Expr::Function { input, function }
1085
}
1086

1087
Product

Resources

Company
