Skip to content

ARROW-13989: [C++] Add support for month-day-nano interval to compute functions #11525

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.

Sign up for GitHub

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

Closed
lidavidm wants to merge 2 commits into from
Closed

ARROW-13989: [C++] Add support for month-day-nano interval to compute functions #11525

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8d938a8
ARROW-13989: [C++] Add support for month-day-nano interval to compute…
lidavidm Oct 22, 2021
39347c3
ARROW-13989: [C++] Adjust replace_with_mask tests
lidavidm Nov 9, 2021
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
2 changes: 0 additions & 2 deletions cpp/src/arrow/compute/kernels/aggregate_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1384,8 +1384,6 @@ template <typename ArrowType>
class TestFloatingMinMaxKernel : public TestPrimitiveMinMaxKernel<ArrowType> {};

class TestBooleanMinMaxKernel : public TestPrimitiveMinMaxKernel<BooleanType> {};
class TestDayTimeIntervalMinMaxKernel
: public TestPrimitiveMinMaxKernel<DayTimeIntervalType> {};

TEST_F(TestBooleanMinMaxKernel, Basics) {
ScalarAggregateOptions options;
Expand Down
2 changes: 2 additions & 0 deletions cpp/src/arrow/compute/kernels/codegen_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1225,6 +1225,8 @@ ArrayKernelExec GenerateTypeAgnosticPrimitive(detail::GetTypeId get_id) {
case Type::DURATION:
case Type::INTERVAL_DAY_TIME:
return Generator<UInt64Type, Args...>::Exec;
case Type::INTERVAL_MONTH_DAY_NANO:
return Generator<MonthDayNanoIntervalType, Args...>::Exec;
default:
DCHECK(false);
return ExecFail;
Expand Down
3 changes: 2 additions & 1 deletion cpp/src/arrow/compute/kernels/hash_aggregate_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,8 @@ TEST(Grouper, SupportedKeys) {
ASSERT_OK(internal::Grouper::Make({timestamp(unit), duration(unit)}));
}

ASSERT_OK(internal::Grouper::Make({day_time_interval(), month_interval()}));
ASSERT_OK(internal::Grouper::Make(
{day_time_interval(), month_interval(), month_day_nano_interval()}));

ASSERT_RAISES(NotImplemented, internal::Grouper::Make({struct_({field("", int64())})}));

Expand Down
7 changes: 5 additions & 2 deletions cpp/src/arrow/compute/kernels/scalar_if_else.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -375,7 +375,9 @@ struct IfElseFunctor {};
// internal::GenerateTypeAgnosticPrimitive forwards types to the corresponding unsigned
// int type
template <typename Type>
struct IfElseFunctor<Type, enable_if_number<Type>> {
struct IfElseFunctor<Type,
enable_if_t<is_number_type<Type>::value ||
std::is_same<Type, MonthDayNanoIntervalType>::value>> {
using T = typename TypeTraits<Type>::CType;
// A - Array, S - Scalar, X = Array/Scalar

Expand Down Expand Up @@ -1295,7 +1297,8 @@ struct CopyFixedWidth<BooleanType> {
};

template <typename Type>
struct CopyFixedWidth<Type, enable_if_number<Type>> {
struct CopyFixedWidth<
Type, enable_if_t<is_number_type<Type>::value || is_interval_type<Type>::value>> {
Copy link
Member

@pitrou pitrou Nov 9, 2021

Choose a reason for hiding this comment

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

Out of curiosity, why is_interval_type here while IfElseFunctor only allows month-day-nano intervals?

Copy link
Member Author

@lidavidm lidavidm Nov 9, 2021

Choose a reason for hiding this comment

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

This helper is used by coalesce/choose/case_when so this is part of adding support for all of those functions.

if_else should support all the interval types. It uses GenerateTypeAgnosticPrimitive, so DayTime and Month intervals are handled as integers.

using CType = typename TypeTraits<Type>::CType;
static void CopyScalar(const Scalar& scalar, const int64_t length,
uint8_t* raw_out_values, const int64_t out_offset) {
Expand Down
77 changes: 77 additions & 0 deletions cpp/src/arrow/compute/kernels/scalar_if_else_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,6 +345,15 @@ TEST_F(TestIfElseKernel, DayTimeInterval) {
ArrayFromJSON(ty, "[[1, 2], [3, -4], [-5, 6], [15, 16]]"));
}

TEST_F(TestIfElseKernel, MonthDayNanoInterval) {
auto ty = month_day_nano_interval();
CheckWithDifferentShapes(
ArrayFromJSON(boolean(), "[true, true, true, false]"),
ArrayFromJSON(ty, "[[1, 2, -3], [3, -4, 5], [-5, 6, 7], [-7, -8, -9]]"),
ArrayFromJSON(ty, "[[-9, -10, 11], [11, -12, 0], [-13, 14, -1], [15, 16, 2]]"),
ArrayFromJSON(ty, "[[1, 2, -3], [3, -4, 5], [-5, 6, 7], [15, 16, 2]]"));
}

TEST_F(TestIfElseKernel, IfElseDispatchBest) {
std::string name = "if_else";
ASSERT_OK_AND_ASSIGN(auto function, GetFunctionRegistry()->GetFunction(name));
Expand Down Expand Up @@ -1368,6 +1377,22 @@ TEST(TestCaseWhen, DayTimeInterval) {
ArrayFromJSON(type, "[null, null, null, [6, 6]]"));
}

TEST(TestCaseWhen, MonthDayNanoInterval) {
auto type = month_day_nano_interval();
auto cond1 = ArrayFromJSON(boolean(), "[true, true, null, null]");
auto cond2 = ArrayFromJSON(boolean(), "[true, false, true, null]");
auto values_null = ArrayFromJSON(type, "[null, null, null, null]");
auto values1 = ArrayFromJSON(type, R"([[0, 1, -2], null, [-3, 4, 5], [-6, -7, -8]])");
auto values2 = ArrayFromJSON(type, R"([[1, 2, 3], [4, 5, 6], null, [0, 2, 4]])");

CheckScalar("case_when", {MakeStruct({cond1, cond2}), values1, values2},
ArrayFromJSON(type, R"([[0, 1, -2], null, null, null])"));
CheckScalar("case_when", {MakeStruct({cond1, cond2}), values1, values2, values1},
ArrayFromJSON(type, R"([[0, 1, -2], null, null, [-6, -7, -8]])"));
CheckScalar("case_when", {MakeStruct({cond1, cond2}), values_null, values2, values1},
ArrayFromJSON(type, R"([null, null, null, [-6, -7, -8]])"));
}

TEST(TestCaseWhen, Decimal) {
for (const auto& type :
std::vector<std::shared_ptr<DataType>>{decimal128(3, 2), decimal256(3, 2)}) {
Expand Down Expand Up @@ -2425,6 +2450,35 @@ TEST(TestCoalesce, DayTimeInterval) {
ArrayFromJSON(type, "[[1, 2], [1, 2], [1, 2], [1, 2]]"));
}

TEST(TestCoalesce, MonthDayNanoInterval) {
auto type = month_day_nano_interval();
auto scalar_null = ScalarFromJSON(type, "null");
auto scalar1 = ScalarFromJSON(type, "[1, 2, 3]");
auto values_null = ArrayFromJSON(type, "[null, null, null, null]");
auto values1 = ArrayFromJSON(type, "[null, [3, 4, 5], [5, 6, 7], [7, 8, 9]]");
auto values2 =
ArrayFromJSON(type, "[[9, 10, 0], [11, 12, 1], [13, 14, 2], [15, 16, 3]]");
auto values3 = ArrayFromJSON(type, "[[17, 18, 4], [19, 20, 5], [21, 22, 6], null]");
CheckScalar("coalesce", {values_null}, values_null);
CheckScalar("coalesce", {values_null, scalar1},
ArrayFromJSON(type, "[[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]"));
CheckScalar("coalesce", {values_null, values1}, values1);
CheckScalar("coalesce", {values_null, values2}, values2);
CheckScalar("coalesce", {values1, values_null}, values1);
CheckScalar("coalesce", {values2, values_null}, values2);
CheckScalar("coalesce", {scalar_null, values1}, values1);
CheckScalar("coalesce", {values1, scalar_null}, values1);
CheckScalar("coalesce", {values2, values1, values_null}, values2);
CheckScalar("coalesce", {values1, scalar1},
ArrayFromJSON(type, "[[1, 2, 3], [3, 4, 5], [5, 6, 7], [7, 8, 9]]"));
CheckScalar("coalesce", {values1, values2},
ArrayFromJSON(type, "[[9, 10, 0], [3, 4, 5], [5, 6, 7], [7, 8, 9]]"));
CheckScalar("coalesce", {values1, values2, values3},
ArrayFromJSON(type, "[[9, 10, 0], [3, 4, 5], [5, 6, 7], [7, 8, 9]]"));
CheckScalar("coalesce", {scalar1, values1},
ArrayFromJSON(type, "[[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]"));
}

TEST(TestCoalesce, Decimal) {
for (const auto& type :
std::vector<std::shared_ptr<DataType>>{decimal128(3, 2), decimal256(3, 2)}) {
Expand Down Expand Up @@ -2834,6 +2888,29 @@ TEST(TestChoose, DayTimeInterval) {
*MakeArrayOfNull(type, 5));
}

TEST(TestChoose, MonthDayNanoInterval) {
auto type = month_day_nano_interval();
auto indices1 = ArrayFromJSON(int64(), "[0, 1, 0, 1, null]");
auto values1 = ArrayFromJSON(type, "[[10, 1, 0], [10, 1, 0], null, null, [10, 1, 0]]");
auto values2 = ArrayFromJSON(type, "[[2, 20, 4], [2, 20, 4], null, null, [2, 20, 4]]");
auto nulls = ArrayFromJSON(type, "[null, null, null, null, null]");
CheckScalar("choose", {indices1, values1, values2},
ArrayFromJSON(type, "[[10, 1, 0], [2, 20, 4], null, null, null]"));
CheckScalar("choose", {indices1, ScalarFromJSON(type, "[1, 2, 3]"), values1},
ArrayFromJSON(type, "[[1, 2, 3], [10, 1, 0], [1, 2, 3], null, null]"));
CheckScalar("choose", {ScalarFromJSON(int64(), "0"), values1, values2}, values1);
CheckScalar("choose", {ScalarFromJSON(int64(), "1"), values1, values2}, values2);
CheckScalar("choose", {ScalarFromJSON(int64(), "null"), values1, values2}, nulls);
auto scalar1 = ScalarFromJSON(type, "[10, 1, 0]");
CheckScalar("choose", {ScalarFromJSON(int64(), "0"), scalar1, values2},
*MakeArrayFromScalar(*scalar1, 5));
CheckScalar("choose", {ScalarFromJSON(int64(), "1"), scalar1, values2}, values2);
CheckScalar("choose", {ScalarFromJSON(int64(), "null"), values1, values2}, nulls);
auto scalar_null = ScalarFromJSON(type, "null");
CheckScalar("choose", {ScalarFromJSON(int64(), "0"), scalar_null, values2},
*MakeArrayOfNull(type, 5));
}

TEST(TestChoose, Decimal) {
for (const auto& type : {decimal128(3, 2), decimal256(3, 2)}) {
auto indices1 = ArrayFromJSON(int64(), "[0, 1, 0, 1, null]");
Expand Down
4 changes: 3 additions & 1 deletion cpp/src/arrow/compute/kernels/vector_replace.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,9 @@ template <typename Type, typename Enable = void>
struct ReplaceWithMask {};

template <typename Type>
struct ReplaceWithMask<Type, enable_if_number<Type>> {
struct ReplaceWithMask<Type,
enable_if_t<is_number_type<Type>::value ||
std::is_same<Type, MonthDayNanoIntervalType>::value>> {
using T = typename TypeTraits<Type>::CType;

static void CopyData(const DataType&, uint8_t* out, const int64_t out_offset,
Expand Down
Loading