jonasdedden commented on code in PR #51611:
URL: https://github.com/apache/arrow/pull/51611#discussion_r4346196791
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +541,406 @@ const FunctionDoc list_element_doc(
"is emitted. Null values emit a null in the output."),
{"lists", "index"});
+// Raise the errors "equal" would when comparing these types
+Status CheckComparable(KernelContext* ctx, const std::shared_ptr<DataType>&
values_type,
+ const std::shared_ptr<DataType>& queries_type) {
+ ARROW_ASSIGN_OR_RAISE(auto values, MakeArrayOfNull(values_type, 1,
ctx->memory_pool()));
+ ARROW_ASSIGN_OR_RAISE(auto queries,
+ MakeArrayOfNull(queries_type, 1, ctx->memory_pool()));
+ return CallFunction("equal", {values, queries},
ctx->exec_context()).status();
+}
+
+// Null queries are matched without "equal", but must be comparable all the
same
+Status CheckNullQueries(KernelContext* ctx, const Datum& values, const Datum&
queries) {
+ if (queries.type()->id() == Type::NA ||
queries.type()->Equals(*values.type())) {
+ return Status::OK();
+ }
+ return CheckComparable(ctx, values.type(), queries.type());
Review Comment:
Fixed in cb5b082: the null-query type check now goes through the same type
resolution as other queries, including the float16 → float32 promotion. All the
reported cases are covered by tests: `list<float16>` with float32 `[null,
null]` and a float32 null scalar, a scalar `list<float16>` with float32 nulls,
and `list<float32>` with float16 nulls.
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +541,406 @@ const FunctionDoc list_element_doc(
"is emitted. Null values emit a null in the output."),
{"lists", "index"});
+// Raise the errors "equal" would when comparing these types
+Status CheckComparable(KernelContext* ctx, const std::shared_ptr<DataType>&
values_type,
+ const std::shared_ptr<DataType>& queries_type) {
+ ARROW_ASSIGN_OR_RAISE(auto values, MakeArrayOfNull(values_type, 1,
ctx->memory_pool()));
+ ARROW_ASSIGN_OR_RAISE(auto queries,
+ MakeArrayOfNull(queries_type, 1, ctx->memory_pool()));
+ return CallFunction("equal", {values, queries},
ctx->exec_context()).status();
+}
+
+// Null queries are matched without "equal", but must be comparable all the
same
+Status CheckNullQueries(KernelContext* ctx, const Datum& values, const Datum&
queries) {
+ if (queries.type()->id() == Type::NA ||
queries.type()->Equals(*values.type())) {
+ return Status::OK();
+ }
+ return CheckComparable(ctx, values.type(), queries.type());
+}
+
+// The types "equal" compares `values` and `queries` as, raising the errors it
would.
+// float16, which "equal" lacks, is compared as float32, and binary views as
large
+// binary.
+Result<std::vector<TypeHolder>> ComparableTypes(KernelContext* ctx, const
Datum& values,
+ const Datum& queries) {
+ std::vector<TypeHolder> types;
+ for (const auto& type : {values.type(), queries.type()}) {
+ types.emplace_back(type->id() == Type::HALF_FLOAT ? float32() : type);
+ }
+ RETURN_NOT_OK(CheckComparable(ctx, types[0].GetSharedPtr(),
types[1].GetSharedPtr()));
+ ARROW_ASSIGN_OR_RAISE(auto equal,
+
ctx->exec_context()->func_registry()->GetFunction("equal"));
+ RETURN_NOT_OK(equal->DispatchBest(&types));
+ for (auto& type : types) {
+ if (is_binary_view_like(type.id())) {
+ type = large_binary();
+ }
+ }
+ return types;
+}
+
+Status CastTo(KernelContext* ctx, const TypeHolder& type, Datum* datum) {
+ if (!datum->type()->Equals(*type)) {
+ ARROW_ASSIGN_OR_RAISE(*datum,
+ Cast(*datum, type, CastOptions::Safe(),
ctx->exec_context()));
+ }
+ return Status::OK();
+}
+
+// Calls `visit` with a function telling whether value j equals query k, both
of the
+// same type. Unlike "equal", NaN equals NaN.
+template <typename Visit>
+Status VisitEqual(const ArraySpan& values, const ArraySpan& queries, Visit&&
visit) {
+ auto numbers = [&](auto zero) {
+ using T = decltype(zero);
+ const T* v = values.GetValues<T>(1);
+ const T* q = queries.GetValues<T>(1);
+ // `x != x` only holds for NaN
+ return visit([v, q](int64_t j, int64_t k) {
+ return v[j] == q[k] || (q[k] != q[k] && v[j] != v[j]);
+ });
+ };
+ auto binary = [&](auto zero) {
+ using T = decltype(zero);
+ const T* vo = values.GetValues<T>(1);
+ const T* qo = queries.GetValues<T>(1);
+ const uint8_t* vd = values.buffers[2].data;
+ const uint8_t* qd = queries.buffers[2].data;
+ return visit([=](int64_t j, int64_t k) {
+ const T length = vo[j + 1] - vo[j];
+ return length == qo[k + 1] - qo[k] &&
+ std::memcmp(vd + vo[j], qd + qo[k], length) == 0;
+ });
+ };
+ switch (values.type->id()) {
+ case Type::BOOL:
+ return visit([&](int64_t j, int64_t k) {
+ return bit_util::GetBit(values.buffers[1].data, values.offset + j) ==
+ bit_util::GetBit(queries.buffers[1].data, queries.offset + k);
+ });
+ case Type::FLOAT:
+ return numbers(float{});
+ case Type::DOUBLE:
+ return numbers(double{});
+ case Type::BINARY:
+ case Type::STRING:
+ return binary(int32_t{});
+ case Type::LARGE_BINARY:
+ case Type::LARGE_STRING:
+ return binary(int64_t{});
+ default:
+ break;
+ }
+ // Integers, temporal types, decimals and fixed-size binary compare bitwise
+ switch (const int width = values.type->bit_width() / 8) {
+ case 0:
+ // Only zero-width fixed-size binary, whose values all equal each other
+ return visit([](int64_t, int64_t) { return true; });
+ case 1:
+ return numbers(uint8_t{});
+ case 2:
+ return numbers(uint16_t{});
+ case 4:
+ return numbers(uint32_t{});
+ case 8:
+ return numbers(uint64_t{});
+ default: {
+ if (width < 0) {
+ return Status::NotImplemented("list_contains for ", *values.type);
+ }
+ const uint8_t* v = values.buffers[1].data + values.offset * width;
+ const uint8_t* q = queries.buffers[1].data + queries.offset * width;
+ return visit([=](int64_t j, int64_t k) {
+ return std::memcmp(v + j * width, q + k * width, width) == 0;
+ });
+ }
+ }
+}
+
+// Returns a function giving the (child offset, length) of list i. With
`kGathered`,
+// list views' offsets are into their gathered values instead, so the function
must
+// be called for i = 0, 1, ...
+template <typename Type, bool kGathered = false>
+auto GetListRanges(const ArraySpan& list) {
+ using Range = std::pair<int64_t, int64_t>;
+ if constexpr (std::is_same_v<Type, FixedSizeListType>) {
+ const int64_t width = checked_cast<const
FixedSizeListType&>(*list.type).list_size();
+ return [&list, width](int64_t i) { return Range((list.offset + i) * width,
width); };
+ } else {
+ const auto* offsets = list.GetValues<typename Type::offset_type>(1);
+ if constexpr (is_list_view_type<Type>::value) {
+ const auto* sizes = list.GetValues<typename Type::offset_type>(2);
+ // Reading the validity bit directly keeps this cheap enough to be
inlined
+ const uint8_t* validity = list.MayHaveNulls() ? list.buffers[0].data :
nullptr;
+ auto size = [sizes, validity, offset = list.offset](int64_t i) ->
int64_t {
+ return validity == nullptr || bit_util::GetBit(validity, offset + i) ?
sizes[i]
+ :
0;
+ };
+ if constexpr (kGathered) {
+ return [size, start = int64_t{0}](int64_t i) mutable {
+ const int64_t length = size(i);
+ start += length;
+ return Range(start - length, length);
+ };
+ } else {
+ return [offsets, size](int64_t i) { return Range(offsets[i], size(i));
};
+ }
+ } else {
+ return
+ [offsets](int64_t i) { return Range(offsets[i], offsets[i + 1] -
offsets[i]); };
+ }
+ }
+}
+
+// Gather the child values referenced by list views, one view after the other
+template <typename Type>
+Result<Datum> GatherListViewValues(KernelContext* ctx, const ArraySpan& list) {
+ auto list_range = GetListRanges<Type>(list);
+ Int64Builder indices(ctx->memory_pool());
+ for (int64_t i = 0; i < list.length; ++i) {
+ const auto [start, length] = list_range(i);
+ RETURN_NOT_OK(indices.Reserve(length));
+ for (int64_t j = start; j < start + length; ++j) {
+ indices.UnsafeAppend(j);
+ }
+ }
+ ARROW_ASSIGN_OR_RAISE(auto gather_indices, indices.Finish());
+ return CallFunction("take", {list.child_data[0].ToArrayData(),
gather_indices},
+ ctx->exec_context());
+}
+
+// Search each list for its query: `matches(j, i)` tells whether child value j
matches
+// the query of list i, relative to the first child value referenced. If given,
+// `confirm` is the exact check, and `matches` a cheaper superset of it that
long lists
+// are searched with first.
+template <typename ListRanges, typename Matches, typename Confirm =
std::nullptr_t>
+void SearchLists(ListRanges list_range, int64_t num_lists, int64_t
values_start,
+ int64_t values_length, ArraySpan* out_arr, Matches&& matches,
+ Confirm&& confirm = nullptr) {
+ constexpr bool kConfirm = !std::is_same_v<std::decay_t<Confirm>,
std::nullptr_t>;
+ auto emit = [&](auto&& found_in) {
+ int64_t i = 0;
+ arrow::internal::GenerateBitsUnrolled(out_arr->buffers[1].data,
out_arr->offset,
+ out_arr->length, [&] {
+ auto [start, length] =
list_range(i);
+ start -= values_start;
+ return found_in(start, start +
length, i++);
+ });
+ };
+ if (values_length < 32 * num_lists) {
+ // Short lists are faster to check in full, without branching on the data
+ auto exact = [&](int64_t j, int64_t i) {
+ if constexpr (kConfirm) {
+ return confirm(j, i);
+ } else {
+ return matches(j, i);
+ }
+ };
+ emit([&](int64_t start, int64_t end, int64_t i) {
+ bool found = false;
+ for (int64_t j = start; j < end; ++j) found |= exact(j, i);
+ return found;
+ });
+ } else {
+ // Stop after the first chunk of 16 values holding a match
+ auto any = [](auto&& predicate, int64_t start, int64_t end, int64_t i) {
+ bool found = false;
+ for (int64_t chunk = start; chunk < end && !found; chunk += 16) {
+ for (int64_t j = chunk; j < std::min(chunk + 16, end); ++j) {
+ found |= predicate(j, i);
+ }
+ }
+ return found;
+ };
+ emit([&](int64_t start, int64_t end, int64_t i) {
+ if constexpr (kConfirm) {
+ return any(matches, start, end, i) && any(confirm, start, end, i);
+ } else {
+ return any(matches, start, end, i);
+ }
+ });
+ }
+}
+
+// Look up each query in the values of a single list
+Status ScalarListContains(KernelContext* ctx, const BaseListScalar& list,
+ const ArraySpan& query_span, ArraySpan* out_arr) {
+ bit_util::SetBitsTo(out_arr->buffers[0].data, out_arr->offset,
out_arr->length,
+ list.is_valid);
+ out_arr->null_count = list.is_valid ? 0 : out_arr->length;
+ if (!list.is_valid) {
Review Comment:
Fixed in b946a72: the types are now checked against the list type's declared
value type before the null-list shortcut, so a null `list<int32>` scalar with
string queries (null or not) raises like the equivalent array. Added tests for
both.
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +541,406 @@ const FunctionDoc list_element_doc(
"is emitted. Null values emit a null in the output."),
{"lists", "index"});
+// Raise the errors "equal" would when comparing these types
+Status CheckComparable(KernelContext* ctx, const std::shared_ptr<DataType>&
values_type,
+ const std::shared_ptr<DataType>& queries_type) {
+ ARROW_ASSIGN_OR_RAISE(auto values, MakeArrayOfNull(values_type, 1,
ctx->memory_pool()));
+ ARROW_ASSIGN_OR_RAISE(auto queries,
+ MakeArrayOfNull(queries_type, 1, ctx->memory_pool()));
+ return CallFunction("equal", {values, queries},
ctx->exec_context()).status();
Review Comment:
Done in 54532f4, since it was small: the time zone check moved into an
internal `ValidateComparison` helper in `codegen_internal`, which both the
comparison kernels (`CompareTimestamps::Exec`) and `list_contains` call.
`list_contains` now resolves types with `DispatchBest` plus
`ValidateComparison`, so the dummy `equal` execution is gone. Errors are
unchanged (all comparison tests pass), and small calls got 1.3–3.3 µs cheaper.
I left `is_in`'s own time zone check alone, as it raises `TypeError` rather
than `Invalid`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
