Skip to content

Commit

Permalink
refactor: *_horizontal dependent on reduce_expr to expression archite…
Browse files Browse the repository at this point in the history
…cture
  • Loading branch information
reswqa committed Oct 12, 2023
1 parent 98ee5a3 commit c6dc59f
Show file tree
Hide file tree
Showing 8 changed files with 55 additions and 12 deletions.
1 change: 1 addition & 0 deletions crates/polars-ops/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ string_from_radix = ["polars-core/strings"]
extract_jsonpath = ["serde_json", "jsonpath_lib", "polars-json"]
log = []
hash = []
zip_with = ["polars-core/zip_with"]
group_by_list = ["polars-core/group_by_list"]
rolling_window = ["polars-core/rolling_window"]
moment = ["polars-core/moment"]
Expand Down
12 changes: 12 additions & 0 deletions crates/polars-ops/src/series/ops/horizontal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,3 +66,15 @@ pub fn all_horizontal(s: &[Series]) -> PolarsResult<Series> {
.with_name("all");
Ok(out.into_series())
}

#[cfg(feature = "zip_with")]
pub fn max_horizontal(s: &[Series]) -> PolarsResult<Option<Series>> {
let df = DataFrame::new_no_checks(Vec::from(s));
df.hmax().map(|opt_s| opt_s.map(|s| s.with_name("max")))
}

#[cfg(feature = "zip_with")]
pub fn min_horizontal(s: &[Series]) -> PolarsResult<Option<Series>> {
let df = DataFrame::new_no_checks(Vec::from(s));
df.hmin().map(|opt_s| opt_s.map(|s| s.with_name("min")))
}
2 changes: 1 addition & 1 deletion crates/polars-plan/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ polars-arrow = { workspace = true }
polars-core = { workspace = true, features = ["lazy", "zip_with", "random"], default-features = false }
polars-ffi = { workspace = true, optional = true }
polars-io = { workspace = true, features = ["lazy"], default-features = false }
polars-ops = { workspace = true, default-features = false }
polars-ops = { workspace = true, features = ["zip_with"], default-features = false }
polars-time = { workspace = true, optional = true }
polars-utils = { workspace = true }

Expand Down
8 changes: 8 additions & 0 deletions crates/polars-plan/src/dsl/function_expr/dispatch.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,3 +63,11 @@ pub(super) fn forward_fill(s: &Series, limit: FillNullLimit) -> PolarsResult<Ser
pub(super) fn sum_horizontal(s: &[Series]) -> PolarsResult<Series> {
polars_ops::prelude::sum_horizontal(s)
}

pub(super) fn max_horizontal(s: &mut [Series]) -> PolarsResult<Option<Series>> {
polars_ops::prelude::max_horizontal(s)
}

pub(super) fn min_horizontal(s: &mut [Series]) -> PolarsResult<Option<Series>> {
polars_ops::prelude::min_horizontal(s)
}
6 changes: 6 additions & 0 deletions crates/polars-plan/src/dsl/function_expr/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -257,6 +257,8 @@ pub enum FunctionExpr {
limit: FillNullLimit,
},
SumHorizontal,
MaxHorizontal,
MinHorizontal,
}

impl Hash for FunctionExpr {
Expand Down Expand Up @@ -429,6 +431,8 @@ impl Display for FunctionExpr {
BackwardFill { .. } => "backward_fill",
ForwardFill { .. } => "forward_fill",
SumHorizontal => "sum_horizontal",
MaxHorizontal => "max_horizontal",
MinHorizontal => "min_horizontal",
};
write!(f, "{s}")
}
Expand Down Expand Up @@ -749,6 +753,8 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
BackwardFill { limit } => map!(dispatch::backward_fill, limit),
ForwardFill { limit } => map!(dispatch::forward_fill, limit),
SumHorizontal => map_as_slice!(dispatch::sum_horizontal),
MaxHorizontal => wrap!(dispatch::max_horizontal),
MinHorizontal => wrap!(dispatch::min_horizontal),
}
}
}
Expand Down
2 changes: 2 additions & 0 deletions crates/polars-plan/src/dsl/function_expr/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,8 @@ impl FunctionExpr {
BackwardFill { .. } => mapper.with_same_dtype(),
ForwardFill { .. } => mapper.with_same_dtype(),
SumHorizontal => mapper.map_to_supertype(),
MaxHorizontal => mapper.map_to_supertype(),
MinHorizontal => mapper.map_to_supertype(),
}
}
}
Expand Down
34 changes: 24 additions & 10 deletions crates/polars-plan/src/dsl/functions/horizontal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -236,11 +236,18 @@ pub fn max_horizontal<E: AsRef<[Expr]>>(exprs: E) -> Expr {
if exprs.is_empty() {
return Expr::Columns(Vec::new());
}
let func = |s1, s2| {
let df = DataFrame::new_no_checks(vec![s1, s2]);
df.hmax()
};
reduce_exprs(func, exprs).alias("max")

Expr::Function {
input: exprs,
function: FunctionExpr::MaxHorizontal,
options: FunctionOptions {
collect_groups: ApplyOptions::ApplyFlat,
input_wildcard_expansion: true,
auto_explode: true,
allow_rename: true,
..Default::default()
},
}
}

/// Create a new column with the the minimum value per row.
Expand All @@ -251,11 +258,18 @@ pub fn min_horizontal<E: AsRef<[Expr]>>(exprs: E) -> Expr {
if exprs.is_empty() {
return Expr::Columns(Vec::new());
}
let func = |s1, s2| {
let df = DataFrame::new_no_checks(vec![s1, s2]);
df.hmin()
};
reduce_exprs(func, exprs).alias("min")

Expr::Function {
input: exprs,
function: FunctionExpr::MinHorizontal,
options: FunctionOptions {
collect_groups: ApplyOptions::ApplyFlat,
input_wildcard_expansion: true,
auto_explode: true,
allow_rename: true,
..Default::default()
},
}
}

/// Create a new column with the the sum of the values in each row.
Expand Down
2 changes: 1 addition & 1 deletion crates/polars/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ sort_multiple = ["polars-core/sort_multiple"]
# extra operations
approx_unique = ["polars-lazy?/approx_unique", "polars-ops/approx_unique"]
is_in = ["polars-lazy?/is_in"]
zip_with = ["polars-core/zip_with"]
zip_with = ["polars-core/zip_with", "polars-ops/zip_with"]
round_series = ["polars-core/round_series", "polars-lazy?/round_series", "polars-ops/round_series"]
checked_arithmetic = ["polars-core/checked_arithmetic"]
repeat_by = ["polars-ops/repeat_by", "polars-lazy?/repeat_by"]
Expand Down

0 comments on commit c6dc59f

Please sign in to comment.