Take care of nulls correctly.

cpp/src/arrow/array.cc

@@ -25,6 +25,7 @@
 #include <type_traits>
 #include <utility>
 
+#include "arrow/array/concatenate.h"
 #include "arrow/array/validate.h"
 #include "arrow/buffer.h"
 #include "arrow/buffer_builder.h"
@@ -304,6 +305,55 @@ Status ListArrayFromArrays(const Array& offsets, const Array& values, MemoryPool
   return Status::OK();
 }
 
+std::shared_ptr<Array> SliceArrayWithOffsets(const Array& array, int64_t begin,
+                                             int64_t end) {
+  return array.Slice(begin, end - begin);
+}
+
+template <typename ListArrayT>
+Result<std::shared_ptr<Array>> FlattenListArray(const ListArrayT& list_array,
+                                                MemoryPool* memory_pool) {
+  const int64_t list_array_length = list_array.length();
+  std::shared_ptr<arrow::Array> value_array = list_array.values();
+
+  // Shortcut: if a ListArray does not contain nulls, then simply slice its
+  // value array with the first and the last offsets. We don't use
+  // Array::null_count() because it's potentially O(N) (thus not faster than not
+  // taking the shortcut).
+  if (list_array_length == 0 || !list_array.null_bitmap_data() ||
+      list_array.data()->null_count == 0) {
+    return SliceArrayWithOffsets(*value_array, list_array.value_offset(0),
+                                 list_array.value_offset(list_array_length));
+  }
+
+  // If a ListArray contains nulls, because there may be a non-empty sub-list
+  // behind a null and it must not be contained in the result.
+  std::vector<std::shared_ptr<Array>> non_null_fragments;
+  int64_t valid_begin = 0, valid_end;
+  while (valid_begin < list_array_length) {
+    valid_end = valid_begin;
+    while (valid_end < list_array_length &&
+           (list_array.IsValid(valid_end) || list_array.value_length(valid_end) == 0)) {
+      ++valid_end;
+    }
+    if (valid_begin < valid_end) {
+      non_null_fragments.push_back(
+          SliceArrayWithOffsets(*value_array, list_array.value_offset(valid_begin),
+                                list_array.value_offset(valid_end)));
+    }
+    valid_begin = valid_end + 1;
+  }
+
+  // Final attempt to avoid invoking Concatenate().
+  if (non_null_fragments.size() == 1) {
+    return non_null_fragments[0];
+  }
+
+  std::shared_ptr<Array> result;
+  RETURN_NOT_OK(Concatenate(non_null_fragments, memory_pool, &result));
+  return result;
+}
+
 }  // namespace
 
 ListArray::ListArray(const std::shared_ptr<ArrayData>& data) { SetData(data); }
@@ -379,6 +429,13 @@ Status LargeListArray::FromArrays(const Array& offsets, const Array& values,
   return ListArrayFromArrays<LargeListType>(offsets, values, pool, out);
 }
 
+Result<std::shared_ptr<Array>> ListArray::Flatten(MemoryPool* memory_pool) const {
+  return FlattenListArray(*this, memory_pool);
+}
+
+Result<std::shared_ptr<Array>> LargeListArray::Flatten(MemoryPool* memory_pool) const {
+  return FlattenListArray(*this, memory_pool);
+}
 // ----------------------------------------------------------------------
 // MapArray
 

cpp/src/arrow/array.h

