[arrow] branch master updated: ARROW-2453: [Python] Improve Table column access

[uw...@apache.org]

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

uwe 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 a6c9d30  ARROW-2453: [Python] Improve Table column access
a6c9d30 is described below

commit a6c9d30c0f28d3beff238ced876746b957b86109
Author: Kee Chong Tan <keechong.tan@>
AuthorDate: Sat Apr 21 21:22:32 2018 +0100

    ARROW-2453: [Python] Improve Table column access
    
    Author: Kee Chong Tan <keechong.tan@>
    
    Closes #1923 from keechongtan/ARROW-2453 and squashes the following commits:
    
    c65033c2 <Kee Chong Tan> Remove whitespace
    0527926e <Kee Chong Tan> Improve Table column access
---
 python/pyarrow/table.pxi           | 25 ++++++++++++++++++++++++-
 python/pyarrow/tests/test_table.py | 17 +++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
index cbf2a69..cbbfe7d 100644
--- a/python/pyarrow/table.pxi
+++ b/python/pyarrow/table.pxi
@@ -1152,7 +1152,30 @@ cdef class Table:
         self._check_nullptr()
         return pyarrow_wrap_schema(self.table.schema())
 
-    def column(self, int i):
+    def column(self, i):
+        """
+        Select a column by its column name, or numeric index.
+
+        Parameters
+        ----------
+        i : int or string
+
+        Returns
+        -------
+        pyarrow.Column
+        """
+        if isinstance(i, six.string_types):
+            field_index = self.schema.get_field_index(i)
+            if field_index < 0:
+                raise KeyError("Column {} does not exist in table".format(i))
+            else:
+                return self._column(field_index)
+        elif isinstance(i, six.integer_types):
+            return self._column(i)
+        else:
+            raise TypeError("Index must either be string or integer")
+
+    def _column(self, int i):
         """
         Select a column by its numeric index.
 
diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py
index 57be9b2..11555c6 100644
--- a/python/pyarrow/tests/test_table.py
+++ b/python/pyarrow/tests/test_table.py
@@ -301,6 +301,23 @@ def test_table_from_arrays_invalid_names():
         pa.Table.from_arrays(data, names=['a'])
 
 
+def test_table_select_column():
+    data = [
+        pa.array(range(5)),
+        pa.array([-10, -5, 0, 5, 10]),
+        pa.array(range(5, 10))
+    ]
+    table = pa.Table.from_arrays(data, names=('a', 'b', 'c'))
+
+    assert table.column('a').equals(table.column(0))
+
+    with pytest.raises(KeyError):
+        table.column('d')
+
+    with pytest.raises(TypeError):
+        table.column(None)
+
+
 def test_table_add_column():
     data = [
         pa.array(range(5)),

-- 
To stop receiving notification emails like this one, please contact
uwe@apache.org.