feat: add info and memory_usage methods to dataframe (#219)

Author: TrevorBergeron

.pre-commit-config.yaml

@@ -38,4 +38,4 @@ repos:
  rev: v1.1.1
  hooks:
  - id: mypy
- additional_dependencies: [types-requests]
+ additional_dependencies: [types-requests, types-tabulate]

bigframes/_config/display_options.py

@@ -32,6 +32,10 @@ class DisplayOptions:
  progress_bar: Optional[str] = "auto"
  repr_mode: Literal["head", "deferred"] = "head"
 
+ max_info_columns: int = 100
+ max_info_rows: Optional[int] = 200000
+ memory_usage: bool = True
+
 
 @contextlib.contextmanager
 def pandas_repr(display_options: DisplayOptions):

bigframes/core/indexes/index.py

@@ -155,6 +155,14 @@ def _block(self) -> blocks.Block:
  def T(self) -> Index:
  return self.transpose()
 
+ def _memory_usage(self) -> int:
+ (n_rows,) = self.shape
+ return sum(
+ self.dtypes.map(
+ lambda dtype: bigframes.dtypes.DTYPE_BYTE_SIZES.get(dtype, 8) * n_rows
+ )
+ )
+
  def transpose(self) -> Index:
  return self
 
@@ -326,7 +334,10 @@ def _apply_aggregation(self, op: agg_ops.AggregateOp) -> typing.Any:
 
  def __getitem__(self, key: int) -> typing.Any:
  if isinstance(key, int):
- result_pd_df, _ = self._block.slice(key, key + 1, 1).to_pandas()
+ if key != -1:
+ result_pd_df, _ = self._block.slice(key, key + 1, 1).to_pandas()
+ else: # special case, want [-1:] instead of [-1:0]
+ result_pd_df, _ = self._block.slice(key).to_pandas()
  if result_pd_df.empty:
  raise IndexError("single positional indexer is out-of-bounds")
  return result_pd_df.index[0]

bigframes/dataframe.py

@@ -18,6 +18,7 @@
 
 import datetime
 import re
+import sys
 import textwrap
 import typing
 from typing import (
@@ -36,6 +37,7 @@
 import google.cloud.bigquery as bigquery
 import numpy
 import pandas
+import tabulate
 
 import bigframes
 import bigframes._config.display_options as display_options
@@ -350,6 +352,88 @@ def query_job(self) -> Optional[bigquery.QueryJob]:
  self._set_internal_query_job(self._compute_dry_run())
  return self._query_job
 
+ def memory_usage(self, index: bool = True):
+ n_rows, _ = self.shape
+ # like pandas, treat all variable-size objects as just 8-byte pointers, ignoring actual object
+ column_sizes = self.dtypes.map(
+ lambda dtype: bigframes.dtypes.DTYPE_BYTE_SIZES.get(dtype, 8) * n_rows
+ )
+ if index:
+ index_size = pandas.Series([self.index._memory_usage()], index=["Index"])
+ column_sizes = pandas.concat([index_size, column_sizes])
+ return column_sizes
+
+ def info(
+ self,
+ verbose: Optional[bool] = None,
+ buf=None,
+ max_cols: Optional[int] = None,
+ memory_usage: Optional[bool] = None,
+ show_counts: Optional[bool] = None,
+ ):
+ obuf = buf or sys.stdout
+
+ n_rows, n_columns = self.shape
+
+ max_cols = (
+ max_cols
+ if max_cols is not None
+ else bigframes.options.display.max_info_columns
+ )
+
+ show_all_columns = verbose if verbose is not None else (n_columns < max_cols)
+
+ obuf.write(f"{type(self)}\n")
+
+ index_type = "MultiIndex" if self.index.nlevels > 1 else "Index"
+
+ # These accessses are kind of expensive, maybe should try to skip?
+ first_indice = self.index[0]
+ last_indice = self.index[-1]
+ obuf.write(f"{index_type}: {n_rows} entries, {first_indice} to {last_indice}\n")
+
+ dtype_strings = self.dtypes.astype("string")
+ if show_all_columns:
+ obuf.write(f"Data columns (total {n_columns} columns):\n")
+ column_info = self.columns.to_frame(name="Column")
+
+ max_rows = bigframes.options.display.max_info_rows
+ too_many_rows = n_rows > max_rows if max_rows is not None else False
+
+ if show_counts if show_counts is not None else (not too_many_rows):
+ non_null_counts = self.count().to_pandas()
+ column_info["Non-Null Count"] = non_null_counts.map(
+ lambda x: f"{int(x)} non-null"
+ )
+
+ column_info["Dtype"] = dtype_strings
+
+ column_info = column_info.reset_index(drop=True)
+ column_info.index.name = "#"
+
+ column_info_formatted = tabulate.tabulate(column_info, headers="keys") # type: ignore
+ obuf.write(column_info_formatted)
+ obuf.write("\n")
+
+ else: # Just number of columns and first, last
+ obuf.write(
+ f"Columns: {n_columns} entries, {self.columns[0]} to {self.columns[-1]}\n"
+ )
+ dtype_counts = dtype_strings.value_counts().sort_index(ascending=True).items()
+ dtype_counts_formatted = ", ".join(
+ f"{dtype}({count})" for dtype, count in dtype_counts
+ )
+ obuf.write(f"dtypes: {dtype_counts_formatted}\n")
+
+ show_memory = (
+ memory_usage
+ if memory_usage is not None
+ else bigframes.options.display.memory_usage
+ )
+ if show_memory:
+ # TODO: Convert to different units (kb, mb, etc.)
+ obuf.write(f"memory usage: {self.memory_usage().sum()} bytes\n")
+
  def _set_internal_query_job(self, query_job: bigquery.QueryJob):
  self._query_job = query_job
 

bigframes/dtypes.py

@@ -143,6 +143,19 @@
 # "string" and "string[pyarrow] are accepted"
 BIGFRAMES_STRING_TO_BIGFRAMES["string[pyarrow]"] = pd.StringDtype(storage="pyarrow")
 
+# For the purposes of dataframe.memory_usage
+# https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types#data_type_sizes
+DTYPE_BYTE_SIZES = {
+ pd.BooleanDtype(): 1,
+ pd.Int64Dtype(): 8,
+ pd.Float32Dtype(): 8,
+ pd.StringDtype(): 8,
+ pd.ArrowDtype(pa.time64("us")): 8,
+ pd.ArrowDtype(pa.timestamp("us")): 8,
+ pd.ArrowDtype(pa.timestamp("us", tz="UTC")): 8,
+ pd.ArrowDtype(pa.date32()): 8,
+}
+
 
 def ibis_dtype_to_bigframes_dtype(
  ibis_dtype: ibis_dtypes.DataType,

noxfile.py

@@ -228,6 +228,7 @@ def mypy(session):
  "types-python-dateutil",
  "types-requests",
  "types-setuptools",
+ "types-tabulate",
  ]
  )
  | set(SYSTEM_TEST_STANDARD_DEPENDENCIES)

setup.py

@@ -50,6 +50,7 @@
  "requests >=2.27.1",
  "scikit-learn >=1.2.2",
  "sqlalchemy >=1.4,<3.0dev",
+ "tabulate >= 0.9",
  "ipywidgets >=7.7.1",
  "humanize >= 4.6.0",
 ]

