Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

GH-45167: [C++] Implement Compute Equals for List Types #45272

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 49 additions & 0 deletions cpp/src/arrow/compute/kernels/codegen_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,13 @@ struct GetViewType<Type, enable_if_t<is_base_binary_type<Type>::value ||
static T LogicalValue(PhysicalType value) { return value; }
};

template <typename Type>
struct GetViewType<Type, enable_if_list_type<Type>> {
using T = typename TypeTraits<Type>::ScalarType;

static T LogicalValue(T value) { return value; }
};

template <>
struct GetViewType<Decimal32Type> {
using T = Decimal32;
Expand Down Expand Up @@ -322,6 +329,26 @@ struct ArrayIterator<Type, enable_if_base_binary<Type>> {
}
};

template <typename Type>
struct ArrayIterator<Type, enable_if_list_type<Type>> {
using T = typename TypeTraits<Type>::ScalarType;
using ArrayT = typename TypeTraits<Type>::ArrayType;
using offset_type = typename Type::offset_type;

const ArraySpan& arr;
int64_t position;

explicit ArrayIterator(const ArraySpan& arr) : arr(arr), position(0) {}

T operator()() {
const auto array_ptr = arr.ToArray();
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The alternative to calling ToArray with the cast would be to implement something like value_slice on the ArraySpan directly, although I'm not sure if the ArraySpan is supposed to return anything but pointers to primitives (as is currently implemeted)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is going to be slow, so we probably want to avoid this IMHO.

You may want to run a crude benchmark from Python to check this.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about:

  1. Get offset, length
  2. Subslice the value array
  3. Build ListScalar / LargeListScalar from the child array?

Or materialize the child array, and using the sub array. arr.ToArray() every call is too expansive?

const auto array = checked_cast<const ArrayT*>(array_ptr.get());

T result{array->value_slice(position++)};
return result;
}
};

template <>
struct ArrayIterator<FixedSizeBinaryType> {
const ArraySpan& arr;
Expand Down Expand Up @@ -390,6 +417,12 @@ struct UnboxScalar<Type, enable_if_has_string_view<Type>> {
}
};