@@ -551,17 +551,6 @@ class BaseListArray : public Array {
   /// \brief Return array object containing the list's values
   std::shared_ptr<Array> values() const { return values_; }
 
-  /// \brief Return an Array that is a concatenation of the lists in this array.
-  ///
-  /// Note that it's different from `values()` in that it takes into
-  /// consideration of this array's offsets and slices its child Array
-  /// accordingly.
-  std::shared_ptr<Array> Flatten() const {
-    const offset_type begin_offset = value_offset(0);
-    const offset_type end_offset = value_offset(length());
-    return values()->Slice(begin_offset, end_offset - begin_offset);
-  }
-
   /// Note that this buffer does not account for any slice offset
   std::shared_ptr<Buffer> value_offsets() const { return data_->buffers[1]; }
 
@@ -617,6 +606,14 @@ class ARROW_EXPORT ListArray : public BaseListArray<ListType> {
   static Status FromArrays(const Array& offsets, const Array& values, MemoryPool* pool,
                            std::shared_ptr<Array>* out);
 
+  /// \brief Return an Array that is a concatenation of the lists in this array.
+  ///
+  /// Note that it's different from `values()` in that it takes into
+  /// consideration of this array's offsets as well as null elements backed
+  /// by non-empty lists (they are skipped, thus copying may be needed).
+  Result<std::shared_ptr<Array>> Flatten(
+      MemoryPool* memory_pool = default_memory_pool()) const;
+
  protected:
   // This constructor defers SetData to a derived array class
   ListArray() = default;
@@ -651,6 +648,14 @@ class ARROW_EXPORT LargeListArray : public BaseListArray<LargeListType> {
   static Status FromArrays(const Array& offsets, const Array& values, MemoryPool* pool,
                            std::shared_ptr<Array>* out);
 
+  /// \brief Return an Array that is a concatenation of the lists in this array.
+  ///
+  /// Note that it's different from `values()` in that it takes into
+  /// consideration of this array's offsets as well as null elements backed
+  /// by non-empty lists (they are skipped, thus copying may be needed).
+  Result<std::shared_ptr<Array>> Flatten(
+      MemoryPool* memory_pool = default_memory_pool()) const;
+
  protected:
   void SetData(const std::shared_ptr<ArrayData>& data);
 };

cpp/src/arrow/array/validate.cc

@@ -407,15 +407,9 @@ struct ValidateArrayDataVisitor {
     auto prev_offset = array.value_offset(0);
     for (int64_t i = 1; i <= array.length(); ++i) {
       auto current_offset = array.value_offset(i);
-      if (array.IsNull(i - 1) && current_offset != prev_offset) {
-        return Status::Invalid("Offset invariant failure at: ", i,
-                               " inconsistent value_offsets for null slot",
-                               current_offset, "!=", prev_offset);
-      }
       if (current_offset < prev_offset) {
-        return Status::Invalid("Offset invariant failure: ", i,
-                               " inconsistent offset for non-null slot: ", current_offset,
-                               "<", prev_offset);
+        return Status::Invalid("Offset invariant failure: inconsistent offset for slot ",
+                               i, ": ", current_offset, "<", prev_offset);
       }
       prev_offset = current_offset;
     }

cpp/src/arrow/array_list_test.cc

@@ -349,32 +349,53 @@ class TestListArray : public TestBuilder {
 
   void TestFlattenZeroLength() {
     Done();
-    auto flattened = result_->Flatten();
+    ASSERT_OK_AND_ASSIGN(auto flattened, result_->Flatten());
     ASSERT_OK(flattened->ValidateFull());
     ASSERT_EQ(0, flattened->length());
   }
 
-  void TestFlatten() {
+  void TestFlattenSimple() {
     auto type = std::make_shared<T>(int32());
     auto list_array = std::dynamic_pointer_cast<ArrayType>(
         ArrayFromJSON(type, "[[1, 2], [3], [4], null, [5], [], [6]]"));
-    auto flattened = list_array->Flatten();
+    ASSERT_OK_AND_ASSIGN(auto flattened, list_array->Flatten());
     ASSERT_OK(flattened->ValidateFull());
-    EXPECT_TRUE(
-        list_array->Flatten()->Equals(ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6]")));
-    EXPECT_TRUE(
-        list_array->values()->Equals(ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6]")));
+    EXPECT_TRUE(flattened->Equals(ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6]")));
+  }
 
+  void TestFlattenSliced() {
+    auto type = std::make_shared<T>(int32());
+    auto list_array = std::dynamic_pointer_cast<ArrayType>(
+        ArrayFromJSON(type, "[[1, 2], [3], [4], null, [5], [], [6]]"));
     auto sliced_list_array =
         std::dynamic_pointer_cast<ArrayType>(list_array->Slice(3, 4));
-    auto sliced_flattened = sliced_list_array->Flatten();
-    ASSERT_OK(sliced_flattened->ValidateFull());
+    ASSERT_OK_AND_ASSIGN(auto flattened, list_array->Flatten());
+    ASSERT_OK(flattened->ValidateFull());
     // Note the difference between values() and Flatten().
-    EXPECT_TRUE(sliced_flattened->Equals(ArrayFromJSON(int32(), "[5, 6]")));
+    EXPECT_TRUE(flattened->Equals(ArrayFromJSON(int32(), "[5, 6]")));
     EXPECT_TRUE(sliced_list_array->values()->Equals(
         ArrayFromJSON(int32(), "[1, 2, 3, 4, 5, 6]")));
   }
 
+  void TestFlattenNonEmptyBackingNulls() {
+    auto type = std::make_shared<T>(int32());
+    auto array_data =
+        std::dynamic_pointer_cast<ArrayType>(
+            ArrayFromJSON(type, "[[1, 2], [3], null, [5, 6], [7, 8], [], [9]]"))
+            ->data();
+    ASSERT_EQ(2, array_data->buffers.size());
+    auto null_bitmap_buffer = array_data->buffers[0];
+    ASSERT_NE(nullptr, null_bitmap_buffer);
+    BitUtil::ClearBit(null_bitmap_buffer->mutable_data(), 1);
+    BitUtil::ClearBit(null_bitmap_buffer->mutable_data(), 3);
+    BitUtil::ClearBit(null_bitmap_buffer->mutable_data(), 4);
+    auto list_array = std::dynamic_pointer_cast<ArrayType>(MakeArray(array_data));
+    ASSERT_OK(list_array->ValidateFull());
+    ASSERT_OK_AND_ASSIGN(auto flattened, list_array->Flatten());
+    EXPECT_TRUE(flattened->Equals(ArrayFromJSON(int32(), "[1, 2, 9]")))
+        << flattened->ToString();
+  }
+
  protected:
   std::shared_ptr<DataType> value_type_;
 
@@ -406,8 +427,11 @@ TYPED_TEST(TestListArray, BuilderPreserveFieldName) {
   this->TestBuilderPreserveFieldName();
 }
 
-TYPED_TEST(TestListArray, Flatten) { this->TestFlatten(); }
+TYPED_TEST(TestListArray, FlattenSimple) { this->TestFlattenSimple(); }
 TYPED_TEST(TestListArray, FlattenZeroLength) { this->TestFlattenZeroLength(); }
+TYPED_TEST(TestListArray, TestFlattenNonEmptyBackingNulls) {
+  this->TestFlattenNonEmptyBackingNulls();
+}
 
 // ----------------------------------------------------------------------
 // Map tests

python/pyarrow/array.pxi

@@ -1320,16 +1320,36 @@ cdef class ListArray(Array):
             int32(), len(self) + 1, [None, self.buffers()[1]],
             offset=self.offset)
 
-    def flatten(self):
+    def flatten(self, MemoryPool memory_pool=None):
         """
         Unnest this ListArray by one level
 
+        The returned Array is logically a concatenation of all the sub-lists
+        in this Array.
+
+        Note that is different from ``self.values()`` in that it takes care of
+        the slicing offset as well as null elements backed by non-empty
+        sub-lists.
+
+        Parameters
+        ----------
+        memory_pool : pyarrow.MemoryPool, optional
+            If not passed, will allocate memory from the currently-set default
+            memory pool
+
         Returns
         -------
         result : Array
         """
-        cdef CListArray* arr = <CListArray*> self.ap
-        return pyarrow_wrap_array(arr.Flatten())
+        cdef:
+            shared_ptr[CArray] c_result_array
+            CMemoryPool* cpool = maybe_unbox_memory_pool(memory_pool)
+            CListArray* arr = <CListArray*> self.ap
+
+        with nogil:
+            c_result_array = GetResultValue(arr.Flatten(cpool))
+
+        return pyarrow_wrap_array(c_result_array)
 
 
 cdef class LargeListArray(Array):
@@ -1382,16 +1402,36 @@ cdef class LargeListArray(Array):
             int64(), len(self) + 1, [None, self.buffers()[1]],
             offset=self.offset)
 
-    def flatten(self):
+    def flatten(self, MemoryPool memory_pool=None):
         """
-        Unnest this LargeListArray by one level
+        Unnest this ListArray by one level.
+
+        The returned Array is logically a concatenation of all the sub-lists
+        in this Array.
+
+        Note that is different from ``self.values()`` in that it takes care of
+        the slicing offset as well as null elements backed by non-empty
+        sub-lists.
+
+        Parameters
+        ----------
+        memory_pool : pyarrow.MemoryPool, optional
+            If not passed, will allocate memory from the currently-set default
+            memory pool
 
         Returns
         -------
         result : Array
         """
-        cdef CLargeListArray* arr = <CLargeListArray*> self.ap
-        return pyarrow_wrap_array(arr.Flatten())
+        cdef:
+            shared_ptr[CArray] c_result_array
+            CMemoryPool* cpool = maybe_unbox_memory_pool(memory_pool)
+            CLargeListArray* arr = <CLargeListArray*> self.ap
+
+        with nogil:
+            c_result_array = GetResultValue(arr.Flatten(cpool))
+
+        return pyarrow_wrap_array(c_result_array)
 
 
 cdef class MapArray(Array):

python/pyarrow/includes/libarrow.pxd

@@ -475,7 +475,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         int32_t value_offset(int i)
         int32_t value_length(int i)
         shared_ptr[CArray] values()
-        shared_ptr[CArray] Flatten()
+        CResult[shared_ptr[CArray]] Flatten(CMemoryPool* memory_pool)
         shared_ptr[CDataType] value_type()
 
     cdef cppclass CLargeListArray" arrow::LargeListArray"(CArray):
@@ -496,7 +496,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         int64_t value_offset(int i)
         int64_t value_length(int i)
         shared_ptr[CArray] values()
-        shared_ptr[CArray] Flatten()
+        CResult[shared_ptr[CArray]] Flatten(CMemoryPool* memory_pool)
         shared_ptr[CDataType] value_type()
 
     cdef cppclass CMapArray" arrow::MapArray"(CArray):

python/pyarrow/tests/test_array.py

@@ -1710,68 +1710,14 @@ def test_invalid_tensor_construction():
     with pytest.raises(TypeError):
         pa.Tensor()
 
-
-def test_list_array_flatten():
-    typ2 = pa.list_(
-        pa.list_(
+@pytest.mark.parametrize(('offset_type', 'list_type_factory'),
+                         [(pa.int32(), pa.list_), (pa.int64(), pa.large_list)])
+def test_list_array_flatten(offset_type, list_type_factory):
+    typ2 = list_type_factory(
+        list_type_factory(
             pa.int64()
         )
     )
-    arr2 = pa.array([
-        None,
-        [
-            [1, None, 2],
-            None,
-            [3, 4]
-        ],
-        [],
-        [
-            [],
-            [5, 6],
-            None
-        ],
-        [
-            [7, 8]
-        ]
-    ])
-    offsets2 = pa.array([0, 0, 3, 3, 6, 7], type=pa.int32())
-    assert arr2.type.equals(typ2)
-
-    typ1 = pa.list_(pa.int64())
-    arr1 = pa.array([
-        [1, None, 2],
-        None,
-        [3, 4],
-        [],
-        [5, 6],
-        None,
-        [7, 8]
-    ])
-    offsets1 = pa.array([0, 3, 3, 5, 5, 7, 7, 9], type=pa.int32())
-    assert arr1.type.equals(typ1)
-
-    typ0 = pa.int64()
-    arr0 = pa.array([
-        1, None, 2,
-        3, 4,
-        5, 6,
-        7, 8
-    ])
-    assert arr0.type.equals(typ0)
-
-    assert arr2.flatten().equals(arr1)
-    assert arr2.offsets.equals(offsets2)
-    assert arr1.flatten().equals(arr0)
-    assert arr1.offsets.equals(offsets1)
-    assert arr2.flatten().flatten().equals(arr0)
-
-
-def test_large_list_array_flatten():
-    typ2 = pa.large_list(
-        pa.large_list(
-            pa.int16()
-        )
-    )
     arr2 = pa.array([
         None,
         [
@@ -1789,10 +1735,9 @@ def test_large_list_array_flatten():
             [7, 8]
         ]
     ], type=typ2)
-    offsets2 = pa.array([0, 0, 3, 3, 6, 7], type=pa.int64())
+    offsets2 = pa.array([0, 0, 3, 3, 6, 7], type=offset_type)
 
-    typ1 = pa.large_list(pa.int16())
-    assert typ1 == typ2.value_type
+    typ1 = list_type_factory(pa.int64())
     arr1 = pa.array([
         [1, None, 2],
         None,
@@ -1802,9 +1747,20 @@ def test_large_list_array_flatten():
         None,
         [7, 8]
     ], type=typ1)
+    offsets1 = pa.array([0, 3, 3, 5, 5, 7, 7, 9], type=offset_type)
+
+    arr0 = pa.array([
+        1, None, 2,
+        3, 4,
+        5, 6,
+        7, 8
+    ], type=pa.int64())
 
     assert arr2.flatten().equals(arr1)
     assert arr2.offsets.equals(offsets2)
+    assert arr1.flatten().equals(arr0)
+    assert arr1.offsets.equals(offsets1)
+    assert arr2.flatten().flatten().equals(arr0)
 
 
 @pytest.mark.parametrize('klass', [pa.ListArray, pa.LargeListArray])
@@ -1816,9 +1772,11 @@ def test_list_array_values_offsets_sliced(klass):
 
     # sliced -> values keeps referring to full values buffer, but offsets is
     # sliced as well so the offsets correctly point into the full values array
+    # sliced -> flatten() will return the sliced value array.
     arr2 = arr[1:]
     assert arr2.values.to_pylist() == [1, 2, 3, 4, 5, 6]
     assert arr2.offsets.to_pylist() == [3, 4, 6]
+    assert arr2.flatten().to_pylist() == [4, 5, 6]
     i = arr2.offsets[0].as_py()
     j = arr2.offsets[1].as_py()
     assert arr2[0].as_py() == arr2.values[i:j].to_pylist() == [4]