Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Apr 29, 2024
1 parent cb3c9ad commit c1286b1
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 9 deletions.
2 changes: 1 addition & 1 deletion crates/polars-lazy/src/frame/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1874,7 +1874,7 @@ impl JoinBuilder {
suffix: self.suffix,
slice: None,
join_nulls: self.join_nulls,
..Default::default()
coalesce: self.coalesce,
};

let lp = self
Expand Down
13 changes: 13 additions & 0 deletions crates/polars-ops/src/frame/join/args.rs
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,19 @@ impl Debug for JoinType {
}
}

impl JoinType {
pub fn is_asof(&self) -> bool {
#[cfg(feature = "asof_join")]
{
matches!(self, JoinType::AsOf(_))
}
#[cfg(not(feature = "asof_join"))]
{
false
}
}
}

#[derive(Copy, Clone, PartialEq, Eq, Default, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum JoinValidation {
Expand Down
8 changes: 3 additions & 5 deletions crates/polars-ops/src/frame/join/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -209,9 +209,7 @@ pub trait DataFrameJoinOps: IntoDf {
JoinType::Left => {
left_df._left_join_from_series(other, s_left, s_right, args, _verbose, None)
},
JoinType::Outer { .. } => {
left_df._outer_join_from_series(other, s_left, s_right, args)
},
JoinType::Outer => left_df._outer_join_from_series(other, s_left, s_right, args),
#[cfg(feature = "semi_anti_join")]
JoinType::Anti => left_df._semi_anti_join_from_series(
s_left,
Expand Down Expand Up @@ -278,14 +276,14 @@ pub trait DataFrameJoinOps: IntoDf {
JoinType::Cross => {
unreachable!()
},
JoinType::Outer => {
JoinType::Outer => {
let names_left = selected_left.iter().map(|s| s.name()).collect::<Vec<_>>();
let coalesce = args.coalesce;
args.coalesce = JoinCoalesce::KeepColumns;
let suffix = args.suffix.clone();
let out = left_df._outer_join_from_series(other, &lhs_keys, &rhs_keys, args);

if matches!(coalesce, JoinCoalesce::CoalesceColumns) {
if coalesce.coalesce(&JoinType::Outer) {
Ok(_coalesce_outer_join(
out?,
&names_left,
Expand Down
9 changes: 6 additions & 3 deletions crates/polars-plan/src/logical_plan/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -314,7 +314,7 @@ pub(crate) fn det_join_schema(
arena.clear();
}
let coalesces_join_keys = options.args.coalesce.coalesce(&options.args.how);
// except in asof joins. Asof joins are not equi-joins
// Except in asof joins. Asof joins are not equi-joins
// so the columns that are joined on, may have different
// values so if the right has a different name, it is added to the schema
#[cfg(feature = "asof_join")]
Expand Down Expand Up @@ -343,9 +343,12 @@ pub(crate) fn det_join_schema(
join_on_right.insert(field.name);
}

let are_coalesced = options.args.coalesce.coalesce(&options.args.how);
let is_asof = options.args.how.is_asof();

// Asof joins are special, if the names are equal they will not be coalesced.
for (name, dtype) in schema_right.iter() {
if !coalesces_join_keys || // The names are not merged
!join_on_right.contains(name.as_str())
if !join_on_right.contains(name.as_str()) || (!are_coalesced && !is_asof)
// The names that are joined on are merged
{
if schema_left.contains(name.as_str()) {
Expand Down

0 comments on commit c1286b1

Please sign in to comment.