fix: Fix single row broadcast with null index (#1803)

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -2895,7 +2895,7 @@ def join_with_single_row(
  combined_expr,
  index_columns=index_cols_post_join,
  column_labels=left.column_labels.append(single_row_block.column_labels),
- index_labels=[left.index.name],
+ index_labels=left.index.names,
  )
  return (
  block,

tests/system/small/test_dataframe.py

@@ -2657,16 +2657,16 @@ def test_listlike_binop_axis_1_bf_index(scalars_dfs):
  assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
 
 
-def test_binop_with_self_aggregate(session, scalars_dfs):
- scalars_df, scalars_pandas_df = scalars_dfs
+def test_binop_with_self_aggregate(scalars_dfs_maybe_ordered):
+ scalars_df, scalars_pandas_df = scalars_dfs_maybe_ordered
 
  df_columns = ["int64_col", "float64_col", "int64_too"]
 
  # Ensure that this takes the optimized single-query path by counting executions
- execution_count_before = session._metrics.execution_count
+ execution_count_before = scalars_df._session._metrics.execution_count
  bf_df = scalars_df[df_columns]
  bf_result = (bf_df - bf_df.mean()).to_pandas()
- execution_count_after = session._metrics.execution_count
+ execution_count_after = scalars_df._session._metrics.execution_count
 
  pd_df = scalars_pandas_df[df_columns]
  pd_result = pd_df - pd_df.mean()
@@ -2677,6 +2677,29 @@ def test_binop_with_self_aggregate(session, scalars_dfs):
  assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
 
 
+def test_binop_with_self_aggregate_w_index_reset(scalars_dfs_maybe_ordered):
+ scalars_df, scalars_pandas_df = scalars_dfs_maybe_ordered
+
+ df_columns = ["int64_col", "float64_col", "int64_too"]
+
+ # Ensure that this takes the optimized single-query path by counting executions
+ execution_count_before = scalars_df._session._metrics.execution_count
+ bf_df = scalars_df[df_columns].reset_index(drop=True)
+ bf_result = (bf_df - bf_df.mean()).to_pandas()
+ execution_count_after = scalars_df._session._metrics.execution_count
+
+ pd_df = scalars_pandas_df[df_columns].reset_index(drop=True)
+ pd_result = pd_df - pd_df.mean()
+
+ executions = execution_count_after - execution_count_before
+
+ assert executions == 1
+ pd_result.index = pd_result.index.astype("Int64")
+ assert_pandas_df_equal(
+ bf_result, pd_result, check_dtype=False, check_index_type=False
+ )
+
+
 @pytest.mark.parametrize(
  ("left_labels", "right_labels"),
  [