diff --git a/cpp/src/arrow/compute/kernels/scalar_nested.cc b/cpp/src/arrow/compute/kernels/scalar_nested.cc index e9c65aff1ce1..7525f7c5787e 100644 --- a/cpp/src/arrow/compute/kernels/scalar_nested.cc +++ b/cpp/src/arrow/compute/kernels/scalar_nested.cc @@ -21,6 +21,7 @@ #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 +29,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 +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& values_type, + const std::shared_ptr& 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 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 +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 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 +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); + 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 +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::nullptr_t>; + const auto list_range = GetListRanges(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 +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(*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::value) { + const auto list_range = GetListRanges(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); + + // 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()) { + RETURN_NOT_OK(CheckNullQueries(ctx, values, queries)); + const ArraySpan v(*values.array()); + SearchLists(list, values_start, values_length, out_arr, + [&](int64_t j, int64_t) { return v.IsNull(j); }); + return Status::OK(); + } + RETURN_NOT_OK(CastToComparable(ctx, &values, &queries)); + const ArraySpan v(*values.array()); + const ArraySpan q(*queries.array()); + return VisitEqual(v, q, [&](auto&& equal) { + auto search = [&](auto&&... predicates) { + SearchLists(list, values_start, values_length, out_arr, predicates...); + }; + // 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); + // 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 +1288,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..c7d88f53815e 100644 --- a/cpp/src/arrow/compute/kernels/scalar_nested_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_nested_test.cc @@ -128,6 +128,287 @@ 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]"); +} + +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"); +} + +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]"); +} + +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)); + } + // 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)); + } +} + 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"))