Add tests and update code

AlenkaF · May 27, 2024 · e8ff13d · e8ff13d
1 parent 27f3ae2
commit e8ff13d
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 1 deletion.
diff --git a/cpp/src/arrow/array/array_list_test.cc b/cpp/src/arrow/array/array_list_test.cc
@@ -1368,6 +1368,23 @@ TEST_F(TestMapArray, FromArrays) {
   ASSERT_EQ(keys_with_null->length(), tmp_items->length());
   ASSERT_RAISES(Invalid,
                 MapArray::FromArrays(offsets1, keys_with_null, tmp_items, pool_));
+
+  // With null_bitmap
+  ASSERT_OK_AND_ASSIGN(auto map7, MapArray::FromArrays(offsets1, keys, items, pool_,
+                                                       offsets3->data()->buffers[0]));
+  ASSERT_OK(map7->Validate());
+  MapArray expected7(map_type, length, offsets1->data()->buffers[1], keys, items,
+                     offsets3->data()->buffers[0], 1);
+  AssertArraysEqual(expected7, *map7);
+
+  // Null bitmap and offset with null
+  ASSERT_RAISES(Invalid, MapArray::FromArrays(offsets3, keys, items, pool_,
+                                              offsets3->data()->buffers[0]));
+
+  // Null bitmap and offset with offset
+  ASSERT_RAISES(NotImplemented,
+                MapArray::FromArrays(offsets3->Slice(2), keys, items, pool_,
+                                     offsets3->data()->buffers[0]));
 }
 
 TEST_F(TestMapArray, FromArraysEquality) {

diff --git a/cpp/src/arrow/array/array_nested.cc b/cpp/src/arrow/array/array_nested.cc
@@ -827,15 +827,32 @@ Result<std::shared_ptr<Array>> MapArray::FromArraysInternal(
     return Status::Invalid("Map key and item arrays must be equal length");
   }
 
+  if (null_bitmap != nullptr && offsets->null_count() > 0) {
+    return Status::Invalid(
+        "Ambiguous to specify both validity map and offsets with nulls");
+  }
+
+  if (null_bitmap != nullptr && offsets->offset() != 0) {
+    return Status::NotImplemented("Null bitmap with offsets slice not supported.");
+  }
+
   if (offsets->null_count() > 0) {
     ARROW_ASSIGN_OR_RAISE(auto buffers,
-                          CleanListOffsets<MapType>(null_bitmap, *offsets, pool));
+                          CleanListOffsets<MapType>(NULLPTR, *offsets, pool));
     return std::make_shared<MapArray>(type, offsets->length() - 1, std::move(buffers),
                                       keys, items, offsets->null_count(), 0);
   }
 
   using OffsetArrayType = typename TypeTraits<OffsetArrowType>::ArrayType;
   const auto& typed_offsets = checked_cast<const OffsetArrayType&>(*offsets);
+
+  if (null_bitmap != nullptr) {
+    auto buffers = BufferVector({std::move(null_bitmap), typed_offsets.values()});
+    return std::make_shared<MapArray>(type, offsets->length() - 1, std::move(buffers),
+                                      keys, items, /*null_count=*/null_bitmap->size(),
+                                      offsets->offset());
+  }
+
   auto buffers = BufferVector({null_bitmap, typed_offsets.values()});
   return std::make_shared<MapArray>(type, offsets->length() - 1, std::move(buffers), keys,
                                     items, /*null_count=*/0, offsets->offset());

diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py
@@ -1079,6 +1079,31 @@ def test_map_from_arrays():
             pa.int64()
         ))
 
+    # pass in null bitmap
+    result = pa.MapArray.from_arrays([0, 2, 2, 6], keys, items, pa.map_(
+        keys.type,
+        items.type),
+        mask=pa.array([False, True, False], type=pa.bool_())
+    )
+    assert result.equals(expected)
+
+    # error if null bitmap and offsets with nulls passed
+    with pytest.raises(pa.ArrowInvalid, match='Ambiguous to specify both validity map and offsets with nulls'):
+        pa.MapArray.from_arrays(offsets, keys, items, pa.map_(
+        keys.type,
+        items.type),
+        mask=pa.array([False, True, False], type=pa.bool_())
+    )
+
+    # error if null bitmap passed to sliced offset
+    offsets = pa.array(offsets, pa.int32())
+    with pytest.raises(pa.ArrowNotImplementedError, match='Null bitmap with offsets slice not supported.'):
+        pa.MapArray.from_arrays(offsets.slice(2), keys, items, pa.map_(
+        keys.type,
+        items.type),
+        mask=pa.array([False, True, False], type=pa.bool_())
+    )
+
     # check invalid usage
     offsets = [0, 1, 3, 5]
     keys = np.arange(5)