ENH: copy ndarrays if they are encountered in setArrayFromSequence

By avoiding the elementwise copy array([a,b,c]) is now as
fast as vstack([a,b,c])

2 files changed, 96 insertions, 20 deletions

diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index 70f63269e..ef0dca414 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -24,6 +24,7 @@
 #include "_datetime.h"
 #include "datetime_strings.h"
 #include "array_assign.h"
+#include "mapping.h" /* for array_item_asarray */
 
 /*
  * Reading from a file or a string.
@@ -396,14 +397,26 @@ copy_and_swap(void *dst, void *src, int itemsize, npy_intp numitems,
     }
 }
 
-/* adapted from Numarray */
+/*
+ * adapted from Numarray,
+ * a: destination array
+ * s: source object, array or sequence
+ * dim: current recursion dimension, must be 0 on first call
+ * dst: must be NULL on first call
+ * it is a view on the destination array viewing the place where to put the
+ * data of the current recursion
+ */
 static int
 setArrayFromSequence(PyArrayObject *a, PyObject *s,
-                        int dim, npy_intp offset)
+                        int dim, PyArrayObject * dst)
 {
     Py_ssize_t i, slen;
     int res = 0;
 
+    /* first recursion, view equal destination */
+    if (dst == NULL)
+        dst = a;
+
     /*
      * This code is to ensure that the sequence access below will
      * return a lower-dimensional sequence.
@@ -412,17 +425,26 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
     /* INCREF on entry DECREF on exit */
     Py_INCREF(s);
 
-    if (PyArray_Check(s) && !(PyArray_CheckExact(s))) {
-      /*
-       * FIXME:  This could probably copy the entire subarray at once here using
-       * a faster algorithm.  Right now, just make sure a base-class array is
-       * used so that the dimensionality reduction assumption is correct.
-       */
-        /* This will DECREF(s) if replaced */
-        s = PyArray_EnsureArray(s);
-        if (s == NULL) {
+    if (PyArray_Check(s)) {
+        if (!(PyArray_CheckExact(s))) {
+            /*
+             * make sure a base-class array is used so that the dimensionality
+             * reduction assumption is correct.
+             */
+            /* This will DECREF(s) if replaced */
+            s = PyArray_EnsureArray(s);
+            if (s == NULL) {
+                goto fail;
+            }
+        }
+
+        /* dst points to correct array subsection */
+        if (PyArray_CopyInto(dst, (PyArrayObject *)s) < 0) {
             goto fail;
         }
+
+        Py_DECREF(s);
+        return 0;
     }
 
     if (dim > PyArray_NDIM(a)) {
@@ -458,17 +480,23 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
 
         for (i = 0; i < alen; i++) {
             if ((PyArray_NDIM(a) - dim) > 1) {
-                res = setArrayFromSequence(a, o, dim+1, offset);
+                PyArrayObject * tmp =
+                    (PyArrayObject *)array_item_asarray(dst, i);
+                if (tmp == NULL) {
+                    goto fail;
+                }
+
+                res = setArrayFromSequence(a, o, dim+1, tmp);
+                Py_DECREF(tmp);
             }
             else {
-                res = PyArray_DESCR(a)->f->setitem(o,
-                                        (PyArray_BYTES(a) + offset), a);
+                char * b = (PyArray_BYTES(dst) + i * PyArray_STRIDES(dst)[0]);
+                res = PyArray_DESCR(dst)->f->setitem(o, b, dst);
             }
             if (res < 0) {
                 Py_DECREF(o);
                 goto fail;
             }
-            offset += PyArray_STRIDES(a)[dim];
         }
         Py_DECREF(o);
     }
@@ -480,17 +508,23 @@ setArrayFromSequence(PyArrayObject *a, PyObject *s,
                 goto fail;
             }
             if ((PyArray_NDIM(a) - dim) > 1) {
-                res = setArrayFromSequence(a, o, dim+1, offset);
+                PyArrayObject * tmp =
+                    (PyArrayObject *)array_item_asarray(dst, i);
+                if (tmp == NULL) {
+                    goto fail;
+                }
+
+                res = setArrayFromSequence(a, o, dim+1, tmp);
+                Py_DECREF(tmp);
             }
             else {
-                res = PyArray_DESCR(a)->f->setitem(o,
-                                 (PyArray_BYTES(a) + offset), a);
+                char * b = (PyArray_BYTES(dst) + i * PyArray_STRIDES(dst)[0]);
+                res = PyArray_DESCR(dst)->f->setitem(o, b, dst);
             }
             Py_DECREF(o);
             if (res < 0) {
                 goto fail;
             }
-            offset += PyArray_STRIDES(a)[dim];
         }
     }
 
@@ -515,7 +549,7 @@ PyArray_AssignFromSequence(PyArrayObject *self, PyObject *v)
                         "assignment to 0-d array");
         return -1;
     }
-    return setArrayFromSequence(self, v, 0, 0);
+    return setArrayFromSequence(self, v, 0, NULL);
 }
 
 /*
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 0e7f6ec61..19d2810a2 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -176,6 +176,48 @@ class TestAttributes(TestCase):
         assert_array_equal(x['a'], [3.5, 3.5])
         assert_array_equal(x['b'], [-2, -2])
 
+
+class TestArrayConstruction(TestCase):
+    def test_array(self):
+        d = np.ones(6)
+        r = np.array([d, d])
+        assert_equal(r, np.ones((2, 6)))
+
+        d = np.ones(6)
+        tgt = np.ones((2, 6))
+        r = np.array([d, d])
+        assert_equal(r, tgt)
+        tgt[1] = 2
+        r = np.array([d, d + 1])
+        assert_equal(r, tgt)
+
+        d = np.ones(6)
+        r = np.array([[d, d]])
+        assert_equal(r, np.ones((1, 2, 6)))
+
+        d = np.ones(6)
+        r = np.array([[d, d], [d, d]])
+        assert_equal(r, np.ones((2, 2, 6)))
+
+        d = np.ones((6, 6))
+        r = np.array([d, d])
+        assert_equal(r, np.ones((2, 6, 6)))
+
+        d = np.ones((6, ))
+        r = np.array([[d, d + 1], d + 2])
+        assert_equal(len(r), 2)
+        assert_equal(r[0], [d, d + 1])
+        assert_equal(r[1], d + 2)
+
+        tgt = np.ones((2, 3), dtype=np.bool)
+        tgt[0, 2] = False
+        tgt[1, 0:2] = False
+        r = np.array([[True, True, False], [False, False, True]])
+        assert_equal(r, tgt)
+        r = np.array([[True, False], [True, False], [False, True]])
+        assert_equal(r, tgt.T)
+
+
 class TestAssignment(TestCase):
     def test_assignment_broadcasting(self):
         a = np.arange(6).reshape(2, 3)