dmitry-chirkov-dremio commented on code in PR #51611:
URL: https://github.com/apache/arrow/pull/51611#discussion_r4372712791
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +540,455 @@ const FunctionDoc list_element_doc(
"is emitted. Null values emit a null in the output."),
{"lists", "index"});
+// Dictionaries are compared as their values, like by "equal"
+const std::shared_ptr<DataType>& Decoded(const std::shared_ptr<DataType>&
type) {
+ return type->id() == Type::DICTIONARY
+ ? checked_cast<const DictionaryType&>(*type).value_type()
+ : 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
their
+// large types, which "equal" compares with other binary types unlike views.
+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 : {Decoded(values_type), Decoded(queries_type)}) {
+ switch (type->id()) {
+ case Type::HALF_FLOAT:
+ types.emplace_back(float32());
+ break;
+ case Type::STRING_VIEW:
+ types.emplace_back(large_utf8());
+ break;
+ case Type::BINARY_VIEW:
+ types.emplace_back(large_binary());
+ break;
+ default:
+ types.emplace_back(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()};
+ }
+ // Raise also where there is nothing to cast, e.g. for null queries
+ for (const auto& [from, to] : {std::pair(Decoded(values_type), types[0]),
+ std::pair(Decoded(queries_type), types[1])}) {
+ 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. An error is kept rather than raised for
queries
+ // that are null-typed or of the values' type: all-null queries are matched
by
+ // validity alone, which works even for values "equal" can't compare, like
structs.
+ // The error is raised once a batch holds a non-null query.
+ 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 (after decoding dictionaries).
+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);
+ const auto& decoded_queries_type = Decoded(queries_type);
+ if (decoded_queries_type->id() != Type::NA &&
+ !decoded_queries_type->Equals(*Decoded(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 no
values for
+// null list views
+template <typename Type>
+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;
+ return [offsets, sizes, validity, offset = list.offset](int64_t i) {
+ return Range(
+ offsets[i],
+ validity == nullptr || bit_util::GetBit(validity, offset + i) ?
sizes[i] : 0);
+ };
+ } else {
+ return
+ [offsets](int64_t i) { return Range(offsets[i], offsets[i + 1] -
offsets[i]); };
+ }
+ }
+}
+
+// Cast `values`, the child values from `values_start` on, to `type`. Values
that no list
+// references, like those of null lists, may not convert, so then they are
cast as nulls.
+template <typename Type>
+Status CastListValues(KernelContext* ctx, const ArraySpan& list, int64_t
values_start,
+ const TypeHolder& type, Datum* values) {
+ if (CastTo(ctx, type, values).ok()) {
+ return Status::OK();
+ }
+ auto data = values->array()->Copy();
+ ARROW_ASSIGN_OR_RAISE(auto validity, AllocateEmptyBitmap(data->offset +
data->length,
+
ctx->memory_pool()));
+ const auto list_range = GetListRanges<Type>(list);
+ for (int64_t i = 0; i < list.length; ++i) {
+ if (list.IsValid(i)) {
+ const auto [start, length] = list_range(i);
+ bit_util::SetBitsTo(validity->mutable_data(), data->offset + start -
values_start,
+ length, true);
+ }
+ }
+ if (data->MayHaveNulls()) {
+ arrow::internal::BitmapAnd(validity->data(), data->offset,
data->buffers[0]->data(),
+ data->offset, data->length, data->offset,
+ validity->mutable_data());
+ }
+ data->buffers[0] = std::move(validity);
+ data->null_count = kUnknownNullCount;
+ *values = std::move(data);
+ return CastTo(ctx, type, values);
+}
+
+// Whether lists are searched in full or in chunks is decided once per batch,
from the
+// mean number of values per list: below kShortListLength, scanning all values
without
+// branching on the data is fastest. Longer lists are searched
kSearchChunkLength values
+// at a time, which compile to vector compares, stopping after the first chunk
holding a
+// match. Both were tuned on int64 values, with lists of 3 to 1000 values and
matches at
+// their start, middle or end (see GH-51611).
+constexpr int64_t kShortListLength = 32;
+constexpr int64_t kSearchChunkLength = 16;
+
+// 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 < kShortListLength * 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 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 +=
kSearchChunkLength) {
+ for (int64_t j = chunk; j < std::min(chunk + kSearchChunkLength, 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 Result<std::vector<TypeHolder>>& types,
+ const BaseListScalar& list, const ArraySpan&
query_span,
+ ArraySpan* out_arr) {
+ Datum queries = query_span.ToArrayData();
+ // Nulls in a dictionary of queries count too
+ const bool null_queries =
+ queries.array()->ComputeLogicalNullCount() == queries.length();
+ if (!null_queries) {
+ // Raise even for a null list, like for list arrays
+ RETURN_NOT_OK(types.status());
+ }
+ 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) {
+ return Status::OK();
+ }
+ Datum values = list.value;
+ if (null_queries) {
+ bit_util::SetBitsTo(out_arr->buffers[1].data, out_arr->offset,
out_arr->length,
+ values.array()->ComputeLogicalNullCount() > 0);
+ return Status::OK();
+ }
+ RETURN_NOT_OK(CastTo(ctx, (*types)[0], &values));
+ RETURN_NOT_OK(CastTo(ctx, (*types)[1], &queries));
+ ExecContext* exec_ctx = ctx->exec_context();
+ if (is_floating(values.type()->id())) {
+ // "is_in" tells -0.0 from 0.0 and NaNs with different payloads apart, so
turn
+ // -0.0 into 0.0 and all NaNs into one
+ ARROW_ASSIGN_OR_RAISE(auto zero, MakeScalar(values.type(), 0.0));
+ ARROW_ASSIGN_OR_RAISE(auto nan, MakeScalar(values.type(), std::nan("")));
+ for (Datum* datum : {&values, &queries}) {
+ ARROW_ASSIGN_OR_RAISE(*datum, CallFunction("add", {*datum, zero},
exec_ctx));
+ ARROW_ASSIGN_OR_RAISE(Datum is_nan, CallFunction("is_nan", {*datum},
exec_ctx));
+ ARROW_ASSIGN_OR_RAISE(*datum,
+ CallFunction("if_else", {is_nan, nan, *datum},
exec_ctx));
+ }
+ }
+ const SetLookupOptions options(values);
+ ARROW_ASSIGN_OR_RAISE(Datum found,
+ CallFunction("is_in", {queries}, &options, exec_ctx));
+ const ArrayData& found_data = *found.array();
+ arrow::internal::CopyBitmap(found_data.buffers[1]->data(), found_data.offset,
+ found_data.length, out_arr->buffers[1].data,
+ out_arr->offset);
+ return Status::OK();
+}
+
+template <typename Type>
+Status ListContains(KernelContext* ctx, const ExecSpan& batch, ExecResult*
out) {
+ ArraySpan* out_arr = out->array_span_mutable();
+ if (batch.length == 0) {
+ return Status::OK();
+ }
+ const auto& state_types = checked_cast<const
ListContainsState&>(*ctx->state()).types;
+ if (batch[0].is_scalar()) {
+ // A scalar list is only paired with an array of queries here
+ return ScalarListContains(ctx, state_types,
+ checked_cast<const
BaseListScalar&>(*batch[0].scalar),
+ batch[1].array, out_arr);
+ }
+ const ArraySpan& list = batch[0].array;
+
+ // Only null lists emit a null
+ if (list.MayHaveNulls()) {
+ arrow::internal::CopyBitmap(list.buffers[0].data, list.offset, list.length,
+ out_arr->buffers[0].data, out_arr->offset);
+ } else {
+ bit_util::SetBitsTo(out_arr->buffers[0].data, out_arr->offset,
out_arr->length, true);
+ }
+ out_arr->null_count = list.null_count;
+
+ // The child values referenced by the lists, in any order for list views
+ auto list_range = GetListRanges<Type>(list);
+ int64_t values_start = 0;
+ int64_t values_end = list.child_data[0].length;
+ if constexpr (!is_list_view_type<Type>::value) {
+ values_start = list_range(0).first;
+ values_end = list_range(list.length - 1).first + list_range(list.length -
1).second;
+ }
+ Datum values =
+ list.child_data[0].ToArrayData()->Slice(values_start, values_end -
values_start);
+ auto search = [&](auto&&... predicates) {
+ SearchLists(list_range, list.length, values_start, values.length(),
out_arr,
+ predicates...);
+ };
+
+ // A scalar query is a single query for all lists
+ const bool scalar_query = batch[1].is_scalar();
+ Datum queries;
+ if (scalar_query) {
+ ARROW_ASSIGN_OR_RAISE(queries,
+ MakeArrayFromScalar(*batch[1].scalar, 1,
ctx->memory_pool()));
+ } else {
+ queries = batch[1].array.ToArrayData();
+ }
+ const int64_t query_step = scalar_query ? 0 : 1;
+
+ if (queries.array()->ComputeLogicalNullCount() == queries.length()) {
+ // Values are also null where their dictionary is
+ RETURN_NOT_OK(CastTo(ctx, Decoded(values.type()), &values));
+ const ArraySpan v(*values.array());
+ search([&](int64_t j, int64_t) { return v.IsNull(j); });
+ return Status::OK();
+ }
+ RETURN_NOT_OK(state_types.status());
+ const auto& types = *state_types;
+ if (!values.type()->Equals(*types[0])) {
+ if constexpr (is_list_view_type<Type>::value) {
+ // Only cast the span of values that views reference
+ values_start = std::numeric_limits<int64_t>::max();
+ values_end = 0;
+ for (int64_t i = 0; i < list.length; ++i) {
+ const auto [start, length] = list_range(i);
+ if (length > 0) {
+ values_start = std::min(values_start, start);
+ values_end = std::max(values_end, start + length);
+ }
+ }
+ values_start = std::min(values_start, values_end);
+ values = list.child_data[0].ToArrayData()->Slice(values_start,
+ values_end -
values_start);
+ }
+ RETURN_NOT_OK(CastListValues<Type>(ctx, list, values_start, types[0],
&values));
+ }
+ RETURN_NOT_OK(CastTo(ctx, types[1], &queries));
Review Comment:
The null scalar-list path returns before casting queries, but this path
still casts them. Please test both forms with `list<float64>` and an int64
query containing `2**53 + 1`. The safe cast should expose the difference.
Should these two agree?
```python
def test_list_contains_null_list_forms_agree():
ty = pa.list_(pa.float64())
queries = pa.array([2**53 + 1], pa.int64())
def outcome(lists):
try:
return pc.list_contains(lists, queries).to_pylist()
except pa.ArrowException as exc:
return type(exc).__name__, str(exc)
scalar_result = outcome(pa.scalar(None, ty))
array_result = outcome(pa.array([None], ty))
assert scalar_result == array_result
```
--
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]
