perf: inline read_pandas for small data (#383)

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #312729021 🦕

Author: GarrettWu

bigframes/dataframe.py

@@ -69,10 +69,6 @@
  import bigframes.session
 
 
-# BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type.
-# TODO(tbergeron): Convert to bytes-based limit
-MAX_INLINE_DF_SIZE = 5000
-
 LevelType = typing.Hashable
 LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]
 SingleItemValue = Union[bigframes.series.Series, int, float, Callable]
@@ -170,17 +166,7 @@ def __init__(
  columns=columns, # type:ignore
  dtype=dtype, # type:ignore
  )
- if (
- pd_dataframe.size < MAX_INLINE_DF_SIZE
- # TODO(swast): Workaround data types limitation in inline data.
- and not any(
- dt.pyarrow_dtype
- for dt in pd_dataframe.dtypes
- if isinstance(dt, pandas.ArrowDtype)
- )
- ):
- self._block = blocks.Block.from_local(pd_dataframe)
- elif session:
+ if session:
  self._block = session.read_pandas(pd_dataframe)._get_block()
  else:
  self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block()

bigframes/operations/base.py

@@ -30,10 +30,6 @@
 import bigframes.session
 import third_party.bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing
 
-# BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type.
-# TODO(tbergeron): Convert to bytes-based limit
-MAX_INLINE_SERIES_SIZE = 5000
-
 
 class SeriesMethods:
  def __init__(
@@ -104,17 +100,7 @@ def __init__(
  if pd_series.name is None:
  # to_frame will set default numeric column label if unnamed, but we do not support int column label, so must rename
  pd_dataframe = pd_dataframe.set_axis(["unnamed_col"], axis=1)
- if (
- pd_dataframe.size < MAX_INLINE_SERIES_SIZE
- # TODO(swast): Workaround data types limitation in inline data.
- and not any(
- dt.pyarrow_dtype
- for dt in pd_dataframe.dtypes
- if isinstance(dt, pd.ArrowDtype)
- )
- ):
- block = blocks.Block.from_local(pd_dataframe)
- elif session:
+ if session:
  block = session.read_pandas(pd_dataframe)._get_block()
  else:
  # Uses default global session

bigframes/session/__init__.py

@@ -108,6 +108,10 @@
  "UTF-32LE",
 }
 
+# BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type.
+# TODO(tbergeron): Convert to bytes-based limit
+MAX_INLINE_DF_SIZE = 5000
+
 logger = logging.getLogger(__name__)
 
 
@@ -882,6 +886,29 @@ def read_pandas(self, pandas_dataframe: pandas.DataFrame) -> dataframe.DataFrame
 
  def _read_pandas(
  self, pandas_dataframe: pandas.DataFrame, api_name: str
+ ) -> dataframe.DataFrame:
+ if (
+ pandas_dataframe.size < MAX_INLINE_DF_SIZE
+ # TODO(swast): Workaround data types limitation in inline data.
+ and not any(
+ (
+ isinstance(s.dtype, pandas.ArrowDtype)
+ or (len(s) > 0 and pandas.api.types.is_list_like(s.iloc[0]))
+ or pandas.api.types.is_datetime64_any_dtype(s)
+ )
+ for _, s in pandas_dataframe.items()
+ )
+ ):
+ return self._read_pandas_inline(pandas_dataframe)
+ return self._read_pandas_load_job(pandas_dataframe, api_name)
+
+ def _read_pandas_inline(
+ self, pandas_dataframe: pandas.DataFrame
+ ) -> dataframe.DataFrame:
+ return dataframe.DataFrame(blocks.Block.from_local(pandas_dataframe))
+
+ def _read_pandas_load_job(
+ self, pandas_dataframe: pandas.DataFrame, api_name: str
  ) -> dataframe.DataFrame:
  col_labels, idx_labels = (
  pandas_dataframe.columns.to_list(),
@@ -1079,7 +1106,7 @@ def read_csv(
  encoding=encoding,
  **kwargs,
  )
- return self.read_pandas(pandas_df) # type: ignore
+ return self._read_pandas(pandas_df, "read_csv") # type: ignore
 
  def read_pickle(
  self,
@@ -1096,7 +1123,7 @@ def read_pickle(
  if isinstance(pandas_obj, pandas.Series):
  if pandas_obj.name is None:
  pandas_obj.name = "0"
- bigframes_df = self.read_pandas(pandas_obj.to_frame())
+ bigframes_df = self._read_pandas(pandas_obj.to_frame(), "read_pickle")
  return bigframes_df[bigframes_df.columns[0]]
  return self._read_pandas(pandas_obj, "read_pickle")
 
@@ -1196,7 +1223,7 @@ def read_json(
  engine=engine,
  **kwargs,
  )
- return self.read_pandas(pandas_df)
+ return self._read_pandas(pandas_df, "read_json")
 
  def _check_file_size(self, filepath: str):
  max_size = 1024 * 1024 * 1024 # 1 GB in bytes

tests/system/small/test_progress_bar.py

@@ -15,10 +15,12 @@
 import re
 import tempfile
 
+import numpy as np
 import pandas as pd
 
 import bigframes as bf
 import bigframes.formatting_helpers as formatting_helpers
+from bigframes.session import MAX_INLINE_DF_SIZE
 
 job_load_message_regex = r"\w+ job [\w-]+ is \w+\."
 
@@ -66,10 +68,15 @@ def test_progress_bar_extract_jobs(
 def test_progress_bar_load_jobs(
  session: bf.Session, penguins_pandas_df_default_index: pd.DataFrame, capsys
 ):
+ # repeat the DF to be big enough to trigger the load job.
+ df = penguins_pandas_df_default_index
+ while len(df) < MAX_INLINE_DF_SIZE:
+ df = pd.DataFrame(np.repeat(df.values, 2, axis=0))
+
  bf.options.display.progress_bar = "terminal"
  with tempfile.TemporaryDirectory() as dir:
  path = dir + "/test_read_csv_progress_bar*.csv"
- penguins_pandas_df_default_index.to_csv(path, index=False)
+ df.to_csv(path, index=False)
  capsys.readouterr() # clear output
  session.read_csv(path)
 

tests/unit/session/test_io_bigquery.py

@@ -23,6 +23,7 @@
 from bigframes.core import log_adapter
 import bigframes.pandas as bpd
 import bigframes.session._io.bigquery as io_bq
+from tests.unit import resources
 
 
 def test_create_job_configs_labels_is_none():
@@ -64,7 +65,9 @@ def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit():
  "bigframes-api": "read_pandas",
  "source": "bigquery-dataframes-temp",
  }
- df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+ df = bpd.DataFrame(
+ {"col1": [1, 2], "col2": [3, 4]}, session=resources.create_bigquery_session()
+ )
  # Test running two methods
  df.head()
  df.max()
@@ -81,15 +84,16 @@ def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit():
  "recent-bigframes-api-2": "dataframe-__init__",
  "recent-bigframes-api-3": "dataframe-head",
  "recent-bigframes-api-4": "dataframe-__init__",
+ "recent-bigframes-api-5": "dataframe-__init__",
  }
- assert labels is not None
- assert len(labels) == 7
  assert labels == expected_dict
 
 
 def test_create_job_configs_labels_length_limit_met_and_labels_is_none():
  log_adapter.get_and_reset_api_methods()
- df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+ df = bpd.DataFrame(
+ {"col1": [1, 2], "col2": [3, 4]}, session=resources.create_bigquery_session()
+ )
  # Test running methods more than the labels' length limit
  for i in range(66):
  df.head()
@@ -114,7 +118,9 @@ def test_create_job_configs_labels_length_limit_met():
  value = f"test{i}"
  cur_labels[key] = value
  # If cur_labels length is 62, we can only add one label from api_methods
- df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+ df = bpd.DataFrame(
+ {"col1": [1, 2], "col2": [3, 4]}, session=resources.create_bigquery_session()
+ )
  # Test running two methods
  df.head()
  df.max()