Add IsNullFast and IsValidFast to ArrayData/Array/ArraySpan

apache · Apr 25, 2023 · c8d7891 · c8d7891
1 parent 84d4716
commit c8d7891
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 0 deletions.
diff --git a/cpp/src/arrow/array/array_base.h b/cpp/src/arrow/array/array_base.h
@@ -77,6 +77,29 @@ class ARROW_EXPORT Array {
     return data_->null_count != data_->length;
   }
 
+  template <typename ArrowType>
+  bool IsNullFast(int64_t i) const {
+    return !IsValidFast<ArrowType>(i);
+  }
+
+  template <typename ArrowType>
+  bool IsValidFast(int64_t i) const {
+    if constexpr (ArrowType::type_id == Type::NA) {
+      return false;
+    } else if constexpr (ArrowType::type_id == Type::SPARSE_UNION) {
+      return !internal::IsNullSparseUnion(*data_, i);
+    } else if constexpr (ArrowType::type_id == Type::DENSE_UNION) {
+      return !internal::IsNullDenseUnion(*data_, i);
+    } else if constexpr (ArrowType::type_id == Type::RUN_END_ENCODED) {
+      return !internal::IsNullRunEndEncoded(*data_, i);
+    } else {
+      if (null_bitmap_data_ != NULLPTR) {
+        return bit_util::GetBit(null_bitmap_data_, i + data_->offset);
+      }
+      return data_->null_count != data_->length;
+    }
+  }
+
   /// \brief Return a Scalar containing the value of this array at i
   Result<std::shared_ptr<Scalar>> GetScalar(int64_t i) const;
 

diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc
@@ -329,6 +329,8 @@ TEST_F(TestArray, TestIsNullIsValid) {
   for (size_t i = 0; i < null_bitmap.size(); ++i) {
     EXPECT_EQ(null_bitmap[i] != 0, !arr->IsNull(i)) << i;
     EXPECT_EQ(null_bitmap[i] != 0, arr->IsValid(i)) << i;
+    EXPECT_EQ(null_bitmap[i] != 0, !arr->IsNullFast<Int32Type>(i)) << i;
+    EXPECT_EQ(null_bitmap[i] != 0, arr->IsValidFast<Int32Type>(i)) << i;
   }
 }
 
@@ -341,6 +343,8 @@ TEST_F(TestArray, TestIsNullIsValidNoNulls) {
   for (size_t i = 0; i < size; ++i) {
     EXPECT_TRUE(arr->IsValid(i));
     EXPECT_FALSE(arr->IsNull(i));
+    EXPECT_TRUE(arr->IsValidFast<Int32Type>(i));
+    EXPECT_FALSE(arr->IsNullFast<Int32Type>(i));
   }
 }
 
@@ -428,6 +432,25 @@ TEST_F(TestArray, TestMakeArrayOfNull) {
       for (int64_t i = 0; i < length; ++i) {
         ASSERT_TRUE(array->IsNull(i));
         ASSERT_FALSE(array->IsValid(i));
+        switch (type->id()) {
+          case Type::NA:
+            ASSERT_TRUE(array->IsNullFast<NullType>(i));
+            break;
+          case Type::SPARSE_UNION:
+            ASSERT_TRUE(array->IsNullFast<SparseUnionType>(i));
+            break;
+          case Type::DENSE_UNION:
+            ASSERT_TRUE(array->IsNullFast<DenseUnionType>(i));
+            break;
+          case Type::RUN_END_ENCODED:
+            ASSERT_TRUE(array->IsNullFast<RunEndEncodedType>(i));
+            break;
+          case Type::INT32:  // a non-special type for IsNullFast
+            ASSERT_TRUE(array->IsNullFast<Int32Type>(i));
+            break;
+          default:
+            break;
+        }
       }
     }
   }
@@ -1788,6 +1811,7 @@ TEST(TestBooleanBuilder, AppendNullsAdvanceBuilder) {
   ASSERT_TRUE(barr.Value(0));
   ASSERT_FALSE(barr.Value(1));
   ASSERT_TRUE(barr.IsNull(2));
+  ASSERT_TRUE(barr.IsNullFast<BooleanType>(2));
   ASSERT_TRUE(barr.Value(3));
 }
 
@@ -1821,9 +1845,11 @@ TEST(TestBooleanBuilder, TestStdBoolVectorAppend) {
   for (int i = 0; i < length; ++i) {
     if (is_valid[i]) {
       ASSERT_FALSE(arr.IsNull(i));
+      ASSERT_FALSE(arr.IsNullFast<BooleanType>(i));
       ASSERT_EQ(values[i], arr.Value(i));
     } else {
       ASSERT_TRUE(arr.IsNull(i));
+      ASSERT_TRUE(arr.IsNullFast<BooleanType>(i));
     }
     ASSERT_EQ(values[i], arr_nn.Value(i));
   }

diff --git a/cpp/src/arrow/array/data.h b/cpp/src/arrow/array/data.h
@@ -197,6 +197,29 @@ struct ARROW_EXPORT ArrayData {
     return null_count.load() != length;
   }
 
+  template <typename ArrowType>
+  bool IsNullFast(int64_t i) const {
+    return !IsValidFast<ArrowType>(i);
+  }
+
+  template <typename ArrowType>
+  bool IsValidFast(int64_t i) const {
+    if constexpr (ArrowType::type_id == Type::NA) {
+      return false;
+    } else if constexpr (ArrowType::type_id == Type::SPARSE_UNION) {
+      return !internal::IsNullSparseUnion(*this, i);
+    } else if constexpr (ArrowType::type_id == Type::DENSE_UNION) {
+      return !internal::IsNullDenseUnion(*this, i);
+    } else if constexpr (ArrowType::type_id == Type::RUN_END_ENCODED) {
+      return !internal::IsNullRunEndEncoded(*this, i);
+    } else {
+      if (buffers[0] != NULLPTR) {
+        return bit_util::GetBit(buffers[0]->data(), i + offset);
+      }
+      return null_count.load() != length;
+    }
+  }
+
   // Access a buffer's data as a typed C pointer
   template <typename T>
   inline const T* GetValues(int i, int64_t absolute_offset) const {
@@ -444,6 +467,29 @@ struct ARROW_EXPORT ArraySpan {
     return this->null_count != this->length;
   }
 
+  template <typename ArrowType>
+  inline bool IsNullFast(int64_t i) const {
+    return !IsValidFast<ArrowType>(i);
+  }
+
+  template <typename ArrowType>
+  inline bool IsValidFast(int64_t i) const {
+    if constexpr (ArrowType::type_id == Type::NA) {
+      return false;
+    } else if constexpr (ArrowType::type_id == Type::SPARSE_UNION) {
+      return !IsNullSparseUnion(i);
+    } else if constexpr (ArrowType::type_id == Type::DENSE_UNION) {
+      return !IsNullDenseUnion(i);
+    } else if constexpr (ArrowType::type_id == Type::RUN_END_ENCODED) {
+      return !IsNullRunEndEncoded(i);
+    } else {
+      if (this->buffers[0].data != NULLPTR) {
+        return bit_util::GetBit(this->buffers[0].data, i + this->offset);
+      }
+      return this->null_count != this->length;
+    }
+  }
+
   std::shared_ptr<ArrayData> ToArrayData() const;
 
   std::shared_ptr<Array> ToArray() const;