feat: add DataFrame.struct.explode to add struct subfields to a DataFrame

bigframes/core/explode.py

@@ -0,0 +1,36 @@
+# 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.
+
+"""Utility functions for implementing 'explode' functions."""
+
+from typing import cast, Sequence, Union
+
+import bigframes.core.blocks as blocks
+import bigframes.core.utils as utils
+
+
+def check_column(
+ column: Union[blocks.Label, Sequence[blocks.Label]],
+) -> Sequence[blocks.Label]:
+ if not utils.is_list_like(column):
+ column_labels = cast(Sequence[blocks.Label], (column,))
+ else:
+ column_labels = cast(Sequence[blocks.Label], tuple(column))
+
+ if not column_labels:
+ raise ValueError("column must be nonempty")
+ if len(column_labels) > len(set(column_labels)):
+ raise ValueError("column must be unique")
+
+ return column_labels

bigframes/dataframe.py

@@ -55,6 +55,7 @@
 import bigframes.core.block_transforms as block_ops
 import bigframes.core.blocks as blocks
 import bigframes.core.convert
+import bigframes.core.explode
 import bigframes.core.expression as ex
 import bigframes.core.groupby as groupby
 import bigframes.core.guid
@@ -71,6 +72,7 @@
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
 import bigframes.operations.plotting as plotting
+import bigframes.operations.structs
 import bigframes.series
 import bigframes.series as bf_series
 import bigframes.session._io.bigquery
@@ -2875,15 +2877,7 @@ def explode(
  *,
  ignore_index: Optional[bool] = False,
  ) -> DataFrame:
- if not utils.is_list_like(column):
- column_labels = typing.cast(typing.Sequence[blocks.Label], (column,))
- else:
- column_labels = typing.cast(typing.Sequence[blocks.Label], tuple(column))
-
- if not column_labels:
- raise ValueError("column must be nonempty")
- if len(column_labels) > len(set(column_labels)):
- raise ValueError("column must be unique")
+ column_labels = bigframes.core.explode.check_column(column)
 
  column_ids = [self._resolve_label_exact(label) for label in column_labels]
  missing = [
@@ -3751,6 +3745,10 @@ def __matmul__(self, other) -> DataFrame:
 
  __matmul__.__doc__ = inspect.getdoc(vendored_pandas_frame.DataFrame.__matmul__)
 
+ @property
+ def struct(self):
+ return bigframes.operations.structs.StructFrameAccessor(self)
+
  def _throw_if_null_index(self, opname: str):
  if not self._has_index:
  raise bigframes.exceptions.NullIndexError(

bigframes/operations/structs.py

@@ -57,3 +57,26 @@ def dtypes(self) -> pd.Series:
  ],
  index=[pa_type.field(i).name for i in range(pa_type.num_fields)],
  )
+
+
+@log_adapter.class_logger
+class StructFrameAccessor(vendoracessors.StructFrameAccessor):
+ __doc__ = vendoracessors.StructAccessor.__doc__
+
+ def __init__(self, data: bigframes.dataframe.DataFrame) -> None:
+ self._parent = data
+
+ def explode(self, column, *, separator: str = ".") -> bigframes.dataframe.DataFrame:
+ df = self._parent
+ column_labels = bigframes.core.explode.check_column(column)
+
+ for label in column_labels:
+ position = df.columns.to_list().index(label)
+ df = df.drop(columns=label)
+ subfields = self._parent[label].struct.explode()
+ for subfield in reversed(subfields.columns):
+ df.insert(
+ position, f"{label}{separator}{subfield}", subfields[subfield]
+ )
+
+ return df

tests/data/nested.jsonl

@@ -1,4 +1,4 @@
-{"rowindex":0,"customer_id":"jkl","day":"2023-12-18","flag":1,"event_sequence":[{"category":"B","timestamp":"2023-12-18 03:43:58","data":[{"key":"x","value":20.2533015856},{"key":"y","value":42.8363462389}]},{"category":"D","timestamp":"2023-12-18 07:15:37","data":[{"key":"x","value":62.0762664928},{"key":"z","value":83.6655402432}]}]}
+{"rowindex":0,"customer_id":"jkl","day":"2023-12-18","flag":1,"label":{"key": "my-key","value":"my-value"},"event_sequence":[{"category":"B","timestamp":"2023-12-18 03:43:58","data":[{"key":"x","value":20.2533015856},{"key":"y","value":42.8363462389}]},{"category":"D","timestamp":"2023-12-18 07:15:37","data":[{"key":"x","value":62.0762664928},{"key":"z","value":83.6655402432}]}],"address":{"street":"123 Test Lane","city":"Testerchon"}}
 {"rowindex":1,"customer_id":"def","day":"2023-12-18","flag":2,"event_sequence":[{"category":"D","timestamp":"2023-12-18 23:11:11","data":[{"key":"w","value":36.1388065179}]},{"category":"B","timestamp":"2023-12-18 07:12:50","data":[{"key":"z","value":68.7673488304}]},{"category":"D","timestamp":"2023-12-18 09:09:03","data":[{"key":"x","value":57.4139647019}]},{"category":"C","timestamp":"2023-12-18 13:05:30","data":[{"key":"z","value":36.087871201}]}]}
 {"rowindex":2,"customer_id":"abc","day":"2023-12-6","flag":0,"event_sequence":[{"category":"C","timestamp":"2023-12-06 10:37:11","data":[]},{"category":"A","timestamp":"2023-12-06 03:35:44","data":[]},{"category":"D","timestamp":"2023-12-06 13:10:57","data":[{"key":"z","value":21.8487807658}]},{"category":"B","timestamp":"2023-12-06 01:39:16","data":[{"key":"y","value":1.6380505139}]}]}
 {"rowindex":3,"customer_id":"mno","day":"2023-12-16","flag":2,"event_sequence":[]}