tests/system/small/test_dataframe.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import io
 import operator
 import tempfile
 import typing
@@ -255,6 +256,47 @@ def test_drop_with_custom_column_labels(scalars_dfs):
  assert_pandas_df_equal(bf_result, pd_result)
 
 
+def test_df_memory_usage(scalars_dfs):
+ scalars_df, scalars_pandas_df = scalars_dfs
+
+ pd_result = scalars_pandas_df.memory_usage()
+ bf_result = scalars_df.memory_usage()
+
+ pd.testing.assert_series_equal(pd_result, bf_result, rtol=1.5)
+
+
+def test_df_info(scalars_dfs):
+ expected = (
+ "<class 'bigframes.dataframe.DataFrame'>\n"
+ "Index: 9 entries, 0 to 8\n"
+ "Data columns (total 13 columns):\n"
+ " # Column Non-Null Count Dtype\n"
+ "--- ------------- ---------------- ------------------------------\n"
+ " 0 bool_col 8 non-null boolean\n"
+ " 1 bytes_col 6 non-null object\n"
+ " 2 date_col 7 non-null date32[day][pyarrow]\n"
+ " 3 datetime_col 6 non-null timestamp[us][pyarrow]\n"
+ " 4 geography_col 4 non-null geometry\n"
+ " 5 int64_col 8 non-null Int64\n"
+ " 6 int64_too 9 non-null Int64\n"
+ " 7 numeric_col 6 non-null object\n"
+ " 8 float64_col 7 non-null Float64\n"
+ " 9 rowindex_2 9 non-null Int64\n"
+ " 10 string_col 8 non-null string\n"
+ " 11 time_col 6 non-null time64[us][pyarrow]\n"
+ " 12 timestamp_col 6 non-null timestamp[us, tz=UTC][pyarrow]\n"
+ "dtypes: Float64(1), Int64(3), boolean(1), date32[day][pyarrow](1), geometry(1), object(2), string(1), time64[us][pyarrow](1), timestamp[us, tz=UTC][pyarrow](1), timestamp[us][pyarrow](1)\n"
+ "memory usage: 945 bytes\n"
+ )
+
+ scalars_df, _ = scalars_dfs
+ bf_result = io.StringIO()
+
+ scalars_df.info(buf=bf_result)
+
+ assert expected == bf_result.getvalue()
+
+
 def test_drop_index(scalars_dfs):
  scalars_df, scalars_pandas_df = scalars_dfs
 

