fix: read_csv with both index_col and use_cols inconsistent with pandas

Author: chelsea-lin

bigframes/session/__init__.py

@@ -1165,7 +1165,11 @@ def _read_csv_w_bigquery_engine(
 
         table_id = self._loader.load_file(filepath_or_buffer, job_config=job_config)
         df = self._loader.read_gbq_table(
-            table_id, index_col=index_col, columns=columns, names=names
+            table_id,
+            index_col=index_col,
+            columns=columns,
+            names=names,
+            is_index_in_columns=True,
         )
 
         if dtype is not None:

bigframes/session/loader.py

@@ -93,7 +93,9 @@ def _to_index_cols(
     return index_cols
 
 
-def _check_column_duplicates(index_cols: Iterable[str], columns: Iterable[str]):
+def _check_column_duplicates(
+    index_cols: Iterable[str], columns: Iterable[str], is_index_in_columns: bool
+) -> List[str]:
     index_cols_list = list(index_cols) if index_cols is not None else []
     columns_list = list(columns) if columns is not None else []
     set_index = set(index_cols_list)
@@ -105,17 +107,29 @@ def _check_column_duplicates(index_cols: Iterable[str], columns: Iterable[str]):
             "All column names specified in 'index_col' must be unique."
         )
 
+    if columns is None:
+        return columns_list
+
     if len(columns_list) > len(set_columns):
         raise ValueError(
             "The 'columns' argument contains duplicate names. "
             "All column names specified in 'columns' must be unique."
         )
 
-    if not set_index.isdisjoint(set_columns):
-        raise ValueError(
-            "Found column names that exist in both 'index_col' and 'columns' arguments. "
-            "These arguments must specify distinct sets of columns."
-        )
+    if is_index_in_columns:
+        if not set_index.issubset(set_columns):
+            raise ValueError(
+                f"The specified index column(s) were not found: {set_index - set_columns}. "
+                f"Available columns are: {set_columns}"
+            )
+        return list(set_columns - set_index)
+    else:
+        if not set_index.isdisjoint(set_columns):
+            raise ValueError(
+                "Found column names that exist in both 'index_col' and 'columns' arguments. "
+                "These arguments must specify distinct sets of columns."
+            )
+        return columns_list
 
 
 @dataclasses.dataclass
@@ -387,6 +401,7 @@ def read_gbq_table(  # type: ignore[overload-overlap]
         enable_snapshot: bool = ...,
         dry_run: Literal[False] = ...,
         force_total_order: Optional[bool] = ...,
+        is_index_in_columns: bool = False,
     ) -> dataframe.DataFrame:
         ...
 
@@ -408,6 +423,7 @@ def read_gbq_table(
         enable_snapshot: bool = ...,
         dry_run: Literal[True] = ...,
         force_total_order: Optional[bool] = ...,
+        is_index_in_columns: bool = False,
     ) -> pandas.Series:
         ...
 
@@ -428,6 +444,7 @@ def read_gbq_table(
         enable_snapshot: bool = True,
         dry_run: bool = False,
         force_total_order: Optional[bool] = None,
+        is_index_in_columns: bool = False,
     ) -> dataframe.DataFrame | pandas.Series:
         import bigframes._tools.strings
         import bigframes.dataframe as dataframe
@@ -510,7 +527,7 @@ def read_gbq_table(
             index_col=index_col,
             names=names,
         )
-        _check_column_duplicates(index_cols, columns)
+        columns = _check_column_duplicates(index_cols, columns, is_index_in_columns)
 
         for key in index_cols:
             if key not in table_column_names:
@@ -790,7 +807,9 @@ def read_gbq_query(
             )
 
         index_cols = _to_index_cols(index_col)
-        _check_column_duplicates(index_cols, columns)
+        columns = _check_column_duplicates(
+            index_cols, columns, is_index_in_columns=False
+        )
 
         filters_copy1, filters_copy2 = itertools.tee(filters)
         has_filters = len(list(filters_copy1)) != 0

tests/system/small/test_session.py

@@ -1479,41 +1479,73 @@ def test_read_csv_for_gcs_file_w_header(session, df_and_gcs_csv, header):
 def test_read_csv_w_usecols(session, df_and_local_csv):
     # Compares results for pandas and bigframes engines
     scalars_df, path = df_and_local_csv
+    usecols = ["rowindex", "bool_col"]
     with open(path, "rb") as buffer:
         bf_df = session.read_csv(
             buffer,
             engine="bigquery",
-            usecols=["bool_col"],
+            usecols=usecols,
         )
     with open(path, "rb") as buffer:
         # Convert default pandas dtypes to match BigQuery DataFrames dtypes.
         pd_df = session.read_csv(
             buffer,
-            usecols=["bool_col"],
+            usecols=usecols,
             dtype=scalars_df[["bool_col"]].dtypes.to_dict(),
         )
 
-    # Cannot compare two dataframe due to b/408499371.
-    assert len(bf_df.columns) == 1
-    assert len(pd_df.columns) == 1
+    assert bf_df.shape == pd_df.shape
+    assert bf_df.columns.tolist() == pd_df.columns.tolist()
 
+    # BigFrames requires `sort_index()` because BigQuery doesn't preserve row IDs
+    # (b/280889935) or guarantee row ordering.
+    bf_df = bf_df.set_index("rowindex").sort_index()
+    pd_df = pd_df.set_index("rowindex")
+    pd.testing.assert_frame_equal(bf_df.to_pandas(), pd_df.to_pandas())
 
-@pytest.mark.parametrize(
-    "engine",
-    [
-        pytest.param("bigquery", id="bq_engine"),
-        pytest.param(None, id="default_engine"),
-    ],
-)
-def test_read_csv_local_w_usecols(session, scalars_pandas_df_index, engine):
-    with tempfile.TemporaryDirectory() as dir:
-        path = dir + "/test_read_csv_local_w_usecols.csv"
-        # Using the pandas to_csv method because the BQ one does not support local write.
-        scalars_pandas_df_index.to_csv(path, index=False)
 
-        # df should only have 1 column which is bool_col.
-        df = session.read_csv(path, usecols=["bool_col"], engine=engine)
-        assert len(df.columns) == 1
+def test_read_csv_w_usecols_and_indexcol(session, df_and_local_csv):
+    # Compares results for pandas and bigframes engines
+    scalars_df, path = df_and_local_csv
+    usecols = ["rowindex", "bool_col"]
+    with open(path, "rb") as buffer:
+        bf_df = session.read_csv(
+            buffer,
+            engine="bigquery",
+            usecols=usecols,
+            index_col="rowindex",
+        )
+    with open(path, "rb") as buffer:
+        # Convert default pandas dtypes to match BigQuery DataFrames dtypes.
+        pd_df = session.read_csv(
+            buffer,
+            usecols=usecols,
+            index_col="rowindex",
+            dtype=scalars_df[["bool_col"]].dtypes.to_dict(),
+        )
+
+    assert bf_df.shape == pd_df.shape
+    assert bf_df.columns.tolist() == pd_df.columns.tolist()
+
+    # BigFrames requires `sort_index()` because BigQuery doesn't preserve row IDs
+    # (b/280889935) or guarantee row ordering.
+    bf_df = bf_df.sort_index()
+    pd.testing.assert_frame_equal(bf_df.to_pandas(), pd_df.to_pandas())
+
+
+def test_read_csv_w_indexcol_not_in_usecols(session, df_and_local_csv):
+    _, path = df_and_local_csv
+    with open(path, "rb") as buffer:
+        with pytest.raises(
+            ValueError,
+            match=re.escape("The specified index column(s) were not found"),
+        ):
+            session.read_csv(
+                buffer,
+                engine="bigquery",
+                usecols=["bool_col"],
+                index_col="rowindex",
+            )
 
 
 @pytest.mark.parametrize(