feat: support names parameter in read_csv for bigquery engine (#1659)

Author: chelsea-lin

bigframes/core/utils.py

@@ -41,8 +41,10 @@ def get_axis_number(axis: typing.Union[str, int]) -> typing.Literal[0, 1]:
  raise ValueError(f"Not a valid axis: {axis}")
 
 
-def is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]:
- return pd.api.types.is_list_like(obj)
+def is_list_like(
+ obj: typing.Any, allow_sets: bool = True
+) -> typing_extensions.TypeGuard[typing.Sequence]:
+ return pd.api.types.is_list_like(obj, allow_sets=allow_sets)
 
 
 def is_dict_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Mapping]:

bigframes/session/__init__.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+from collections import abc
 import datetime
 import logging
 import os
@@ -569,7 +570,7 @@ def read_gbq_table(
  columns = col_order
 
  return self._loader.read_gbq_table(
- query=query,
+ table_id=query,
  index_col=index_col,
  columns=columns,
  max_results=max_results,
@@ -953,14 +954,21 @@ def _read_csv_w_bigquery_engine(
  native CSV loading capabilities, making it suitable for large datasets
  that may not fit into local memory.
  """
-
- if any(param is not None for param in (dtype, names)):
- not_supported = ("dtype", "names")
+ if dtype is not None:
  raise NotImplementedError(
- f"BigQuery engine does not support these arguments: {not_supported}. "
+ f"BigQuery engine does not support the `dtype` argument."
  f"{constants.FEEDBACK_LINK}"
  )
 
+ if names is not None:
+ if len(names) != len(set(names)):
+ raise ValueError("Duplicated names are not allowed.")
+ if not (
+ bigframes.core.utils.is_list_like(names, allow_sets=False)
+ or isinstance(names, abc.KeysView)
+ ):
+ raise ValueError("Names should be an ordered collection.")
+
  if index_col is True:
  raise ValueError("The value of index_col couldn't be 'True'")
 
@@ -1004,11 +1012,9 @@ def _read_csv_w_bigquery_engine(
  elif header > 0:
  job_config.skip_leading_rows = header + 1
 
- return self._loader.read_bigquery_load_job(
- filepath_or_buffer,
- job_config=job_config,
- index_col=index_col,
- columns=columns,
+ table_id = self._loader.load_file(filepath_or_buffer, job_config=job_config)
+ return self._loader.read_gbq_table(
+ table_id, index_col=index_col, columns=columns, names=names
  )
 
  def read_pickle(
@@ -1049,8 +1055,8 @@ def read_parquet(
  job_config = bigquery.LoadJobConfig()
  job_config.source_format = bigquery.SourceFormat.PARQUET
  job_config.labels = {"bigframes-api": "read_parquet"}
-
- return self._loader.read_bigquery_load_job(path, job_config=job_config)
+ table_id = self._loader.load_file(path, job_config=job_config)
+ return self._loader.read_gbq_table(table_id)
  else:
  if "*" in path:
  raise ValueError(
@@ -1121,10 +1127,8 @@ def read_json(
  job_config.encoding = encoding
  job_config.labels = {"bigframes-api": "read_json"}
 
- return self._loader.read_bigquery_load_job(
- path_or_buf,
- job_config=job_config,
- )
+ table_id = self._loader.load_file(path_or_buf, job_config=job_config)
+ return self._loader.read_gbq_table(table_id)
  else:
  if any(arg in kwargs for arg in ("chunksize", "iterator")):
  raise NotImplementedError(

bigframes/session/_io/bigquery/read_gbq_table.py

@@ -235,6 +235,8 @@ def get_index_cols(
  | Iterable[int]
  | int
  | bigframes.enums.DefaultIndexKind,
+ *,
+ names: Optional[Iterable[str]] = None,
 ) -> List[str]:
  """
  If we can get a total ordering from the table, such as via primary key
@@ -245,6 +247,14 @@ def get_index_cols(
  # Transform index_col -> index_cols so we have a variable that is
  # always a list of column names (possibly empty).
  schema_len = len(table.schema)
+
+ # If the `names` is provided, the index_col provided by the user is the new
+ # name, so we need to rename it to the original name in the table schema.
+ renamed_schema: Optional[Dict[str, str]] = None
+ if names is not None:
+ assert len(list(names)) == schema_len
+ renamed_schema = {name: field.name for name, field in zip(names, table.schema)}
+
  index_cols: List[str] = []
  if isinstance(index_col, bigframes.enums.DefaultIndexKind):
  if index_col == bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64:
@@ -261,6 +271,8 @@ def get_index_cols(
  f"Got unexpected index_col {repr(index_col)}. {constants.FEEDBACK_LINK}"
  )
  elif isinstance(index_col, str):
+ if renamed_schema is not None:
+ index_col = renamed_schema.get(index_col, index_col)
  index_cols = [index_col]
  elif isinstance(index_col, int):
  if not 0 <= index_col < schema_len:
@@ -272,6 +284,8 @@ def get_index_cols(
  elif isinstance(index_col, Iterable):
  for item in index_col:
  if isinstance(item, str):
+ if renamed_schema is not None:
+ item = renamed_schema.get(item, item)
  index_cols.append(item)
  elif isinstance(item, int):
  if not 0 <= item < schema_len:

bigframes/session/loader.py

@@ -348,14 +348,15 @@ def _start_generic_job(self, job: formatting_helpers.GenericJob):
 
  def read_gbq_table(
  self,
- query: str,
+ table_id: str,
  *,
  index_col: Iterable[str]
  | str
  | Iterable[int]
  | int
  | bigframes.enums.DefaultIndexKind = (),
  columns: Iterable[str] = (),
+ names: Optional[Iterable[str]] = None,
  max_results: Optional[int] = None,
  api_name: str = "read_gbq_table",
  use_cache: bool = True,
@@ -375,7 +376,7 @@ def read_gbq_table(
  )
 
  table_ref = google.cloud.bigquery.table.TableReference.from_string(
- query, default_project=self._bqclient.project
+ table_id, default_project=self._bqclient.project
  )
 
  columns = list(columns)
@@ -411,12 +412,37 @@ def read_gbq_table(
  f"Column '{key}' of `columns` not found in this table. Did you mean '{possibility}'?"
  )
 
+ # TODO(b/408499371): check `names` work with `use_cols` for read_csv method.
+ if names is not None:
+ len_names = len(list(names))
+ len_columns = len(table.schema)
+ if len_names > len_columns:
+ raise ValueError(
+ f"Too many columns specified: expected {len_columns}"
+ f" and found {len_names}"
+ )
+ elif len_names < len_columns:
+ if (
+ isinstance(index_col, bigframes.enums.DefaultIndexKind)
+ or index_col != ()
+ ):
+ raise KeyError(
+ "When providing both `index_col` and `names`, ensure the "
+ "number of `names` matches the number of columns in your "
+ "data."
+ )
+ index_col = range(len_columns - len_names)
+ names = [
+ field.name for field in table.schema[: len_columns - len_names]
+ ] + list(names)
+
  # Converting index_col into a list of column names requires
  # the table metadata because we might use the primary keys
  # when constructing the index.
  index_cols = bf_read_gbq_table.get_index_cols(
  table=table,
  index_col=index_col,
+ names=names,
  )
  _check_column_duplicates(index_cols, columns)
 
@@ -443,15 +469,15 @@ def read_gbq_table(
  # TODO(b/338419730): We don't need to fallback to a query for wildcard
  # tables if we allow some non-determinism when time travel isn't supported.
  if max_results is not None or bf_io_bigquery.is_table_with_wildcard_suffix(
- query
+ table_id
  ):
  # TODO(b/338111344): If we are running a query anyway, we might as
  # well generate ROW_NUMBER() at the same time.
  all_columns: Iterable[str] = (
  itertools.chain(index_cols, columns) if columns else ()
  )
  query = bf_io_bigquery.to_query(
- query,
+ table_id,
  columns=all_columns,
  sql_predicate=bf_io_bigquery.compile_filters(filters)
  if filters
@@ -561,6 +587,15 @@ def read_gbq_table(
  index_names = [None]
 
  value_columns = [col for col in array_value.column_ids if col not in index_cols]
+ if names is not None:
+ renamed_cols: Dict[str, str] = {
+ col: new_name for col, new_name in zip(array_value.column_ids, names)
+ }
+ index_names = [
+ renamed_cols.get(index_col, index_col) for index_col in index_cols
+ ]
+ value_columns = [renamed_cols.get(col, col) for col in value_columns]
+
  block = blocks.Block(
  array_value,
  index_columns=index_cols,
@@ -576,18 +611,12 @@ def read_gbq_table(
  df.sort_index()
  return df
 
- def read_bigquery_load_job(
+ def load_file(
  self,
  filepath_or_buffer: str | IO["bytes"],
  *,
  job_config: bigquery.LoadJobConfig,
- index_col: Iterable[str]
- | str
- | Iterable[int]
- | int
- | bigframes.enums.DefaultIndexKind = (),
- columns: Iterable[str] = (),
- ) -> dataframe.DataFrame:
+ ) -> str:
  # Need to create session table beforehand
  table = self._storage_manager.create_temp_table(_PLACEHOLDER_SCHEMA)
  # but, we just overwrite the placeholder schema immediately with the load job
@@ -615,16 +644,7 @@ def read_bigquery_load_job(
 
  self._start_generic_job(load_job)
  table_id = f"{table.project}.{table.dataset_id}.{table.table_id}"
-
- # The BigQuery REST API for tables.get doesn't take a session ID, so we
- # can't get the schema for a temp table that way.
-
- return self.read_gbq_table(
- query=table_id,
- index_col=index_col,
- columns=columns,
- api_name="read_gbq_table",
- )
+ return table_id
 
  def read_gbq_query(
  self,