Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix(rust, python): pow return type evaluation #15506

Merged
merged 7 commits into from
Apr 6, 2024
Merged

fix(rust, python): pow return type evaluation #15506

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
972b17f
pow int float
CanglongCl Apr 6, 2024
500605b
avoid pow overflow in test
CanglongCl Apr 6, 2024
62f0a83
format
CanglongCl Apr 6, 2024
3e035bc
lint
CanglongCl Apr 6, 2024
a81fc25
fix typo
CanglongCl Apr 6, 2024
2dacf75
fix: features check
CanglongCl Apr 6, 2024
9dd04d9
update docs
CanglongCl Apr 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions crates/polars-core/src/utils/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,18 +319,35 @@ macro_rules! with_match_physical_integer_type {(
macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
use $crate::datatypes::DataType::*;
match $dtype {
#[cfg(feature = "dtype-i8")]
Int8 => __with_ty__! { i8 },
#[cfg(feature = "dtype-i16")]
Int16 => __with_ty__! { i16 },
Int32 => __with_ty__! { i32 },
Int64 => __with_ty__! { i64 },
#[cfg(feature = "dtype-u8")]
UInt8 => __with_ty__! { u8 },
#[cfg(feature = "dtype-u16")]
UInt16 => __with_ty__! { u16 },
UInt32 => __with_ty__! { u32 },
UInt64 => __with_ty__! { u64 },
dt => panic!("not implemented for dtype {:?}", dt),
}
})}

#[macro_export]
macro_rules! with_match_physical_float_type {(
$dtype:expr, | $_:tt $T:ident | $($body:tt)*
) => ({
macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
use $crate::datatypes::DataType::*;
match $dtype {
Float32 => __with_ty__! { f32 },
Float64 => __with_ty__! { f64 },
dt => panic!("not implemented for dtype {:?}", dt),
}
})}

