Support strided ndarrays in pandas conversion from nested lists

Change-Id: I31257fcf9231f88be37c6364f0e8b450256bc983

cpp/src/arrow/python/numpy-internal.h

@@ -51,7 +51,12 @@ class Ndarray1DIndexer {
 
   int64_t size() const { return PyArray_SIZE(arr_); }
 
+  T* data() const { return data_; }
+
+  bool is_strided() const { return stride_ == 1; }
+
   T& operator[](size_type index) { return *(data_ + index * stride_); }
+  T& operator[](size_type index) const { return *(data_ + index * stride_); }
 
  private:
   PyArrayObject* arr_;

cpp/src/arrow/python/pandas_to_arrow.cc

@@ -125,22 +125,27 @@ static int64_t MaskToBitmap(PyArrayObject* mask, int64_t length, uint8_t* bitmap
   return null_count;
 }
 
-template <int TYPE>
-static int64_t ValuesToValidBytes(const void* data, int64_t length,
-                                  uint8_t* valid_bytes) {
+template <int TYPE, typename BuilderType>
+static Status AppendNdarrayToBuilder(PyArrayObject* array, BuilderType* builder) {
   typedef internal::npy_traits<TYPE> traits;
   typedef typename traits::value_type T;
 
-  int64_t null_count = 0;
-  const T* values = reinterpret_cast<const T*>(data);
-
-  // TODO(wesm): striding
-  for (int i = 0; i < length; ++i) {
-    valid_bytes[i] = !traits::isnull(values[i]);
-    if (traits::isnull(values[i])) null_count++;
+  // TODO(wesm): Vector append when not strided
+  Ndarray1DIndexer<T> values(array);
+  if (traits::supports_nulls) {
+    for (int64_t i = 0; i < values.size(); ++i) {
+      if (traits::isnull(values[i])) {
+        RETURN_NOT_OK(builder->AppendNull());
+      } else {
+        RETURN_NOT_OK(builder->Append(values[i]));
+      }
+    }
+  } else {
+    for (int64_t i = 0; i < values.size(); ++i) {
+      RETURN_NOT_OK(builder->Append(values[i]));
+    }
   }
-
-  return null_count;
+  return Status::OK();
 }
 
 Status CheckFlatNumpyArray(PyArrayObject* numpy_array, int np_type) {
@@ -577,7 +582,7 @@ Status PandasConverter::ConvertDecimals() {
   RETURN_NOT_OK(ImportModule("decimal", &decimal));
   RETURN_NOT_OK(ImportFromModule(decimal, "Decimal", &Decimal));
 
-  PyObject** objects = reinterpret_cast<PyObject**>(PyArray_DATA(arr_));
+  Ndarray1DIndexer<PyObject*> objects(arr_);
   PyObject* object = objects[0];
 
   int precision;
@@ -618,7 +623,7 @@ Status PandasConverter::ConvertTimes() {
   PyAcquireGIL lock;
   PyDateTime_IMPORT;
 
-  PyObject** objects = reinterpret_cast<PyObject**>(PyArray_DATA(arr_));
+  Ndarray1DIndexer<PyObject*> objects(arr_);
 
   // datetime.time stores microsecond resolution
   Time64Builder builder(::arrow::time64(TimeUnit::MICRO), pool_);
@@ -906,7 +911,7 @@ Status LoopPySequence(PyObject* sequence, T func) {
     Py_ssize_t size = PySequence_Size(sequence);
     if (PyArray_Check(sequence)) {
       auto array = reinterpret_cast<PyArrayObject*>(sequence);
-      PyObject** objects = reinterpret_cast<PyObject**>(PyArray_DATA(array));
+      Ndarray1DIndexer<PyObject*> objects(array);
       for (int64_t i = 0; i < size; ++i) {
         RETURN_NOT_OK(func(objects[i]));
       }
@@ -934,7 +939,6 @@ template <int ITEM_TYPE, typename ArrowType>
 inline Status PandasConverter::ConvertTypedLists(const std::shared_ptr<DataType>& type,
                                                  ListBuilder* builder, PyObject* list) {
   typedef internal::npy_traits<ITEM_TYPE> traits;
-  typedef typename traits::value_type T;
   typedef typename traits::BuilderClass BuilderT;
 
   PyAcquireGIL lock;
@@ -956,18 +960,7 @@ inline Status PandasConverter::ConvertTypedLists(const std::shared_ptr<DataType>
       // TODO(uwe): Support more complex numpy array structures
       RETURN_NOT_OK(CheckFlatNumpyArray(numpy_array, ITEM_TYPE));
 
-      int64_t size = PyArray_DIM(numpy_array, 0);
-      auto data = reinterpret_cast<const T*>(PyArray_DATA(numpy_array));
-      if (traits::supports_nulls) {
-        RETURN_NOT_OK(null_bitmap_->Resize(size, false));
-        // TODO(uwe): A bitmap would be more space-efficient but the Builder API doesn't
-        // currently support this.
-        // ValuesToBitmap<ITEM_TYPE>(data, size, null_bitmap_->mutable_data());
-        ValuesToValidBytes<ITEM_TYPE>(data, size, null_bitmap_->mutable_data());
-        return value_builder->Append(data, size, null_bitmap_->data());
-      } else {
-        return value_builder->Append(data, size);
-      }
+      return AppendNdarrayToBuilder<ITEM_TYPE, BuilderT>(numpy_array, value_builder);
     } else if (PyList_Check(object)) {
       int64_t size;
       std::shared_ptr<DataType> inferred_type;

cpp/src/arrow/table.cc

@@ -301,8 +301,8 @@ Table::Table(const std::shared_ptr<Schema>& schema,
 
   columns_.resize(columns.size());
   for (size_t i = 0; i < columns.size(); ++i) {
-    columns_[i] = std::make_shared<Column>(schema->field(static_cast<int>(i)),
-                                           columns[i]);
+    columns_[i] =
+        std::make_shared<Column>(schema->field(static_cast<int>(i)), columns[i]);
   }
 }
 

python/pyarrow/tests/test_convert_pandas.py

@@ -534,6 +534,15 @@ class TestPandasConversion(unittest.TestCase):
             field = schema.field_by_name(column)
             self._check_array_roundtrip(df[column], type=field.type)
 
+    def test_column_of_lists_strided(self):
+        df, schema = dataframe_with_lists()
+        df = pd.concat([df] * 6, ignore_index=True)
+
+        arr = df['int64'].values[::3]
+        assert arr.strides[0] != 8
+
+        self._check_array_roundtrip(arr)
+
     def test_nested_lists_all_none(self):
         data = np.array([[None, None], None], dtype=object)