fix: Fix bug selecting column repeatedly (#1858)

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -330,12 +330,27 @@ def create_constant(
 
         return self.project_to_id(ex.const(value, dtype))
 
-    def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue:
+    def select_columns(
+        self, column_ids: typing.Sequence[str], allow_renames: bool = False
+    ) -> ArrayValue:
         # This basically just drops and reorders columns - logically a no-op except as a final step
-        selections = (
-            bigframes.core.nodes.AliasedRef.identity(ids.ColumnId(col_id))
-            for col_id in column_ids
-        )
+        selections = []
+        seen = set()
+
+        for id in column_ids:
+            if id not in seen:
+                ref = nodes.AliasedRef.identity(ids.ColumnId(id))
+            elif allow_renames:
+                ref = nodes.AliasedRef(
+                    ex.deref(id), ids.ColumnId(bigframes.core.guid.generate_guid())
+                )
+            else:
+                raise ValueError(
+                    "Must set allow_renames=True to select columns repeatedly"
+                )
+            selections.append(ref)
+            seen.add(id)
+
         return ArrayValue(
             nodes.SelectionNode(
                 child=self.node,

bigframes/core/blocks.py

@@ -1210,7 +1210,10 @@ def select_column(self, id: str) -> Block:
         return self.select_columns([id])
 
     def select_columns(self, ids: typing.Sequence[str]) -> Block:
-        expr = self._expr.select_columns([*self.index_columns, *ids])
+        # Allow renames as may end up selecting same columns multiple times
+        expr = self._expr.select_columns(
+            [*self.index_columns, *ids], allow_renames=True
+        )
         col_labels = self._get_labels_for_columns(ids)
         return Block(expr, self.index_columns, col_labels, self.index.names)
 

tests/system/small/test_dataframe.py

@@ -3408,6 +3408,15 @@ def test__dir__with_rename(scalars_dfs):
     assert "drop" in results
 
 
+def test_loc_select_columns_w_repeats(scalars_df_index, scalars_pandas_df_index):
+    bf_result = scalars_df_index[["int64_col", "int64_col", "int64_too"]].to_pandas()
+    pd_result = scalars_pandas_df_index[["int64_col", "int64_col", "int64_too"]]
+    pd.testing.assert_frame_equal(
+        bf_result,
+        pd_result,
+    )
+
+
 @pytest.mark.parametrize(
     ("start", "stop", "step"),
     [