Skip to content

Commit

Permalink
feat: DataFrame.join supports Series other (#1303)
Browse files Browse the repository at this point in the history
  • Loading branch information
TrevorBergeron authored Jan 21, 2025
1 parent 774e56b commit ee37a0a
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 2 deletions.
9 changes: 8 additions & 1 deletion bigframes/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -3033,8 +3033,15 @@ def merge(
return DataFrame(block)

def join(
self, other: DataFrame, *, on: Optional[str] = None, how: str = "left"
self,
other: Union[DataFrame, bigframes.series.Series],
*,
on: Optional[str] = None,
how: str = "left",
) -> DataFrame:
if isinstance(other, bigframes.series.Series):
other = other.to_frame()

left, right = self, other

if not left.columns.intersection(right.columns).empty:
Expand Down
21 changes: 21 additions & 0 deletions tests/system/small/test_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -2554,6 +2554,27 @@ def test_join_param_on(scalars_dfs, how):
assert_pandas_df_equal(bf_result, pd_result, ignore_order=True)


@all_joins
def test_df_join_series(scalars_dfs, how):
bf_df, pd_df = scalars_dfs

bf_df_a = bf_df[["string_col", "int64_col", "rowindex_2"]]
bf_df_a = bf_df_a.assign(rowindex_2=bf_df_a["rowindex_2"] + 2)
bf_series_b = bf_df["float64_col"]

if how == "cross":
with pytest.raises(ValueError):
bf_df_a.join(bf_series_b, on="rowindex_2", how=how)
else:
bf_result = bf_df_a.join(bf_series_b, on="rowindex_2", how=how).to_pandas()

pd_df_a = pd_df[["string_col", "int64_col", "rowindex_2"]]
pd_df_a = pd_df_a.assign(rowindex_2=pd_df_a["rowindex_2"] + 2)
pd_series_b = pd_df["float64_col"]
pd_result = pd_df_a.join(pd_series_b, on="rowindex_2", how=how)
assert_pandas_df_equal(bf_result, pd_result, ignore_order=True)


@pytest.mark.parametrize(
("by", "ascending", "na_position"),
[
Expand Down
2 changes: 1 addition & 1 deletion third_party/bigframes_vendored/pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -4384,7 +4384,7 @@ def join(self, other, *, on: Optional[str] = None, how: str) -> DataFrame:
Args:
other:
DataFrame with an Index similar to the Index of this one.
DataFrame or Series with an Index similar to the Index of this one.
on:
Column in the caller to join on the index in other, otherwise
joins index-on-index. Like an Excel VLOOKUP operation.
Expand Down

0 comments on commit ee37a0a

Please sign in to comment.