jonasdedden commented on code in PR #51611:
URL: https://github.com/apache/arrow/pull/51611#discussion_r4335776309
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +539,332 @@ 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());
+}
+
+// Cast `values` and `queries` to the types "equal" compares them as, raising
the
+// errors it would. float16, which "equal" lacks, is compared as float32, and
binary
+// views as large binary.
+Status CastToComparable(KernelContext* ctx, Datum* values, Datum* queries) {
+ ExecContext* exec_ctx = ctx->exec_context();
+ for (Datum* datum : {values, queries}) {
+ if (datum->type()->id() == Type::HALF_FLOAT) {
+ ARROW_ASSIGN_OR_RAISE(*datum,
+ Cast(*datum, float32(), CastOptions::Safe(),
exec_ctx));
+ }
+ }
+ RETURN_NOT_OK(CheckComparable(ctx, values->type(), queries->type()));
+ ARROW_ASSIGN_OR_RAISE(auto equal,
exec_ctx->func_registry()->GetFunction("equal"));
+ std::vector<TypeHolder> types{values->type(), queries->type()};
+ RETURN_NOT_OK(equal->DispatchBest(&types));
+ for (Datum* datum : {values, queries}) {
+ TypeHolder type = types[datum == values ? 0 : 1];
+ if (is_binary_view_like(type.id())) {
+ type = large_binary();
+ }
+ if (!datum->type()->Equals(*type)) {
+ ARROW_ASSIGN_OR_RAISE(*datum, Cast(*datum, type, CastOptions::Safe(),
exec_ctx));
+ }
+ }
+ 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 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);
+ return [&list, offsets, sizes](int64_t i) {
+ return Range(offsets[i], list.IsValid(i) ? sizes[i] : 0);
+ };
+ } else {
+ return
+ [offsets](int64_t i) { return Range(offsets[i], offsets[i + 1] -
offsets[i]); };
+ }
+ }
+}
+
+// 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 Type, typename Matches, typename Confirm = std::nullptr_t>
+void SearchLists(const ArraySpan& list, 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>;
+ const auto list_range = GetListRanges<Type>(list);
+ 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 * list.length) {
+ // 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) {
+ return Status::OK();
+ }
+ Datum values = list.value;
+ Datum queries = query_span.ToArrayData();
+ if (queries.null_count() == queries.length()) {
+ RETURN_NOT_OK(CheckNullQueries(ctx, values, queries));
+ bit_util::SetBitsTo(out_arr->buffers[1].data, out_arr->offset,
out_arr->length,
+ values.null_count() > 0);
+ return Status::OK();
+ }
+ RETURN_NOT_OK(CastToComparable(ctx, &values, &queries));
+ ExecContext* exec_ctx = ctx->exec_context();
+ if (is_floating(values.type()->id())) {
+ // "is_in" tells -0.0 from 0.0, so turn -0.0 into 0.0
+ ARROW_ASSIGN_OR_RAISE(auto zero, MakeScalar(values.type(), 0.0));
+ for (Datum* datum : {&values, &queries}) {
+ ARROW_ASSIGN_OR_RAISE(*datum, CallFunction("add", {*datum, zero},
exec_ctx));
+ }
+ }
+ const SetLookupOptions options(values);
+ ARROW_ASSIGN_OR_RAISE(Datum found,
+ CallFunction("is_in", {queries}, &options, exec_ctx));
Review Comment:
Fixed in 6fc4996: before `is_in`, NaNs in both the list and the queries are
replaced by one canonical NaN (next to the existing -0.0 → 0.0 step). Added a
regression test with payloads `0x7ff8000000000001` / `0x7ff8000000000002` for
list arrays with scalar and array queries, and for a scalar list. Cost:
+0.4–1.4 ms for 1M float queries.
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +539,332 @@ 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());
+}
+
+// Cast `values` and `queries` to the types "equal" compares them as, raising
the
+// errors it would. float16, which "equal" lacks, is compared as float32, and
binary
+// views as large binary.
+Status CastToComparable(KernelContext* ctx, Datum* values, Datum* queries) {
+ ExecContext* exec_ctx = ctx->exec_context();
+ for (Datum* datum : {values, queries}) {
+ if (datum->type()->id() == Type::HALF_FLOAT) {
+ ARROW_ASSIGN_OR_RAISE(*datum,
+ Cast(*datum, float32(), CastOptions::Safe(),
exec_ctx));
+ }
+ }
+ RETURN_NOT_OK(CheckComparable(ctx, values->type(), queries->type()));
+ ARROW_ASSIGN_OR_RAISE(auto equal,
exec_ctx->func_registry()->GetFunction("equal"));
+ std::vector<TypeHolder> types{values->type(), queries->type()};
+ RETURN_NOT_OK(equal->DispatchBest(&types));
+ for (Datum* datum : {values, queries}) {
+ TypeHolder type = types[datum == values ? 0 : 1];
+ if (is_binary_view_like(type.id())) {
+ type = large_binary();
+ }
+ if (!datum->type()->Equals(*type)) {
+ ARROW_ASSIGN_OR_RAISE(*datum, Cast(*datum, type, CastOptions::Safe(),
exec_ctx));
+ }
+ }
+ 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 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) {
Review Comment:
Fixed in 0110702: width 0 now matches every value (without `memcmp`, which
could otherwise get null data pointers). Tests cover list arrays with scalar
and array queries and a scalar list.
##########
cpp/src/arrow/compute/kernels/scalar_nested.cc:
##########
@@ -537,6 +539,332 @@ 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());
+}
+
+// Cast `values` and `queries` to the types "equal" compares them as, raising
the
+// errors it would. float16, which "equal" lacks, is compared as float32, and
binary
+// views as large binary.
+Status CastToComparable(KernelContext* ctx, Datum* values, Datum* queries) {
+ ExecContext* exec_ctx = ctx->exec_context();
+ for (Datum* datum : {values, queries}) {
+ if (datum->type()->id() == Type::HALF_FLOAT) {
+ ARROW_ASSIGN_OR_RAISE(*datum,
+ Cast(*datum, float32(), CastOptions::Safe(),
exec_ctx));
+ }
+ }
+ RETURN_NOT_OK(CheckComparable(ctx, values->type(), queries->type()));
+ ARROW_ASSIGN_OR_RAISE(auto equal,
exec_ctx->func_registry()->GetFunction("equal"));
+ std::vector<TypeHolder> types{values->type(), queries->type()};
+ RETURN_NOT_OK(equal->DispatchBest(&types));
+ for (Datum* datum : {values, queries}) {
+ TypeHolder type = types[datum == values ? 0 : 1];
+ if (is_binary_view_like(type.id())) {
+ type = large_binary();
+ }
+ if (!datum->type()->Equals(*type)) {
+ ARROW_ASSIGN_OR_RAISE(*datum, Cast(*datum, type, CastOptions::Safe(),
exec_ctx));
+ }
+ }
+ 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 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);
+ return [&list, offsets, sizes](int64_t i) {
+ return Range(offsets[i], list.IsValid(i) ? sizes[i] : 0);
+ };
+ } else {
+ return
+ [offsets](int64_t i) { return Range(offsets[i], offsets[i + 1] -
offsets[i]); };
+ }
+ }
+}
+
+// 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 Type, typename Matches, typename Confirm = std::nullptr_t>
+void SearchLists(const ArraySpan& list, 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>;
+ const auto list_range = GetListRanges<Type>(list);
+ 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 * list.length) {
+ // 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) {
+ return Status::OK();
+ }
+ Datum values = list.value;
+ Datum queries = query_span.ToArrayData();
+ if (queries.null_count() == queries.length()) {
+ RETURN_NOT_OK(CheckNullQueries(ctx, values, queries));
+ bit_util::SetBitsTo(out_arr->buffers[1].data, out_arr->offset,
out_arr->length,
+ values.null_count() > 0);
+ return Status::OK();
+ }
+ RETURN_NOT_OK(CastToComparable(ctx, &values, &queries));
+ ExecContext* exec_ctx = ctx->exec_context();
+ if (is_floating(values.type()->id())) {
+ // "is_in" tells -0.0 from 0.0, so turn -0.0 into 0.0
+ ARROW_ASSIGN_OR_RAISE(auto zero, MakeScalar(values.type(), 0.0));
+ for (Datum* datum : {&values, &queries}) {
+ ARROW_ASSIGN_OR_RAISE(*datum, CallFunction("add", {*datum, zero},
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();
+ }
+ if (batch[0].is_scalar()) {
+ // A scalar list is only paired with an array of queries here
+ return ScalarListContains(ctx, 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
+ int64_t values_start = 0;
+ int64_t values_length = list.child_data[0].length;
+ if constexpr (!is_list_view_type<Type>::value) {
+ const auto list_range = GetListRanges<Type>(list);
+ values_start = list_range(0).first;
+ values_length = list_range(list.length - 1).first +
+ list_range(list.length - 1).second - values_start;
+ }
+ Datum values = list.child_data[0].ToArrayData()->Slice(values_start,
values_length);
Review Comment:
Fixed in d9dd833. Without a cast, nothing changed (the child is searched in
place, nothing is converted). With a cast, only the span from the first to the
last referenced value is cast, as a zero-copy slice, so your example no longer
converts the unused value. If converting the span fails, which can only happen
in a gap between views, just the referenced values are gathered and cast.
Regression tests cover the value outside every view and one in a gap, with
scalar and array queries.
I measured the alternatives first. Flattening the views was up to ~100×
slower for short out-of-order views (it concatenates one slice per view), and
always gathering made casted views 2–4× slower with twice the memory. With this
approach, casted views take the same time and memory as before (e.g. 131 ms /
404 MB for 16.7M 3-element views), and views without a cast got 3–13% faster.
--
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]
