Repository: lincc-frameworks/nested-pandas
Pull request #257: Change NestedExtArray.__arrow__ to return list_struct
Author: login: hombit
Created: 2025-05-09T18:37:21Z | Merged: 2025-05-09T19:20:20Z

It is a small change which would allow us to implicitly track multiple-nested arrays: each field which is a list-struct array could be cast to a nested array.


Diff:
diff --git a/src/nested_pandas/nestedframe/core.py b/src/nested_pandas/nestedframe/core.py
index 4edbeae6..906d2711 100644
--- a/src/nested_pandas/nestedframe/core.py
+++ b/src/nested_pandas/nestedframe/core.py
@@ -1242,6 +1242,28 @@ def reduce(self, func, *args, infer_nesting=True, **kwargs) -> NestedFrame:  # t
 
         return results_nf
 
+    def to_pandas(self, list_struct=False) -> pd.DataFrame:
+        """Convert to an ordinal pandas DataFrame, with no NestedDtype series.
+
+        NestedDtype is cast to pd.ArrowDtype
+
+        Parameters
+        ----------
+        list_struct: bool
+            If True, cast nested columns to pandas struct-list arrow extension
+            array columns. If False (default), cast nested columns to
+            list-struct array columns.
+
+        Returns
+        -------
+        pd.DataFrame
+            Ordinal pandas DataFrame.
+        """
+        df = pd.DataFrame(self)
+        for col in self.nested_columns:
+            df[col] = df[col].array.to_arrow_ext_array(list_struct=list_struct)
+        return df
+
     def to_parquet(self, path, **kwargs) -> None:
         """Creates parquet file(s) with the data of a NestedFrame, either
         as a single parquet file where each nested dataset is packed into its
@@ -1269,13 +1291,14 @@ def to_parquet(self, path, **kwargs) -> None:
         >>> nf = generate_data(5,5, seed=1)
         >>> nf.to_parquet("nestedframe.parquet")
         """
+        df = self.to_pandas(list_struct=False)
 
         # Write through pyarrow
         # This is potentially not zero-copy
         # Note: Without pandas metadata, index writing is not as robust set
         # preserve_index=None for best behavior but index will generally
         # need to be set manually on load
-        table = pa.Table.from_pandas(self, preserve_index=None)
+        table = pa.Table.from_pandas(df, preserve_index=None)
 
         # Drop pandas metadata to make sure nesteddtypes are not preserved
         # Do this by rebuilding the schema
diff --git a/src/nested_pandas/series/ext_array.py b/src/nested_pandas/series/ext_array.py
index beccfac0..1fedd36b 100644
--- a/src/nested_pandas/series/ext_array.py
+++ b/src/nested_pandas/series/ext_array.py
@@ -577,10 +577,10 @@ def __arrow_array__(self, type=None):
         """Convert the extension array to a PyArrow array."""
         # struct_array is the default "external" representation
         if type is None:
-            return self.struct_array
-        if is_pa_type_a_list(type):
-            return self.list_array.cast(type)
-        return self.struct_array.cast(type)
+            return self.list_array
+        if pa.types.is_struct(type):
+            return self.struct_array.cast(type)
+        return self.list_array.cast(type)
 
     def __array__(self, dtype=None):
         """Convert the extension array to a numpy array."""
diff --git a/tests/nested_pandas/series/test_ext_array.py b/tests/nested_pandas/series/test_ext_array.py
index 84cbb4eb..e1264808 100644
--- a/tests/nested_pandas/series/test_ext_array.py
+++ b/tests/nested_pandas/series/test_ext_array.py
@@ -1063,17 +1063,13 @@ def test_dropna():
 
 def test___arrow_array__():
     """Test that the extension array can be converted to a pyarrow array."""
-    struct_array = pa.StructArray.from_arrays(
-        arrays=[
-            pa.array([np.array([1, 2, 3]), np.array([1, 2, 1])]),
-            pa.array([-np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0])]),
-        ],
-        names=["a", "b"],
+    list_array = pa.array(
+        [[{"a": 1, "b": 2}, {"a": 3, "b": 4}], [{"a": 5, "b": 6}], [{"a": -1, "b": -2}], []]
     )
-    ext_array = NestedExtensionArray(struct_array)
+    ext_array = NestedExtensionArray(list_array)
 
     arrow_array = pa.array(ext_array)
-    assert arrow_array == struct_array
+    assert arrow_array == list_array
 
 
 def test___arrow_array___with_type_cast():
@@ -1891,14 +1887,7 @@ def test___init___with_list_struct_array():
     ext_array = NestedExtensionArray(list_array)
     assert ext_array.field_names == ["a", "b"]
     assert ext_array.flat_length == 4
-    struct_array = pa.StructArray.from_arrays(
-        arrays=[
-            pa.array([[1, 3], [5], [-1], []]),
-            pa.array([[2, 4], [6], [-2], []]),
-        ],
-        names=["a", "b"],
-    )
-    assert pa.array(ext_array) == struct_array
+    assert pa.array(ext_array) == list_array
 
 
 def test__struct_array():


