Skip to content

Commit 551e1cd

Browse files
PHILO-HEPHILO-HE
authored
Enable date type for kPreceeding & kFollowing window range bound (#291)
1 parent dba81cb commit 551e1cd

File tree

3 files changed

+89
-105
lines changed

3 files changed

+89
-105
lines changed

‎velox/exec/Window.cpp‎

Lines changed: 50 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -602,52 +602,47 @@ void Window::updateFrameBounds(
602602
updateKRowsFrameBounds(
603603
true, frameArg.value(), startRow, numRows, rawFrameBounds);
604604
} else {
605+
#define VELOX_DYNAMIC_LIMITED_SCALAR_TYPE_DISPATCH( \
606+
TEMPLATE_FUNC, typeKind, ...) \
607+
[&]() { \
608+
switch (typeKind) { \
609+
case ::facebook::velox::TypeKind::INTEGER: { \
610+
return TEMPLATE_FUNC<::facebook::velox::TypeKind::INTEGER>( \
611+
__VA_ARGS__); \
612+
} \
613+
case ::facebook::velox::TypeKind::TINYINT: { \
614+
return TEMPLATE_FUNC<::facebook::velox::TypeKind::TINYINT>( \
615+
__VA_ARGS__); \
616+
} \
617+
case ::facebook::velox::TypeKind::SMALLINT: { \
618+
return TEMPLATE_FUNC<::facebook::velox::TypeKind::SMALLINT>( \
619+
__VA_ARGS__); \
620+
} \
621+
case ::facebook::velox::TypeKind::BIGINT: { \
622+
return TEMPLATE_FUNC<::facebook::velox::TypeKind::BIGINT>( \
623+
__VA_ARGS__); \
624+
} \
625+
case ::facebook::velox::TypeKind::DATE: { \
626+
return TEMPLATE_FUNC<::facebook::velox::TypeKind::DATE>(__VA_ARGS__); \
627+
} \
628+
default: \
629+
VELOX_FAIL( \
630+
"Not supported type for sort key!: {}", \
631+
mapTypeKindToName(typeKind)); \
632+
} \
633+
}()
605634
// Sort key type.
606635
auto sortKeyTypePtr = outputType_->childAt(sortKeyInfo_[0].first);
607-
switch (sortKeyTypePtr->kind()) {
608-
case TypeKind::TINYINT:
609-
updateKRangeFrameBounds<TypeKind::TINYINT>(
610-
true,
611-
isStartBound,
612-
frameArg.value(),
613-
numRows,
614-
rawFrameBounds,
615-
rawPeerStarts,
616-
rawPeerEnds);
617-
break;
618-
case TypeKind::SMALLINT:
619-
updateKRangeFrameBounds<TypeKind::SMALLINT>(
620-
true,
621-
isStartBound,
622-
frameArg.value(),
623-
numRows,
624-
rawFrameBounds,
625-
rawPeerStarts,
626-
rawPeerEnds);
627-
break;
628-
case TypeKind::INTEGER:
629-
updateKRangeFrameBounds<TypeKind::INTEGER>(
630-
true,
631-
isStartBound,
632-
frameArg.value(),
633-
numRows,
634-
rawFrameBounds,
635-
rawPeerStarts,
636-
rawPeerEnds);
637-
break;
638-
case TypeKind::BIGINT:
639-
updateKRangeFrameBounds<TypeKind::BIGINT>(
640-
true,
641-
isStartBound,
642-
frameArg.value(),
643-
numRows,
644-
rawFrameBounds,
645-
rawPeerStarts,
646-
rawPeerEnds);
647-
break;
648-
default:
649-
VELOX_USER_FAIL("Not supported type for sort key!");
650-
}
636+
VELOX_DYNAMIC_LIMITED_SCALAR_TYPE_DISPATCH(
637+
updateKRangeFrameBounds,
638+
sortKeyTypePtr->kind(),
639+
true,
640+
isStartBound,
641+
frameArg.value(),
642+
numRows,
643+
rawFrameBounds,
644+
rawPeerStarts,
645+
rawPeerEnds);
651646
}
652647
break;
653648
}
@@ -658,50 +653,17 @@ void Window::updateFrameBounds(
658653
} else {
659654
// Sort key type.
660655
auto sortKeyTypePtr = outputType_->childAt(sortKeyInfo_[0].first);
661-
switch (sortKeyTypePtr->kind()) {
662-
case TypeKind::TINYINT:
663-
updateKRangeFrameBounds<TypeKind::TINYINT>(
664-
false,
665-
isStartBound,
666-
frameArg.value(),
667-
numRows,
668-
rawFrameBounds,
669-
rawPeerStarts,
670-
rawPeerEnds);
671-
break;
672-
case TypeKind::SMALLINT:
673-
updateKRangeFrameBounds<TypeKind::SMALLINT>(
674-
false,
675-
isStartBound,
676-
frameArg.value(),
677-
numRows,
678-
rawFrameBounds,
679-
rawPeerStarts,
680-
rawPeerEnds);
681-
break;
682-
case TypeKind::INTEGER:
683-
updateKRangeFrameBounds<TypeKind::INTEGER>(
684-
false,
685-
isStartBound,
686-
frameArg.value(),
687-
numRows,
688-
rawFrameBounds,
689-
rawPeerStarts,
690-
rawPeerEnds);
691-
break;
692-
case TypeKind::BIGINT:
693-
updateKRangeFrameBounds<TypeKind::BIGINT>(
694-
false,
695-
isStartBound,
696-
frameArg.value(),
697-
numRows,
698-
rawFrameBounds,
699-
rawPeerStarts,
700-
rawPeerEnds);
701-
break;
702-
default:
703-
VELOX_USER_FAIL("Not supported type for sort key!");
704-
}
656+
VELOX_DYNAMIC_LIMITED_SCALAR_TYPE_DISPATCH(
657+
updateKRangeFrameBounds,
658+
sortKeyTypePtr->kind(),
659+
false,
660+
isStartBound,
661+
frameArg.value(),
662+
numRows,
663+
rawFrameBounds,
664+
rawPeerStarts,
665+
rawPeerEnds);
666+
#undef VELOX_DYNAMIC_LIMITED_SCALAR_TYPE_DISPATCH
705667
}
706668
break;
707669
}

