Merge pull request #56 from martindurant/auto_to_output

Allow ak.to_output()
intake · Jun 11, 2024 · 2480e3d · 2480e3d
2 parents 621a65c + d966a9e
commit 2480e3d
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/src/awkward_pandas/mixin.py b/src/awkward_pandas/mixin.py
@@ -150,7 +150,8 @@ def is_dataframe(cls, data):
     def _to_output(cls, data):
         raise NotImplementedError
 
-    def to_output(self, data):
+    def to_output(self, data=None):
+        data = data if data is not None else self.array
         return self._to_output(data)
 
     def apply(self, fn: Callable):

diff --git a/src/awkward_pandas/pandas.py b/src/awkward_pandas/pandas.py
@@ -23,8 +23,9 @@ def _to_output(cls, data):
             pd.arrays.ArrowExtensionArray(ak.to_arrow(data, extensionarray=False))
         )
 
-    def to_output(self, data):
+    def to_output(self, data=None):
         # override to apply index
+        data = data if data is not None else self.array
         arr = pd.arrays.ArrowExtensionArray(ak.to_arrow(data, extensionarray=False))
         if self._obj is not None and len(arr) == len(self._obj.index):
             return pd.Series(arr, index=self._obj.index)

diff --git a/tests/test_pandas.py b/tests/test_pandas.py
@@ -48,3 +48,11 @@ def test_ufunc():
 
     assert (s.ak + s.ak).tolist() == [[2, 4, 6], [8, 10], [12]]
     assert (s.ak + s).tolist() == [[2, 4, 6], [8, 10], [12]]
+
+
+def test_to_autoarrow():
+    a = [[1, 2, 3], [4, 5], [6]]
+    s = pd.Series(a)
+    s2 = s.ak.to_output()
+    assert s2.tolist() == a
+    assert "pyarrow" in str(s2.dtype)