#[macro_export]
macro_rules! with_match_physical_float_polars_type {(
$key_type:expr, | $_:tt $T:ident | $($body:tt)*
Expand Down
107 changes: 48 additions & 59 deletions crates/polars-plan/src/dsl/function_expr/pow.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use arrow::legacy::kernels::pow::pow as pow_kernel;
use num::pow::Pow;
use polars_core::export::num;
use polars_core::export::num::{Float, ToPrimitive};
use polars_core::with_match_physical_integer_type;

use super::*;

Expand Down Expand Up @@ -128,65 +129,53 @@ where

fn pow_on_series(base: &Series, exponent: &Series) -> PolarsResult<Option<Series>> {
use DataType::*;
match (base.dtype(), exponent.dtype()) {
#[cfg(feature = "dtype-u8")]
(UInt8, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.u8().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
#[cfg(feature = "dtype-i8")]
(Int8, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.i8().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
#[cfg(feature = "dtype-u16")]
(UInt16, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.u16().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
#[cfg(feature = "dtype-i16")]
(Int16, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.i16().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
(UInt32, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.u32().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
(Int32, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.i32().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
(UInt64, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.u64().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
(Int64, UInt8 | UInt16 | UInt32 | UInt64) => {
let ca = base.i64().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
},
(Float32, _) => {
let ca = base.f32().unwrap();
let exponent = exponent.strict_cast(&DataType::Float32)?;
pow_on_floats(ca, exponent.f32().unwrap())
},
(Float64, _) => {
let ca = base.f64().unwrap();
let exponent = exponent.strict_cast(&DataType::Float64)?;
pow_on_floats(ca, exponent.f64().unwrap())
},
_ => {
let base = base.cast(&DataType::Float64)?;
pow_on_series(&base, exponent)
},

let base_dtype = base.dtype();
polars_ensure!(
base_dtype.is_numeric(),
InvalidOperation: "`pow` operation not supported for dtype `{}` as base", base_dtype
);
let exponent_dtype = exponent.dtype();
polars_ensure!(
exponent_dtype.is_numeric(),
InvalidOperation: "`pow` operation not supported for dtype `{}` as exponent", exponent_dtype
);

// if false, dtype is float
if base_dtype.is_integer() {
with_match_physical_integer_type!(base_dtype, |$native_type| {
if exponent_dtype.is_float() {
match exponent_dtype {
Float32 => {
let ca = base.cast(&DataType::Float32)?;
pow_on_floats(ca.f32().unwrap(), exponent.f32().unwrap())
},
Float64 => {
let ca = base.cast(&DataType::Float64)?;
pow_on_floats(ca.f64().unwrap(), exponent.f64().unwrap())
},
_ => unreachable!(),
}
} else {
let ca = base.$native_type().unwrap();
let exponent = exponent.strict_cast(&DataType::UInt32)?;
pow_to_uint_dtype(ca, exponent.u32().unwrap())
}
})
} else {
match base_dtype {
Float32 => {
let ca = base.f32().unwrap();
let exponent = exponent.strict_cast(&DataType::Float32)?;
pow_on_floats(ca, exponent.f32().unwrap())
},
Float64 => {
let ca = base.f64().unwrap();
let exponent = exponent.strict_cast(&DataType::Float64)?;
pow_on_floats(ca, exponent.f64().unwrap())
},
_ => unreachable!(),
}
}
}

Expand Down
18 changes: 10 additions & 8 deletions crates/polars-plan/src/dsl/function_expr/schema.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -466,14 +466,16 @@ impl<'a> FieldsMapper<'a> {
}

pub(super) fn pow_dtype(&self) -> PolarsResult<Field> {
// base, exponent
match (self.fields[0].data_type(), self.fields[1].data_type()) {
(
base_dtype,
DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64,
) => Ok(Field::new(self.fields[0].name(), base_dtype.clone())),
(DataType::Float32, _) => Ok(Field::new(self.fields[0].name(), DataType::Float32)),
(_, _) => Ok(Field::new(self.fields[0].name(), DataType::Float64)),
let base_dtype = self.fields[0].data_type();
let exponent_dtype = self.fields[1].data_type();
if base_dtype.is_integer() {
if exponent_dtype.is_float() {
Ok(Field::new(self.fields[0].name(), exponent_dtype.clone()))
} else {
Ok(Field::new(self.fields[0].name(), base_dtype.clone()))
}
} else {
Ok(Field::new(self.fields[0].name(), base_dtype.clone()))
}
}

Expand Down
60 changes: 30 additions & 30 deletions py-polars/polars/dataframe/frame.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8192,16 +8192,16 @@ def with_columns(
... )
>>> df.with_columns((pl.col("a") ** 2).alias("a^2"))
shape: (4, 4)
┌─────┬──────┬───────┬──────┐
│ a ┆ b ┆ c ┆ a^2 │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ bool ┆ f64 │
╞═════╪══════╪═══════╪══════╡
│ 1 ┆ 0.5 ┆ true ┆ 1.0 │
│ 2 ┆ 4.0 ┆ true ┆ 4.0 │
│ 3 ┆ 10.0 ┆ false ┆ 9.0 │
│ 4 ┆ 13.0 ┆ true ┆ 16.0 │
└─────┴──────┴───────┴──────┘
┌─────┬──────┬───────┬─────┐
│ a ┆ b ┆ c ┆ a^2 │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ bool ┆ i64 │
╞═════╪══════╪═══════╪═════╡
│ 1 ┆ 0.5 ┆ true ┆ 1 │
│ 2 ┆ 4.0 ┆ true ┆ 4 │
│ 3 ┆ 10.0 ┆ false ┆ 9 │
│ 4 ┆ 13.0 ┆ true ┆ 16 │
└─────┴──────┴───────┴─────┘

Added columns will replace existing columns with the same name.

Expand All @@ -8228,16 +8228,16 @@ def with_columns(
... ]
... )
shape: (4, 6)
┌─────┬──────┬───────┬──────┬──────┬───────┐
│ a ┆ b ┆ c ┆ a^2 ┆ b/2 ┆ not c │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ bool ┆ f64 ┆ f64 ┆ bool │
╞═════╪══════╪═══════╪══════╪══════╪═══════╡
│ 1 ┆ 0.5 ┆ true ┆ 1.0 ┆ 0.25 ┆ false │
│ 2 ┆ 4.0 ┆ true ┆ 4.0 ┆ 2.0 ┆ false │
│ 3 ┆ 10.0 ┆ false ┆ 9.0 ┆ 5.0 ┆ true │
│ 4 ┆ 13.0 ┆ true ┆ 16.0 ┆ 6.5 ┆ false │
└─────┴──────┴───────┴──────┴──────┴───────┘
┌─────┬──────┬───────┬─────┬──────┬───────┐
│ a ┆ b ┆ c ┆ a^2 ┆ b/2 ┆ not c │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ bool ┆ i64 ┆ f64 ┆ bool │
╞═════╪══════╪═══════╪═════╪══════╪═══════╡
│ 1 ┆ 0.5 ┆ true ┆ 1 ┆ 0.25 ┆ false │
│ 2 ┆ 4.0 ┆ true ┆ 4 ┆ 2.0 ┆ false │
│ 3 ┆ 10.0 ┆ false ┆ 9 ┆ 5.0 ┆ true │
│ 4 ┆ 13.0 ┆ true ┆ 16 ┆ 6.5 ┆ false │
└─────┴──────┴───────┴─────┴──────┴───────┘

Multiple columns also can be added using positional arguments instead of a list.

Expand All @@ -8247,16 +8247,16 @@ def with_columns(
... (pl.col("c").not_()).alias("not c"),
... )
shape: (4, 6)
┌─────┬──────┬───────┬──────┬──────┬───────┐
│ a ┆ b ┆ c ┆ a^2 ┆ b/2 ┆ not c │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ bool ┆ f64 ┆ f64 ┆ bool │
╞═════╪══════╪═══════╪══════╪══════╪═══════╡
│ 1 ┆ 0.5 ┆ true ┆ 1.0 ┆ 0.25 ┆ false │
│ 2 ┆ 4.0 ┆ true ┆ 4.0 ┆ 2.0 ┆ false │
│ 3 ┆ 10.0 ┆ false ┆ 9.0 ┆ 5.0 ┆ true │
│ 4 ┆ 13.0 ┆ true ┆ 16.0 ┆ 6.5 ┆ false │
└─────┴──────┴───────┴──────┴──────┴───────┘
┌─────┬──────┬───────┬─────┬──────┬───────┐
│ a ┆ b ┆ c ┆ a^2 ┆ b/2 ┆ not c │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ bool ┆ i64 ┆ f64 ┆ bool │
╞═════╪══════╪═══════╪═════╪══════╪═══════╡
│ 1 ┆ 0.5 ┆ true ┆ 1 ┆ 0.25 ┆ false │
│ 2 ┆ 4.0 ┆ true ┆ 4 ┆ 2.0 ┆ false │
│ 3 ┆ 10.0 ┆ false ┆ 9 ┆ 5.0 ┆ true │
│ 4 ┆ 13.0 ┆ true ┆ 16 ┆ 6.5 ┆ false │
└─────┴──────┴───────┴─────┴──────┴───────┘

Use keyword arguments to easily name your expression inputs.

Expand Down
32 changes: 16 additions & 16 deletions py-polars/polars/expr/expr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5330,16 +5330,16 @@ def pow(self, exponent: IntoExprColumn | int | float) -> Self:
... pl.col("x").pow(pl.col("x").log(2)).alias("x ** xlog2"),
... )
shape: (4, 3)
┌─────┬───────┬────────────┐
│ x ┆ cube ┆ x ** xlog2 │
│ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ f64 │
╞═════╪═══════╪════════════╡
│ 1 ┆ 1.0 ┆ 1.0 │
│ 2 ┆ 8.0 ┆ 2.0 │
│ 4 ┆ 64.0 ┆ 16.0 │
│ 8 ┆ 512.0 ┆ 512.0 │
└─────┴───────┴────────────┘
┌─────┬──────┬────────────┐
│ x ┆ cube ┆ x ** xlog2 │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ f64 │
╞═════╪══════╪════════════╡
│ 1 ┆ 1 ┆ 1.0 │
│ 2 ┆ 8 ┆ 2.0 │
│ 4 ┆ 64 ┆ 16.0 │
│ 8 ┆ 512 ┆ 512.0 │
└─────┴──────┴────────────┘
"""
return self.__pow__(exponent)

Expand Down Expand Up @@ -9185,13 +9185,13 @@ def cumulative_eval(
┌────────┐
│ values │
│ --- │
│ f64 │
│ i64 │
╞════════╡
│ 0.0 │
│ -3.0 │
│ -8.0 │
│ -15.0 │
│ -24.0 │
│ 0 │
│ -3 │
│ -8 │
│ -15 │
│ -24 │
└────────┘
"""
return self._from_pyexpr(
Expand Down
6 changes: 3 additions & 3 deletions py-polars/polars/functions/lazy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2210,10 +2210,10 @@ def sql_expr(sql: str | Sequence[str]) -> Expr | list[Expr]:
┌─────┬─────┬───────┐
│ a ┆ a_a ┆ a_txt │
│ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ str │
│ i64 ┆ i64 ┆ str │
╞═════╪═════╪═══════╡
│ 2 ┆ 4.0 ┆ 2 │
│ 1 ┆ 1.0 ┆ 1 │
│ 2 ┆ 4 ┆ 2 │
│ 1 ┆ 1 ┆ 1 │
└─────┴─────┴───────┘
"""
if isinstance(sql, str):
Expand Down
Loading
Loading