GH-33295: [C++][Python] Add list_contains compute function

cpp/src/arrow/compute/kernels/scalar_nested.cc

@@ -21,13 +21,15 @@
 #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"
 #include "arrow/type_fwd.h"
 #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<DataType>& values_type,
+                       const std::shared_ptr<DataType>& queries_type) {
+  ARROW_ASSIGN_OR_RAISE(auto values, MakeArrayOfNull(values_type, 1, ctx->memory_pool()));
+  ARROW_ASSIGN_OR_RAISE(auto queries,
+                        MakeArrayOfNull(queries_type, 1, ctx->memory_pool()));
+  return CallFunction("equal", {values, queries}, ctx->exec_context()).status();
+}
+
+// Null queries are matched without "equal", but must be comparable all the same
+Status CheckNullQueries(KernelContext* ctx, const Datum& values, const Datum& queries) {
+  if (queries.type()->id() == Type::NA || queries.type()->Equals(*values.type())) {
+    return Status::OK();
+  }
+  return CheckComparable(ctx, values.type(), queries.type());
+}
+
+// Cast `values` and `queries` to the types "equal" compares them as, raising the
+// errors it would. float16, which "equal" lacks, is compared as float32, and binary
+// views as large binary.
+Status CastToComparable(KernelContext* ctx, Datum* values, Datum* queries) {
+  ExecContext* exec_ctx = ctx->exec_context();
+  for (Datum* datum : {values, queries}) {
+    if (datum->type()->id() == Type::HALF_FLOAT) {
+      ARROW_ASSIGN_OR_RAISE(*datum,
+                            Cast(*datum, float32(), CastOptions::Safe(), exec_ctx));
+    }
+  }
+  RETURN_NOT_OK(CheckComparable(ctx, values->type(), queries->type()));
+  ARROW_ASSIGN_OR_RAISE(auto equal, exec_ctx->func_registry()->GetFunction("equal"));
+  std::vector<TypeHolder> types{values->type(), queries->type()};
+  RETURN_NOT_OK(equal->DispatchBest(&types));
+  for (Datum* datum : {values, queries}) {
+    TypeHolder type = types[datum == values ? 0 : 1];
+    if (is_binary_view_like(type.id())) {
+      type = large_binary();
+    }
+    if (!datum->type()->Equals(*type)) {
+      ARROW_ASSIGN_OR_RAISE(*datum, Cast(*datum, type, CastOptions::Safe(), exec_ctx));
+    }
+  }
+  return Status::OK();
+}
+
+// Calls `visit` with a function telling whether value j equals query k, both of the
+// same type. Unlike "equal", NaN equals NaN.
+template <typename Visit>
+Status VisitEqual(const ArraySpan& values, const ArraySpan& queries, Visit&& visit) {
+  auto numbers = [&](auto zero) {
+    using T = decltype(zero);
+    const T* v = values.GetValues<T>(1);
+    const T* q = queries.GetValues<T>(1);
+    // `x != x` only holds for NaN
+    return visit([v, q](int64_t j, int64_t k) {
+      return v[j] == q[k] || (q[k] != q[k] && v[j] != v[j]);
+    });
+  };
+  auto binary = [&](auto zero) {
+    using T = decltype(zero);
+    const T* vo = values.GetValues<T>(1);
+    const T* qo = queries.GetValues<T>(1);
+    const uint8_t* vd = values.buffers[2].data;
+    const uint8_t* qd = queries.buffers[2].data;
+    return visit([=](int64_t j, int64_t k) {
+      const T length = vo[j + 1] - vo[j];
+      return length == qo[k + 1] - qo[k] &&
+             std::memcmp(vd + vo[j], qd + qo[k], length) == 0;
+    });
+  };
+  switch (values.type->id()) {
+    case Type::BOOL:
+      return visit([&](int64_t j, int64_t k) {
+        return bit_util::GetBit(values.buffers[1].data, values.offset + j) ==
+               bit_util::GetBit(queries.buffers[1].data, queries.offset + k);
+      });
+    case Type::FLOAT:
+      return numbers(float{});
+    case Type::DOUBLE:
+      return numbers(double{});
+    case Type::BINARY:
+    case Type::STRING:
+      return binary(int32_t{});
+    case Type::LARGE_BINARY:
+    case Type::LARGE_STRING:
+      return binary(int64_t{});
+    default:
+      break;
+  }
+  // Integers, temporal types, decimals and fixed-size binary compare bitwise
+  switch (const int width = values.type->bit_width() / 8) {
+    case 1:
+      return numbers(uint8_t{});
+    case 2:
+      return numbers(uint16_t{});
+    case 4:
+      return numbers(uint32_t{});
+    case 8:
+      return numbers(uint64_t{});
+    default: {
+      if (width <= 0) {
+        return Status::NotImplemented("list_contains for ", *values.type);
+      }
+      const uint8_t* v = values.buffers[1].data + values.offset * width;
+      const uint8_t* q = queries.buffers[1].data + queries.offset * width;
+      return visit([=](int64_t j, int64_t k) {
+        return std::memcmp(v + j * width, q + k * width, width) == 0;
+      });
+    }
+  }
+}
+
+// Returns a function giving the (child offset, length) of list i, with no values for
+// null list views
+template <typename Type>
+auto GetListRanges(const ArraySpan& list) {
+  using Range = std::pair<int64_t, int64_t>;
+  if constexpr (std::is_same_v<Type, FixedSizeListType>) {
+    const int64_t width = checked_cast<const FixedSizeListType&>(*list.type).list_size();
+    return [&list, width](int64_t i) { return Range((list.offset + i) * width, width); };
+  } else {
+    const auto* offsets = list.GetValues<typename Type::offset_type>(1);
+    if constexpr (is_list_view_type<Type>::value) {
+      const auto* sizes = list.GetValues<typename Type::offset_type>(2);
+      return [&list, offsets, sizes](int64_t i) {
+        return Range(offsets[i], list.IsValid(i) ? sizes[i] : 0);
+      };
+    } else {
+      return
+          [offsets](int64_t i) { return Range(offsets[i], offsets[i + 1] - offsets[i]); };
+    }
+  }
+}
+
+// Search each list for its query: `matches(j, i)` tells whether child value j matches
+// the query of list i, relative to the first child value referenced. If given,
+// `confirm` is the exact check, and `matches` a cheaper superset of it that long lists
+// are searched with first.
+template <typename Type, typename Matches, typename Confirm = std::nullptr_t>
+void SearchLists(const ArraySpan& list, int64_t values_start, int64_t values_length,
+                 ArraySpan* out_arr, Matches&& matches, Confirm&& confirm = nullptr) {
+  constexpr bool kConfirm = !std::is_same_v<std::decay_t<Confirm>, std::nullptr_t>;
+  const auto list_range = GetListRanges<Type>(list);
+  auto emit = [&](auto&& found_in) {
+    int64_t i = 0;
+    arrow::internal::GenerateBitsUnrolled(out_arr->buffers[1].data, out_arr->offset,
+                                          out_arr->length, [&] {
+                                            auto [start, length] = list_range(i);
+                                            start -= values_start;
+                                            return found_in(start, start + length, i++);
+                                          });
+  };
+  if (values_length < 32 * list.length) {
+    // Short lists are faster to check in full, without branching on the data
+    auto exact = [&](int64_t j, int64_t i) {
+      if constexpr (kConfirm) {
+        return confirm(j, i);
+      } else {
+        return matches(j, i);
+      }
+    };
+    emit([&](int64_t start, int64_t end, int64_t i) {
+      bool found = false;
+      for (int64_t j = start; j < end; ++j) found |= exact(j, i);
+      return found;
+    });
+  } else {
+    // Stop after the first chunk of 16 values holding a match
+    auto any = [](auto&& predicate, int64_t start, int64_t end, int64_t i) {
+      bool found = false;
+      for (int64_t chunk = start; chunk < end && !found; chunk += 16) {
+        for (int64_t j = chunk; j < std::min(chunk + 16, end); ++j) {
+          found |= predicate(j, i);
+        }
+      }
+      return found;
+    };
+    emit([&](int64_t start, int64_t end, int64_t i) {
+      if constexpr (kConfirm) {
+        return any(matches, start, end, i) && any(confirm, start, end, i);
+      } else {
+        return any(matches, start, end, i);
+      }
+    });
+  }
+}
+
+// Look up each query in the values of a single list
+Status ScalarListContains(KernelContext* ctx, const BaseListScalar& list,
+                          const ArraySpan& query_span, ArraySpan* out_arr) {
+  bit_util::SetBitsTo(out_arr->buffers[0].data, out_arr->offset, out_arr->length,
+                      list.is_valid);
+  out_arr->null_count = list.is_valid ? 0 : out_arr->length;
+  if (!list.is_valid) {
+    return Status::OK();
+  }
+  Datum values = list.value;
+  Datum queries = query_span.ToArrayData();
+  if (queries.null_count() == queries.length()) {
+    RETURN_NOT_OK(CheckNullQueries(ctx, values, queries));
+    bit_util::SetBitsTo(out_arr->buffers[1].data, out_arr->offset, out_arr->length,
+                        values.null_count() > 0);
+    return Status::OK();
+  }
+  RETURN_NOT_OK(CastToComparable(ctx, &values, &queries));
+  ExecContext* exec_ctx = ctx->exec_context();
+  if (is_floating(values.type()->id())) {
+    // "is_in" tells -0.0 from 0.0, so turn -0.0 into 0.0
+    ARROW_ASSIGN_OR_RAISE(auto zero, MakeScalar(values.type(), 0.0));
+    for (Datum* datum : {&values, &queries}) {
+      ARROW_ASSIGN_OR_RAISE(*datum, CallFunction("add", {*datum, zero}, exec_ctx));
+    }
+  }
+  const SetLookupOptions options(values);
+  ARROW_ASSIGN_OR_RAISE(Datum found,
+                        CallFunction("is_in", {queries}, &options, exec_ctx));
+  const ArrayData& found_data = *found.array();
+  arrow::internal::CopyBitmap(found_data.buffers[1]->data(), found_data.offset,
+                              found_data.length, out_arr->buffers[1].data,
+                              out_arr->offset);
+  return Status::OK();
+}
+
+template <typename Type>
+Status ListContains(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
+  ArraySpan* out_arr = out->array_span_mutable();
+  if (batch.length == 0) {
+    return Status::OK();
+  }
+  if (batch[0].is_scalar()) {
+    // A scalar list is only paired with an array of queries here
+    return ScalarListContains(ctx, checked_cast<const BaseListScalar&>(*batch[0].scalar),
+                              batch[1].array, out_arr);
+  }
+  const ArraySpan& list = batch[0].array;
+
+  // Only null lists emit a null
+  if (list.MayHaveNulls()) {
+    arrow::internal::CopyBitmap(list.buffers[0].data, list.offset, list.length,
+                                out_arr->buffers[0].data, out_arr->offset);
+  } else {
+    bit_util::SetBitsTo(out_arr->buffers[0].data, out_arr->offset, out_arr->length, true);
+  }
+  out_arr->null_count = list.null_count;
+
+  // The child values referenced by the lists, in any order for list views
+  int64_t values_start = 0;
+  int64_t values_length = list.child_data[0].length;
+  if constexpr (!is_list_view_type<Type>::value) {
+    const auto list_range = GetListRanges<Type>(list);
+    values_start = list_range(0).first;
+    values_length = list_range(list.length - 1).first +
+                    list_range(list.length - 1).second - values_start;
+  }
+  Datum values = list.child_data[0].ToArrayData()->Slice(values_start, values_length);
+
+  // 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<Type>(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<Type>(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<int>(bit_util::GetBit(validity, offset + j)) &
+                       static_cast<int>(equal(j, i * query_step))) != 0;
+             });
+    } else {
+      search([&](int64_t j, int64_t i) { return equal(j, i * query_step); });
+    }
+    return Status::OK();
+  });
+}
+
+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<ListType>);
+  add_kernel(Type::LARGE_LIST, ListContains<LargeListType>);
+  add_kernel(Type::LIST_VIEW, ListContains<ListViewType>);
+  add_kernel(Type::LARGE_LIST_VIEW, ListContains<LargeListViewType>);
+  add_kernel(Type::FIXED_SIZE_LIST, ListContains<FixedSizeListType>);
+}
+
+const FunctionDoc list_contains_doc(
+    "Check whether lists contain a given value",
+    ("`lists` must have a list-like type and `value` must be comparable\n"
+     "with the list value type.\n"
+     "For each list in `lists`, true is emitted if any of its values is equal\n"
+     "to the corresponding `value`, false otherwise. A null `value` matches\n"
+     "null list values, and a NaN `value` matches NaN list values; otherwise\n"
+     "null list values never match. Null lists emit a null in the output."),
+    {"lists", "value"});
+
 struct StructFieldFunctor {
   static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
     const auto& options = OptionsWrapper<StructFieldOptions>::Get(ctx);
@@ -960,6 +1288,11 @@ void RegisterScalarNested(FunctionRegistry* registry) {
   AddListElementKernels(list_element.get());
   DCHECK_OK(registry->AddFunction(std::move(list_element)));
 
+  auto list_contains = std::make_shared<ScalarFunction>("list_contains", Arity::Binary(),
+                                                        list_contains_doc);
+  AddListContainsKernels(list_contains.get());
+  DCHECK_OK(registry->AddFunction(std::move(list_contains)));
+
   auto list_slice =
       std::make_shared<ScalarFunction>("list_slice", Arity::Unary(), list_slice_doc);
   AddListSliceKernels(list_slice.get());