Reviews:
user: login: dougbrn
body: Just one performance note, but no changes requested
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-09T18:43:02Z
user: login: hombit
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T19:17:45Z
user: login: dougbrn
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-09T19:19:17Z

Review comments:
diff_hunk: @@ -1269,13 +1291,14 @@ def to_parquet(self, path, **kwargs) -> None:
         >>> nf = generate_data(5,5, seed=1)
         >>> nf.to_parquet("nestedframe.parquet")
         """
+        df = self.to_pandas(list_struct=False)
path: src/nested_pandas/nestedframe/core.py
user: login: dougbrn
body: Not neccesarily a blocking issue, but just want to acknowledge that we're adding an additional copy operation to our parquet writing here. At scale if we run into memory issues this can be a good thing to come back to.
created_at: 2025-05-09T18:42:29Z
author_association: COLLABORATOR
diff_hunk: @@ -1269,13 +1291,14 @@ def to_parquet(self, path, **kwargs) -> None:
         >>> nf = generate_data(5,5, seed=1)
         >>> nf.to_parquet("nestedframe.parquet")
         """
+        df = self.to_pandas(list_struct=False)
path: src/nested_pandas/nestedframe/core.py
user: login: hombit
body: I believe it is a shallow copy

```python
In [1]: import os
   ...: import psutil
   ...:
   ...: process = psutil.Process(os.getpid())
   ...: print(f"{process.memory_info().rss / 1024 ** 2:.2f} MiB")
175.52 MiB

In [2]: from nested_pandas.datasets import generate_data

In [3]: nf = generate_data(100_000, 100)

In [4]: print(f"{process.memory_info().rss / 1024 ** 2:.2f} MiB")
1460.98 MiB

In [5]: %time df = nf.to_pandas()
CPU times: user 527 μs, sys: 56 μs, total: 583 μs
Wall time: 516 μs

In [6]: print(f"{process.memory_info().rss / 1024 ** 2:.2f} MiB")
1461.58 MiB
```
created_at: 2025-05-09T19:17:45Z
author_association: COLLABORATOR
diff_hunk: @@ -1269,13 +1291,14 @@ def to_parquet(self, path, **kwargs) -> None:
         >>> nf = generate_data(5,5, seed=1)
         >>> nf.to_parquet("nestedframe.parquet")
         """
+        df = self.to_pandas(list_struct=False)
path: src/nested_pandas/nestedframe/core.py
user: login: dougbrn
body: Ah, wonderful thank you for verifying that!
created_at: 2025-05-09T19:19:17Z
author_association: COLLABORATOR

PR comments:
user: login: codecov[bot]
created_at: 2025-05-09T18:39:28Z
body: ## [Codecov](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/pull/257?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks) Report
:white_check_mark: All modified and coverable lines are covered by tests.
:white_check_mark: Project coverage is 97.90%. Comparing base ([`5f43b2d`](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/commit/5f43b2d43d28c72b316b168e13546edf41fe313c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks)) to head ([`0bad737`](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/commit/0bad737cf4991ce7c815cf1cb16d8b07c695eb3b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks)).
:warning: Report is 95 commits behind head on main.

<details><summary>Additional details and impacted files</summary>



```diff
@@           Coverage Diff           @@
##             main     #257   +/-   ##
=======================================
  Coverage   97.89%   97.90%           
=======================================
  Files          18       18           
  Lines        1473     1479    +6     
=======================================
+ Hits         1442     1448    +6     
  Misses         31       31           
```
</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/lincc-frameworks/nested-pandas/pull/257?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=lincc-frameworks).
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-05-09T18:39:33Z
body: | Before [5f43b2d4]    | After [90ce4ef9]    |   Ratio | Benchmark (Parameter)                               |
|----------------------|---------------------|---------|-----------------------------------------------------|
| 46.3±0.6ms           | 47.5±0.5ms          |    1.03 | benchmarks.ReassignHalfOfNestedSeries.time_run      |
| 25.6±0.4ms           | 25.8±0.1ms          |    1.01 | benchmarks.AssignSingleDfToNestedSeries.time_run    |
| 275M                 | 275M                |    1    | benchmarks.AssignSingleDfToNestedSeries.peakmem_run |
| 99.8M                | 99.8M               |    1    | benchmarks.NestedFrameAddNested.peakmem_run         |
| 103M                 | 104M                |    1    | benchmarks.NestedFrameQuery.peakmem_run             |
| 103M                 | 104M                |    1    | benchmarks.NestedFrameReduce.peakmem_run            |
| 1.34±0.01ms          | 1.34±0.02ms         |    1    | benchmarks.NestedFrameReduce.time_run               |
| 293M                 | 294M                |    1    | benchmarks.ReassignHalfOfNestedSeries.peakmem_run   |
| 10.4±0.07ms          | 10.3±0.05ms         |    0.99 | benchmarks.NestedFrameAddNested.time_run            |
| 11.1±0.08ms          | 11.1±0.04ms         |    0.99 | benchmarks.NestedFrameQuery.time_run                |


Click [here](https://github.com/lincc-frameworks/nested-pandas/actions/runs/14935471657/job/41961747489#step:10:1) to view all benchmarks.
author_association: NONE