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.

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

chore(rust): Feature gate iejoin #18646

Merged
merged 3 commits into from
Sep 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
1 change: 1 addition & 0 deletions crates/polars-lazy/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,7 @@ is_between = ["polars-plan/is_between", "polars-expr/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", "polars-mem-engine/asof_join"]
iejoin = ["polars-plan/iejoin"]
business = ["polars-plan/business"]
concat_str = ["polars-plan/concat_str"]
range = ["polars-plan/range"]
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 @@ -117,6 +117,7 @@ pivot = ["polars-core/reinterpret", "polars-core/dtype-struct"]
cross_join = []
chunked_ids = []
asof_join = []
iejoin = []
semi_anti_join = []
array_any_all = ["dtype-array"]
array_count = ["dtype-array"]
Expand Down
3 changes: 3 additions & 0 deletions crates/polars-ops/src/frame/join/args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ impl JoinCoalesce {
},
#[cfg(feature = "asof_join")]
AsOf(_) => matches!(self, JoinSpecific | CoalesceColumns),
#[cfg(feature = "iejoin")]
IEJoin(_) => false,
Cross => false,
#[cfg(feature = "semi_anti_join")]
Expand Down Expand Up @@ -121,6 +122,7 @@ pub enum JoinType {
Semi,
#[cfg(feature = "semi_anti_join")]
Anti,
#[cfg(feature = "iejoin")]
IEJoin(IEJoinOptions),
}

Expand All @@ -140,6 +142,7 @@ impl Display for JoinType {
Full { .. } => "FULL",
#[cfg(feature = "asof_join")]
AsOf(_) => "ASOF",
#[cfg(feature = "iejoin")]
IEJoin(_) => "IEJOIN",
Cross => "CROSS",
#[cfg(feature = "semi_anti_join")]
Expand Down
5 changes: 5 additions & 0 deletions crates/polars-ops/src/frame/join/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ mod cross_join;
mod dispatch_left_right;
mod general;
mod hash_join;
#[cfg(feature = "iejoin")]
mod iejoin;
#[cfg(feature = "merge_sorted")]
mod merge_sorted;
Expand All @@ -29,6 +30,7 @@ use general::create_chunked_index_mapping;
pub use general::{_coalesce_full_join, _finish_join, _join_suffix_name};
pub use hash_join::*;
use hashbrown::hash_map::{Entry, RawEntryMut};
#[cfg(feature = "iejoin")]
pub use iejoin::{IEJoinOptions, InequalityOperator};
#[cfg(feature = "merge_sorted")]
pub use merge_sorted::_merge_sorted_dfs;
Expand Down Expand Up @@ -199,6 +201,7 @@ pub trait DataFrameJoinOps: IntoDf {
}
}

#[cfg(feature = "iejoin")]
if let JoinType::IEJoin(options) = args.how {
let func = if POOL.current_num_threads() > 1 && !left_df.is_empty() && !other.is_empty()
{
Expand Down Expand Up @@ -289,6 +292,7 @@ pub trait DataFrameJoinOps: IntoDf {
panic!("expected by arguments on both sides")
},
},
#[cfg(feature = "iejoin")]
JoinType::IEJoin(_) => {
unreachable!()
},
Expand Down Expand Up @@ -316,6 +320,7 @@ pub trait DataFrameJoinOps: IntoDf {
JoinType::AsOf(_) => polars_bail!(
ComputeError: "asof join not supported for join on multiple keys"
),
#[cfg(feature = "iejoin")]
JoinType::IEJoin(_) => {
unreachable!()
},
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 @@ -116,6 +116,7 @@ is_unique = ["polars-ops/is_unique"]
is_between = ["polars-ops/is_between"]
cross_join = ["polars-ops/cross_join"]
asof_join = ["polars-time", "polars-ops/asof_join"]
iejoin = ["polars-ops/iejoin"]
concat_str = []
business = ["polars-ops/business"]
range = []
Expand Down
22 changes: 14 additions & 8 deletions crates/polars-plan/src/plans/conversion/join.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
use arrow::legacy::error::PolarsResult;
use either::Either;
use polars_core::error::feature_gated;

use super::*;
use crate::dsl::Expr;
#[cfg(feature = "iejoin")]
use crate::plans::AExpr;

fn check_join_keys(keys: &[Expr]) -> PolarsResult<()> {
Expand All @@ -26,14 +28,16 @@ pub fn resolve_join(
ctxt: &mut DslConversionContext,
) -> PolarsResult<Node> {
if !predicates.is_empty() {
debug_assert!(left_on.is_empty() && right_on.is_empty());
return resolve_join_where(
input_left.unwrap_left(),
input_right.unwrap_left(),
predicates,
options,
ctxt,
);
feature_gated!("iejoin", {
debug_assert!(left_on.is_empty() && right_on.is_empty());
return resolve_join_where(
input_left.unwrap_left(),
input_right.unwrap_left(),
predicates,
options,
ctxt,
);
})
}

let owned = Arc::unwrap_or_clone;
Expand Down Expand Up @@ -119,6 +123,7 @@ pub fn resolve_join(
run_conversion(lp, ctxt, "join")
}

#[cfg(feature = "iejoin")]
impl From<InequalityOperator> for Operator {
fn from(value: InequalityOperator) -> Self {
match value {
Expand All @@ -130,6 +135,7 @@ impl From<InequalityOperator> for Operator {
}
}

#[cfg(feature = "iejoin")]
fn resolve_join_where(
input_left: Arc<DslPlan>,
input_right: Arc<DslPlan>,
Expand Down
1 change: 1 addition & 0 deletions crates/polars-python/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@ json = ["polars/serde", "serde_json", "polars/json", "polars-utils/serde"]
trigonometry = ["polars/trigonometry"]
sign = ["polars/sign"]
asof_join = ["polars/asof_join"]
iejoin = ["polars/iejoin"]
cross_join = ["polars/cross_join"]
pct_change = ["polars/pct_change"]
repeat_by = ["polars/repeat_by"]
Expand Down
1 change: 1 addition & 0 deletions crates/polars-python/src/lazyframe/visitor/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -481,6 +481,7 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<PyObject> {
JoinType::Cross => "cross",
JoinType::Semi => "leftsemi",
JoinType::Anti => "leftanti",
#[cfg(feature = "iejoin")]
JoinType::IEJoin(_) => return Err(PyNotImplementedError::new_err("IEJoin")),
},
options.args.join_nulls,
Expand Down
1 change: 1 addition & 0 deletions crates/polars/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ approx_unique = ["polars-lazy?/approx_unique", "polars-ops/approx_unique"]
arg_where = ["polars-lazy?/arg_where"]
array_any_all = ["polars-lazy?/array_any_all", "dtype-array"]
asof_join = ["polars-lazy?/asof_join", "polars-ops/asof_join"]
iejoin = ["polars-lazy?/iejoin"]
binary_encoding = ["polars-ops/binary_encoding", "polars-lazy?/binary_encoding", "polars-sql?/binary_encoding"]
business = ["polars-lazy?/business", "polars-ops/business"]
checked_arithmetic = ["polars-core/checked_arithmetic"]
Expand Down
2 changes: 1 addition & 1 deletion py-polars/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ crate-type = ["cdylib"]

[dependencies]
libc = { workspace = true }
polars-python = { workspace = true, features = ["pymethods"] }
polars-python = { workspace = true, features = ["pymethods", "iejoin"] }
pyo3 = { workspace = true, features = ["abi3-py38", "chrono", "extension-module", "multiple-pymethods"] }

[build-dependencies]
Expand Down