add tests for multiple columns and custom separator

Author: tswast

bigframes/operations/structs.py

@@ -66,7 +66,7 @@ class StructFrameAccessor(vendoracessors.StructFrameAccessor):
     def __init__(self, data: bigframes.dataframe.DataFrame) -> None:
         self._parent = data
 
-    def explode(self, column, separator: str = ".") -> bigframes.dataframe.DataFrame:
+    def explode(self, column, *, separator: str = ".") -> bigframes.dataframe.DataFrame:
         df = self._parent
         column_labels = bigframes.core.explode.check_column(column)
 

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",
+    ]

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

@@ -124,7 +124,11 @@ def dtypes(self):
 
 
 class StructFrameAccessor:
-    def explode(self, column, separator: str = "."):
+    """
+    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.