diff --git a/cpp/src/arrow/array/array_test.cc b/cpp/src/arrow/array/array_test.cc index 2cdf94bfab95..838dc81be388 100644 --- a/cpp/src/arrow/array/array_test.cc +++ b/cpp/src/arrow/array/array_test.cc @@ -179,6 +179,10 @@ TEST_F(TestArray, TestNullCount) { true, true); CheckDictionaryNullCount(dict_type, "[null, true, null, false]", "[null, 1, 0, 2, 3]", 1, 3, true, true); + + // null-typed dictionary, without a validity bitmap + CheckDictionaryNullCount(dictionary(index_type, null()), "[null]", "[0, null, 0]", 1, + 3, true, true); } } diff --git a/cpp/src/arrow/compute/kernels/codegen_internal.cc b/cpp/src/arrow/compute/kernels/codegen_internal.cc index 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..41050fa30a09 100644 --- a/cpp/src/arrow/compute/kernels/scalar_nested.cc +++ b/cpp/src/arrow/compute/kernels/scalar_nested.cc @@ -18,9 +18,11 @@ // Vector kernels involving nested types #include +#include #include "arrow/array/array_base.h" #include "arrow/array/builder_nested.h" #include "arrow/compute/api_scalar.h" +#include "arrow/compute/cast.h" #include "arrow/compute/kernels/common_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/result.h" @@ -28,6 +30,7 @@ #include "arrow/util/bit_block_counter.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_generate.h" +#include "arrow/util/bitmap_ops.h" #include "arrow/util/logging_internal.h" #include "arrow/util/string.h" #include "arrow/util/unreachable.h" @@ -537,6 +540,455 @@ const FunctionDoc list_element_doc( "is emitted. Null values emit a null in the output."), {"lists", "index"}); +// Dictionaries are compared as their values, like by "equal" +const std::shared_ptr& Decoded(const std::shared_ptr& type) { + return type->id() == Type::DICTIONARY + ? checked_cast(*type).value_type() + : type; +} + +// The types "equal" compares `values` and `queries` as, raising the errors it would. +// float16, which "equal" lacks, is compared as float32, and binary views as their +// large types, which "equal" compares with other binary types unlike views. +Result> ComparableTypes( + KernelContext* ctx, const std::shared_ptr& values_type, + const std::shared_ptr& queries_type) { + std::vector types; + for (const auto& type : {Decoded(values_type), Decoded(queries_type)}) { + switch (type->id()) { + case Type::HALF_FLOAT: + types.emplace_back(float32()); + break; + case Type::STRING_VIEW: + types.emplace_back(large_utf8()); + break; + case Type::BINARY_VIEW: + types.emplace_back(large_binary()); + break; + default: + types.emplace_back(type); + } + } + ARROW_ASSIGN_OR_RAISE(auto equal, + ctx->exec_context()->func_registry()->GetFunction("equal")); + RETURN_NOT_OK(equal->DispatchBest(&types)); + RETURN_NOT_OK(ValidateComparison(*types[0], *types[1])); + if (types[0].id() == Type::FIXED_SIZE_BINARY && *types[0] != *types[1]) { + // Fixed-size binary values of different widths differ, as binary values do + types = {binary(), binary()}; + } + // Raise also where there is nothing to cast, e.g. for null queries + for (const auto& [from, to] : {std::pair(Decoded(values_type), types[0]), + std::pair(Decoded(queries_type), types[1])}) { + if (!from->Equals(*to) && !CanCast(*from, *to)) { + return Status::NotImplemented("Unsupported cast from ", *from, " to ", *to); + } + } + return types; +} + +struct ListContainsState : public KernelState { + explicit ListContainsState(Result> types) + : types(std::move(types)) {} + + // The result of ComparableTypes. An error is kept rather than raised for queries + // that are null-typed or of the values' type: all-null queries are matched by + // validity alone, which works even for values "equal" can't compare, like structs. + // The error is raised once a batch holds a non-null query. + Result> types; +}; + +// Compare types once, raising even if there is nothing to search, like for null lists. +// Null queries are matched without "equal", but must be comparable too unless they +// are null-typed or of the values' type (after decoding dictionaries). +Result> 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); + const auto& decoded_queries_type = Decoded(queries_type); + if (decoded_queries_type->id() != Type::NA && + !decoded_queries_type->Equals(*Decoded(values_type))) { + RETURN_NOT_OK(types.status()); + } + return std::make_unique(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 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); + // Reading the validity bit directly keeps this cheap enough to be inlined + const uint8_t* validity = list.MayHaveNulls() ? list.buffers[0].data : nullptr; + return [offsets, sizes, validity, offset = list.offset](int64_t i) { + return Range( + offsets[i], + validity == nullptr || bit_util::GetBit(validity, offset + i) ? sizes[i] : 0); + }; + } else { + return + [offsets](int64_t i) { return Range(offsets[i], offsets[i + 1] - offsets[i]); }; + } + } +} + +// Cast `values`, the child values from `values_start` on, to `type`. Values that no list +// references, like those of null lists, may not convert, so then they are cast as nulls. +template +Status CastListValues(KernelContext* ctx, const ArraySpan& list, int64_t values_start, + const TypeHolder& type, Datum* values) { + if (CastTo(ctx, type, values).ok()) { + return Status::OK(); + } + auto data = values->array()->Copy(); + ARROW_ASSIGN_OR_RAISE(auto validity, AllocateEmptyBitmap(data->offset + data->length, + ctx->memory_pool())); + const auto list_range = GetListRanges(list); + for (int64_t i = 0; i < list.length; ++i) { + if (list.IsValid(i)) { + const auto [start, length] = list_range(i); + bit_util::SetBitsTo(validity->mutable_data(), data->offset + start - values_start, + length, true); + } + } + if (data->MayHaveNulls()) { + arrow::internal::BitmapAnd(validity->data(), data->offset, data->buffers[0]->data(), + data->offset, data->length, data->offset, + validity->mutable_data()); + } + data->buffers[0] = std::move(validity); + data->null_count = kUnknownNullCount; + *values = std::move(data); + return CastTo(ctx, type, values); +} + +// Whether lists are searched in full or in chunks is decided once per batch, from the +// mean number of values per list: below kShortListLength, scanning all values without +// branching on the data is fastest. Longer lists are searched kSearchChunkLength values +// at a time, which compile to vector compares, stopping after the first chunk holding a +// match. Both were tuned on int64 values, with lists of 3 to 1000 values and matches at +// their start, middle or end (see GH-51611). +constexpr int64_t kShortListLength = 32; +constexpr int64_t kSearchChunkLength = 16; + +// Search each list for its query: `matches(j, i)` tells whether child value j matches +// the query of list i, relative to the first child value referenced. If given, +// `confirm` is the exact check, and `matches` a cheaper superset of it that long lists +// are searched with first. +template +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 < kShortListLength * num_lists) { + // Short lists are faster to check in full, without branching on the data + auto exact = [&](int64_t j, int64_t i) { + if constexpr (kConfirm) { + return confirm(j, i); + } else { + return matches(j, i); + } + }; + emit([&](int64_t start, int64_t end, int64_t i) { + bool found = false; + for (int64_t j = start; j < end; ++j) found |= exact(j, i); + return found; + }); + } else { + // Stop after the first chunk holding a match + auto any = [](auto&& predicate, int64_t start, int64_t end, int64_t i) { + bool found = false; + for (int64_t chunk = start; chunk < end && !found; chunk += kSearchChunkLength) { + for (int64_t j = chunk; j < std::min(chunk + kSearchChunkLength, end); ++j) { + found |= predicate(j, i); + } + } + return found; + }; + emit([&](int64_t start, int64_t end, int64_t i) { + if constexpr (kConfirm) { + return any(matches, start, end, i) && any(confirm, start, end, i); + } else { + return any(matches, start, end, i); + } + }); + } +} + +// Look up each query in the values of a single list +Status ScalarListContains(KernelContext* ctx, + const Result>& types, + const BaseListScalar& list, const ArraySpan& query_span, + ArraySpan* out_arr) { + Datum queries = query_span.ToArrayData(); + // Nulls in a dictionary of queries count too + const bool null_queries = + queries.array()->ComputeLogicalNullCount() == queries.length(); + if (!null_queries) { + // Raise even for a null list, like for list arrays + RETURN_NOT_OK(types.status()); + } + bit_util::SetBitsTo(out_arr->buffers[0].data, out_arr->offset, out_arr->length, + list.is_valid); + out_arr->null_count = list.is_valid ? 0 : out_arr->length; + if (!list.is_valid) { + return Status::OK(); + } + Datum values = list.value; + if (null_queries) { + bit_util::SetBitsTo(out_arr->buffers[1].data, out_arr->offset, out_arr->length, + values.array()->ComputeLogicalNullCount() > 0); + return Status::OK(); + } + RETURN_NOT_OK(CastTo(ctx, (*types)[0], &values)); + RETURN_NOT_OK(CastTo(ctx, (*types)[1], &queries)); + ExecContext* exec_ctx = ctx->exec_context(); + if (is_floating(values.type()->id())) { + // "is_in" tells -0.0 from 0.0 and NaNs with different payloads apart, so turn + // -0.0 into 0.0 and all NaNs into one + ARROW_ASSIGN_OR_RAISE(auto zero, MakeScalar(values.type(), 0.0)); + ARROW_ASSIGN_OR_RAISE(auto nan, MakeScalar(values.type(), std::nan(""))); + for (Datum* datum : {&values, &queries}) { + ARROW_ASSIGN_OR_RAISE(*datum, CallFunction("add", {*datum, zero}, exec_ctx)); + ARROW_ASSIGN_OR_RAISE(Datum is_nan, CallFunction("is_nan", {*datum}, exec_ctx)); + ARROW_ASSIGN_OR_RAISE(*datum, + CallFunction("if_else", {is_nan, nan, *datum}, exec_ctx)); + } + } + const SetLookupOptions options(values); + ARROW_ASSIGN_OR_RAISE(Datum found, + CallFunction("is_in", {queries}, &options, exec_ctx)); + const ArrayData& found_data = *found.array(); + arrow::internal::CopyBitmap(found_data.buffers[1]->data(), found_data.offset, + found_data.length, out_arr->buffers[1].data, + out_arr->offset); + return Status::OK(); +} + +template +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); + auto search = [&](auto&&... predicates) { + SearchLists(list_range, list.length, values_start, values.length(), out_arr, + predicates...); + }; + + // A scalar query is a single query for all lists + const bool scalar_query = batch[1].is_scalar(); + Datum queries; + if (scalar_query) { + ARROW_ASSIGN_OR_RAISE(queries, + MakeArrayFromScalar(*batch[1].scalar, 1, ctx->memory_pool())); + } else { + queries = batch[1].array.ToArrayData(); + } + const int64_t query_step = scalar_query ? 0 : 1; + + if (queries.array()->ComputeLogicalNullCount() == queries.length()) { + // Values are also null where their dictionary is + RETURN_NOT_OK(CastTo(ctx, Decoded(values.type()), &values)); + const ArraySpan v(*values.array()); + search([&](int64_t j, int64_t) { return v.IsNull(j); }); + return Status::OK(); + } + RETURN_NOT_OK(state_types.status()); + const auto& types = *state_types; + if (!values.type()->Equals(*types[0])) { + if constexpr (is_list_view_type::value) { + // Only cast the span of values that views reference + 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); + } + RETURN_NOT_OK(CastListValues(ctx, list, values_start, 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 +1412,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..ec049ade828d 100644 --- a/cpp/src/arrow/compute/kernels/scalar_nested_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_nested_test.cc @@ -128,6 +128,431 @@ TEST(TestScalarNested, ListElementInvalid) { Raises(StatusCode::Invalid)); } +void CheckListContains(Datum lists, Datum value, const std::string& expected) { + CheckScalar("list_contains", {std::move(lists), std::move(value)}, + ArrayFromJSON(boolean(), expected)); +} + +TEST(TestScalarNested, ListContains) { + auto sample = "[[7, 5, 81], [6, null, 4, 7, 8], [], [5], null, [null]]"; + for (auto ty : NumericTypes()) { + for (auto list_type : + {list(ty), large_list(ty), list_view(ty), large_list_view(ty)}) { + auto input = ArrayFromJSON(list_type, sample); + CheckListContains(input, ScalarFromJSON(ty, "5"), + "[true, false, false, true, null, false]"); + CheckListContains(input, ScalarFromJSON(ty, "7"), + "[true, true, false, false, null, false]"); + CheckListContains(input, ScalarFromJSON(ty, "null"), + "[false, true, false, false, null, true]"); + CheckListContains(ArrayFromJSON(list_type, "[]"), ScalarFromJSON(ty, "5"), "[]"); + } + } + + auto input = + ArrayFromJSON(list(utf8()), R"([["a", "b"], ["a", "c"], ["b", "c", "d"]])"); + CheckListContains(input, ScalarFromJSON(utf8(), R"("a")"), "[true, true, false]"); + CheckListContains(input->Slice(1), ScalarFromJSON(utf8(), R"("b")"), "[false, true]"); + + // All-null lists contain no non-null value + CheckListContains(ArrayFromJSON(list(int64()), "[[null, null]]"), + ScalarFromJSON(int64(), "1"), "[false]"); +} + +TEST(TestScalarNested, ListContainsLongLists) { + // Matches and nulls beyond the first bitmap words, at unaligned offsets + std::string values = "["; + for (int i = 0; i < 300; ++i) { + values += i == 0 ? "" : ", "; + values += i == 100 ? "null" : i == 290 ? "7" : "1"; + } + values += "]"; + ASSERT_OK_AND_ASSIGN(auto input, + ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 3, 150, 300]"), + *ArrayFromJSON(int32(), values))); + CheckListContains(input, ScalarFromJSON(int32(), "7"), "[false, false, true]"); + CheckListContains(input, ScalarFromJSON(int32(), "null"), "[false, true, false]"); + CheckListContains(input, ArrayFromJSON(int32(), "[7, 7, null]"), + "[false, false, false]"); + CheckListContains(input, ArrayFromJSON(int32(), "[1, null, 7]"), "[true, true, true]"); + // The null value's slot holds 0, which must not match + CheckListContains(input, ScalarFromJSON(int32(), "0"), "[false, false, false]"); +} + +TEST(TestScalarNested, ListContainsManyShortLists) { + // Short lists at every bit alignment, with null lists and null values + std::string lists = "["; + std::string contains_5 = "["; + std::string contains_null = "["; + std::string contains_0 = "["; + for (int i = 0; i < 200; ++i) { + const std::string sep = i == 0 ? "" : ", "; + const bool is_null = i % 11 == 0; + lists += sep + (is_null ? "null" + : i % 7 == 0 ? "[5, 1]" + : i % 3 == 0 ? "[null, 2, 3]" + : "[1]"); + contains_5 += sep + (is_null ? "null" : i % 7 == 0 ? "true" : "false"); + contains_null += sep + (is_null ? "null" + : i % 7 != 0 && i % 3 == 0 ? "true" + : "false"); + contains_0 += sep + (is_null ? "null" : "false"); + } + auto input = ArrayFromJSON(list(int32()), lists + "]"); + CheckListContains(input, ScalarFromJSON(int32(), "5"), contains_5 + "]"); + CheckListContains(input, ScalarFromJSON(int32(), "null"), contains_null + "]"); + // Null values' slots hold 0, which must not match + CheckListContains(input, ScalarFromJSON(int32(), "0"), contains_0 + "]"); +} + +TEST(TestScalarNested, ListContainsNull) { + // A null value matches lists holding a null + auto input = ArrayFromJSON( + list(int32()), "[[1, null], [null], [1, 2], [1, 1], [2, 1, null, 1], [], null]"); + for (auto value : {ScalarFromJSON(int32(), "null"), MakeNullScalar(null()), + ScalarFromJSON(int64(), "null")}) { + CheckListContains(input, value, "[true, true, false, false, true, false, null]"); + } + CheckListContains(ArrayFromJSON(list(int32()), "[[]]"), ScalarFromJSON(int32(), "null"), + "[false]"); + + input = ArrayFromJSON(list(utf8()), R"([["x", null], [null], ["x", "y"], [], null])"); + CheckListContains(input, ScalarFromJSON(utf8(), "null"), + "[true, true, false, false, null]"); + + input = ArrayFromJSON(list(list(int32())), + "[[[1], null], [null], [[1], [null]], [], null]"); + CheckListContains(input, MakeNullScalar(list(int32())), + "[true, true, false, false, null]"); + + input = ArrayFromJSON(list(null()), "[[null], [], null]"); + CheckListContains(input, MakeNullScalar(null()), "[true, false, null]"); + + // Dictionary values are also null where their dictionary is + ASSERT_OK_AND_ASSIGN( + input, ListArray::FromArrays( + *ArrayFromJSON(int32(), "[0, 1, 2]"), + *DictArrayFromJSON(dictionary(int8(), int64()), "[0, 1]", "[null, 1]"))); + for (auto value : {ScalarFromJSON(int64(), "null"), MakeNullScalar(null())}) { + CheckListContains(input, value, "[true, false]"); + } + CheckListContains(input, ArrayFromJSON(int64(), "[null, null]"), "[true, false]"); + ASSERT_OK_AND_ASSIGN(auto list_scalar, input->GetScalar(0)); + CheckListContains(list_scalar, ArrayFromJSON(int64(), "[null, 1]"), "[true, false]"); + // Also of a null-typed dictionary, both of values and of queries + auto null_dict = DictArrayFromJSON(dictionary(int8(), null()), "[0, 0]", "[null]"); + ASSERT_OK_AND_ASSIGN( + input, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 2, 2]"), *null_dict)); + CheckListContains(input, MakeNullScalar(null()), "[true, false]"); + ASSERT_OK_AND_ASSIGN(list_scalar, input->GetScalar(0)); + CheckListContains(list_scalar, ArrayFromJSON(int64(), "[null, null]"), "[true, true]"); + CheckListContains(ArrayFromJSON(list(struct_({field("a", int64())})), "[[null], []]"), + null_dict, "[true, false]"); +} + +TEST(TestScalarNested, ListContainsNaN) { + // Unlike "equal", a NaN value matches NaN list values + for (auto ty : {float32(), float64()}) { + auto input = ArrayFromJSON(list(ty), "[[1.5, null], [NaN], [1.5, NaN], [], null]"); + for (auto value_ty : {float32(), float64()}) { + CheckListContains(input, ScalarFromJSON(value_ty, "NaN"), + "[false, true, true, false, null]"); + CheckListContains(input, ScalarFromJSON(value_ty, "1.5"), + "[true, false, true, false, null]"); + } + } + CheckListContains(ArrayFromJSON(list(int64()), "[[1], []]"), + ScalarFromJSON(float64(), "NaN"), "[false, false]"); + + auto input = + ArrayFromJSON(list(float16()), "[[1.5, null], [NaN], [1.5, NaN], [], null]"); + CheckListContains(input, ScalarFromJSON(float16(), "NaN"), + "[false, true, true, false, null]"); + CheckListContains(ArrayFromJSON(list(float32()), "[[1.5], [NaN]]"), + ScalarFromJSON(float16(), "NaN"), "[false, true]"); + // Null queries compare float16 as float32 too + input = ArrayFromJSON(list(float16()), "[[1], [null]]"); + CheckListContains(input, ArrayFromJSON(float32(), "[null, null]"), "[false, true]"); + CheckListContains(input, ScalarFromJSON(float32(), "null"), "[false, true]"); + CheckListContains(ArrayFromJSON(list(float32()), "[[1], [null]]"), + ArrayFromJSON(float16(), "[null, null]"), "[false, true]"); + CheckListContains(ScalarFromJSON(list(float16()), "[1, null]"), + ArrayFromJSON(float32(), "[null, null]"), "[true, true]"); + // Also in dictionaries, both of values and of queries + auto f16_dict = + DictArrayFromJSON(dictionary(int8(), float16()), "[0, 1, 2]", "[NaN, 1, null]"); + ASSERT_OK_AND_ASSIGN( + input, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 2, 3]"), *f16_dict)); + CheckListContains(input, ScalarFromJSON(float32(), "NaN"), "[true, false, false]"); + CheckListContains(input, ArrayFromJSON(float32(), "[1, 1, null]"), + "[false, true, true]"); + CheckListContains(ArrayFromJSON(list(float32()), "[[NaN], [1], [null]]"), f16_dict, + "[true, true, true]"); + ASSERT_OK_AND_ASSIGN(auto f16_nan, f16_dict->GetScalar(0)); + CheckListContains(ArrayFromJSON(list(float32()), "[[NaN], [1], [null]]"), f16_nan, + "[true, false, false]"); + + // NaNs with different payloads (0x7ff8000000000001, 0x7ff8000000000002) match too + ASSERT_OK_AND_ASSIGN(auto nans, + ArrayFromJSON(uint64(), "[9221120237041090561]")->View(float64())); + ASSERT_OK_AND_ASSIGN( + auto queries, ArrayFromJSON(uint64(), "[9221120237041090562, 4607182418800017408]") + ->View(float64())); // [NaN, 1.0] + ASSERT_OK_AND_ASSIGN( + input, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 1]"), *nans)); + CheckListContains(input, queries, "[true, false]"); + ASSERT_OK_AND_ASSIGN(auto nan_query, queries->GetScalar(0)); + CheckListContains(input, nan_query, "[true, false]"); + ASSERT_OK_AND_ASSIGN(auto nan_list, input->GetScalar(0)); + CheckListContains(nan_list, queries, "[true, false]"); +} + +TEST(TestScalarNested, ListContainsValueTypes) { + auto check = [](const std::shared_ptr& ty, const std::string& x, + const std::string& y) { + auto input = ArrayFromJSON( + list(ty), "[[" + x + ", null], [" + y + "], [" + x + ", " + y + "], [], null]"); + CheckListContains(input, ScalarFromJSON(ty, x), "[true, false, true, false, null]"); + CheckListContains(input, ScalarFromJSON(ty, y), "[false, true, true, false, null]"); + CheckListContains( + input, + ArrayFromJSON(ty, "[" + x + ", " + y + ", " + y + ", " + x + ", " + x + "]"), + "[true, true, true, false, null]"); + // A single list is searched for each query + CheckListContains(ScalarFromJSON(list(ty), "[" + x + ", null]"), + ArrayFromJSON(ty, "[" + x + ", " + y + ", null]"), + "[true, false, true]"); + }; + check(utf8(), R"("x")", R"("y")"); + check(large_binary(), R"("x")", R"("y")"); + check(binary_view(), R"("x")", R"("y")"); + check(fixed_size_binary(1), R"("x")", R"("y")"); + // String views compare with other string layouts, unlike in "equal" + CheckListContains(ArrayFromJSON(list(utf8()), R"([["x"], ["y"]])"), + ScalarFromJSON(utf8_view(), R"("x")"), "[true, false]"); + CheckListContains(ArrayFromJSON(list(binary_view()), R"([["x"], ["y"]])"), + ArrayFromJSON(large_binary(), R"(["y", "y"])"), "[false, true]"); + check(boolean(), "true", "false"); + check(int16(), "1", "2"); + check(float16(), "1.5", "2.5"); + check(date32(), "18262", "18628"); + check(decimal128(38, 2), R"("1.50")", R"("2.50")"); + check(decimal256(40, 2), R"("1.50")", R"("2.50")"); + check(timestamp(TimeUnit::NANO), "1", "2"); + check(duration(TimeUnit::MILLI), "1", "2"); + + // Zero-width values all equal each other + auto input = + ArrayFromJSON(list(fixed_size_binary(0)), R"([["", null], [null], [], null])"); + CheckListContains(input, ScalarFromJSON(fixed_size_binary(0), R"("")"), + "[true, false, false, null]"); + CheckListContains(input, ArrayFromJSON(fixed_size_binary(0), R"(["", "", "", ""])"), + "[true, false, false, null]"); + CheckListContains(ScalarFromJSON(list(fixed_size_binary(0)), R"(["", null])"), + ArrayFromJSON(fixed_size_binary(0), R"(["", null])"), "[true, true]"); + + // Values of other widths never match, even with a common prefix + for (const auto& list_type : + {list(fixed_size_binary(1)), list_view(fixed_size_binary(1))}) { + input = ArrayFromJSON(list_type, R"([["a"], ["b", null], null])"); + CheckListContains(input, ScalarFromJSON(fixed_size_binary(2), R"("ab")"), + "[false, false, null]"); + CheckListContains(input, ArrayFromJSON(fixed_size_binary(2), R"(["ab", null, "ab"])"), + "[false, true, null]"); + } + CheckListContains(ScalarFromJSON(list(fixed_size_binary(2)), R"(["ab", null])"), + ArrayFromJSON(fixed_size_binary(1), R"(["a", null])"), + "[false, true]"); +} + +TEST(TestScalarNested, ListContainsSignedZero) { + // Like "equal", -0.0 equals 0.0 + auto input = ArrayFromJSON(list(float64()), "[[0.0], [-0.0], [1.0], null]"); + CheckListContains(input, ScalarFromJSON(float64(), "-0.0"), + "[true, true, false, null]"); + CheckListContains(input, ArrayFromJSON(float64(), "[-0.0, 0.0, 0.0, 0.0]"), + "[true, true, false, null]"); + CheckListContains(ScalarFromJSON(list(float64()), "[-0.0, NaN]"), + ArrayFromJSON(float64(), "[0.0, -0.0, NaN, 1.0]"), + "[true, true, true, false]"); +} + +TEST(TestScalarNested, ListContainsImplicitCast) { + auto input = ArrayFromJSON(list(int64()), + "[[2, 2, 3, null, null], null, [], [null], [1], [0, -1]]"); + CheckListContains(input, ScalarFromJSON(float64(), "2.0"), + "[true, null, false, false, false, false]"); + CheckListContains(input, ScalarFromJSON(float64(), "1.5"), + "[false, null, false, false, false, false]"); + + // No overflow of the list values + CheckListContains(ArrayFromJSON(list(int8()), "[[44, 1], null, [], [null]]"), + ScalarFromJSON(int64(), "300"), "[false, null, false, false]"); + CheckListContains(ArrayFromJSON(list(float64()), "[[1.0, 2.0], [3.0], null, []]"), + ScalarFromJSON(int64(), "1"), "[true, false, null, false]"); + CheckListContains(ArrayFromJSON(list(decimal128(38, 2)), R"([["1.50"], ["2.50"]])"), + ScalarFromJSON(decimal128(3, 2), R"("1.50")"), "[true, false]"); + + // Values of null lists are not cast: 2^53 + 1 isn't a float64 + auto values = ArrayFromJSON(int64(), "[9007199254740993, 1]"); + auto validity = ArrayFromJSON(boolean(), "[false, true]")->data()->buffers[1]; + ASSERT_OK_AND_ASSIGN( + auto lists, ListArray::FromArrays(*ArrayFromJSON(int32(), "[0, 1, 2]"), *values, + default_memory_pool(), validity)); + ASSERT_OK_AND_ASSIGN(auto fixed_size_lists, + FixedSizeListArray::FromArrays(values, 1, validity)); + for (Datum input : {Datum(lists), Datum(fixed_size_lists)}) { + CheckListContains(input, ScalarFromJSON(float64(), "1"), "[null, true]"); + CheckListContains(input, ArrayFromJSON(float64(), "[1, null]"), "[null, false]"); + } + // Unlike values of non-null lists + EXPECT_THAT(CallFunction("list_contains", + {ArrayFromJSON(list(int64()), "[[9007199254740993], null]"), + ScalarFromJSON(float64(), "1")}), + Raises(StatusCode::Invalid)); +} + +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: 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, + // also in dictionaries + auto timestamps = list(timestamp(TimeUnit::SECOND)); + for (Datum lists : {Datum(ArrayFromJSON(timestamps, "[[1], null]")), + Datum(ScalarFromJSON(timestamps, "null"))}) { + for (auto queries : {"[1, null]", "[null, null]"}) { + auto durations = ArrayFromJSON(duration(TimeUnit::SECOND), queries); + ASSERT_OK_AND_ASSIGN(Datum encoded, DictionaryEncode(durations)); + for (Datum value : {Datum(durations), encoded}) { + EXPECT_THAT(CallFunction("list_contains", {lists, value}), + Raises(StatusCode::NotImplemented)); + } + } + } +} + using VarLenListLikeTypeFactory = std::shared_ptr (*)(std::shared_ptr); static const VarLenListLikeTypeFactory kVarLenListTypeFactories[] = { diff --git a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc index 0a041e7e47f7..b72ad9cbe85f 100644 --- a/cpp/src/arrow/compute/kernels/scalar_validity_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_validity_test.cc @@ -202,6 +202,13 @@ TEST(TestValidityKernels, IsValidIsNullNullType) { ArrayFromJSON(boolean(), "[false, false, false, false, false]")); } +TEST(TestValidityKernels, IsValidIsNullDictionaryOfNullType) { + // The dictionary has no validity bitmap + auto arr = DictArrayFromJSON(dictionary(int8(), null()), "[0, null, 0]", "[null]"); + CheckScalarUnary("is_null", arr, ArrayFromJSON(boolean(), "[true, true, true]")); + CheckScalarUnary("is_valid", arr, ArrayFromJSON(boolean(), "[false, false, false]")); +} + TEST(TestValidityKernels, IsNullSetsZeroNullCount) { auto arr = ArrayFromJSON(int32(), "[1, 2, 3, 4, null]"); ASSERT_OK_AND_ASSIGN(Datum out, IsNull(arr)); diff --git a/cpp/src/arrow/util/dict_util.cc b/cpp/src/arrow/util/dict_util.cc index abcfd240e67f..00f72746604b 100644 --- a/cpp/src/arrow/util/dict_util.cc +++ b/cpp/src/arrow/util/dict_util.cc @@ -46,7 +46,8 @@ int64_t LogicalNullCount(const ArraySpan& span) { // in the dictionary array to see if they are null for (int64_t i = position; i < position + length; i++) { CType index = indices_data[i]; - if (!bit_util::GetBit(dictionary_valid_bitmap, + if (!dictionary_valid_bitmap || + !bit_util::GetBit(dictionary_valid_bitmap, index + dictionary_span.offset)) { null_count++; } @@ -77,7 +78,8 @@ void SetLogicalNullBits(const ArraySpan& span, uint8_t* out_bitmap, int64_t out_ // in the dictionary array to see if they are null for (int64_t i = position; i < position + length; i++) { CType index = indices_data[i]; - bool is_null = !bit_util::GetBit(dictionary_valid_bitmap, + bool is_null = !dictionary_valid_bitmap || + !bit_util::GetBit(dictionary_valid_bitmap, index + dictionary_span.offset); bit_util::SetBitTo(out_bitmap, out_offset + i, is_null == set_on_null); } diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index f24fea20a9d8..05566259218f 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,14 @@ Structural transforms index *n* and the type code at index *n* is 2. * The indices ``2`` and ``7`` are invalid. +* \(7) Output is true for each list containing a value equal to the second + argument. If the second argument is an array, each list is searched for the + value at the same index, and a scalar list is searched for each value: the + list ``[1, 2]`` and the array ``[2, 3]`` give ``[true, false]``. Unlike + ``equal``, a null second argument matches null list values and a NaN second + argument matches NaN list values; otherwise null list values never match. + Null lists emit a null in the output. + .. _cpp-compute-vector-replace-functions: Replace functions diff --git a/docs/source/python/api/compute.rst b/docs/source/python/api/compute.rst index 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..da0396cf35f2 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -3899,6 +3899,114 @@ 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] + + +@pytest.mark.parametrize("value_type, values, query_type, queries", [ + # 2**53 + 1 doesn't cast to float64, so it's only in null lists + (pa.int64(), [0, 1, -1, 2**53 + 1], pa.float64(), [0.0, 1.0, 1.5]), + (pa.float32(), [0.0, -0.0, 1.5, math.nan], pa.float64(), [0.0, 1.5, math.nan]), + (pa.string(), ["", "a", "ab"], pa.large_string(), ["", "a", "b"]), + (pa.string_view(), ["", "a", "ab"], pa.string(), ["", "a", "b"]), + (pa.binary(1), [b"a", b"b"], pa.binary(2), [b"ab"]), + # Dictionary values are null for null indices and for nulls in the dictionary + (pa.dictionary(pa.int8(), pa.int64()), [0, 1, None], pa.int64(), [0, 1, 2]), + (pa.dictionary(pa.int8(), pa.float16()), [0.0, math.nan, None], pa.float32(), + [0.0, 1.0, math.nan]), +]) +def test_list_contains_matches_decoded(value_type, values, query_type, queries): + # Compare list_contains to searching the decoded lists, for each input form + def search(lists, query): + if lists is None: + return None + if query is None: + return None in lists + if query != query: + return any(v != v for v in lists if v is not None) + return any(v == query for v in lists if v is not None) + + rng = random.Random(42) + n = 100 + sizes = [rng.randint(0, 4) for _ in range(n)] + offsets = list(itertools.accumulate(sizes, initial=0)) + mask = [rng.random() < 0.1 for _ in range(n)] + hidden = set() + if pa.types.is_dictionary(value_type): + indices = [rng.choice([0, 1, 2, None]) for _ in range(offsets[-1])] + child = pa.DictionaryArray.from_arrays( + pa.array(indices, value_type.index_type), + pa.array(values, value_type.value_type)) + else: + child = [rng.choice(values[:3] + [None]) for _ in range(offsets[-1])] + for i in itertools.compress(range(n), mask): + child[offsets[i]:offsets[i + 1]] = values[-1:] * sizes[i] + hidden.update(range(offsets[i], offsets[i + 1])) + child = pa.array(child, value_type) + lists = pa.ListArray.from_arrays(offsets, child, mask=pa.array(mask)) + forms = [ + lists, lists.slice(7, 50), lists.cast(pa.large_list(value_type)), + # Views of the same lists in reverse order + pa.ListViewArray.from_arrays(offsets[-2::-1], sizes[::-1], child, + mask=pa.array(mask[::-1])), + pa.FixedSizeListArray.from_arrays( + child.slice(0, n), 1, mask=pa.array([i in hidden for i in range(n)])), + ] + [pa.chunked_array([lists.slice(i, size) for i in range(0, n, size)]) + for size in (1, 7, n)] + query_array = pa.array([rng.choice(queries + [None]) for _ in range(n)], + query_type) + for form in forms: + decoded = form.to_pylist() + for query in queries + [None]: + result = pc.list_contains(form, pa.scalar(query, query_type)) + assert result.to_pylist() == [search(x, query) for x in decoded] + expected = [search(x, q) for x, q in zip(decoded, query_array.to_pylist())] + q = query_array.slice(0, len(form)) + assert pc.list_contains(form, q).to_pylist() == expected + assert pc.list_contains(form, pa.chunked_array( + [q.slice(0, 3), q.slice(3)])).to_pylist() == expected + assert pc.list_contains(form, q.dictionary_encode()).to_pylist() == expected + for x in lists[:10]: + result = pc.list_contains(x, query_array) + assert result.to_pylist() == [search(x.as_py(), q) + for q in query_array.to_pylist()] + + +def test_list_contains_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"))