template <typename Type>
struct UnboxScalar<Type, enable_if_list_type<Type>> {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So fixed_size_list is also declared as scalar, but not being registered in compute? ( It's ok to me, just to make sure this)

using T = typename TypeTraits<Type>::ScalarType;
static const T& Unbox(const Scalar& val) { return checked_cast<const T&>(val); }
};

template <>
struct UnboxScalar<Decimal32Type> {
using T = Decimal32;
Expand Down Expand Up @@ -1383,6 +1416,22 @@ ArrayKernelExec GenerateDecimal(detail::GetTypeId get_id) {
}
}

// Generate a kernel given a templated functor for list types
//
// See "Numeric" above for description of the generator functor
template <template <typename...> class Generator, typename Type0, typename... Args>
ArrayKernelExec GenerateList(detail::GetTypeId get_id) {
switch (get_id.id) {
case Type::LIST:
return Generator<Type0, ListType, Args...>::Exec;
case Type::LARGE_LIST:
return Generator<Type0, LargeListType, Args...>::Exec;
default:
DCHECK(false);
return nullptr;
}
}

// END of kernel generator-dispatchers
// ----------------------------------------------------------------------
// BEGIN of DispatchBest helpers
Expand Down
8 changes: 8 additions & 0 deletions cpp/src/arrow/compute/kernels/scalar_compare.cc
Original file line number Diff line number Diff line change
Expand Up @@ -445,6 +445,14 @@ std::shared_ptr<ScalarFunction> MakeCompareFunction(std::string name, FunctionDo
DCHECK_OK(func->AddKernel({ty, ty}, boolean(), std::move(exec)));
}

if constexpr (std::is_same_v<Op, Equal> || std::is_same_v<Op, NotEqual>) {
for (const auto id : {Type::LIST, Type::LARGE_LIST}) {
auto exec = GenerateList<applicator::ScalarBinaryEqualTypes, BooleanType, Op>(id);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another approach with perhaps a better performance potential would be to leverage the existing RangeDataEqualsImpl in arrow/compare.cc.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the heads up - I will give that a look. So I see all of the functions right now in the compare module are registered via RegisterScalarComparison. With what you are suggesting, I'm guessing I should be creating a new registry function along with that like RegisterRangeComparison right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also I'm guessing the RangeDataEqualsImpl is supposed to work when comparing two arrays, but not when comparing an array with a scalar

FWIW though I did benchmark the current implementation and it was definitely slow. Seemed about 1000x slower than an equivalent comparison using primitive types

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With what you are suggesting, I'm guessing I should be creating a new registry function along with that like RegisterRangeComparison right?

I think we can avoid that by directly calling into RangeDataEqualsImpl.

Also I'm guessing the RangeDataEqualsImpl is supposed to work when comparing two arrays, but not when comparing an array with a scalar

A list scalar's value is actually an array, so that should not necessarily be a problem.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK took a closer look at this. So AFAICT the RangeDataEqualsImpl returns a scalar bool value, rather than an array of booleans like we would need in the result here. That class is also private to the compare.cc module and doesn't expose any suitable entrypoint in compare.h that I think would work here.

Are you thinking we should refactor the RangeDataEqualsImpl to support vector functions and move it to make it accessible to the compute module, or do you think we should just create a dedicated class drawing some inspiration from it in scalar_compute.cc?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks again for the guidance and patience here! Trying to wrap my head around the structure of the compute modules

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So AFAICT the RangeDataEqualsImpl returns a scalar bool value, rather than an array of booleans like we would need in the result here.

That's right, so it would need to be called once for each list element (which is admittedly non optimal, but probably better than using GetScalar anyway?).

That class is also private to the compare.cc module and doesn't expose any suitable entrypoint in compare.h that I think would work here.

Well, we could add a suitable entrypoint in compare_internal.h if that's useful.

Another possible approach would be to leverage the comparison kernel for the child type, but that would probably be even more involved. So that's up to how much work you want to put into this :)

DCHECK_OK(
func->AddKernel({InputType(id), InputType(id)}, boolean(), std::move(exec)));
}
}

return func;
}

Expand Down
93 changes: 93 additions & 0 deletions cpp/src/arrow/compute/kernels/scalar_compare_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -680,6 +680,99 @@ TYPED_TEST(TestCompareDecimal, DifferentParameters) {
}
}

template <typename ArrowType>
class TestCompareList : public ::testing::Test {};
TYPED_TEST_SUITE(TestCompareList, ListArrowTypes);

TYPED_TEST(TestCompareList, ArrayScalar) {
const auto int_value_typ = std::make_shared<Int32Type>();
const auto int_ty = std::make_shared<TypeParam>(std::move(int_value_typ));
const auto bin_value_typ = std::make_shared<StringType>();
const auto bin_ty = std::make_shared<TypeParam>(std::move(bin_value_typ));

const std::vector<std::pair<std::string, std::string>> cases = {
{"equal", "[1, 0, 0, null]"},
{"not_equal", "[0, 1, 1, null]"},
};
const auto lhs_int = ArrayFromJSON(int_ty, R"([[1, 2, 3], [4, 5, 6], [42], null])");
const auto lhs_bin = ArrayFromJSON(
bin_ty, R"([["a", "b", "c"], ["foo", "bar", "baz"], ["hello"], null])");
const auto rhs_int = ScalarFromJSON(int_ty, R"([1, 2, 3])");
const auto rhs_bin = ScalarFromJSON(bin_ty, R"(["a", "b", "c"])");
for (const auto& op : cases) {
const auto& function = op.first;
const auto& expected = op.second;

SCOPED_TRACE(function);
CheckScalarBinary(function, lhs_int, rhs_int, ArrayFromJSON(boolean(), expected));
CheckScalarBinary(function, lhs_bin, rhs_bin, ArrayFromJSON(boolean(), expected));
}
}

