Skip to content

Commit 6a71375

Browse files
denis0x0Ddenis0x0D
committed
[OLAP Predicate pushdown] Add block date arithmetics
1 parent 8f4bd11 commit 6a71375

File tree

6 files changed

+79
-28
lines changed

6 files changed

+79
-28
lines changed

‎ydb/core/kqp/expr_nodes/kqp_expr_nodes.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -642,7 +642,8 @@
642642
"Children": [
643643
{"Index": 0, "Name": "Operator", "Type": "TCoAtom"},
644644
{"Index": 1, "Name": "Left", "Type": "TExprBase"},
645-
{"Index": 2, "Name": "Right", "Type": "TExprBase"}
645+
{"Index": 2, "Name": "Right", "Type": "TExprBase"},
646+
{"Index": 3, "Name": "OpType", "Type": "TExprBase", "Optional": true}
646647
]
647648
},
648649
{

‎ydb/core/kqp/opt/kqp_statistics_transformer.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -552,7 +552,7 @@ class TKqpOlapPredicateSelectivityComputer: public TPredicateSelectivityComputer
552552
}
553553

554554
size_t listSize = listPtr->ChildrenSize();
555-
if (listSize == 3) {
555+
if (listSize == 3 || listSize == 4/*OpType optional field*/) {
556556
TString compSign = TString(listPtr->Child(0)->Content());
557557
auto left = listPtr->ChildPtr(1);
558558
auto right = listPtr->ChildPtr(2);

‎ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,13 +180,13 @@ std::vector<std::pair<TExprBase, TExprBase>> ExtractComparisonParameters(const T
180180

181181
TMaybeNode<TExprBase> ComparisonPushdown(const std::vector<std::pair<TExprBase, TExprBase>>& parameters, const TCoCompare& predicate, TExprContext& ctx, TPositionHandle pos);
182182

183-
[[maybe_unused]]
184183
TMaybeNode<TExprBase> CoalescePushdown(const TCoCoalesce& coalesce, const TExprNode& argument, TExprContext& ctx, bool allowApply) {
185184
if (const auto params = ExtractBinaryFunctionParameters(coalesce, argument, ctx, coalesce.Pos(), allowApply)) {
186185
return Build<TKqpOlapFilterBinaryOp>(ctx, coalesce.Pos())
187186
.Operator().Value("??", TNodeFlags::Default).Build()
188187
.Left(params->first)
189188
.Right(params->second)
189+
.OpType(ExpandType(coalesce.Pos(), *(coalesce.Ptr()->GetTypeAnn()), ctx))
190190
.Done();
191191
}
192192

@@ -409,6 +409,7 @@ std::vector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, const TExp
409409
.Operator().Value(arithmetic.Ref().Content(), TNodeFlags::Default).Build()
410410
.Left(UnwrapOptionalTKqpOlapApplyColumnArg(params->first))
411411
.Right(UnwrapOptionalTKqpOlapApplyColumnArg(params->second))
412+
.OpType(ExpandType(node.Pos(), *(arithmetic.Ptr()->GetTypeAnn()), ctx))
412413
.Done();
413414
}
414415
}
@@ -552,6 +553,7 @@ TExprBase BuildOneElementComparison(const std::pair<TExprBase, TExprBase>& param
552553
.Operator().Value(compareOperator, TNodeFlags::Default).Build()
553554
.Left(UnwrapOptionalTKqpOlapApplyColumnArg(parameter.first))
554555
.Right(UnwrapOptionalTKqpOlapApplyColumnArg(parameter.second))
556+
.OpType(ExpandType(predicate.Pos(), *(predicate.Ptr()->GetTypeAnn()), ctx))
555557
.Done();
556558
}
557559

‎ydb/core/kqp/opt/physical/predicate_collector.cpp

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,9 +79,17 @@ bool IsMemberColumn(const TExprBase& node, const TExprNode* lambdaArg) {
7979
return false;
8080
}
8181

82-
bool IsGoodTypeForArithmeticPushdown(const TTypeAnnotationNode& type, bool allowOlapApply) {
82+
bool IsGoodTypeForUnaryArithmeticPushdown(const TTypeAnnotationNode& type, bool allowOlapApply) {
8383
const auto features = NUdf::GetDataTypeInfo(RemoveOptionality(type).Cast<TDataExprType>()->GetSlot()).Features;
84-
return (NUdf::EDataTypeFeatures::NumericType & features)
84+
return ((NUdf::EDataTypeFeatures::NumericType) & features)
85+
|| (allowOlapApply && ((NUdf::EDataTypeFeatures::ExtDateType |
86+
NUdf::EDataTypeFeatures::DateType |
87+
NUdf::EDataTypeFeatures::TimeIntervalType) & features) && !(NUdf::EDataTypeFeatures::TzDateType & features));
88+
}
89+
90+
bool IsGoodTypeForBinaryArithmeticPushdown(const TTypeAnnotationNode& type, bool allowOlapApply) {
91+
const auto features = NUdf::GetDataTypeInfo(RemoveOptionality(type).Cast<TDataExprType>()->GetSlot()).Features;
92+
return ((NUdf::EDataTypeFeatures::NumericType | NUdf::EDataTypeFeatures::DateType | NUdf::EDataTypeFeatures::TimeIntervalType) & features)
8593
|| (allowOlapApply && ((NUdf::EDataTypeFeatures::ExtDateType |
8694
NUdf::EDataTypeFeatures::DateType |
8795
NUdf::EDataTypeFeatures::TimeIntervalType) & features) && !(NUdf::EDataTypeFeatures::TzDateType & features));
@@ -118,7 +126,6 @@ bool CanPushdownStringUdf(const TExprNode& udf, bool pushdownSubstring) {
118126
return substringMatchUdfs.contains(name);
119127
}
120128

121-
[[maybe_unused]]
122129
bool AbstractTreeCanBePushed(const TExprBase& expr, const TExprNode*, bool pushdownSubstring) {
123130
if (!expr.Ref().IsCallable({"Apply", "Coalesce", "NamedApply", "IfPresent", "Visit"})) {
124131
return false;
@@ -194,10 +201,13 @@ bool CheckExpressionNodeForPushdown(const TExprBase& node, const TExprNode* lamb
194201
}
195202

196203
if (const auto op = node.Maybe<TCoUnaryArithmetic>()) {
197-
return CheckExpressionNodeForPushdown(op.Cast().Arg(), lambdaArg, options) && IsGoodTypeForArithmeticPushdown(*op.Cast().Ref().GetTypeAnn(), options.AllowOlapApply);
204+
return CheckExpressionNodeForPushdown(op.Cast().Arg(), lambdaArg, options) &&
205+
IsGoodTypeForUnaryArithmeticPushdown(*op.Cast().Ref().GetTypeAnn(), options.AllowOlapApply);
198206
} else if (const auto op = node.Maybe<TCoBinaryArithmetic>()) {
199-
return CheckExpressionNodeForPushdown(op.Cast().Left(), lambdaArg, options) && CheckExpressionNodeForPushdown(op.Cast().Right(), lambdaArg, options)
200-
&& IsGoodTypeForArithmeticPushdown(*op.Cast().Ref().GetTypeAnn(), options.AllowOlapApply) && !op.Cast().Maybe<TCoAggrAdd>();
207+
return CheckExpressionNodeForPushdown(op.Cast().Left(), lambdaArg, options) &&
208+
CheckExpressionNodeForPushdown(op.Cast().Right(), lambdaArg, options) &&
209+
IsGoodTypeForBinaryArithmeticPushdown(*op.Cast().Ref().GetTypeAnn(), options.AllowOlapApply) &&
210+
!op.Cast().Maybe<TCoAggrAdd>();
201211
}
202212

203213
if (options.AllowOlapApply) {

‎ydb/core/kqp/query_compiler/kqp_olap_compiler.cpp

Lines changed: 37 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -115,19 +115,27 @@ class TKqpOlapCompileContext {
115115
return ExprContext.MakeType<TBlockExprType>(resultItemType);
116116
}
117117

118-
const TTypeAnnotationNode* GetReturnType(TPositionHandle pos, const TTypeAnnotationNode& left, const TTypeAnnotationNode& right, const TTypeAnnotationNode* resultItemType, bool optionalityFromRight) const {
118+
const TTypeAnnotationNode *GetReturnType(TPositionHandle pos, const TTypeAnnotationNode &left,
119+
const TTypeAnnotationNode &right, const TTypeAnnotationNode *resultItemType,
120+
bool optionalityFromRight) const {
119121
bool isScalarLeft, isScalarRight;
120122
const auto leftItemType = GetBlockItemType(left, isScalarLeft);
121123
const auto rightItemType = GetBlockItemType(right, isScalarRight);
122124

123125
if (!resultItemType) {
124-
const auto& leftCleanType = RemoveOptionality(*leftItemType);
125-
const auto& rightCleanType = RemoveOptionality(*rightItemType);
126+
const auto &leftCleanType = RemoveOptionality(*leftItemType);
127+
const auto &rightCleanType = RemoveOptionality(*rightItemType);
126128
resultItemType = CommonType<true>(pos, &leftCleanType, &rightCleanType, ExprContext);
129+
} else {
130+
if (resultItemType->GetKind() == ETypeAnnotationKind::Optional) {
131+
resultItemType = resultItemType->Cast<TOptionalExprType>()->GetItemType();
132+
}
127133
}
128-
YQL_ENSURE(resultItemType);
129134

130-
if ((ETypeAnnotationKind::Optional == leftItemType->GetKind() && !optionalityFromRight) || ETypeAnnotationKind::Optional == rightItemType->GetKind()) {
135+
Y_ENSURE(resultItemType, "KqpOlapCompiler: Result type is nullptr.");
136+
137+
if ((ETypeAnnotationKind::Optional == leftItemType->GetKind() && !optionalityFromRight) ||
138+
ETypeAnnotationKind::Optional == rightItemType->GetKind()) {
131139
resultItemType = ExprContext.MakeType<TOptionalExprType>(resultItemType);
132140
}
133141

@@ -137,17 +145,23 @@ class TKqpOlapCompileContext {
137145
return ExprContext.MakeType<TBlockExprType>(resultItemType);
138146
}
139147

140-
std::pair<ui32, const TTypeAnnotationNode*> AddYqlKernelIfFunc(TPositionHandle pos, const TTypeAnnotationNode& conditionType, const TTypeAnnotationNode& thenType, const TTypeAnnotationNode& elseType) const {
148+
std::pair<ui32, const TTypeAnnotationNode *> AddYqlKernelIfFunc(TPositionHandle pos, const TTypeAnnotationNode &conditionType,
149+
const TTypeAnnotationNode &thenType,
150+
const TTypeAnnotationNode &elseType) const {
141151
const auto retBlockType = GetReturnType(pos, thenType, elseType, nullptr, false);
142152
return std::make_pair(YqlKernelRequestBuilder->AddIf(&conditionType, &thenType, &elseType), retBlockType);
143153
}
144154

145-
std::pair<ui32, const TTypeAnnotationNode*> AddYqlKernelBinaryFunc(TPositionHandle pos, TKernelRequestBuilder::EBinaryOp op, const TTypeAnnotationNode& argTypeOne, const TTypeAnnotationNode& argTypeTwo, const TTypeAnnotationNode* retType) const {
155+
std::pair<ui32, const TTypeAnnotationNode *> AddYqlKernelBinaryFunc(TPositionHandle pos, TKernelRequestBuilder::EBinaryOp op,
156+
const TTypeAnnotationNode &argTypeOne,
157+
const TTypeAnnotationNode &argTypeTwo,
158+
const TTypeAnnotationNode *retType) const {
146159
const auto retBlockType = GetReturnType(pos, argTypeOne, argTypeTwo, retType, TKernelRequestBuilder::EBinaryOp::Coalesce == op);
147160
return std::make_pair(YqlKernelRequestBuilder->AddBinaryOp(op, &argTypeOne, &argTypeTwo, retBlockType), retBlockType);
148161
}
149162

150-
ui32 AddYqlKernelBinaryFunc(TPositionHandle pos, TKernelRequestBuilder::EBinaryOp op, const TExprBase& arg1, const TExprBase& arg2, const TTypeAnnotationNode* retType) const {
163+
ui32 AddYqlKernelBinaryFunc(TPositionHandle pos, TKernelRequestBuilder::EBinaryOp op, const TExprBase &arg1, const TExprBase &arg2,
164+
const TTypeAnnotationNode *retType) const {
151165
const auto arg1Type = GetArgType(arg1);
152166
const auto arg2Type = GetArgType(arg2);
153167
return AddYqlKernelBinaryFunc(pos, op, *arg1Type, *arg2Type, retType).first;
@@ -665,9 +679,15 @@ TTypedColumn CompileYqlKernelUnaryOperation(const TKqpOlapFilterUnaryOp& operati
665679
return {command->GetColumn().GetId(), resultType};
666680
}
667681

668-
[[maybe_unused]]
669-
TTypedColumn CompileYqlKernelBinaryOperation(const TKqpOlapFilterBinaryOp& operation, TKqpOlapCompileContext& ctx)
670-
{
682+
const TTypeAnnotationNode *TryToGetType(const TKqpOlapFilterBinaryOp &operation) {
683+
const auto opPtr = operation.Ptr();
684+
if (opPtr->ChildrenSize() > TKqpOlapFilterBinaryOp::idx_OpType) {
685+
return opPtr->Child(TKqpOlapFilterBinaryOp::idx_OpType)->GetTypeAnn()->Cast<TTypeExprType>()->GetType();
686+
}
687+
return nullptr;
688+
}
689+
690+
TTypedColumn CompileYqlKernelBinaryOperation(const TKqpOlapFilterBinaryOp &operation, TKqpOlapCompileContext &ctx) {
671691
// Columns should be created before operation, otherwise operation fail to find columns
672692
const auto leftColumn = GetOrCreateColumnIdAndType(operation.Left(), ctx);
673693
const auto rightColumn = GetOrCreateColumnIdAndType(operation.Right(), ctx);
@@ -697,22 +717,22 @@ TTypedColumn CompileYqlKernelBinaryOperation(const TKqpOlapFilterBinaryOp& opera
697717
op = TKernelRequestBuilder::EBinaryOp::GreaterOrEqual;
698718
} else if (oper == "+"sv) {
699719
op = TKernelRequestBuilder::EBinaryOp::Add;
700-
type = nullptr;
720+
type = TryToGetType(operation);
701721
} else if (oper == "-"sv) {
702722
op = TKernelRequestBuilder::EBinaryOp::Sub;
703-
type = nullptr;
723+
type = TryToGetType(operation);
704724
} else if (oper == "*"sv) {
705725
op = TKernelRequestBuilder::EBinaryOp::Mul;
706-
type = nullptr;
726+
type = TryToGetType(operation);
707727
} else if (oper == "/"sv) {
708728
op = TKernelRequestBuilder::EBinaryOp::Div;
709-
type = nullptr;
729+
type = TryToGetType(operation);
710730
} else if (oper == "%"sv) {
711731
op = TKernelRequestBuilder::EBinaryOp::Mod;
712-
type = nullptr;
732+
type = TryToGetType(operation);
713733
} else if (oper == "??"sv) {
714734
op = TKernelRequestBuilder::EBinaryOp::Coalesce;
715-
type = nullptr;
735+
type = TryToGetType(operation);
716736
} else {
717737
YQL_ENSURE(false, "Unknown binary OLAP operation: " << oper);
718738
}

‎ydb/core/kqp/ut/olap/kqp_olap_ut.cpp

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1880,10 +1880,28 @@ Y_UNIT_TEST_SUITE(KqpOlap) {
18801880
R"(dtm64 >= dtm)",
18811881
R"(dt >= ts)",
18821882
R"(dt >= ts64)",
1883-
//R"(dt >= dt)", Failed in runtime
1883+
//R"(dt >= dt)", Failed in runtime, needs support from CS, to support to build an execution graph with same nodes.
18841884

18851885
// 2. Arithmetic
1886-
// R"(dt <= dt - inter64)", KqpOlapCompiler cannot deduce result type
1886+
R"(dt <= dt - inter64)",
1887+
R"(dt <= dt - Interval("P100D"))",
1888+
R"(dt32 <= dt32 - inter64)",
1889+
R"(dt32 <= dt32 - Interval("P100D"))",
1890+
R"(dtm <= dtm - inter64)",
1891+
R"(dtm <= dtm - Interval("P100D"))",
1892+
R"(dtm64 <= dtm64 - inter64)",
1893+
R"(dtm64 <= dtm64 - Interval("P100D"))",
1894+
R"(ts <= ts - inter64)",
1895+
R"(ts <= ts - Interval("P100D"))",
1896+
R"(ts64 <= ts64 - inter64)",
1897+
R"(ts64 <= ts64 - Interval("P100D"))",
1898+
1899+
R"(inter64 <= dt - Date('2001-01-01'))",
1900+
R"(inter64 <= dt32 - Date32('2001-01-01'))",
1901+
R"(inter64 <= dtm - DateTime('1998-12-01T15:30:00Z'))",
1902+
R"(inter64 <= dtm64 - DateTime64('1998-12-01T15:30:00Z'))",
1903+
R"(inter64 <= ts - Timestamp("1970-01-01T00:00:03.000001Z"))",
1904+
R"(inter64 <= ts64 - Timestamp64("1970-01-01T00:00:03.000001Z"))",
18871905
};
18881906

18891907
auto queryPrefix = R"(

0 commit comments

Comments
 (0)