Skip to content

Commit

Permalink
feat(rust): implement is_between on the rust side
Browse files Browse the repository at this point in the history
  • Loading branch information
petrosbar committed Oct 24, 2023
1 parent 15ec106 commit 3e6ba17
Show file tree
Hide file tree
Showing 16 changed files with 105 additions and 17 deletions.
1 change: 1 addition & 0 deletions crates/polars-lazy/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ repeat_by = ["polars-plan/repeat_by"]
round_series = ["polars-plan/round_series", "polars-ops/round_series"]
is_first_distinct = ["polars-plan/is_first_distinct"]
is_last_distinct = ["polars-plan/is_last_distinct"]
is_between = ["polars-plan/is_between"]
is_unique = ["polars-plan/is_unique"]
cross_join = ["polars-plan/cross_join", "polars-pipe?/cross_join", "polars-ops/cross_join"]
asof_join = ["polars-plan/asof_join", "polars-time", "polars-ops/asof_join"]
Expand Down
1 change: 1 addition & 0 deletions crates/polars-ops/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ is_first_distinct = []
is_last_distinct = []
is_unique = []
unique_counts = []
is_between = []
approx_unique = []
fused = []
cutqcut = ["dtype-categorical", "dtype-struct"]
Expand Down
34 changes: 34 additions & 0 deletions crates/polars-ops/src/series/ops/is_between.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
use std::ops::BitAnd;

use polars_core::prelude::*;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum ClosedInterval {
#[default]
Both,
Left,
Right,
None,
}

