jonasdedden commented on code in PR #51611:
URL: https://github.com/apache/arrow/pull/51611#discussion_r4366743484
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +541,430 @@ const FunctionDoc list_element_doc(
"is emitted. Null values emit a null in the output."),
{"lists", "index"});
+// 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 std::shared_ptr<DataType>& values_type,
+ const std::shared_ptr<DataType>& queries_type) {
+ std::vector<TypeHolder> types;
+ for (const auto& type : {values_type, queries_type}) {
+ types.emplace_back(type->id() == Type::HALF_FLOAT ? float32() : type);
+ }
+ ARROW_ASSIGN_OR_RAISE(auto equal,
+
ctx->exec_context()->func_registry()->GetFunction("equal"));
+ RETURN_NOT_OK(equal->DispatchBest(&types));
+ RETURN_NOT_OK(ValidateComparison(*types[0], *types[1]));
+ if (types[0].id() == Type::FIXED_SIZE_BINARY && *types[0] != *types[1]) {
+ // Fixed-size binary values of different widths differ, as binary values do
+ types = {binary(), binary()};
+ }
+ for (const auto& [from, to] :
+ {std::pair(values_type, &types[0]), std::pair(queries_type,
&types[1])}) {
+ if (is_binary_view_like(to->id())) {
+ *to = large_binary();
+ }
+ // Raise also where there is nothing to cast, e.g. for null queries
+ if (!from->Equals(**to) && !CanCast(*from, **to)) {
+ return Status::NotImplemented("Unsupported cast from ", *from, " to ",
**to);
+ }
+ }
+ return types;
+}
+
+struct ListContainsState : public KernelState {
+ explicit ListContainsState(Result<std::vector<TypeHolder>> types)
+ : types(std::move(types)) {}
+
+ // The result of ComparableTypes, raised for non-null queries only
+ Result<std::vector<TypeHolder>> types;
+};
+
+// Compare types once, raising even if there is nothing to search, like for
null lists.
+// Null queries are matched without "equal", but must be comparable too unless
they
+// are null-typed or of the values' type.
+Result<std::unique_ptr<KernelState>> InitListContains(KernelContext* ctx,
+ const KernelInitArgs&
args) {
+ const auto& values_type =
+ checked_cast<const BaseListType&>(*args.inputs[0]).value_type();
+ const auto queries_type = args.inputs[1].GetSharedPtr();
+ auto types = ComparableTypes(ctx, values_type, queries_type);
+ if (queries_type->id() != Type::NA && !queries_type->Equals(*values_type)) {
+ RETURN_NOT_OK(types.status());
+ }
+ return std::make_unique<ListContainsState>(std::move(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) {
Review Comment:
Done in 7b65617, thanks!
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +541,430 @@ const FunctionDoc list_element_doc(
"is emitted. Null values emit a null in the output."),
{"lists", "index"});
+// 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 std::shared_ptr<DataType>& values_type,
+ const std::shared_ptr<DataType>& queries_type) {
+ std::vector<TypeHolder> types;
+ for (const auto& type : {values_type, queries_type}) {
+ types.emplace_back(type->id() == Type::HALF_FLOAT ? float32() : type);
+ }
+ ARROW_ASSIGN_OR_RAISE(auto equal,
+
ctx->exec_context()->func_registry()->GetFunction("equal"));
+ RETURN_NOT_OK(equal->DispatchBest(&types));
+ RETURN_NOT_OK(ValidateComparison(*types[0], *types[1]));
+ if (types[0].id() == Type::FIXED_SIZE_BINARY && *types[0] != *types[1]) {
+ // Fixed-size binary values of different widths differ, as binary values do
+ types = {binary(), binary()};
+ }
+ for (const auto& [from, to] :
+ {std::pair(values_type, &types[0]), std::pair(queries_type,
&types[1])}) {
+ if (is_binary_view_like(to->id())) {
+ *to = large_binary();
+ }
+ // Raise also where there is nothing to cast, e.g. for null queries
+ if (!from->Equals(**to) && !CanCast(*from, **to)) {
+ return Status::NotImplemented("Unsupported cast from ", *from, " to ",
**to);
+ }
+ }
+ return types;
+}
+
+struct ListContainsState : public KernelState {
+ explicit ListContainsState(Result<std::vector<TypeHolder>> types)
+ : types(std::move(types)) {}
+
+ // The result of ComparableTypes, raised for non-null queries only
Review Comment:
Done in 7b65617, thanks!
##########
docs/source/cpp/compute.rst:
##########
@@ -1982,6 +1984,12 @@ Structural transforms
index *n* and the type code at index *n* is 2.
* The indices ``2`` and ``7`` are invalid.
+* \(7) Output is true for each list containing a value equal to the second
+ argument. If the second argument is an array, each list is searched for the
Review Comment:
Done in 7b65617, thanks!
--
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]
