diff --git a/cpp/src/arrow/compute/kernels/scalar_compare.cc b/cpp/src/arrow/compute/kernels/scalar_compare.cc index 65987c47ef2b..404b45d2b6aa 100644 --- a/cpp/src/arrow/compute/kernels/scalar_compare.cc +++ b/cpp/src/arrow/compute/kernels/scalar_compare.cc @@ -24,6 +24,7 @@ #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/type.h" +#include "arrow/util/binary_view_util.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" #include "arrow/util/logging_internal.h" @@ -380,6 +381,225 @@ struct VarArgsCompareFunction : ScalarFunction { } }; +// Shared equality comparison between a BinaryView array and any BaseBinaryScalar. +// Pins scalar length and 4-byte prefix in registers for early rejection before +// dereferencing view variadic buffers. +template +Status ViewArrayCompareBinaryScalarEqual(KernelContext* ctx, const ArraySpan& arg0, + const Scalar& arg1, ExecResult* out) { + constexpr bool kIsEqual = std::is_same_v; + if (!arg1.is_valid) { + return applicator::OutputAdapter::Write(ctx, out->array_span_mutable(), + []() -> bool { return false; }); + } + const auto* views0 = arg0.GetValues(1); + const auto* buffers0 = arg0.GetVariadicBuffers().data(); + const auto& scalar = checked_cast(arg1); + std::string_view s = scalar.view(); + BinaryViewType::c_type s_val = util::ToBinaryView(s, 0, 0); + const uint8_t* s_data = reinterpret_cast(s.data()); + int32_t s_len = s_val.size(); + + int64_t s_sp; + memcpy(&s_sp, &s_val, sizeof(int64_t)); + bool s_inline = s_val.is_inline(); + int64_t s_inlined = 0; + if (s_inline) { + memcpy(&s_inlined, s_val.inline_data() + BinaryViewType::kPrefixSize, + sizeof(int64_t)); + } + + int64_t i = 0; + return applicator::OutputAdapter::Write( + ctx, out->array_span_mutable(), [&]() -> bool { + const auto& v = views0[i++]; + int64_t v_sp; + memcpy(&v_sp, &v, sizeof(int64_t)); + if (v_sp != s_sp) { + return !kIsEqual; + } + if (s_inline) { + int64_t v_inlined; + memcpy(&v_inlined, v.inline_data() + BinaryViewType::kPrefixSize, + sizeof(int64_t)); + bool eq = (v_inlined == s_inlined); + return kIsEqual ? eq : !eq; + } + const uint8_t* v_data = buffers0[v.ref.buffer_index]->data() + v.ref.offset; + bool eq = (memcmp(v_data + BinaryViewType::kPrefixSize, + s_data + BinaryViewType::kPrefixSize, + s_len - BinaryViewType::kPrefixSize) == 0); + return kIsEqual ? eq : !eq; + }); +} + +template +struct ViewCompareViewEqual { + static constexpr bool kIsEqual = std::is_same_v; + + static Status ArrayArray(KernelContext* ctx, const ArraySpan& arg0, + const ArraySpan& arg1, ExecResult* out) { + const auto* views0 = arg0.GetValues(1); + const auto* views1 = arg1.GetValues(1); + const auto* buffers0 = arg0.GetVariadicBuffers().data(); + const auto* buffers1 = arg1.GetVariadicBuffers().data(); + int64_t i = 0; + + return applicator::OutputAdapter::Write( + ctx, out->array_span_mutable(), [&]() -> bool { + bool eq = util::EqualBinaryView(views0[i], views1[i], buffers0, buffers1); + ++i; + return kIsEqual ? eq : !eq; + }); + } + + static Status ArrayScalar(KernelContext* ctx, const ArraySpan& arg0, const Scalar& arg1, + ExecResult* out) { + return ViewArrayCompareBinaryScalarEqual(ctx, arg0, arg1, out); + } + + static Status ScalarArray(KernelContext* ctx, const Scalar& arg0, const ArraySpan& arg1, + ExecResult* out) { + return ArrayScalar(ctx, arg1, arg0, out); + } + + static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { + if (batch[0].is_array()) { + if (batch[1].is_array()) { + return ArrayArray(ctx, batch[0].array, batch[1].array, out); + } else { + return ArrayScalar(ctx, batch[0].array, *batch[1].scalar, out); + } + } else { + if (batch[1].is_array()) { + return ScalarArray(ctx, *batch[0].scalar, batch[1].array, out); + } else { + return applicator::ScalarBinary::Exec(ctx, batch, out); + } + } + } +}; + +template +struct ViewCompareStringEqual { + using offset_type = typename StringType::offset_type; + static constexpr bool kIsEqual = std::is_same_v; + + static Status ArrayArray(KernelContext* ctx, const ArraySpan& arg0, + const ArraySpan& arg1, ExecResult* out) { + const auto* views = arg0.GetValues(1); + const auto* v_buffers = arg0.GetVariadicBuffers().data(); + const auto* offsets = + reinterpret_cast(arg1.buffers[1].data) + arg1.offset; + const uint8_t* str_base = arg1.buffers[2].data; + int64_t i = 0; + + return applicator::OutputAdapter::Write( + ctx, out->array_span_mutable(), [&]() -> bool { + offset_type offset = offsets[i]; + offset_type str_len = offsets[i + 1] - offset; + const uint8_t* str_data = str_base + offset; + + bool eq = util::EqualBinaryViewString(views[i], str_data, str_len, v_buffers); + ++i; + return kIsEqual ? eq : !eq; + }); + } + + static Status ArrayScalar(KernelContext* ctx, const ArraySpan& arg0, const Scalar& arg1, + ExecResult* out) { + return ViewArrayCompareBinaryScalarEqual(ctx, arg0, arg1, out); + } + + static Status ScalarArray(KernelContext* ctx, const Scalar& arg0, const ArraySpan& arg1, + ExecResult* out) { + if (!arg0.is_valid) { + return applicator::OutputAdapter::Write( + ctx, out->array_span_mutable(), []() -> bool { return false; }); + } + const auto& scalar = checked_cast(arg0); + std::string_view s = scalar.view(); + int32_t s_len = static_cast(s.size()); + const uint8_t* s_data = reinterpret_cast(s.data()); + + uint32_t s_prefix = 0; + if (s_len >= BinaryViewType::kPrefixSize) { + memcpy(&s_prefix, s_data, sizeof(uint32_t)); + } + + const auto* offsets = + reinterpret_cast(arg1.buffers[1].data) + arg1.offset; + const uint8_t* str_base = arg1.buffers[2].data; + int64_t i = 0; + + return applicator::OutputAdapter::Write( + ctx, out->array_span_mutable(), [&]() -> bool { + offset_type offset = offsets[i]; + offset_type str_len = offsets[i + 1] - offset; + const uint8_t* str_data = str_base + offset; + ++i; + + if (str_len != s_len) return !kIsEqual; + if (s_len == 0) return kIsEqual; + + if (s_len < BinaryViewType::kPrefixSize) { + bool eq = (memcmp(s_data, str_data, s_len) == 0); + return kIsEqual ? eq : !eq; + } + + uint32_t str_prefix; + memcpy(&str_prefix, str_data, sizeof(uint32_t)); + if (str_prefix != s_prefix) return !kIsEqual; + + bool eq = (memcmp(s_data + BinaryViewType::kPrefixSize, + str_data + BinaryViewType::kPrefixSize, + s_len - BinaryViewType::kPrefixSize) == 0); + return kIsEqual ? eq : !eq; + }); + } + + static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { + if (batch[0].is_array()) { + if (batch[1].is_array()) { + return ArrayArray(ctx, batch[0].array, batch[1].array, out); + } else { + return ArrayScalar(ctx, batch[0].array, *batch[1].scalar, out); + } + } else { + if (batch[1].is_array()) { + return ScalarArray(ctx, *batch[0].scalar, batch[1].array, out); + } else { + return applicator::ScalarBinary::Exec(ctx, batch, out); + } + } + } +}; + +template +struct StringCompareViewEqual { + static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { + if (batch[0].is_array()) { + if (batch[1].is_array()) { + return ViewCompareStringEqual::ArrayArray(ctx, batch[1].array, + batch[0].array, out); + } else { + return ViewCompareStringEqual::ScalarArray(ctx, *batch[1].scalar, + batch[0].array, out); + } + } else { + if (batch[1].is_array()) { + return ViewCompareStringEqual::ArrayScalar(ctx, batch[1].array, + *batch[0].scalar, out); + } else { + return applicator::ScalarBinary::Exec(ctx, batch, out); + } + } + } +}; + template std::shared_ptr MakeCompareFunction(std::string name, FunctionDoc doc) { auto func = std::make_shared(name, Arity::Binary(), std::move(doc)); @@ -427,19 +647,32 @@ std::shared_ptr MakeCompareFunction(std::string name, FunctionDo DCHECK_OK(func->AddKernel({ty, ty}, boolean(), std::move(exec))); } for (const auto& ty : BinaryViewTypes()) { - auto exec = GenerateVarBinaryViewBase, Op>(*ty); + ArrayKernelExec exec; + if constexpr (std::is_same_v || std::is_same_v) { + exec = ViewCompareViewEqual::Exec; + } else { + exec = GenerateVarBinaryViewBase, Op>(*ty); + } DCHECK_OK(func->AddKernel({ty, ty}, boolean(), std::move(exec))); } - auto exec_base_view = - applicator::ScalarBinary::Exec; - auto exec_view_base = - applicator::ScalarBinary::Exec; - auto exec_large_view = - applicator::ScalarBinary::Exec; - auto exec_view_large = - applicator::ScalarBinary::Exec; + ArrayKernelExec exec_base_view, exec_view_base, exec_large_view, exec_view_large; + if constexpr (std::is_same_v || std::is_same_v) { + exec_base_view = StringCompareViewEqual::Exec; + exec_view_base = ViewCompareStringEqual::Exec; + exec_large_view = StringCompareViewEqual::Exec; + exec_view_large = ViewCompareStringEqual::Exec; + } else { + exec_base_view = + applicator::ScalarBinary::Exec; + exec_view_base = + applicator::ScalarBinary::Exec; + exec_large_view = + applicator::ScalarBinary::Exec; + exec_view_large = + applicator::ScalarBinary::Exec; + } DCHECK_OK(func->AddKernel({utf8(), utf8_view()}, boolean(), exec_base_view)); DCHECK_OK(func->AddKernel({utf8_view(), utf8()}, boolean(), exec_view_base)); diff --git a/cpp/src/arrow/compute/kernels/scalar_compare_test.cc b/cpp/src/arrow/compute/kernels/scalar_compare_test.cc index d389cca377ad..84087d6a708a 100644 --- a/cpp/src/arrow/compute/kernels/scalar_compare_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_compare_test.cc @@ -1216,20 +1216,27 @@ TEST(TestBinaryViewCompareKernel, ArrayArray) { {large_binary(), binary_view()}, {binary_view(), large_binary()}, }; const auto expected = std::vector>{ - {"equal", "[true, false, false, false, false, false, null]"}, - {"not_equal", "[false, true, true, true, true, true, null]"}, - {"greater", "[false, false, false, false, false, true, null]"}, - {"greater_equal", "[true, false, false, false, false, true, null]"}, - {"less", "[false, true, true, true, true, false, null]"}, - {"less_equal", "[true, true, true, true, true, false, null]"}}; + {"equal", + "[true, false, false, false, false, false, null, true, false, true, false]"}, + {"not_equal", + "[false, true, true, true, true, true, null, false, true, false, true]"}, + {"greater", + "[false, false, false, false, false, true, null, false, false, false, false]"}, + {"greater_equal", + "[true, false, false, false, false, true, null, true, false, true, false]"}, + {"less", "[false, true, true, true, true, false, null, false, true, false, true]"}, + {"less_equal", + "[true, true, true, true, true, false, null, true, true, true, true]"}}; for (const auto& [lhs_type, rhs_type] : cases) { auto lhs = ArrayFromJSON( lhs_type, R"(["", "abc", "abcdefghijkl", "abcdefghijklm", "prefix_same_A", - "samepref_size", null])"); + "samepref_size", null, "equal_long_string_12345", + "alpha_long_string_123", "a\u0000b", "a\u0000b"])"); auto rhs = ArrayFromJSON( rhs_type, R"(["", "abd", "abcdefghijklm", "abcdefghijklz", "prefix_same_B", - "samepref", null])"); + "samepref", null, "equal_long_string_12345", + "omega_long_string_123", "a\u0000b", "a\u0000c"])"); CheckScalarBinary("equal", ArrayFromJSON(lhs_type, R"([])"), ArrayFromJSON(rhs_type, R"([])"), @@ -1255,20 +1262,36 @@ TEST(TestBinaryViewCompareKernel, ArrayScalar) { }; for (const auto& [arr_type, scalar_type] : cases) { - auto arr = ArrayFromJSON(arr_type, R"(["", "abc", "abcdefghijklmnop", null])"); + auto arr = ArrayFromJSON( + arr_type, R"(["", "abc", "abcdefghijklmnop", "equal_long_string_12345", null])"); auto scalar = ScalarFromJSON(scalar_type, R"("abc")"); + auto long_scalar = ScalarFromJSON(scalar_type, R"("equal_long_string_12345")"); auto null_scalar = ScalarFromJSON(scalar_type, "null"); CheckScalarBinary("equal", arr, scalar, - ArrayFromJSON(boolean(), R"([false, true, false, null])")); + ArrayFromJSON(boolean(), R"([false, true, false, false, null])")); CheckScalarBinary("equal", scalar, arr, - ArrayFromJSON(boolean(), R"([false, true, false, null])")); + ArrayFromJSON(boolean(), R"([false, true, false, false, null])")); + CheckScalarBinary("not_equal", arr, scalar, + ArrayFromJSON(boolean(), R"([true, false, true, true, null])")); + CheckScalarBinary("not_equal", scalar, arr, + ArrayFromJSON(boolean(), R"([true, false, true, true, null])")); + CheckScalarBinary("equal", arr, long_scalar, + ArrayFromJSON(boolean(), R"([false, false, false, true, null])")); + CheckScalarBinary("equal", long_scalar, arr, + ArrayFromJSON(boolean(), R"([false, false, false, true, null])")); + CheckScalarBinary("not_equal", arr, long_scalar, + ArrayFromJSON(boolean(), R"([true, true, true, false, null])")); + CheckScalarBinary("not_equal", long_scalar, arr, + ArrayFromJSON(boolean(), R"([true, true, true, false, null])")); CheckScalarBinary("greater", arr, scalar, - ArrayFromJSON(boolean(), R"([false, false, true, null])")); + ArrayFromJSON(boolean(), R"([false, false, true, true, null])")); CheckScalarBinary("less", scalar, arr, - ArrayFromJSON(boolean(), R"([false, false, true, null])")); + ArrayFromJSON(boolean(), R"([false, false, true, true, null])")); CheckScalarBinary("equal", arr, null_scalar, - ArrayFromJSON(boolean(), R"([null, null, null, null])")); + ArrayFromJSON(boolean(), R"([null, null, null, null, null])")); + CheckScalarBinary("equal", null_scalar, arr, + ArrayFromJSON(boolean(), R"([null, null, null, null, null])")); } } @@ -1279,6 +1302,10 @@ TEST(TestBinaryViewCompareKernel, Sliced) { ArrayFromJSON(boolean(), "[false, null]")); CheckScalarBinary("equal", view_arr->Slice(1, 2), str_arr->Slice(1, 2), ArrayFromJSON(boolean(), "[false, null]")); + CheckScalarBinary("equal", str_arr->Slice(0, 2), view_arr->Slice(1, 2), + ArrayFromJSON(boolean(), "[false, false]")); + CheckScalarBinary("equal", view_arr->Slice(1, 2), str_arr->Slice(0, 2), + ArrayFromJSON(boolean(), "[false, false]")); } TEST(TestBinaryViewCompareKernel, Dictionary) { diff --git a/cpp/src/arrow/util/binary_view_util.h b/cpp/src/arrow/util/binary_view_util.h index 3eae84bd0f2a..ff6764e88cba 100644 --- a/cpp/src/arrow/util/binary_view_util.h +++ b/cpp/src/arrow/util/binary_view_util.h @@ -98,6 +98,45 @@ bool EqualBinaryView(BinaryViewType::c_type l, BinaryViewType::c_type r, l.size() - BinaryViewType::kPrefixSize) == 0; } +/// \brief Compare a BinaryView value with contiguous string data. +/// +/// Performs prefix-accelerated early rejection by checking string length and the +/// 4-byte inline prefix before dereferencing variadic buffers. +template +bool EqualBinaryViewString(BinaryViewType::c_type v, const uint8_t* str_data, + int64_t str_len, const BufferPtr* v_buffers) { + if (v.size() != str_len) return false; + if (v.size() == 0) return true; + + if (v.size() < BinaryViewType::kPrefixSize) { + return memcmp(v.inline_data(), str_data, v.size()) == 0; + } + + uint32_t v_prefix, str_prefix; + memcpy(&v_prefix, v.inline_data(), sizeof(uint32_t)); + memcpy(&str_prefix, str_data, sizeof(uint32_t)); + if (v_prefix != str_prefix) return false; + + if (v.is_inline()) { + if (v.size() == BinaryViewType::kPrefixSize) return true; + return memcmp(v.inline_data() + BinaryViewType::kPrefixSize, + str_data + BinaryViewType::kPrefixSize, + v.size() - BinaryViewType::kPrefixSize) == 0; + } + + const uint8_t* v_payload = v_buffers[v.ref.buffer_index]->data() + v.ref.offset; + return memcmp(v_payload + BinaryViewType::kPrefixSize, + str_data + BinaryViewType::kPrefixSize, + v.size() - BinaryViewType::kPrefixSize) == 0; +} + +template +bool EqualBinaryViewString(BinaryViewType::c_type v, std::string_view s, + const BufferPtr* v_buffers) { + return EqualBinaryViewString(v, reinterpret_cast(s.data()), + static_cast(s.size()), v_buffers); +} + /// \brief Compute the total size of a list of binary views including null /// views. ///