feat: add to_parquet_row_groups

src/awkward/operations/ak_to_parquet.py

@@ -229,7 +229,7 @@ def _impl(
     parquet_byte_stream_split,
     parquet_coerce_timestamps,
     parquet_old_int96_timestamps,
-    parquet_compliant_nested,  # https://issues.apache.org/jira/browse/ARROW-16348
+    parquet_compliant_nested,
     parquet_extra_options,
     storage_options,
     write_iteratively,
@@ -242,9 +242,9 @@ def _impl(
     pyarrow_parquet = awkward._connect.pyarrow.import_pyarrow_parquet("ak.to_parquet")
     fsspec = awkward._connect.pyarrow.import_fsspec("ak.to_parquet")
 
-    if not write_iteratively:
+    def get_layout_and_table(x):
         layout = ak.operations.ak_to_layout._impl(
-            data,
+            x,
             allow_record=True,
             allow_unknown=False,
             none_policy="error",
@@ -253,7 +253,6 @@ def _impl(
             primitive_policy="error",
             string_policy="as-characters",
         )
-
         table = ak.operations.ak_to_arrow_table._impl(
             layout,
             list_to32,
@@ -264,27 +263,44 @@ def _impl(
             extensionarray,
             count_nulls,
         )
+        return layout, table
+
+    if not write_iteratively:
+        layout, table = get_layout_and_table(data)
+
     else:
-        layout = ak.operations.ak_to_layout._impl(
-            next(data),
-            allow_record=True,
-            allow_unknown=False,
-            none_policy="error",
-            regulararray=True,
-            use_from_iter=True,
-            primitive_policy="error",
-            string_policy="as-characters",
-        )
-        table = ak.operations.ak_to_arrow_table._impl(
-            layout,
-            list_to32,
-            string_to32,
-            bytestring_to32,
-            emptyarray_to,
-            categorical_as_dictionary,
-            extensionarray,
-            count_nulls,
-        )
+        if (
+            isinstance(
+                data,
+                (
+                    ak.highlevel.Array,
+                    ak.highlevel.Record,
+                    ak.highlevel.ArrayBuilder,
+                    ak.contents.Content,
+                    ak.record.Record,
+                ),
+            )
+            or hasattr(data, "dtype")
+            or hasattr(data, "shape")
+        ):
+            raise TypeError(
+                "The first argument of ak.to_parquet_row_groups should be an "
+                "iterable collection of arrays to put into separate row groups, "
+                "not a single array. If this was intended, wrap the first argument "
+                "with the `iter` function."
+            )
+
+        data = iter(data)  # necessary for iterables; no-op for iterators
+
+        try:
+            first_array = next(data)
+        except StopIteration:
+            raise ValueError(
+                "The first argument of ak.to_parquet_row_groups must contain at least one array."
+            ) from None
+
+        # if `write_iteratively`, then `layout` and `table` represent the FIRST, not all
+        layout, table = get_layout_and_table(first_array)
 
     if parquet_compliant_nested:
         list_indicator = "list.element"
@@ -415,38 +431,24 @@ def parquet_columns(specifier, only=None):
         metadata_collector=metalist,
         **parquet_extra_options,
     ) as writer:
-        if write_iteratively:
-            writer.write_table(table, row_group_size=row_group_size)
-            try:
-                while True:
-                    try:
-                        layout = ak.operations.ak_to_layout._impl(
-                            next(data),
-                            allow_record=True,
-                            allow_unknown=False,
-                            none_policy="error",
-                            regulararray=True,
-                            use_from_iter=True,
-                            primitive_policy="error",
-                            string_policy="as-characters",
-                        )
-                        table = ak.operations.ak_to_arrow_table._impl(
-                            layout,
-                            list_to32,
-                            string_to32,
-                            bytestring_to32,
-                            emptyarray_to,
-                            categorical_as_dictionary,
-                            extensionarray,
-                            count_nulls,
-                        )
-                        writer.write_table(table, row_group_size=row_group_size)
-                    except StopIteration:
-                        break
-            finally:
-                writer.close()
-        else:
-            writer.write_table(table, row_group_size=row_group_size)
+        try:
+            if not write_iteratively:
+                writer.write_table(table, row_group_size=row_group_size)
+
+            else:
+                # this `table` is JUST for the first array
+                writer.write_table(table, row_group_size=row_group_size)
+
+                # this `data` is an iterator (not iterable), starting AFTER the first array
+                # a `for` loop implicitly calls `next` and stops at `StopIteration`
+                for item in data:
+                    layout, table = get_layout_and_table(item)
+                    writer.write_table(table, row_group_size=row_group_size)
+
+        finally:
+            # ensure that the ParquetWriter is closed for iterative and non-iterative cases
+            writer.close()
+
     meta = metalist[0]
     meta.set_file_path(destination.rsplit("/", 1)[-1])
     return meta