From 96cdc1c4afb52f31e656c352b4becc769ba349e7 Mon Sep 17 00:00:00 2001 From: Weijie Guo Date: Sun, 24 Sep 2023 00:52:46 +0800 Subject: [PATCH] refactor: remove default auto-explode for man_many_private --- crates/polars-plan/src/dsl/arithmetic.rs | 11 ++--------- crates/polars-plan/src/dsl/binary.rs | 3 +++ crates/polars-plan/src/dsl/dt.rs | 5 ++++- crates/polars-plan/src/dsl/list.rs | 15 ++++++++++++--- crates/polars-plan/src/dsl/mod.rs | 5 +++-- crates/polars-plan/src/dsl/string.rs | 24 +++++++++++++++++++----- 6 files changed, 43 insertions(+), 20 deletions(-) diff --git a/crates/polars-plan/src/dsl/arithmetic.rs b/crates/polars-plan/src/dsl/arithmetic.rs index 4e3924b79be0f..874d42d7192a6 100644 --- a/crates/polars-plan/src/dsl/arithmetic.rs +++ b/crates/polars-plan/src/dsl/arithmetic.rs @@ -55,11 +55,8 @@ impl Expr { FunctionExpr::Pow(PowFunction::Generic), &[exponent.into()], false, + false, ) - .with_function_options(|mut options| { - options.auto_explode = false; - options - }) } /// Compute the square root of the given expression @@ -117,11 +114,7 @@ impl Expr { /// Compute the inverse tangent of the given expression, with the angle expressed as the argument of a complex number #[cfg(feature = "trigonometry")] pub fn arctan2(self, x: Self) -> Self { - self.map_many_private(FunctionExpr::Atan2, &[x], false) - .with_function_options(|mut options| { - options.auto_explode = false; - options - }) + self.map_many_private(FunctionExpr::Atan2, &[x], false, false) } /// Compute the hyperbolic cosine of the given expression diff --git a/crates/polars-plan/src/dsl/binary.rs b/crates/polars-plan/src/dsl/binary.rs index 9a1d6f2fd81c4..5c8aab05c5796 100644 --- a/crates/polars-plan/src/dsl/binary.rs +++ b/crates/polars-plan/src/dsl/binary.rs @@ -10,6 +10,7 @@ impl BinaryNameSpace { FunctionExpr::BinaryExpr(BinaryFunction::Contains), &[pat], true, + true, ) } @@ -19,6 +20,7 @@ impl BinaryNameSpace { FunctionExpr::BinaryExpr(BinaryFunction::EndsWith), &[sub], true, + true, ) } @@ -28,6 +30,7 @@ impl BinaryNameSpace { FunctionExpr::BinaryExpr(BinaryFunction::StartsWith), &[sub], true, + true, ) } } diff --git a/crates/polars-plan/src/dsl/dt.rs b/crates/polars-plan/src/dsl/dt.rs index 25c37906014c7..873838686a8d2 100644 --- a/crates/polars-plan/src/dsl/dt.rs +++ b/crates/polars-plan/src/dsl/dt.rs @@ -219,6 +219,7 @@ impl DateLikeNameSpace { self.0.map_many_private( FunctionExpr::TemporalExpr(TemporalFunction::Truncate(offset)), &[every, ambiguous], + true, false, ) } @@ -265,7 +266,7 @@ impl DateLikeNameSpace { #[cfg(feature = "date_offset")] pub fn offset_by(self, by: Expr) -> Expr { self.0 - .map_many_private(FunctionExpr::DateOffset, &[by], false) + .map_many_private(FunctionExpr::DateOffset, &[by], true, false) } #[cfg(feature = "timezones")] @@ -273,6 +274,7 @@ impl DateLikeNameSpace { self.0.map_many_private( FunctionExpr::TemporalExpr(TemporalFunction::ReplaceTimeZone(time_zone)), &[ambiguous], + true, false, ) } @@ -281,6 +283,7 @@ impl DateLikeNameSpace { self.0.map_many_private( FunctionExpr::TemporalExpr(TemporalFunction::Combine(tu)), &[time], + true, false, ) } diff --git a/crates/polars-plan/src/dsl/list.rs b/crates/polars-plan/src/dsl/list.rs index 6d7b1c10b9c3a..b9e584b4a1916 100644 --- a/crates/polars-plan/src/dsl/list.rs +++ b/crates/polars-plan/src/dsl/list.rs @@ -84,8 +84,12 @@ impl ListNameSpace { /// Get items in every sublist by index. pub fn get(self, index: Expr) -> Expr { - self.0 - .map_many_private(FunctionExpr::ListExpr(ListFunction::Get), &[index], false) + self.0.map_many_private( + FunctionExpr::ListExpr(ListFunction::Get), + &[index], + true, + false, + ) } /// Get items in every sublist by multiple indexes. @@ -98,6 +102,7 @@ impl ListNameSpace { self.0.map_many_private( FunctionExpr::ListExpr(ListFunction::Take(null_on_oob)), &[index], + true, false, ) } @@ -119,6 +124,7 @@ impl ListNameSpace { self.0.map_many_private( FunctionExpr::ListExpr(ListFunction::Join), &[separator], + true, false, ) } @@ -169,6 +175,7 @@ impl ListNameSpace { self.0.map_many_private( FunctionExpr::ListExpr(ListFunction::Slice), &[offset, length], + true, false, ) } @@ -248,6 +255,7 @@ impl ListNameSpace { .map_many_private( FunctionExpr::ListExpr(ListFunction::Contains), &[other], + true, false, ) .with_function_options(|mut options| { @@ -264,6 +272,7 @@ impl ListNameSpace { .map_many_private( FunctionExpr::ListExpr(ListFunction::CountMatches), &[other], + true, false, ) .with_function_options(|mut options| { @@ -278,11 +287,11 @@ impl ListNameSpace { .map_many_private( FunctionExpr::ListExpr(ListFunction::SetOperation(set_operation)), &[other], + false, true, ) .with_function_options(|mut options| { options.input_wildcard_expansion = true; - options.auto_explode = false; options }) } diff --git a/crates/polars-plan/src/dsl/mod.rs b/crates/polars-plan/src/dsl/mod.rs index 955409be0baf2..1e46c101335ff 100644 --- a/crates/polars-plan/src/dsl/mod.rs +++ b/crates/polars-plan/src/dsl/mod.rs @@ -690,6 +690,7 @@ impl Expr { self, function_expr: FunctionExpr, arguments: &[Expr], + auto_explode: bool, cast_to_supertypes: bool, ) -> Self { let mut input = Vec::with_capacity(arguments.len() + 1); @@ -701,7 +702,7 @@ impl Expr { function: function_expr, options: FunctionOptions { collect_groups: ApplyOptions::ApplyFlat, - auto_explode: true, + auto_explode, cast_to_supertypes, ..Default::default() }, @@ -1053,7 +1054,7 @@ impl Expr { let arguments = &[other]; // we don't have to apply on groups, so this is faster if has_literal { - self.map_many_private(BooleanFunction::IsIn.into(), arguments, true) + self.map_many_private(BooleanFunction::IsIn.into(), arguments, true, true) } else { self.apply_many_private(BooleanFunction::IsIn.into(), arguments, true, true) } diff --git a/crates/polars-plan/src/dsl/string.rs b/crates/polars-plan/src/dsl/string.rs index 5145e86ebd00d..cb321fcfe5926 100644 --- a/crates/polars-plan/src/dsl/string.rs +++ b/crates/polars-plan/src/dsl/string.rs @@ -14,6 +14,7 @@ impl StringNameSpace { }), &[pat], true, + true, ) } @@ -28,6 +29,7 @@ impl StringNameSpace { }), &[pat], true, + true, ) } @@ -37,6 +39,7 @@ impl StringNameSpace { FunctionExpr::StringExpr(StringFunction::EndsWith), &[sub], true, + true, ) } @@ -46,6 +49,7 @@ impl StringNameSpace { FunctionExpr::StringExpr(StringFunction::StartsWith), &[sub], true, + true, ) } @@ -119,13 +123,17 @@ impl StringNameSpace { /// Extract each successive non-overlapping match in an individual string as an array pub fn extract_all(self, pat: Expr) -> Expr { self.0 - .map_many_private(StringFunction::ExtractAll.into(), &[pat], false) + .map_many_private(StringFunction::ExtractAll.into(), &[pat], true, false) } /// Count all successive non-overlapping regex matches. pub fn count_matches(self, pat: Expr, literal: bool) -> Expr { - self.0 - .map_many_private(StringFunction::CountMatches(literal).into(), &[pat], false) + self.0.map_many_private( + StringFunction::CountMatches(literal).into(), + &[pat], + true, + false, + ) } /// Convert a Utf8 column into a Date/Datetime/Time column. @@ -134,6 +142,7 @@ impl StringNameSpace { self.0.map_many_private( StringFunction::Strptime(dtype, options).into(), &[ambiguous], + true, false, ) } @@ -205,13 +214,13 @@ impl StringNameSpace { /// Split the string by a substring. The resulting dtype is `List`. pub fn split(self, by: Expr) -> Expr { self.0 - .map_many_private(StringFunction::Split(false).into(), &[by], false) + .map_many_private(StringFunction::Split(false).into(), &[by], true, false) } /// Split the string by a substring and keep the substring. The resulting dtype is `List`. pub fn split_inclusive(self, by: Expr) -> Expr { self.0 - .map_many_private(StringFunction::Split(true).into(), &[by], false) + .map_many_private(StringFunction::Split(true).into(), &[by], true, false) } #[cfg(feature = "dtype-struct")] @@ -261,6 +270,7 @@ impl StringNameSpace { FunctionExpr::StringExpr(StringFunction::Replace { n: 1, literal }), &[pat, value], true, + true, ) } @@ -271,6 +281,7 @@ impl StringNameSpace { FunctionExpr::StringExpr(StringFunction::Replace { n, literal }), &[pat, value], true, + true, ) } @@ -281,6 +292,7 @@ impl StringNameSpace { FunctionExpr::StringExpr(StringFunction::Replace { n: -1, literal }), &[pat, value], true, + true, ) } @@ -313,6 +325,7 @@ impl StringNameSpace { self.0.map_many_private( FunctionExpr::StringExpr(StringFunction::StripPrefix), &[prefix], + true, false, ) } @@ -322,6 +335,7 @@ impl StringNameSpace { self.0.map_many_private( FunctionExpr::StringExpr(StringFunction::StripSuffix), &[suffix], + true, false, ) }