Draft Python wrapper classes for DictionaryType, DictionaryArray. Avoid

segfault when conversion to UTF8 fails. Starting on CategoricalBlock
implementation

cpp/src/arrow/array.h

@@ -553,6 +553,8 @@ class ARROW_EXPORT DictionaryArray : public Array {
   std::shared_ptr<Array> indices() const { return indices_; }
   std::shared_ptr<Array> dictionary() const;
 
+  const DictionaryType* dict_type() { return dict_type_; }
+
   bool EqualsExact(const DictionaryArray& other) const;
   bool Equals(const std::shared_ptr<Array>& arr) const override;
   bool RangeEquals(int32_t start_idx, int32_t end_idx, int32_t other_start_idx,

cpp/src/arrow/builder.cc

@@ -421,9 +421,12 @@ Status MakeBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type,
     BUILDER_CASE(FLOAT, FloatBuilder);
     BUILDER_CASE(DOUBLE, DoubleBuilder);
 
-    BUILDER_CASE(STRING, StringBuilder);
-    BUILDER_CASE(BINARY, BinaryBuilder);
-
+    case Type::STRING:
+      out->reset(new StringBuilder(pool));
+      return Status::OK();
+    case Type::BINARY:
+      out->reset(new BinaryBuilder(pool, type));
+      return Status::OK();
     case Type::LIST: {
       std::shared_ptr<ArrayBuilder> value_builder;
       std::shared_ptr<DataType> value_type =

cpp/src/arrow/builder.h

@@ -24,6 +24,7 @@
 #include <vector>
 
 #include "arrow/buffer.h"
+#include "arrow/memory_pool.h"
 #include "arrow/status.h"
 #include "arrow/type.h"
 #include "arrow/util/bit-util.h"
@@ -33,7 +34,6 @@
 namespace arrow {
 
 class Array;
-class MemoryPool;
 
 static constexpr int32_t kMinBuilderCapacity = 1 << 5;
 
@@ -378,8 +378,11 @@ class ARROW_EXPORT BinaryBuilder : public ListBuilder {
 // String builder
 class ARROW_EXPORT StringBuilder : public BinaryBuilder {
  public:
-  explicit StringBuilder(MemoryPool* pool, const TypePtr& type)
-      : BinaryBuilder(pool, type) {}
+  explicit StringBuilder(MemoryPool* pool = default_memory_pool())
+    : BinaryBuilder(pool, utf8()) {}
+
+  explicit StringBuilder(MemoryPool* pool, const std::shared_ptr<DataType>& type)
+    : BinaryBuilder(pool, type) {}
 
   using BinaryBuilder::Append;
 

python/CMakeLists.txt

@@ -55,6 +55,10 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}")
     OFF)
 endif()
 
+if(NOT PYARROW_BUILD_TESTS)
+  set(NO_TESTS 1)
+endif()
+
 find_program(CCACHE_FOUND ccache)
 if(CCACHE_FOUND)
   set_property(GLOBAL PROPERTY RULE_LAUNCH_COMPILE ccache)
@@ -339,6 +343,12 @@ set(PYARROW_MIN_TEST_LIBS
   pyarrow
   ${PYARROW_BASE_LIBS})
 
+if(NOT APPLE)
+  ADD_THIRDPARTY_LIB(python
+    SHARED_LIB "${PYTHON_LIBRARIES}")
+  list(APPEND PYARROW_MIN_TEST_LIBS python)
+endif()
+
 set(PYARROW_TEST_LINK_LIBS ${PYARROW_MIN_TEST_LIBS})
 
 ############################################################

python/pyarrow/__init__.py

@@ -31,9 +31,14 @@ from pyarrow.config import cpu_count, set_cpu_count
 from pyarrow.array import (Array,
                            from_pandas_series, from_pylist,
                            total_allocated_bytes,
-                           BooleanArray, NumericArray,
+                           NumericArray, IntegerArray, FloatingPointArray,
+                           BooleanArray,
                            Int8Array, UInt8Array,
-                           ListArray, StringArray)
+                           Int16Array, UInt16Array,
+                           Int32Array, UInt32Array,
+                           Int64Array, UInt64Array,
+                           ListArray, StringArray,
+                           DictionaryArray, dictionary)
 
 from pyarrow.error import ArrowException
 

python/pyarrow/array.pxd

@@ -47,35 +47,50 @@ cdef class NumericArray(Array):
     pass
 
 
-cdef class Int8Array(NumericArray):
+cdef class IntegerArray(NumericArray):
+    pass
+
+cdef class FloatingPointArray(NumericArray):
     pass
 
 
-cdef class UInt8Array(NumericArray):
+cdef class Int8Array(IntegerArray):
     pass
 
 
-cdef class Int16Array(NumericArray):
+cdef class UInt8Array(IntegerArray):
     pass
 
 
-cdef class UInt16Array(NumericArray):
+cdef class Int16Array(IntegerArray):
     pass
 
 
-cdef class Int32Array(NumericArray):
+cdef class UInt16Array(IntegerArray):
     pass
 
 
-cdef class UInt32Array(NumericArray):
+cdef class Int32Array(IntegerArray):
     pass
 
 
-cdef class Int64Array(NumericArray):
+cdef class UInt32Array(IntegerArray):
     pass
 
 
-cdef class UInt64Array(NumericArray):
+cdef class Int64Array(IntegerArray):
+    pass
+
+
+cdef class UInt64Array(IntegerArray):
+    pass
+
+
+cdef class FloatArray(FloatingPointArray):
+    pass
+
+
+cdef class DoubleArray(FloatingPointArray):
     pass
 
 
@@ -85,3 +100,11 @@ cdef class ListArray(Array):
 
 cdef class StringArray(Array):
     pass
+
+
+cdef class BinaryArray(Array):
+    pass
+
+
+cdef class DictionaryArray(Array):
+    pass

python/pyarrow/array.pyx

@@ -33,12 +33,17 @@ from pyarrow.error cimport check_status
 cimport pyarrow.scalar as scalar
 from pyarrow.scalar import NA
 
-from pyarrow.schema cimport Field, Schema
+from pyarrow.schema cimport Field, Schema, DictionaryType
 import pyarrow.schema as schema
 
 cimport cpython
 
 
+cdef _pandas():
+    import pandas as pd
+    return pd
+
+
 def total_allocated_bytes():
     cdef MemoryPool* pool = pyarrow.get_memory_pool()
     return pool.bytes_allocated()
@@ -53,20 +58,22 @@ cdef class Array:
         self.type.init(self.sp_array.get().type())
 
     @staticmethod
-    def from_pandas(obj, mask=None):
+    def from_pandas(obj, mask=None, timestamps_to_ms=False, Field field=None):
         """
-        Create an array from a pandas.Series
+        Convert pandas.Series to an Arrow Array.
 
         Parameters
         ----------
-        obj : pandas.Series or numpy.ndarray
-            vector holding the data
-        mask : numpy.ndarray, optional
+        series : pandas.Series or numpy.ndarray
+
+        mask : pandas.Series or numpy.ndarray, optional
             boolean mask if the object is valid or null
 
-        Returns
-        -------
-        pyarrow.Array
+        timestamps_to_ms : bool, optional
+            Convert datetime columns to ms resolution. This is needed for
+            compability with other functionality like Parquet I/O which
+            only supports milliseconds.
+
 
         Examples
         --------
@@ -80,16 +87,47 @@ cdef class Array:
           2
         ]
 
-
         >>> import numpy as np
-        >>> pa.Array.from_pandas(pd.Series([1, 2]), np.array([0, 1], dtype=bool))
+        >>> pa.Array.from_pandas(pd.Series([1, 2]), np.array([0, 1],
+        ... dtype=bool))
         <pyarrow.array.Int64Array object at 0x7f9019e11208>
         [
           1,
           NA
         ]
+
+        Returns
+        -------
+        pyarrow.array.Array
         """
-        return from_pandas_series(obj, mask)
+        cdef:
+            shared_ptr[CArray] out
+            shared_ptr[CField] c_field
+
+        pd = _pandas()
+
+        series_values = get_series_values(obj)
+        if series_values.dtype.type == np.datetime64 and timestamps_to_ms:
+            series_values = series_values.astype('datetime64[ms]')
+
+        if field is not None:
+            c_field = field.sp_field
+
+        if mask is not None:
+            mask = get_series_values(mask)
+
+        if isinstance(series_values, pd.Categorical):
+            indices = Array.from_pandas(series_values.codes)
+            dictionary = Array.from_pandas(series_values.categories.values)
+            return DictionaryArray.from_arrays(indices, dictionary,
+                                               mask=mask)
+        else:
+            with nogil:
+                check_status(pyarrow.PandasToArrow(
+                    pyarrow.get_memory_pool(), series_values, mask,
+                    c_field, &out))
+
+            return box_arrow_array(out)
 
     @staticmethod
     def from_list(object list_obj, DataType type=None):
@@ -209,35 +247,43 @@ cdef class NumericArray(Array):
     pass
 
 
-cdef class Int8Array(NumericArray):
+cdef class IntegerArray(NumericArray):
     pass
 
 
-cdef class UInt8Array(NumericArray):
+cdef class FloatingPointArray(NumericArray):
     pass
 
 
-cdef class Int16Array(NumericArray):
+cdef class Int8Array(IntegerArray):
     pass
 
 
-cdef class UInt16Array(NumericArray):
+cdef class UInt8Array(IntegerArray):
     pass
 
 
-cdef class Int32Array(NumericArray):
+cdef class Int16Array(IntegerArray):
     pass
 
 
-cdef class UInt32Array(NumericArray):
+cdef class UInt16Array(IntegerArray):
     pass
 
 
-cdef class Int64Array(NumericArray):
+cdef class Int32Array(IntegerArray):
     pass
 
 
-cdef class UInt64Array(NumericArray):
+cdef class UInt32Array(IntegerArray):
+    pass
+
+
+cdef class Int64Array(IntegerArray):
+    pass
+
+
+cdef class UInt64Array(IntegerArray):
     pass
 
 
@@ -245,11 +291,11 @@ cdef class DateArray(NumericArray):
     pass
 
 
-cdef class FloatArray(NumericArray):
+cdef class FloatArray(FloatingPointArray):
     pass
 
 
-cdef class DoubleArray(NumericArray):
+cdef class DoubleArray(FloatingPointArray):
     pass
 
 
@@ -265,6 +311,46 @@ cdef class BinaryArray(Array):
     pass
 
 
+cdef class DictionaryArray(Array):
+
+    @staticmethod
+    def from_arrays(indices, dictionary, mask=None):
+        """
+        Construct Arrow DictionaryArray from array of indices (must be
+        non-negative integers) and corresponding array of dictionary values
+
+        Parameters
+        ----------
+        indices : ndarray or pandas.Series, integer type
+        dictionary : ndarray or pandas.Series
+        mask : ndarray or pandas.Series, boolean type
+            True values indicate that indices are actually null
+
+        Returns
+        -------
+        dict_array : DictionaryArray
+        """
+        cdef:
+            Array arrow_indices, arrow_dictionary
+            DictionaryArray result
+            shared_ptr[CDataType] c_type
+            shared_ptr[CArray] c_result
+
+        arrow_indices = Array.from_pandas(indices, mask=mask)
+        arrow_dictionary = Array.from_pandas(dictionary)
+
+        if not isinstance(arrow_indices, IntegerArray):
+            raise ValueError('Indices must be integer type')
+
+        c_type.reset(new CDictionaryType(arrow_indices.type.sp_type,
+                                         arrow_dictionary.sp_array))
+        c_result.reset(new CDictionaryArray(c_type, arrow_indices.sp_array))
+
+        result = DictionaryArray()
+        result.init(c_result)
+        return result
+
+
 cdef dict _array_classes = {
     Type_NA: NullArray,
     Type_BOOL: BooleanArray,
@@ -283,6 +369,7 @@ cdef dict _array_classes = {
     Type_BINARY: BinaryArray,
     Type_STRING: StringArray,
     Type_TIMESTAMP: Int64Array,
+    Type_DICTIONARY: DictionaryArray
 }
 
 cdef object box_arrow_array(const shared_ptr[CArray]& sp_array):
@@ -299,76 +386,7 @@ cdef object box_arrow_array(const shared_ptr[CArray]& sp_array):
     return arr
 
 
-def from_pylist(object list_obj, DataType type=None):
-    """
-    Convert Python list to Arrow array
-
-    Parameters
-    ----------
-    list_obj : array_like
-
-    Returns
-    -------
-    pyarrow.array.Array
-    """
-    cdef:
-        shared_ptr[CArray] sp_array
-
-    if type is None:
-        check_status(pyarrow.ConvertPySequence(list_obj, &sp_array))
-    else:
-        raise NotImplementedError()
-
-    return box_arrow_array(sp_array)
-
-
-def from_pandas_series(object series, object mask=None, timestamps_to_ms=False, Field field=None):
-    """
-    Convert pandas.Series to an Arrow Array.
-
-    Parameters
-    ----------
-    series : pandas.Series or numpy.ndarray
-
-    mask : pandas.Series or numpy.ndarray, optional
-        array to mask null entries in the series
-
-    timestamps_to_ms : bool, optional
-        Convert datetime columns to ms resolution. This is needed for
-        compability with other functionality like Parquet I/O which
-        only supports milliseconds.
-
-    field: pyarrow.Field
-        Schema indicator to what type this column should render in Arrow
-
-    Returns
-    -------
-    pyarrow.array.Array
-    """
-    cdef:
-        shared_ptr[CArray] out
-        shared_ptr[CField] c_field
-
-    series_values = series_as_ndarray(series)
-    if series_values.dtype.type == np.datetime64 and timestamps_to_ms:
-        series_values = series_values.astype('datetime64[ms]')
-    if field is not None:
-        c_field = field.sp_field
-
-    if mask is None:
-        with nogil:
-            check_status(pyarrow.PandasToArrow(pyarrow.get_memory_pool(),
-                                               series_values, c_field, &out))
-    else:
-        mask = series_as_ndarray(mask)
-        with nogil:
-            check_status(pyarrow.PandasMaskedToArrow(
-                pyarrow.get_memory_pool(), series_values, mask, c_field, &out))
-
-    return box_arrow_array(out)
-
-
-cdef object series_as_ndarray(object obj):
+cdef object get_series_values(object obj):
     import pandas as pd
 
     if isinstance(obj, pd.Series):
@@ -378,4 +396,17 @@ cdef object series_as_ndarray(object obj):
 
     return result
 
+
 from_pylist = Array.from_list
+from_pandas_series = Array.from_pandas
+
+
+def dictionary(DataType index_type, Array dictionary):
+    """
+    Dictionary (categorical, or simply encoded) type
+    """
+    cdef DictionaryType out = schema.DictionaryType()
+    cdef shared_ptr[CDataType] dict_type
+    dict_type.reset(new CDictionaryType(index_type.sp_type, dictionary.sp_array))
+    out.init(dict_type)
+    return out

python/pyarrow/includes/libarrow.pxd

@@ -45,6 +45,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
         Type_LIST" arrow::Type::LIST"
         Type_STRUCT" arrow::Type::STRUCT"
+        Type_DICTIONARY" arrow::Type::DICTIONARY"
 
     enum TimeUnit" arrow::TimeUnit":
         TimeUnit_SECOND" arrow::TimeUnit::SECOND"
@@ -60,6 +61,33 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
 
         c_string ToString()
 
+    cdef cppclass CArray" arrow::Array":
+        shared_ptr[CDataType] type()
+
+        int32_t length()
+        int32_t null_count()
+        Type type_enum()
+
+        c_bool Equals(const shared_ptr[CArray]& arr)
+        c_bool IsNull(int i)
+
+    cdef cppclass CFixedWidthType" arrow::FixedWidthType"(CDataType):
+        int bit_width()
+
+    cdef cppclass CDictionaryArray" arrow::DictionaryArray"(CArray):
+        CDictionaryArray(const shared_ptr[CDataType]& type,
+                         const shared_ptr[CArray]& indices)
+
+        shared_ptr[CArray] indices()
+        shared_ptr[CArray] dictionary()
+
+    cdef cppclass CDictionaryType" arrow::DictionaryType"(CFixedWidthType):
+        CDictionaryType(const shared_ptr[CDataType]& index_type,
+                        const shared_ptr[CArray]& dictionary)
+
+        shared_ptr[CDataType] index_type()
+        shared_ptr[CArray] dictionary()
+
     shared_ptr[CDataType] timestamp(TimeUnit unit)
 
     cdef cppclass MemoryPool" arrow::MemoryPool":
@@ -111,16 +139,6 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         int num_fields()
         c_string ToString()
 
-    cdef cppclass CArray" arrow::Array":
-        shared_ptr[CDataType] type()
-
-        int32_t length()
-        int32_t null_count()
-        Type type_enum()
-
-        c_bool Equals(const shared_ptr[CArray]& arr)
-        c_bool IsNull(int i)
-
     cdef cppclass CBooleanArray" arrow::BooleanArray"(CArray):
         c_bool Value(int i)
 

python/pyarrow/includes/pyarrow.pxd

@@ -30,11 +30,9 @@ cdef extern from "pyarrow/api.h" namespace "pyarrow" nogil:
     shared_ptr[CDataType] GetTimestampType(TimeUnit unit)
     CStatus ConvertPySequence(object obj, shared_ptr[CArray]* out)
 
-    CStatus PandasToArrow(MemoryPool* pool, object ao, shared_ptr[CField] field,
+    CStatus PandasToArrow(MemoryPool* pool, object ao, object mo,
+                          shared_ptr[CField] field,
                           shared_ptr[CArray]* out)
-    CStatus PandasMaskedToArrow(MemoryPool* pool, object ao, object mo,
-                                shared_ptr[CField] field,
-                                shared_ptr[CArray]* out)
 
     CStatus ConvertArrayToPandas(const shared_ptr[CArray]& arr,
                                  PyObject* py_ref, PyObject** out)

python/pyarrow/schema.pxd

@@ -16,7 +16,8 @@
 # under the License.
 
 from pyarrow.includes.common cimport *
-from pyarrow.includes.libarrow cimport CDataType, CField, CSchema
+from pyarrow.includes.libarrow cimport (CDataType, CDictionaryType,
+                                        CField, CSchema)
 
 cdef class DataType:
     cdef:
@@ -25,6 +26,11 @@ cdef class DataType:
 
     cdef init(self, const shared_ptr[CDataType]& type)
 
+
+cdef class DictionaryType(DataType):
+    cdef:
+        const CDictionaryType* dict_type
+
 cdef class Field:
     cdef:
         shared_ptr[CField] sp_field

python/pyarrow/schema.pyx

@@ -66,6 +66,19 @@ cdef class DataType:
             raise TypeError('Invalid comparison')
 
 
+cdef class DictionaryType(DataType):
+
+    cdef init(self, const shared_ptr[CDataType]& type):
+        DataType.init(self, type)
+        self.dict_type = <const CDictionaryType*> type.get()
+
+    def __str__(self):
+        return frombytes(self.type.ToString())
+
+    def __repr__(self):
+        return 'DictionaryType({0})'.format(str(self))
+
+
 cdef class Field:
 
     def __cinit__(self):

python/pyarrow/table.pyx

@@ -39,6 +39,11 @@ cimport cpython
 from collections import OrderedDict
 
 
+cdef _pandas():
+    import pandas as pd
+    return pd
+
+
 cdef class ChunkedArray:
     """
     Array backed via one or more memory chunks.
@@ -148,12 +153,10 @@ cdef class Column:
         cdef:
             PyObject* arr
 
-        import pandas as pd
-
         check_status(pyarrow.ConvertColumnToPandas(self.sp_column,
                                                    <PyObject*> self, &arr))
 
-        return pd.Series(PyObject_to_object(arr), name=self.name)
+        return _pandas().Series(PyObject_to_object(arr), name=self.name)
 
     def equals(self, Column other):
         """
@@ -278,8 +281,6 @@ cdef _schema_from_arrays(arrays, names, shared_ptr[CSchema]* schema):
 
 
 cdef _dataframe_to_arrays(df, name, timestamps_to_ms, Schema schema):
-    from pyarrow.array import from_pandas_series
-
     cdef:
         list names = []
         list arrays = []
@@ -289,9 +290,8 @@ cdef _dataframe_to_arrays(df, name, timestamps_to_ms, Schema schema):
         col = df[name]
         if schema is not None:
             field = schema.field_by_name(name)
-        arr = from_pandas_series(col, timestamps_to_ms=timestamps_to_ms,
-                                 field=field)
-
+        arr = Array.from_pandas(col, timestamps_to_ms=timestamps_to_ms,
+                                field=field)
         names.append(name)
         arrays.append(arr)
 
@@ -304,7 +304,8 @@ cdef class RecordBatch:
 
     Warning
     -------
-    Do not call this class's constructor directly, use one of the ``from_*`` methods instead.
+    Do not call this class's constructor directly, use one of the ``from_*``
+    methods instead.
     """
 
     def __cinit__(self):
@@ -401,7 +402,7 @@ cdef class RecordBatch:
         return OrderedDict(entries)
 
 
-    def to_pandas(self):
+    def to_pandas(self, nthreads=None):
         """
         Convert the arrow::RecordBatch to a pandas DataFrame
 
@@ -409,23 +410,7 @@ cdef class RecordBatch:
         -------
         pandas.DataFrame
         """
-        cdef:
-            PyObject* np_arr
-            shared_ptr[CArray] arr
-            Column column
-
-        import pandas as pd
-
-        names = []
-        data = []
-        for i in range(self.batch.num_columns()):
-            arr = self.batch.column(i)
-            check_status(pyarrow.ConvertArrayToPandas(arr, <PyObject*> self,
-                                                      &np_arr))
-            names.append(frombytes(self.batch.column_name(i)))
-            data.append(PyObject_to_object(np_arr))
-
-        return pd.DataFrame(dict(zip(names, data)), columns=names)
+        return Table.from_batches([self]).to_pandas(nthreads=nthreads)
 
     @classmethod
     def from_pandas(cls, df, schema=None):
@@ -518,7 +503,8 @@ cdef class Table:
 
     Warning
     -------
-    Do not call this class's constructor directly, use one of the ``from_*`` methods instead.
+    Do not call this class's constructor directly, use one of the ``from_*``
+    methods instead.
     """
 
     def __cinit__(self):
@@ -688,13 +674,11 @@ cdef class Table:
         -------
         pandas.DataFrame
         """
-        import pandas as pd
-
         if nthreads is None:
             nthreads = pyarrow.config.cpu_count()
 
         mgr = table_to_blockmanager(self.sp_table, nthreads)
-        return pd.DataFrame(mgr)
+        return _pandas().DataFrame(mgr)
 
     def to_pydict(self):
         """
@@ -835,6 +819,7 @@ cdef api object table_from_ctable(const shared_ptr[CTable]& ctable):
     table.init(ctable)
     return table
 
+
 cdef api object batch_from_cbatch(const shared_ptr[CRecordBatch]& cbatch):
     cdef RecordBatch batch = RecordBatch()
     batch.init(cbatch)

python/setup.py

@@ -128,6 +128,7 @@ class build_ext(_build_ext):
 
         cmake_options = [
             '-DPYTHON_EXECUTABLE=%s' % sys.executable,
+            '-DPYARROW_BUILD_TESTS=off',
             static_lib_option,
             build_tests_option,
         ]

python/src/pyarrow/CMakeLists.txt

@@ -18,3 +18,5 @@
 #######################################
 # Unit tests
 #######################################
+
+ADD_PYARROW_TEST(adapters/pandas-test)

python/src/pyarrow/adapters/pandas-test.cc

@@ -0,0 +1,66 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "gtest/gtest.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "arrow/array.h"
+#include "arrow/builder.h"
+#include "arrow/test-util.h"
+#include "arrow/schema.h"
+#include "arrow/table.h"
+#include "arrow/type.h"
+#include "pyarrow/adapters/pandas.h"
+
+using namespace arrow;
+
+namespace pyarrow {
+
+TEST(PandasConversionTest, TestObjectBlockWriteFails) {
+  StringBuilder builder;
+  const char value[] = {'\xf1', '\0'};
+
+  for (int i = 0; i < 1000; ++i) {
+    builder.Append(value, strlen(value));
+  }
+
+  std::shared_ptr<Array> arr;
+  ASSERT_OK(builder.Finish(&arr));
+
+  auto f1 = field("f1", utf8());
+  auto f2 = field("f2", utf8());
+  auto f3 = field("f3", utf8());
+  std::vector<std::shared_ptr<Field>> fields = {f1, f2, f3};
+  std::vector<std::shared_ptr<Column>> cols = {
+    std::make_shared<Column>(f1, arr),
+    std::make_shared<Column>(f2, arr),
+    std::make_shared<Column>(f3, arr)};
+
+  auto schema = std::make_shared<Schema>(fields);
+  auto table = std::make_shared<Table>("", schema, cols);
+
+  PyObject* out;
+  Py_BEGIN_ALLOW_THREADS;
+  ASSERT_RAISES(UnknownError, ConvertTableToPandas(table, 2, &out));
+  Py_END_ALLOW_THREADS;
+}
+
+}  // namespace arrow

python/src/pyarrow/adapters/pandas.cc

@@ -634,12 +634,12 @@ inline Status ArrowSerializer<NPY_OBJECT>::ConvertData() {
     RETURN_NOT_OK(converter.Convert(out));                  \
   } break;
 
-Status PandasMaskedToArrow(arrow::MemoryPool* pool, PyObject* ao, PyObject* mo,
+Status PandasToArrow(arrow::MemoryPool* pool, PyObject* ao, PyObject* mo,
     const std::shared_ptr<Field>& field, std::shared_ptr<Array>* out) {
   PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(ao);
   PyArrayObject* mask = nullptr;
 
-  if (mo != nullptr) { mask = reinterpret_cast<PyArrayObject*>(mo); }
+  if (mo != nullptr and mo != Py_None) { mask = reinterpret_cast<PyArrayObject*>(mo); }
 
   if (PyArray_NDIM(arr) != 1) {
     return Status::Invalid("only handle 1-dimensional arrays");
@@ -671,11 +671,6 @@ Status PandasMaskedToArrow(arrow::MemoryPool* pool, PyObject* ao, PyObject* mo,
   return Status::OK();
 }
 
-Status PandasToArrow(arrow::MemoryPool* pool, PyObject* ao,
-    const std::shared_ptr<Field>& field, std::shared_ptr<Array>* out) {
-  return PandasMaskedToArrow(pool, ao, nullptr, field, out);
-}
-
 // ----------------------------------------------------------------------
 // Deserialization
 
@@ -909,7 +904,11 @@ inline Status ConvertBinaryLike(const ChunkedArray& data, PyObject** out_values)
         data_ptr = arr->GetValue(i, &length);
         *out_values = WrapBytes<ArrayType>::Wrap(data_ptr, length);
         if (*out_values == nullptr) {
-          return Status::UnknownError("String initialization failed");
+          PyErr_Clear();
+          std::stringstream ss;
+          ss << "Wrapping " << std::string(reinterpret_cast<const char*>(data_ptr), length)
+             << " failed";
+          return Status::UnknownError(ss.str());
         }
       }
       ++out_values;
@@ -921,9 +920,6 @@ inline Status ConvertBinaryLike(const ChunkedArray& data, PyObject** out_values)
 template <typename ArrowType>
 inline Status ConvertListsLike(
     const std::shared_ptr<Column>& col, PyObject** out_values) {
-  typedef arrow_traits<ArrowType::type_id> traits;
-  typedef typename ::arrow::TypeTraits<ArrowType>::ArrayType ArrayType;
-
   const ChunkedArray& data = *col->data().get();
   auto list_type = std::static_pointer_cast<ListType>(col->type());
 
@@ -1551,7 +1547,44 @@ class DatetimeBlock : public PandasBlock {
   }
 };
 
-// class CategoricalBlock : public PandasBlock {};
+// template <int ARROW_INDEX_TYPE>
+// class CategoricalBlock : public PandasBlock {
+//  public:
+//   using PandasBlock::PandasBlock;
+
+//   Status Allocate() override {
+//     constexpr int npy_type = arrow_traits<ARROW_INDEX_TYPE>::npy_type;
+
+//     if (!(npy_type == NPY_INT8 || npy_type == NPY_INT16 ||
+//             npy_type == NPY_INT32 || npy_type == NPY_INT64)) {
+//       return Status::Invalid("Category indices must be signed integers");
+//     }
+//     return AllocateNDArray(npy_type);
+//   }
+
+//   Status Write(const std::shared_ptr<Column>& col, int64_t abs_placement,
+//       int64_t rel_placement) override {
+//     using T = typename arrow_traits<ARROW_INDEX_TYPE>::T;
+
+//     T* out_values = reinterpret_cast<T*>(block_data_) + rel_placement * num_rows_;
+
+//     const ChunkedArray& data = *col->data().get();
+
+//     for (int c = 0; c < data.num_chunks(); c++) {
+//       const std::shared_ptr<Array> arr = data.chunk(c);
+//       auto prim_arr = static_cast<arrow::PrimitiveArray*>(arr.get());
+//       auto in_values = reinterpret_cast<const T*>(prim_arr->data()->data());
+
+//       // Null is -1 in CategoricalBlock
+//       for (int i = 0; i < arr->length(); ++i) {
+//         *out_values++ = prim_arr->IsNull(i) ? -1 : in_values[i];
+//       }
+//     }
+
+//     placement_data_[rel_placement] = abs_placement;
+//     return Status::OK();
+//   }
+// };
 
 Status MakeBlock(PandasBlock::type type, int64_t num_rows, int num_columns,
     std::shared_ptr<PandasBlock>* block) {
@@ -1575,7 +1608,8 @@ Status MakeBlock(PandasBlock::type type, int64_t num_rows, int num_columns,
     BLOCK_CASE(BOOL, BoolBlock);
     BLOCK_CASE(DATETIME, DatetimeBlock);
     case PandasBlock::CATEGORICAL:
-      return Status::NotImplemented("categorical");
+      // *block = std::make_shared<CategoricalBlock>(num_rows, num_columns);
+      break;
   }
 
 #undef BLOCK_CASE

python/src/pyarrow/adapters/pandas.h

@@ -63,11 +63,7 @@ arrow::Status ConvertTableToPandas(
     const std::shared_ptr<arrow::Table>& table, int nthreads, PyObject** out);
 
 PYARROW_EXPORT
-arrow::Status PandasMaskedToArrow(arrow::MemoryPool* pool, PyObject* ao, PyObject* mo,
-    const std::shared_ptr<arrow::Field>& field, std::shared_ptr<arrow::Array>* out);
-
-PYARROW_EXPORT
-arrow::Status PandasToArrow(arrow::MemoryPool* pool, PyObject* ao,
+arrow::Status PandasToArrow(arrow::MemoryPool* pool, PyObject* ao, PyObject* mo,
     const std::shared_ptr<arrow::Field>& field, std::shared_ptr<arrow::Array>* out);
 
 }  // namespace pyarrow

python/src/pyarrow/common.cc

@@ -93,7 +93,7 @@ PyBytesBuffer::PyBytesBuffer(PyObject* obj)
 }
 
 PyBytesBuffer::~PyBytesBuffer() {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   Py_DECREF(obj_);
 }
 

python/src/pyarrow/common.h

@@ -30,6 +30,17 @@ class MemoryPool;
 
 namespace pyarrow {
 
+class PyAcquireGIL {
+ public:
+  PyAcquireGIL() { state_ = PyGILState_Ensure(); }
+
+  ~PyAcquireGIL() { PyGILState_Release(state_); }
+
+ private:
+  PyGILState_STATE state_;
+  DISALLOW_COPY_AND_ASSIGN(PyAcquireGIL);
+};
+
 #define PYARROW_IS_PY2 PY_MAJOR_VERSION <= 2
 
 class OwnedRef {
@@ -38,7 +49,10 @@ class OwnedRef {
 
   OwnedRef(PyObject* obj) : obj_(obj) {}
 
-  ~OwnedRef() { Py_XDECREF(obj_); }
+  ~OwnedRef() {
+    PyAcquireGIL lock;
+    Py_XDECREF(obj_);
+  }
 
   void reset(PyObject* obj) {
     if (obj_ != nullptr) { Py_XDECREF(obj_); }
@@ -69,17 +83,6 @@ struct PyObjectStringify {
   }
 };
 
-class PyGILGuard {
- public:
-  PyGILGuard() { state_ = PyGILState_Ensure(); }
-
-  ~PyGILGuard() { PyGILState_Release(state_); }
-
- private:
-  PyGILState_STATE state_;
-  DISALLOW_COPY_AND_ASSIGN(PyGILGuard);
-};
-
 // TODO(wesm): We can just let errors pass through. To be explored later
 #define RETURN_IF_PYERROR()                         \
   if (PyErr_Occurred()) {                           \
@@ -88,8 +91,9 @@ class PyGILGuard {
     PyObjectStringify stringified(exc_value);       \
     std::string message(stringified.bytes);         \
     Py_DECREF(exc_type);                            \
-    Py_DECREF(exc_value);                           \
-    Py_DECREF(traceback);                           \
+    Py_XDECREF(exc_value);                          \
+    Py_XDECREF(traceback);                          \
+    PyErr_Clear();                                  \
     return Status::UnknownError(message);           \
   }
 
@@ -122,17 +126,6 @@ class PYARROW_EXPORT PyBytesBuffer : public arrow::Buffer {
   PyObject* obj_;
 };
 
-class PyAcquireGIL {
- public:
-  PyAcquireGIL() { state_ = PyGILState_Ensure(); }
-
-  ~PyAcquireGIL() { PyGILState_Release(state_); }
-
- private:
-  PyGILState_STATE state_;
-  DISALLOW_COPY_AND_ASSIGN(PyAcquireGIL);
-};
-
 }  // namespace pyarrow
 
 #endif  // PYARROW_COMMON_H

python/src/pyarrow/io.cc

@@ -114,22 +114,22 @@ PyReadableFile::PyReadableFile(PyObject* file) {
 PyReadableFile::~PyReadableFile() {}
 
 Status PyReadableFile::Close() {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   return file_->Close();
 }
 
 Status PyReadableFile::Seek(int64_t position) {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   return file_->Seek(position, 0);
 }
 
 Status PyReadableFile::Tell(int64_t* position) {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   return file_->Tell(position);
 }
 
 Status PyReadableFile::Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out) {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   PyObject* bytes_obj;
   ARROW_RETURN_NOT_OK(file_->Read(nbytes, &bytes_obj));
 
@@ -141,7 +141,7 @@ Status PyReadableFile::Read(int64_t nbytes, int64_t* bytes_read, uint8_t* out) {
 }
 
 Status PyReadableFile::Read(int64_t nbytes, std::shared_ptr<arrow::Buffer>* out) {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
 
   PyObject* bytes_obj;
   ARROW_RETURN_NOT_OK(file_->Read(nbytes, &bytes_obj));
@@ -153,7 +153,7 @@ Status PyReadableFile::Read(int64_t nbytes, std::shared_ptr<arrow::Buffer>* out)
 }
 
 Status PyReadableFile::GetSize(int64_t* size) {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
 
   int64_t current_position;
   ;
@@ -185,17 +185,17 @@ PyOutputStream::PyOutputStream(PyObject* file) {
 PyOutputStream::~PyOutputStream() {}
 
 Status PyOutputStream::Close() {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   return file_->Close();
 }
 
 Status PyOutputStream::Tell(int64_t* position) {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   return file_->Tell(position);
 }
 
 Status PyOutputStream::Write(const uint8_t* data, int64_t nbytes) {
-  PyGILGuard lock;
+  PyAcquireGIL lock;
   return file_->Write(data, nbytes);
 }
 

python/src/pyarrow/util/CMakeLists.txt

@@ -20,7 +20,7 @@
 #######################################
 
 if (PYARROW_BUILD_TESTS)
-  add_library(pyarrow_test_main
+  add_library(pyarrow_test_main STATIC
 	test_main.cc)
 
   if (APPLE)

python/src/pyarrow/util/test_main.cc

@@ -15,12 +15,22 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include <Python.h>
+
 #include <gtest/gtest.h>
 
+#include "pyarrow/do_import_numpy.h"
+#include "pyarrow/numpy_interop.h"
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
 
+  Py_Initialize();
+  pyarrow::import_numpy();
+
   int ret = RUN_ALL_TESTS();
 
+  Py_Finalize();
+
   return ret;
 }