‎velox/functions/prestosql/window/tests/SimpleAggregatesTest.cpp‎

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,21 @@ TEST_F(StringAggregatesTest, nonFixedWidthAggregate) {
127127
testWindowFunction(input, "max(c2)", kOverClauses);
128128
}
129129

130-
class KPreceedingFollowingTest : public WindowTestBase {};
130+
class KPrecedingFollowingTest : public WindowTestBase {
131+
public:
132+
const std::vector<std::string> kRangeFrames = {
133+
"range between unbounded preceding and 1 following",
134+
"range between unbounded preceding and 2 following",
135+
"range between unbounded preceding and 3 following",
136+
"range between 1 preceding and unbounded following",
137+
"range between 2 preceding and unbounded following",
138+
"range between 3 preceding and unbounded following",
139+
"range between 1 preceding and 3 following",
140+
"range between 3 preceding and 1 following",
141+
"range between 2 preceding and 2 following"};
142+
};
131143

132-
TEST_F(KPreceedingFollowingTest, rangeFrames1) {
144+
TEST_F(KPrecedingFollowingTest, rangeFrames1) {
133145
auto vectors = makeRowVector({
134146
makeFlatVector<int64_t>({1, 1, 2147483650, 3, 2, 2147483650}),
135147
makeFlatVector<std::string>({"1", "1", "1", "2", "1", "2"}),
@@ -147,7 +159,7 @@ TEST_F(KPreceedingFollowingTest, rangeFrames1) {
147159
testWindowFunction({vectors}, "count(c0)", {overClause}, kRangeFrames2);
148160
}
149161

150-
TEST_F(KPreceedingFollowingTest, rangeFrames2) {
162+
TEST_F(KPrecedingFollowingTest, rangeFrames2) {
151163
const std::vector<RowVectorPtr> vectors = {
152164
makeRowVector(
153165
{makeFlatVector<int64_t>({5, 6, 8, 9, 10, 2, 8, 9, 3}),
@@ -164,36 +176,42 @@ TEST_F(KPreceedingFollowingTest, rangeFrames2) {
164176
makeRowVector(
165177
{makeFlatVector<int64_t>({5, 5, 4, 6, 3, 2}),
166178
makeFlatVector<std::string>({"1", "2", "2", "2", "1", "2"})}),
167-
// Uses int32 for sort column.
179+
// Uses int32 type for sort column.
168180
makeRowVector(
169181
{makeFlatVector<int32_t>({5, 5, 4, 6, 3, 2}),
170182
makeFlatVector<std::string>({"1", "2", "2", "2", "1", "2"})}),
171183
};
172-
173184
const std::string overClause = "partition by c1 order by c0";
174-
const std::vector<std::string> kRangeFrames = {
175-
"range between unbounded preceding and 1 following",
176-
"range between unbounded preceding and 2 following",
177-
"range between unbounded preceding and 3 following",
178-
"range between 1 preceding and unbounded following",
179-
"range between 2 preceding and unbounded following",
180-
"range between 3 preceding and unbounded following",
181-
"range between 1 preceding and 3 following",
182-
"range between 3 preceding and 1 following",
183-
"range between 2 preceding and 2 following"};
184185
for (int i = 0; i < vectors.size(); i++) {
185186
testWindowFunction({vectors[i]}, "avg(c0)", {overClause}, kRangeFrames);
186187
testWindowFunction({vectors[i]}, "sum(c0)", {overClause}, kRangeFrames);
187188
testWindowFunction({vectors[i]}, "count(c0)", {overClause}, kRangeFrames);
188189
}
189190
}
190191

191-
TEST_F(KPreceedingFollowingTest, rowsFrames) {
192+
TEST_F(KPrecedingFollowingTest, rangeFrames3) {
193+
const std::vector<RowVectorPtr> vectors = {
194+
// Uses date type for sort column.
195+
makeRowVector(
196+
{makeFlatVector<Date>(
197+
{Date(6), Date(1), Date(5), Date(0), Date(7), Date(1)}),
198+
makeFlatVector<std::string>({"1", "2", "2", "2", "1", "2"})}),
199+
makeRowVector(
200+
{makeFlatVector<Date>(
201+
{Date(5), Date(5), Date(4), Date(6), Date(3), Date(2)}),
202+
makeFlatVector<std::string>({"1", "2", "2", "2", "1", "2"})}),
203+
};
204+
const std::string overClause = "partition by c1 order by c0";
205+
for (int i = 0; i < vectors.size(); i++) {
206+
testWindowFunction({vectors[i]}, "count(c0)", {overClause}, kRangeFrames);
207+
}
208+
}
209+
210+
TEST_F(KPrecedingFollowingTest, rowsFrames) {
192211
auto vectors = makeRowVector({
193212
makeFlatVector<int64_t>({1, 1, 2147483650, 3, 2, 2147483650}),
194213
makeFlatVector<std::string>({"1", "1", "1", "2", "1", "2"}),
195214
});
196-
197215
const std::string overClause = "partition by c1 order by c0";
198216
const std::vector<std::string> kRangeFrames = {
199217
"rows between current row and 2147483647 following",

‎velox/type/Date.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,10 @@ struct Date {
3939
days_ += days;
4040
}
4141

42+
Date operator+(const int32_t days) const {
43+
return Date(days_ + days);
44+
}
45+
4246
bool operator==(const Date& other) const {
4347
return days_ == other.days_;
4448
}

0 commit comments

Comments
 (0)