[arrow] branch master updated: ARROW-3261: [Python] Add "field" method to select fields from StructArray

[we...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wesm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/master by this push:
     new cdd458e  ARROW-3261: [Python] Add "field" method to select fields from StructArray
cdd458e is described below

commit cdd458e6ab74a1bdcc22d56bff72af9875271e0f
Author: Krisztián Szűcs <sz...@gmail.com>
AuthorDate: Thu Sep 20 17:17:58 2018 -0400

    ARROW-3261: [Python] Add "field" method to select fields from StructArray
    
    Author: Krisztián Szűcs <sz...@gmail.com>
    
    Closes #2586 from kszucs/ARROW-3261 and squashes the following commits:
    
    39750520a <Krisztián Szűcs> docstring
    248ae7ae4 <Krisztián Szűcs> StructArray.field
---
 python/pyarrow/array.pxi           | 18 ++++++++++++++++++
 python/pyarrow/tests/test_array.py | 27 +++++++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi
index 8962ac4..7de2340 100644
--- a/python/pyarrow/array.pxi
+++ b/python/pyarrow/array.pxi
@@ -1128,6 +1128,24 @@ cdef class DictionaryArray(Array):
 
 cdef class StructArray(Array):
 
+    def field(self, index):
+        """
+        Retrieves the child array belonging to field
+
+        Parameters
+        ----------
+        index : int
+            Index / position of the field
+
+        Returns
+        -------
+        result : Array
+        """
+        cdef:
+            int ix = <int> _normalize_index(index, self.ap.num_fields())
+            CStructArray* sarr = <CStructArray*> self.ap
+        return pyarrow_wrap_array(sarr.field(ix))
+
     def flatten(self, MemoryPool memory_pool=None):
         """
         Flatten this StructArray, returning one individual array for each
diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py
index 89f4d3d..410c9d0 100644
--- a/python/pyarrow/tests/test_array.py
+++ b/python/pyarrow/tests/test_array.py
@@ -1121,6 +1121,33 @@ def test_struct_array_flatten():
     assert ys.to_pylist() == [None, 2.5]
 
 
+def test_struct_array_field():
+    ty = pa.struct([pa.field('x', pa.int16()),
+                    pa.field('y', pa.float32())])
+    a = pa.array([(1, 2.5), (3, 4.5), (5, 6.5)], type=ty)
+
+    x = a.field(0)
+    y = a.field(1)
+    x_ = a.field(-2)
+    y_ = a.field(-1)
+
+    assert isinstance(x, pa.lib.Int16Array)
+    assert isinstance(y, pa.lib.FloatArray)
+
+    assert x.equals(pa.array([1, 3, 5], type=pa.int16()))
+    assert y.equals(pa.array([2.5, 4.5, 6.5], type=pa.float32()))
+    assert x.equals(x_)
+    assert y.equals(y_)
+
+    for invalid_index in [None, 'x']:
+        with pytest.raises(TypeError):
+            a.field(invalid_index)
+
+    for invalid_index in [3, -3]:
+        with pytest.raises(IndexError):
+            a.field(invalid_index)
+
+
 def test_nested_dictionary_array():
     dict_arr = pa.DictionaryArray.from_arrays([0, 1, 0], ['a', 'b'])
     list_arr = pa.ListArray.from_arrays([0, 2, 3], dict_arr)