diff --git a/.kokoro/continuous/e2e.cfg b/.kokoro/continuous/e2e.cfg index 774b63313e..3dbd0b47f0 100644 --- a/.kokoro/continuous/e2e.cfg +++ b/.kokoro/continuous/e2e.cfg @@ -3,7 +3,7 @@ # Only run this nox session. env_vars: { key: "NOX_SESSION" - value: "unit_prerelease system_prerelease system_noextras e2e notebook" + value: "e2e doctest notebook unit_prerelease system_prerelease system_noextras" } env_vars: { diff --git a/.kokoro/presubmit/e2e.cfg b/.kokoro/presubmit/e2e.cfg index 774b63313e..3dbd0b47f0 100644 --- a/.kokoro/presubmit/e2e.cfg +++ b/.kokoro/presubmit/e2e.cfg @@ -3,7 +3,7 @@ # Only run this nox session. env_vars: { key: "NOX_SESSION" - value: "unit_prerelease system_prerelease system_noextras e2e notebook" + value: "e2e doctest notebook unit_prerelease system_prerelease system_noextras" } env_vars: { diff --git a/CHANGELOG.md b/CHANGELOG.md index 34ab012fd7..cad061ce05 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,42 @@ [1]: https://pypi.org/project/bigframes/#history +## [1.8.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v1.7.0...v1.8.0) (2024-05-31) + + +### Features + +* `merge` only generates a default index if both inputs already have an index ([#733](https://github.com/googleapis/python-bigquery-dataframes/issues/733)) ([25d049c](https://github.com/googleapis/python-bigquery-dataframes/commit/25d049c078693466905a19cc0954fafcac6c414c)) +* Add `+`, `-` as unary ops, `^` binary op ([#724](https://github.com/googleapis/python-bigquery-dataframes/issues/724)) ([968d825](https://github.com/googleapis/python-bigquery-dataframes/commit/968d8257edbfcb6d437c6203c7c0078ba782cfed)) +* Add `GroupBy.size()` to get number of rows in each group ([#479](https://github.com/googleapis/python-bigquery-dataframes/issues/479)) ([1fca588](https://github.com/googleapis/python-bigquery-dataframes/commit/1fca588e4398baa0dae61bdea0d3bff17e3971b5)) +* Add DataFrame `~` operator ([#721](https://github.com/googleapis/python-bigquery-dataframes/issues/721)) ([354abc1](https://github.com/googleapis/python-bigquery-dataframes/commit/354abc17b5bd55d70d47f893cfccd7cd0ac9794a)) +* Add GeminiText 1.5 Preview models ([#737](https://github.com/googleapis/python-bigquery-dataframes/issues/737)) ([56cbd3b](https://github.com/googleapis/python-bigquery-dataframes/commit/56cbd3b6f17c5ac22572e872b270ac7e3636675a)) +* Add slot_millis and add stats to session object ([#725](https://github.com/googleapis/python-bigquery-dataframes/issues/725)) ([72e9583](https://github.com/googleapis/python-bigquery-dataframes/commit/72e95834f8755760f3529d38f340703f3b971f0a)) +* Adds bigframes.bigquery.array_to_string to convert array elements to delimited strings ([#731](https://github.com/googleapis/python-bigquery-dataframes/issues/731)) ([f12c906](https://github.com/googleapis/python-bigquery-dataframes/commit/f12c90611adb4741069ec32840ebbf2aea83a9f3)) +* Allow functions decorated with `bpd.remote_function()` to execute locally ([#704](https://github.com/googleapis/python-bigquery-dataframes/issues/704)) ([d850da6](https://github.com/googleapis/python-bigquery-dataframes/commit/d850da6364b98c4e01120725e1e609ad8f6c1263)) +* Ensure `"bigframes-api"` label is always set on jobs, even if the API is unknown ([#722](https://github.com/googleapis/python-bigquery-dataframes/issues/722)) ([1832778](https://github.com/googleapis/python-bigquery-dataframes/commit/1832778cfc4f29fdab1b22380f03b192eb8aebb9)) +* Support `ml.SimpleImputer` in bigframes ([#708](https://github.com/googleapis/python-bigquery-dataframes/issues/708)) ([4c4415f](https://github.com/googleapis/python-bigquery-dataframes/commit/4c4415fb137e3baedc4b2d77ec146827b003557e)) +* Support type annotations to supply input and output types to `bpd.remote_function()` decorator ([#717](https://github.com/googleapis/python-bigquery-dataframes/issues/717)) ([4a12e3c](https://github.com/googleapis/python-bigquery-dataframes/commit/4a12e3c6d49d78fc2b51d783cc8de5d09e7c9995)) +* Support type annotations with `bpd.remote_function()` and `axis=1` (a preview feature) ([#730](https://github.com/googleapis/python-bigquery-dataframes/issues/730)) ([e5a2992](https://github.com/googleapis/python-bigquery-dataframes/commit/e5a299271e3bcf94c66fb6ef70393071c1b7dc69)) + + +### Bug Fixes + +* Correct index labels in multiple aggregations for DataFrameGroupBy ([#723](https://github.com/googleapis/python-bigquery-dataframes/issues/723)) ([6a78c89](https://github.com/googleapis/python-bigquery-dataframes/commit/6a78c89a3a766b747b03c8a739760db1c79f533f)) +* Fix Null index assign series to column ([#711](https://github.com/googleapis/python-bigquery-dataframes/issues/711)) ([ffb4b57](https://github.com/googleapis/python-bigquery-dataframes/commit/ffb4b5712a1a07c703ea88f66ba3f43dd2f98197)) +* Set `bpd.remote_function()`s `input_types` and `output_types` default to `None` to allow omitting them when type annotations are present ([#729](https://github.com/googleapis/python-bigquery-dataframes/issues/729)) ([0e25a3b](https://github.com/googleapis/python-bigquery-dataframes/commit/0e25a3b3ae704bf75b752c57f613e778af58bac3)) +* Warn and disable time travel for linked datasets ([#712](https://github.com/googleapis/python-bigquery-dataframes/issues/712)) ([085fa9d](https://github.com/googleapis/python-bigquery-dataframes/commit/085fa9d8fe1ea4cd02a3d25d443beaa697e10784)) + + +### Performance Improvements + +* Optimize dataframe-series alignment on axis=1 ([#732](https://github.com/googleapis/python-bigquery-dataframes/issues/732)) ([3d39221](https://github.com/googleapis/python-bigquery-dataframes/commit/3d39221526df82617a8560fd2ab7ea13bc3c03d9)) + + +### Documentation + +* Add examples to DataFrameGroupBy and SeriesGroupBy ([#701](https://github.com/googleapis/python-bigquery-dataframes/issues/701)) ([e7da0f0](https://github.com/googleapis/python-bigquery-dataframes/commit/e7da0f085eb9b9cec06e5de972f07d9c1d545ac7)) + ## [1.7.0](https://github.com/googleapis/python-bigquery-dataframes/compare/v1.6.0...v1.7.0) (2024-05-20) diff --git a/bigframes/_config/__init__.py b/bigframes/_config/__init__.py index bf33420e60..4729532e98 100644 --- a/bigframes/_config/__init__.py +++ b/bigframes/_config/__init__.py @@ -61,10 +61,12 @@ def _init_bigquery_thread_local(self): @property def bigquery(self) -> bigquery_options.BigQueryOptions: """Options to use with the BigQuery engine.""" - if self._local.bigquery_options is not None: + if ( + bigquery_options := getattr(self._local, "bigquery_options", None) + ) is not None: # The only way we can get here is if someone called # _init_bigquery_thread_local. - return self._local.bigquery_options + return bigquery_options return self._bigquery_options diff --git a/bigframes/bigquery/__init__.py b/bigframes/bigquery/__init__.py index 6c9c04dca7..5808aa28bf 100644 --- a/bigframes/bigquery/__init__.py +++ b/bigframes/bigquery/__init__.py @@ -57,8 +57,7 @@ def array_length(series: series.Series) -> series.Series: dtype: Int64 Args: - series (bigframes.series.Series): - A Series with array columns. + series (bigframes.series.Series): A Series with array columns. Returns: bigframes.series.Series: A Series of integer values indicating @@ -104,7 +103,7 @@ def array_agg( Args: obj (groupby.SeriesGroupBy | groupby.DataFrameGroupBy): - A GroupBy object to be applied the function. + A GroupBy object to be applied the function. Returns: bigframes.series.Series | bigframes.dataframe.DataFrame: A Series or @@ -119,3 +118,33 @@ def array_agg( raise ValueError( f"Unsupported type {type(obj)} to apply `array_agg` function. {constants.FEEDBACK_LINK}" ) + + +def array_to_string(series: series.Series, delimiter: str) -> series.Series: + """Converts array elements within a Series into delimited strings. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import bigframes.bigquery as bbq + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> s = bpd.Series([["H", "i", "!"], ["Hello", "World"], np.nan, [], ["Hi"]]) + >>> bbq.array_to_string(s, delimiter=", ") + 0 H, i, ! + 1 Hello, World + 2 + 3 + 4 Hi + dtype: string + + Args: + series (bigframes.series.Series): A Series containing arrays. + delimiter (str): The string used to separate array elements. + + Returns: + bigframes.series.Series: A Series containing delimited strings. + + """ + return series._apply_unary_op(ops.ArrayToStringOp(delimiter=delimiter)) diff --git a/bigframes/constants.py b/bigframes/constants.py index 4778eb9c9e..9591297956 100644 --- a/bigframes/constants.py +++ b/bigframes/constants.py @@ -12,19 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -import datetime - """Constants used across BigQuery DataFrames. This module should not depend on any others in the package. """ -FEEDBACK_LINK = ( - "Share your usecase with the BigQuery DataFrames team at the " - "https://bit.ly/bigframes-feedback survey." -) +import datetime + +import bigframes_vendored.constants -ABSTRACT_METHOD_ERROR_MESSAGE = f"Abstract method. You have likely encountered a bug. Please share this stacktrace and how you reached it with the BigQuery DataFrames team. {FEEDBACK_LINK}" +FEEDBACK_LINK = bigframes_vendored.constants.FEEDBACK_LINK +ABSTRACT_METHOD_ERROR_MESSAGE = ( + bigframes_vendored.constants.ABSTRACT_METHOD_ERROR_MESSAGE +) DEFAULT_EXPIRATION = datetime.timedelta(days=7) diff --git a/bigframes/core/__init__.py b/bigframes/core/__init__.py index 0a2936419f..133d271fed 100644 --- a/bigframes/core/__init__.py +++ b/bigframes/core/__init__.py @@ -14,12 +14,15 @@ from __future__ import annotations from dataclasses import dataclass +import datetime import functools import io import itertools import typing -from typing import Iterable, Sequence +from typing import Iterable, Optional, Sequence +import warnings +import google.cloud.bigquery import ibis.expr.types as ibis_types import pandas import pyarrow as pa @@ -92,6 +95,35 @@ def from_pyarrow(cls, arrow_table: pa.Table, session: Session): ) return cls(node) + @classmethod + def from_table( + cls, + table: google.cloud.bigquery.Table, + schema: schemata.ArraySchema, + session: Session, + *, + predicate: Optional[str] = None, + at_time: Optional[datetime.datetime] = None, + primary_key: Sequence[str] = (), + ): + if any(i.field_type == "JSON" for i in table.schema if i.name in schema.names): + warnings.warn( + "Interpreting JSON column(s) as StringDtype. This behavior may change in future versions.", + bigframes.exceptions.PreviewWarning, + ) + node = nodes.ReadTableNode( + project_id=table.reference.project, + dataset_id=table.reference.dataset_id, + table_id=table.reference.table_id, + physical_schema=tuple(table.schema), + total_order_cols=tuple(primary_key), + columns=schema, + at_time=at_time, + table_session=session, + sql_predicate=predicate, + ) + return cls(node) + @property def column_ids(self) -> typing.Sequence[str]: return self.schema.names @@ -184,7 +216,7 @@ def project_to_id(self, expression: ex.Expression, output_id: str): child=self.node, assignments=tuple(exprs), ) - ).merge_projections() + ) def assign(self, source_id: str, destination_id: str) -> ArrayValue: if destination_id in self.column_ids: # Mutate case @@ -209,7 +241,7 @@ def assign(self, source_id: str, destination_id: str) -> ArrayValue: child=self.node, assignments=tuple(exprs), ) - ).merge_projections() + ) def assign_constant( self, @@ -243,7 +275,7 @@ def assign_constant( child=self.node, assignments=tuple(exprs), ) - ).merge_projections() + ) def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue: selections = ((ex.free_var(col_id), col_id) for col_id in column_ids) @@ -252,7 +284,7 @@ def select_columns(self, column_ids: typing.Sequence[str]) -> ArrayValue: child=self.node, assignments=tuple(selections), ) - ).merge_projections() + ) def drop_columns(self, columns: Iterable[str]) -> ArrayValue: new_projection = ( @@ -265,7 +297,7 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue: child=self.node, assignments=tuple(new_projection), ) - ).merge_projections() + ) def aggregate( self, diff --git a/bigframes/core/blocks.py b/bigframes/core/blocks.py index 0bbb8a0b61..ea063669d5 100644 --- a/bigframes/core/blocks.py +++ b/bigframes/core/blocks.py @@ -124,7 +124,7 @@ def __init__( if len(index_columns) == 0: warnings.warn( "Creating object with Null Index. Null Index is a preview feature.", - category=bigframes.exceptions.PreviewWarning, + category=bigframes.exceptions.NullIndexPreviewWarning, ) self._index_columns = tuple(index_columns) # Index labels don't need complicated hierarchical access so can store as tuple @@ -155,7 +155,13 @@ def __init__( self._transpose_cache: Optional[Block] = transpose_cache @classmethod - def from_local(cls, data: pd.DataFrame, session: bigframes.Session) -> Block: + def from_local( + cls, + data: pd.DataFrame, + session: bigframes.Session, + *, + cache_transpose: bool = True, + ) -> Block: # Assumes caller has already converted datatypes to bigframes ones. pd_data = data column_labels = pd_data.columns @@ -169,12 +175,21 @@ def from_local(cls, data: pd.DataFrame, session: bigframes.Session) -> Block: pd_data = pd_data.reset_index(names=index_ids) as_pyarrow = pa.Table.from_pandas(pd_data, preserve_index=False) array_value = core.ArrayValue.from_pyarrow(as_pyarrow, session=session) - return cls( + block = cls( array_value, column_labels=column_labels, index_columns=index_ids, index_labels=index_labels, ) + if cache_transpose: + try: + # this cache will help when aligning on axis=1 + block = block.with_transpose_cache( + cls.from_local(data.T, session, cache_transpose=False) + ) + except Exception: + pass + return block @property def index(self) -> BlockIndexProperties: @@ -724,12 +739,18 @@ def with_column_labels( f"The column labels size `{len(label_list)} ` should equal to the value" + f"columns size: {len(self.value_columns)}." ) - return Block( + block = Block( self._expr, index_columns=self.index_columns, column_labels=label_list, index_labels=self.index.names, ) + singleton_label = len(list(value)) == 1 and list(value)[0] + if singleton_label is not None and self._transpose_cache is not None: + new_cache, label_id = self._transpose_cache.create_constant(singleton_label) + new_cache = new_cache.set_index([label_id]) + block = block.with_transpose_cache(new_cache) + return block def with_transpose_cache(self, transposed: Block): return Block( @@ -1014,6 +1035,34 @@ def aggregate_all_and_stack( index_labels=self.index.names, ) + def aggregate_size( + self, + by_column_ids: typing.Sequence[str] = (), + *, + dropna: bool = True, + ): + """Returns a block object to compute the size(s) of groups.""" + agg_specs = [ + (ex.NullaryAggregation(agg_ops.SizeOp()), guid.generate_guid()), + ] + output_col_ids = [agg_spec[1] for agg_spec in agg_specs] + result_expr = self.expr.aggregate(agg_specs, by_column_ids, dropna=dropna) + names: typing.List[Label] = [] + for by_col_id in by_column_ids: + if by_col_id in self.value_columns: + names.append(self.col_id_to_label[by_col_id]) + else: + names.append(self.col_id_to_index_name[by_col_id]) + return ( + Block( + result_expr, + index_columns=by_column_ids, + column_labels=["size"], + index_labels=names, + ), + output_col_ids, + ) + def select_column(self, id: str) -> Block: return self.select_columns([id]) @@ -1902,10 +1951,169 @@ def merge( coalesce_labels=matching_join_labels, suffixes=suffixes, ) - # Constructs default index - offset_index_id = guid.generate_guid() - expr = joined_expr.promote_offsets(offset_index_id) - return Block(expr, index_columns=[offset_index_id], column_labels=labels) + + # Construct a default index only if this object and the other both have + # indexes. In other words, joining anything to a NULL index object + # keeps everything as a NULL index. + # + # This keeps us from generating an index if the user joins a large + # BigQuery table against small local data, for example. + if len(self._index_columns) > 0 and len(other._index_columns) > 0: + offset_index_id = guid.generate_guid() + expr = joined_expr.promote_offsets(offset_index_id) + index_columns = [offset_index_id] + else: + expr = joined_expr + index_columns = [] + + return Block(expr, index_columns=index_columns, column_labels=labels) + + def _align_both_axes( + self, other: Block, how: str + ) -> Tuple[Block, pd.Index, Sequence[Tuple[ex.Expression, ex.Expression]]]: + # Join rows + aligned_block, (get_column_left, get_column_right) = self.join(other, how=how) + # join columns schema + # indexers will be none for exact match + if self.column_labels.equals(other.column_labels): + columns, lcol_indexer, rcol_indexer = self.column_labels, None, None + else: + columns, lcol_indexer, rcol_indexer = self.column_labels.join( + other.column_labels, how="outer", return_indexers=True + ) + lcol_indexer = ( + lcol_indexer if (lcol_indexer is not None) else range(len(columns)) + ) + rcol_indexer = ( + rcol_indexer if (rcol_indexer is not None) else range(len(columns)) + ) + + left_input_lookup = ( + lambda index: ex.free_var(get_column_left[self.value_columns[index]]) + if index != -1 + else ex.const(None) + ) + righ_input_lookup = ( + lambda index: ex.free_var(get_column_right[other.value_columns[index]]) + if index != -1 + else ex.const(None) + ) + + left_inputs = [left_input_lookup(i) for i in lcol_indexer] + right_inputs = [righ_input_lookup(i) for i in rcol_indexer] + return aligned_block, columns, tuple(zip(left_inputs, right_inputs)) + + def _align_axis_0( + self, other: Block, how: str + ) -> Tuple[Block, pd.Index, Sequence[Tuple[ex.Expression, ex.Expression]]]: + assert len(other.value_columns) == 1 + aligned_block, (get_column_left, get_column_right) = self.join(other, how=how) + + series_column_id = other.value_columns[0] + inputs = tuple( + ( + ex.free_var(get_column_left[col]), + ex.free_var(get_column_right[series_column_id]), + ) + for col in self.value_columns + ) + return aligned_block, self.column_labels, inputs + + def _align_series_block_axis_1( + self, other: Block, how: str + ) -> Tuple[Block, pd.Index, Sequence[Tuple[ex.Expression, ex.Expression]]]: + assert len(other.value_columns) == 1 + if other._transpose_cache is None: + raise ValueError( + "Wrong align method, this approach requires transpose cache" + ) + + # Join rows + aligned_block, (get_column_left, get_column_right) = join_with_single_row( + self, other.transpose() + ) + # join columns schema + # indexers will be none for exact match + if self.column_labels.equals(other.transpose().column_labels): + columns, lcol_indexer, rcol_indexer = self.column_labels, None, None + else: + columns, lcol_indexer, rcol_indexer = self.column_labels.join( + other.transpose().column_labels, how=how, return_indexers=True + ) + lcol_indexer = ( + lcol_indexer if (lcol_indexer is not None) else range(len(columns)) + ) + rcol_indexer = ( + rcol_indexer if (rcol_indexer is not None) else range(len(columns)) + ) + + left_input_lookup = ( + lambda index: ex.free_var(get_column_left[self.value_columns[index]]) + if index != -1 + else ex.const(None) + ) + righ_input_lookup = ( + lambda index: ex.free_var( + get_column_right[other.transpose().value_columns[index]] + ) + if index != -1 + else ex.const(None) + ) + + left_inputs = [left_input_lookup(i) for i in lcol_indexer] + right_inputs = [righ_input_lookup(i) for i in rcol_indexer] + return aligned_block, columns, tuple(zip(left_inputs, right_inputs)) + + def _align_pd_series_axis_1( + self, other: pd.Series, how: str + ) -> Tuple[Block, pd.Index, Sequence[Tuple[ex.Expression, ex.Expression]]]: + if self.column_labels.equals(other.index): + columns, lcol_indexer, rcol_indexer = self.column_labels, None, None + else: + if not (self.column_labels.is_unique and other.index.is_unique): + raise ValueError("Cannot align non-unique indices") + columns, lcol_indexer, rcol_indexer = self.column_labels.join( + other.index, how=how, return_indexers=True + ) + lcol_indexer = ( + lcol_indexer if (lcol_indexer is not None) else range(len(columns)) + ) + rcol_indexer = ( + rcol_indexer if (rcol_indexer is not None) else range(len(columns)) + ) + + left_input_lookup = ( + lambda index: ex.free_var(self.value_columns[index]) + if index != -1 + else ex.const(None) + ) + righ_input_lookup = ( + lambda index: ex.const(other.iloc[index]) if index != -1 else ex.const(None) + ) + + left_inputs = [left_input_lookup(i) for i in lcol_indexer] + right_inputs = [righ_input_lookup(i) for i in rcol_indexer] + return self, columns, tuple(zip(left_inputs, right_inputs)) + + def _apply_binop( + self, + op: ops.BinaryOp, + inputs: Sequence[Tuple[ex.Expression, ex.Expression]], + labels: pd.Index, + reverse: bool = False, + ) -> Block: + block = self + binop_result_ids = [] + for left_input, right_input in inputs: + expr = ( + op.as_expr(right_input, left_input) + if reverse + else op.as_expr(left_input, right_input) + ) + block, result_col_id = block.project_expr(expr) + binop_result_ids.append(result_col_id) + + return block.select_columns(binop_result_ids).with_column_labels(labels) def join( self, @@ -2026,27 +2234,17 @@ def to_sql_query( idx_labels, ) - def cached(self, *, optimize_offsets=False, force: bool = False) -> Block: - """Write the block to a session table and create a new block object that references it.""" + def cached(self, *, optimize_offsets=False, force: bool = False) -> None: + """Write the block to a session table.""" # use a heuristic for whether something needs to be cached if (not force) and self.session._is_trivially_executable(self.expr): - return self + return if optimize_offsets: - expr = self.session._cache_with_offsets(self.expr) + self.session._cache_with_offsets(self.expr) else: - expr = self.session._cache_with_cluster_cols( + self.session._cache_with_cluster_cols( self.expr, cluster_cols=self.index_columns ) - return self.swap_array_expr(expr) - - def swap_array_expr(self, expr: core.ArrayValue) -> Block: - # TODO: Validate schema unchanged - return Block( - expr, - index_columns=self.index_columns, - column_labels=self.column_labels, - index_labels=self.index.names, - ) def _is_monotonic( self, column_ids: typing.Union[str, Sequence[str]], increasing: bool @@ -2116,8 +2314,8 @@ def _get_rows_as_json_values(self) -> Block: # TODO(shobs): Replace direct SQL manipulation by structured expression # manipulation ordering_column_name = guid.generate_guid() - expr = self.session._cache_with_offsets(self.expr) - expr = expr.promote_offsets(ordering_column_name) + self.session._cache_with_offsets(self.expr) + expr = self.expr.promote_offsets(ordering_column_name) expr_sql = self.session._to_sql(expr) # Names of the columns to serialize for the row. @@ -2238,15 +2436,6 @@ def column_ids(self) -> Sequence[str]: """Column(s) to use as row labels.""" return self._block._index_columns - def __repr__(self) -> str: - """Converts an Index to a string.""" - # TODO(swast): Add a timeout here? If the query is taking a long time, - # maybe we just print the job metadata that we have so far? - # TODO(swast): Avoid downloading the whole index by using job - # metadata, like we do with DataFrame. - preview = self.to_pandas() - return repr(preview) - def to_pandas(self) -> pd.Index: """Executes deferred operations and downloads the results.""" if len(self.column_ids) == 0: @@ -2341,6 +2530,61 @@ def join_indexless( ) +def join_with_single_row( + left: Block, + single_row_block: Block, +) -> Tuple[Block, Tuple[Mapping[str, str], Mapping[str, str]],]: + """ + Special join case where other is a single row block. + This property is not validated, caller responsible for not passing multi-row block. + Preserves index of the left block, ignoring label of other. + """ + left_expr = left.expr + # ignore index columns by dropping them + right_expr = single_row_block.expr.select_columns(single_row_block.value_columns) + left_mappings = [ + join_defs.JoinColumnMapping( + source_table=join_defs.JoinSide.LEFT, + source_id=id, + destination_id=guid.generate_guid(), + ) + for id in left_expr.column_ids + ] + right_mappings = [ + join_defs.JoinColumnMapping( + source_table=join_defs.JoinSide.RIGHT, + source_id=id, + destination_id=guid.generate_guid(), + ) + for id in right_expr.column_ids # skip index column + ] + + join_def = join_defs.JoinDefinition( + conditions=(), + mappings=(*left_mappings, *right_mappings), + type="cross", + ) + combined_expr = left_expr.join( + right_expr, + join_def=join_def, + ) + get_column_left = join_def.get_left_mapping() + get_column_right = join_def.get_right_mapping() + # Drop original indices from each side. and used the coalesced combination generated by the join. + index_cols_post_join = [get_column_left[id] for id in left.index_columns] + + block = Block( + combined_expr, + index_columns=index_cols_post_join, + column_labels=left.column_labels.append(single_row_block.column_labels), + index_labels=[left.index.name], + ) + return ( + block, + (get_column_left, get_column_right), + ) + + def join_mono_indexed( left: Block, right: Block, @@ -2528,7 +2772,7 @@ def coalesce_columns( ) -> Tuple[core.ArrayValue, Sequence[str]]: result_ids = [] for left_id, right_id in zip(left_ids, right_ids): - if how == "left" or how == "inner": + if how == "left" or how == "inner" or how == "cross": result_ids.append(left_id) expr = expr.drop_columns([right_id]) elif how == "right": diff --git a/bigframes/core/compile/aggregate_compiler.py b/bigframes/core/compile/aggregate_compiler.py index c0b0562a54..fada4ebbd8 100644 --- a/bigframes/core/compile/aggregate_compiler.py +++ b/bigframes/core/compile/aggregate_compiler.py @@ -36,6 +36,8 @@ def compile_aggregate( bindings: typing.Dict[str, ibis_types.Value], order_by: typing.Sequence[ibis_types.Value] = [], ) -> ibis_types.Value: + if isinstance(aggregate, ex.NullaryAggregation): + return compile_nullary_agg(aggregate.op) if isinstance(aggregate, ex.UnaryAggregation): input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings) if aggregate.op.can_order_by: @@ -55,7 +57,9 @@ def compile_analytic( window: window_spec.WindowSpec, bindings: typing.Dict[str, ibis_types.Value], ) -> ibis_types.Value: - if isinstance(aggregate, ex.UnaryAggregation): + if isinstance(aggregate, ex.NullaryAggregation): + return compile_nullary_agg(aggregate.op, window) + elif isinstance(aggregate, ex.UnaryAggregation): input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings) return compile_unary_agg(aggregate.op, input, window) elif isinstance(aggregate, ex.BinaryAggregation): @@ -93,6 +97,14 @@ def compile_ordered_unary_agg( raise ValueError(f"Can't compile unrecognized operation: {op}") +@functools.singledispatch +def compile_nullary_agg( + op: agg_ops.WindowOp, + window: Optional[window_spec.WindowSpec] = None, +) -> ibis_types.Value: + raise ValueError(f"Can't compile unrecognized operation: {op}") + + def numeric_op(operation): @functools.wraps(operation) def constrained_op( @@ -118,6 +130,11 @@ def constrained_op( ### Specific Op implementations Below +@compile_nullary_agg.register +def _(op: agg_ops.SizeOp, window=None) -> ibis_types.NumericValue: + return _apply_window_if_present(vendored_ibis_ops.count(1), window) + + @compile_unary_agg.register @numeric_op def _( diff --git a/bigframes/core/compile/compiler.py b/bigframes/core/compile/compiler.py index a9908192f3..f948d10a5b 100644 --- a/bigframes/core/compile/compiler.py +++ b/bigframes/core/compile/compiler.py @@ -17,12 +17,20 @@ import io import typing +import ibis +import ibis.backends +import ibis.backends.bigquery +import ibis.expr.types import pandas as pd import bigframes.core.compile.compiled as compiled import bigframes.core.compile.concat as concat_impl +import bigframes.core.compile.default_ordering as default_ordering +import bigframes.core.compile.schema_translator import bigframes.core.compile.single_column import bigframes.core.nodes as nodes +import bigframes.core.ordering as bf_ordering +import bigframes.dtypes as bigframes_dtypes if typing.TYPE_CHECKING: import bigframes.core @@ -88,6 +96,87 @@ def compile_readlocal(node: nodes.ReadLocalNode, ordered: bool = True): return ordered_ir.to_unordered() +@_compile_node.register +def compile_readtable(node: nodes.ReadTableNode, ordered: bool = True): + if ordered: + return compile_read_table_ordered(node) + else: + return compile_read_table_unordered(node) + + +def read_table_as_unordered_ibis(node: nodes.ReadTableNode) -> ibis.expr.types.Table: + full_table_name = f"{node.project_id}.{node.dataset_id}.{node.table_id}" + used_columns = ( + *node.schema.names, + *[i for i in node.total_order_cols if i not in node.schema.names], + ) + # Physical schema might include unused columns, unsupported datatypes like JSON + physical_schema = ibis.backends.bigquery.BigQuerySchema.to_ibis( + list(i for i in node.physical_schema if i.name in used_columns) + ) + if node.at_time is not None or node.sql_predicate is not None: + import bigframes.session._io.bigquery + + sql = bigframes.session._io.bigquery.to_query( + full_table_name, + columns=used_columns, + sql_predicate=node.sql_predicate, + time_travel_timestamp=node.at_time, + ) + return ibis.backends.bigquery.Backend().sql(schema=physical_schema, query=sql) + else: + return ibis.table(physical_schema, full_table_name) + + +def compile_read_table_unordered(node: nodes.ReadTableNode): + ibis_table = read_table_as_unordered_ibis(node) + return compiled.UnorderedIR( + ibis_table, + tuple( + bigframes_dtypes.ibis_value_to_canonical_type(ibis_table[col]) + for col in node.schema.names + ), + ) + + +def compile_read_table_ordered(node: nodes.ReadTableNode): + ibis_table = read_table_as_unordered_ibis(node) + if node.total_order_cols: + ordering_value_columns = tuple( + bf_ordering.ascending_over(col) for col in node.total_order_cols + ) + if node.order_col_is_sequential: + integer_encoding = bf_ordering.IntegerEncoding( + is_encoded=True, is_sequential=True + ) + else: + integer_encoding = bf_ordering.IntegerEncoding() + ordering = bf_ordering.ExpressionOrdering( + ordering_value_columns, + integer_encoding=integer_encoding, + total_ordering_columns=frozenset(node.total_order_cols), + ) + hidden_columns = () + else: + ibis_table, ordering = default_ordering.gen_default_ordering( + ibis_table, use_double_hash=True + ) + hidden_columns = tuple( + ibis_table[col] + for col in ibis_table.columns + if col not in node.schema.names + ) + return compiled.OrderedIR( + ibis_table, + columns=tuple( + bigframes_dtypes.ibis_value_to_canonical_type(ibis_table[col]) + for col in node.schema.names + ), + ordering=ordering, + hidden_ordering_columns=hidden_columns, + ) + + @_compile_node.register def compile_readgbq(node: nodes.ReadGbqNode, ordered: bool = True): if ordered: diff --git a/bigframes/core/compile/default_ordering.py b/bigframes/core/compile/default_ordering.py new file mode 100644 index 0000000000..d8bdc0546b --- /dev/null +++ b/bigframes/core/compile/default_ordering.py @@ -0,0 +1,91 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Private helpers for loading a BigQuery table as a BigQuery DataFrames DataFrame. +""" + +from __future__ import annotations + +import itertools +from typing import cast + +import bigframes_vendored.ibis.expr.operations as vendored_ibis_ops +import ibis +import ibis.backends +import ibis.expr.datatypes as ibis_dtypes +import ibis.expr.types as ibis_types + +import bigframes.core.guid as guid +import bigframes.core.ordering as order + + +def _convert_to_nonnull_string(column: ibis_types.Column) -> ibis_types.StringValue: + col_type = column.type() + if ( + col_type.is_numeric() + or col_type.is_boolean() + or col_type.is_binary() + or col_type.is_temporal() + ): + result = column.cast(ibis_dtypes.String(nullable=True)) + elif col_type.is_geospatial(): + result = cast(ibis_types.GeoSpatialColumn, column).as_text() + elif col_type.is_string(): + result = column + else: + # TO_JSON_STRING works with all data types, but isn't the most efficient + # Needed for JSON, STRUCT and ARRAY datatypes + result = vendored_ibis_ops.ToJsonString(column).to_expr() # type: ignore + # Escape backslashes and use backslash as delineator + escaped = cast(ibis_types.StringColumn, result.fillna("")).replace("\\", "\\\\") # type: ignore + return cast(ibis_types.StringColumn, ibis.literal("\\")).concat(escaped) + + +def gen_default_ordering(table: ibis.table, use_double_hash: bool = True): + ordering_hash_part = guid.generate_guid("bigframes_ordering_") + ordering_hash_part2 = guid.generate_guid("bigframes_ordering_") + ordering_rand_part = guid.generate_guid("bigframes_ordering_") + + # All inputs into hash must be non-null or resulting hash will be null + str_values = list( + map(lambda col: _convert_to_nonnull_string(table[col]), table.columns) + ) + full_row_str = ( + str_values[0].concat(*str_values[1:]) if len(str_values) > 1 else str_values[0] + ) + full_row_hash = full_row_str.hash().name(ordering_hash_part) + # By modifying value slightly, we get another hash uncorrelated with the first + full_row_hash_p2 = (full_row_str + "_").hash().name(ordering_hash_part2) + # Used to disambiguate between identical rows (which will have identical hash) + random_value = ibis.random().name(ordering_rand_part) + + order_values = ( + [full_row_hash, full_row_hash_p2, random_value] + if use_double_hash + else [full_row_hash, random_value] + ) + + original_column_ids = table.columns + table_with_ordering = table.select( + itertools.chain(original_column_ids, order_values) + ) + + ordering = order.ExpressionOrdering( + ordering_value_columns=tuple( + order.ascending_over(col.get_name()) for col in order_values + ), + total_ordering_columns=frozenset(col.get_name() for col in order_values), + ) + return table_with_ordering, ordering diff --git a/bigframes/core/compile/scalar_op_compiler.py b/bigframes/core/compile/scalar_op_compiler.py index e8e5a1f3ac..000c4a4c09 100644 --- a/bigframes/core/compile/scalar_op_compiler.py +++ b/bigframes/core/compile/scalar_op_compiler.py @@ -375,6 +375,16 @@ def abs_op_impl(x: ibis_types.Value): return typing.cast(ibis_types.NumericValue, x).abs() +@scalar_op_compiler.register_unary_op(ops.pos_op) +def pos_op_impl(x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x) + + +@scalar_op_compiler.register_unary_op(ops.neg_op) +def neg_op_impl(x: ibis_types.Value): + return typing.cast(ibis_types.NumericValue, x).negate() + + @scalar_op_compiler.register_unary_op(ops.sqrt_op) def sqrt_op_impl(x: ibis_types.Value): numeric_value = typing.cast(ibis_types.NumericValue, x) @@ -856,11 +866,12 @@ def to_timestamp_op_impl(x: ibis_types.Value, op: ops.ToTimestampOp): @scalar_op_compiler.register_unary_op(ops.RemoteFunctionOp, pass_op=True) def remote_function_op_impl(x: ibis_types.Value, op: ops.RemoteFunctionOp): - if not hasattr(op.func, "bigframes_remote_function"): + ibis_node = getattr(op.func, "ibis_node", None) + if ibis_node is None: raise TypeError( f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}" ) - x_transformed = op.func(x) + x_transformed = ibis_node(x) if not op.apply_on_null: x_transformed = ibis.case().when(x.isnull(), x).else_(x_transformed).end() return x_transformed @@ -874,6 +885,12 @@ def map_op_impl(x: ibis_types.Value, op: ops.MapOp): return case.else_(x).end() +# Array Ops +@scalar_op_compiler.register_unary_op(ops.ArrayToStringOp, pass_op=True) +def array_to_string_op_impl(x: ibis_types.Value, op: ops.ArrayToStringOp): + return typing.cast(ibis_types.ArrayValue, x).join(op.delimiter) + + ### Binary Ops def short_circuit_nulls(type_override: typing.Optional[ibis_dtypes.DataType] = None): """Wraps a binary operator to generate nulls of the expected type if either input is a null scalar.""" @@ -978,6 +995,16 @@ def or_op( ) +@scalar_op_compiler.register_binary_op(ops.xor_op) +def xor_op( + x: ibis_types.Value, + y: ibis_types.Value, +): + return typing.cast(ibis_types.BooleanValue, x) ^ typing.cast( + ibis_types.BooleanValue, y + ) + + @scalar_op_compiler.register_binary_op(ops.add_op) @short_circuit_nulls() def add_op( @@ -1342,11 +1369,12 @@ def minimum_impl( def binary_remote_function_op_impl( x: ibis_types.Value, y: ibis_types.Value, op: ops.BinaryRemoteFunctionOp ): - if not hasattr(op.func, "bigframes_remote_function"): + ibis_node = getattr(op.func, "ibis_node", None) + if ibis_node is None: raise TypeError( f"only a bigframes remote function is supported as a callable. {constants.FEEDBACK_LINK}" ) - x_transformed = op.func(x, y) + x_transformed = ibis_node(x, y) return x_transformed diff --git a/bigframes/core/compile/schema_translator.py b/bigframes/core/compile/schema_translator.py new file mode 100644 index 0000000000..03e9691af6 --- /dev/null +++ b/bigframes/core/compile/schema_translator.py @@ -0,0 +1,34 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from __future__ import annotations + +import ibis +import ibis.expr.schema + +import bigframes.core.schema as bf_schema +import bigframes.dtypes + + +def convert_bf_schema(schema: bf_schema.ArraySchema) -> ibis.expr.schema.Schema: + """ + Convert bigframes schema to ibis schema. This is unambigous as every bigframes type is backed by a specific SQL/ibis dtype. + """ + names = schema.names + types = [ + bigframes.dtypes.bigframes_dtype_to_ibis_dtype(bf_type) + for bf_type in schema.dtypes + ] + return ibis.schema(names=names, types=types) diff --git a/bigframes/core/expression.py b/bigframes/core/expression.py index 70eb519a1b..c216c29717 100644 --- a/bigframes/core/expression.py +++ b/bigframes/core/expression.py @@ -46,6 +46,16 @@ def output_type( ... +@dataclasses.dataclass(frozen=True) +class NullaryAggregation(Aggregation): + op: agg_ops.NullaryWindowOp = dataclasses.field() + + def output_type( + self, input_types: dict[str, bigframes.dtypes.Dtype] + ) -> dtypes.ExpressionType: + return self.op.output_type() + + @dataclasses.dataclass(frozen=True) class UnaryAggregation(Aggregation): op: agg_ops.UnaryWindowOp = dataclasses.field() diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 41d0750030..ee120635d3 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -104,6 +104,20 @@ def __getitem__( dropna=self._dropna, ) + def size(self) -> typing.Union[df.DataFrame, series.Series]: + agg_block, _ = self._block.aggregate_size( + by_column_ids=self._by_col_ids, + dropna=self._dropna, + ) + agg_block = agg_block.with_column_labels(pd.Index(["size"])) + dataframe = df.DataFrame(agg_block) + + if self._as_index: + series = dataframe["size"] + return series.rename(None) + else: + return self._convert_index(dataframe) + def sum(self, numeric_only: bool = False, *args) -> df.DataFrame: if not numeric_only: self._raise_on_non_numeric("sum") @@ -325,9 +339,30 @@ def _agg_list(self, func: typing.Sequence) -> df.DataFrame: for col_id in self._aggregated_columns() for f in func ] - column_labels = [ - (col_id, f) for col_id in self._aggregated_columns() for f in func - ] + + if self._block.column_labels.nlevels > 1: + # Restructure MultiIndex for proper format: (idx1, idx2, func) + # rather than ((idx1, idx2), func). + aggregated_columns = pd.MultiIndex.from_tuples( + [ + self._block.col_id_to_label[col_id] + for col_id in self._aggregated_columns() + ], + names=[*self._block.column_labels.names], + ).to_frame(index=False) + + column_labels = [ + tuple(col_id) + (f,) + for col_id in aggregated_columns.to_numpy() + for f in func + ] + else: + column_labels = [ + (self._block.col_id_to_label[col_id], f) + for col_id in self._aggregated_columns() + for f in func + ] + agg_block, _ = self._block.aggregate( by_column_ids=self._by_col_ids, aggregations=aggregations, @@ -520,6 +555,13 @@ def std(self, *args, **kwargs) -> series.Series: def var(self, *args, **kwargs) -> series.Series: return self._aggregate(agg_ops.var_op) + def size(self) -> series.Series: + agg_block, _ = self._block.aggregate_size( + by_column_ids=self._by_col_ids, + dropna=self._dropna, + ) + return series.Series(agg_block.with_column_labels([self._value_name])) + def skew(self, *args, **kwargs) -> series.Series: block = block_ops.skew(self._block, [self._value_column], self._by_col_ids) return series.Series(block) diff --git a/bigframes/core/indexes/base.py b/bigframes/core/indexes/base.py index 7f2c56c20a..e40e20b0cb 100644 --- a/bigframes/core/indexes/base.py +++ b/bigframes/core/indexes/base.py @@ -95,6 +95,7 @@ def __new__( result = typing.cast(Index, object.__new__(klass)) # type: ignore result._query_job = None result._block = block + block.session._register_object(result) return result @classmethod diff --git a/bigframes/core/log_adapter.py b/bigframes/core/log_adapter.py index 877e4a9fa1..4afa6037de 100644 --- a/bigframes/core/log_adapter.py +++ b/bigframes/core/log_adapter.py @@ -99,9 +99,12 @@ def add_api_method(api_method_name): _api_methods = _api_methods[:MAX_LABELS_COUNT] -def get_and_reset_api_methods(): +def get_and_reset_api_methods(dry_run: bool = False): global _lock with _lock: previous_api_methods = list(_api_methods) - _api_methods.clear() + + # dry_run might not make a job resource, so only reset the log on real queries. + if not dry_run: + _api_methods.clear() return previous_api_methods diff --git a/bigframes/core/nodes.py b/bigframes/core/nodes.py index 688e165732..1af7c5bd17 100644 --- a/bigframes/core/nodes.py +++ b/bigframes/core/nodes.py @@ -16,11 +16,14 @@ import abc from dataclasses import dataclass, field, fields, replace +import datetime import functools import itertools import typing from typing import Callable, Tuple +import google.cloud.bigquery as bq + import bigframes.core.expression as ex import bigframes.core.guid from bigframes.core.join_def import JoinColumnMapping, JoinDefinition, JoinSide @@ -345,6 +348,67 @@ def transform_children( return self +## Put ordering in here or just add order_by node above? +@dataclass(frozen=True) +class ReadTableNode(BigFrameNode): + project_id: str = field() + dataset_id: str = field() + table_id: str = field() + + physical_schema: Tuple[bq.SchemaField, ...] = field() + # Subset of physical schema columns, with chosen BQ types + columns: schemata.ArraySchema = field() + + table_session: bigframes.session.Session = field() + # Empty tuple if no primary key (primary key can be any set of columns that together form a unique key) + # Empty if no known unique key + total_order_cols: Tuple[str, ...] = field() + # indicates a primary key that is exactly offsets 0, 1, 2, ..., N-2, N-1 + order_col_is_sequential: bool = False + at_time: typing.Optional[datetime.datetime] = None + # Added for backwards compatibility, not validated + sql_predicate: typing.Optional[str] = None + + def __post_init__(self): + # enforce invariants + physical_names = set(map(lambda i: i.name, self.physical_schema)) + if not set(self.columns.names).issubset(physical_names): + raise ValueError( + f"Requested schema {self.columns} cannot be derived from table schemal {self.physical_schema}" + ) + if self.order_col_is_sequential and len(self.total_order_cols) == 1: + raise ValueError("Sequential primary key must have only one component") + + @property + def session(self): + return self.table_session + + def __hash__(self): + return self._node_hash + + @property + def roots(self) -> typing.Set[BigFrameNode]: + return {self} + + @property + def schema(self) -> schemata.ArraySchema: + return self.columns + + @property + def relation_ops_created(self) -> int: + # Assume worst case, where readgbq actually has baked in analytic operation to generate index + return 3 + + @functools.cached_property + def variables_introduced(self) -> int: + return len(self.schema.items) + 1 + + def transform_children( + self, t: Callable[[BigFrameNode], BigFrameNode] + ) -> BigFrameNode: + return self + + # Unary nodes @dataclass(frozen=True) class PromoteOffsetsNode(UnaryNode): diff --git a/bigframes/core/schema.py b/bigframes/core/schema.py index 3629778aaf..30a2a6593a 100644 --- a/bigframes/core/schema.py +++ b/bigframes/core/schema.py @@ -18,6 +18,8 @@ import functools import typing +import google.cloud.bigquery + import bigframes.core.guid import bigframes.dtypes @@ -34,6 +36,16 @@ class SchemaItem: class ArraySchema: items: typing.Tuple[SchemaItem, ...] + @classmethod + def from_bq_table(cls, table: google.cloud.bigquery.Table): + items = tuple( + SchemaItem(name, dtype) + for name, dtype in bigframes.dtypes.bf_type_from_type_kind( + table.schema + ).items() + ) + return ArraySchema(items) + @property def names(self) -> typing.Tuple[str, ...]: return tuple(item.column for item in self.items) @@ -51,6 +63,11 @@ def drop(self, columns: typing.Iterable[str]) -> ArraySchema: tuple(item for item in self.items if item.column not in columns) ) + def select(self, columns: typing.Iterable[str]) -> ArraySchema: + return ArraySchema( + tuple(SchemaItem(name, self.get_type(name)) for name in columns) + ) + def append(self, item: SchemaItem): return ArraySchema(tuple([*self.items, item])) diff --git a/bigframes/core/tree_properties.py b/bigframes/core/tree_properties.py index 125a7e6bff..2847a8f7f1 100644 --- a/bigframes/core/tree_properties.py +++ b/bigframes/core/tree_properties.py @@ -15,7 +15,7 @@ import functools import itertools -from typing import Dict +from typing import Callable, Dict, Optional import bigframes.core.nodes as nodes @@ -40,46 +40,66 @@ def peekable(node: nodes.BigFrameNode) -> bool: return children_peekable and self_peekable -def count_complex_nodes( - root: nodes.BigFrameNode, min_complexity: float, max_complexity: float -) -> Dict[nodes.BigFrameNode, int]: +# Replace modified_cost(node) = cost(apply_cache(node)) +def select_cache_target( + root: nodes.BigFrameNode, + min_complexity: float, + max_complexity: float, + cache: dict[nodes.BigFrameNode, nodes.BigFrameNode], + heuristic: Callable[[int, int], float], +) -> Optional[nodes.BigFrameNode]: + """Take tree, and return candidate nodes with (# of occurences, post-caching planning complexity). + + heurstic takes two args, node complexity, and node occurence count, in that order + """ + + @functools.cache + def _with_caching(subtree: nodes.BigFrameNode) -> nodes.BigFrameNode: + return replace_nodes(subtree, cache) + + def _combine_counts( + left: Dict[nodes.BigFrameNode, int], right: Dict[nodes.BigFrameNode, int] + ) -> Dict[nodes.BigFrameNode, int]: + return { + key: left.get(key, 0) + right.get(key, 0) + for key in itertools.chain(left.keys(), right.keys()) + } + @functools.cache def _node_counts_inner( subtree: nodes.BigFrameNode, ) -> Dict[nodes.BigFrameNode, int]: """Helper function to count occurences of duplicate nodes in a subtree. Considers only nodes in a complexity range""" empty_counts: Dict[nodes.BigFrameNode, int] = {} - if subtree.planning_complexity >= min_complexity: + subtree_complexity = _with_caching(subtree).planning_complexity + if subtree_complexity >= min_complexity: child_counts = [_node_counts_inner(child) for child in subtree.child_nodes] node_counts = functools.reduce(_combine_counts, child_counts, empty_counts) - if subtree.planning_complexity <= max_complexity: + if subtree_complexity <= max_complexity: return _combine_counts(node_counts, {subtree: 1}) else: return node_counts return empty_counts - return _node_counts_inner(root) + node_counts = _node_counts_inner(root) + + return max( + node_counts.keys(), + key=lambda node: heuristic( + _with_caching(node).planning_complexity, node_counts[node] + ), + ) def replace_nodes( root: nodes.BigFrameNode, - to_replace: nodes.BigFrameNode, - replacemenet: nodes.BigFrameNode, + replacements: dict[nodes.BigFrameNode, nodes.BigFrameNode], ): @functools.cache - def apply_substition(n: nodes.BigFrameNode) -> nodes.BigFrameNode: - if n == to_replace: - return replacemenet + def apply_substition(node: nodes.BigFrameNode) -> nodes.BigFrameNode: + if node in replacements.keys(): + return replacements[node] else: - return n.transform_children(apply_substition) - - return root.transform_children(apply_substition) - + return node.transform_children(apply_substition) -def _combine_counts( - left: Dict[nodes.BigFrameNode, int], right: Dict[nodes.BigFrameNode, int] -) -> Dict[nodes.BigFrameNode, int]: - return { - key: left.get(key, 0) + right.get(key, 0) - for key in itertools.chain(left.keys(), right.keys()) - } + return apply_substition(root) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 105588de2f..e404e439ab 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -191,6 +191,7 @@ def __init__( else: self._block = bigframes.pandas.read_pandas(pd_dataframe)._get_block() self._query_job: Optional[bigquery.QueryJob] = None + self._block.session._register_object(self) def __dir__(self): return dir(type(self)) + [ @@ -399,11 +400,12 @@ def memory_usage(self, index: bool = True): column_sizes = self.dtypes.map( lambda dtype: bigframes.dtypes.DTYPE_BYTE_SIZES.get(dtype, 8) * n_rows ) - if index: + if index and self._has_index: index_size = pandas.Series([self.index._memory_usage()], index=["Index"]) column_sizes = pandas.concat([index_size, column_sizes]) return column_sizes + @requires_index def info( self, verbose: Optional[bool] = None, @@ -712,13 +714,9 @@ def _apply_binop( DataFrame(other), op, how=how, reverse=reverse ) elif utils.get_axis_number(axis) == 0: - bf_series = bigframes.core.convert.to_bf_series( - other, self.index, self._session - ) - return self._apply_series_binop_axis_0(bf_series, op, how, reverse) + return self._apply_series_binop_axis_0(other, op, how, reverse) elif utils.get_axis_number(axis) == 1: - pd_series = bigframes.core.convert.to_pd_series(other, self.columns) - return self._apply_series_binop_axis_1(pd_series, op, how, reverse) + return self._apply_series_binop_axis_1(other, op, how, reverse) raise NotImplementedError( f"binary operation is not implemented on the second operand of type {type(other).__name__}." f"{constants.FEEDBACK_LINK}" @@ -743,89 +741,49 @@ def _apply_scalar_binop( def _apply_series_binop_axis_0( self, - other: bigframes.series.Series, + other, op: ops.BinaryOp, how: str = "outer", reverse: bool = False, ) -> DataFrame: - block, (get_column_left, get_column_right) = self._block.join( - other._block, how=how + bf_series = bigframes.core.convert.to_bf_series( + other, self.index, self._session ) - - series_column_id = other._value_column - series_col = get_column_right[series_column_id] - for column_id, label in zip( - self._block.value_columns, self._block.column_labels - ): - self_col = get_column_left[column_id] - expr = ( - op.as_expr(series_col, self_col) - if reverse - else op.as_expr(self_col, series_col) - ) - block, _ = block.project_expr(expr, label) - block = block.drop_columns([get_column_left[column_id]]) - - block = block.drop_columns([series_col]) - block = block.with_index_labels(self.index.names) - return DataFrame(block) + aligned_block, columns, expr_pairs = self._block._align_axis_0( + bf_series._block, how=how + ) + result = aligned_block._apply_binop( + op, inputs=expr_pairs, labels=columns, reverse=reverse + ) + return DataFrame(result) def _apply_series_binop_axis_1( self, - other: pandas.Series, + other, op: ops.BinaryOp, how: str = "outer", reverse: bool = False, ) -> DataFrame: - # Somewhat different alignment than df-df so separate codepath for now. - if self.columns.equals(other.index): - columns, lcol_indexer, rcol_indexer = self.columns, None, None + """Align dataframe with pandas series by inlining series values as literals.""" + # If we already know the transposed schema (from the transpose cache), we don't need to materialize rows from other + # Instead, can fully defer execution (as a cross-join) + if ( + isinstance(other, bigframes.series.Series) + and other._block._transpose_cache is not None + ): + aligned_block, columns, expr_pairs = self._block._align_series_block_axis_1( + other._block, how=how + ) else: - if not (self.columns.is_unique and other.index.is_unique): - raise ValueError("Cannot align non-unique indices") - columns, lcol_indexer, rcol_indexer = self.columns.join( - other.index, how=how, return_indexers=True + # Fallback path, materialize `other` locally + pd_series = bigframes.core.convert.to_pd_series(other, self.columns) + aligned_block, columns, expr_pairs = self._block._align_pd_series_axis_1( + pd_series, how=how ) - - binop_result_ids = [] - - column_indices = zip( - lcol_indexer if (lcol_indexer is not None) else range(len(columns)), - rcol_indexer if (rcol_indexer is not None) else range(len(columns)), + result = aligned_block._apply_binop( + op, inputs=expr_pairs, labels=columns, reverse=reverse ) - - block = self._block - for left_index, right_index in column_indices: - if left_index >= 0 and right_index >= 0: # -1 indices indicate missing - self_col_id = self._block.value_columns[left_index] - other_scalar = other.iloc[right_index] - expr = ( - op.as_expr(ex.const(other_scalar), self_col_id) - if reverse - else op.as_expr(self_col_id, ex.const(other_scalar)) - ) - elif left_index >= 0: - self_col_id = self._block.value_columns[left_index] - expr = ( - op.as_expr(ex.const(None), self_col_id) - if reverse - else op.as_expr(self_col_id, ex.const(None)) - ) - elif right_index >= 0: - other_scalar = other.iloc[right_index] - expr = ( - op.as_expr(ex.const(other_scalar), ex.const(None)) - if reverse - else op.as_expr(ex.const(None), ex.const(other_scalar)) - ) - else: - # Should not be possible - raise ValueError("No right or left index.") - block, result_col_id = block.project_expr(expr) - binop_result_ids.append(result_col_id) - - block = block.select_columns(binop_result_ids) - return DataFrame(block.with_column_labels(columns)) + return DataFrame(result) def _apply_dataframe_binop( self, @@ -834,57 +792,13 @@ def _apply_dataframe_binop( how: str = "outer", reverse: bool = False, ) -> DataFrame: - # Join rows - block, (get_column_left, get_column_right) = self._block.join( + aligned_block, columns, expr_pairs = self._block._align_both_axes( other._block, how=how ) - # join columns schema - # indexers will be none for exact match - if self.columns.equals(other.columns): - columns, lcol_indexer, rcol_indexer = self.columns, None, None - else: - columns, lcol_indexer, rcol_indexer = self.columns.join( - other.columns, how=how, return_indexers=True - ) - - binop_result_ids = [] - - column_indices = zip( - lcol_indexer if (lcol_indexer is not None) else range(len(columns)), - rcol_indexer if (rcol_indexer is not None) else range(len(columns)), + result = aligned_block._apply_binop( + op, inputs=expr_pairs, labels=columns, reverse=reverse ) - - for left_index, right_index in column_indices: - if left_index >= 0 and right_index >= 0: # -1 indices indicate missing - self_col_id = get_column_left[self._block.value_columns[left_index]] - other_col_id = get_column_right[other._block.value_columns[right_index]] - expr = ( - op.as_expr(other_col_id, self_col_id) - if reverse - else op.as_expr(self_col_id, other_col_id) - ) - elif left_index >= 0: - self_col_id = get_column_left[self._block.value_columns[left_index]] - expr = ( - op.as_expr(ex.const(None), self_col_id) - if reverse - else op.as_expr(self_col_id, ex.const(None)) - ) - elif right_index >= 0: - other_col_id = get_column_right[other._block.value_columns[right_index]] - expr = ( - op.as_expr(other_col_id, ex.const(None)) - if reverse - else op.as_expr(ex.const(None), other_col_id) - ) - else: - # Should not be possible - raise ValueError("No right or left index.") - block, result_col_id = block.project_expr(expr) - binop_result_ids.append(result_col_id) - - block = block.select_columns(binop_result_ids).with_column_labels(columns) - return DataFrame(block) + return DataFrame(result) def eq(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: return self._apply_binop(other, ops.eq_op, axis=axis) @@ -902,6 +816,11 @@ def __ne__(self, other) -> DataFrame: # type: ignore __ne__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__ne__) + def __invert__(self) -> DataFrame: + return self._apply_unary_op(ops.invert_op) + + __invert__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__invert__) + def le(self, other: typing.Any, axis: str | int = "columns") -> DataFrame: return self._apply_binop(other, ops.le_op, axis=axis) @@ -1105,6 +1024,33 @@ def __rpow__(self, other): __rpow__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__rpow__) + def __and__(self, other: bool | int | bigframes.series.Series) -> DataFrame: + return self._apply_binop(other, ops.and_op) + + __and__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__and__) + + __rand__ = __and__ + + def __or__(self, other: bool | int | bigframes.series.Series) -> DataFrame: + return self._apply_binop(other, ops.or_op) + + __or__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__or__) + + __ror__ = __or__ + + def __xor__(self, other: bool | int | bigframes.series.Series) -> DataFrame: + return self._apply_binop(other, ops.xor_op) + + __xor__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__xor__) + + __rxor__ = __xor__ + + def __pos__(self) -> DataFrame: + return self._apply_unary_op(ops.pos_op) + + def __neg__(self) -> DataFrame: + return self._apply_unary_op(ops.neg_op) + def align( self, other: typing.Union[DataFrame, bigframes.series.Series], @@ -1610,7 +1556,7 @@ def _assign_series_join_on_index( # Update case, remove after copying into columns block = block.drop_columns([source_column]) - return DataFrame(block.with_index_labels(self.index.names)) + return DataFrame(block.with_index_labels(self._block.index.names)) def reset_index(self, *, drop: bool = False) -> DataFrame: block = self._block.reset_index(drop) @@ -2905,7 +2851,9 @@ def to_csv( field_delimiter=sep, header=header, ) - _, query_job = self._block.expr.session._start_query(export_data_statement) + _, query_job = self._block.expr.session._start_query( + export_data_statement, api_name="dataframe-to_csv" + ) self._set_internal_query_job(query_job) def to_json( @@ -2947,7 +2895,9 @@ def to_json( format="JSON", export_options={}, ) - _, query_job = self._block.expr.session._start_query(export_data_statement) + _, query_job = self._block.expr.session._start_query( + export_data_statement, api_name="dataframe-to_json" + ) self._set_internal_query_job(query_job) def to_gbq( @@ -3079,7 +3029,9 @@ def to_parquet( format="PARQUET", export_options=export_options, ) - _, query_job = self._block.expr.session._start_query(export_data_statement) + _, query_job = self._block.expr.session._start_query( + export_data_statement, api_name="dataframe-to_parquet" + ) self._set_internal_query_job(query_job) def to_dict( @@ -3282,7 +3234,7 @@ def _prepare_export( array_value = self._block.expr new_col_labels, new_idx_labels = utils.get_standardized_ids( - self._block.column_labels, self.index.names + self._block.column_labels, self._block.index.names ) columns = list(self._block.value_columns) @@ -3319,7 +3271,7 @@ def _run_io_query( session = self._block.expr.session self._optimize_query_complexity() export_array, id_overrides = self._prepare_export( - index=index, ordering_id=ordering_id + index=index and self._has_index, ordering_id=ordering_id ) _, query_job = session._execute( @@ -3515,7 +3467,7 @@ def _cached(self, *, force: bool = False) -> DataFrame: No-op if the dataframe represents a trivial transformation of an existing materialization. Force=True is used for BQML integration where need to copy data rather than use snapshot. """ - self._set_block(self._block.cached(force=force)) + self._block.cached(force=force) return self def _optimize_query_complexity(self): @@ -3523,8 +3475,7 @@ def _optimize_query_complexity(self): May generate many queries and take substantial time to execute. """ # TODO: Move all this to session - new_expr = self._session._simplify_with_caching(self._block.expr) - self._set_block(self._block.swap_array_expr(new_expr)) + self._session._simplify_with_caching(self._block.expr) _DataFrameOrSeries = typing.TypeVar("_DataFrameOrSeries") diff --git a/bigframes/dtypes.py b/bigframes/dtypes.py index 2a344aff2d..3df67ed9e4 100644 --- a/bigframes/dtypes.py +++ b/bigframes/dtypes.py @@ -19,6 +19,7 @@ import textwrap import typing from typing import Any, Dict, Iterable, Literal, Tuple, Union +import warnings import bigframes_vendored.ibis.backends.bigquery.datatypes as third_party_ibis_bqtypes import bigframes_vendored.ibis.expr.operations as vendored_ibis_ops @@ -33,6 +34,7 @@ import pyarrow as pa import bigframes.constants as constants +import bigframes.exceptions # Type hints for Pandas dtypes supported by BigQuery DataFrame Dtype = Union[ @@ -132,6 +134,14 @@ def is_array_like(type: ExpressionType) -> bool: ) +def is_array_string_like(type: ExpressionType) -> bool: + return ( + isinstance(type, pd.ArrowDtype) + and isinstance(type.pyarrow_dtype, pa.ListType) + and pa.types.is_string(type.pyarrow_dtype.value_type) + ) + + def is_struct_like(type: ExpressionType) -> bool: return isinstance(type, pd.ArrowDtype) and isinstance( type.pyarrow_dtype, pa.StructType @@ -181,6 +191,10 @@ def is_bool_coercable(type: ExpressionType) -> bool: ibis_dtypes.Decimal(precision=76, scale=38, nullable=True), pd.ArrowDtype(pa.decimal256(76, 38)), ), + ( + ibis_dtypes.GeoSpatial(geotype="geography", srid=4326, nullable=True), + gpd.array.GeometryDtype(), + ), ) BIGFRAMES_TO_IBIS: Dict[Dtype, ibis_dtypes.DataType] = { @@ -212,9 +226,6 @@ def is_bool_coercable(type: ExpressionType) -> bool: ) IBIS_TO_BIGFRAMES.update( { - ibis_dtypes.GeoSpatial( - geotype="geography", srid=4326, nullable=True - ): gpd.array.GeometryDtype(), # TODO: Interval } ) @@ -280,6 +291,14 @@ def ibis_dtype_to_bigframes_dtype( if isinstance(ibis_dtype, ibis_dtypes.Integer): return pd.Int64Dtype() + # Temporary: Will eventually support an explicit json type instead of casting to string. + if isinstance(ibis_dtype, ibis_dtypes.JSON): + warnings.warn( + "Interpreting JSON as string. This behavior may change in future versions.", + bigframes.exceptions.PreviewWarning, + ) + return STRING_DTYPE + if ibis_dtype in IBIS_TO_BIGFRAMES: return IBIS_TO_BIGFRAMES[ibis_dtype] elif isinstance(ibis_dtype, ibis_dtypes.Decimal): diff --git a/bigframes/exceptions.py b/bigframes/exceptions.py index 027b3a4236..bae239b6da 100644 --- a/bigframes/exceptions.py +++ b/bigframes/exceptions.py @@ -39,5 +39,13 @@ class PreviewWarning(Warning): """The feature is in preview.""" +class NullIndexPreviewWarning(PreviewWarning): + """Null index feature is in preview.""" + + class NullIndexError(ValueError): """Object has no index.""" + + +class TimeTravelDisabledWarning(Warning): + """A query was reattempted without time travel.""" diff --git a/bigframes/functions/remote_function.py b/bigframes/functions/remote_function.py index 9eff802cc7..7be252406c 100644 --- a/bigframes/functions/remote_function.py +++ b/bigframes/functions/remote_function.py @@ -14,6 +14,7 @@ from __future__ import annotations +import collections.abc import hashlib import inspect import logging @@ -24,7 +25,17 @@ import sys import tempfile import textwrap -from typing import List, NamedTuple, Optional, Sequence, TYPE_CHECKING, Union +from typing import ( + Any, + cast, + List, + Mapping, + NamedTuple, + Optional, + Sequence, + TYPE_CHECKING, + Union, +) import warnings import ibis @@ -133,6 +144,8 @@ def __init__( cloud_function_service_account, cloud_function_kms_key_name, cloud_function_docker_repository, + *, + session: Session, ): self._gcp_project_id = gcp_project_id self._cloud_function_region = cloud_function_region @@ -145,6 +158,7 @@ def __init__( self._cloud_function_service_account = cloud_function_service_account self._cloud_function_kms_key_name = cloud_function_kms_key_name self._cloud_function_docker_repository = cloud_function_docker_repository + self._session = session def create_bq_remote_function( self, @@ -216,10 +230,8 @@ def create_bq_remote_function( # This requires bigquery.datasets.create IAM permission self._bq_client.create_dataset(dataset, exists_ok=True) - # TODO: Use session._start_query() so we get progress bar - query_job = self._bq_client.query(create_function_ddl) # Make an API request. - query_job.result() # Wait for the job to complete. - + # TODO(swast): plumb through the original, user-facing api_name. + _, query_job = self._session._start_query(create_function_ddl) logger.info(f"Created remote function {query_job.ddl_target_routine}") def get_cloud_function_fully_qualified_parent(self): @@ -735,8 +747,8 @@ def get_routine_reference( # which has moved as @js to the ibis package # https://github.com/ibis-project/ibis/blob/master/ibis/backends/bigquery/udf/__init__.py def remote_function( - input_types: Union[type, Sequence[type]], - output_type: type, + input_types: Union[None, type, Sequence[type]] = None, + output_type: Optional[type] = None, session: Optional[Session] = None, bigquery_client: Optional[bigquery.Client] = None, bigquery_connection_client: Optional[ @@ -800,11 +812,11 @@ def remote_function( `$ gcloud projects add-iam-policy-binding PROJECT_ID --member="serviceAccount:CONNECTION_SERVICE_ACCOUNT_ID" --role="roles/run.invoker"`. Args: - input_types (type or sequence(type)): + input_types (None, type, or sequence(type)): For scalar user defined function it should be the input type or sequence of input types. For row processing user defined function, type `Series` should be specified. - output_type (type): + output_type (Optional[type]): Data type of the output in the user defined function. session (bigframes.Session, Optional): BigQuery DataFrames session to use for getting default project, @@ -907,28 +919,12 @@ def remote_function( service(s) that are on a VPC network. See for more details https://cloud.google.com/functions/docs/networking/connecting-vpc. """ - is_row_processor = False - - import bigframes.series - - if input_types == bigframes.series.Series: - warnings.warn( - "input_types=Series scenario is in preview.", - stacklevel=1, - category=bigframes.exceptions.PreviewWarning, - ) - - # we will model the row as a json serialized string containing the data - # and the metadata representing the row - input_types = [str] - is_row_processor = True - elif isinstance(input_types, type): - input_types = [input_types] - # Some defaults may be used from the session if not provided otherwise import bigframes.pandas as bpd + import bigframes.series + import bigframes.session - session = session or bpd.get_global_session() + session = cast(bigframes.session.Session, session or bpd.get_global_session()) # A BigQuery client is required to perform BQ operations if not bigquery_client: @@ -1018,11 +1014,67 @@ def remote_function( bq_connection_manager = None if session is None else session.bqconnectionmanager - def wrapper(f): - if not callable(f): - raise TypeError("f must be callable, got {}".format(f)) + def wrapper(func): + nonlocal input_types, output_type + + if not callable(func): + raise TypeError("f must be callable, got {}".format(func)) + + if sys.version_info >= (3, 10): + # Add `eval_str = True` so that deferred annotations are turned into their + # corresponding type objects. Need Python 3.10 for eval_str parameter. + # https://docs.python.org/3/library/inspect.html#inspect.signature + signature_kwargs: Mapping[str, Any] = {"eval_str": True} + else: + signature_kwargs = {} + + signature = inspect.signature( + func, + **signature_kwargs, + ) + + # Try to get input types via type annotations. + if input_types is None: + input_types = [] + for parameter in signature.parameters.values(): + if (param_type := parameter.annotation) is inspect.Signature.empty: + raise ValueError( + "'input_types' was not set and parameter " + f"'{parameter.name}' is missing a type annotation. " + "Types are required to use @remote_function." + ) + input_types.append(param_type) + elif not isinstance(input_types, collections.abc.Sequence): + input_types = [input_types] + + if output_type is None: + if (output_type := signature.return_annotation) is inspect.Signature.empty: + raise ValueError( + "'output_type' was not set and function is missing a " + "return type annotation. Types are required to use " + "@remote_function." + ) + + # The function will actually be receiving a pandas Series, but allow both + # BigQuery DataFrames and pandas object types for compatibility. + is_row_processor = False + if len(input_types) == 1 and ( + (input_type := input_types[0]) == bigframes.series.Series + or input_type == pandas.Series + ): + warnings.warn( + "input_types=Series is in preview.", + stacklevel=1, + category=bigframes.exceptions.PreviewWarning, + ) + + # we will model the row as a json serialized string containing the data + # and the metadata representing the row + input_types = [str] + is_row_processor = True + elif isinstance(input_types, type): + input_types = [input_types] - signature = inspect.signature(f) # TODO(b/340898611): fix type error ibis_signature = ibis_signature_from_python_signature( signature, input_types, output_type # type: ignore @@ -1040,10 +1092,26 @@ def wrapper(f): cloud_function_service_account, cloud_function_kms_key_name, cloud_function_docker_repository, + session=session, # type: ignore ) + # In the unlikely case where the user is trying to re-deploy the same + # function, cleanup the attributes we add below, first. This prevents + # the pickle from having dependencies that might not otherwise be + # present such as ibis or pandas. + def try_delattr(attr): + try: + delattr(func, attr) + except AttributeError: + pass + + try_delattr("bigframes_cloud_function") + try_delattr("bigframes_remote_function") + try_delattr("output_dtype") + try_delattr("ibis_node") + rf_name, cf_name = remote_function_client.provision_bq_remote_function( - f, + func, ibis_signature.input_types, ibis_signature.output_type, reuse, @@ -1058,19 +1126,20 @@ def wrapper(f): # TODO: Move ibis logic to compiler step node = ibis.udf.scalar.builtin( - f, + func, name=rf_name, schema=f"{dataset_ref.project}.{dataset_ref.dataset_id}", signature=(ibis_signature.input_types, ibis_signature.output_type), ) - node.bigframes_cloud_function = ( + func.bigframes_cloud_function = ( remote_function_client.get_cloud_function_fully_qualified_name(cf_name) ) - node.bigframes_remote_function = str(dataset_ref.routine(rf_name)) # type: ignore - node.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype( + func.bigframes_remote_function = str(dataset_ref.routine(rf_name)) # type: ignore + func.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype( ibis_signature.output_type ) - return node + func.ibis_node = node + return func return wrapper @@ -1121,19 +1190,23 @@ def read_gbq_function( # The name "args" conflicts with the Ibis operator, so we use # non-standard names for the arguments here. - def node(*ignored_args, **ignored_kwargs): + def func(*ignored_args, **ignored_kwargs): f"""Remote function {str(routine_ref)}.""" + # TODO(swast): Construct an ibis client from bigquery_client and + # execute node via a query. # TODO: Move ibis logic to compiler step - node.__name__ = routine_ref.routine_id + func.__name__ = routine_ref.routine_id + node = ibis.udf.scalar.builtin( - node, + func, name=routine_ref.routine_id, schema=f"{routine_ref.project}.{routine_ref.dataset_id}", signature=(ibis_signature.input_types, ibis_signature.output_type), ) - node.bigframes_remote_function = str(routine_ref) # type: ignore - node.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype( # type: ignore + func.bigframes_remote_function = str(routine_ref) # type: ignore + func.output_dtype = bigframes.dtypes.ibis_dtype_to_bigframes_dtype( # type: ignore ibis_signature.output_type ) - return node + func.ibis_node = node # type: ignore + return func diff --git a/bigframes/ml/compose.py b/bigframes/ml/compose.py index 77bfd76bde..abf1a662b9 100644 --- a/bigframes/ml/compose.py +++ b/bigframes/ml/compose.py @@ -28,7 +28,7 @@ from bigframes import constants from bigframes.core import log_adapter -from bigframes.ml import base, core, globals, preprocessing, utils +from bigframes.ml import base, core, globals, impute, preprocessing, utils import bigframes.pandas as bpd _BQML_TRANSFROM_TYPE_MAPPING = types.MappingProxyType( @@ -40,6 +40,7 @@ "ML.BUCKETIZE": preprocessing.KBinsDiscretizer, "ML.QUANTILE_BUCKETIZE": preprocessing.KBinsDiscretizer, "ML.LABEL_ENCODER": preprocessing.LabelEncoder, + "ML.IMPUTER": impute.SimpleImputer, } ) @@ -58,7 +59,7 @@ def __init__( transformers: List[ Tuple[ str, - preprocessing.PreprocessingType, + Union[preprocessing.PreprocessingType, impute.SimpleImputer], Union[str, List[str]], ] ], @@ -73,12 +74,14 @@ def __init__( @property def transformers_( self, - ) -> List[Tuple[str, preprocessing.PreprocessingType, str,]]: + ) -> List[ + Tuple[str, Union[preprocessing.PreprocessingType, impute.SimpleImputer], str] + ]: """The collection of transformers as tuples of (name, transformer, column).""" result: List[ Tuple[ str, - preprocessing.PreprocessingType, + Union[preprocessing.PreprocessingType, impute.SimpleImputer], str, ] ] = [] @@ -107,7 +110,7 @@ def _extract_from_bq_model( transformers: List[ Tuple[ str, - preprocessing.PreprocessingType, + Union[preprocessing.PreprocessingType, impute.SimpleImputer], Union[str, List[str]], ] ] = [] @@ -152,7 +155,9 @@ def camel_to_snake(name): def _merge( self, bq_model: bigquery.Model - ) -> Union[ColumnTransformer, preprocessing.PreprocessingType,]: + ) -> Union[ + ColumnTransformer, Union[preprocessing.PreprocessingType, impute.SimpleImputer] + ]: """Try to merge the column transformer to a simple transformer. Depends on all the columns in bq_model are transformed with the same transformer.""" transformers = self.transformers_ diff --git a/bigframes/ml/impute.py b/bigframes/ml/impute.py new file mode 100644 index 0000000000..d21fcbb1ad --- /dev/null +++ b/bigframes/ml/impute.py @@ -0,0 +1,123 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transformers for missing value imputation. This module is styled after +scikit-learn's preprocessing module: https://scikit-learn.org/stable/modules/impute.html.""" + +from __future__ import annotations + +import typing +from typing import Any, List, Literal, Optional, Tuple, Union + +import bigframes_vendored.sklearn.impute._base + +from bigframes.core import log_adapter +from bigframes.ml import base, core, globals, utils +import bigframes.pandas as bpd + + +@log_adapter.class_logger +class SimpleImputer( + base.Transformer, + bigframes_vendored.sklearn.impute._base.SimpleImputer, +): + + __doc__ = bigframes_vendored.sklearn.impute._base.SimpleImputer.__doc__ + + def __init__( + self, + strategy: Literal["mean", "median", "most_frequent"] = "mean", + ): + self.strategy = strategy + self._bqml_model: Optional[core.BqmlModel] = None + self._bqml_model_factory = globals.bqml_model_factory() + self._base_sql_generator = globals.base_sql_generator() + + # TODO(garrettwu): implement __hash__ + def __eq__(self, other: Any) -> bool: + return ( + type(other) is SimpleImputer + and self.strategy == other.strategy + and self._bqml_model == other._bqml_model + ) + + def _compile_to_sql( + self, + columns: List[str], + X=None, + ) -> List[Tuple[str, str]]: + """Compile this transformer to a list of SQL expressions that can be included in + a BQML TRANSFORM clause + + Args: + columns: + A list of column names to transform. + X: + The Dataframe with training data. + + Returns: a list of tuples of (sql_expression, output_name)""" + return [ + ( + self._base_sql_generator.ml_imputer( + column, self.strategy, f"imputer_{column}" + ), + f"imputer_{column}", + ) + for column in columns + ] + + @classmethod + def _parse_from_sql(cls, sql: str) -> tuple[SimpleImputer, str]: + """Parse SQL to tuple(SimpleImputer, column_label). + + Args: + sql: SQL string of format "ML.IMPUTER({col_label}, {strategy}) OVER()" + + Returns: + tuple(SimpleImputer, column_label)""" + s = sql[sql.find("(") + 1 : sql.find(")")] + col_label, strategy = s.split(", ") + return cls(strategy[1:-1]), col_label # type: ignore[arg-type] + + def fit( + self, + X: Union[bpd.DataFrame, bpd.Series], + y=None, # ignored + ) -> SimpleImputer: + (X,) = utils.convert_to_dataframe(X) + + compiled_transforms = self._compile_to_sql(X.columns.tolist(), X) + transform_sqls = [transform_sql for transform_sql, _ in compiled_transforms] + + self._bqml_model = self._bqml_model_factory.create_model( + X, + options={"model_type": "transform_only"}, + transforms=transform_sqls, + ) + + # The schema of TRANSFORM output is not available in the model API, so save it during fitting + self._output_names = [name for _, name in compiled_transforms] + return self + + def transform(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame: + if not self._bqml_model: + raise RuntimeError("Must be fitted before transform") + + (X,) = utils.convert_to_dataframe(X) + + df = self._bqml_model.transform(X) + return typing.cast( + bpd.DataFrame, + df[self._output_names], + ) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 77dc1d2b0f..7fa0e236eb 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -46,6 +46,14 @@ ) _GEMINI_PRO_ENDPOINT = "gemini-pro" +_GEMINI_1P5_PRO_PREVIEW_ENDPOINT = "gemini-1.5-pro-preview-0514" +_GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT = "gemini-1.5-flash-preview-0514" +_GEMINI_ENDPOINTS = ( + _GEMINI_PRO_ENDPOINT, + _GEMINI_1P5_PRO_PREVIEW_ENDPOINT, + _GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT, +) + _ML_GENERATE_TEXT_STATUS = "ml_generate_text_status" _ML_EMBED_TEXT_STATUS = "ml_embed_text_status" @@ -547,13 +555,16 @@ def to_gbq( class GeminiTextGenerator(base.BaseEstimator): """Gemini text generator LLM model. - .. note:: - This product or feature is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the - Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is" - and might have limited support. For more information, see the launch stage descriptions - (https://cloud.google.com/products#product-launch-stages). - Args: + model_name (str, Default to "gemini-pro"): + The model for natural language tasks. Accepted values are "gemini-pro", "gemini-1.5-pro-preview-0514" and "gemini-1.5-flash-preview-0514". Default to "gemini-pro". + + .. note:: + "gemini-1.5-pro-preview-0514" and "gemini-1.5-flash-preview-0514" is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the + Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is" + and might have limited support. For more information, see the launch stage descriptions + (https://cloud.google.com/products#product-launch-stages). + session (bigframes.Session or None): BQ session to create the model. If None, use the global default session. connection_name (str or None): @@ -565,9 +576,13 @@ class GeminiTextGenerator(base.BaseEstimator): def __init__( self, *, + model_name: Literal[ + "gemini-pro", "gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514" + ] = "gemini-pro", session: Optional[bigframes.Session] = None, connection_name: Optional[str] = None, ): + self.model_name = model_name self.session = session or bpd.get_global_session() self._bq_connection_manager = self.session.bqconnectionmanager @@ -601,7 +616,12 @@ def _create_bqml_model(self): iam_role="aiplatform.user", ) - options = {"endpoint": _GEMINI_PRO_ENDPOINT} + if self.model_name not in _GEMINI_ENDPOINTS: + raise ValueError( + f"Model name {self.model_name} is not supported. We only support {', '.join(_GEMINI_ENDPOINTS)}." + ) + + options = {"endpoint": self.model_name} return self._bqml_model_factory.create_remote_model( session=self.session, connection_name=self.connection_name, options=options @@ -613,12 +633,17 @@ def _from_bq( ) -> GeminiTextGenerator: assert bq_model.model_type == "MODEL_TYPE_UNSPECIFIED" assert "remoteModelInfo" in bq_model._properties + assert "endpoint" in bq_model._properties["remoteModelInfo"] assert "connection" in bq_model._properties["remoteModelInfo"] # Parse the remote model endpoint + bqml_endpoint = bq_model._properties["remoteModelInfo"]["endpoint"] model_connection = bq_model._properties["remoteModelInfo"]["connection"] + model_endpoint = bqml_endpoint.split("/")[-1] - model = cls(session=session, connection_name=model_connection) + model = cls( + model_name=model_endpoint, session=session, connection_name=model_connection + ) model._bqml_model = core.BqmlModel(session, bq_model) return model diff --git a/bigframes/ml/loader.py b/bigframes/ml/loader.py index 916949077f..66f207929a 100644 --- a/bigframes/ml/loader.py +++ b/bigframes/ml/loader.py @@ -29,6 +29,7 @@ ensemble, forecasting, imported, + impute, linear_model, llm, pipeline, @@ -60,6 +61,8 @@ llm._EMBEDDING_GENERATOR_GECKO_ENDPOINT: llm.PaLM2TextEmbeddingGenerator, llm._EMBEDDING_GENERATOR_GECKO_MULTILINGUAL_ENDPOINT: llm.PaLM2TextEmbeddingGenerator, llm._GEMINI_PRO_ENDPOINT: llm.GeminiTextGenerator, + llm._GEMINI_1P5_PRO_PREVIEW_ENDPOINT: llm.GeminiTextGenerator, + llm._GEMINI_1P5_PRO_FLASH_PREVIEW_ENDPOINT: llm.GeminiTextGenerator, } ) @@ -84,6 +87,7 @@ def from_bq( pipeline.Pipeline, compose.ColumnTransformer, preprocessing.PreprocessingType, + impute.SimpleImputer, ]: """Load a BQML model to BigQuery DataFrames ML. diff --git a/bigframes/ml/pipeline.py b/bigframes/ml/pipeline.py index 61b5a22da9..03e5688453 100644 --- a/bigframes/ml/pipeline.py +++ b/bigframes/ml/pipeline.py @@ -26,7 +26,15 @@ import bigframes import bigframes.constants as constants from bigframes.core import log_adapter -from bigframes.ml import base, compose, forecasting, loader, preprocessing, utils +from bigframes.ml import ( + base, + compose, + forecasting, + impute, + loader, + preprocessing, + utils, +) import bigframes.pandas as bpd @@ -56,6 +64,7 @@ def __init__(self, steps: List[Tuple[str, base.BaseEstimator]]): preprocessing.MinMaxScaler, preprocessing.KBinsDiscretizer, preprocessing.LabelEncoder, + impute.SimpleImputer, ), ): self._transform = transform diff --git a/bigframes/ml/preprocessing.py b/bigframes/ml/preprocessing.py index 954d5adff0..4b1a3fb7b7 100644 --- a/bigframes/ml/preprocessing.py +++ b/bigframes/ml/preprocessing.py @@ -305,6 +305,7 @@ def __eq__(self, other: Any) -> bool: return ( type(other) is KBinsDiscretizer and self.n_bins == other.n_bins + and self.strategy == other.strategy and self._bqml_model == other._bqml_model ) diff --git a/bigframes/ml/sql.py b/bigframes/ml/sql.py index b701ab301c..f060584a11 100644 --- a/bigframes/ml/sql.py +++ b/bigframes/ml/sql.py @@ -103,6 +103,15 @@ def ml_min_max_scaler(self, numeric_expr_sql: str, name: str) -> str: """Encode ML.MIN_MAX_SCALER for BQML""" return f"""ML.MIN_MAX_SCALER({numeric_expr_sql}) OVER() AS {name}""" + def ml_imputer( + self, + expr_sql: str, + strategy: str, + name: str, + ) -> str: + """Encode ML.IMPUTER for BQML""" + return f"""ML.IMPUTER({expr_sql}, '{strategy}') OVER() AS {name}""" + def ml_bucketize( self, numeric_expr_sql: str, diff --git a/bigframes/operations/__init__.py b/bigframes/operations/__init__.py index fe9fe6df20..42f83913ee 100644 --- a/bigframes/operations/__init__.py +++ b/bigframes/operations/__init__.py @@ -186,7 +186,7 @@ def create_binary_op( dtypes.is_binary_like, description="binary-like", ), -) # numeric +) isnull_op = create_unary_op( name="isnull", type_signature=op_typing.FixedOutputType( @@ -311,6 +311,8 @@ def create_binary_op( floor_op = create_unary_op(name="floor", type_signature=op_typing.UNARY_REAL_NUMERIC) ceil_op = create_unary_op(name="ceil", type_signature=op_typing.UNARY_REAL_NUMERIC) abs_op = create_unary_op(name="abs", type_signature=op_typing.UNARY_NUMERIC) +pos_op = create_unary_op(name="pos", type_signature=op_typing.UNARY_NUMERIC) +neg_op = create_unary_op(name="neg", type_signature=op_typing.UNARY_NUMERIC) exp_op = create_unary_op(name="exp", type_signature=op_typing.UNARY_REAL_NUMERIC) expm1_op = create_unary_op(name="expm1", type_signature=op_typing.UNARY_REAL_NUMERIC) ln_op = create_unary_op(name="log", type_signature=op_typing.UNARY_REAL_NUMERIC) @@ -578,6 +580,19 @@ def output_type(self, *input_types): return input_types[0] +## Array Ops +@dataclasses.dataclass(frozen=True) +class ArrayToStringOp(UnaryOp): + name: typing.ClassVar[str] = "array_to_string" + delimiter: str + + def output_type(self, *input_types): + input_type = input_types[0] + if not dtypes.is_array_string_like(input_type): + raise TypeError("Input type must be an array of string type.") + return dtypes.STRING_DTYPE + + # Binary Ops fillna_op = create_binary_op(name="fillna", type_signature=op_typing.COERCE) maximum_op = create_binary_op(name="maximum", type_signature=op_typing.COERCE) @@ -650,6 +665,7 @@ def output_type(self, *input_types): # Logical Ops and_op = create_binary_op(name="and", type_signature=op_typing.LOGICAL) or_op = create_binary_op(name="or", type_signature=op_typing.LOGICAL) +xor_op = create_binary_op(name="xor", type_signature=op_typing.LOGICAL) ## Comparison Ops eq_op = create_binary_op(name="eq", type_signature=op_typing.COMPARISON) diff --git a/bigframes/operations/aggregations.py b/bigframes/operations/aggregations.py index 3b5310554b..783abfd788 100644 --- a/bigframes/operations/aggregations.py +++ b/bigframes/operations/aggregations.py @@ -47,6 +47,13 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT ... +@dataclasses.dataclass(frozen=True) +class NullaryWindowOp(WindowOp): + @property + def arguments(self) -> int: + return 0 + + @dataclasses.dataclass(frozen=True) class UnaryWindowOp(WindowOp): @property @@ -72,6 +79,13 @@ def arguments(self) -> int: ... +@dataclasses.dataclass(frozen=True) +class NullaryAggregateOp(AggregateOp, NullaryWindowOp): + @property + def arguments(self) -> int: + return 0 + + @dataclasses.dataclass(frozen=True) class UnaryAggregateOp(AggregateOp, UnaryWindowOp): @property @@ -86,6 +100,14 @@ def arguments(self) -> int: return 2 +@dataclasses.dataclass(frozen=True) +class SizeOp(NullaryAggregateOp): + name: ClassVar[str] = "size" + + def output_type(self, *input_types: dtypes.ExpressionType): + return dtypes.INT_DTYPE + + @dataclasses.dataclass(frozen=True) class SumOp(UnaryAggregateOp): name: ClassVar[str] = "sum" @@ -446,6 +468,7 @@ def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionT ) +size_op = SizeOp() sum_op = SumOp() mean_op = MeanOp() median_op = MedianOp() diff --git a/bigframes/operations/base.py b/bigframes/operations/base.py index 49ef7f76ee..f339345971 100644 --- a/bigframes/operations/base.py +++ b/bigframes/operations/base.py @@ -70,27 +70,35 @@ def __init__( ) block: typing.Optional[blocks.Block] = None + if (name is not None) and not isinstance(name, typing.Hashable): + raise ValueError( + f"BigQuery DataFrames only supports hashable series names. {constants.FEEDBACK_LINK}" + ) if copy is not None and not copy: raise ValueError( f"Series constructor only supports copy=True. {constants.FEEDBACK_LINK}" ) if isinstance(data, blocks.Block): + # Constructing from block is for internal use only - shouldn't use parameters, block encompasses all state assert len(data.value_columns) == 1 assert len(data.column_labels) == 1 assert index is None + assert name is None + assert dtype is None block = data # interpret these cases as both index and data - elif ( - isinstance(data, SeriesMethods) - or isinstance(data, pd.Series) - or pd.api.types.is_dict_like(data) - ): - if isinstance(data, pd.Series): - data = read_pandas_func(data) - elif pd.api.types.is_dict_like(data): - data = read_pandas_func(pd.Series(data, dtype=dtype)) # type: ignore - dtype = None + elif isinstance(data, bigframes.pandas.Series) or pd.api.types.is_dict_like( + data + ): # includes pd.Series + if isinstance(data, bigframes.pandas.Series): + data = data.copy() + if name is not None: + data.name = name + if dtype is not None: + data = data.astype(dtype) + else: # local dict-like data + data = read_pandas_func(pd.Series(data, name=name, dtype=dtype)) # type: ignore data_block = data._block if index is not None: # reindex @@ -103,10 +111,8 @@ def __init__( # list-like data that will get default index elif isinstance(data, indexes.Index) or pd.api.types.is_list_like(data): - data = indexes.Index(data, dtype=dtype, session=session) - dtype = ( - None # set to none as it has already been applied, avoid re-cast later - ) + data = indexes.Index(data, dtype=dtype, name=name, session=session) + # set to none as it has already been applied, avoid re-cast later if data.nlevels != 1: raise NotImplementedError("Cannot interpret multi-index as Series.") # Reset index to promote index columns to value columns, set default index @@ -135,20 +141,9 @@ def __init__( dtype=bigframes.dtypes.INT_DTYPE, ) block, _ = bf_index._block.create_constant(data, dtype) - dtype = None block = block.with_column_labels([name]) assert block is not None - if name: - if not isinstance(name, typing.Hashable): - raise ValueError( - f"BigQuery DataFrames only supports hashable series names. {constants.FEEDBACK_LINK}" - ) - block = block.with_column_labels([name]) - if dtype: - block = block.multi_apply_unary_op( - block.value_columns, ops.AsTypeOp(to_type=dtype) - ) self._block: blocks.Block = block @property diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 3628ecf67b..603f6678a5 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -399,6 +399,9 @@ def _set_default_session_location_if_possible(query): bqclient = clients_provider.bqclient if bigframes.session._io.bigquery.is_query(query): + # Intentionally run outside of the session so that we can detect the + # location before creating the session. Since it's a dry_run, labels + # aren't necessary. job = bqclient.query(query, bigquery.QueryJobConfig(dry_run=True)) options.bigquery.location = job.location else: @@ -649,8 +652,8 @@ def read_parquet( def remote_function( - input_types: Union[type, Sequence[type]], - output_type: type, + input_types: Union[None, type, Sequence[type]] = None, + output_type: Optional[type] = None, dataset: Optional[str] = None, bigquery_connection: Optional[str] = None, reuse: bool = True, @@ -773,7 +776,10 @@ def clean_up_by_session_id( dataset = session._anonymous_dataset else: dataset = bigframes.session._io.bigquery.create_bq_dataset_reference( - client, location=location, project=project + client, + location=location, + project=project, + api_name="clean_up_by_session_id", ) bigframes.session._io.bigquery.delete_tables_matching_session_id( diff --git a/bigframes/series.py b/bigframes/series.py index 4595164e80..367301f08e 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -72,6 +72,7 @@ class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Ser def __init__(self, *args, **kwargs): self._query_job: Optional[bigquery.QueryJob] = None super().__init__(*args, **kwargs) + self._block.session._register_object(self) @property def dt(self) -> dt.DatetimeMethods: @@ -667,6 +668,13 @@ def __or__(self, other: bool | int | Series) -> Series: __ror__ = __or__ + def __xor__(self, other: bool | int | Series) -> Series: + return self._apply_binary_op(other, ops.xor_op) + + __or__.__doc__ = inspect.getdoc(vendored_pandas_series.Series.__xor__) + + __rxor__ = __xor__ + def __add__(self, other: float | int | Series) -> Series: return self.add(other) @@ -1035,6 +1043,12 @@ def __invert__(self) -> Series: __invert__.__doc__ = inspect.getdoc(vendored_pandas_series.Series.__invert__) + def __pos__(self) -> Series: + return self._apply_unary_op(ops.pos_op) + + def __neg__(self) -> Series: + return self._apply_unary_op(ops.neg_op) + def eq(self, other: object) -> Series: # TODO: enforce stricter alignment return self._apply_binary_op(other, ops.eq_op) @@ -1777,7 +1791,7 @@ def cache(self): return self._cached(force=True) def _cached(self, *, force: bool = True) -> Series: - self._set_block(self._block.cached(force=force)) + self._block.cached(force=force) return self def _optimize_query_complexity(self): @@ -1785,8 +1799,7 @@ def _optimize_query_complexity(self): May generate many queries and take substantial time to execute. """ # TODO: Move all this to session - new_expr = self._block.session._simplify_with_caching(self._block.expr) - self._set_block(self._block.swap_array_expr(new_expr)) + self._block.session._simplify_with_caching(self._block.expr) def _is_list_like(obj: typing.Any) -> typing_extensions.TypeGuard[typing.Sequence]: diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index ccdc3c5eeb..2919b2d77f 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -16,9 +16,12 @@ from __future__ import annotations +import collections.abc import copy import datetime +import itertools import logging +import math import os import secrets import typing @@ -40,9 +43,8 @@ ) import uuid import warnings +import weakref -# Even though the ibis.backends.bigquery import is unused, it's needed -# to register new and replacement ops with the Ibis BigQuery backend. import bigframes_vendored.ibis.backends.bigquery # noqa import bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq import bigframes_vendored.pandas.io.parquet as third_party_pandas_parquet @@ -84,9 +86,14 @@ import bigframes.core.nodes as nodes from bigframes.core.ordering import IntegerEncoding import bigframes.core.ordering as order +import bigframes.core.schema as schemata import bigframes.core.tree_properties as traversals import bigframes.core.tree_properties as tree_properties import bigframes.core.utils as utils + +# Even though the ibis.backends.bigquery import is unused, it's needed +# to register new and replacement ops with the Ibis BigQuery backend. +import bigframes.dataframe import bigframes.dtypes import bigframes.exceptions import bigframes.formatting_helpers as formatting_helpers @@ -230,7 +237,9 @@ def __init__( self._anonymous_dataset = ( bigframes.session._io.bigquery.create_bq_dataset_reference( - self.bqclient, location=self._location + self.bqclient, + location=self._location, + api_name="session-__init__", ) ) @@ -269,6 +278,24 @@ def __init__( self._table_ids: List[str] = [] # store table ids and delete them when the session is closed + self._objects: list[ + weakref.ReferenceType[ + Union[ + bigframes.core.indexes.Index, + bigframes.series.Series, + dataframe.DataFrame, + ] + ] + ] = [] + self._cached_executions: weakref.WeakKeyDictionary[ + nodes.BigFrameNode, nodes.BigFrameNode + ] = weakref.WeakKeyDictionary() + + # performance logging + self._bytes_processed_sum = 0 + self._slot_millis_sum = 0 + self._execution_count = 0 + @property def bqclient(self): return self._clients_provider.bqclient @@ -303,10 +330,43 @@ def bqconnectionmanager(self): def session_id(self): return self._session_id + @property + def objects( + self, + ) -> collections.abc.Set[ + Union[ + bigframes.core.indexes.Index, bigframes.series.Series, dataframe.DataFrame + ] + ]: + # Create a set with strong references, be careful not to hold onto this needlessly, as will prevent garbage collection. + return set(i() for i in self._objects if i() is not None) # type: ignore + @property def _project(self): return self.bqclient.project + @property + def bytes_processed_sum(self): + """The sum of all bytes processed by bigquery jobs using this session.""" + return self._bytes_processed_sum + + @property + def slot_millis_sum(self): + """The sum of all slot time used by bigquery jobs in this session.""" + return self._slot_millis_sum + + def _add_bytes_processed(self, amount: int): + """Increment bytes_processed_sum by amount.""" + self._bytes_processed_sum += amount + + def _add_slot_millis(self, amount: int): + """Increment slot_millis_sum by amount.""" + self._slot_millis_sum += amount + + def _add_execution(self, amount: int = 1): + """Increment slot_millis_sum by amount.""" + self._execution_count += amount + def __hash__(self): # Stable hash needed to use in expression tree return hash(str(self._anonymous_dataset)) @@ -371,6 +431,14 @@ def read_gbq( filters=filters, ) + def _register_object( + self, + object: Union[ + bigframes.core.indexes.Index, bigframes.series.Series, dataframe.DataFrame + ], + ): + self._objects.append(weakref.ref(object)) + def _query_to_destination( self, query: str, @@ -379,13 +447,16 @@ def _query_to_destination( configuration: dict = {"query": {"useQueryCache": True}}, do_clustering=True, ) -> Tuple[Optional[bigquery.TableReference], bigquery.QueryJob]: + self._add_execution(1) # If a dry_run indicates this is not a query type job, then don't # bother trying to do a CREATE TEMP TABLE ... AS SELECT ... statement. dry_run_config = bigquery.QueryJobConfig() dry_run_config.dry_run = True - _, dry_run_job = self._start_query(query, job_config=dry_run_config) + _, dry_run_job = self._start_query( + query, job_config=dry_run_config, api_name=api_name + ) if dry_run_job.statement_type != "SELECT": - _, query_job = self._start_query(query) + _, query_job = self._start_query(query, api_name=api_name) return query_job.destination, query_job # Create a table to workaround BigQuery 10 GB query results limit. See: @@ -414,7 +485,6 @@ def _query_to_destination( bigquery.QueryJobConfig, bigquery.QueryJobConfig.from_api_repr(configuration), ) - job_config.labels["bigframes-api"] = api_name job_config.destination = temp_table try: @@ -422,7 +492,10 @@ def _query_to_destination( # limit. See: internal issue 303057336. job_config.labels["error_caught"] = "true" _, query_job = self._start_query( - query, job_config=job_config, timeout=timeout + query, + job_config=job_config, + timeout=timeout, + api_name=api_name, ) return query_job.destination, query_job except google.api_core.exceptions.BadRequest: @@ -430,7 +503,7 @@ def _query_to_destination( # tables as the destination. For example, if the query has a # top-level ORDER BY, this conflicts with our ability to cluster # the table by the index column(s). - _, query_job = self._start_query(query, timeout=timeout) + _, query_job = self._start_query(query, timeout=timeout, api_name=api_name) return query_job.destination, query_job def read_gbq_query( @@ -556,11 +629,11 @@ def _read_gbq_query( if len(filters) != 0 or max_results is not None: # TODO(b/338111344): If we are running a query anyway, we might as # well generate ROW_NUMBER() at the same time. + all_columns = itertools.chain(index_cols, columns) if columns else () query = bf_io_bigquery.to_query( query, - index_cols, - columns, - filters, + all_columns, + bf_io_bigquery.compile_filters(filters) if filters else None, max_results=max_results, # We're executing the query, so we don't need time travel for # determinism. @@ -676,7 +749,7 @@ def _read_gbq_table( # Fetch table metadata and validate # --------------------------------- - (time_travel_timestamp, table,) = bf_read_gbq_table.get_table_metadata( + time_travel_timestamp, table = bf_read_gbq_table.get_table_metadata( self.bqclient, table_ref=table_ref, api_name=api_name, @@ -733,11 +806,13 @@ def _read_gbq_table( ): # TODO(b/338111344): If we are running a query anyway, we might as # well generate ROW_NUMBER() at the same time. + all_columns = itertools.chain(index_cols, columns) if columns else () query = bf_io_bigquery.to_query( query, - index_cols=index_cols, - columns=columns, - filters=filters, + columns=all_columns, + sql_predicate=bf_io_bigquery.compile_filters(filters) + if filters + else None, max_results=max_results, # We're executing the query, so we don't need time travel for # determinism. @@ -753,21 +828,24 @@ def _read_gbq_table( ) # ----------------------------------------- - # Create Ibis table expression and validate + # Validate table access and features # ----------------------------------------- # Use a time travel to make sure the DataFrame is deterministic, even # if the underlying table changes. - # TODO(b/340540991): If a dry run query fails with time travel but + + # If a dry run query fails with time travel but # succeeds without it, omit the time travel clause and raise a warning # about potential non-determinism if the underlying tables are modified. - table_expression = bf_read_gbq_table.get_ibis_time_travel_table( - ibis_client=self.ibis_client, - table_ref=table_ref, - index_cols=index_cols, - columns=columns, - filters=filters, - time_travel_timestamp=time_travel_timestamp, + filter_str = bf_io_bigquery.compile_filters(filters) if filters else None + all_columns = ( + () + if len(columns) == 0 + else (*columns, *[col for col in index_cols if col not in columns]) + ) + + supports_snapshot = bf_read_gbq_table.validate_table( + self.bqclient, table_ref, all_columns, time_travel_timestamp, filter_str ) # ---------------------------- @@ -785,20 +863,17 @@ def _read_gbq_table( index_cols=index_cols, api_name=api_name, ) - - if is_index_unique: - array_value = bf_read_gbq_table.to_array_value_with_total_ordering( - session=self, - table_expression=table_expression, - total_ordering_cols=index_cols, - ) - else: - # Note: Even though we're adding a default ordering here, that's - # just so we have a deterministic total ordering. If the user - # specified a non-unique index, we still sort by that later. - array_value = bf_read_gbq_table.to_array_value_with_default_ordering( - session=self, table=table_expression, table_rows=table.num_rows - ) + schema = schemata.ArraySchema.from_bq_table(table) + if columns: + schema = schema.select(index_cols + columns) + array_value = core.ArrayValue.from_table( + table, + schema=schema, + predicate=filter_str, + at_time=time_travel_timestamp if supports_snapshot else None, + primary_key=index_cols if is_index_unique else (), + session=self, + ) # ---------------------------------------------------- # Create Default Sequential Index if still have no index @@ -966,7 +1041,7 @@ def read_pandas( # Try to handle non-dataframe pandas objects as well if isinstance(pandas_dataframe, pandas.Series): bf_df = self._read_pandas(pandas.DataFrame(pandas_dataframe), "read_pandas") - bf_series = typing.cast(series.Series, bf_df[bf_df.columns[0]]) + bf_series = series.Series(bf_df._block) # wrapping into df can set name to 0 so reset to original object name bf_series.name = pandas_dataframe.name return bf_series @@ -1011,9 +1086,8 @@ def _read_pandas_inline( return None try: - inline_df = dataframe.DataFrame( - blocks.Block.from_local(pandas_dataframe, self) - ) + local_block = blocks.Block.from_local(pandas_dataframe, self) + inline_df = dataframe.DataFrame(local_block) except pa.ArrowInvalid as e: raise pa.ArrowInvalid( f"Could not convert with a BigQuery type: `{e}`. " @@ -1458,8 +1532,8 @@ def _ibis_to_temp_table( def remote_function( self, - input_types: Union[type, Sequence[type]], - output_type: type, + input_types: Union[None, type, Sequence[type]] = None, + output_type: Optional[type] = None, dataset: Optional[str] = None, bigquery_connection: Optional[str] = None, reuse: bool = True, @@ -1714,12 +1788,6 @@ def _prepare_query_job_config( bigframes.options.compute.maximum_bytes_billed ) - current_labels = job_config.labels if job_config.labels else {} - for key, value in bigframes.options.compute.extra_query_labels.items(): - if key not in current_labels: - current_labels[key] = value - job_config.labels = current_labels - if self._bq_kms_key_name: job_config.destination_encryption_configuration = ( bigquery.EncryptionConfiguration(kms_key_name=self._bq_kms_key_name) @@ -1755,13 +1823,19 @@ def _start_query( job_config: Optional[bigquery.job.QueryJobConfig] = None, max_results: Optional[int] = None, timeout: Optional[float] = None, + api_name: Optional[str] = None, ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: """ Starts BigQuery query job and waits for results. """ job_config = self._prepare_query_job_config(job_config) return bigframes.session._io.bigquery.start_query_with_client( - self.bqclient, sql, job_config, max_results, timeout + self, + sql, + job_config, + max_results, + timeout, + api_name=api_name, ) def _start_query_ml_ddl( @@ -1780,12 +1854,12 @@ def _start_query_ml_ddl( job_config.destination_encryption_configuration = None return bigframes.session._io.bigquery.start_query_with_client( - self.bqclient, sql, job_config + self, sql, job_config ) def _cache_with_cluster_cols( self, array_value: core.ArrayValue, cluster_cols: typing.Sequence[str] - ) -> core.ArrayValue: + ): """Executes the query and uses the resulting table to rewrite future executions.""" # TODO: Use this for all executions? Problem is that caching materializes extra # ordering columns @@ -1807,16 +1881,16 @@ def _cache_with_cluster_cols( table_expression[column] for column in compiled_value._hidden_ordering_column_names ] - # TODO: Instead, keep session-wide map of cached results and automatically reuse - return core.ArrayValue.from_ibis( + cached_replacement = core.ArrayValue.from_ibis( self, table_expression, columns=new_columns, hidden_ordering_columns=new_hidden_columns, ordering=compiled_value._ordering, - ) + ).node + self._cached_executions[array_value.node] = cached_replacement - def _cache_with_offsets(self, array_value: core.ArrayValue) -> core.ArrayValue: + def _cache_with_offsets(self, array_value: core.ArrayValue): """Executes the query and uses the resulting table to rewrite future executions.""" # TODO: Use this for all executions? Problem is that caching materializes extra # ordering columns @@ -1835,60 +1909,50 @@ def _cache_with_offsets(self, array_value: core.ArrayValue) -> core.ArrayValue: ) new_columns = [table_expression[column] for column in compiled_value.column_ids] new_hidden_columns = [table_expression["bigframes_offsets"]] - # TODO: Instead, keep session-wide map of cached results and automatically reuse - return core.ArrayValue.from_ibis( + cached_replacement = core.ArrayValue.from_ibis( self, table_expression, columns=new_columns, hidden_ordering_columns=new_hidden_columns, ordering=order.ExpressionOrdering.from_offset_col("bigframes_offsets"), - ) + ).node + self._cached_executions[array_value.node] = cached_replacement - def _simplify_with_caching(self, array_value: core.ArrayValue) -> core.ArrayValue: + def _simplify_with_caching(self, array_value: core.ArrayValue): """Attempts to handle the complexity by caching duplicated subtrees and breaking the query into pieces.""" + # Apply existing caching first if not bigframes.options.compute.enable_multi_query_execution: - return array_value - node = array_value.node - if node.planning_complexity < QUERY_COMPLEXITY_LIMIT: - return array_value + return for _ in range(MAX_SUBTREE_FACTORINGS): - updated = self._cache_most_complex_subtree(node) - if updated is None: - return core.ArrayValue(node) - else: - node = updated + node_with_cache = self._with_cached_executions(array_value.node) + if node_with_cache.planning_complexity < QUERY_COMPLEXITY_LIMIT: + return - return core.ArrayValue(node) + did_cache = self._cache_most_complex_subtree(array_value.node) + if not did_cache: + return - def _cache_most_complex_subtree( - self, node: nodes.BigFrameNode - ) -> Optional[nodes.BigFrameNode]: + def _cache_most_complex_subtree(self, node: nodes.BigFrameNode) -> bool: # TODO: If query fails, retry with lower complexity limit - valid_candidates = traversals.count_complex_nodes( + selection = traversals.select_cache_target( node, min_complexity=(QUERY_COMPLEXITY_LIMIT / 500), max_complexity=QUERY_COMPLEXITY_LIMIT, - ).items() - # Heuristic: subtree_compleixty * (copies of subtree)^2 - best_candidate = max( - valid_candidates, - key=lambda i: i[0].planning_complexity + (i[1] ** 2), - default=None, + cache=dict(self._cached_executions), + # Heuristic: subtree_compleixty * (copies of subtree)^2 + heuristic=lambda complexity, count: math.log(complexity) + + 2 * math.log(count), ) - - if best_candidate is None: + if selection is None: # No good subtrees to cache, just return original tree - return None + return False - # TODO: Add clustering columns based on access patterns - materialized = self._cache_with_cluster_cols( - core.ArrayValue(best_candidate[0]), [] - ).node + self._cache_with_cluster_cols(core.ArrayValue(selection), []) + return True - return traversals.replace_nodes( - node, to_replace=best_candidate[0], replacemenet=materialized - ) + def _with_cached_executions(self, node: nodes.BigFrameNode) -> nodes.BigFrameNode: + return traversals.replace_nodes(node, (dict(self._cached_executions))) def _is_trivially_executable(self, array_value: core.ArrayValue): """ @@ -1897,7 +1961,9 @@ def _is_trivially_executable(self, array_value: core.ArrayValue): """ # Once rewriting is available, will want to rewrite before # evaluating execution cost. - return traversals.is_trivially_executable(array_value.node) + return traversals.is_trivially_executable( + self._with_cached_executions(array_value.node) + ) def _execute( self, @@ -1908,6 +1974,8 @@ def _execute( dry_run=False, col_id_overrides: Mapping[str, str] = {}, ) -> tuple[bigquery.table.RowIterator, bigquery.QueryJob]: + if not dry_run: + self._add_execution(1) sql = self._to_sql( array_value, sorted=sorted, col_id_overrides=col_id_overrides ) # type:ignore @@ -1915,6 +1983,9 @@ def _execute( job_config = bigquery.QueryJobConfig(dry_run=dry_run) else: job_config.dry_run = dry_run + + # TODO(swast): plumb through the api_name of the user-facing api that + # caused this query. return self._start_query( sql=sql, job_config=job_config, @@ -1924,9 +1995,12 @@ def _peek( self, array_value: core.ArrayValue, n_rows: int ) -> tuple[bigquery.table.RowIterator, bigquery.QueryJob]: """A 'peek' efficiently accesses a small number of rows in the dataframe.""" - if not tree_properties.peekable(array_value.node): + if not tree_properties.peekable(self._with_cached_executions(array_value.node)): warnings.warn("Peeking this value cannot be done efficiently.") sql = self._compile_unordered(array_value).peek_sql(n_rows) + + # TODO(swast): plumb through the api_name of the user-facing api that + # caused this query. return self._start_query( sql=sql, ) @@ -1951,12 +2025,16 @@ def _to_sql( def _compile_ordered( self, array_value: core.ArrayValue ) -> bigframes.core.compile.OrderedIR: - return bigframes.core.compile.compile_ordered_ir(array_value.node) + return bigframes.core.compile.compile_ordered_ir( + self._with_cached_executions(array_value.node) + ) def _compile_unordered( self, array_value: core.ArrayValue ) -> bigframes.core.compile.UnorderedIR: - return bigframes.core.compile.compile_unordered_ir(array_value.node) + return bigframes.core.compile.compile_unordered_ir( + self._with_cached_executions(array_value.node) + ) def _get_table_size(self, destination_table): table = self.bqclient.get_table(destination_table) diff --git a/bigframes/session/_io/bigquery/__init__.py b/bigframes/session/_io/bigquery/__init__.py index 28eed47965..f26ca26c2a 100644 --- a/bigframes/session/_io/bigquery/__init__.py +++ b/bigframes/session/_io/bigquery/__init__.py @@ -28,6 +28,7 @@ import bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq import google.api_core.exceptions import google.cloud.bigquery as bigquery +import google.cloud.bigquery.table import bigframes from bigframes.core import log_adapter @@ -40,19 +41,34 @@ # will be limited to this many tables LOGGING_NAME_ENV_VAR = "BIGFRAMES_PERFORMANCE_LOG_NAME" +CHECK_DRIVE_PERMISSIONS = "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions." def create_job_configs_labels( job_configs_labels: Optional[Dict[str, str]], api_methods: typing.List[str], + api_name: Optional[str] = None, ) -> Dict[str, str]: if job_configs_labels is None: job_configs_labels = {} - if api_methods: + # If the user has labels they wish to set, make sure we set those first so + # they are preserved. + for key, value in bigframes.options.compute.extra_query_labels.items(): + job_configs_labels[key] = value + + if api_name is not None: + job_configs_labels["bigframes-api"] = api_name + + if api_methods and "bigframes-api" not in job_configs_labels: job_configs_labels["bigframes-api"] = api_methods[0] del api_methods[0] + # Make sure we always populate bigframes-api with _something_, even if we + # have a code path which doesn't populate the list of api_methods. See + # internal issue 336521938. + job_configs_labels.setdefault("bigframes-api", "unknown") + labels = list( itertools.chain( job_configs_labels.keys(), @@ -193,27 +209,34 @@ def format_option(key: str, value: Union[bool, str]) -> str: return f"{key}={repr(value)}" +def add_labels(job_config, api_name: Optional[str] = None): + api_methods = log_adapter.get_and_reset_api_methods(dry_run=job_config.dry_run) + job_config.labels = create_job_configs_labels( + job_configs_labels=job_config.labels, + api_methods=api_methods, + api_name=api_name, + ) + + def start_query_with_client( - bq_client: bigquery.Client, + session: bigframes.session.Session, sql: str, job_config: bigquery.job.QueryJobConfig, max_results: Optional[int] = None, timeout: Optional[float] = None, + api_name: Optional[str] = None, ) -> Tuple[bigquery.table.RowIterator, bigquery.QueryJob]: """ Starts query job and waits for results. """ - if not job_config.dry_run: - api_methods = log_adapter.get_and_reset_api_methods() - job_config.labels = create_job_configs_labels( - job_configs_labels=job_config.labels, api_methods=api_methods - ) + bq_client: bigquery.Client = session.bqclient + add_labels(job_config, api_name=api_name) try: query_job = bq_client.query(sql, job_config=job_config, timeout=timeout) except google.api_core.exceptions.Forbidden as ex: if "Drive credentials" in ex.message: - ex.message += "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions." + ex.message += CHECK_DRIVE_PERMISSIONS raise opts = bigframes.options.display @@ -224,14 +247,41 @@ def start_query_with_client( else: results_iterator = query_job.result(max_results=max_results) - if LOGGING_NAME_ENV_VAR in os.environ: - # when running notebooks via pytest nbmake - pytest_log_job(query_job) + stats = get_performance_stats(query_job) + if stats is not None: + bytes_processed, slot_millis = stats + session._add_bytes_processed(bytes_processed) + session._add_slot_millis(slot_millis) + if LOGGING_NAME_ENV_VAR in os.environ: + # when running notebooks via pytest nbmake + write_stats_to_disk(bytes_processed, slot_millis) return results_iterator, query_job -def pytest_log_job(query_job: bigquery.QueryJob): +def get_performance_stats(query_job: bigquery.QueryJob) -> Optional[Tuple[int, int]]: + """Parse the query job for performance stats. + + Return None if the stats do not reflect real work done in bigquery. + """ + bytes_processed = query_job.total_bytes_processed + if not isinstance(bytes_processed, int): + return None # filter out mocks + if query_job.configuration.dry_run: + # dry run stats are just predictions of the real run + bytes_processed = 0 + + slot_millis = query_job.slot_millis + if not isinstance(slot_millis, int): + return None # filter out mocks + if query_job.configuration.dry_run: + # dry run stats are just predictions of the real run + slot_millis = 0 + + return bytes_processed, slot_millis + + +def write_stats_to_disk(bytes_processed: int, slot_millis: int): """For pytest runs only, log information about the query job to a file in order to create a performance report. """ @@ -243,16 +293,17 @@ def pytest_log_job(query_job: bigquery.QueryJob): ) test_name = os.environ[LOGGING_NAME_ENV_VAR] current_directory = os.getcwd() - bytes_processed = query_job.total_bytes_processed - if not isinstance(bytes_processed, int): - return # filter out mocks - if query_job.configuration.dry_run: - # dry runs don't process their total_bytes_processed - bytes_processed = 0 + + # store bytes processed bytes_file = os.path.join(current_directory, test_name + ".bytesprocessed") with open(bytes_file, "a") as f: f.write(str(bytes_processed) + "\n") + # store slot milliseconds + bytes_file = os.path.join(current_directory, test_name + ".slotmillis") + with open(bytes_file, "a") as f: + f.write(str(slot_millis) + "\n") + def delete_tables_matching_session_id( client: bigquery.Client, dataset: bigquery.DatasetReference, session_id: str @@ -286,7 +337,10 @@ def delete_tables_matching_session_id( def create_bq_dataset_reference( - bq_client: bigquery.Client, location=None, project=None + bq_client: bigquery.Client, + location=None, + project=None, + api_name: str = "unknown", ) -> bigquery.DatasetReference: """Create and identify dataset(s) for temporary BQ resources. @@ -307,7 +361,11 @@ def create_bq_dataset_reference( Returns: bigquery.DatasetReference: The constructed reference to the anonymous dataset. """ - query_job = bq_client.query("SELECT 1", location=location, project=project) + job_config = google.cloud.bigquery.QueryJobConfig() + add_labels(job_config, api_name=api_name) + query_job = bq_client.query( + "SELECT 1", location=location, project=project, job_config=job_config + ) query_job.result() # blocks until finished # The anonymous dataset is used by BigQuery to write query results and @@ -334,14 +392,12 @@ def is_table_with_wildcard_suffix(query_or_table: str) -> bool: def to_query( query_or_table: str, - index_cols: Iterable[str], columns: Iterable[str], - filters: third_party_pandas_gbq.FiltersType, - max_results: Optional[int], - time_travel_timestamp: Optional[datetime.datetime], + sql_predicate: Optional[str], + max_results: Optional[int] = None, + time_travel_timestamp: Optional[datetime.datetime] = None, ) -> str: """Compile query_or_table with conditions(filters, wildcards) to query.""" - filters = list(filters) sub_query = ( f"({query_or_table})" if is_query(query_or_table) else f"`{query_or_table}`" ) @@ -351,8 +407,7 @@ def to_query( if columns: # We only reduce the selection if columns is set, but we always # want to make sure index_cols is also included. - all_columns = itertools.chain(index_cols, columns) - select_clause = "SELECT " + ", ".join(f"`{column}`" for column in all_columns) + select_clause = "SELECT " + ", ".join(f"`{column}`" for column in columns) else: select_clause = "SELECT *" @@ -365,77 +420,84 @@ def to_query( if max_results is not None: limit_clause = f" LIMIT {bigframes.core.sql.simple_literal(max_results)}" - filter_string = "" - if filters: - valid_operators: Mapping[third_party_pandas_gbq.FilterOps, str] = { - "in": "IN", - "not in": "NOT IN", - "LIKE": "LIKE", - "==": "=", - ">": ">", - "<": "<", - ">=": ">=", - "<=": "<=", - "!=": "!=", - } - - # If single layer filter, add another pseudo layer. So the single layer represents "and" logic. - if isinstance(filters[0], tuple) and ( - len(filters[0]) == 0 or not isinstance(list(filters[0])[0], tuple) - ): - filters = typing.cast(third_party_pandas_gbq.FiltersType, [filters]) - - for group in filters: - if not isinstance(group, Iterable): - group = [group] - - and_expression = "" - for filter_item in group: - if not isinstance(filter_item, tuple) or (len(filter_item) != 3): - raise ValueError( - f"Elements of filters must be tuples of length 3, but got {repr(filter_item)}.", - ) - - column, operator, value = filter_item - - if not isinstance(column, str): - raise ValueError( - f"Column name should be a string, but received '{column}' of type {type(column).__name__}." - ) - - if operator not in valid_operators: - raise ValueError(f"Operator {operator} is not valid.") - - operator_str = valid_operators[operator] - - column_ref = bigframes.core.sql.identifier(column) - if operator_str in ["IN", "NOT IN"]: - value_literal = bigframes.core.sql.multi_literal(*value) - else: - value_literal = bigframes.core.sql.simple_literal(value) - expression = bigframes.core.sql.infix_op( - operator_str, column_ref, value_literal - ) - if and_expression: - and_expression = bigframes.core.sql.infix_op( - "AND", and_expression, expression - ) - else: - and_expression = expression - - if filter_string: - filter_string = bigframes.core.sql.infix_op( - "OR", filter_string, and_expression - ) - else: - filter_string = and_expression - - where_clause = "" - if filter_string: - where_clause = f" WHERE {filter_string}" + where_clause = f" WHERE {sql_predicate}" if sql_predicate else "" return ( f"{select_clause} " f"FROM {sub_query}" f"{time_travel_clause}{where_clause}{limit_clause}" ) + + +def compile_filters(filters: third_party_pandas_gbq.FiltersType) -> str: + """Compiles a set of filters into a boolean sql expression""" + if not filters: + return "" + filter_string = "" + valid_operators: Mapping[third_party_pandas_gbq.FilterOps, str] = { + "in": "IN", + "not in": "NOT IN", + "LIKE": "LIKE", + "==": "=", + ">": ">", + "<": "<", + ">=": ">=", + "<=": "<=", + "!=": "!=", + } + + # If single layer filter, add another pseudo layer. So the single layer represents "and" logic. + filters_list: list = list(filters) + if isinstance(filters_list[0], tuple) and ( + len(filters_list[0]) == 0 or not isinstance(list(filters_list[0])[0], tuple) + ): + filter_items = [filters_list] + else: + filter_items = filters_list + + for group in filter_items: + if not isinstance(group, Iterable): + group = [group] + + and_expression = "" + for filter_item in group: + if not isinstance(filter_item, tuple) or (len(filter_item) != 3): + raise ValueError( + f"Elements of filters must be tuples of length 3, but got {repr(filter_item)}.", + ) + + column, operator, value = filter_item + + if not isinstance(column, str): + raise ValueError( + f"Column name should be a string, but received '{column}' of type {type(column).__name__}." + ) + + if operator not in valid_operators: + raise ValueError(f"Operator {operator} is not valid.") + + operator_str = valid_operators[operator] + + column_ref = bigframes.core.sql.identifier(column) + if operator_str in ["IN", "NOT IN"]: + value_literal = bigframes.core.sql.multi_literal(*value) + else: + value_literal = bigframes.core.sql.simple_literal(value) + expression = bigframes.core.sql.infix_op( + operator_str, column_ref, value_literal + ) + if and_expression: + and_expression = bigframes.core.sql.infix_op( + "AND", and_expression, expression + ) + else: + and_expression = expression + + if filter_string: + filter_string = bigframes.core.sql.infix_op( + "OR", filter_string, and_expression + ) + else: + filter_string = and_expression + + return filter_string diff --git a/bigframes/session/_io/bigquery/read_gbq_table.py b/bigframes/session/_io/bigquery/read_gbq_table.py index 063dde2a24..879a8ba44c 100644 --- a/bigframes/session/_io/bigquery/read_gbq_table.py +++ b/bigframes/session/_io/bigquery/read_gbq_table.py @@ -19,27 +19,18 @@ from __future__ import annotations import datetime -import itertools import typing -from typing import Dict, Iterable, List, Optional, Tuple +from typing import Dict, Iterable, List, Optional, Sequence, Tuple import warnings -import bigframes_vendored.ibis.expr.operations as vendored_ibis_ops -import bigframes_vendored.pandas.io.gbq as third_party_pandas_gbq import google.api_core.exceptions import google.cloud.bigquery as bigquery -import ibis -import ibis.backends -import ibis.expr.datatypes as ibis_dtypes -import ibis.expr.types as ibis_types import bigframes import bigframes.clients import bigframes.constants -import bigframes.core as core import bigframes.core.compile -import bigframes.core.guid as guid -import bigframes.core.ordering as order +import bigframes.core.compile.default_ordering import bigframes.core.sql import bigframes.dtypes import bigframes.session._io.bigquery @@ -51,28 +42,6 @@ import bigframes.session -def _convert_to_nonnull_string(column: ibis_types.Column) -> ibis_types.StringValue: - col_type = column.type() - if ( - col_type.is_numeric() - or col_type.is_boolean() - or col_type.is_binary() - or col_type.is_temporal() - ): - result = column.cast(ibis_dtypes.String(nullable=True)) - elif col_type.is_geospatial(): - result = typing.cast(ibis_types.GeoSpatialColumn, column).as_text() - elif col_type.is_string(): - result = column - else: - # TO_JSON_STRING works with all data types, but isn't the most efficient - # Needed for JSON, STRUCT and ARRAY datatypes - result = vendored_ibis_ops.ToJsonString(column).to_expr() # type: ignore - # Escape backslashes and use backslash as delineator - escaped = typing.cast(ibis_types.StringColumn, result.fillna("")).replace("\\", "\\\\") # type: ignore - return typing.cast(ibis_types.StringColumn, ibis.literal("\\")).concat(escaped) - - def get_table_metadata( bqclient: bigquery.Client, table_ref: google.cloud.bigquery.table.TableReference, @@ -113,8 +82,10 @@ def get_table_metadata( # atomically. table = bqclient.get_table(table_ref) + # TODO(swast): Use session._start_query instead? + # TODO(swast): Use query_and_wait since we know these are small results. job_config = bigquery.QueryJobConfig() - job_config.labels["bigframes-api"] = api_name + bigframes.session._io.bigquery.add_labels(job_config, api_name=api_name) snapshot_timestamp = list( bqclient.query( "SELECT CURRENT_TIMESTAMP() AS `current_timestamp`", @@ -126,39 +97,55 @@ def get_table_metadata( return cached_table -def get_ibis_time_travel_table( - ibis_client: ibis.BaseBackend, - table_ref: bigquery.TableReference, - index_cols: Iterable[str], - columns: Iterable[str], - filters: third_party_pandas_gbq.FiltersType, - time_travel_timestamp: Optional[datetime.datetime], -) -> ibis_types.Table: - # If we have an anonymous query results table, it can't be modified and - # there isn't any BigQuery time travel. - if table_ref.dataset_id.startswith("_"): - time_travel_timestamp = None - +def validate_table( + bqclient: bigquery.Client, + table_ref: bigquery.table.TableReference, + columns: Optional[Sequence[str]], + snapshot_time: datetime.datetime, + filter_str: Optional[str] = None, +) -> bool: + """Validates that the table can be read, returns True iff snapshot is supported.""" + # First run without snapshot to verify table can be read + sql = bigframes.session._io.bigquery.to_query( + query_or_table=f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}", + columns=columns or (), + sql_predicate=filter_str, + ) + dry_run_config = bigquery.QueryJobConfig() + dry_run_config.dry_run = True try: - return ibis_client.sql( - bigframes.session._io.bigquery.to_query( - f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}", - index_cols=index_cols, - columns=columns, - filters=filters, - time_travel_timestamp=time_travel_timestamp, - # If we've made it this far, we know we don't have any - # max_results to worry about, because in that case we will - # have executed a query with a LIMI clause. - max_results=None, - ) - ) + bqclient.query_and_wait(sql, job_config=dry_run_config) except google.api_core.exceptions.Forbidden as ex: - # Ibis does a dry run to get the types of the columns from the SQL. if "Drive credentials" in ex.message: ex.message += "\nCheck https://cloud.google.com/bigquery/docs/query-drive-data#Google_Drive_permissions." raise + # Anonymous dataset, does not support snapshot ever + if table_ref.dataset_id.startswith("_"): + return False + + # Second, try with snapshot to verify table supports this feature + snapshot_sql = bigframes.session._io.bigquery.to_query( + query_or_table=f"{table_ref.project}.{table_ref.dataset_id}.{table_ref.table_id}", + columns=columns or (), + sql_predicate=filter_str, + time_travel_timestamp=snapshot_time, + ) + try: + bqclient.query_and_wait(snapshot_sql, job_config=dry_run_config) + return True + except google.api_core.exceptions.NotFound: + # note that a notfound caused by a simple typo will be + # caught above when the metadata is fetched, not here + warnings.warn( + "NotFound error when reading table with time travel." + " Attempting query without time travel. Warning: Without" + " time travel, modifications to the underlying table may" + " result in errors or unexpected behavior.", + category=bigframes.exceptions.TimeTravelDisabledWarning, + ) + return False + def are_index_cols_unique( bqclient: bigquery.Client, @@ -334,83 +321,3 @@ def get_time_travel_datetime_and_table_metadata( cached_table = (snapshot_timestamp, table) cache[table_ref] = cached_table return cached_table - - -def to_array_value_with_total_ordering( - session: bigframes.session.Session, - table_expression: ibis_types.Table, - total_ordering_cols: List[str], -) -> core.ArrayValue: - """Create an ArrayValue, assuming we already have a total ordering.""" - ordering = order.ExpressionOrdering( - ordering_value_columns=tuple( - order.ascending_over(column_id) for column_id in total_ordering_cols - ), - total_ordering_columns=frozenset(total_ordering_cols), - ) - column_values = [table_expression[col] for col in table_expression.columns] - return core.ArrayValue.from_ibis( - session, - table_expression, - columns=column_values, - hidden_ordering_columns=[], - ordering=ordering, - ) - - -def to_array_value_with_default_ordering( - session: bigframes.session.Session, - table: ibis_types.Table, - table_rows: Optional[int], -) -> core.ArrayValue: - """Create an ArrayValue with a deterministic default ordering.""" - # Since this might also be used as the index, don't use the default - # "ordering ID" name. - - # For small tables, 64 bits is enough to avoid collisions, 128 bits will never ever collide no matter what - # Assume table is large if table row count is unknown - use_double_hash = (table_rows is None) or (table_rows == 0) or (table_rows > 100000) - - ordering_hash_part = guid.generate_guid("bigframes_ordering_") - ordering_hash_part2 = guid.generate_guid("bigframes_ordering_") - ordering_rand_part = guid.generate_guid("bigframes_ordering_") - - # All inputs into hash must be non-null or resulting hash will be null - str_values = list( - map(lambda col: _convert_to_nonnull_string(table[col]), table.columns) - ) - full_row_str = ( - str_values[0].concat(*str_values[1:]) if len(str_values) > 1 else str_values[0] - ) - full_row_hash = full_row_str.hash().name(ordering_hash_part) - # By modifying value slightly, we get another hash uncorrelated with the first - full_row_hash_p2 = (full_row_str + "_").hash().name(ordering_hash_part2) - # Used to disambiguate between identical rows (which will have identical hash) - random_value = ibis.random().name(ordering_rand_part) - - order_values = ( - [full_row_hash, full_row_hash_p2, random_value] - if use_double_hash - else [full_row_hash, random_value] - ) - - original_column_ids = table.columns - table_with_ordering = table.select( - itertools.chain(original_column_ids, order_values) - ) - - ordering = order.ExpressionOrdering( - ordering_value_columns=tuple( - order.ascending_over(col.get_name()) for col in order_values - ), - total_ordering_columns=frozenset(col.get_name() for col in order_values), - ) - columns = [table_with_ordering[col] for col in original_column_ids] - hidden_columns = [table_with_ordering[col.get_name()] for col in order_values] - return core.ArrayValue.from_ibis( - session, - table_with_ordering, - columns, - hidden_ordering_columns=hidden_columns, - ordering=ordering, - ) diff --git a/bigframes/version.py b/bigframes/version.py index da33be63fc..74a30e35b7 100644 --- a/bigframes/version.py +++ b/bigframes/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.7.0" +__version__ = "1.8.0" diff --git a/docs/reference/bigframes.ml/impute.rst b/docs/reference/bigframes.ml/impute.rst new file mode 100644 index 0000000000..3796e287ef --- /dev/null +++ b/docs/reference/bigframes.ml/impute.rst @@ -0,0 +1,7 @@ +bigframes.ml.impute +========================== + +.. automodule:: bigframes.ml.impute + :members: + :inherited-members: + :undoc-members: diff --git a/docs/reference/bigframes.ml/index.rst b/docs/reference/bigframes.ml/index.rst index 37504b0830..c14efaede6 100644 --- a/docs/reference/bigframes.ml/index.rst +++ b/docs/reference/bigframes.ml/index.rst @@ -19,6 +19,8 @@ API Reference imported + impute + linear_model llm diff --git a/docs/templates/toc.yml b/docs/templates/toc.yml index 67e628eb7d..aee3058ecc 100644 --- a/docs/templates/toc.yml +++ b/docs/templates/toc.yml @@ -134,6 +134,12 @@ - name: XGBoostModel uid: bigframes.ml.imported.XGBoostModel name: imported + - items: + - name: Overview + uid: bigframes.ml.impute + - name: SimpleImputer + uid: bigframes.ml.impute.SimpleImputer + name: impute - items: - name: Overview uid: bigframes.ml.linear_model diff --git a/notebooks/dataframes/index_col_null.ipynb b/notebooks/dataframes/index_col_null.ipynb new file mode 100644 index 0000000000..de373050fe --- /dev/null +++ b/notebooks/dataframes/index_col_null.ipynb @@ -0,0 +1,1409 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "eeec3428", + "metadata": {}, + "outputs": [], + "source": [ + "# Copyright 2023 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "47439dbd-4e54-4954-8b16-edc4bcd4f855", + "metadata": {}, + "source": [ + "# Operations with an \"NULL index\" DataFrame\n", + "\n", + "**Note**: This notebook describes a feature that is currently in [preview](https://cloud.google.com/blog/products/gcp/google-cloud-gets-simplified-product-launch-stages). There may be breaking changes to the functionality when using \"NULL index\" objects.\n", + "\n", + "Use the \"NULL\" index for more efficient query generation, but\n", + "some pandas-compatible methods may not be possible without an index." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "96757c59-fc22-420e-a42f-c6cb956110ec", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "import bigframes.enums\n", + "import bigframes.exceptions\n", + "import bigframes.pandas as bpd\n", + "\n", + "# Explicitly opt-in to the NULL index preview feature.\n", + "warnings.simplefilter(\n", + " \"ignore\",\n", + " bigframes.exceptions.NullIndexPreviewWarning,\n", + ")\n", + "\n", + "df = bpd.read_gbq(\n", + " \"bigquery-public-data.baseball.schedules\",\n", + " index_col=bigframes.enums.DefaultIndexKind.NULL,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "d15688e1", + "metadata": {}, + "source": [ + "Use `peek()` to view an arbitrary selection of rows from the DataFrame. This is much more efficient than `head()`, which requires a total ordering for determinism." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "c93949fb", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 1b8726ce-c4ea-47fe-a47c-d6fae50d8fb0 is DONE. 582.8 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdgameNumberseasonIdyeartypedayNightdurationduration_minuteshomeTeamIdhomeTeamNameawayTeamIdawayTeamNamestartTimeattendancestatuscreated
0e14b6493-9e7f-404f-840a-8a680cc364bf1565de4be-dc80-4849-a7e1-54bc79156cc82016REGD3:0718703556285-bdbb-4576-a06d-42f71f46ddc5Marlins55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-26 17:10:00+00:0027318closed2016-10-06 06:25:15+00:00
11f32b347-cbcb-4c31-a145-0e685306d1681565de4be-dc80-4849-a7e1-54bc79156cc82016REGD3:0918903556285-bdbb-4576-a06d-42f71f46ddc5Marlins55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-25 20:10:00+00:0029457closed2016-10-06 06:25:15+00:00
20c2292d1-7398-48be-bf8e-b41dad5e1a431565de4be-dc80-4849-a7e1-54bc79156cc82016REGD2:4516512079497-e414-450a-8bf2-29f91de646bfBraves55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-11 20:10:00+00:0043114closed2016-10-06 06:25:15+00:00
38fbec734-a15a-42ab-8d51-60790de7750b1565de4be-dc80-4849-a7e1-54bc79156cc82016REGD3:4222212079497-e414-450a-8bf2-29f91de646bfBraves55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-12 17:35:00+00:0031625closed2016-10-06 06:25:15+00:00
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd1565de4be-dc80-4849-a7e1-54bc79156cc82016REGD2:441642142e1ba-3b40-445c-b8bb-f1f8b1054220Phillies55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-08 17:05:00+00:0028650closed2016-10-06 06:25:15+00:00
\n", + "
" + ], + "text/plain": [ + " gameId gameNumber \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 1 \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 1 \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 1 \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 1 \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 1 \n", + "\n", + " seasonId year type dayNight duration \\\n", + "0 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG D 3:07 \n", + "1 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG D 3:09 \n", + "2 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG D 2:45 \n", + "3 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG D 3:42 \n", + "4 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG D 2:44 \n", + "\n", + " duration_minutes homeTeamId homeTeamName \\\n", + "0 187 03556285-bdbb-4576-a06d-42f71f46ddc5 Marlins \n", + "1 189 03556285-bdbb-4576-a06d-42f71f46ddc5 Marlins \n", + "2 165 12079497-e414-450a-8bf2-29f91de646bf Braves \n", + "3 222 12079497-e414-450a-8bf2-29f91de646bf Braves \n", + "4 164 2142e1ba-3b40-445c-b8bb-f1f8b1054220 Phillies \n", + "\n", + " awayTeamId awayTeamName \\\n", + "0 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "1 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "2 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "3 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "4 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "\n", + " startTime attendance status created \n", + "0 2016-06-26 17:10:00+00:00 27318 closed 2016-10-06 06:25:15+00:00 \n", + "1 2016-06-25 20:10:00+00:00 29457 closed 2016-10-06 06:25:15+00:00 \n", + "2 2016-06-11 20:10:00+00:00 43114 closed 2016-10-06 06:25:15+00:00 \n", + "3 2016-06-12 17:35:00+00:00 31625 closed 2016-10-06 06:25:15+00:00 \n", + "4 2016-06-08 17:05:00+00:00 28650 closed 2016-10-06 06:25:15+00:00 " + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.peek()" + ] + }, + { + "cell_type": "markdown", + "id": "78e3d27d", + "metadata": {}, + "source": [ + "# Inspect the properties of the DataFrame\n", + "\n", + "Some properties, such as `dtypes`, can be retrieved without executing a query job." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "38f566c5", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "gameId string[pyarrow]\n", + "gameNumber Int64\n", + "seasonId string[pyarrow]\n", + "year Int64\n", + "type string[pyarrow]\n", + "dayNight string[pyarrow]\n", + "duration string[pyarrow]\n", + "duration_minutes Int64\n", + "homeTeamId string[pyarrow]\n", + "homeTeamName string[pyarrow]\n", + "awayTeamId string[pyarrow]\n", + "awayTeamName string[pyarrow]\n", + "startTime timestamp[us, tz=UTC][pyarrow]\n", + "attendance Int64\n", + "status string[pyarrow]\n", + "created timestamp[us, tz=UTC][pyarrow]\n", + "dtype: object" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.dtypes" + ] + }, + { + "cell_type": "markdown", + "id": "38a59ecc", + "metadata": {}, + "source": [ + "Other properties, such as `shape` require a query. In this case, `shape` runs a `COUNT(1)` query." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "e3b43d37", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 0f85f12c-227c-4001-b851-6e9b9087ab7e is DONE. 0 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "(2431, 16)" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df.shape" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "13861abc-120c-4db6-ad0c-e414b85d3443", + "metadata": {}, + "source": [ + "### Select a subset of the DataFrame\n", + "\n", + "Filter columns by selecting a list of columns from the DataFrame.\n", + "\n", + "**Note**: Even with `index_col=bigframes.enums.DefaultIndexKind.NULL`, it is more efficient to do this selection in `read_gbq` / `read_gbq_table` except in cases where the total ordering ID columns can be pruned." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "05cb36e9-bb75-4f6f-8eb6-e4219df6e1d2", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job efa6b4be-cf60-4951-9125-7d77fb6b6b44 is DONE. 174.4 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdyearhomeTeamNameawayTeamNameduration_minutes
0e14b6493-9e7f-404f-840a-8a680cc364bf2016MarlinsCubs187
11f32b347-cbcb-4c31-a145-0e685306d1682016MarlinsCubs189
20c2292d1-7398-48be-bf8e-b41dad5e1a432016BravesCubs165
38fbec734-a15a-42ab-8d51-60790de7750b2016BravesCubs222
489e514d5-fbf5-4b9d-bdac-6ca45bfd18dd2016PhilliesCubs164
\n", + "
" + ], + "text/plain": [ + " gameId year homeTeamName awayTeamName \\\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf 2016 Marlins Cubs \n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 2016 Marlins Cubs \n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 2016 Braves Cubs \n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b 2016 Braves Cubs \n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd 2016 Phillies Cubs \n", + "\n", + " duration_minutes \n", + "0 187 \n", + "1 189 \n", + "2 165 \n", + "3 222 \n", + "4 164 " + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "column_filtered = df[[\"gameId\", \"year\", \"homeTeamName\", \"awayTeamName\", \"duration_minutes\"]]\n", + "column_filtered.peek()" + ] + }, + { + "cell_type": "markdown", + "id": "d4d52c41", + "metadata": {}, + "source": [ + "Filter by rows using a boolean Series. This Series must be derived from the DataFrame being filtered so that the NULL index can still align correctly." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "a6b8b3ac-1df8-46ff-ac4f-d6e7657fc80c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 0be8e44d-854a-45ca-950b-269280e3de41 is DONE. 582.8 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdgameNumberseasonIdyeartypedayNightdurationduration_minuteshomeTeamIdhomeTeamNameawayTeamIdawayTeamNamestartTimeattendancestatuscreated
063f14670-c28e-432b-84ee-1a2c6ac295271565de4be-dc80-4849-a7e1-54bc79156cc82016REGN2:4316303556285-bdbb-4576-a06d-42f71f46ddc5Marlins55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-23 23:10:00+00:0025291closed2016-10-06 06:25:15+00:00
1bf4e80d1-3125-44fa-8a89-de93d039d4651565de4be-dc80-4849-a7e1-54bc79156cc82016REGN3:2420403556285-bdbb-4576-a06d-42f71f46ddc5Marlins55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-24 23:10:00+00:0024385closed2016-10-06 06:25:15+00:00
2e8af534c-36ed-4ff9-8511-780825fdd0411565de4be-dc80-4849-a7e1-54bc79156cc82016REGN2:5117112079497-e414-450a-8bf2-29f91de646bfBraves55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-10 23:35:00+00:0030547closed2016-10-06 06:25:15+00:00
3e599c525-ac42-4b54-928d-7ee5fbe67dd91565de4be-dc80-4849-a7e1-54bc79156cc82016REGN2:451652142e1ba-3b40-445c-b8bb-f1f8b1054220Phillies55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-07 23:05:00+00:0027381closed2016-10-06 06:25:15+00:00
4d80ffb65-57a4-42c9-ae1c-2c51d06503361565de4be-dc80-4849-a7e1-54bc79156cc82016REGN3:051852142e1ba-3b40-445c-b8bb-f1f8b1054220Phillies55714da8-fcaf-4574-8443-59bfb511a524Cubs2016-06-06 23:05:00+00:0022162closed2016-10-06 06:25:15+00:00
\n", + "
" + ], + "text/plain": [ + " gameId gameNumber \\\n", + "0 63f14670-c28e-432b-84ee-1a2c6ac29527 1 \n", + "1 bf4e80d1-3125-44fa-8a89-de93d039d465 1 \n", + "2 e8af534c-36ed-4ff9-8511-780825fdd041 1 \n", + "3 e599c525-ac42-4b54-928d-7ee5fbe67dd9 1 \n", + "4 d80ffb65-57a4-42c9-ae1c-2c51d0650336 1 \n", + "\n", + " seasonId year type dayNight duration \\\n", + "0 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG N 2:43 \n", + "1 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG N 3:24 \n", + "2 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG N 2:51 \n", + "3 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG N 2:45 \n", + "4 565de4be-dc80-4849-a7e1-54bc79156cc8 2016 REG N 3:05 \n", + "\n", + " duration_minutes homeTeamId homeTeamName \\\n", + "0 163 03556285-bdbb-4576-a06d-42f71f46ddc5 Marlins \n", + "1 204 03556285-bdbb-4576-a06d-42f71f46ddc5 Marlins \n", + "2 171 12079497-e414-450a-8bf2-29f91de646bf Braves \n", + "3 165 2142e1ba-3b40-445c-b8bb-f1f8b1054220 Phillies \n", + "4 185 2142e1ba-3b40-445c-b8bb-f1f8b1054220 Phillies \n", + "\n", + " awayTeamId awayTeamName \\\n", + "0 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "1 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "2 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "3 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "4 55714da8-fcaf-4574-8443-59bfb511a524 Cubs \n", + "\n", + " startTime attendance status created \n", + "0 2016-06-23 23:10:00+00:00 25291 closed 2016-10-06 06:25:15+00:00 \n", + "1 2016-06-24 23:10:00+00:00 24385 closed 2016-10-06 06:25:15+00:00 \n", + "2 2016-06-10 23:35:00+00:00 30547 closed 2016-10-06 06:25:15+00:00 \n", + "3 2016-06-07 23:05:00+00:00 27381 closed 2016-10-06 06:25:15+00:00 \n", + "4 2016-06-06 23:05:00+00:00 22162 closed 2016-10-06 06:25:15+00:00 " + ] + }, + "execution_count": 7, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "night_games = df[df['dayNight'] == 'N']\n", + "night_games.peek()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "80e9a2e2-c4c9-4c17-bbd0-06882d7657fe", + "metadata": {}, + "source": [ + "### Join two DataFrames\n", + "\n", + "Even though pandas usually joins by the index, NULL index objects can still be manually joined by a column using the `on` parameter in `merge`." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "3f09ff32-ef43-4fab-a86b-8868afc34363", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 5d2c69d2-33fe-4513-923b-fd64f4da098b is DONE. 113.9 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamName
0e14b6493-9e7f-404f-840a-8a680cc364bfMarlins
11f32b347-cbcb-4c31-a145-0e685306d168Marlins
20c2292d1-7398-48be-bf8e-b41dad5e1a43Braves
38fbec734-a15a-42ab-8d51-60790de7750bBraves
489e514d5-fbf5-4b9d-bdac-6ca45bfd18ddPhillies
\n", + "
" + ], + "text/plain": [ + " gameId homeTeamName\n", + "0 e14b6493-9e7f-404f-840a-8a680cc364bf Marlins\n", + "1 1f32b347-cbcb-4c31-a145-0e685306d168 Marlins\n", + "2 0c2292d1-7398-48be-bf8e-b41dad5e1a43 Braves\n", + "3 8fbec734-a15a-42ab-8d51-60790de7750b Braves\n", + "4 89e514d5-fbf5-4b9d-bdac-6ca45bfd18dd Phillies" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df1 = df[[\"gameId\", \"homeTeamName\"]]\n", + "df1.peek()" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "5331d2c8-7912-4d96-8da1-f64b57374df3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job b6b70d6d-a490-44d6-ba74-0ee32b4f0a1a is DONE. 582.8 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 68acd168-8b42-44f8-8702-99618935991e is DONE. 94 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdawayTeamName
0af72a0b9-65f7-49fb-9b30-d505068bdf6dBrewers
1d60c6036-0ce1-4c90-8dd9-de3b403c92a8Brewers
\n", + "
" + ], + "text/plain": [ + " gameId awayTeamName\n", + "0 af72a0b9-65f7-49fb-9b30-d505068bdf6d Brewers\n", + "1 d60c6036-0ce1-4c90-8dd9-de3b403c92a8 Brewers" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "df2 = df[[\"gameId\", \"awayTeamName\"]].head(2)\n", + "df2.peek()" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "a574ad3e-a219-454c-8bb5-c5ed6627f2c6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 0ac171dd-3859-4589-b7ff-59fd81ec3c3a is DONE. 582.9 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 034f8807-c128-444a-8033-0c95f34b0e32 is DONE. 111 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0af72a0b9-65f7-49fb-9b30-d505068bdf6dRedsBrewers
1d60c6036-0ce1-4c90-8dd9-de3b403c92a8NationalsBrewers
\n", + "
" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 af72a0b9-65f7-49fb-9b30-d505068bdf6d Reds Brewers\n", + "1 d60c6036-0ce1-4c90-8dd9-de3b403c92a8 Nationals Brewers" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "merged = df1.merge(df2, on=\"gameId\", how=\"inner\")\n", + "merged.peek()" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "288e7a95-a077-46c4-8fe6-802474c01f8b", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 30fd5a60-772c-4ef0-a151-5ab390ff4322 is DONE. 582.9 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 701fa9a8-1ec6-49b9-ac41-228cb34d4c8c is DONE. 114.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0039bb40e-7613-4674-a653-584b93e9b21bAmerican League<NA>
178000e12-2ef3-4246-adc1-c8a4d157631cAngels<NA>
2de5555dc-9228-4f7c-88ae-4451e3ffb980Angels<NA>
3f29a2754-004b-436c-91fe-3d86c0bb17a8Angels<NA>
48e5af008-8a07-4f9a-90cb-336ca4c84c71Angels<NA>
\n", + "
" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 039bb40e-7613-4674-a653-584b93e9b21b American League \n", + "1 78000e12-2ef3-4246-adc1-c8a4d157631c Angels \n", + "2 de5555dc-9228-4f7c-88ae-4451e3ffb980 Angels \n", + "3 f29a2754-004b-436c-91fe-3d86c0bb17a8 Angels \n", + "4 8e5af008-8a07-4f9a-90cb-336ca4c84c71 Angels " + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "merged = df1.merge(df2, on=\"gameId\", how=\"outer\")\n", + "merged.peek()" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "7ee87a01-2ff5-4021-855d-44b71cf2a225", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job e3d8168c-48e9-4ba9-a916-10259ad9c0ea is DONE. 582.9 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job 82d2a5e4-66a8-4478-92de-57d3f806aa76 is DONE. 114.0 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0039bb40e-7613-4674-a653-584b93e9b21bAmerican League<NA>
1f6fcd83c-e130-487c-a0cc-d00b2712d08bAngels<NA>
2fe401dd2-089c-4822-8657-4d510d460f38Angels<NA>
3c894bdee-5dda-49f4-87c8-53b9b9bfcd3bAngels<NA>
4bbda59d9-fd52-4bed-bcfb-2ceed4be997cAngels<NA>
\n", + "
" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 039bb40e-7613-4674-a653-584b93e9b21b American League \n", + "1 f6fcd83c-e130-487c-a0cc-d00b2712d08b Angels \n", + "2 fe401dd2-089c-4822-8657-4d510d460f38 Angels \n", + "3 c894bdee-5dda-49f4-87c8-53b9b9bfcd3b Angels \n", + "4 bbda59d9-fd52-4bed-bcfb-2ceed4be997c Angels " + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "merged = df1.merge(df2, on=\"gameId\", how=\"left\")\n", + "merged.peek()" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "330ed69c-f122-4af9-bf5e-96e309d3fa0c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 518ed511-606a-42b2-a28d-61a601eccfa7 is DONE. 582.9 kB processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "Query job bc381640-74e0-4885-9c32-87805a49f357 is DONE. 111 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0af72a0b9-65f7-49fb-9b30-d505068bdf6dRedsBrewers
1d60c6036-0ce1-4c90-8dd9-de3b403c92a8NationalsBrewers
\n", + "
" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 af72a0b9-65f7-49fb-9b30-d505068bdf6d Reds Brewers\n", + "1 d60c6036-0ce1-4c90-8dd9-de3b403c92a8 Nationals Brewers" + ] + }, + "execution_count": 13, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "merged = df1.merge(df2, on=\"gameId\", how=\"right\")\n", + "merged.peek()" + ] + }, + { + "cell_type": "markdown", + "id": "162eede7", + "metadata": {}, + "source": [ + "### Download the result as (in-memory) pandas DataFrame\n", + "\n", + "Use the `ordered=False` argument for more efficient query execution." + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "ab429fa5", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "Query job 2d4fbd55-ba6a-46d2-87ae-5da416ad3642 is DONE. 159 Bytes processed. Open Job" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
gameIdhomeTeamNameawayTeamName
0d60c6036-0ce1-4c90-8dd9-de3b403c92a8NationalsBrewers
1af72a0b9-65f7-49fb-9b30-d505068bdf6dRedsBrewers
\n", + "
" + ], + "text/plain": [ + " gameId homeTeamName awayTeamName\n", + "0 d60c6036-0ce1-4c90-8dd9-de3b403c92a8 Nationals Brewers\n", + "1 af72a0b9-65f7-49fb-9b30-d505068bdf6d Reds Brewers" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "dfp = merged.to_pandas(ordered=False)\n", + "dfp" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "896212ab", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.1" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/noxfile.py b/noxfile.py index c816ec5f51..52583bbf1a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -91,7 +91,6 @@ "unit", "unit_noextras", "system", - "doctest", "cover", ] @@ -801,33 +800,57 @@ def notebook(session: nox.Session): for process in processes: process.join() - # when run via pytest, notebooks output a .bytesprocessed report + # when the environment variable is set as it is above, + # notebooks output a .bytesprocessed and .slotmillis report # collect those reports and print a summary - _print_bytes_processed_report() + _print_performance_report() -def _print_bytes_processed_report(): - """Add an informational report about http queries and bytes - processed to the testlog output for purposes of measuring - bigquery-related performance changes. +def _print_performance_report(): + """Add an informational report about http queries, bytes + processed, and slot time to the testlog output for purposes + of measuring bigquery-related performance changes. """ print("---BIGQUERY USAGE REPORT---") - cumulative_queries = 0 - cumulative_bytes = 0 - for report in Path("notebooks/").glob("*/*.bytesprocessed"): - with open(report, "r") as f: - filename = report.stem - lines = f.read().splitlines() + results_dict = {} + for bytes_report in Path("notebooks/").glob("*/*.bytesprocessed"): + with open(bytes_report, "r") as bytes_file: + filename = bytes_report.stem + lines = bytes_file.read().splitlines() query_count = len(lines) total_bytes = sum([int(line) for line in lines]) - format_string = f"{filename} - query count: {query_count}, bytes processed sum: {total_bytes}" - print(format_string) - cumulative_bytes += total_bytes - cumulative_queries += query_count - print( - "---total queries: {total_queries}, total bytes: {total_bytes}---".format( - total_queries=cumulative_queries, total_bytes=cumulative_bytes + results_dict[filename] = [query_count, total_bytes] + for millis_report in Path("notebooks/").glob("*/*.slotmillis"): + with open(millis_report, "r") as millis_file: + filename = millis_report.stem + lines = millis_file.read().splitlines() + total_slot_millis = sum([int(line) for line in lines]) + results_dict[filename] += [total_slot_millis] + + cumulative_queries = 0 + cumulative_bytes = 0 + cumulative_slot_millis = 0 + for results in results_dict.values(): + if len(results) != 3: + raise IOError( + "Mismatch in performance logging output. " + "Expected one .bytesprocessed and one .slotmillis " + "file for each notebook." + ) + query_count, total_bytes, total_slot_millis = results + cumulative_queries += query_count + cumulative_bytes += total_bytes + cumulative_slot_millis += total_slot_millis + print( + f"{filename} - query count: {query_count}," + f" bytes processed sum: {total_bytes}," + f" slot millis sum: {total_slot_millis}" ) + + print( + f"---total queries: {cumulative_queries}, " + f"total bytes: {cumulative_bytes}, " + f"total slot millis: {cumulative_slot_millis}---" ) diff --git a/samples/snippets/requirements-test.txt b/samples/snippets/requirements-test.txt index 62b0c02e79..beca2e44d9 100644 --- a/samples/snippets/requirements-test.txt +++ b/samples/snippets/requirements-test.txt @@ -1,3 +1,3 @@ # samples/snippets should be runnable with no "extras" google-cloud-testutils==1.4.0 -pytest==8.1.1 +pytest==8.2.0 diff --git a/tests/system/conftest.py b/tests/system/conftest.py index ecf633b27f..2501693084 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -18,6 +18,7 @@ import math import pathlib import textwrap +import traceback import typing from typing import Dict, Generator, Optional @@ -29,6 +30,7 @@ import google.cloud.resourcemanager_v3 as resourcemanager_v3 import google.cloud.storage as storage # type: ignore import ibis.backends.base +import numpy as np import pandas as pd import pytest import pytz @@ -36,6 +38,7 @@ import bigframes import bigframes.dataframe +import bigframes.pandas as bpd import tests.system.utils # Use this to control the number of cloud functions being deleted in a single @@ -392,7 +395,7 @@ def scalars_df_index( @pytest.fixture(scope="session") -def scalars_df_empty_index( +def scalars_df_null_index( scalars_table_id: str, session: bigframes.Session ) -> bigframes.dataframe.DataFrame: """DataFrame pointing at test data.""" @@ -624,6 +627,18 @@ def new_penguins_pandas_df(): ).set_index("tag_number") +@pytest.fixture(scope="session") +def missing_values_penguins_df(): + """Additional data matching the missing values penguins dataset""" + return bpd.DataFrame( + { + "culmen_length_mm": [39.5, 38.5, 37.9], + "culmen_depth_mm": [np.nan, 17.2, 18.1], + "flipper_length_mm": [np.nan, 181.0, 188.0], + } + ) + + @pytest.fixture(scope="session") def new_penguins_df(session, new_penguins_pandas_df): return session.read_pandas(new_penguins_pandas_df) @@ -1082,46 +1097,54 @@ def cleanup_cloud_functions(session, cloudfunctions_client, dataset_id_permanent session.bqclient, dataset_id_permanent ) delete_count = 0 - for cloud_function in tests.system.utils.get_cloud_functions( - cloudfunctions_client, - session.bqclient.project, - session.bqclient.location, - name_prefix="bigframes-", - ): - # Ignore bigframes cloud functions referred by the remote functions in - # the permanent dataset - if cloud_function.service_config.uri in permanent_endpoints: - continue - - # Ignore the functions less than one day old - age = datetime.now() - datetime.fromtimestamp( - cloud_function.update_time.timestamp() - ) - if age.days <= 0: - continue - - # Go ahead and delete - try: - tests.system.utils.delete_cloud_function( - cloudfunctions_client, cloud_function.name + try: + for cloud_function in tests.system.utils.get_cloud_functions( + cloudfunctions_client, + session.bqclient.project, + session.bqclient.location, + name_prefix="bigframes-", + ): + # Ignore bigframes cloud functions referred by the remote functions in + # the permanent dataset + if cloud_function.service_config.uri in permanent_endpoints: + continue + + # Ignore the functions less than one day old + age = datetime.now() - datetime.fromtimestamp( + cloud_function.update_time.timestamp() ) - delete_count += 1 - if delete_count >= MAX_NUM_FUNCTIONS_TO_DELETE_PER_SESSION: - break - except google.api_core.exceptions.NotFound: - # This can happen when multiple pytest sessions are running in - # parallel. Two or more sessions may discover the same cloud - # function, but only one of them would be able to delete it - # successfully, while the other instance will run into this - # exception. Ignore this exception. - pass - except google.api_core.exceptions.ResourceExhausted: - # This can happen if we are hitting GCP limits, e.g. - # google.api_core.exceptions.ResourceExhausted: 429 Quota exceeded - # for quota metric 'Per project mutation requests' and limit - # 'Per project mutation requests per minute per region' of service - # 'cloudfunctions.googleapis.com' for consumer - # 'project_number:1084210331973'. - # [reason: "RATE_LIMIT_EXCEEDED" domain: "googleapis.com" ... - # Let's stop further clean up and leave it to later. - break + if age.days <= 0: + continue + + # Go ahead and delete + try: + tests.system.utils.delete_cloud_function( + cloudfunctions_client, cloud_function.name + ) + delete_count += 1 + if delete_count >= MAX_NUM_FUNCTIONS_TO_DELETE_PER_SESSION: + break + except google.api_core.exceptions.NotFound: + # This can happen when multiple pytest sessions are running in + # parallel. Two or more sessions may discover the same cloud + # function, but only one of them would be able to delete it + # successfully, while the other instance will run into this + # exception. Ignore this exception. + pass + except Exception as exc: + # Don't fail the tests for unknown exceptions. + # + # This can happen if we are hitting GCP limits, e.g. + # google.api_core.exceptions.ResourceExhausted: 429 Quota exceeded + # for quota metric 'Per project mutation requests' and limit + # 'Per project mutation requests per minute per region' of service + # 'cloudfunctions.googleapis.com' for consumer + # 'project_number:1084210331973'. + # [reason: "RATE_LIMIT_EXCEEDED" domain: "googleapis.com" ... + # + # It can also happen occasionally with + # google.api_core.exceptions.ServiceUnavailable when there is some + # backend flakiness. + # + # Let's stop further clean up and leave it to later. + traceback.print_exception(exc) diff --git a/tests/system/large/ml/test_ensemble.py b/tests/system/large/ml/test_ensemble.py index e00f7fa665..706cbfdfaf 100644 --- a/tests/system/large/ml/test_ensemble.py +++ b/tests/system/large/ml/test_ensemble.py @@ -395,7 +395,7 @@ def test_randomforestclassifier_multiple_params(penguins_df_default_index, datas f"{dataset_id}.temp_configured_randomforestclassifier_model" in reloaded_model._bqml_model.model_name ) - assert reloaded_model.tree_method == "auto" + assert reloaded_model.tree_method.casefold() == "auto" assert reloaded_model.colsample_bytree == 0.95 assert reloaded_model.colsample_bylevel == 0.95 assert reloaded_model.colsample_bynode == 0.95 diff --git a/tests/system/large/ml/test_pipeline.py b/tests/system/large/ml/test_pipeline.py index 3d7eb2e426..6e18248e0f 100644 --- a/tests/system/large/ml/test_pipeline.py +++ b/tests/system/large/ml/test_pipeline.py @@ -20,6 +20,7 @@ compose, decomposition, ensemble, + impute, linear_model, pipeline, preprocessing, @@ -476,6 +477,11 @@ def test_pipeline_columntransformer_fit_predict(session, penguins_df_default_ind preprocessing.KBinsDiscretizer(strategy="uniform"), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "simple_imputer", + impute.SimpleImputer(strategy="mean"), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -556,6 +562,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id preprocessing.KBinsDiscretizer(strategy="uniform"), ["culmen_length_mm", "flipper_length_mm"], ), + ( + "simple_imputer", + impute.SimpleImputer(), + ["culmen_length_mm", "flipper_length_mm"], + ), ( "label", preprocessing.LabelEncoder(), @@ -600,6 +611,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id preprocessing.KBinsDiscretizer(strategy="uniform"), "culmen_length_mm", ), + ( + "simple_imputer", + impute.SimpleImputer(), + "culmen_length_mm", + ), ("standard_scaler", preprocessing.StandardScaler(), "flipper_length_mm"), ("max_abs_scaler", preprocessing.MaxAbsScaler(), "flipper_length_mm"), ("min_max_scaler", preprocessing.MinMaxScaler(), "flipper_length_mm"), @@ -608,6 +624,11 @@ def test_pipeline_columntransformer_to_gbq(penguins_df_default_index, dataset_id preprocessing.KBinsDiscretizer(strategy="uniform"), "flipper_length_mm", ), + ( + "simple_imputer", + impute.SimpleImputer(), + "flipper_length_mm", + ), ] assert transformers == expected @@ -795,3 +816,36 @@ def test_pipeline_label_encoder_to_gbq(penguins_df_default_index, dataset_id): assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) assert pl_loaded._estimator.fit_intercept is False + + +def test_pipeline_simple_imputer_to_gbq(penguins_df_default_index, dataset_id): + pl = pipeline.Pipeline( + [ + ( + "transform", + impute.SimpleImputer(strategy="most_frequent"), + ), + ("estimator", linear_model.LinearRegression(fit_intercept=False)), + ] + ) + + df = penguins_df_default_index.dropna() + X_train = df[ + [ + "sex", + "species", + ] + ] + y_train = df[["body_mass_g"]] + pl.fit(X_train, y_train) + + pl_loaded = pl.to_gbq( + f"{dataset_id}.test_penguins_pipeline_simple_imputer", replace=True + ) + assert isinstance(pl_loaded._transform, impute.SimpleImputer) + + simple_imputer = pl_loaded._transform + assert simple_imputer.strategy == "most_frequent" + + assert isinstance(pl_loaded._estimator, linear_model.LinearRegression) + assert pl_loaded._estimator.fit_intercept is False diff --git a/tests/system/large/test_remote_function.py b/tests/system/large/test_remote_function.py index cac8483b5b..cce49ea187 100644 --- a/tests/system/large/test_remote_function.py +++ b/tests/system/large/test_remote_function.py @@ -21,7 +21,7 @@ import tempfile import textwrap -from google.api_core.exceptions import BadRequest, InvalidArgument, NotFound +import google.api_core.exceptions from google.cloud import bigquery, storage import pandas import pytest @@ -194,6 +194,9 @@ def test_remote_function_stringify_with_ibis( def stringify(x): return f"I got {x}" + # Function should work locally. + assert stringify(42) == "I got 42" + _, dataset_name, table_name = scalars_table_id.split(".") if not ibis_client.dataset: ibis_client.dataset = dataset_name @@ -205,7 +208,7 @@ def stringify(x): pandas_df_orig = bigquery_client.query(sql).to_dataframe() col = table[col_name] - col_2x = stringify(col).name("int64_str_col") + col_2x = stringify.ibis_node(col).name("int64_str_col") table = table.mutate([col_2x]) sql = table.compile() pandas_df_new = bigquery_client.query(sql).to_dataframe() @@ -846,7 +849,7 @@ def square(x): expected_remote_function = f"{dataset_id}.{rf_name}" # Initially the expected BQ remote function should not exist - with pytest.raises(NotFound): + with pytest.raises(google.api_core.exceptions.NotFound): session.bqclient.get_routine(expected_remote_function) # Create the remote function with the name provided explicitly @@ -978,7 +981,7 @@ def test_internal(rf, udf): expected_remote_function = f"{dataset_id}.{rf_name}" # Initially the expected BQ remote function should not exist - with pytest.raises(NotFound): + with pytest.raises(google.api_core.exceptions.NotFound): session.bqclient.get_routine(expected_remote_function) # Create a new remote function with the name provided explicitly @@ -1195,7 +1198,8 @@ def square(x): scalars_df, _ = scalars_dfs with pytest.raises( - BadRequest, match="400.*errorMessage.*unsupported operand type" + google.api_core.exceptions.BadRequest, + match="400.*errorMessage.*unsupported operand type", ): # int64_col has nulls which should cause error in square scalars_df["int64_col"].apply(square).to_pandas() @@ -1428,20 +1432,6 @@ def square_num(x): ) -def test_remote_function_via_session_vpc_invalid(session): - with pytest.raises( - InvalidArgument, match="400.*Serverless VPC Access connector is not found" - ): - - @session.remote_function( - [int], int, reuse=False, cloud_function_vpc_connector="does-not-exist" - ) - def square_num(x): - if x is None: - return x - return x * x - - @pytest.mark.parametrize( ("max_batching_rows"), [ diff --git a/tests/system/small/bigquery/test_array.py b/tests/system/small/bigquery/test_array.py index 0664c31a3c..d6823a3a54 100644 --- a/tests/system/small/bigquery/test_array.py +++ b/tests/system/small/bigquery/test_array.py @@ -139,3 +139,16 @@ def test_array_agg_matches_after_explode(): result.to_pandas(), # type: ignore df.to_pandas(), ) + + +@pytest.mark.parametrize( + ("data"), + [ + pytest.param([[1, 2], [3, 4], [5]], id="int_array"), + pytest.param(["hello", "world"], id="string"), + ], +) +def test_array_to_string_w_type_checks(data): + series = bpd.Series(data) + with pytest.raises(TypeError): + bbq.array_to_string(series, delimiter=", ") diff --git a/tests/system/small/ml/conftest.py b/tests/system/small/ml/conftest.py index 33351afe45..ee96646687 100644 --- a/tests/system/small/ml/conftest.py +++ b/tests/system/small/ml/conftest.py @@ -275,11 +275,6 @@ def palm2_embedding_generator_multilingual_model( ) -@pytest.fixture(scope="session") -def gemini_text_generator_model(session, bq_connection) -> llm.GeminiTextGenerator: - return llm.GeminiTextGenerator(session=session, connection_name=bq_connection) - - @pytest.fixture(scope="session") def linear_remote_model_params() -> dict: # Pre-deployed endpoint of linear reg model in Vertex. diff --git a/tests/system/small/ml/test_impute.py b/tests/system/small/ml/test_impute.py new file mode 100644 index 0000000000..46a614d703 --- /dev/null +++ b/tests/system/small/ml/test_impute.py @@ -0,0 +1,89 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pandas as pd + +from bigframes.ml import impute + + +def test_simple_imputer_fit_transform_default_params(missing_values_penguins_df): + imputer = impute.SimpleImputer(strategy="mean") + result = imputer.fit_transform( + missing_values_penguins_df[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ).to_pandas() + + expected = pd.DataFrame( + { + "imputer_culmen_length_mm": [39.5, 38.5, 37.9], + "imputer_culmen_depth_mm": [17.65, 17.2, 18.1], + "imputer_flipper_length_mm": [184.5, 181.0, 188.0], + }, + dtype="Float64", + index=pd.Index([0, 1, 2], dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected) + + +def test_simple_imputer_series(missing_values_penguins_df): + imputer = impute.SimpleImputer(strategy="mean") + imputer.fit(missing_values_penguins_df["culmen_depth_mm"]) + + result = imputer.transform( + missing_values_penguins_df["culmen_depth_mm"] + ).to_pandas() + + expected = pd.DataFrame( + { + "imputer_culmen_depth_mm": [17.65, 17.2, 18.1], + }, + dtype="Float64", + index=pd.Index([0, 1, 2], dtype="Int64"), + ) + + pd.testing.assert_frame_equal(result, expected, rtol=0.1) + + +def test_simple_imputer_save_load_mean(missing_values_penguins_df, dataset_id): + transformer = impute.SimpleImputer(strategy="mean") + transformer.fit( + missing_values_penguins_df[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ) + + reloaded_transformer = transformer.to_gbq( + f"{dataset_id}.temp_configured_model", replace=True + ) + assert isinstance(reloaded_transformer, impute.SimpleImputer) + assert reloaded_transformer.strategy == transformer.strategy + assert reloaded_transformer._bqml_model is not None + + +def test_simple_imputer_save_load_most_frequent(missing_values_penguins_df, dataset_id): + transformer = impute.SimpleImputer(strategy="most_frequent") + transformer.fit( + missing_values_penguins_df[ + ["culmen_length_mm", "culmen_depth_mm", "flipper_length_mm"] + ] + ) + + reloaded_transformer = transformer.to_gbq( + f"{dataset_id}.temp_configured_model", replace=True + ) + assert isinstance(reloaded_transformer, impute.SimpleImputer) + assert reloaded_transformer.strategy == transformer.strategy + assert reloaded_transformer._bqml_model is not None diff --git a/tests/system/small/ml/test_llm.py b/tests/system/small/ml/test_llm.py index 8a6874b178..20e8dd0c19 100644 --- a/tests/system/small/ml/test_llm.py +++ b/tests/system/small/ml/test_llm.py @@ -303,10 +303,16 @@ def test_embedding_generator_predict_series_success( assert len(value) == 768 -def test_create_gemini_text_generator_model( - gemini_text_generator_model, dataset_id, bq_connection +@pytest.mark.parametrize( + "model_name", + ("gemini-pro", "gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514"), +) +def test_create_load_gemini_text_generator_model( + dataset_id, model_name, session, bq_connection ): - # Model creation doesn't return error + gemini_text_generator_model = llm.GeminiTextGenerator( + model_name=model_name, connection_name=bq_connection, session=session + ) assert gemini_text_generator_model is not None assert gemini_text_generator_model._bqml_model is not None @@ -316,12 +322,25 @@ def test_create_gemini_text_generator_model( ) assert f"{dataset_id}.temp_text_model" == reloaded_model._bqml_model.model_name assert reloaded_model.connection_name == bq_connection - - + assert reloaded_model.model_name == model_name + + +@pytest.mark.parametrize( + "model_name", + ( + "gemini-pro", + "gemini-1.5-pro-preview-0514", + # TODO(garrrettwu): enable when cl/637028077 is in prod. + # "gemini-1.5-flash-preview-0514" + ), +) @pytest.mark.flaky(retries=2) def test_gemini_text_generator_predict_default_params_success( - gemini_text_generator_model, llm_text_df + llm_text_df, model_name, session, bq_connection ): + gemini_text_generator_model = llm.GeminiTextGenerator( + model_name=model_name, connection_name=bq_connection, session=session + ) df = gemini_text_generator_model.predict(llm_text_df).to_pandas() assert df.shape == (3, 4) assert "ml_generate_text_llm_result" in df.columns @@ -329,10 +348,17 @@ def test_gemini_text_generator_predict_default_params_success( assert all(series.str.len() > 20) +@pytest.mark.parametrize( + "model_name", + ("gemini-pro", "gemini-1.5-pro-preview-0514", "gemini-1.5-flash-preview-0514"), +) @pytest.mark.flaky(retries=2) def test_gemini_text_generator_predict_with_params_success( - gemini_text_generator_model, llm_text_df + llm_text_df, model_name, session, bq_connection ): + gemini_text_generator_model = llm.GeminiTextGenerator( + model_name=model_name, connection_name=bq_connection, session=session + ) df = gemini_text_generator_model.predict( llm_text_df, temperature=0.5, max_output_tokens=100, top_k=20, top_p=0.5 ).to_pandas() diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 4b50922c0d..eae25bb027 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -1699,6 +1699,32 @@ def test_df_abs(scalars_dfs): assert_pandas_df_equal(bf_result, pd_result) +def test_df_pos(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = (+scalars_df[["int64_col", "numeric_col"]]).to_pandas() + pd_result = +scalars_pandas_df[["int64_col", "numeric_col"]] + + assert_pandas_df_equal(pd_result, bf_result) + + +def test_df_neg(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = (-scalars_df[["int64_col", "numeric_col"]]).to_pandas() + pd_result = -scalars_pandas_df[["int64_col", "numeric_col"]] + + assert_pandas_df_equal(pd_result, bf_result) + + +def test_df_invert(scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + columns = ["int64_col", "bool_col"] + + bf_result = (~scalars_df[columns]).to_pandas() + pd_result = ~scalars_pandas_df[columns] + + assert_pandas_df_equal(bf_result, pd_result) + + def test_df_isnull(scalars_dfs): scalars_df, scalars_pandas_df = scalars_dfs @@ -2104,14 +2130,12 @@ def test_series_binop_axis_index( (pd.Index([1000, 2000, 3000])), (bf_indexes.Index([1000, 2000, 3000])), (pd.Series((1000, 2000), index=["int64_too", "float64_col"])), - (series.Series((1000, 2000), index=["int64_too", "float64_col"])), ], ids=[ "tuple", "pd_index", "bf_index", "pd_series", - "bf_series", ], ) def test_listlike_binop_axis_1(scalars_dfs, input): @@ -2127,6 +2151,26 @@ def test_listlike_binop_axis_1(scalars_dfs, input): assert_pandas_df_equal(bf_result, pd_result, check_dtype=False) +def test_binop_with_self_aggregate(session, scalars_dfs): + scalars_df, scalars_pandas_df = scalars_dfs + + df_columns = ["int64_col", "float64_col", "int64_too"] + + # Ensure that this takes the optimized single-query path by counting executions + execution_count_before = session._execution_count + bf_df = scalars_df[df_columns] + bf_result = (bf_df - bf_df.mean()).to_pandas() + execution_count_after = session._execution_count + + pd_df = scalars_pandas_df[df_columns] + pd_result = pd_df - pd_df.mean() + + executions = execution_count_after - execution_count_before + + assert executions == 1 + assert_pandas_df_equal(bf_result, pd_result, check_dtype=False) + + @pytest.mark.parametrize( ("left_labels", "right_labels"), [ diff --git a/tests/system/small/test_empty_index.py b/tests/system/small/test_empty_index.py deleted file mode 100644 index 7a1715e3d1..0000000000 --- a/tests/system/small/test_empty_index.py +++ /dev/null @@ -1,212 +0,0 @@ -# Copyright 2024 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -import pandas as pd -import pytest - -import bigframes.exceptions -import bigframes.pandas as bpd -from tests.system.utils import skip_legacy_pandas - - -def test_empty_index_materialize( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = scalars_df_empty_index.to_pandas() - pd.testing.assert_frame_equal( - bf_result, scalars_pandas_df_default_index, check_index_type=False - ) - - -def test_empty_index_series_repr( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = scalars_df_empty_index["int64_too"].head(5).__repr__() - pd_result = ( - scalars_pandas_df_default_index["int64_too"] - .head(5) - .to_string(dtype=True, index=False, length=False, name=True) - ) - assert bf_result == pd_result - - -def test_empty_index_dataframe_repr( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = scalars_df_empty_index[["int64_too", "int64_col"]].head(5).__repr__() - pd_result = ( - scalars_pandas_df_default_index[["int64_too", "int64_col"]] - .head(5) - .to_string(index=False) - ) - assert bf_result == pd_result + "\n\n[5 rows x 2 columns]" - - -def test_empty_index_reset_index( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = scalars_df_empty_index.reset_index().to_pandas() - pd_result = scalars_pandas_df_default_index.reset_index(drop=True) - pd.testing.assert_frame_equal(bf_result, pd_result, check_index_type=False) - - -def test_empty_index_set_index(scalars_df_empty_index, scalars_pandas_df_default_index): - bf_result = scalars_df_empty_index.set_index("int64_col").to_pandas() - pd_result = scalars_pandas_df_default_index.set_index("int64_col") - pd.testing.assert_frame_equal(bf_result, pd_result) - - -def test_empty_index_concat(scalars_df_empty_index, scalars_pandas_df_default_index): - bf_result = bpd.concat( - [scalars_df_empty_index, scalars_df_empty_index], axis=0 - ).to_pandas() - pd_result = pd.concat( - [scalars_pandas_df_default_index, scalars_pandas_df_default_index], axis=0 - ) - pd.testing.assert_frame_equal(bf_result, pd_result.reset_index(drop=True)) - - -def test_empty_index_aggregate(scalars_df_empty_index, scalars_pandas_df_default_index): - bf_result = scalars_df_empty_index.count().to_pandas() - pd_result = scalars_pandas_df_default_index.count() - - pd_result.index = pd_result.index.astype("string[pyarrow]") - - pd.testing.assert_series_equal( - bf_result, pd_result, check_dtype=False, check_index_type=False - ) - - -def test_empty_index_groupby_aggregate( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = scalars_df_empty_index.groupby("int64_col").count().to_pandas() - pd_result = scalars_pandas_df_default_index.groupby("int64_col").count() - - pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) - - -@skip_legacy_pandas -def test_empty_index_analytic(scalars_df_empty_index, scalars_pandas_df_default_index): - bf_result = scalars_df_empty_index["int64_col"].cumsum().to_pandas() - pd_result = scalars_pandas_df_default_index["int64_col"].cumsum() - pd.testing.assert_series_equal( - bf_result, pd_result.reset_index(drop=True), check_dtype=False - ) - - -def test_empty_index_groupby_analytic( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = ( - scalars_df_empty_index.groupby("bool_col")["int64_col"].cummax().to_pandas() - ) - pd_result = scalars_pandas_df_default_index.groupby("bool_col")[ - "int64_col" - ].cummax() - pd.testing.assert_series_equal( - bf_result, pd_result.reset_index(drop=True), check_dtype=False - ) - - -@skip_legacy_pandas -def test_empty_index_stack(scalars_df_empty_index, scalars_pandas_df_default_index): - stacking_cols = ["int64_col", "int64_too"] - bf_result = scalars_df_empty_index[stacking_cols].stack().to_pandas() - pd_result = ( - scalars_pandas_df_default_index[stacking_cols] - .stack(future_stack=True) - .droplevel(level=0, axis=0) - ) - pd_result.index = pd_result.index.astype(bf_result.index.dtype) - pd.testing.assert_series_equal( - bf_result, - pd_result, - check_dtype=False, - ) - - -def test_empty_index_series_self_aligns( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = ( - scalars_df_empty_index["int64_col"] + scalars_df_empty_index["int64_too"] - ) - pd_result = ( - scalars_pandas_df_default_index["int64_col"] - + scalars_pandas_df_default_index["int64_too"] - ) - pd.testing.assert_series_equal( - bf_result.to_pandas(), pd_result.reset_index(drop=True), check_dtype=False - ) - - -def test_empty_index_df_self_aligns( - scalars_df_empty_index, scalars_pandas_df_default_index -): - bf_result = ( - scalars_df_empty_index[["int64_col", "float64_col"]] - + scalars_df_empty_index[["int64_col", "float64_col"]] - ) - pd_result = ( - scalars_pandas_df_default_index[["int64_col", "float64_col"]] - + scalars_pandas_df_default_index[["int64_col", "float64_col"]] - ) - pd.testing.assert_frame_equal( - bf_result.to_pandas(), pd_result.reset_index(drop=True), check_dtype=False - ) - - -def test_empty_index_df_concat(scalars_df_empty_index, scalars_pandas_df_default_index): - bf_result = bpd.concat([scalars_df_empty_index, scalars_df_empty_index]) - pd_result = pd.concat( - [scalars_pandas_df_default_index, scalars_pandas_df_default_index] - ) - pd.testing.assert_frame_equal( - bf_result.to_pandas(), pd_result.reset_index(drop=True), check_dtype=False - ) - - -def test_empty_index_align_error(scalars_df_empty_index): - with pytest.raises(bigframes.exceptions.NullIndexError): - _ = ( - scalars_df_empty_index["int64_col"] - + scalars_df_empty_index["int64_col"].cumsum() - ) - - -def test_empty_index_loc_error(scalars_df_empty_index): - with pytest.raises(bigframes.exceptions.NullIndexError): - scalars_df_empty_index["int64_col"].loc[1] - - -def test_empty_index_at_error(scalars_df_empty_index): - with pytest.raises(bigframes.exceptions.NullIndexError): - scalars_df_empty_index["int64_col"].at[1] - - -def test_empty_index_idxmin_error(scalars_df_empty_index): - with pytest.raises(bigframes.exceptions.NullIndexError): - scalars_df_empty_index[["int64_col", "int64_too"]].idxmin() - - -def test_empty_index_index_property(scalars_df_empty_index): - with pytest.raises(bigframes.exceptions.NullIndexError): - _ = scalars_df_empty_index.index - - -def test_empty_index_transpose(scalars_df_empty_index): - with pytest.raises(bigframes.exceptions.NullIndexError): - _ = scalars_df_empty_index.T diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 7b36a06f49..b332d48574 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -18,6 +18,10 @@ import bigframes.pandas as bpd from tests.system.utils import assert_pandas_df_equal +# ================= +# DataFrame.groupby +# ================= + @pytest.mark.parametrize( ("operator"), @@ -140,6 +144,23 @@ def test_dataframe_groupby_agg_list(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) +def test_dataframe_groupby_agg_list_w_column_multi_index( + scalars_df_index, scalars_pandas_df_index +): + columns = ["int64_too", "string_col", "bool_col"] + multi_columns = pd.MultiIndex.from_tuples(zip(["a", "b", "a"], columns)) + bf_df = scalars_df_index[columns].copy() + bf_df.columns = multi_columns + pd_df = scalars_pandas_df_index[columns].copy() + pd_df.columns = multi_columns + + bf_result = bf_df.groupby(level=0).agg(["count", "min"]) + pd_result = pd_df.groupby(level=0).agg(["count", "min"]) + + bf_result_computed = bf_result.to_pandas() + pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) + + @pytest.mark.parametrize( ("as_index"), [ @@ -269,21 +290,26 @@ def test_dataframe_groupby_analytic( pd.testing.assert_frame_equal(pd_result, bf_result_computed, check_dtype=False) -def test_series_groupby_skew(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.groupby("bool_col")["int64_too"].skew().to_pandas() - pd_result = scalars_pandas_df_index.groupby("bool_col")["int64_too"].skew() +def test_dataframe_groupby_size_as_index_false( + scalars_df_index, scalars_pandas_df_index +): + bf_result = scalars_df_index.groupby("string_col", as_index=False).size() + bf_result_computed = bf_result.to_pandas() + pd_result = scalars_pandas_df_index.groupby("string_col", as_index=False).size() - pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + pd.testing.assert_frame_equal( + pd_result, bf_result_computed, check_dtype=False, check_index_type=False + ) -def test_series_groupby_kurt(scalars_df_index, scalars_pandas_df_index): - bf_result = scalars_df_index.groupby("bool_col")["int64_too"].kurt().to_pandas() - # Pandas doesn't have groupby.kurt yet: https://github.com/pandas-dev/pandas/issues/40139 - pd_result = scalars_pandas_df_index.groupby("bool_col")["int64_too"].apply( - pd.Series.kurt - ) +def test_dataframe_groupby_size_as_index_true( + scalars_df_index, scalars_pandas_df_index +): + bf_result = scalars_df_index.groupby("string_col", as_index=True).size() + pd_result = scalars_pandas_df_index.groupby("string_col", as_index=True).size() + bf_result_computed = bf_result.to_pandas() - pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + pd.testing.assert_series_equal(pd_result, bf_result_computed, check_dtype=False) def test_dataframe_groupby_skew(scalars_df_index, scalars_pandas_df_index): @@ -356,6 +382,30 @@ def test_dataframe_groupby_getitem_list( pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False) +def test_dataframe_groupby_nonnumeric_with_mean(): + df = pd.DataFrame( + { + "key1": ["a", "a", "a", "b"], + "key2": ["a", "a", "c", "c"], + "key3": [1, 2, 3, 4], + "key4": [1.6, 2, 3, 4], + } + ) + pd_result = df.groupby(["key1", "key2"]).mean() + + with bpd.option_context("bigquery.location", "US"): + bf_result = bpd.DataFrame(df).groupby(["key1", "key2"]).mean().to_pandas() + + pd.testing.assert_frame_equal( + pd_result, bf_result, check_index_type=False, check_dtype=False + ) + + +# ============== +# Series.groupby +# ============== + + def test_series_groupby_agg_string(scalars_df_index, scalars_pandas_df_index): bf_result = ( scalars_df_index["int64_col"] @@ -392,21 +442,49 @@ def test_series_groupby_agg_list(scalars_df_index, scalars_pandas_df_index): ) -def test_dataframe_groupby_nonnumeric_with_mean(): - df = pd.DataFrame( - { - "key1": ["a", "a", "a", "b"], - "key2": ["a", "a", "c", "c"], - "key3": [1, 2, 3, 4], - "key4": [1.6, 2, 3, 4], - } +def test_series_groupby_kurt(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_too"] + .groupby(scalars_df_index["bool_col"]) + .kurt() + .to_pandas() + ) + # Pandas doesn't have groupby.kurt yet: https://github.com/pandas-dev/pandas/issues/40139 + pd_result = scalars_pandas_df_index.groupby("bool_col")["int64_too"].apply( + pd.Series.kurt ) - pd_result = df.groupby(["key1", "key2"]).mean() - bf_result = bpd.DataFrame(df).groupby(["key1", "key2"]).mean().to_pandas() - pd.testing.assert_frame_equal( - pd_result, bf_result, check_index_type=False, check_dtype=False + pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) + + +def test_series_groupby_size(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_too"].groupby(scalars_df_index["bool_col"]).size() + ) + pd_result = ( + scalars_pandas_df_index["int64_too"] + .groupby(scalars_pandas_df_index["bool_col"]) + .size() ) + bf_result_computed = bf_result.to_pandas() + + pd.testing.assert_series_equal(pd_result, bf_result_computed, check_dtype=False) + + +def test_series_groupby_skew(scalars_df_index, scalars_pandas_df_index): + bf_result = ( + scalars_df_index["int64_too"] + .groupby(scalars_df_index["bool_col"]) + .skew() + .to_pandas() + ) + pd_result = ( + scalars_pandas_df_index["int64_too"] + .groupby(scalars_pandas_df_index["bool_col"]) + .skew() + ) + + pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) @pytest.mark.parametrize( diff --git a/tests/system/small/test_null_index.py b/tests/system/small/test_null_index.py new file mode 100644 index 0000000000..27a3d8dffe --- /dev/null +++ b/tests/system/small/test_null_index.py @@ -0,0 +1,288 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pandas as pd +import pytest + +import bigframes.exceptions +import bigframes.pandas as bpd +from tests.system.utils import skip_legacy_pandas + + +def test_null_index_materialize(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = scalars_df_null_index.to_pandas() + pd.testing.assert_frame_equal( + bf_result, scalars_pandas_df_default_index, check_index_type=False + ) + + +def test_null_index_series_repr(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = scalars_df_null_index["int64_too"].head(5).__repr__() + pd_result = ( + scalars_pandas_df_default_index["int64_too"] + .head(5) + .to_string(dtype=True, index=False, length=False, name=True) + ) + assert bf_result == pd_result + + +def test_null_index_dataframe_repr( + scalars_df_null_index, scalars_pandas_df_default_index +): + bf_result = scalars_df_null_index[["int64_too", "int64_col"]].head(5).__repr__() + pd_result = ( + scalars_pandas_df_default_index[["int64_too", "int64_col"]] + .head(5) + .to_string(index=False) + ) + assert bf_result == pd_result + "\n\n[5 rows x 2 columns]" + + +def test_null_index_reset_index(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = scalars_df_null_index.reset_index().to_pandas() + pd_result = scalars_pandas_df_default_index.reset_index(drop=True) + pd.testing.assert_frame_equal(bf_result, pd_result, check_index_type=False) + + +def test_null_index_set_index(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = scalars_df_null_index.set_index("int64_col").to_pandas() + pd_result = scalars_pandas_df_default_index.set_index("int64_col") + pd.testing.assert_frame_equal(bf_result, pd_result) + + +def test_null_index_concat(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = bpd.concat( + [scalars_df_null_index, scalars_df_null_index], axis=0 + ).to_pandas() + pd_result = pd.concat( + [scalars_pandas_df_default_index, scalars_pandas_df_default_index], axis=0 + ) + pd.testing.assert_frame_equal(bf_result, pd_result.reset_index(drop=True)) + + +def test_null_index_aggregate(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = scalars_df_null_index.count().to_pandas() + pd_result = scalars_pandas_df_default_index.count() + + pd_result.index = pd_result.index.astype("string[pyarrow]") + + pd.testing.assert_series_equal( + bf_result, pd_result, check_dtype=False, check_index_type=False + ) + + +def test_null_index_groupby_aggregate( + scalars_df_null_index, scalars_pandas_df_default_index +): + bf_result = scalars_df_null_index.groupby("int64_col").count().to_pandas() + pd_result = scalars_pandas_df_default_index.groupby("int64_col").count() + + pd.testing.assert_frame_equal(bf_result, pd_result, check_dtype=False) + + +@skip_legacy_pandas +def test_null_index_analytic(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = scalars_df_null_index["int64_col"].cumsum().to_pandas() + pd_result = scalars_pandas_df_default_index["int64_col"].cumsum() + pd.testing.assert_series_equal( + bf_result, pd_result.reset_index(drop=True), check_dtype=False + ) + + +def test_null_index_groupby_analytic( + scalars_df_null_index, scalars_pandas_df_default_index +): + bf_result = ( + scalars_df_null_index.groupby("bool_col")["int64_col"].cummax().to_pandas() + ) + pd_result = scalars_pandas_df_default_index.groupby("bool_col")[ + "int64_col" + ].cummax() + pd.testing.assert_series_equal( + bf_result, pd_result.reset_index(drop=True), check_dtype=False + ) + + +def test_null_index_merge_left_null_index_object( + scalars_df_null_index, scalars_df_default_index, scalars_pandas_df_default_index +): + df1 = scalars_df_null_index[scalars_df_null_index["int64_col"] > 0] + df1_pd = scalars_pandas_df_default_index[ + scalars_pandas_df_default_index["int64_col"] > 0 + ] + assert not df1._has_index + df2 = scalars_df_default_index[scalars_df_default_index["int64_col"] <= 55555] + df2_pd = scalars_pandas_df_default_index[ + scalars_pandas_df_default_index["int64_col"] <= 55555 + ] + assert df2._has_index + + got = df1.merge(df2, how="inner", on="bool_col") + expected = df1_pd.merge(df2_pd, how="inner", on="bool_col") + + # Combining any NULL index object should result in a NULL index. + # This keeps us from generating an index if the user joins a large + # BigQuery table against small local data, for example. + assert not got._has_index + assert got.shape == expected.shape + + +def test_null_index_merge_right_null_index_object( + scalars_df_null_index, scalars_df_default_index, scalars_pandas_df_default_index +): + df1 = scalars_df_default_index[scalars_df_default_index["int64_col"] > 0] + df1_pd = scalars_pandas_df_default_index[ + scalars_pandas_df_default_index["int64_col"] > 0 + ] + assert df1._has_index + df2 = scalars_df_null_index[scalars_df_null_index["int64_col"] <= 55555] + df2_pd = scalars_pandas_df_default_index[ + scalars_pandas_df_default_index["int64_col"] <= 55555 + ] + assert not df2._has_index + + got = df1.merge(df2, how="left", on="bool_col") + expected = df1_pd.merge(df2_pd, how="left", on="bool_col") + + # Combining any NULL index object should result in a NULL index. + # This keeps us from generating an index if the user joins a large + # BigQuery table against small local data, for example. + assert not got._has_index + assert got.shape == expected.shape + + +def test_null_index_merge_two_null_index_objects( + scalars_df_null_index, scalars_pandas_df_default_index +): + df1 = scalars_df_null_index[scalars_df_null_index["int64_col"] > 0] + df1_pd = scalars_pandas_df_default_index[ + scalars_pandas_df_default_index["int64_col"] > 0 + ] + assert not df1._has_index + df2 = scalars_df_null_index[scalars_df_null_index["int64_col"] <= 55555] + df2_pd = scalars_pandas_df_default_index[ + scalars_pandas_df_default_index["int64_col"] <= 55555 + ] + assert not df2._has_index + + got = df1.merge(df2, how="outer", on="bool_col") + expected = df1_pd.merge(df2_pd, how="outer", on="bool_col") + + assert not got._has_index + assert got.shape == expected.shape + + +@skip_legacy_pandas +def test_null_index_stack(scalars_df_null_index, scalars_pandas_df_default_index): + stacking_cols = ["int64_col", "int64_too"] + bf_result = scalars_df_null_index[stacking_cols].stack().to_pandas() + pd_result = ( + scalars_pandas_df_default_index[stacking_cols] + .stack(future_stack=True) + .droplevel(level=0, axis=0) + ) + pd_result.index = pd_result.index.astype(bf_result.index.dtype) + pd.testing.assert_series_equal( + bf_result, + pd_result, + check_dtype=False, + ) + + +def test_null_index_series_self_aligns( + scalars_df_null_index, scalars_pandas_df_default_index +): + bf_result = scalars_df_null_index["int64_col"] + scalars_df_null_index["int64_too"] + pd_result = ( + scalars_pandas_df_default_index["int64_col"] + + scalars_pandas_df_default_index["int64_too"] + ) + pd.testing.assert_series_equal( + bf_result.to_pandas(), pd_result.reset_index(drop=True), check_dtype=False + ) + + +def test_null_index_df_self_aligns( + scalars_df_null_index, scalars_pandas_df_default_index +): + bf_result = ( + scalars_df_null_index[["int64_col", "float64_col"]] + + scalars_df_null_index[["int64_col", "float64_col"]] + ) + pd_result = ( + scalars_pandas_df_default_index[["int64_col", "float64_col"]] + + scalars_pandas_df_default_index[["int64_col", "float64_col"]] + ) + pd.testing.assert_frame_equal( + bf_result.to_pandas(), pd_result.reset_index(drop=True), check_dtype=False + ) + + +def test_null_index_setitem(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = scalars_df_null_index.copy() + bf_result["new_col"] = ( + scalars_df_null_index["int64_col"] + scalars_df_null_index["float64_col"] + ) + pd_result = scalars_pandas_df_default_index.copy() + pd_result["new_col"] = ( + scalars_pandas_df_default_index["int64_col"] + + scalars_pandas_df_default_index["float64_col"] + ) + pd.testing.assert_frame_equal( + bf_result.to_pandas(), pd_result.reset_index(drop=True), check_dtype=False + ) + + +def test_null_index_df_concat(scalars_df_null_index, scalars_pandas_df_default_index): + bf_result = bpd.concat([scalars_df_null_index, scalars_df_null_index]) + pd_result = pd.concat( + [scalars_pandas_df_default_index, scalars_pandas_df_default_index] + ) + pd.testing.assert_frame_equal( + bf_result.to_pandas(), pd_result.reset_index(drop=True), check_dtype=False + ) + + +def test_null_index_align_error(scalars_df_null_index): + with pytest.raises(bigframes.exceptions.NullIndexError): + _ = ( + scalars_df_null_index["int64_col"] + + scalars_df_null_index["int64_col"].cumsum() + ) + + +def test_null_index_loc_error(scalars_df_null_index): + with pytest.raises(bigframes.exceptions.NullIndexError): + scalars_df_null_index["int64_col"].loc[1] + + +def test_null_index_at_error(scalars_df_null_index): + with pytest.raises(bigframes.exceptions.NullIndexError): + scalars_df_null_index["int64_col"].at[1] + + +def test_null_index_idxmin_error(scalars_df_null_index): + with pytest.raises(bigframes.exceptions.NullIndexError): + scalars_df_null_index[["int64_col", "int64_too"]].idxmin() + + +def test_null_index_index_property(scalars_df_null_index): + with pytest.raises(bigframes.exceptions.NullIndexError): + _ = scalars_df_null_index.index + + +def test_null_index_transpose(scalars_df_null_index): + with pytest.raises(bigframes.exceptions.NullIndexError): + _ = scalars_df_null_index.T diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index 4a39e75ff9..d2ee4411f4 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -67,7 +67,7 @@ def bq_cf_connection_location_project(bigquery_client) -> str: @pytest.fixture(scope="module") def bq_cf_connection_location_project_mismatched() -> str: - """Pre-created BQ connection in the migframes-metrics project in US location, + """Pre-created BQ connection in the bigframes-metrics project in US location, in format PROJECT_ID.LOCATION.CONNECTION_NAME, used to invoke cloud function. $ bq show --connection --location=us --project_id=PROJECT_ID bigframes-rf-conn @@ -108,11 +108,15 @@ def test_remote_function_direct_no_session_param( reuse=True, ) def square(x): - # This executes on a remote function, where coverage isn't tracked. - return x * x # pragma: NO COVER + return x * x - assert square.bigframes_remote_function - assert square.bigframes_cloud_function + # Function should still work normally. + assert square(2) == 4 + + # Function should have extra metadata attached for remote execution. + assert hasattr(square, "bigframes_remote_function") + assert hasattr(square, "bigframes_cloud_function") + assert hasattr(square, "ibis_node") scalars_df, scalars_pandas_df = scalars_dfs @@ -161,8 +165,10 @@ def test_remote_function_direct_no_session_param_location_specified( reuse=True, ) def square(x): - # This executes on a remote function, where coverage isn't tracked. - return x * x # pragma: NO COVER + return x * x + + # Function should still work normally. + assert square(2) == 4 scalars_df, scalars_pandas_df = scalars_dfs @@ -197,7 +203,10 @@ def test_remote_function_direct_no_session_param_location_mismatched( dataset_id_permanent, bq_cf_connection_location_mismatched, ): - with pytest.raises(ValueError): + with pytest.raises( + ValueError, + match=re.escape("The location does not match BigQuery connection location:"), + ): @rf.remote_function( [int], @@ -212,7 +221,8 @@ def test_remote_function_direct_no_session_param_location_mismatched( reuse=True, ) def square(x): - # This executes on a remote function, where coverage isn't tracked. + # Not expected to reach this code, as the location of the + # connection doesn't match the location of the dataset. return x * x # pragma: NO COVER @@ -239,8 +249,10 @@ def test_remote_function_direct_no_session_param_location_project_specified( reuse=True, ) def square(x): - # This executes on a remote function, where coverage isn't tracked. - return x * x # pragma: NO COVER + return x * x + + # Function should still work normally. + assert square(2) == 4 scalars_df, scalars_pandas_df = scalars_dfs @@ -275,7 +287,12 @@ def test_remote_function_direct_no_session_param_project_mismatched( dataset_id_permanent, bq_cf_connection_location_project_mismatched, ): - with pytest.raises(ValueError): + with pytest.raises( + ValueError, + match=re.escape( + "The project_id does not match BigQuery connection gcp_project_id:" + ), + ): @rf.remote_function( [int], @@ -290,7 +307,8 @@ def test_remote_function_direct_no_session_param_project_mismatched( reuse=True, ) def square(x): - # This executes on a remote function, where coverage isn't tracked. + # Not expected to reach this code, as the project of the + # connection doesn't match the project of the dataset. return x * x # pragma: NO COVER @@ -302,8 +320,10 @@ def test_remote_function_direct_session_param(session_with_bq_connection, scalar session=session_with_bq_connection, ) def square(x): - # This executes on a remote function, where coverage isn't tracked. - return x * x # pragma: NO COVER + return x * x + + # Function should still work normally. + assert square(2) == 4 scalars_df, scalars_pandas_df = scalars_dfs @@ -340,8 +360,10 @@ def test_remote_function_via_session_default(session_with_bq_connection, scalars # cloud function would be common and quickly reused. @session_with_bq_connection.remote_function([int], int) def square(x): - # This executes on a remote function, where coverage isn't tracked. - return x * x # pragma: NO COVER + return x * x + + # Function should still work normally. + assert square(2) == 4 scalars_df, scalars_pandas_df = scalars_dfs @@ -380,8 +402,10 @@ def test_remote_function_via_session_with_overrides( reuse=True, ) def square(x): - # This executes on a remote function, where coverage isn't tracked. - return x * x # pragma: NO COVER + return x * x + + # Function should still work normally. + assert square(2) == 4 scalars_df, scalars_pandas_df = scalars_dfs @@ -508,7 +532,7 @@ def test_skip_bq_connection_check(dataset_id_permanent): @session.remote_function([int], int, dataset=dataset_id_permanent) def add_one(x): - # This executes on a remote function, where coverage isn't tracked. + # Not expected to reach this code, as the connection doesn't exist. return x + 1 # pragma: NO COVER @@ -546,8 +570,10 @@ def test_read_gbq_function_like_original( reuse=True, ) def square1(x): - # This executes on a remote function, where coverage isn't tracked. - return x * x # pragma: NO COVER + return x * x + + # Function should still work normally. + assert square1(2) == 4 square2 = rf.read_gbq_function( function_name=square1.bigframes_remote_function, @@ -701,7 +727,7 @@ def add_ints(row): with pytest.warns( bigframes.exceptions.PreviewWarning, - match="input_types=Series scenario is in preview.", + match="input_types=Series is in preview.", ): add_ints_remote = session.remote_function(bigframes.series.Series, int)( add_ints diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index dbc8ddec6f..3e21418f2f 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -354,6 +354,36 @@ def test_abs(scalars_dfs, col_name): assert_series_equal(pd_result, bf_result) +@pytest.mark.parametrize( + ("col_name",), + ( + ("float64_col",), + ("int64_too",), + ), +) +def test_series_pos(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = (+scalars_df[col_name]).to_pandas() + pd_result = +scalars_pandas_df[col_name] + + assert_series_equal(pd_result, bf_result) + + +@pytest.mark.parametrize( + ("col_name",), + ( + ("float64_col",), + ("int64_too",), + ), +) +def test_series_neg(scalars_dfs, col_name): + scalars_df, scalars_pandas_df = scalars_dfs + bf_result = (-scalars_df[col_name]).to_pandas() + pd_result = -scalars_pandas_df[col_name] + + assert_series_equal(pd_result, bf_result) + + @pytest.mark.parametrize( ("col_name",), ( @@ -678,10 +708,12 @@ def test_series_pow_scalar_reverse(scalars_dfs): [ (lambda x, y: x & y), (lambda x, y: x | y), + (lambda x, y: x ^ y), ], ids=[ "and", "or", + "xor", ], ) @pytest.mark.parametrize(("other_scalar"), [True, False, pd.NA]) @@ -714,6 +746,7 @@ def test_series_bool_bool_operators_scalar( (lambda x, y: x // y), (lambda x, y: x & y), (lambda x, y: x | y), + (lambda x, y: x ^ y), ], ids=[ "add", @@ -728,6 +761,7 @@ def test_series_bool_bool_operators_scalar( "floordivide", "bitwise_and", "bitwise_or", + "bitwise_xor", ], ) def test_series_int_int_operators_series(scalars_dfs, operator): diff --git a/tests/system/small/test_session.py b/tests/system/small/test_session.py index 2b7c6178ff..9631e0c7ab 100644 --- a/tests/system/small/test_session.py +++ b/tests/system/small/test_session.py @@ -20,6 +20,7 @@ import time import typing from typing import List, Optional, Sequence +import warnings import google import google.cloud.bigquery as bigquery @@ -386,6 +387,13 @@ def test_read_gbq_twice_with_same_timestamp(session, penguins_table_id): assert df3 is not None +def test_read_gbq_on_linked_dataset_warns(session): + with warnings.catch_warnings(record=True) as warned: + session.read_gbq("bigframes-dev.thelook_ecommerce.orders") + assert len(warned) == 1 + assert warned[0].category == bigframes.exceptions.TimeTravelDisabledWarning + + def test_read_gbq_table_clustered_with_filter(session: bigframes.Session): df = session.read_gbq_table( "bigquery-public-data.cloud_storage_geo_index.landsat_index", @@ -537,18 +545,22 @@ def test_read_gbq_with_configuration( def test_read_gbq_with_custom_global_labels( session: bigframes.Session, scalars_table_id: str ): - bigframes.options.compute.assign_extra_query_labels(test1=1, test2="abc") - bigframes.options.compute.extra_query_labels["test3"] = False + # Ensure we use thread-local variables to avoid conflicts with parallel tests. + with bigframes.option_context("compute.extra_query_labels", {}): + bigframes.options.compute.assign_extra_query_labels(test1=1, test2="abc") + bigframes.options.compute.extra_query_labels["test3"] = False - job_labels = session.read_gbq(scalars_table_id).query_job.labels # type:ignore - expected_labels = {"test1": "1", "test2": "abc", "test3": "false"} + job_labels = session.read_gbq(scalars_table_id).query_job.labels # type:ignore + expected_labels = {"test1": "1", "test2": "abc", "test3": "false"} - assert all(job_labels.get(key) == value for key, value in expected_labels.items()) + # All jobs should include a bigframes-api key. See internal issue 336521938. + assert "bigframes-api" in job_labels - del bigframes.options.compute.extra_query_labels["test1"] - del bigframes.options.compute.extra_query_labels["test2"] - del bigframes.options.compute.extra_query_labels["test3"] + assert all( + job_labels.get(key) == value for key, value in expected_labels.items() + ) + # No labels outside of the option_context. assert len(bigframes.options.compute.extra_query_labels) == 0 diff --git a/tests/unit/ml/test_sql.py b/tests/unit/ml/test_sql.py index 07b247fb41..01f173812c 100644 --- a/tests/unit/ml/test_sql.py +++ b/tests/unit/ml/test_sql.py @@ -106,6 +106,13 @@ def test_min_max_scaler_correct( assert sql == "ML.MIN_MAX_SCALER(col_a) OVER() AS scaled_col_a" +def test_imputer_correct( + base_sql_generator: ml_sql.BaseSqlGenerator, +): + sql = base_sql_generator.ml_imputer("col_a", "mean", "scaled_col_a") + assert sql == "ML.IMPUTER(col_a, 'mean') OVER() AS scaled_col_a" + + def test_k_bins_discretizer_correct( base_sql_generator: ml_sql.BaseSqlGenerator, ): diff --git a/tests/unit/resources.py b/tests/unit/resources.py index 623af93f65..84699459e6 100644 --- a/tests/unit/resources.py +++ b/tests/unit/resources.py @@ -23,6 +23,7 @@ import pytest import bigframes +import bigframes.clients import bigframes.core as core import bigframes.core.ordering import bigframes.dataframe @@ -39,6 +40,7 @@ def create_bigquery_session( session_id: str = "abcxyz", table_schema: Sequence[google.cloud.bigquery.SchemaField] = TEST_SCHEMA, anonymous_dataset: Optional[google.cloud.bigquery.DatasetReference] = None, + location: str = "test-region", ) -> bigframes.Session: credentials = mock.create_autospec( google.auth.credentials.Credentials, instance=True @@ -53,11 +55,12 @@ def create_bigquery_session( if bqclient is None: bqclient = mock.create_autospec(google.cloud.bigquery.Client, instance=True) bqclient.project = "test-project" + bqclient.location = location # Mock the location. table = mock.create_autospec(google.cloud.bigquery.Table, instance=True) table._properties = {} - type(table).location = mock.PropertyMock(return_value="test-region") + type(table).location = mock.PropertyMock(return_value=location) type(table).schema = mock.PropertyMock(return_value=table_schema) type(table).reference = mock.PropertyMock( return_value=anonymous_dataset.table("test_table") @@ -93,10 +96,11 @@ def query_mock(query, *args, **kwargs): type(clients_provider).bqclient = mock.PropertyMock(return_value=bqclient) clients_provider._credentials = credentials - bqoptions = bigframes.BigQueryOptions( - credentials=credentials, location="test-region" - ) + bqoptions = bigframes.BigQueryOptions(credentials=credentials, location=location) session = bigframes.Session(context=bqoptions, clients_provider=clients_provider) + session._bq_connection_manager = mock.create_autospec( + bigframes.clients.BqConnectionManager, instance=True + ) return session diff --git a/tests/unit/session/test_io_bigquery.py b/tests/unit/session/test_io_bigquery.py index 8ba13a7276..1d6be3dff8 100644 --- a/tests/unit/session/test_io_bigquery.py +++ b/tests/unit/session/test_io_bigquery.py @@ -36,6 +36,25 @@ def test_create_job_configs_labels_is_none(): assert labels == expected_dict +def test_create_job_configs_labels_always_includes_bigframes_api(): + labels = io_bq.create_job_configs_labels(None, []) + assert labels == { + "bigframes-api": "unknown", + } + + +def test_create_job_configs_labels_includes_extra_query_labels(): + user_labels = {"my-label-1": "my-value-1", "my-label-2": "my-value-2"} + + with bigframes.option_context("compute.extra_query_labels", user_labels): + labels = io_bq.create_job_configs_labels(None, []) + assert labels == { + "my-label-1": "my-value-1", + "my-label-2": "my-value-2", + "bigframes-api": "unknown", + } + + def test_create_job_configs_labels_length_limit_not_met(): cur_labels = { "source": "bigquery-dataframes-temp", @@ -105,7 +124,7 @@ def test_create_job_configs_labels_length_limit_met(): "bigframes-api": "read_pandas", "source": "bigquery-dataframes-temp", } - for i in range(100): + for i in range(61): key = f"bigframes-api-test-{i}" value = f"test{i}" cur_labels[key] = value @@ -208,7 +227,6 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) @pytest.mark.parametrize( ( "query_or_table", - "index_cols", "columns", "filters", "max_results", @@ -218,8 +236,7 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) [ pytest.param( "test_table", - ["row_index"], - ["string_col"], + ["row_index", "string_col"], [ (("rowindex", "not in", [0, 6]),), (("string_col", "in", ["Hello, World!", "こんにちは"]),), @@ -242,8 +259,7 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) FROM `test_table` AS t """ ), - ["rowindex"], - ["string_col"], + ["rowindex", "string_col"], [ ("rowindex", "<", 4), ("string_col", "==", "Hello, World!"), @@ -264,7 +280,6 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) ), pytest.param( "test_table", - [], ["col_a", "col_b"], [], None, # max_results @@ -275,7 +290,6 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) pytest.param( "test_table", [], - [], [("date_col", ">", "2022-10-20")], None, # max_results None, # time_travel_timestampe @@ -286,7 +300,6 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) "test_table*", [], [], - [], None, # max_results None, # time_travel_timestampe "SELECT * FROM `test_table*`", @@ -295,7 +308,6 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) pytest.param( "test_table*", [], - [], [("_TABLE_SUFFIX", ">", "2022-10-20")], None, # max_results None, # time_travel_timestampe @@ -306,7 +318,6 @@ def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str) ) def test_to_query( query_or_table, - index_cols, columns, filters, max_results, @@ -315,9 +326,8 @@ def test_to_query( ): query = io_bq.to_query( query_or_table, - index_cols=index_cols, columns=columns, - filters=filters, + sql_predicate=io_bq.compile_filters(filters), max_results=max_results, time_travel_timestamp=time_travel_timestamp, ) @@ -337,9 +347,8 @@ def test_to_query_fails_with_bad_filters(filters, expected_message): with pytest.raises(ValueError, match=re.escape(expected_message)): io_bq.to_query( "test_table", - index_cols=(), columns=(), - filters=filters, + sql_predicate=io_bq.compile_filters(filters), max_results=None, time_travel_timestamp=None, ) diff --git a/tests/unit/session/test_read_gbq_table.py b/tests/unit/session/test_read_gbq_table.py index 52c86cd1e4..6933957e53 100644 --- a/tests/unit/session/test_read_gbq_table.py +++ b/tests/unit/session/test_read_gbq_table.py @@ -14,11 +14,9 @@ """Unit tests for read_gbq_table helper functions.""" -import datetime import unittest.mock as mock import google.cloud.bigquery -import google.cloud.bigquery as bigquery import pytest import bigframes.session._io.bigquery.read_gbq_table as bf_read_gbq_table @@ -26,30 +24,6 @@ from .. import resources -def test_get_ibis_time_travel_table_doesnt_timetravel_anonymous_datasets(): - bqsession = resources.create_bigquery_session() - - table_ref = bigquery.TableReference.from_string( - "my-test-project._e8166e0cdb.anonbb92cd" - ) - - table_expression = bf_read_gbq_table.get_ibis_time_travel_table( - bqsession.ibis_client, - table_ref, - index_cols=(), - columns=(), - filters=(), - time_travel_timestamp=datetime.datetime.now(datetime.timezone.utc), - ) - sql = table_expression.compile() - - # Anonymous query results tables don't support time travel. - assert "SYSTEM_TIME" not in sql - - # Need fully-qualified table name. - assert "my-test-project" in sql - - @pytest.mark.parametrize( ("index_cols", "primary_keys", "values_distinct", "expected"), ( diff --git a/tests/unit/session/test_session.py b/tests/unit/session/test_session.py index 4f5daebc87..31029abd67 100644 --- a/tests/unit/session/test_session.py +++ b/tests/unit/session/test_session.py @@ -35,6 +35,14 @@ "datasetId": "my_dataset", "tableId": "my_table", } +SCHEMA = { + "fields": [ + {"name": "col1", "type": "INTEGER"}, + {"name": "col2", "type": "INTEGER"}, + {"name": "col3", "type": "INTEGER"}, + {"name": "col4", "type": "INTEGER"}, + ] +} CLUSTERED_OR_PARTITIONED_TABLES = [ pytest.param( google.cloud.bigquery.Table.from_api_repr( @@ -43,6 +51,7 @@ "clustering": { "fields": ["col1", "col2"], }, + "schema": SCHEMA, }, ), id="clustered", @@ -59,6 +68,7 @@ "interval": 1, }, }, + "schema": SCHEMA, }, ), id="range-partitioned", @@ -71,6 +81,7 @@ "type": "MONTH", "field": "col1", }, + "schema": SCHEMA, }, ), id="time-partitioned", @@ -86,6 +97,7 @@ "type": "MONTH", "field": "col1", }, + "schema": SCHEMA, }, ), id="time-partitioned-and-clustered", @@ -163,22 +175,18 @@ def test_read_gbq_cached_table(): google.cloud.bigquery.DatasetReference("my-project", "my_dataset"), "my_table", ) - table = google.cloud.bigquery.Table(table_ref) + table = google.cloud.bigquery.Table( + table_ref, (google.cloud.bigquery.SchemaField("col", "INTEGER"),) + ) + table._properties["location"] = session._location + table._properties["numRows"] = "1000000000" table._properties["location"] = session._location session._df_snapshot[table_ref] = ( datetime.datetime(1999, 1, 2, 3, 4, 5, 678901, tzinfo=datetime.timezone.utc), table, ) - def get_table_mock(table_ref): - table = google.cloud.bigquery.Table( - table_ref, (google.cloud.bigquery.SchemaField("col", "INTEGER"),) - ) - table._properties["numRows"] = "1000000000" - table._properties["location"] = session._location - return table - - session.bqclient.get_table = get_table_mock + session.bqclient.get_table.return_value = table session.bqclient.query_and_wait.return_value = ( {"total_count": 3, "distinct_count": 2}, ) @@ -373,7 +381,7 @@ def query_mock(query, *args, **kwargs): return session_query_mock(query, *args, **kwargs) - session.bqclient.query = query_mock + session.bqclient.query_and_wait = query_mock def get_table_mock(table_ref): table = google.cloud.bigquery.Table( diff --git a/tests/unit/test_pandas.py b/tests/unit/test_pandas.py index 70639315be..54a7a79d3c 100644 --- a/tests/unit/test_pandas.py +++ b/tests/unit/test_pandas.py @@ -50,7 +50,7 @@ def all_session_methods(): [(method_name,) for method_name in all_session_methods()], ) def test_method_matches_session(method_name: str): - if sys.version_info <= (3, 10): + if sys.version_info < (3, 10): pytest.skip( "Need Python 3.10 to reconcile deferred annotations." ) # pragma: no cover diff --git a/tests/unit/test_remote_function.py b/tests/unit/test_remote_function.py index 1acff27c7f..1bd3f3b14f 100644 --- a/tests/unit/test_remote_function.py +++ b/tests/unit/test_remote_function.py @@ -12,10 +12,51 @@ # See the License for the specific language governing permissions and # limitations under the License. +import re + import bigframes_vendored.ibis.backends.bigquery.datatypes as third_party_ibis_bqtypes from ibis.expr import datatypes as ibis_types +import pandas +import pytest import bigframes.dtypes +import bigframes.functions.remote_function +import bigframes.series +from tests.unit import resources + + +@pytest.mark.parametrize( + "series_type", + ( + pytest.param( + pandas.Series, + id="pandas.Series", + ), + pytest.param( + bigframes.series.Series, + id="bigframes.series.Series", + ), + ), +) +def test_series_input_types_to_str(series_type): + """Check that is_row_processor=True uses str as the input type to serialize a row.""" + session = resources.create_bigquery_session() + remote_function_decorator = bigframes.functions.remote_function.remote_function( + session=session + ) + + with pytest.warns( + bigframes.exceptions.PreviewWarning, + match=re.escape("input_types=Series is in preview."), + ): + + @remote_function_decorator + def axis_1_function(myparam: series_type) -> str: # type: ignore + return "Hello, " + myparam["str_col"] + "!" # type: ignore + + # Still works as a normal function. + assert axis_1_function(pandas.Series({"str_col": "World"})) == "Hello, World!" + assert axis_1_function.ibis_node is not None def test_supported_types_correspond(): @@ -29,3 +70,39 @@ def test_supported_types_correspond(): } assert ibis_types_from_python == ibis_types_from_bigquery + + +def test_missing_input_types(): + session = resources.create_bigquery_session() + remote_function_decorator = bigframes.functions.remote_function.remote_function( + session=session + ) + + def function_without_parameter_annotations(myparam) -> str: + return str(myparam) + + assert function_without_parameter_annotations(42) == "42" + + with pytest.raises( + ValueError, + match="'input_types' was not set .* 'myparam' is missing a type annotation", + ): + remote_function_decorator(function_without_parameter_annotations) + + +def test_missing_output_type(): + session = resources.create_bigquery_session() + remote_function_decorator = bigframes.functions.remote_function.remote_function( + session=session + ) + + def function_without_return_annotation(myparam: int): + return str(myparam) + + assert function_without_return_annotation(42) == "42" + + with pytest.raises( + ValueError, + match="'output_type' was not set .* missing a return type annotation", + ): + remote_function_decorator(function_without_return_annotation) diff --git a/third_party/bigframes_vendored/constants.py b/third_party/bigframes_vendored/constants.py new file mode 100644 index 0000000000..0d4a7d1df6 --- /dev/null +++ b/third_party/bigframes_vendored/constants.py @@ -0,0 +1,29 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Constants used across BigQuery DataFrames and bigframes_vendored. + +This module should not depend on any others in the package. +""" + +FEEDBACK_LINK = ( + "Share your usecase with the BigQuery DataFrames team at the " + "https://bit.ly/bigframes-feedback survey." +) + +ABSTRACT_METHOD_ERROR_MESSAGE = ( + "Abstract method. You have likely encountered a bug. " + "Please share this stacktrace and how you reached it with the BigQuery DataFrames team. " + f"{FEEDBACK_LINK}" +) diff --git a/third_party/bigframes_vendored/ibis/expr/operations/analytic.py b/third_party/bigframes_vendored/ibis/expr/operations/analytic.py index 3d6a3b37b1..2e1373533e 100644 --- a/third_party/bigframes_vendored/ibis/expr/operations/analytic.py +++ b/third_party/bigframes_vendored/ibis/expr/operations/analytic.py @@ -2,10 +2,19 @@ from __future__ import annotations +import ibis import ibis.expr.operations as ops import ibis.expr.rules as rlz +# TODO(swast): We can remove this if ibis adds aggregates over scalar values. +# See: https://github.com/ibis-project/ibis/issues/8698 +@ibis.udf.agg.builtin +def count(value: int) -> int: + """Count of a scalar.""" + return 0 # pragma: NO COVER + + class FirstNonNullValue(ops.Analytic): """Retrieve the first element.""" @@ -21,6 +30,7 @@ class LastNonNullValue(ops.Analytic): __all__ = [ + "count", "FirstNonNullValue", "LastNonNullValue", ] diff --git a/third_party/bigframes_vendored/pandas/core/frame.py b/third_party/bigframes_vendored/pandas/core/frame.py index 31d5e88c7e..d46fa4cfc7 100644 --- a/third_party/bigframes_vendored/pandas/core/frame.py +++ b/third_party/bigframes_vendored/pandas/core/frame.py @@ -13,12 +13,11 @@ from typing import Hashable, Iterable, Literal, Mapping, Optional, Sequence, Union +from bigframes_vendored import constants import bigframes_vendored.pandas.core.generic as generic import numpy as np import pandas as pd -from bigframes import constants - # ----------------------------------------------------------------------- # DataFrame class @@ -2003,6 +2002,30 @@ def __eq__(self, other): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def __invert__(self) -> DataFrame: + """ + Returns the bitwise inversion of the DataFrame, element-wise + using operator `~`. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'a':[True, False, True], 'b':[-1, 0, 1]}) + >>> ~df + a b + 0 False 0 + 1 True -1 + 2 False -2 + + [3 rows x 2 columns] + + Returns: + DataFrame: The result of inverting elements in the input. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def ne(self, other, axis: str | int = "columns") -> DataFrame: """ Get not equal to of DataFrame and other, element-wise (binary operator `ne`). @@ -3532,6 +3555,42 @@ def __rpow__(self, other): """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def __and__(self, other): + """Get bitwise AND of DataFrame and other, element-wise, using operator `&`. + + Args: + other (scalar, Series or DataFrame): + Object to bitwise AND with the DataFrame. + + Returns: + bigframes.dataframe.DataFrame: The result of the operation. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def __or__(self, other): + """Get bitwise OR of DataFrame and other, element-wise, using operator `|`. + + Args: + other (scalar, Series or DataFrame): + Object to bitwise OR with the DataFrame. + + Returns: + bigframes.dataframe.DataFrame: The result of the operation. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def __xor__(self, other): + """Get bitwise XOR of DataFrame and other, element-wise, using operator `^`. + + Args: + other (scalar, Series or DataFrame): + Object to bitwise XOR with the DataFrame. + + Returns: + bigframes.dataframe.DataFrame: The result of the operation. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + def combine( self, other, func, fill_value=None, overwrite: bool = True ) -> DataFrame: @@ -3892,8 +3951,8 @@ def map(self, func, na_action: Optional[str] = None) -> DataFrame: to potentially reuse a previously deployed ``remote_function`` from the same user defined function. - >>> @bpd.remote_function(int, float, reuse=False) - ... def minutes_to_hours(x): + >>> @bpd.remote_function(reuse=False) + ... def minutes_to_hours(x: int) -> float: ... return x/60 >>> df_minutes = bpd.DataFrame( @@ -4214,6 +4273,7 @@ def apply(self, func, *, axis=0, args=(), **kwargs): **Examples:** >>> import bigframes.pandas as bpd + >>> import pandas as pd >>> bpd.options.display.progress_bar = None >>> df = bpd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) @@ -4235,16 +4295,19 @@ def apply(self, func, *, axis=0, args=(), **kwargs): [2 rows x 2 columns] You could apply a user defined function to every row of the DataFrame by - creating a remote function out of it, and using it with `axis=1`. + creating a remote function out of it, and using it with `axis=1`. Within + the function, each row is passed as a ``pandas.Series``. It is recommended + to select only the necessary columns before calling `apply()`. Note: This + feature is currently in **preview**. - >>> @bpd.remote_function(bpd.Series, int, reuse=False) - ... def foo(row): + >>> @bpd.remote_function(reuse=False) + ... def foo(row: pd.Series) -> int: ... result = 1 ... result += row["col1"] ... result += row["col2"]*row["col2"] ... return result - >>> df.apply(foo, axis=1) + >>> df[["col1", "col2"]].apply(foo, axis=1) 0 11 1 19 dtype: Int64 diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py index 54c876ef3c..95302e51b2 100644 --- a/third_party/bigframes_vendored/pandas/core/generic.py +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -3,11 +3,10 @@ from typing import Callable, Iterator, Literal, Optional, TYPE_CHECKING +import bigframes_vendored.constants as constants from bigframes_vendored.pandas.core import indexing import bigframes_vendored.pandas.core.common as common -import bigframes.constants as constants - if TYPE_CHECKING: from bigframes_vendored.pandas.pandas._typing import T diff --git a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py index f3f7748e34..6011dbfe5b 100644 --- a/third_party/bigframes_vendored/pandas/core/groupby/__init__.py +++ b/third_party/bigframes_vendored/pandas/core/groupby/__init__.py @@ -21,6 +21,33 @@ def any(self): """ Return True if any value in the group is true, else False. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b'] + >>> ser = bpd.Series([1, 2, 0], index=lst) + >>> ser.groupby(level=0).any() + a True + b False + dtype: boolean + + For DataFrameGroupBy: + + >>> data = [[1, 0, 3], [1, 0, 6], [7, 1, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["ostrich", "penguin", "parrot"]) + >>> df.groupby(by=["a"]).any() + b c + a + 1 False True + 7 True True + + [2 rows x 2 columns] + Returns: Series or DataFrame: DataFrame or Series of boolean values, where a value is True if any element is True within its @@ -32,6 +59,33 @@ def all(self): """ Return True if all values in the group are true, else False. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b'] + >>> ser = bpd.Series([1, 2, 0], index=lst) + >>> ser.groupby(level=0).all() + a True + b False + dtype: boolean + + For DataFrameGroupBy: + + >>> data = [[1, 0, 3], [1, 5, 6], [7, 8, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["ostrich", "penguin", "parrot"]) + >>> df.groupby(by=["a"]).all() + b c + a + 1 False True + 7 True True + + [2 rows x 2 columns] + Returns: Series or DataFrame: DataFrame or Series of boolean values, where a value is True if all elements are True within its @@ -43,6 +97,34 @@ def count(self): """ Compute count of group, excluding missing values. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b'] + >>> ser = bpd.Series([1, 2, np.nan], index=lst) + >>> ser.groupby(level=0).count() + a 2 + b 0 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, np.nan, 3], [1, np.nan, 6], [7, 8, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["cow", "horse", "bull"]) + >>> df.groupby(by=["a"]).count() + b c + a + 1 0 2 + 7 1 1 + + [2 rows x 2 columns] + Returns: Series or DataFrame: Count of values within each group. """ @@ -55,6 +137,45 @@ def mean( """ Compute mean of groups, excluding missing values. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + >>> df = bpd.DataFrame({'A': [1, 1, 2, 1, 2], + ... 'B': [np.nan, 2, 3, 4, 5], + ... 'C': [1, 2, 1, 1, 2]}, columns=['A', 'B', 'C']) + + Groupby one column and return the mean of the remaining columns in each group. + + >>> df.groupby('A').mean() + B C + A + 1 3.0 1.333333 + 2 4.0 1.5 + + [2 rows x 2 columns] + + Groupby two columns and return the mean of the remaining column. + + >>> df.groupby(['A', 'B']).mean() + C + A B + 1 2.0 2.0 + 4.0 1.0 + 2 3.0 1.0 + 5.0 2.0 + + [4 rows x 1 columns] + + Groupby one column and return the mean of only particular column in the group. + + >>> df.groupby('A')['B'].mean() + A + 1 3.0 + 2 4.0 + Name: B, dtype: Float64 + Args: numeric_only (bool, default False): Include only float, int, boolean columns. @@ -73,6 +194,33 @@ def median( """ Compute median of groups, excluding missing values. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'a', 'b', 'b', 'b'] + >>> ser = bpd.Series([7, 2, 8, 4, 3, 3], index=lst) + >>> ser.groupby(level=0).median() + a 7.0 + b 3.0 + dtype: Float64 + + For DataFrameGroupBy: + + >>> data = {'a': [1, 3, 5, 7, 7, 8, 3], 'b': [1, 4, 8, 4, 4, 2, 1]} + >>> df = bpd.DataFrame(data, index=['dog', 'dog', 'dog', + ... 'mouse', 'mouse', 'mouse', 'mouse']) + >>> df.groupby(level=0).median() + a b + dog 3.0 4.0 + mouse 7.0 3.0 + + [2 rows x 2 columns] + Args: numeric_only (bool, default False): Include only float, int, boolean columns. @@ -125,6 +273,33 @@ def std( For multiple groupings, the result index will be a MultiIndex. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'a', 'b', 'b', 'b'] + >>> ser = bpd.Series([7, 2, 8, 4, 3, 3], index=lst) + >>> ser.groupby(level=0).std() + a 3.21455 + b 0.57735 + dtype: Float64 + + For DataFrameGroupBy: + + >>> data = {'a': [1, 3, 5, 7, 7, 8, 3], 'b': [1, 4, 8, 4, 4, 2, 1]} + >>> df = bpd.DataFrame(data, index=['dog', 'dog', 'dog', + ... 'mouse', 'mouse', 'mouse', 'mouse']) + >>> df.groupby(level=0).std() + a b + dog 2.0 3.511885 + mouse 2.217356 1.5 + + [2 rows x 2 columns] + Args: numeric_only (bool, default False): Include only `float`, `int` or `boolean` data. @@ -144,6 +319,33 @@ def var( For multiple groupings, the result index will be a MultiIndex. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'a', 'b', 'b', 'b'] + >>> ser = bpd.Series([7, 2, 8, 4, 3, 3], index=lst) + >>> ser.groupby(level=0).var() + a 10.333333 + b 0.333333 + dtype: Float64 + + For DataFrameGroupBy: + + >>> data = {'a': [1, 3, 5, 7, 7, 8, 3], 'b': [1, 4, 8, 4, 4, 2, 1]} + >>> df = bpd.DataFrame(data, index=['dog', 'dog', 'dog', + ... 'mouse', 'mouse', 'mouse', 'mouse']) + >>> df.groupby(level=0).var() + a b + dog 4.0 12.333333 + mouse 4.916667 2.25 + + [2 rows x 2 columns] + Args: numeric_only (bool, default False): Include only `float`, `int` or `boolean` data. @@ -164,6 +366,23 @@ def skew( Normalized by N-1. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> ser = bpd.Series([390., 350., 357., np.nan, 22., 20., 30.], + ... index=['Falcon', 'Falcon', 'Falcon', 'Falcon', + ... 'Parrot', 'Parrot', 'Parrot'], + ... name="Max Speed") + >>> ser.groupby(level=0).skew() + Falcon 1.525174 + Parrot 1.457863 + Name: Max Speed, dtype: Float64 + Args: numeric_only (bool, default False): Include only `float`, `int` or `boolean` data. @@ -185,6 +404,51 @@ def kurt( Kurtosis obtained using Fisher's definition of kurtosis (kurtosis of normal == 0.0). Normalized by N-1. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b'] + >>> ser = bpd.Series([0, 1, 1, 0, 0, 1, 2, 4, 5], index=lst) + >>> ser.groupby(level=0).kurt() + a -6.0 + b -1.963223 + dtype: Float64 + + Args: + numeric_only (bool, default False): + Include only `float`, `int` or `boolean` data. + + Returns: + Series or DataFrame + Variance of values within each group. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def kurtosis( + self, + *, + numeric_only: bool = False, + ): + """ + Return unbiased kurtosis over requested axis. + + Kurtosis obtained using Fisher's definition of + kurtosis (kurtosis of normal == 0.0). Normalized by N-1. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b'] + >>> ser = bpd.Series([0, 1, 1, 0, 0, 1, 2, 4, 5], index=lst) + >>> ser.groupby(level=0).kurtosis() + a -6.0 + b -1.963223 + dtype: Float64 + Args: numeric_only (bool, default False): Include only `float`, `int` or `boolean` data. @@ -203,6 +467,33 @@ def sum( """ Compute sum of group values. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b', 'b'] + >>> ser = bpd.Series([1, 2, 3, 4], index=lst) + >>> ser.groupby(level=0).sum() + a 3 + b 7 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, 8, 2], [1, 2, 5], [2, 5, 8], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["tiger", "leopard", "cheetah", "lion"]) + >>> df.groupby("a").sum() + b c + a + 1 10 7 + 2 11 17 + + [2 rows x 2 columns] + Args: numeric_only (bool, default False): Include only float, int, boolean columns. @@ -218,6 +509,22 @@ def sum( def prod(self, numeric_only: bool = False, min_count: int = 0): """ Compute prod of group values. + (DataFrameGroupBy functionality is not yet available.) + + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b', 'b'] + >>> ser = bpd.Series([1, 2, 3, 4], index=lst) + >>> ser.groupby(level=0).prod() + a 2.0 + b 12.0 + dtype: Float64 Args: numeric_only (bool, default False): @@ -239,6 +546,34 @@ def min( """ Compute min of group values. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b', 'b'] + >>> ser = bpd.Series([1, 2, 3, 4], index=lst) + >>> ser.groupby(level=0).min() + a 1 + b 3 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, 8, 2], [1, 2, 5], [2, 5, 8], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["tiger", "leopard", "cheetah", "lion"]) + >>> df.groupby(by=["a"]).min() + b c + a + 1 2 2 + 2 5 8 + + [2 rows x 2 columns] + Args: numeric_only (bool, default False): Include only float, int, boolean columns. @@ -259,6 +594,33 @@ def max( """ Compute max of group values. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b', 'b'] + >>> ser = bpd.Series([1, 2, 3, 4], index=lst) + >>> ser.groupby(level=0).max() + a 2 + b 4 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, 8, 2], [1, 2, 5], [2, 5, 8], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["tiger", "leopard", "cheetah", "lion"]) + >>> df.groupby(by=["a"]).max() + b c + a + 1 8 5 + 2 6 9 + + [2 rows x 2 columns] + Args: numeric_only (bool, default False): Include only float, int, boolean columns. @@ -274,6 +636,31 @@ def max( def cumcount(self, ascending: bool = True): """ Number each item in each group from 0 to the length of that group - 1. + (DataFrameGroupBy functionality is not yet available.) + + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b', 'b', 'c'] + >>> ser = bpd.Series([5, 1, 2, 3, 4], index=lst) + >>> ser.groupby(level=0).cumcount() + a 0 + a 1 + b 0 + b 1 + c 0 + dtype: Int64 + >>> ser.groupby(level=0).cumcount(ascending=False) + a 0 + a 1 + b 0 + b 1 + c 0 + dtype: Int64 Args: ascending (bool, default True): @@ -288,6 +675,35 @@ def cumprod(self, *args, **kwargs): """ Cumulative product for each group. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b'] + >>> ser = bpd.Series([6, 2, 0], index=lst) + >>> ser.groupby(level=0).cumprod() + a 6.0 + a 12.0 + b 0.0 + dtype: Float64 + + For DataFrameGroupBy: + + >>> data = [[1, 8, 2], [1, 2, 5], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["cow", "horse", "bull"]) + >>> df.groupby("a").cumprod() + b c + cow 8.0 2.0 + horse 16.0 10.0 + bull 6.0 9.0 + + [3 rows x 2 columns] + Returns: Series or DataFrame: Cumulative product for each group. """ @@ -297,6 +713,35 @@ def cumsum(self, *args, **kwargs): """ Cumulative sum for each group. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b'] + >>> ser = bpd.Series([6, 2, 0], index=lst) + >>> ser.groupby(level=0).cumsum() + a 6 + a 8 + b 0 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, 8, 2], [1, 2, 5], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["fox", "gorilla", "lion"]) + >>> df.groupby("a").cumsum() + b c + fox 8 2 + gorilla 10 7 + lion 6 9 + + [3 rows x 2 columns] + Returns: Series or DataFrame: Cumulative sum for each group. """ @@ -306,6 +751,35 @@ def cummin(self, *args, numeric_only: bool = False, **kwargs): """ Cumulative min for each group. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b'] + >>> ser = bpd.Series([6, 2, 0], index=lst) + >>> ser.groupby(level=0).cummin() + a 6 + a 2 + b 0 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, 8, 2], [1, 2, 5], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["fox", "gorilla", "lion"]) + >>> df.groupby("a").cummin() + b c + fox 8 2 + gorilla 2 2 + lion 6 9 + + [3 rows x 2 columns] + Returns: Series or DataFrame: Cumulative min for each group. """ @@ -315,6 +789,35 @@ def cummax(self, *args, numeric_only: bool = False, **kwargs): """ Cumulative max for each group. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b'] + >>> ser = bpd.Series([6, 2, 0], index=lst) + >>> ser.groupby(level=0).cummax() + a 6 + a 6 + b 0 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, 8, 2], [1, 2, 5], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["fox", "gorilla", "lion"]) + >>> df.groupby("a").cummax() + b c + fox 8 2 + gorilla 8 5 + lion 6 9 + + [3 rows x 2 columns] + Returns: Series or DataFrame: Cumulative max for each group. """ @@ -326,6 +829,42 @@ def diff(self): Calculates the difference of each element compared with another element in the group (default is element in previous row). + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'a', 'b', 'b', 'b'] + >>> ser = bpd.Series([7, 2, 8, 4, 3, 3], index=lst) + >>> ser.groupby(level=0).diff() + a + a -5 + a 6 + b + b -1 + b 0 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = {'a': [1, 3, 5, 7, 7, 8, 3], 'b': [1, 4, 8, 4, 4, 2, 1]} + >>> df = bpd.DataFrame(data, index=['dog', 'dog', 'dog', + ... 'mouse', 'mouse', 'mouse', 'mouse']) + >>> df.groupby(level=0).diff() + a b + dog + dog 2 3 + dog 2 4 + mouse + mouse 0 0 + mouse 1 -2 + mouse -5 -1 + + [7 rows x 2 columns] + Returns: Series or DataFrame: First differences. """ @@ -335,6 +874,37 @@ def shift(self, periods: int = 1): """ Shift each group by periods observations. + **Examples:** + + For SeriesGroupBy: + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b', 'b'] + >>> ser = bpd.Series([1, 2, 3, 4], index=lst) + >>> ser.groupby(level=0).shift(1) + a + a 1 + b + b 3 + dtype: Int64 + + For DataFrameGroupBy: + + >>> data = [[1, 2, 3], [1, 5, 6], [2, 5, 8], [2, 6, 9]] + >>> df = bpd.DataFrame(data, columns=["a", "b", "c"], + ... index=["tuna", "salmon", "catfish", "goldfish"]) + >>> df.groupby("a").shift(1) + b c + tuna + salmon 2 3 + catfish + goldfish 5 8 + + [4 rows x 2 columns] + Args: periods (int, default 1): Number of periods to shift. @@ -348,6 +918,23 @@ def rolling(self, *args, **kwargs): """ Returns a rolling grouper, providing rolling functionality per group. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'a', 'a', 'e'] + >>> ser = bpd.Series([1, 0, -2, -1, 2], index=lst) + >>> ser.groupby(level=0).rolling(2).min() + index index + a a + a 0 + a -2 + a -2 + e e + dtype: Int64 + Args: min_periods (int, default None): Minimum number of observations in window required to have a value; @@ -368,6 +955,23 @@ def expanding(self, *args, **kwargs): """ Provides expanding functionality. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'c', 'c', 'e'] + >>> ser = bpd.Series([1, 0, -2, -1, 2], index=lst) + >>> ser.groupby(level=0).expanding().min() + index index + a a 1 + a 0 + c c -2 + c -2 + e e 2 + dtype: Int64 + Returns: Series or DataFrame: An expanding grouper, providing expanding functionality per group. """ @@ -379,6 +983,52 @@ def agg(self, func): """ Aggregate using one or more operations. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> s = bpd.Series([1, 2, 3, 4], index=[1, 1, 2, 2]) + >>> s.groupby(level=0).agg(['min', 'max']) + min max + 1 1 2 + 2 3 4 + + [2 rows x 2 columns] + + Args: + func : function, str, list, dict or None + Function to use for aggregating the data. + + Accepted combinations are: + + - string function name + - list of function names, e.g. ``['sum', 'mean']`` + + Returns: + Series or DataFrame + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def aggregate(self, func): + """ + Aggregate using one or more operations. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> s = bpd.Series([1, 2, 3, 4], index=[1, 1, 2, 2]) + >>> s.groupby(level=0).aggregate(['min', 'max']) + min max + 1 1 2 + 2 3 4 + + [2 rows x 2 columns] + Args: func : function, str, list, dict or None Function to use for aggregating the data. @@ -397,6 +1047,19 @@ def nunique(self): """ Return number of unique elements in the group. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> lst = ['a', 'a', 'b', 'b'] + >>> ser = bpd.Series([1, 2, 3, 3], index=lst) + >>> ser.groupby(level=0).nunique() + a 2 + b 1 + dtype: Int64 + Returns: Series: Number of unique values within each group. """ @@ -408,6 +1071,97 @@ def agg(self, func, **kwargs): """ Aggregate using one or more operations. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> data = {"A": [1, 1, 2, 2], + ... "B": [1, 2, 3, 4], + ... "C": [0.362838, 0.227877, 1.267767, -0.562860]} + >>> df = bpd.DataFrame(data) + + The aggregation is for each column. + + >>> df.groupby('A').agg('min') + B C + A + 1 1 0.227877 + 2 3 -0.56286 + + [2 rows x 2 columns] + + Multiple aggregations + + >>> df.groupby('A').agg(['min', 'max']) + B C + min max min max + A + 1 1 2 0.227877 0.362838 + 2 3 4 -0.56286 1.267767 + + [2 rows x 4 columns] + + Args: + func (function, str, list, dict or None): + Function to use for aggregating the data. + + Accepted combinations are: + + - string function name + - list of function names, e.g. ``['sum', 'mean']`` + - dict of axis labels -> function names or list of such. + - None, in which case ``**kwargs`` are used with Named Aggregation. Here the + output has one column for each element in ``**kwargs``. The name of the + column is keyword, whereas the value determines the aggregation used to compute + the values in the column. + + kwargs + If ``func`` is None, ``**kwargs`` are used to define the output names and + aggregations via Named Aggregation. See ``func`` entry. + + Returns: + DataFrame + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def aggregate(self, func, **kwargs): + """ + Aggregate using one or more operations. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> data = {"A": [1, 1, 2, 2], + ... "B": [1, 2, 3, 4], + ... "C": [0.362838, 0.227877, 1.267767, -0.562860]} + >>> df = bpd.DataFrame(data) + + The aggregation is for each column. + + >>> df.groupby('A').aggregate('min') + B C + A + 1 1 0.227877 + 2 3 -0.56286 + + [2 rows x 2 columns] + + Multiple aggregations + + >>> df.groupby('A').agg(['min', 'max']) + B C + min max min max + A + 1 1 2 0.227877 0.362838 + 2 3 4 -0.56286 1.267767 + + [2 rows x 4 columns] + Args: func (function, str, list, dict or None): Function to use for aggregating the data. @@ -435,6 +1189,25 @@ def nunique(self): """ Return DataFrame with counts of unique elements in each position. + **Examples:** + + >>> import bigframes.pandas as bpd + >>> import numpy as np + >>> bpd.options.display.progress_bar = None + + >>> df = bpd.DataFrame({'id': ['spam', 'egg', 'egg', 'spam', + ... 'ham', 'ham'], + ... 'value1': [1, 5, 5, 2, 5, 5], + ... 'value2': list('abbaxy')}) + >>> df.groupby('id').nunique() + value1 value2 + id + egg 1 1 + ham 1 2 + spam 2 1 + + [3 rows x 2 columns] + Returns: DataFrame """ diff --git a/third_party/bigframes_vendored/pandas/core/indexing.py b/third_party/bigframes_vendored/pandas/core/indexing.py index 3c7f8a6c9f..a188c7197e 100644 --- a/third_party/bigframes_vendored/pandas/core/indexing.py +++ b/third_party/bigframes_vendored/pandas/core/indexing.py @@ -1,6 +1,6 @@ # Contains code from https://github.com/pandas-dev/pandas/blob/main/pandas/core/indexing.py -import bigframes.constants as constants +import bigframes_vendored.constants as constants class IndexingMixin: diff --git a/third_party/bigframes_vendored/pandas/core/series.py b/third_party/bigframes_vendored/pandas/core/series.py index 585e20275c..a430c3375f 100644 --- a/third_party/bigframes_vendored/pandas/core/series.py +++ b/third_party/bigframes_vendored/pandas/core/series.py @@ -1183,8 +1183,8 @@ def apply( to potentially reuse a previously deployed `remote_function` from the same user defined function. - >>> @bpd.remote_function(int, float, reuse=False) - ... def minutes_to_hours(x): + >>> @bpd.remote_function(reuse=False) + ... def minutes_to_hours(x: int) -> float: ... return x/60 >>> minutes = bpd.Series([0, 30, 60, 90, 120]) @@ -1210,12 +1210,10 @@ def apply( `packages` param. >>> @bpd.remote_function( - ... str, - ... str, ... reuse=False, ... packages=["cryptography"], ... ) - ... def get_hash(input): + ... def get_hash(input: str) -> str: ... from cryptography.fernet import Fernet ... ... # handle missing value @@ -3452,8 +3450,8 @@ def mask(self, cond, other): condition is evaluated based on a complicated business logic which cannot be expressed in form of a Series. - >>> @bpd.remote_function(str, bool, reuse=False) - ... def should_mask(name): + >>> @bpd.remote_function(reuse=False) + ... def should_mask(name: str) -> bool: ... hash = 0 ... for char_ in name: ... hash += ord(char_) @@ -3971,8 +3969,8 @@ def map( It also accepts a remote function: - >>> @bpd.remote_function(str, str) - ... def my_mapper(val): + >>> @bpd.remote_function() + ... def my_mapper(val: str) -> str: ... vowels = ["a", "e", "i", "o", "u"] ... if val: ... return "".join([ @@ -4226,7 +4224,7 @@ def __and__(self, other): Object to bitwise AND with the Series. Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) @@ -4264,7 +4262,45 @@ def __or__(self, other): Object to bitwise OR with the Series. Returns: - Series: The result of the operation. + bigframes.series.Series: The result of the operation. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def __xor__(self, other): + """Get bitwise XOR of Series and other, element-wise, using operator `^`. + + **Examples:** + + >>> import bigframes.pandas as bpd + >>> bpd.options.display.progress_bar = None + + >>> s = bpd.Series([0, 1, 2, 3]) + + You can operate with a scalar. + + >>> s ^ 6 + 0 6 + 1 7 + 2 4 + 3 5 + dtype: Int64 + + You can operate with another Series. + + >>> s1 = bpd.Series([5, 6, 7, 8]) + >>> s ^ s1 + 0 5 + 1 7 + 2 5 + 3 11 + dtype: Int64 + + Args: + other (scalar or Series): + Object to bitwise XOR with the Series. + + Returns: + bigframes.series.Series: The result of the operation. """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) diff --git a/third_party/bigframes_vendored/sklearn/impute/_base.py b/third_party/bigframes_vendored/sklearn/impute/_base.py new file mode 100644 index 0000000000..3064e8a118 --- /dev/null +++ b/third_party/bigframes_vendored/sklearn/impute/_base.py @@ -0,0 +1,53 @@ +# Authors: Nicolas Tresegnie +# Sergey Feldman +# License: BSD 3 clause + +from bigframes_vendored.sklearn.base import BaseEstimator, TransformerMixin + +from bigframes import constants + + +class _BaseImputer(TransformerMixin, BaseEstimator): + """Base class for all imputers.""" + + +class SimpleImputer(_BaseImputer): + """ + Univariate imputer for completing missing values with simple strategies. + + Replace missing values using a descriptive statistic (e.g. mean, median, or + most frequent) along each column. + + Args: + strategy ({'mean', 'median', 'most_frequent'}, default='mean'): + The imputation strategy. 'mean': replace missing values using the mean along + the axis. 'median':replace missing values using the median along + the axis. 'most_frequent', replace missing using the most frequent + value along the axis. + """ + + def fit(self, X, y=None): + """Fit the imputer on X. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The Dataframe or Series with training data. + + y (default None): + Ignored. + + Returns: + SimpleImputer: Fitted scaler. + """ + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE) + + def transform(self, X): + """Impute all missing values in X. + + Args: + X (bigframes.dataframe.DataFrame or bigframes.series.Series): + The DataFrame or Series to be transformed. + + Returns: + bigframes.dataframe.DataFrame: Transformed result.""" + raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)