tests/data/nested_schema.json

@@ -19,6 +19,20 @@
  "name": "flag",
  "type": "INTEGER"
  },
+ {
+ "fields": [
+ {
+ "name": "key",
+ "type": "STRING"
+ },
+ {
+ "name": "value",
+ "type": "STRING"
+ }
+ ],
+ "name": "label",
+ "type": "RECORD"
+ },
  {
  "fields": [
  {
@@ -52,5 +66,19 @@
  "mode": "REPEATED",
  "name": "event_sequence",
  "type": "RECORD"
+ },
+ {
+ "fields": [
+ {
+ "name": "street",
+ "type": "STRING"
+ },
+ {
+ "name": "city",
+ "type": "STRING"
+ }
+ ],
+ "name": "address",
+ "type": "RECORD"
  }
 ]

tests/system/small/operations/test_struct.py

@@ -0,0 +1,40 @@
+# 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.
+
+
+def test_dataframe_struct_explode_multiple_columns(nested_df):
+ got = nested_df.struct.explode(["label", "address"])
+ assert got.columns.to_list() == [
+ "customer_id",
+ "day",
+ "flag",
+ "label.key",
+ "label.value",
+ "event_sequence",
+ "address.street",
+ "address.city",
+ ]
+
+
+def test_dataframe_struct_explode_separator(nested_df):
+ got = nested_df.struct.explode("label", separator="__sep__")
+ assert got.columns.to_list() == [
+ "customer_id",
+ "day",
+ "flag",
+ "label__sep__key",
+ "label__sep__value",
+ "event_sequence",
+ "address",
+ ]

tests/system/small/test_dataframe.py

@@ -1436,6 +1436,14 @@ def test_get_dtypes_array_struct_table(nested_df):
  "customer_id": pd.StringDtype(storage="pyarrow"),
  "day": pd.ArrowDtype(pa.date32()),
  "flag": pd.Int64Dtype(),
+ "label": pd.ArrowDtype(
+ pa.struct(
+ [
+ ("key", pa.string()),
+ ("value", pa.string()),
+ ]
+ ),
+ ),
  "event_sequence": pd.ArrowDtype(
  pa.list_(
  pa.struct(
@@ -1457,6 +1465,14 @@ def test_get_dtypes_array_struct_table(nested_df):
  ),
  ),
  ),
+ "address": pd.ArrowDtype(
+ pa.struct(
+ [
+ ("street", pa.string()),
+ ("city", pa.string()),
+ ]
+ ),
+ ),
  }
  ),
  )

third_party/bigframes_vendored/pandas/core/arrays/arrow/accessors.py

@@ -121,3 +121,56 @@ def dtypes(self):
  A *pandas* Series with the data type of all child fields.
  """
  raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)
+
+
+class StructFrameAccessor:
+ """
+ Accessor object for structured data properties of the DataFrame values.
+ """
+
+ def explode(self, column, *, separator: str = "."):
+ """
+ Extract all child fields of struct column(s) and add to the DataFrame.
+
+ **Examples:**
+
+ >>> import bigframes.pandas as bpd
+ >>> import pyarrow as pa
+ >>> bpd.options.display.progress_bar = None
+ >>> countries = bpd.Series(["cn", "es", "us"])
+ >>> files = bpd.Series(
+ ... [
+ ... {"version": 1, "project": "pandas"},
+ ... {"version": 2, "project": "pandas"},
+ ... {"version": 1, "project": "numpy"},
+ ... ],
+ ... dtype=bpd.ArrowDtype(pa.struct(
+ ... [("version", pa.int64()), ("project", pa.string())]
+ ... ))
+ ... )
+ >>> downloads = bpd.Series([100, 200, 300])
+ >>> df = bpd.DataFrame({"country": countries, "file": files, "download_count": downloads})
+ >>> df.struct.explode("file")
+ country file.version file.project download_count
+ 0 cn 1 pandas 100
+ 1 es 2 pandas 200
+ 2 us 1 numpy 300
+ <BLANKLINE>
+ [3 rows x 4 columns]
+
+ Args:
+ column:
+ Column(s) to explode. For multiple columns, specify a non-empty
+ list with each element be str or tuple, and all specified
+ columns their list-like data on same row of the frame must
+ have matching length.
+ separator:
+ Separator/delimiter to use to separate the original column name
+ from the sub-field column name.
+
+
+ Returns:
+ DataFrame:
+ Original DataFrame with exploded struct column(s).
+ """
+ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)