third_party/bigframes_vendored/pandas/core/config_init.py

@@ -33,6 +33,17 @@
  Instead estimated bytes processed will be shown. Dataframe and Series
  objects can still be computed with methods that explicitly execute and
  download results.
+ max_info_columns (int):
+ max_info_columns is used in DataFrame.info method to decide if
+ per column information will be printed.
+ max_info_rows (int or None):
+ df.info() will usually show null-counts for each column.
+ For large frames this can be quite slow. max_info_rows and max_info_cols
+ limit this null check only to frames with smaller dimensions than
+ specified.
+ memory_usage (bool):
+ This specifies if the memory usage of a DataFrame should be displayed when
+ df.info() is called. Valid values True,False,
 """
 
 sampling_options_doc = """

third_party/bigframes_vendored/pandas/core/frame.py

@@ -92,6 +92,72 @@ def values(self) -> np.ndarray:
  """
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
 
+ def info(
+ self,
+ verbose: bool | None = None,
+ buf=None,
+ max_cols: int | None = None,
+ memory_usage: bool | None = None,
+ show_counts: bool | None = None,
+ ) -> None:
+ """
+ Print a concise summary of a DataFrame.
+
+ This method prints information about a DataFrame including
+ the index dtypeand columns, non-null values and memory usage.
+
+ Args:
+ verbose (bool, optional):
+ Whether to print the full summary. By default, the setting in
+ ``pandas.options.display.max_info_columns`` is followed.
+ buf (writable buffer, defaults to sys.stdout):
+ Where to send the output. By default, the output is printed to
+ sys.stdout. Pass a writable buffer if you need to further process
+ the output.
+ max_cols (int, optional):
+ When to switch from the verbose to the truncated output. If the
+ DataFrame has more than `max_cols` columns, the truncated output
+ is used. By default, the setting in
+ ``pandas.options.display.max_info_columns`` is used.
+ memory_usage (bool, optional):
+ Specifies whether total memory usage of the DataFrame
+ elements (including the index) should be displayed. By default,
+ this follows the ``pandas.options.display.memory_usage`` setting.
+ True always show memory usage. False never shows memory usage.
+ Memory estimation is made based in column dtype and number of rows
+ assuming values consume the same memory amount for corresponding dtypes.
+ show_counts (bool, optional):
+ Whether to show the non-null counts. By default, this is shown
+ only if the DataFrame is smaller than
+ ``pandas.options.display.max_info_rows`` and
+ ``pandas.options.display.max_info_columns``. A value of True always
+ shows the counts, and False never shows the counts.
+
+ Returns:
+ None: This method prints a summary of a DataFrame and returns None."""
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+ def memory_usage(self, index: bool = True):
+ """
+ Return the memory usage of each column in bytes.
+
+ The memory usage can optionally include the contribution of
+ the index and elements of `object` dtype.
+
+ This value is displayed in `DataFrame.info` by default. This can be
+ suppressed by setting ``pandas.options.display.memory_usage`` to False.
+
+ Args:
+ index (bool, default True):
+ Specifies whether to include the memory usage of the DataFrame's
+ index in returned Series. If ``index=True``, the memory usage of
+ the index is the first item in the output.
+
+ Returns:
+ Series: A Series whose index is the original column names and whose values is the memory usage of each column in bytes.
+ """
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
  # ----------------------------------------------------------------------
  # IO methods (to / from other formats)
  def to_numpy(self, dtype=None, copy=False, na_value=None, **kwargs) -> np.ndarray: