feat: add level param to DataFrame.stack (#88)

* feat: add level param to DataFrame.stack

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -53,7 +53,7 @@ def equals(block1: blocks.Block, block2: blocks.Block) -> bool:
     joined_block = joined_block.select_columns(equality_ids).with_column_labels(
         list(range(len(equality_ids)))
     )
-    stacked_block = joined_block.stack(dropna=False, sort=False)
+    stacked_block = joined_block.stack()
     result = stacked_block.get_stat(stacked_block.value_columns[0], agg_ops.all_op)
     return typing.cast(bool, result)
 

bigframes/core/blocks.py

@@ -1284,20 +1284,20 @@ def pivot(
 
         return result_block.with_column_labels(column_index)
 
-    def stack(self, how="left", dropna=True, sort=True, levels: int = 1):
+    def stack(self, how="left", levels: int = 1):
         """Unpivot last column axis level into row axis"""
+        if levels == 0:
+            return self
+
         # These are the values that will be turned into rows
 
         col_labels, row_labels = utils.split_index(self.column_labels, levels=levels)
-        if dropna:
-            row_labels = row_labels.drop_duplicates()
-        if sort:
-            row_labels = row_labels.sort_values()
+        row_labels = row_labels.drop_duplicates()
 
         row_label_tuples = utils.index_as_tuples(row_labels)
 
         if col_labels is not None:
-            result_index = col_labels.drop_duplicates().sort_values().dropna(how="all")
+            result_index = col_labels.drop_duplicates().dropna(how="all")
             result_col_labels = utils.index_as_tuples(result_index)
         else:
             result_index = pd.Index([None])

bigframes/dataframe.py

@@ -1741,24 +1741,49 @@ def pivot(
         )
         return DataFrame(pivot_block)
 
-    def stack(self):
-        # TODO: support 'level' param by simply reordering levels such that selected level is last before passing to Block.stack.
-        # TODO: match impl to pandas future_stack as described in pandas 2.1 release notes
-        stack_block = self._block.stack()
-        result_block = block_ops.dropna(
-            stack_block, stack_block.value_columns, how="all"
-        )
+    def stack(self, level: LevelsType = -1):
         if not isinstance(self.columns, pandas.MultiIndex):
-            return bigframes.series.Series(result_block)
-        return DataFrame(result_block)
+            if level not in [0, -1, self.columns.name]:
+                raise IndexError(f"Invalid level {level} for single-level index")
+            return self._stack_mono()
+        return self._stack_multi(level)
+
+    def _stack_mono(self):
+        result_block = self._block.stack()
+        return bigframes.series.Series(result_block)
+
+    def _stack_multi(self, level: LevelsType = -1):
+        n_levels = self.columns.nlevels
+        if isinstance(level, int) or isinstance(level, str):
+            level = [level]
+        level_indices = []
+        for level_ref in level:
+            if isinstance(level_ref, int):
+                if level_ref < 0:
+                    level_indices.append(n_levels + level_ref)
+                else:
+                    level_indices.append(level_ref)
+            else:  # str
+                level_indices.append(self.columns.names.index(level_ref))
+
+        new_order = [
+            *[i for i in range(n_levels) if i not in level_indices],
+            *level_indices,
+        ]
+
+        original_columns = typing.cast(pandas.MultiIndex, self.columns)
+        new_columns = original_columns.reorder_levels(new_order)
+
+        block = self._block.with_column_labels(new_columns)
+
+        block = block.stack(levels=len(level))
+        return DataFrame(block)
 
     def unstack(self):
         block = self._block
         # Special case, unstack with mono-index transpose into a series
         if self.index.nlevels == 1:
-            block = block.stack(
-                how="right", dropna=False, sort=False, levels=self.columns.nlevels
-            )
+            block = block.stack(how="right", levels=self.columns.nlevels)
             return bigframes.series.Series(block)
 
         # Pivot by last level of index

tests/system/small/test_dataframe.py

@@ -1885,6 +1885,8 @@ def test_df_describe(scalars_dfs):
 
 
 def test_df_stack(scalars_dfs):
+    if pandas.__version__.startswith("1.") or pandas.__version__.startswith("2.0"):
+        pytest.skip("pandas <2.1 uses different stack implementation")
     scalars_df, scalars_pandas_df = scalars_dfs
     # To match bigquery dataframes
     scalars_pandas_df = scalars_pandas_df.copy()
@@ -1893,7 +1895,7 @@ def test_df_stack(scalars_dfs):
     columns = ["int64_col", "int64_too", "rowindex_2"]
 
     bf_result = scalars_df[columns].stack().to_pandas()
-    pd_result = scalars_pandas_df[columns].stack()
+    pd_result = scalars_pandas_df[columns].stack(future_stack=True)
 
     # Pandas produces NaN, where bq dataframes produces pd.NA
     pd.testing.assert_series_equal(bf_result, pd_result, check_dtype=False)

tests/system/small/test_multiindex.py

@@ -718,25 +718,37 @@ def test_column_multi_index_cumsum(scalars_df_index, scalars_pandas_df_index):
     pandas.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False)
 
 
-def test_column_multi_index_stack(scalars_df_index, scalars_pandas_df_index):
-    columns = ["int64_too", "int64_col", "rowindex_2"]
+@pytest.mark.parametrize(
+    ("level",),
+    [(["l3", "l1"],), ([-2, -1],), (["l3"],), ("l2",), (-3,)],
+)
+def test_column_multi_index_stack(level):
+    if pandas.__version__.startswith("1.") or pandas.__version__.startswith("2.0"):
+        pytest.skip("pandas <2.1 uses different stack implementation")
+
     level1 = pandas.Index(["b", "a", "b"])
-    # Need resulting column to be pyarrow string rather than object dtype
-    level2 = pandas.Index(["a", "b", "b"], dtype="string[pyarrow]")
-    multi_columns = pandas.MultiIndex.from_arrays([level1, level2])
-    bf_df = scalars_df_index[columns].copy()
-    bf_df.columns = multi_columns
-    pd_df = scalars_pandas_df_index[columns].copy()
-    pd_df.columns = multi_columns
+    level2 = pandas.Index(["a", "b", "b"])
+    level3 = pandas.Index(["b", "b", "a"])
 
-    bf_result = bf_df.stack().to_pandas()
-    # Shifting sort behavior in stack
-    pd_result = pd_df.stack()
+    multi_columns = pandas.MultiIndex.from_arrays(
+        [level1, level2, level3], names=["l1", "l2", "l3"]
+    )
+    pd_df = pandas.DataFrame(
+        [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+        index=[5, 2, None],
+        columns=multi_columns,
+        dtype="Int64",
+    )
+    bf_df = bpd.DataFrame(pd_df)
+
+    bf_result = bf_df.stack(level=level).to_pandas()
+    # BigFrames emulates future_stack impl
+    pd_result = pd_df.stack(level=level, future_stack=True)
 
     # Pandas produces NaN, where bq dataframes produces pd.NA
     # Column ordering seems to depend on pandas version
     pandas.testing.assert_frame_equal(
-        bf_result.sort_index(axis=1), pd_result.sort_index(axis=1), check_dtype=False
+        bf_result, pd_result, check_dtype=False, check_index_type=False
     )