zanmato1984 commented on code in PR #51611:
URL: https://github.com/apache/arrow/pull/51611#discussion_r4339748905
##########
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:
Could we validate the declared list value type against the query type before
this early return?
A null `list<int32>` scalar with string queries `["a", "b"]` currently
returns `[null, null]`, whereas the equivalent `list<int32>` array `[null,
null]` rejects the incompatible types. The same inconsistency occurs with
typed-null string queries.
Type acceptance should not depend on whether the list scalar is valid.
##########
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:
Non-blocking design question: could the execution-time comparison type
checks, such as the timestamp timezone compatibility check, be extracted into
an internal helper shared by the comparison kernels and `list_contains`, while
retaining `DispatchBest` for type selection?
Currently, `CheckComparable` constructs two length-one null arrays and
executes `equal` solely to obtain its validation status. A shared validator
would avoid that dummy execution and make the validation contract explicit,
without introducing a public API or changing `equal`’s supported types.
This could also be a follow-up if it would broaden the scope here.
--
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]