pub fn is_between(
s: &Series,
lower: &Series,
upper: &Series,
closed: ClosedInterval,
) -> PolarsResult<BooleanChunked> {
let left_cmp_op = match closed {
ClosedInterval::None | ClosedInterval::Right => Series::gt,
ClosedInterval::Both | ClosedInterval::Left => Series::gt_eq,
};
let right_cmp_op = match closed {
ClosedInterval::None | ClosedInterval::Left => Series::lt,
ClosedInterval::Both | ClosedInterval::Right => Series::lt_eq,
};
let left = left_cmp_op(s, lower)?;
let right = right_cmp_op(s, upper)?;
Ok(left.bitand(right))
}
4 changes: 4 additions & 0 deletions crates/polars-ops/src/series/ops/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ mod fused;
mod horizontal;
#[cfg(feature = "convert_index")]
mod index;
#[cfg(feature = "is_between")]
mod is_between;
#[cfg(feature = "is_first_distinct")]
mod is_first_distinct;
#[cfg(feature = "is_in")]
Expand Down Expand Up @@ -68,6 +70,8 @@ pub use fused::*;
pub use horizontal::*;
#[cfg(feature = "convert_index")]
pub use index::*;
#[cfg(feature = "is_between")]
pub use is_between::*;
#[cfg(feature = "is_first_distinct")]
pub use is_first_distinct::*;
#[cfg(feature = "is_in")]
Expand Down
1 change: 1 addition & 0 deletions crates/polars-plan/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ round_series = ["polars-ops/round_series"]
is_first_distinct = ["polars-core/is_first_distinct", "polars-ops/is_first_distinct"]
is_last_distinct = ["polars-core/is_last_distinct", "polars-ops/is_last_distinct"]
is_unique = ["polars-ops/is_unique"]
is_between = ["polars-ops/is_between"]
cross_join = ["polars-ops/cross_join"]
asof_join = ["polars-core/asof_join", "polars-time", "polars-ops/asof_join"]
concat_str = []
Expand Down
16 changes: 16 additions & 0 deletions crates/polars-plan/src/dsl/function_expr/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@ pub enum BooleanFunction {
IsUnique,
#[cfg(feature = "is_unique")]
IsDuplicated,
#[cfg(feature = "is_between")]
IsBetween {
closed: ClosedInterval,
},
#[cfg(feature = "is_in")]
IsIn,
AllHorizontal,
Expand Down Expand Up @@ -64,6 +68,8 @@ impl Display for BooleanFunction {
IsUnique => "is_unique",
#[cfg(feature = "is_unique")]
IsDuplicated => "is_duplicated",
#[cfg(feature = "is_between")]
IsBetween { .. } => "is_between",
#[cfg(feature = "is_in")]
IsIn => "is_in",
AnyHorizontal => "any_horizontal",
Expand Down Expand Up @@ -94,6 +100,8 @@ impl From<BooleanFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
IsUnique => map!(is_unique),
#[cfg(feature = "is_unique")]
IsDuplicated => map!(is_duplicated),
#[cfg(feature = "is_between")]
IsBetween { closed } => map_as_slice!(is_between, closed),
#[cfg(feature = "is_in")]
IsIn => wrap!(is_in),
AllHorizontal => map_as_slice!(all_horizontal),
Expand Down Expand Up @@ -171,6 +179,14 @@ fn is_duplicated(s: &Series) -> PolarsResult<Series> {
polars_ops::prelude::is_duplicated(s).map(|ca| ca.into_series())
}

#[cfg(feature = "is_between")]
fn is_between(s: &[Series], closed: ClosedInterval) -> PolarsResult<Series> {
let ser = &s[0];
let lower = &s[1];
let upper = &s[2];
polars_ops::prelude::is_between(ser, lower, upper, closed).map(|ca| ca.into_series())
}

#[cfg(feature = "is_in")]
fn is_in(s: &mut [Series]) -> PolarsResult<Option<Series>> {
let left = &s[0];
Expand Down
11 changes: 11 additions & 0 deletions crates/polars-plan/src/dsl/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1017,6 +1017,17 @@ impl Expr {
self.apply_private(BooleanFunction::IsDuplicated.into())
}

#[allow(clippy::wrong_self_convention)]
#[cfg(feature = "is_between")]
pub fn is_between<E: Into<Expr>>(self, lower: E, upper: E, closed: ClosedInterval) -> Self {
self.map_many_private(
BooleanFunction::IsBetween { closed }.into(),
&[lower.into(), upper.into()],
true,
true,
)
}

/// Get a mask of unique values.
#[allow(clippy::wrong_self_convention)]
#[cfg(feature = "is_unique")]
Expand Down
2 changes: 2 additions & 0 deletions crates/polars/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@ repeat_by = ["polars-ops/repeat_by", "polars-lazy?/repeat_by"]
is_first_distinct = ["polars-lazy?/is_first_distinct", "polars-ops/is_first_distinct"]
is_last_distinct = ["polars-lazy?/is_last_distinct", "polars-ops/is_last_distinct"]
is_unique = ["polars-lazy?/is_unique", "polars-ops/is_unique"]
is_between = ["polars-lazy?/is_between", "polars-ops/is_between"]
asof_join = ["polars-core/asof_join", "polars-lazy?/asof_join", "polars-ops/asof_join"]
cross_join = ["polars-lazy?/cross_join", "polars-ops/cross_join"]
dot_product = ["polars-core/dot_product"]
Expand Down Expand Up @@ -308,6 +309,7 @@ docs-selection = [
"checked_arithmetic",
"ndarray",
"repeat_by",
"is_between",
"is_first_distinct",
"is_last_distinct",
"asof_join",
Expand Down
1 change: 1 addition & 0 deletions crates/polars/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,7 @@
//! - `repeat_by` - [Repeat element in an Array N times, where N is given by another array.
//! - `is_first_distinct` - Check if element is first unique value.
//! - `is_last_distinct` - Check if element is last unique value.
//! - `is_between` - Check if this expression is between the given lower and upper bounds.
//! - `checked_arithmetic` - checked arithmetic/ returning [`None`] on invalid operations.
//! - `dot_product` - Dot/inner product on [`Series`] and [`Expr`].
//! - `concat_str` - Concat string data in linear time.
Expand Down
1 change: 1 addition & 0 deletions py-polars/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ features = [
"is_first_distinct",
"is_last_distinct",
"is_unique",
"is_between",
"lazy",
"list_eval",
"list_to_struct",
Expand Down
1 change: 1 addition & 0 deletions py-polars/docs/source/reference/expressions/boolean.rst
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ Boolean
Expr.is_infinite
Expr.is_last
Expr.is_last_distinct
Expr.is_between
Expr.is_nan
Expr.is_not
Expr.is_not_nan
Expand Down
1 change: 1 addition & 0 deletions py-polars/docs/source/reference/series/descriptive.rst
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ Descriptive
Series.is_integer
Series.is_last
Series.is_last_distinct
Series.is_between
Series.is_nan
Series.is_not_nan
Series.is_not_null
Expand Down
22 changes: 6 additions & 16 deletions py-polars/polars/expr/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -5204,7 +5204,7 @@ def is_between(
closed: ClosedInterval = "both",
) -> Self:
"""
Check if this expression is between the given start and end values.
Check if this expression is between the given lower and upper bounds.
Parameters
----------
Expand Down Expand Up @@ -5281,22 +5281,12 @@ def is_between(
└─────┴────────────┘
"""
lower_bound = self._from_pyexpr(parse_as_expression(lower_bound))
upper_bound = self._from_pyexpr(parse_as_expression(upper_bound))
lower_bound = parse_as_expression(lower_bound)
upper_bound = parse_as_expression(upper_bound)

if closed == "none":
return (self > lower_bound) & (self < upper_bound)
elif closed == "both":
return (self >= lower_bound) & (self <= upper_bound)
elif closed == "right":
return (self > lower_bound) & (self <= upper_bound)
elif closed == "left":
return (self >= lower_bound) & (self < upper_bound)
else:
raise ValueError(
"`closed` must be one of {'left', 'right', 'both', 'none'},"
f" got {closed!r}"
)
return self._from_pyexpr(
self._pyexpr.is_between(lower_bound, upper_bound, closed)
)

def hash(
self,
Expand Down
2 changes: 1 addition & 1 deletion py-polars/polars/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -3894,7 +3894,7 @@ def is_between(
closed: ClosedInterval = "both",
) -> Series:
"""
Get a boolean mask of the values that fall between the given start/end values.
Get a boolean mask of the values that are between the given lower/upper bounds.
Parameters
----------
Expand Down
17 changes: 17 additions & 0 deletions py-polars/src/conversion.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1424,6 +1424,23 @@ impl FromPyObject<'_> for Wrap<SearchSortedSide> {
}
}

impl FromPyObject<'_> for Wrap<ClosedInterval> {
fn extract(ob: &PyAny) -> PyResult<Self> {
let parsed = match ob.extract::<&str>()? {
"both" => ClosedInterval::Both,
"left" => ClosedInterval::Left,
"right" => ClosedInterval::Right,
"none" => ClosedInterval::None,
v => {
return Err(PyValueError::new_err(format!(
"`closed` must be one of {{'both', 'left', 'right', 'none'}}, got {v}",
)))
},
};
Ok(Wrap(parsed))
}
}

impl FromPyObject<'_> for Wrap<WindowMapping> {
fn extract(ob: &PyAny) -> PyResult<Self> {
let parsed = match ob.extract::<&str>()? {
Expand Down
7 changes: 7 additions & 0 deletions py-polars/src/expr/general.rs
Original file line number Diff line number Diff line change
Expand Up @@ -397,6 +397,13 @@ impl PyExpr {
self.inner.clone().is_unique().into()
}

fn is_between(&self, lower: Self, upper: Self, closed: Wrap<ClosedInterval>) -> Self {
self.clone()
.inner
.is_between(lower.inner, upper.inner, closed.0)
.into()
}

fn approx_n_unique(&self) -> Self {
self.inner.clone().approx_n_unique().into()
}
Expand Down

0 comments on commit 3e6ba17

Please sign in to comment.