diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index 99a9173041fc..7909971f105c 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(lhs); + const auto& rhs_timestamp = checked_cast(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 3a2bcefab0a1..9da440cc65ab 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 0f2c479b0828..e7f12838c91a 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 struct CompareTimestamps { static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { - const auto& lhs = checked_cast(*batch[0].type()); - const auto& rhs = checked_cast(*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::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 e9c65aff1ce1..63afe9208a9e 100644 --- a/cpp/src/arrow/compute/kernels/scalar_nested.cc +++ b/cpp/src/arrow/compute/kernels/scalar_nested.cc @@ -18,9 +18,12 @@ // Vector kernels involving nested types #include +#include #include "arrow/array/array_base.h" #include "arrow/array/builder_nested.h" +#include "arrow/array/builder_primitive.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 +31,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 +541,430 @@ const FunctionDoc list_element_doc( "is emitted. Null values emit a null in the output."), {"lists", "index"}); +// The types "equal" compares `values` and `queries` as, raising the errors it would. +// float16, which "equal" lacks, is compared as float32, and binary views as large +// binary. +Result> ComparableTypes( + KernelContext* ctx, const std::shared_ptr& values_type, + const std::shared_ptr& queries_type) { + std::vector types; + for (const auto& type : {values_type, queries_type}) { + types.emplace_back(type->id() == Type::HALF_FLOAT ? float32() : type); + } + ARROW_ASSIGN_OR_RAISE(auto equal, + ctx->exec_context()->func_registry()->GetFunction("equal")); + RETURN_NOT_OK(equal->DispatchBest(&types)); + RETURN_NOT_OK(ValidateComparison(*types[0], *types[1])); + if (types[0].id() == Type::FIXED_SIZE_BINARY && *types[0] != *types[1]) { + // Fixed-size binary values of different widths differ, as binary values do + types = {binary(), binary()}; + } + for (const auto& [from, to] : + {std::pair(values_type, &types[0]), std::pair(queries_type, &types[1])}) { + if (is_binary_view_like(to->id())) { + *to = large_binary(); + } + // Raise also where there is nothing to cast, e.g. for null queries + if (!from->Equals(**to) && !CanCast(*from, **to)) { + return Status::NotImplemented("Unsupported cast from ", *from, " to ", **to); + } + } + return types; +} + +struct ListContainsState : public KernelState { + explicit ListContainsState(Result> types) + : types(std::move(types)) {} + + // The result of ComparableTypes, raised for non-null queries only + Result> types; +}; + +// Compare types once, raising even if there is nothing to search, like for null lists. +// Null queries are matched without "equal", but must be comparable too unless they +// are null-typed or of the values' type. +Result> InitListContains(KernelContext* ctx, + const KernelInitArgs& args) { + const auto& values_type = + checked_cast(*args.inputs[0]).value_type(); + const auto queries_type = args.inputs[1].GetSharedPtr(); + auto types = ComparableTypes(ctx, values_type, queries_type); + if (queries_type->id() != Type::NA && !queries_type->Equals(*values_type)) { + RETURN_NOT_OK(types.status()); + } + return std::make_unique(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 +Status VisitEqual(const ArraySpan& values, const ArraySpan& queries, Visit&& visit) { + auto numbers = [&](auto zero) { + using T = decltype(zero); + const T* v = values.GetValues(1); + const T* q = queries.GetValues(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(1); + const T* qo = queries.GetValues(1); + const uint8_t* vd = values.buffers[2].data; + const uint8_t* qd = queries.buffers[2].data; + return visit([=](int64_t j, int64_t k) { + const T length = vo[j + 1] - vo[j]; + return length == qo[k + 1] - qo[k] && + std::memcmp(vd + vo[j], qd + qo[k], length) == 0; + }); + }; + switch (values.type->id()) { + case Type::BOOL: + return visit([&](int64_t j, int64_t k) { + return bit_util::GetBit(values.buffers[1].data, values.offset + j) == + bit_util::GetBit(queries.buffers[1].data, queries.offset + k); + }); + case Type::FLOAT: + return numbers(float{}); + case Type::DOUBLE: + return numbers(double{}); + case Type::BINARY: + case Type::STRING: + return binary(int32_t{}); + case Type::LARGE_BINARY: + case Type::LARGE_STRING: + return binary(int64_t{}); + default: + break; + } + // Integers, temporal types, decimals and fixed-size binary compare bitwise + switch (const int width = values.type->bit_width() / 8) { + case 0: + // Only zero-width fixed-size binary, whose values all equal each other + return visit([](int64_t, int64_t) { return true; }); + case 1: + return numbers(uint8_t{}); + case 2: + return numbers(uint16_t{}); + case 4: + return numbers(uint32_t{}); + case 8: + return numbers(uint64_t{}); + default: { + if (width < 0) { + return Status::NotImplemented("list_contains for ", *values.type); + } + const uint8_t* v = values.buffers[1].data + values.offset * width; + const uint8_t* q = queries.buffers[1].data + queries.offset * width; + return visit([=](int64_t j, int64_t k) { + return std::memcmp(v + j * width, q + k * width, width) == 0; + }); + } + } +} + +// Returns a function giving the (child offset, length) of list i. With `kGathered`, +// list views' offsets are into their gathered values instead, so the function must +// be called for i = 0, 1, ... +template +auto GetListRanges(const ArraySpan& list) { + using Range = std::pair; + if constexpr (std::is_same_v) { + const int64_t width = checked_cast(*list.type).list_size(); + return [&list, width](int64_t i) { return Range((list.offset + i) * width, width); }; + } else { + const auto* offsets = list.GetValues(1); + if constexpr (is_list_view_type::value) { + const auto* sizes = list.GetValues(2); + // Reading the validity bit directly keeps this cheap enough to be inlined + const uint8_t* validity = list.MayHaveNulls() ? list.buffers[0].data : nullptr; + auto size = [sizes, validity, offset = list.offset](int64_t i) -> int64_t { + return validity == nullptr || bit_util::GetBit(validity, offset + i) ? sizes[i] + : 0; + }; + if constexpr (kGathered) { + return [size, start = int64_t{0}](int64_t i) mutable { + const int64_t length = size(i); + start += length; + return Range(start - length, length); + }; + } else { + return [offsets, size](int64_t i) { return Range(offsets[i], size(i)); }; + } + } else { + return + [offsets](int64_t i) { return Range(offsets[i], offsets[i + 1] - offsets[i]); }; + } + } +} + +// Gather the child values referenced by list views, one view after the other +template +Result GatherListViewValues(KernelContext* ctx, const ArraySpan& list) { + auto list_range = GetListRanges(list); + Int64Builder indices(ctx->memory_pool()); + for (int64_t i = 0; i < list.length; ++i) { + const auto [start, length] = list_range(i); + RETURN_NOT_OK(indices.Reserve(length)); + for (int64_t j = start; j < start + length; ++j) { + indices.UnsafeAppend(j); + } + } + ARROW_ASSIGN_OR_RAISE(auto gather_indices, indices.Finish()); + return CallFunction("take", {list.child_data[0].ToArrayData(), gather_indices}, + ctx->exec_context()); +} + +// Search each list for its query: `matches(j, i)` tells whether child value j matches +// the query of list i, relative to the first child value referenced. If given, +// `confirm` is the exact check, and `matches` a cheaper superset of it that long lists +// are searched with first. +template +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::nullptr_t>; + auto emit = [&](auto&& found_in) { + int64_t i = 0; + arrow::internal::GenerateBitsUnrolled(out_arr->buffers[1].data, out_arr->offset, + out_arr->length, [&] { + auto [start, length] = list_range(i); + start -= values_start; + return found_in(start, start + length, i++); + }); + }; + if (values_length < 32 * num_lists) { + // Short lists are faster to check in full, without branching on the data + auto exact = [&](int64_t j, int64_t i) { + if constexpr (kConfirm) { + return confirm(j, i); + } else { + return matches(j, i); + } + }; + emit([&](int64_t start, int64_t end, int64_t i) { + bool found = false; + for (int64_t j = start; j < end; ++j) found |= exact(j, i); + return found; + }); + } else { + // Stop after the first chunk of 16 values holding a match + auto any = [](auto&& predicate, int64_t start, int64_t end, int64_t i) { + bool found = false; + for (int64_t chunk = start; chunk < end && !found; chunk += 16) { + for (int64_t j = chunk; j < std::min(chunk + 16, end); ++j) { + found |= predicate(j, i); + } + } + return found; + }; + emit([&](int64_t start, int64_t end, int64_t i) { + if constexpr (kConfirm) { + return any(matches, start, end, i) && any(confirm, start, end, i); + } else { + return any(matches, start, end, i); + } + }); + } +} + +// Look up each query in the values of a single list +Status ScalarListContains(KernelContext* ctx, + const Result>& types, + const BaseListScalar& list, const ArraySpan& query_span, + ArraySpan* out_arr) { + Datum queries = query_span.ToArrayData(); + const bool null_queries = queries.null_count() == 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.null_count() > 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 +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(*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(*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(list); + int64_t values_start = 0; + int64_t values_end = list.child_data[0].length; + if constexpr (!is_list_view_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); + bool gathered = false; + auto search = [&](auto&&... predicates) { + if constexpr (is_list_view_type::value) { + if (gathered) { + SearchLists(GetListRanges(list), list.length, /*values_start=*/0, + values.length(), out_arr, predicates...); + return; + } + } + 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.null_count() == queries.length()) { + 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 constexpr (is_list_view_type::value) { + if (!values.type()->Equals(*types[0])) { + // Only cast the span of values that views reference. Values in gaps between + // views may not convert, so then only cast the referenced values. + values_start = std::numeric_limits::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); + if (!CastTo(ctx, types[0], &values).ok()) { + ARROW_ASSIGN_OR_RAISE(values, GatherListViewValues(ctx, list)); + gathered = true; + } + } + } + RETURN_NOT_OK(CastTo(ctx, types[0], &values)); + 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(bit_util::GetBit(validity, offset + j)) & + static_cast(equal(j, i * query_step))) != 0; + }); + } else { + search([&](int64_t j, int64_t i) { return equal(j, i * query_step); }); + } + return Status::OK(); + }); +} + +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); + add_kernel(Type::LARGE_LIST, ListContains); + add_kernel(Type::LIST_VIEW, ListContains); + add_kernel(Type::LARGE_LIST_VIEW, ListContains); + add_kernel(Type::FIXED_SIZE_LIST, ListContains); +} + +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::Get(ctx); @@ -960,6 +1388,11 @@ void RegisterScalarNested(FunctionRegistry* registry) { AddListElementKernels(list_element.get()); DCHECK_OK(registry->AddFunction(std::move(list_element))); + auto list_contains = std::make_shared("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("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 b5a68d12cb0c..811a71991794 100644 --- a/cpp/src/arrow/compute/kernels/scalar_nested_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_nested_test.cc @@ -128,6 +128,371 @@ 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]"); +} + +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]"); + + // 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& 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")"); + 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]"); +} + +TEST(TestScalarNested, ListContainsChunked) { + auto array = ArrayFromJSON(list(int32()), "[[1, 2], [3], null, [], [2, null]]"); + auto input = std::make_shared( + 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 either: 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 only compare with null queries + EXPECT_THAT(CallFunction("list_contains", {ScalarFromJSON(list(list(int32())), "null"), + ArrayFromJSON(list(int32()), "[[1]]")}), + Raises(StatusCode::NotImplemented)); + // 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 + 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]"}) { + EXPECT_THAT( + CallFunction("list_contains", + {lists, ArrayFromJSON(duration(TimeUnit::SECOND), queries)}), + Raises(StatusCode::NotImplemented)); + } + } +} + using VarLenListLikeTypeFactory = std::shared_ptr (*)(std::shared_ptr); static const VarLenListLikeTypeFactory kVarLenListTypeFactories[] = { diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index f24fea20a9d8..5e29107a9f1c 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -1924,6 +1924,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) | @@ -1982,6 +1984,12 @@ Structural transforms index *n* and the type code at index *n* is 2. * The indices ``2`` and ``7`` are invalid. +* \(7) Output is true for each list containing a value equal to the second + argument. If the second argument is an array, each list is searched for the + value at the same index. 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 6a4b04468d4c..f83fce2be84e 100644 --- a/docs/source/python/api/compute.rst +++ b/docs/source/python/api/compute.rst @@ -569,6 +569,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 797fbc220ec3..042456a12439 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -3899,6 +3899,45 @@ def test_list_element(): assert result.equals(expected) +@pytest.mark.parametrize("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] + + +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"))