This is an automated email from the ASF dual-hosted git repository.
zanmato1984 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 00c32a69cb0 GH-33295: [C++][Python] Add list_contains compute function
(#51611)
00c32a69cb0 is described below
commit 00c32a69cb0dbe7a45a8c764638611b7bb021410
Author: Jonas Dedden <[email protected]>
AuthorDate: Mon Nov 2 09:01:45 2026 +0100
GH-33295: [C++][Python] Add list_contains compute function (#51611)
### Rationale for this change
There is no direct way to check whether each list in a list array contains
a value. This PR closes https://github.com/apache/arrow/issues/33295.
### What changes are included in this PR?
A binary scalar function `list_contains(lists, value)` returning a boolean
per list, for list, large list, list view, large list view and fixed-size list
inputs.
`value` can be a scalar, or an array of the same length, in which case each
list is searched for the value at the same index.
- **List array:** each list is searched for its value directly, without
materializing comparison bitmaps or repeated values. Short lists (fewer than 32
values on average) are compared in full without data-dependent branches; longer
lists stop after the first 16-value chunk holding a match. With null child
values, long lists are first searched without validity checks, and only lists
found are searched again with them.
- **Scalar list with an array of values:** the values are looked up with
`is_in` instead of broadcasting the list.
- **Comparison semantics follow `equal`:** values and list values are cast
to the types `equal` would compare them as (so implicit casts apply, e.g. an
int8 list and an int64 value), and `equal`'s type errors are raised. float16 is
compared as float32 and binary views as large binary.
Semantics:
- Null lists give null.
- Null list values don't match a non-null value.
- A null value matches lists that hold a null, as in Polars. Its type must
still be comparable with the list value type, like a non-null value.
- A NaN value (including float16) matches NaN list values, as in `is_in`
and Polars (`equal` says NaN != NaN).
- -0.0 matches 0.0, as with `equal`.
### Are these changes tested?
Yes, in `scalar_nested_test.cc` and `test_compute.py`. Cases include
sliced, chunked and out-of-order list-view inputs, arrays of values and scalar
lists with arrays of values for 12 value types, short and long lists, null
values whose slots hold the searched value, null and NaN values (including
float16), -0.0, implicit casts, and type and time zone errors (also for null
values).
### Are there any user-facing changes?
Yes, a new compute function, `pyarrow.compute.list_contains` in Python.
### Was AI used for this PR?
**PR code and description written by:**
- [ ] Human
- [x] AI
**Reviewed before submission by:**
- [X] Human
- [ ] AI
- [ ] Not reviewed
* GitHub Issue: #33295
Authored-by: Jonas Dedden <[email protected]>
Signed-off-by: Rossi Sun <[email protected]>
---
cpp/src/arrow/array/array_test.cc | 4 +
cpp/src/arrow/compute/kernels/codegen_internal.cc | 13 +
cpp/src/arrow/compute/kernels/codegen_internal.h | 5 +
cpp/src/arrow/compute/kernels/scalar_compare.cc | 8 +-
cpp/src/arrow/compute/kernels/scalar_nested.cc | 480 +++++++++++++++++++++
.../arrow/compute/kernels/scalar_nested_test.cc | 465 ++++++++++++++++++++
.../arrow/compute/kernels/scalar_validity_test.cc | 7 +
cpp/src/arrow/util/dict_util.cc | 6 +-
docs/source/cpp/compute.rst | 10 +
docs/source/python/api/compute.rst | 1 +
python/pyarrow/tests/test_compute.py | 118 +++++
11 files changed, 1108 insertions(+), 9 deletions(-)
diff --git a/cpp/src/arrow/array/array_test.cc
b/cpp/src/arrow/array/array_test.cc
index 55ac106e690..6e1bfed0d81 100644
--- a/cpp/src/arrow/array/array_test.cc
+++ b/cpp/src/arrow/array/array_test.cc
@@ -179,6 +179,10 @@ TEST_F(TestArray, TestNullCount) {
true, true);
CheckDictionaryNullCount(dict_type, "[null, true, null, false]", "[null,
1, 0, 2, 3]",
1, 3, true, true);
+
+ // null-typed dictionary, without a validity bitmap
+ CheckDictionaryNullCount(dictionary(index_type, null()), "[null]", "[0,
null, 0]", 1,
+ 3, true, true);
}
}
diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc
b/cpp/src/arrow/compute/kernels/codegen_internal.cc
index 99a9173041f..7909971f105 100644
--- a/cpp/src/arrow/compute/kernels/codegen_internal.cc
+++ b/cpp/src/arrow/compute/kernels/codegen_internal.cc
@@ -393,6 +393,19 @@ TypeHolder CommonBinary(const TypeHolder* begin, size_t
count) {
return large_binary();
}
+Status ValidateComparison(const DataType& lhs, const DataType& rhs) {
+ if (lhs.id() == Type::TIMESTAMP && rhs.id() == Type::TIMESTAMP) {
+ const auto& lhs_timestamp = checked_cast<const TimestampType&>(lhs);
+ const auto& rhs_timestamp = checked_cast<const TimestampType&>(rhs);
+ if (lhs_timestamp.timezone().empty() ^ rhs_timestamp.timezone().empty()) {
+ return Status::Invalid(
+ "Cannot compare timestamp with timezone to timestamp without
timezone, got: ",
+ lhs, " and ", rhs);
+ }
+ }
+ return Status::OK();
+}
+
bool CastableToDecimal(const DataType& type) {
return is_numeric(type.id()) || is_decimal(type.id());
}
diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.h
b/cpp/src/arrow/compute/kernels/codegen_internal.h
index 3a2bcefab0a..9da440cc65a 100644
--- a/cpp/src/arrow/compute/kernels/codegen_internal.h
+++ b/cpp/src/arrow/compute/kernels/codegen_internal.h
@@ -1459,6 +1459,11 @@ bool CommonTemporalResolution(const TypeHolder* begin,
size_t count,
ARROW_EXPORT
TypeHolder CommonBinary(const TypeHolder* begin, size_t count);
+/// Raise the errors comparing values of these types fails with beyond
dispatch,
+/// such as for timestamps where only one has a time zone.
+ARROW_EXPORT
+Status ValidateComparison(const DataType& lhs, const DataType& rhs);
+
/// How to promote decimal precision/scale in CastBinaryDecimalArgs.
enum class DecimalPromotion : uint8_t {
kAdd,
diff --git a/cpp/src/arrow/compute/kernels/scalar_compare.cc
b/cpp/src/arrow/compute/kernels/scalar_compare.cc
index 0f2c479b082..e7f12838c91 100644
--- a/cpp/src/arrow/compute/kernels/scalar_compare.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_compare.cc
@@ -300,13 +300,7 @@ struct CompareKernel {
template <typename Op>
struct CompareTimestamps {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult*
out) {
- const auto& lhs = checked_cast<const TimestampType&>(*batch[0].type());
- const auto& rhs = checked_cast<const TimestampType&>(*batch[1].type());
- if (lhs.timezone().empty() ^ rhs.timezone().empty()) {
- return Status::Invalid(
- "Cannot compare timestamp with timezone to timestamp without
timezone, got: ",
- lhs, " and ", rhs);
- }
+ RETURN_NOT_OK(ValidateComparison(*batch[0].type(), *batch[1].type()));
return CompareKernel<Int64Type>::Exec(ctx, batch, out);
}
};
diff --git a/cpp/src/arrow/compute/kernels/scalar_nested.cc
b/cpp/src/arrow/compute/kernels/scalar_nested.cc
index e9c65aff1ce..0df3dc738b0 100644
--- a/cpp/src/arrow/compute/kernels/scalar_nested.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_nested.cc
@@ -18,9 +18,11 @@
// Vector kernels involving nested types
#include <cmath>
+#include <limits>
#include "arrow/array/array_base.h"
#include "arrow/array/builder_nested.h"
#include "arrow/compute/api_scalar.h"
+#include "arrow/compute/cast.h"
#include "arrow/compute/kernels/common_internal.h"
#include "arrow/compute/registry_internal.h"
#include "arrow/result.h"
@@ -28,6 +30,7 @@
#include "arrow/util/bit_block_counter.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/bitmap_generate.h"
+#include "arrow/util/bitmap_ops.h"
#include "arrow/util/logging_internal.h"
#include "arrow/util/string.h"
#include "arrow/util/unreachable.h"
@@ -537,6 +540,478 @@ 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(
+ FunctionRegistry* registry, 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, 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->exec_context()->func_registry(), 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) {
+ // Only failures that depend on the values are retried, not unsupported casts
+ Status status = CastTo(ctx, type, values);
+ if (!status.IsInvalid()) {
+ return status;
+ }
+ 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) {
+ // Types that don't compare 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) {
+ // Unlike types, values of null lists don't matter
+ 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));
+ }
+ // Already done by ListContainsFunction, unless the kernel is called directly
+ RETURN_NOT_OK(CastTo(ctx, types[1], &queries));
+ const ArraySpan v(*values.array());
+ const ArraySpan q(*queries.array());
+ return VisitEqual(v, q, [&](auto&& equal) {
+ // Null checks are only paid for where there are nulls
+ if (q.MayHaveNulls()) {
+ search([&](int64_t j, int64_t i) {
+ const int64_t k = i * query_step;
+ return q.IsNull(k) ? v.IsNull(j) : v.IsValid(j) && equal(j, k);
+ });
+ } else if (v.MayHaveNulls()) {
+ // Null values may hold anything, so lists found are searched again for
them
+ const uint8_t* validity = v.buffers[0].data;
+ search([&](int64_t j, int64_t i) { return equal(j, i * query_step); },
+ [&, validity, offset = v.offset](int64_t j, int64_t i) {
+ // `&` rather than `&&`, so short lists don't branch on the data
+ return (static_cast<int>(bit_util::GetBit(validity, offset +
j)) &
+ static_cast<int>(equal(j, i * query_step))) != 0;
+ });
+ } else {
+ search([&](int64_t j, int64_t i) { return equal(j, i * query_step); });
+ }
+ return Status::OK();
+ });
+}
+
+// Queries are cast to the type they compare as before the kernel runs, like
for "equal"
+class ListContainsFunction : public ScalarFunction {
+ public:
+ using ScalarFunction::ScalarFunction;
+
+ Result<const Kernel*> DispatchBest(std::vector<TypeHolder>* types) const
override {
+ RETURN_NOT_OK(CheckArity(types->size()));
+ if (const auto* list_type = dynamic_cast<const
BaseListType*>((*types)[0].type)) {
+ // Types that don't compare are left to the kernel, as null queries need
no cast
+ auto comparable = ComparableTypes(GetFunctionRegistry(),
list_type->value_type(),
+ (*types)[1].GetSharedPtr());
+ if (comparable.ok()) {
+ (*types)[1] = (*comparable)[1];
+ }
+ }
+ return DispatchExact(*types);
+ }
+};
+
+void AddListContainsKernels(ScalarFunction* func) {
+ auto add_kernel = [&](Type::type list_type_id, ArrayKernelExec exec) {
+ ScalarKernel kernel({InputType(list_type_id), InputType::Any()},
boolean(), exec,
+ InitListContains);
+ // A null value is searched for rather than propagated
+ kernel.null_handling = NullHandling::COMPUTED_PREALLOCATE;
+ DCHECK_OK(func->AddKernel(std::move(kernel)));
+ };
+ add_kernel(Type::LIST, ListContains<ListType>);
+ add_kernel(Type::LARGE_LIST, ListContains<LargeListType>);
+ add_kernel(Type::LIST_VIEW, ListContains<ListViewType>);
+ add_kernel(Type::LARGE_LIST_VIEW, ListContains<LargeListViewType>);
+ add_kernel(Type::FIXED_SIZE_LIST, ListContains<FixedSizeListType>);
+}
+
+const FunctionDoc list_contains_doc(
+ "Check whether lists contain a given value",
+ ("`lists` must have a list-like type and `value` must be comparable\n"
+ "with the list value type.\n"
+ "For each list in `lists`, true is emitted if any of its values is
equal\n"
+ "to the corresponding `value`, false otherwise. A null `value` matches\n"
+ "null list values, and a NaN `value` matches NaN list values; otherwise\n"
+ "null list values never match. Null lists emit a null in the output."),
+ {"lists", "value"});
+
struct StructFieldFunctor {
static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult*
out) {
const auto& options = OptionsWrapper<StructFieldOptions>::Get(ctx);
@@ -960,6 +1435,11 @@ void RegisterScalarNested(FunctionRegistry* registry) {
AddListElementKernels(list_element.get());
DCHECK_OK(registry->AddFunction(std::move(list_element)));
+ auto list_contains = std::make_shared<ListContainsFunction>(
+ "list_contains", Arity::Binary(), list_contains_doc);
+ AddListContainsKernels(list_contains.get());
+ DCHECK_OK(registry->AddFunction(std::move(list_contains)));
+
auto list_slice =
std::make_shared<ScalarFunction>("list_slice", Arity::Unary(),
list_slice_doc);
AddListSliceKernels(list_slice.get());
diff --git a/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
b/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
index b5a68d12cb0..7078291f6c5 100644
--- a/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_nested_test.cc
@@ -128,6 +128,471 @@ TEST(TestScalarNested, ListElementInvalid) {
Raises(StatusCode::Invalid));
}
+void CheckListContains(Datum lists, Datum value, const std::string& expected) {
+ CheckScalar("list_contains", {std::move(lists), std::move(value)},
+ ArrayFromJSON(boolean(), expected));
+}
+
+TEST(TestScalarNested, ListContains) {
+ auto sample = "[[7, 5, 81], [6, null, 4, 7, 8], [], [5], null, [null]]";
+ for (auto ty : NumericTypes()) {
+ for (auto list_type :
+ {list(ty), large_list(ty), list_view(ty), large_list_view(ty)}) {
+ auto input = ArrayFromJSON(list_type, sample);
+ CheckListContains(input, ScalarFromJSON(ty, "5"),
+ "[true, false, false, true, null, false]");
+ CheckListContains(input, ScalarFromJSON(ty, "7"),
+ "[true, true, false, false, null, false]");
+ CheckListContains(input, ScalarFromJSON(ty, "null"),
+ "[false, true, false, false, null, true]");
+ CheckListContains(ArrayFromJSON(list_type, "[]"), ScalarFromJSON(ty,
"5"), "[]");
+ }
+ }
+
+ auto input =
+ ArrayFromJSON(list(utf8()), R"([["a", "b"], ["a", "c"], ["b", "c",
"d"]])");
+ CheckListContains(input, ScalarFromJSON(utf8(), R"("a")"), "[true, true,
false]");
+ CheckListContains(input->Slice(1), ScalarFromJSON(utf8(), R"("b")"),
"[false, true]");
+
+ // All-null lists contain no non-null value
+ CheckListContains(ArrayFromJSON(list(int64()), "[[null, null]]"),
+ ScalarFromJSON(int64(), "1"), "[false]");
+}
+
+TEST(TestScalarNested, ListContainsLongLists) {
+ // Matches and nulls beyond the first bitmap words, at unaligned offsets
+ std::string values = "[";
+ for (int i = 0; i < 300; ++i) {
+ values += i == 0 ? "" : ", ";
+ values += i == 100 ? "null" : i == 290 ? "7" : "1";
+ }
+ values += "]";
+ ASSERT_OK_AND_ASSIGN(auto input,
+ ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 3,
150, 300]"),
+ *ArrayFromJSON(int32(), values)));
+ CheckListContains(input, ScalarFromJSON(int32(), "7"), "[false, false,
true]");
+ CheckListContains(input, ScalarFromJSON(int32(), "null"), "[false, true,
false]");
+ CheckListContains(input, ArrayFromJSON(int32(), "[7, 7, null]"),
+ "[false, false, false]");
+ CheckListContains(input, ArrayFromJSON(int32(), "[1, null, 7]"), "[true,
true, true]");
+ // The null value's slot holds 0, which must not match
+ CheckListContains(input, ScalarFromJSON(int32(), "0"), "[false, false,
false]");
+}
+
+TEST(TestScalarNested, ListContainsManyShortLists) {
+ // Short lists at every bit alignment, with null lists and null values
+ std::string lists = "[";
+ std::string contains_5 = "[";
+ std::string contains_null = "[";
+ std::string contains_0 = "[";
+ for (int i = 0; i < 200; ++i) {
+ const std::string sep = i == 0 ? "" : ", ";
+ const bool is_null = i % 11 == 0;
+ lists += sep + (is_null ? "null"
+ : i % 7 == 0 ? "[5, 1]"
+ : i % 3 == 0 ? "[null, 2, 3]"
+ : "[1]");
+ contains_5 += sep + (is_null ? "null" : i % 7 == 0 ? "true" : "false");
+ contains_null += sep + (is_null ? "null"
+ : i % 7 != 0 && i % 3 == 0 ? "true"
+ : "false");
+ contains_0 += sep + (is_null ? "null" : "false");
+ }
+ auto input = ArrayFromJSON(list(int32()), lists + "]");
+ CheckListContains(input, ScalarFromJSON(int32(), "5"), contains_5 + "]");
+ CheckListContains(input, ScalarFromJSON(int32(), "null"), contains_null +
"]");
+ // Null values' slots hold 0, which must not match
+ CheckListContains(input, ScalarFromJSON(int32(), "0"), contains_0 + "]");
+}
+
+TEST(TestScalarNested, ListContainsNull) {
+ // A null value matches lists holding a null
+ auto input = ArrayFromJSON(
+ list(int32()), "[[1, null], [null], [1, 2], [1, 1], [2, 1, null, 1], [],
null]");
+ for (auto value : {ScalarFromJSON(int32(), "null"), MakeNullScalar(null()),
+ ScalarFromJSON(int64(), "null")}) {
+ CheckListContains(input, value, "[true, true, false, false, true, false,
null]");
+ }
+ CheckListContains(ArrayFromJSON(list(int32()), "[[]]"),
ScalarFromJSON(int32(), "null"),
+ "[false]");
+
+ input = ArrayFromJSON(list(utf8()), R"([["x", null], [null], ["x", "y"], [],
null])");
+ CheckListContains(input, ScalarFromJSON(utf8(), "null"),
+ "[true, true, false, false, null]");
+
+ input = ArrayFromJSON(list(list(int32())),
+ "[[[1], null], [null], [[1], [null]], [], null]");
+ CheckListContains(input, MakeNullScalar(list(int32())),
+ "[true, true, false, false, null]");
+
+ input = ArrayFromJSON(list(null()), "[[null], [], null]");
+ CheckListContains(input, MakeNullScalar(null()), "[true, false, null]");
+
+ // Dictionary values are also null where their dictionary is
+ ASSERT_OK_AND_ASSIGN(
+ input, ListArray::FromArrays(
+ *ArrayFromJSON(int32(), "[0, 1, 2]"),
+ *DictArrayFromJSON(dictionary(int8(), int64()), "[0, 1]",
"[null, 1]")));
+ for (auto value : {ScalarFromJSON(int64(), "null"), MakeNullScalar(null())})
{
+ CheckListContains(input, value, "[true, false]");
+ }
+ CheckListContains(input, ArrayFromJSON(int64(), "[null, null]"), "[true,
false]");
+ ASSERT_OK_AND_ASSIGN(auto list_scalar, input->GetScalar(0));
+ CheckListContains(list_scalar, ArrayFromJSON(int64(), "[null, 1]"), "[true,
false]");
+ // Also of a null-typed dictionary, both of values and of queries
+ auto null_dict = DictArrayFromJSON(dictionary(int8(), null()), "[0, 0]",
"[null]");
+ ASSERT_OK_AND_ASSIGN(
+ input, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 2, 2]"),
*null_dict));
+ CheckListContains(input, MakeNullScalar(null()), "[true, false]");
+ ASSERT_OK_AND_ASSIGN(list_scalar, input->GetScalar(0));
+ CheckListContains(list_scalar, ArrayFromJSON(int64(), "[null, null]"),
"[true, true]");
+ CheckListContains(ArrayFromJSON(list(struct_({field("a", int64())})),
"[[null], []]"),
+ null_dict, "[true, false]");
+}
+
+TEST(TestScalarNested, ListContainsNaN) {
+ // Unlike "equal", a NaN value matches NaN list values
+ for (auto ty : {float32(), float64()}) {
+ auto input = ArrayFromJSON(list(ty), "[[1.5, null], [NaN], [1.5, NaN], [],
null]");
+ for (auto value_ty : {float32(), float64()}) {
+ CheckListContains(input, ScalarFromJSON(value_ty, "NaN"),
+ "[false, true, true, false, null]");
+ CheckListContains(input, ScalarFromJSON(value_ty, "1.5"),
+ "[true, false, true, false, null]");
+ }
+ }
+ CheckListContains(ArrayFromJSON(list(int64()), "[[1], []]"),
+ ScalarFromJSON(float64(), "NaN"), "[false, false]");
+
+ auto input =
+ ArrayFromJSON(list(float16()), "[[1.5, null], [NaN], [1.5, NaN], [],
null]");
+ CheckListContains(input, ScalarFromJSON(float16(), "NaN"),
+ "[false, true, true, false, null]");
+ CheckListContains(ArrayFromJSON(list(float32()), "[[1.5], [NaN]]"),
+ ScalarFromJSON(float16(), "NaN"), "[false, true]");
+ // Null queries compare float16 as float32 too
+ input = ArrayFromJSON(list(float16()), "[[1], [null]]");
+ CheckListContains(input, ArrayFromJSON(float32(), "[null, null]"), "[false,
true]");
+ CheckListContains(input, ScalarFromJSON(float32(), "null"), "[false, true]");
+ CheckListContains(ArrayFromJSON(list(float32()), "[[1], [null]]"),
+ ArrayFromJSON(float16(), "[null, null]"), "[false, true]");
+ CheckListContains(ScalarFromJSON(list(float16()), "[1, null]"),
+ ArrayFromJSON(float32(), "[null, null]"), "[true, true]");
+ // Also in dictionaries, both of values and of queries
+ auto f16_dict =
+ DictArrayFromJSON(dictionary(int8(), float16()), "[0, 1, 2]", "[NaN, 1,
null]");
+ ASSERT_OK_AND_ASSIGN(
+ input, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 2, 3]"),
*f16_dict));
+ CheckListContains(input, ScalarFromJSON(float32(), "NaN"), "[true, false,
false]");
+ CheckListContains(input, ArrayFromJSON(float32(), "[1, 1, null]"),
+ "[false, true, true]");
+ CheckListContains(ArrayFromJSON(list(float32()), "[[NaN], [1], [null]]"),
f16_dict,
+ "[true, true, true]");
+ ASSERT_OK_AND_ASSIGN(auto f16_nan, f16_dict->GetScalar(0));
+ CheckListContains(ArrayFromJSON(list(float32()), "[[NaN], [1], [null]]"),
f16_nan,
+ "[true, false, false]");
+
+ // NaNs with different payloads (0x7ff8000000000001, 0x7ff8000000000002)
match too
+ ASSERT_OK_AND_ASSIGN(auto nans,
+ ArrayFromJSON(uint64(),
"[9221120237041090561]")->View(float64()));
+ ASSERT_OK_AND_ASSIGN(
+ auto queries, ArrayFromJSON(uint64(), "[9221120237041090562,
4607182418800017408]")
+ ->View(float64())); // [NaN, 1.0]
+ ASSERT_OK_AND_ASSIGN(
+ input, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 1]"),
*nans));
+ CheckListContains(input, queries, "[true, false]");
+ ASSERT_OK_AND_ASSIGN(auto nan_query, queries->GetScalar(0));
+ CheckListContains(input, nan_query, "[true, false]");
+ ASSERT_OK_AND_ASSIGN(auto nan_list, input->GetScalar(0));
+ CheckListContains(nan_list, queries, "[true, false]");
+}
+
+TEST(TestScalarNested, ListContainsValueTypes) {
+ auto check = [](const std::shared_ptr<DataType>& ty, const std::string& x,
+ const std::string& y) {
+ auto input = ArrayFromJSON(
+ list(ty), "[[" + x + ", null], [" + y + "], [" + x + ", " + y + "],
[], null]");
+ CheckListContains(input, ScalarFromJSON(ty, x), "[true, false, true,
false, null]");
+ CheckListContains(input, ScalarFromJSON(ty, y), "[false, true, true,
false, null]");
+ CheckListContains(
+ input,
+ ArrayFromJSON(ty, "[" + x + ", " + y + ", " + y + ", " + x + ", " + x
+ "]"),
+ "[true, true, true, false, null]");
+ // A single list is searched for each query
+ CheckListContains(ScalarFromJSON(list(ty), "[" + x + ", null]"),
+ ArrayFromJSON(ty, "[" + x + ", " + y + ", null]"),
+ "[true, false, true]");
+ };
+ check(utf8(), R"("x")", R"("y")");
+ check(large_binary(), R"("x")", R"("y")");
+ check(binary_view(), R"("x")", R"("y")");
+ check(fixed_size_binary(1), R"("x")", R"("y")");
+ // String views compare with other string layouts, unlike in "equal"
+ CheckListContains(ArrayFromJSON(list(utf8()), R"([["x"], ["y"]])"),
+ ScalarFromJSON(utf8_view(), R"("x")"), "[true, false]");
+ CheckListContains(ArrayFromJSON(list(binary_view()), R"([["x"], ["y"]])"),
+ ArrayFromJSON(large_binary(), R"(["y", "y"])"), "[false,
true]");
+ check(boolean(), "true", "false");
+ check(int16(), "1", "2");
+ check(float16(), "1.5", "2.5");
+ check(date32(), "18262", "18628");
+ check(decimal128(38, 2), R"("1.50")", R"("2.50")");
+ check(decimal256(40, 2), R"("1.50")", R"("2.50")");
+ check(timestamp(TimeUnit::NANO), "1", "2");
+ check(duration(TimeUnit::MILLI), "1", "2");
+
+ // Zero-width values all equal each other
+ auto input =
+ ArrayFromJSON(list(fixed_size_binary(0)), R"([["", null], [null], [],
null])");
+ CheckListContains(input, ScalarFromJSON(fixed_size_binary(0), R"("")"),
+ "[true, false, false, null]");
+ CheckListContains(input, ArrayFromJSON(fixed_size_binary(0), R"(["", "", "",
""])"),
+ "[true, false, false, null]");
+ CheckListContains(ScalarFromJSON(list(fixed_size_binary(0)), R"(["",
null])"),
+ ArrayFromJSON(fixed_size_binary(0), R"(["", null])"),
"[true, true]");
+
+ // Values of other widths never match, even with a common prefix
+ for (const auto& list_type :
+ {list(fixed_size_binary(1)), list_view(fixed_size_binary(1))}) {
+ input = ArrayFromJSON(list_type, R"([["a"], ["b", null], null])");
+ CheckListContains(input, ScalarFromJSON(fixed_size_binary(2), R"("ab")"),
+ "[false, false, null]");
+ CheckListContains(input, ArrayFromJSON(fixed_size_binary(2), R"(["ab",
null, "ab"])"),
+ "[false, true, null]");
+ }
+ CheckListContains(ScalarFromJSON(list(fixed_size_binary(2)), R"(["ab",
null])"),
+ ArrayFromJSON(fixed_size_binary(1), R"(["a", null])"),
+ "[false, true]");
+}
+
+TEST(TestScalarNested, ListContainsSignedZero) {
+ // Like "equal", -0.0 equals 0.0
+ auto input = ArrayFromJSON(list(float64()), "[[0.0], [-0.0], [1.0], null]");
+ CheckListContains(input, ScalarFromJSON(float64(), "-0.0"),
+ "[true, true, false, null]");
+ CheckListContains(input, ArrayFromJSON(float64(), "[-0.0, 0.0, 0.0, 0.0]"),
+ "[true, true, false, null]");
+ CheckListContains(ScalarFromJSON(list(float64()), "[-0.0, NaN]"),
+ ArrayFromJSON(float64(), "[0.0, -0.0, NaN, 1.0]"),
+ "[true, true, true, false]");
+}
+
+TEST(TestScalarNested, ListContainsImplicitCast) {
+ auto input = ArrayFromJSON(list(int64()),
+ "[[2, 2, 3, null, null], null, [], [null], [1],
[0, -1]]");
+ CheckListContains(input, ScalarFromJSON(float64(), "2.0"),
+ "[true, null, false, false, false, false]");
+ CheckListContains(input, ScalarFromJSON(float64(), "1.5"),
+ "[false, null, false, false, false, false]");
+
+ // No overflow of the list values
+ CheckListContains(ArrayFromJSON(list(int8()), "[[44, 1], null, [], [null]]"),
+ ScalarFromJSON(int64(), "300"), "[false, null, false,
false]");
+ CheckListContains(ArrayFromJSON(list(float64()), "[[1.0, 2.0], [3.0], null,
[]]"),
+ ScalarFromJSON(int64(), "1"), "[true, false, null,
false]");
+ CheckListContains(ArrayFromJSON(list(decimal128(38, 2)), R"([["1.50"],
["2.50"]])"),
+ ScalarFromJSON(decimal128(3, 2), R"("1.50")"), "[true,
false]");
+
+ // Values of null lists are not cast: 2^53 + 1 isn't a float64
+ auto values = ArrayFromJSON(int64(), "[9007199254740993, 1]");
+ auto validity = ArrayFromJSON(boolean(), "[false,
true]")->data()->buffers[1];
+ ASSERT_OK_AND_ASSIGN(
+ auto lists, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 2]"),
*values,
+ default_memory_pool(), validity));
+ ASSERT_OK_AND_ASSIGN(auto fixed_size_lists,
+ FixedSizeListArray::FromArrays(values, 1, validity));
+ for (Datum input : {Datum(lists), Datum(fixed_size_lists)}) {
+ CheckListContains(input, ScalarFromJSON(float64(), "1"), "[null, true]");
+ CheckListContains(input, ArrayFromJSON(float64(), "[1, null]"), "[null,
false]");
+ }
+ // Unlike values of non-null lists
+ EXPECT_THAT(CallFunction("list_contains",
+ {ArrayFromJSON(list(int64()), "[[9007199254740993],
null]"),
+ ScalarFromJSON(float64(), "1")}),
+ Raises(StatusCode::Invalid));
+
+ // Queries are cast like for "equal", even those of null lists
+ auto floats = list(float64());
+ for (Datum lists : {Datum(ArrayFromJSON(floats, "[[1], null]")),
+ Datum(ArrayFromJSON(floats, "[null, null]")),
+ Datum(ScalarFromJSON(floats, "null"))}) {
+ for (Datum queries : {Datum(ArrayFromJSON(int64(), "[9007199254740993,
1]")),
+ Datum(ScalarFromJSON(int64(), "9007199254740993"))})
{
+ EXPECT_THAT(CallFunction("list_contains", {lists, queries}),
+ Raises(StatusCode::Invalid));
+ }
+ }
+}
+
+TEST(TestScalarNested, ListContainsChunked) {
+ auto array = ArrayFromJSON(list(int32()), "[[1, 2], [3], null, [], [2,
null]]");
+ auto input = std::make_shared<ChunkedArray>(
+ ArrayVector{array->Slice(3), array->Slice(0, 0), array->Slice(1, 2)});
+ ASSERT_OK_AND_ASSIGN(
+ Datum result, CallFunction("list_contains", {input,
ScalarFromJSON(int32(), "2")}));
+ AssertDatumsEqual(
+ ChunkedArrayFromJSON(boolean(), {"[false, true]", "[]", "[false,
null]"}), result,
+ /*verbose=*/true);
+}
+
+TEST(TestScalarNested, ListContainsFixedSizeList) {
+ auto input = ArrayFromJSON(fixed_size_list(int32(), 2),
+ "[[1, 2], [3, null], null, [2, 4], [null,
null]]");
+ CheckListContains(input, ScalarFromJSON(int32(), "2"),
+ "[true, false, null, true, false]");
+ CheckListContains(input->Slice(2), ScalarFromJSON(int32(), "4"), "[null,
true, false]");
+}
+
+TEST(TestScalarNested, ListContainsListViewOutOfOrder) {
+ // Views: [3, 4], [1, 2], [2, 3], [4], []
+ ASSERT_OK_AND_ASSIGN(
+ auto input, ListViewArray::FromArrays(*ArrayFromJSON(int32(), "[2, 0, 1,
3, 0]"),
+ *ArrayFromJSON(int32(), "[2, 2, 2,
1, 0]"),
+ *ArrayFromJSON(int32(), "[1, 2, 3,
4]")));
+ CheckListContains(input, ScalarFromJSON(int32(), "3"),
+ "[true, false, true, false, false]");
+ CheckListContains(input, ScalarFromJSON(int32(), "1"),
+ "[false, true, false, false, false]");
+ CheckListContains(input, ArrayFromJSON(int32(), "[3, 1, 1, 4, 1]"),
+ "[true, true, false, true, false]");
+ // Only the child values referenced by the (sliced) views are compared
+ CheckListContains(input->Slice(2), ScalarFromJSON(int32(), "3"),
+ "[true, false, false]");
+ CheckListContains(input->Slice(3), ScalarFromJSON(int32(), "1"), "[false,
false]");
+ CheckListContains(input->Slice(4), ScalarFromJSON(int32(), "1"), "[false]");
+
+ // Child values outside every view are not cast: 2^53 + 1 isn't a float64
+ ASSERT_OK_AND_ASSIGN(
+ input, ListViewArray::FromArrays(
+ *ArrayFromJSON(int32(), "[0, 2]"), *ArrayFromJSON(int32(),
"[1, 1]"),
+ *ArrayFromJSON(int64(), "[1, 9007199254740993, 3,
9007199254740993]")));
+ CheckListContains(input, ScalarFromJSON(float64(), "1"), "[true, false]");
+ CheckListContains(input, ArrayFromJSON(float64(), "[3, 3]"), "[false,
true]");
+ auto views = ArrayFromJSON(list_view(int64()), "[[1, 2], null, [], [3,
null]]");
+ CheckListContains(views, ScalarFromJSON(float64(), "3"), "[false, null,
false, true]");
+ CheckListContains(views, ArrayFromJSON(float64(), "[2, 1, 1, null]"),
+ "[true, null, false, true]");
+}
+
+TEST(TestScalarNested, ListContainsArrayOfValues) {
+ // Each list is searched for the value at the same index
+ for (auto list_type : {list(int32()), large_list(int32()),
list_view(int32()),
+ large_list_view(int32())}) {
+ CheckListContains(
+ ArrayFromJSON(list_type, "[[1, 2], [3, null], [], null, [4, 5],
[null], [6]]"),
+ ArrayFromJSON(int32(), "[2, null, 1, 3, 4, null, 7]"),
+ "[true, true, false, null, true, true, false]");
+ }
+ CheckListContains(ArrayFromJSON(fixed_size_list(int32(), 2),
+ "[[1, 2], [3, null], null, [4, 5], [null,
6]]"),
+ ArrayFromJSON(int32(), "[2, null, 1, 5, 7]"),
+ "[true, true, null, true, false]");
+ CheckListContains(ArrayFromJSON(list(int64()), "[[1, 2], [3], [4]]"),
+ ArrayFromJSON(float64(), "[2.0, 3.5, 4.0]"), "[true,
false, true]");
+ CheckListContains(ArrayFromJSON(list(float64()), "[[1.5, NaN], [NaN], [1.5],
[null]]"),
+ ArrayFromJSON(float64(), "[NaN, 1.5, NaN, NaN]"),
+ "[true, false, false, false]");
+ CheckListContains(ArrayFromJSON(list(int32()), "[[1, null], [2], []]"),
+ ArrayFromJSON(null(), "[null, null, null]"), "[true,
false, false]");
+
+ // A scalar list is searched for each value
+ CheckListContains(ScalarFromJSON(list(int32()), "[1, null]"),
+ ArrayFromJSON(int32(), "[1, 2, null]"), "[true, false,
true]");
+ CheckListContains(ScalarFromJSON(fixed_size_list(int32(), 2), "[1, 2]"),
+ ArrayFromJSON(int32(), "[2, 3]"), "[true, false]");
+ CheckListContains(ScalarFromJSON(list(int32()), "null"),
+ ArrayFromJSON(int32(), "[1, 2]"), "[null, null]");
+ CheckListContains(ScalarFromJSON(list(int64()), "[1, 2]"),
+ ArrayFromJSON(float64(), "[2.0, 2.5, null]"), "[true,
false, false]");
+ CheckListContains(ScalarFromJSON(list(utf8()), R"(["a", null])"),
+ ArrayFromJSON(null(), "[null, null]"), "[true, true]");
+ CheckListContains(ScalarFromJSON(list(utf8()), R"(["a"])"),
+ ArrayFromJSON(utf8(), "[null, null]"), "[false, false]");
+}
+
+TEST(TestScalarNested, ListContainsInvalid) {
+ auto input = ArrayFromJSON(list(int32()), "[[1, 2], [3]]");
+ // Null values must be comparable too
+ for (Datum value :
+ {Datum(ArrayFromJSON(utf8(), R"(["a", "b"])")),
+ Datum(ScalarFromJSON(utf8(), R"("a")")),
Datum(ScalarFromJSON(boolean(), "true")),
+ Datum(ScalarFromJSON(utf8(), "null")),
+ Datum(ArrayFromJSON(utf8(), "[null, null]"))}) {
+ EXPECT_THAT(CallFunction("list_contains", {input, value}),
+ Raises(StatusCode::NotImplemented));
+ EXPECT_THAT(
+ CallFunction("list_contains", {ScalarFromJSON(list(int32()), "[1]"),
+ ArrayFromJSON(value.type(), "[null,
null]")}),
+ Raises(StatusCode::NotImplemented));
+ }
+ // Also for null lists, whose value type must be comparable all the same
+ for (Datum lists : {Datum(ArrayFromJSON(list(int32()), "[null, null]")),
+ Datum(ScalarFromJSON(list(int32()), "null"))}) {
+ for (auto queries : {ArrayFromJSON(utf8(), R"(["a", "b"])"),
+ ArrayFromJSON(utf8(), "[null, null]")}) {
+ EXPECT_THAT(CallFunction("list_contains", {lists, queries}),
+ Raises(StatusCode::NotImplemented));
+ }
+ }
+ EXPECT_THAT(CallFunction("list_contains", {ArrayFromJSON(list(int32()),
"[]"),
+ ArrayFromJSON(utf8(), "[]")}),
+ Raises(StatusCode::NotImplemented));
+ // Lists, unions and run-end encoded values only compare with null queries,
which
+ // only need validity. Other queries raise, also for null lists.
+ auto union_type = sparse_union({field("i", int64()), field("s", utf8())},
{0, 1});
+ auto ree_type = run_end_encoded(int32(), int64());
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<Array> ree_queries,
+ RunEndEncodedArray::Make(2, ArrayFromJSON(int32(),
"[2]"),
+ ArrayFromJSON(int64(),
"[1]")));
+ ASSERT_OK_AND_ASSIGN(auto null_ree_queries, MakeArrayOfNull(ree_type, 2));
+ for (const auto& [queries, null_queries] :
+ {std::pair(ArrayFromJSON(list(int32()), "[[1], null]"),
+ ArrayFromJSON(list(int32()), "[null, null]")),
+ std::pair(ArrayFromJSON(union_type, R"([[0, 1], [1, "a"]])"),
+ ArrayFromJSON(union_type, "[[0, null], [1, null]]")),
+ std::pair(ree_queries, null_ree_queries)}) {
+ ARROW_SCOPED_TRACE(*queries->type());
+ ASSERT_OK_AND_ASSIGN(auto null_lists,
MakeArrayOfNull(list(queries->type()), 2));
+ ASSERT_OK_AND_ASSIGN(
+ auto empty_and_null,
+ ListArray::FromArrays(
+ *ArrayFromJSON(int32(), "[0, 0, 0]"), *null_queries->Slice(0, 0),
+ default_memory_pool(),
+ ArrayFromJSON(boolean(), "[true, false]")->data()->buffers[1]));
+ for (Datum lists : {Datum(null_lists),
Datum(MakeNullScalar(list(queries->type()))),
+ Datum(empty_and_null)}) {
+ EXPECT_THAT(CallFunction("list_contains", {lists, queries}),
+ Raises(StatusCode::NotImplemented));
+ }
+ CheckListContains(null_lists, null_queries, "[null, null]");
+ CheckListContains(MakeNullScalar(list(queries->type())), null_queries,
+ "[null, null]");
+ CheckListContains(empty_and_null, null_queries, "[false, null]");
+ }
+ // Like "equal", timestamps with and without a time zone don't compare
+ for (auto value : {"1", "null"}) {
+ EXPECT_THAT(
+ CallFunction("list_contains",
+ {ArrayFromJSON(list(timestamp(TimeUnit::SECOND, "UTC")),
"[[1]]"),
+ ScalarFromJSON(timestamp(TimeUnit::SECOND), value)}),
+ Raises(StatusCode::Invalid));
+ }
+ // Like "equal", timestamps don't compare with durations, which don't cast
to them,
+ // also in dictionaries
+ auto timestamps = list(timestamp(TimeUnit::SECOND));
+ for (Datum lists : {Datum(ArrayFromJSON(timestamps, "[[1], null]")),
+ Datum(ScalarFromJSON(timestamps, "null"))}) {
+ for (auto queries : {"[1, null]", "[null, null]"}) {
+ auto durations = ArrayFromJSON(duration(TimeUnit::SECOND), queries);
+ ASSERT_OK_AND_ASSIGN(Datum encoded, DictionaryEncode(durations));
+ for (Datum value : {Datum(durations), encoded}) {
+ EXPECT_THAT(CallFunction("list_contains", {lists, value}),
+ Raises(StatusCode::NotImplemented));
+ }
+ }
+ }
+}
+
using VarLenListLikeTypeFactory =
std::shared_ptr<DataType> (*)(std::shared_ptr<DataType>);
static const VarLenListLikeTypeFactory kVarLenListTypeFactories[] = {
diff --git a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc
b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc
index 0a041e7e47f..b72ad9cbe85 100644
--- a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc
+++ b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc
@@ -202,6 +202,13 @@ TEST(TestValidityKernels, IsValidIsNullNullType) {
ArrayFromJSON(boolean(), "[false, false, false, false,
false]"));
}
+TEST(TestValidityKernels, IsValidIsNullDictionaryOfNullType) {
+ // The dictionary has no validity bitmap
+ auto arr = DictArrayFromJSON(dictionary(int8(), null()), "[0, null, 0]",
"[null]");
+ CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[true, true,
true]"));
+ CheckScalarUnary("is_valid", arr, ArrayFromJSON(boolean(), "[false, false,
false]"));
+}
+
TEST(TestValidityKernels, IsNullSetsZeroNullCount) {
auto arr = ArrayFromJSON(int32(), "[1, 2, 3, 4, null]");
ASSERT_OK_AND_ASSIGN(Datum out, IsNull(arr));
diff --git a/cpp/src/arrow/util/dict_util.cc b/cpp/src/arrow/util/dict_util.cc
index abcfd240e67..00f72746604 100644
--- a/cpp/src/arrow/util/dict_util.cc
+++ b/cpp/src/arrow/util/dict_util.cc
@@ -46,7 +46,8 @@ int64_t LogicalNullCount(const ArraySpan& span) {
// in the dictionary array to see if they are null
for (int64_t i = position; i < position + length; i++) {
CType index = indices_data[i];
- if (!bit_util::GetBit(dictionary_valid_bitmap,
+ if (!dictionary_valid_bitmap ||
+ !bit_util::GetBit(dictionary_valid_bitmap,
index + dictionary_span.offset)) {
null_count++;
}
@@ -77,7 +78,8 @@ void SetLogicalNullBits(const ArraySpan& span, uint8_t*
out_bitmap, int64_t out_
// in the dictionary array to see if they are null
for (int64_t i = position; i < position + length; i++) {
CType index = indices_data[i];
- bool is_null = !bit_util::GetBit(dictionary_valid_bitmap,
+ bool is_null = !dictionary_valid_bitmap ||
+ !bit_util::GetBit(dictionary_valid_bitmap,
index + dictionary_span.offset);
bit_util::SetBitTo(out_bitmap, out_offset + i, is_null ==
set_on_null);
}
diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst
index 15d426dca6c..c3c182ebddd 100644
--- a/docs/source/cpp/compute.rst
+++ b/docs/source/cpp/compute.rst
@@ -1947,6 +1947,8 @@ Structural transforms
+---------------------+------------+-------------------------------------+------------------+------------------------------+--------+
| Function name | Arity | Input types |
Output type | Options class | Notes |
+=====================+============+=====================================+==================+==============================+========+
+| list_contains | Binary | List-like (Arg 0), Any (Arg 1) |
Boolean | | \(7) |
++---------------------+------------+-------------------------------------+------------------+------------------------------+--------+
| list_element | Binary | List-like (Arg 0), Integral (Arg 1) |
List value type | | \(1) |
+---------------------+------------+-------------------------------------+------------------+------------------------------+--------+
| list_flatten | Unary | List-like |
List value type | | \(2) |
@@ -2005,6 +2007,14 @@ 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
+ value at the same index, and a scalar list is searched for each value: the
+ list ``[1, 2]`` and the array ``[2, 3]`` give ``[true, false]``. Unlike
+ ``equal``, a null second argument matches null list values and a NaN second
+ argument matches NaN list values; otherwise null list values never match.
+ Null lists emit a null in the output.
+
.. _cpp-compute-vector-replace-functions:
Replace functions
diff --git a/docs/source/python/api/compute.rst
b/docs/source/python/api/compute.rst
index ab9ef655b79..43bcd0479ce 100644
--- a/docs/source/python/api/compute.rst
+++ b/docs/source/python/api/compute.rst
@@ -573,6 +573,7 @@ Structural Transforms
fill_null
fill_null_backward
fill_null_forward
+ list_contains
list_element
list_flatten
list_parent_indices
diff --git a/python/pyarrow/tests/test_compute.py
b/python/pyarrow/tests/test_compute.py
index a8c2112d5e5..691d2ec98f3 100644
--- a/python/pyarrow/tests/test_compute.py
+++ b/python/pyarrow/tests/test_compute.py
@@ -3903,6 +3903,124 @@ def test_list_element():
assert result.equals(expected)
[email protected]("list_type", [
+ pa.list_(pa.int64()), pa.large_list(pa.int64()),
+ pa.list_view(pa.int64()), pa.large_list_view(pa.int64()),
+])
+def test_list_contains(list_type):
+ lists = pa.array([[1, 2, None], [3], [], None], list_type)
+ assert pc.list_contains(lists, 2).to_pylist() == [True, False, False, None]
+ assert pc.list_contains(lists, 2.0).to_pylist() == [True, False, False,
None]
+ assert pc.list_contains(lists, 4).to_pylist() == [False, False, False,
None]
+ # A null value matches null list values
+ assert pc.list_contains(lists, None).to_pylist() == [True, False, False,
None]
+
+
+def test_list_contains_nan():
+ lists = pa.array([[1.5, None], [float("nan")], []])
+ result = pc.list_contains(lists, float("nan"))
+ assert result.to_pylist() == [False, True, False]
+
+
+def test_list_contains_fixed_size_list():
+ lists = pa.array([["a", "b"], ["c", None], None], pa.list_(pa.string(), 2))
+ assert pc.list_contains(lists, "c").to_pylist() == [False, True, None]
+
+
+def test_list_contains_array():
+ lists = pa.array([[1, 2], [3, None], [4], None])
+ values = pa.array([2, None, 5, 1])
+ result = pc.list_contains(lists, values)
+ assert result.to_pylist() == [True, True, False, None]
+
+
[email protected]("value_type, values, query_type, queries", [
+ # 2**53 + 1 doesn't cast to float64, so it's only in null lists
+ (pa.int64(), [0, 1, -1, 2**53 + 1], pa.float64(), [0.0, 1.0, 1.5]),
+ (pa.float32(), [0.0, -0.0, 1.5, math.nan], pa.float64(), [0.0, 1.5,
math.nan]),
+ (pa.string(), ["", "a", "ab"], pa.large_string(), ["", "a", "b"]),
+ (pa.string_view(), ["", "a", "ab"], pa.string(), ["", "a", "b"]),
+ (pa.binary(1), [b"a", b"b"], pa.binary(2), [b"ab"]),
+ # Dictionary values are null for null indices and for nulls in the
dictionary
+ (pa.dictionary(pa.int8(), pa.int64()), [0, 1, None], pa.int64(), [0, 1,
2]),
+ (pa.dictionary(pa.int8(), pa.float16()), [0.0, math.nan, None],
pa.float32(),
+ [0.0, 1.0, math.nan]),
+])
+def test_list_contains_matches_decoded(value_type, values, query_type,
queries):
+ # Compare list_contains to searching the decoded lists, for each input form
+ def search(lists, query):
+ if lists is None:
+ return None
+ if query is None:
+ return None in lists
+ if query != query:
+ return any(v != v for v in lists if v is not None)
+ return any(v == query for v in lists if v is not None)
+
+ rng = random.Random(42)
+ n = 100
+ sizes = [rng.randint(0, 4) for _ in range(n)]
+ offsets = list(itertools.accumulate(sizes, initial=0))
+ mask = [rng.random() < 0.1 for _ in range(n)]
+ hidden = set()
+ if pa.types.is_dictionary(value_type):
+ indices = [rng.choice([0, 1, 2, None]) for _ in range(offsets[-1])]
+ child = pa.DictionaryArray.from_arrays(
+ pa.array(indices, value_type.index_type),
+ pa.array(values, value_type.value_type))
+ else:
+ child = [rng.choice(values[:3] + [None]) for _ in range(offsets[-1])]
+ for i in itertools.compress(range(n), mask):
+ child[offsets[i]:offsets[i + 1]] = values[-1:] * sizes[i]
+ hidden.update(range(offsets[i], offsets[i + 1]))
+ child = pa.array(child, value_type)
+ lists = pa.ListArray.from_arrays(offsets, child, mask=pa.array(mask))
+ forms = [
+ lists, lists.slice(7, 50), lists.cast(pa.large_list(value_type)),
+ # Views of the same lists in reverse order
+ pa.ListViewArray.from_arrays(offsets[-2::-1], sizes[::-1], child,
+ mask=pa.array(mask[::-1])),
+ pa.FixedSizeListArray.from_arrays(
+ child.slice(0, n), 1, mask=pa.array([i in hidden for i in
range(n)])),
+ ] + [pa.chunked_array([lists.slice(i, size) for i in range(0, n, size)])
+ for size in (1, 7, n)]
+ query_array = pa.array([rng.choice(queries + [None]) for _ in range(n)],
+ query_type)
+ for form in forms:
+ decoded = form.to_pylist()
+ for query in queries + [None]:
+ result = pc.list_contains(form, pa.scalar(query, query_type))
+ assert result.to_pylist() == [search(x, query) for x in decoded]
+ expected = [search(x, q) for x, q in zip(decoded,
query_array.to_pylist())]
+ q = query_array.slice(0, len(form))
+ assert pc.list_contains(form, q).to_pylist() == expected
+ assert pc.list_contains(form, pa.chunked_array(
+ [q.slice(0, 3), q.slice(3)])).to_pylist() == expected
+ assert pc.list_contains(form, q.dictionary_encode()).to_pylist() ==
expected
+ for x in lists[:10]:
+ result = pc.list_contains(x, query_array)
+ assert result.to_pylist() == [search(x.as_py(), q)
+ for q in query_array.to_pylist()]
+
+
+def test_list_contains_null_list_queries():
+ # Unlike values of null lists, their queries are cast, as for "equal":
+ # 2**53 + 1 isn't a float64
+ ty = pa.list_(pa.float64())
+ queries = pa.array([2**53 + 1], pa.int64())
+ for lists in (pa.scalar(None, ty), pa.array([None], ty)):
+ with pytest.raises(pa.ArrowInvalid):
+ pc.list_contains(lists, queries)
+
+
+def test_list_contains_invalid():
+ lists = pa.array([[1, 2], [3]])
+ with pytest.raises(pa.ArrowNotImplementedError):
+ pc.list_contains(lists, "a")
+ with pytest.raises(pa.ArrowNotImplementedError):
+ pc.list_contains(lists, pa.array(["a", "b"]))
+
+
def test_count_distinct():
samples = [datetime.datetime(year=y, month=1, day=1) for y in range(1992,
2092)]
arr = pa.array(samples, pa.timestamp("ns"))
