Skip to content

Commit 74b451a

Browse files
zhli1142015facebook-github-bot
zhli1142015
authored and
facebook-github-bot
committed
feat: Register Spark array_min/max functions with orderable types (#12576)
Summary: https://github.com/apache/spark/blob/84f5fd99fef1ad3a2b2c41dbee09c60560479c11/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala#L2341 Pull Request resolved: #12576 Reviewed By: bikramSingh91 Differential Revision: D71479672 Pulled By: xiaoxmeng fbshipit-source-id: 8e5471ac471886aeb7fa14446caf2c9d2dd0d517
1 parent abb1eee commit 74b451a

5 files changed

Lines changed: 174 additions & 43 deletions

File tree

‎velox/docs/functions/spark/array.rst‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,9 @@ Array Functions
9494
SELECT array_max(array(-1, -2, NULL)); -- -1
9595
SELECT array_max(array()); -- NULL
9696
SELECT array_max(array(-0.0001, -0.0002, -0.0003, float('nan'))); -- NaN
97+
SELECT array_max(array(array(1), array(NULL))); -- array(1)
98+
SELECT array_max(array(array(1), array(2, 1), array(2))); -- array(2, 1)
99+
SELECT array_max(array(array(1.0), array(1.0, 2.0), array(cast('NaN' as double)))); --array(NaN)
97100

98101
.. spark:function:: array_min(array(E)) -> E
99102
@@ -105,8 +108,11 @@ Array Functions
105108
SELECT array_min(array(-1, -2, NULL)); -- -2
106109
SELECT array_min(array(NULL, NULL)); -- NULL
107110
SELECT array_min(array()); -- NULL
108-
SELECT array_min(array(4.0, float('nan')]); -- 4.0
111+
SELECT array_min(array(4.0, float('nan'))); -- 4.0
109112
SELECT array_min(array(NULL, float('nan'))); -- NaN
113+
SELECT array_min(array(array(1), array(NULL))); -- array(NULL)
114+
SELECT array_min(array(array(1), array(1, 2), array(2))); -- array(1)
115+
SELECT array_min(array(array(1.0), array(1.0, 2.0), array(cast('NaN' as double)))); --array(1.0)
110116

111117
.. spark:function:: array_position(x, element) -> bigint
112118

‎velox/functions/sparksql/ArrayMinMaxFunction.h‎

Lines changed: 82 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,78 @@ struct ArrayMinMaxFunction {
2828
// Results refer to strings in the first argument.
2929
static constexpr int32_t reuse_strings_from_arg = 0;
3030

31+
template <typename TReturn, typename TInput>
32+
bool call(TReturn& out, const TInput& array) {
33+
// Result is null if array is empty.
34+
if (array.size() == 0) {
35+
return false;
36+
}
37+
38+
if (!array.mayHaveNulls()) {
39+
// Input array does not have nulls.
40+
auto currentValue = *array[0];
41+
for (auto i = 1; i < array.size(); i++) {
42+
update(currentValue, array[i].value());
43+
}
44+
assign(out, currentValue);
45+
return true;
46+
}
47+
48+
// Try to find the first non-null element.
49+
auto it = array.begin();
50+
while (it != array.end() && !it->has_value()) {
51+
++it;
52+
}
53+
// If array contains only NULL elements, return NULL.
54+
if (it == array.end()) {
55+
return false;
56+
}
57+
58+
// Now 'it' point to the first non-null element.
59+
auto currentValue = it->value();
60+
++it;
61+
while (it != array.end()) {
62+
if (it->has_value()) {
63+
update(currentValue, it->value());
64+
}
65+
++it;
66+
}
67+
68+
assign(out, currentValue);
69+
return true;
70+
}
71+
72+
bool call(
73+
out_type<Orderable<T1>>& out,
74+
const arg_type<Array<Orderable<T1>>>& array) {
75+
// Result is null if array is empty.
76+
if (array.size() == 0) {
77+
return false;
78+
}
79+
80+
int currentIndex = -1;
81+
for (auto i = 0; i < array.size(); i++) {
82+
if (array[i].has_value()) {
83+
if (currentIndex == -1) {
84+
currentIndex = i;
85+
} else {
86+
auto currentValue = array[currentIndex].value();
87+
auto candidateValue = array[i].value();
88+
if (compare(currentValue, candidateValue)) {
89+
currentIndex = i;
90+
}
91+
}
92+
}
93+
}
94+
if (currentIndex == -1) {
95+
// If array contains only NULL elements, return NULL.
96+
return false;
97+
}
98+
out.copy_from(array[currentIndex].value());
99+
return true;
100+
}
101+
102+
private:
31103
template <typename T>
32104
void update(T& currentValue, const T& candidateValue) {
33105
// NaN is greater than any non-NaN elements for double/float type.
@@ -66,45 +138,18 @@ struct ArrayMinMaxFunction {
66138
out.setNoCopy(value);
67139
}
68140

69-
template <typename TReturn, typename TInput>
70-
bool call(TReturn& out, const TInput& array) {
71-
// Result is null if array is empty.
72-
if (array.size() == 0) {
73-
return false;
74-
}
75-
76-
if (!array.mayHaveNulls()) {
77-
// Input array does not have nulls.
78-
auto currentValue = *array[0];
79-
for (auto i = 1; i < array.size(); i++) {
80-
update(currentValue, array[i].value());
81-
}
82-
assign(out, currentValue);
83-
return true;
84-
}
85-
86-
// Try to find the first non-null element.
87-
auto it = array.begin();
88-
while (it != array.end() && !it->has_value()) {
89-
++it;
90-
}
91-
// If array contains only NULL elements, return NULL.
92-
if (it == array.end()) {
93-
return false;
94-
}
141+
bool compare(
142+
exec::GenericView currentValue,
143+
exec::GenericView candidateValue) {
144+
static constexpr CompareFlags kFlags = {
145+
.nullHandlingMode = CompareFlags::NullHandlingMode::kNullAsValue};
95146

96-
// Now 'it' point to the first non-null element.
97-
auto currentValue = it->value();
98-
++it;
99-
while (it != array.end()) {
100-
if (it->has_value()) {
101-
update(currentValue, it->value());
102-
}
103-
++it;
147+
auto compareResult = candidateValue.compare(currentValue, kFlags).value();
148+
if constexpr (isMax) {
149+
return compareResult > 0;
150+
} else {
151+
return compareResult < 0;
104152
}
105-
106-
assign(out, currentValue);
107-
return true;
108153
}
109154
};
110155

‎velox/functions/sparksql/registration/RegisterArray.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,9 +103,11 @@ inline void registerArrayMinMaxFunctions(const std::string& prefix) {
103103
registerArrayMinMaxFunctions<float>(prefix);
104104
registerArrayMinMaxFunctions<double>(prefix);
105105
registerArrayMinMaxFunctions<bool>(prefix);
106+
registerArrayMinMaxFunctions<Varbinary>(prefix);
106107
registerArrayMinMaxFunctions<Varchar>(prefix);
107108
registerArrayMinMaxFunctions<Timestamp>(prefix);
108109
registerArrayMinMaxFunctions<Date>(prefix);
110+
registerArrayMinMaxFunctions<Orderable<T1>>(prefix);
109111
}
110112

111113
template <typename T>

‎velox/functions/sparksql/tests/ArrayMaxTest.cpp‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,11 @@ namespace {
2929
class ArrayMaxTest : public SparkFunctionBaseTest {
3030
protected:
3131
template <typename T>
32-
std::optional<T> arrayMax(const std::vector<std::optional<T>>& input) {
32+
std::optional<T> arrayMax(
33+
const std::vector<std::optional<T>>& input,
34+
const TypePtr& type = ARRAY(CppToType<T>::create())) {
3335
auto row = makeRowVector({makeNullableArrayVector(
34-
std::vector<std::vector<std::optional<T>>>{input})});
36+
std::vector<std::vector<std::optional<T>>>{input}, type)});
3537
return evaluateOnce<T>("array_max(C0)", row);
3638
}
3739
};
@@ -47,6 +49,17 @@ TEST_F(ArrayMaxTest, boolean) {
4749
EXPECT_EQ(arrayMax<bool>({true, true, true}), true);
4850
}
4951

52+
TEST_F(ArrayMaxTest, varbinary) {
53+
EXPECT_EQ(arrayMax<std::string>({"red", "blue"}, ARRAY(VARBINARY())), "red");
54+
EXPECT_EQ(
55+
arrayMax<std::string>(
56+
{std::nullopt, "blue", "yellow", "orange"}, ARRAY(VARBINARY())),
57+
"yellow");
58+
EXPECT_EQ(arrayMax<std::string>({}, ARRAY(VARBINARY())), std::nullopt);
59+
EXPECT_EQ(
60+
arrayMax<std::string>({std::nullopt}, ARRAY(VARBINARY())), std::nullopt);
61+
}
62+
5063
TEST_F(ArrayMaxTest, varchar) {
5164
EXPECT_EQ(arrayMax<std::string>({"red", "blue"}), "red");
5265
EXPECT_EQ(
@@ -87,6 +100,30 @@ TEST_F(ArrayMaxTest, timestamp) {
87100
EXPECT_EQ(arrayMax<Timestamp>({ts(0), std::nullopt}), ts(0));
88101
}
89102

103+
TEST_F(ArrayMaxTest, complexTypes) {
104+
auto testExpression = [&](const VectorPtr& input, const VectorPtr& expected) {
105+
auto result = evaluate("array_max(c0)", makeRowVector({input}));
106+
assertEqualVectors(expected, result);
107+
};
108+
testExpression(
109+
makeNestedArrayVectorFromJson<int64_t>(
110+
{"[[1, 1, 1], [1, 2, 2], [1, 3, 1]]"}),
111+
makeArrayVectorFromJson<int64_t>({"[1, 3, 1]"}));
112+
113+
testExpression(
114+
makeNestedArrayVectorFromJson<double>(
115+
{"[[2.0, null], [null, 2.0], [NaN, 1.0]]"}),
116+
makeArrayVectorFromJson<double>({"[NaN, 1.0]"}));
117+
118+
testExpression(
119+
makeNestedArrayVectorFromJson<int64_t>({"[[1, null], [null, 2], [1]]"}),
120+
makeArrayVectorFromJson<int64_t>({"[1, null]"}));
121+
122+
testExpression(
123+
makeNestedArrayVectorFromJson<int64_t>({"[null, null]"}),
124+
makeArrayVectorFromJson<int64_t>({"null"}));
125+
}
126+
90127
template <typename Type>
91128
class ArrayMaxIntegralTest : public ArrayMaxTest {
92129
public:

‎velox/functions/sparksql/tests/ArrayMinTest.cpp‎

Lines changed: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,11 @@ namespace {
3030
class ArrayMinTest : public SparkFunctionBaseTest {
3131
protected:
3232
template <typename T>
33-
std::optional<T> arrayMin(const std::vector<std::optional<T>>& input) {
33+
std::optional<T> arrayMin(
34+
const std::vector<std::optional<T>>& input,
35+
const TypePtr& type = ARRAY(CppToType<T>::create())) {
3436
auto row = makeRowVector({makeNullableArrayVector(
35-
std::vector<std::vector<std::optional<T>>>{input})});
37+
std::vector<std::vector<std::optional<T>>>{input}, type)});
3638
return evaluateOnce<T>("array_min(C0)", row);
3739
}
3840
};
@@ -46,7 +48,18 @@ TEST_F(ArrayMinTest, boolean) {
4648
EXPECT_EQ(arrayMin<bool>({std::nullopt, true, false, true}), false);
4749
EXPECT_EQ(arrayMin<bool>({false, false, false}), false);
4850
EXPECT_EQ(arrayMin<bool>({true, true, true}), true);
49-
} // namespace
51+
}
52+
53+
TEST_F(ArrayMinTest, varbinary) {
54+
EXPECT_EQ(arrayMin<std::string>({"red", "blue"}, ARRAY(VARBINARY())), "blue");
55+
EXPECT_EQ(
56+
arrayMin<std::string>(
57+
{std::nullopt, "blue", "yellow", "orange"}, ARRAY(VARBINARY())),
58+
"blue");
59+
EXPECT_EQ(arrayMin<std::string>({}, ARRAY(VARBINARY())), std::nullopt);
60+
EXPECT_EQ(
61+
arrayMin<std::string>({std::nullopt}, ARRAY(VARBINARY())), std::nullopt);
62+
}
5063

5164
TEST_F(ArrayMinTest, varchar) {
5265
EXPECT_EQ(arrayMin<std::string>({"red", "blue"}), "blue");
@@ -88,6 +101,34 @@ TEST_F(ArrayMinTest, timestamp) {
88101
EXPECT_EQ(arrayMin<Timestamp>({ts(0), std::nullopt}), ts(0));
89102
}
90103

104+
TEST_F(ArrayMinTest, complexTypes) {
105+
auto testExpression = [&](const VectorPtr& input, const VectorPtr& expected) {
106+
auto result = evaluate("array_min(c0)", makeRowVector({input}));
107+
assertEqualVectors(expected, result);
108+
};
109+
testExpression(
110+
makeNestedArrayVectorFromJson<int64_t>(
111+
{"[[1, 1, 1], [1, 1, 2], [1, 3, 1]]"}),
112+
makeArrayVectorFromJson<int64_t>({"[1, 1, 1]"}));
113+
114+
testExpression(
115+
makeNestedArrayVectorFromJson<int64_t>(
116+
{"[[1, null], [null, 2], [null, null]]"}),
117+
makeArrayVectorFromJson<int64_t>({"[null, null]"}));
118+
119+
testExpression(
120+
makeNestedArrayVectorFromJson<int64_t>({"[[1, null], [1, 2], [1]]"}),
121+
makeArrayVectorFromJson<int64_t>({"[1]"}));
122+
123+
testExpression(
124+
makeNestedArrayVectorFromJson<int64_t>({"[[1], [null], []]"}),
125+
makeArrayVectorFromJson<int64_t>({"[]"}));
126+
127+
testExpression(
128+
makeNestedArrayVectorFromJson<int64_t>({"[null, null]"}),
129+
makeArrayVectorFromJson<int64_t>({"null"}));
130+
}
131+
91132
template <typename Type>
92133
class ArrayMinIntegralTest : public ArrayMinTest {
93134
public:

0 commit comments

Comments
 (0)