TYPED_TEST(TestCompareList, ScalarArray) {
const auto int_value_typ = std::make_shared<Int32Type>();
const auto int_ty = std::make_shared<TypeParam>(std::move(int_value_typ));
const auto bin_value_typ = std::make_shared<StringType>();
const auto bin_ty = std::make_shared<TypeParam>(std::move(bin_value_typ));

const std::vector<std::pair<std::string, std::string>> cases = {
{"equal", "[1, 0, 0, null]"},
{"not_equal", "[0, 1, 1, null]"},
};
const auto lhs_int = ScalarFromJSON(int_ty, R"([1, 2, 3])");
const auto lhs_bin = ScalarFromJSON(bin_ty, R"(["a", "b", "c"])");
const auto rhs_int = ArrayFromJSON(int_ty, R"([[1, 2, 3], [4, 5, 6], [42], null])");
const auto rhs_bin = ArrayFromJSON(
bin_ty, R"([["a", "b", "c"], ["foo", "bar"], ["baz", "hello", "world"], null])");
for (const auto& op : cases) {
const auto& function = op.first;
const auto& expected = op.second;

SCOPED_TRACE(function);
CheckScalarBinary(function, lhs_int, rhs_int, ArrayFromJSON(boolean(), expected));
CheckScalarBinary(function, lhs_bin, rhs_bin, ArrayFromJSON(boolean(), expected));
}
}

TYPED_TEST(TestCompareList, ArrayArray) {
const auto int_value_typ = std::make_shared<Int32Type>();
const auto int_ty = std::make_shared<TypeParam>(std::move(int_value_typ));
const auto bin_value_typ = std::make_shared<StringType>();
const auto bin_ty = std::make_shared<TypeParam>(std::move(bin_value_typ));

const std::vector<std::pair<std::string, std::string>> cases = {
{"equal", "[1, 0, 0, null]"},
{"not_equal", "[0, 1, 1, null]"},
};
const auto lhs_int = ArrayFromJSON(int_ty, R"([[1, 2, 3], [4, 5, 6], [7], null])");
const auto lhs_bin = ArrayFromJSON(
bin_ty, R"([["a", "b", "c"], ["foo", "bar", "baz"], ["hello"], null])");
const auto rhs_int = ArrayFromJSON(int_ty, R"([[1, 2, 3], [4, 5], [6, 7, 8], null])");
const auto rhs_bin = ArrayFromJSON(
bin_ty, R"([["a", "b", "c"], ["foo", "bar"], ["baz", "hello", "world"], null])");
for (const auto& op : cases) {
const auto& function = op.first;
const auto& expected = op.second;

SCOPED_TRACE(function);
CheckScalarBinary(function, ArrayFromJSON(int_ty, R"([])"),
ArrayFromJSON(int_ty, R"([])"), ArrayFromJSON(boolean(), "[]"));
CheckScalarBinary(function, ArrayFromJSON(int_ty, R"([null])"),
ArrayFromJSON(int_ty, R"([null])"),
ArrayFromJSON(boolean(), "[null]"));

CheckScalarBinary(function, lhs_int, rhs_int, ArrayFromJSON(boolean(), expected));

CheckScalarBinary(function, ArrayFromJSON(bin_ty, R"([])"),
ArrayFromJSON(int_ty, R"([])"), ArrayFromJSON(boolean(), "[]"));
CheckScalarBinary(function, ArrayFromJSON(int_ty, R"([null])"),
ArrayFromJSON(bin_ty, R"([null])"),
ArrayFromJSON(boolean(), "[null]"));

CheckScalarBinary(function, lhs_bin, rhs_bin, ArrayFromJSON(boolean(), expected));
}
}

// Helper to organize tests for fixed size binary comparisons
struct CompareCase {
std::shared_ptr<DataType> lhs_type;
Expand Down
Loading