Add Categorical conversion for single array or column. Required moving code around a little bit in pandas.cc

Change-Id: Id48db372aa14aa98dbd2fdc5de6ffdb4fe6e0032

cpp/src/arrow/type.cc

@@ -148,8 +148,8 @@ bool DictionaryType::Equals(const DataType& other) const {
 
 std::string DictionaryType::ToString() const {
   std::stringstream ss;
-  ss << "dictionary<" << dictionary_->type()->ToString() << ", "
-     << index_type_->ToString() << ">";
+  ss << "dictionary<values=" << dictionary_->type()->ToString()
+     << ", indices=" << index_type_->ToString() << ">";
   return ss.str();
 }
 

python/pyarrow/array.pxd

@@ -22,6 +22,8 @@ from pyarrow.scalar import NA
 
 from pyarrow.schema cimport DataType
 
+from cpython cimport PyObject
+
 cdef extern from "Python.h":
     int PySlice_Check(object)
 
@@ -108,3 +110,6 @@ cdef class BinaryArray(Array):
 
 cdef class DictionaryArray(Array):
     pass
+
+
+cdef wrap_array_output(PyObject* output)

python/pyarrow/array.pyx

@@ -221,12 +221,13 @@ cdef class Array:
         RecordBatch.to_pandas
         """
         cdef:
-            PyObject* np_arr
+            PyObject* out
 
-        check_status(pyarrow.ConvertArrayToPandas(
-            self.sp_array, <PyObject*> self, &np_arr))
-
-        return PyObject_to_object(np_arr)
+        with nogil:
+            check_status(
+                pyarrow.ConvertArrayToPandas(self.sp_array, <PyObject*> self,
+                                             &out))
+        return wrap_array_output(out)
 
     def to_pylist(self):
         """
@@ -235,6 +236,19 @@ cdef class Array:
         return [x.as_py() for x in self]
 
 
+cdef wrap_array_output(PyObject* output):
+    cdef object obj = PyObject_to_object(output)
+
+    if isinstance(obj, dict):
+        return _pandas().Categorical(obj['indices'],
+                                     categories=obj['dictionary'],
+                                     fastpath=True)
+    else:
+        return obj
+
+    return values
+
+
 cdef class NullArray(Array):
     pass
 

python/pyarrow/table.pyx

@@ -27,7 +27,7 @@ cimport pyarrow.includes.pyarrow as pyarrow
 
 import pyarrow.config
 
-from pyarrow.array cimport Array, box_arrow_array
+from pyarrow.array cimport Array, box_arrow_array, wrap_array_output
 from pyarrow.error import ArrowException
 from pyarrow.error cimport check_status
 from pyarrow.schema cimport box_data_type, box_schema, Field
@@ -151,12 +151,12 @@ cdef class Column:
         pandas.Series
         """
         cdef:
-            PyObject* arr
+            PyObject* out
 
         check_status(pyarrow.ConvertColumnToPandas(self.sp_column,
-                                                   <PyObject*> self, &arr))
+                                                   <PyObject*> self, &out))
 
-        return _pandas().Series(PyObject_to_object(arr), name=self.name)
+        return _pandas().Series(wrap_array_output(out), name=self.name)
 
     def equals(self, Column other):
         """
@@ -489,7 +489,7 @@ cdef table_to_blockmanager(const shared_ptr[CTable]& table, int nthreads):
         block_arr = item['block']
         placement = item['placement']
         if 'dictionary' in item:
-            cat = Categorical(block_arr[0],
+            cat = Categorical(block_arr,
                               categories=item['dictionary'],
                               ordered=False, fastpath=True)
             block = _int.make_block(cat, placement=placement,

python/pyarrow/tests/test_column.py

@@ -47,4 +47,3 @@ class TestColumn(unittest.TestCase):
         assert series.name == 'a'
         assert series.shape == (5,)
         assert series.iloc[0] == -10
-

python/pyarrow/tests/test_convert_pandas.py

@@ -73,6 +73,13 @@ class TestPandasConversion(unittest.TestCase):
             expected = df
         tm.assert_frame_equal(result, expected)
 
+    def _check_array_roundtrip(self, values, expected=None,
+                                timestamps_to_ms=False, field=None):
+        arr = A.Array.from_pandas(values, timestamps_to_ms=timestamps_to_ms,
+                                  field=field)
+        result = arr.to_pandas()
+        tm.assert_series_equal(pd.Series(result), pd.Series(values))
+
     def test_float_no_nulls(self):
         data = {}
         fields = []
@@ -360,3 +367,11 @@ class TestPandasConversion(unittest.TestCase):
                            'strings2': v1 * repeats,
                            'strings3': v3 * repeats})
         self._check_pandas_roundtrip(df)
+
+        arrays = [
+            pd.Categorical(v1 * repeats),
+            pd.Categorical(v2 * repeats),
+            pd.Categorical(v3 * repeats)
+        ]
+        for values in arrays:
+            self._check_array_roundtrip(values)