Skip to content

Coalesce pushdown #744

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

Merged
merged 9 commits into from
Jan 9, 2024
Merged

Coalesce pushdown #744

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
fb86065
Pushdown coalesce as YQL kernel. [Draft]
Tony-Romanov Dec 22, 2023
97fb075
Add some nulls into test source table.
Tony-Romanov Dec 24, 2023
045db3d
JsonExist a little fix.
Tony-Romanov Dec 27, 2023
6aedc4b
Merge branch 'ydb-platform:main' into coalesce_pushdown
Tony-Romanov Dec 27, 2023
de9de3b
Fix test
Tony-Romanov Dec 29, 2023
9e2de83
Fix make return type for binary operations.
Tony-Romanov Dec 29, 2023
4307a10
Merge branch 'ydb-platform:main' into coalesce_pushdown
Tony-Romanov Dec 29, 2023
a9ddf8d
Some fixes.
Tony-Romanov Jan 2, 2024
ab1e5ff
Enable tests.
Tony-Romanov Jan 2, 2024
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
28 changes: 28 additions & 0 deletions ydb/core/formats/arrow/program.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,34 @@ class TAssign {
, FuncOpts(nullptr)
{}

explicit TAssign(const TColumnInfo& column, i8 value)
: Column(column)
, Operation(EOperation::Constant)
, Constant(std::make_shared<arrow::Int8Scalar>(value))
, FuncOpts(nullptr)
{}

explicit TAssign(const TColumnInfo& column, ui8 value)
: Column(column)
, Operation(EOperation::Constant)
, Constant(std::make_shared<arrow::UInt8Scalar>(value))
, FuncOpts(nullptr)
{}

explicit TAssign(const TColumnInfo& column, i16 value)
: Column(column)
, Operation(EOperation::Constant)
, Constant(std::make_shared<arrow::Int16Scalar>(value))
, FuncOpts(nullptr)
{}

explicit TAssign(const TColumnInfo& column, ui16 value)
: Column(column)
, Operation(EOperation::Constant)
, Constant(std::make_shared<arrow::UInt16Scalar>(value))
, FuncOpts(nullptr)
{}

explicit TAssign(const TColumnInfo& column, i32 value)
: Column(column)
, Operation(EOperation::Constant)
Expand Down
2 changes: 1 addition & 1 deletion ydb/core/kqp/host/kqp_type_ann.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -844,7 +844,7 @@ bool ValidateOlapFilterConditions(const TExprNode* node, const TStructExprType*
if (!EnsureAtom(*op, ctx)) {
return false;
}
if (!op->IsAtom({"eq", "neq", "lt", "lte", "gt", "gte", "string_contains", "starts_with", "ends_with", "+", "-", "*", "/", "%"})) {
if (!op->IsAtom({"eq", "neq", "lt", "lte", "gt", "gte", "string_contains", "starts_with", "ends_with", "+", "-", "*", "/", "%", "??"})) {
ctx.AddError(TIssue(ctx.GetPosition(node->Pos()),
TStringBuilder() << "Unexpected OLAP binary operation: " << op->Content()
));
Expand Down
83 changes: 58 additions & 25 deletions ydb/core/kqp/opt/physical/kqp_opt_phy_olap_filter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#include <ydb/core/formats/arrow/ssa_runtime_version.h>
#include <ydb/core/kqp/common/kqp_yql.h>
#include <ydb/library/yql/core/extract_predicate/extract_predicate.h>
#include <ydb/library/yql/core/yql_expr_optimize.h>
#include <ydb/library/yql/core/yql_opt_utils.h>
#include <ydb/library/yql/providers/common/pushdown/collection.h>
#include <ydb/library/yql/providers/common/pushdown/predicate_node.h>
Expand Down Expand Up @@ -99,12 +100,27 @@ bool IsFalseLiteral(TExprBase node) {
return node.Maybe<TCoBool>() && !FromString<bool>(node.Cast<TCoBool>().Literal().Value());
}

std::optional<std::pair<TExprBase, TExprBase>> ExtractArithmeticParameters(const TCoBinaryArithmetic& op, TExprContext& ctx, TPositionHandle pos);
std::optional<std::pair<TExprBase, TExprBase>> ExtractBinaryFunctionParameters(const TExprBase& op, TExprContext& ctx, TPositionHandle pos);
std::vector<std::pair<TExprBase, TExprBase>> ExtractComparisonParameters(const TCoCompare& predicate, TExprContext& ctx, TPositionHandle pos);

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

TMaybeNode<TExprBase> YqlCoalescePushdown(const TCoCoalesce& coalesce, TExprContext& ctx) {
if (const auto params = ExtractBinaryFunctionParameters(coalesce, ctx, coalesce.Pos())) {
return Build<TKqpOlapFilterBinaryOp>(ctx, coalesce.Pos())
.Operator().Value("??", TNodeFlags::Default).Build()
.Left(params->first)
.Right(params->second)
.Done();
}

return NullNode;
}

TVector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, TExprContext& ctx, TPositionHandle pos)
{
TVector<TExprBase> out;
auto convertNode = [&ctx, &pos](const TExprBase& node) -> TMaybeNode<TExprBase> {
const auto convertNode = [&ctx, &pos](const TExprBase& node) -> TMaybeNode<TExprBase> {
if (node.Maybe<TCoNull>()) {
return node;
}
Expand Down Expand Up @@ -140,23 +156,34 @@ TVector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, TExprContext&
builder.ReturningType(maybeReturningType.Cast());
} else {
builder.ReturningType<TCoDataType>()
.Type(ctx.NewAtom(node.Pos(), "Utf8"))
.Type().Value("Utf8", TNodeFlags::Default).Build()
.Build();
}
return builder.Done();
}

if (const auto maybeArithmetic = node.Maybe<TCoBinaryArithmetic>()) {
if (const auto arithmetic = maybeArithmetic.Cast(); !arithmetic.Maybe<TCoAggrAdd>()) {
if (const auto params = ExtractArithmeticParameters(arithmetic, ctx, pos)) {
if (const auto params = ExtractBinaryFunctionParameters(arithmetic, ctx, pos)) {
return Build<TKqpOlapFilterBinaryOp>(ctx, pos)
.Operator(ctx.NewAtom(pos, arithmetic.Ref().Content(), TNodeFlags::Default))
.Operator().Value(arithmetic.Ref().Content(), TNodeFlags::Default).Build()
.Left(params->first)
.Right(params->second)
.Done();
}
}
}

if (const auto maybeCoalesce = node.Maybe<TCoCoalesce>()) {
return YqlCoalescePushdown(maybeCoalesce.Cast(), ctx);
}

if (const auto maybeCompare = node.Maybe<TCoCompare>()) {
if (const auto params = ExtractComparisonParameters(maybeCompare.Cast(), ctx, pos); !params.empty()) {
return ComparisonPushdown(params, maybeCompare.Cast(), ctx, pos);
}
}

return NullNode;
};

Expand Down Expand Up @@ -192,24 +219,24 @@ TVector<TExprBase> ConvertComparisonNode(const TExprBase& nodeIn, TExprContext&
return out;
}

std::optional<std::pair<TExprBase, TExprBase>> ExtractArithmeticParameters(const TCoBinaryArithmetic& op, TExprContext& ctx, TPositionHandle pos)
std::optional<std::pair<TExprBase, TExprBase>> ExtractBinaryFunctionParameters(const TExprBase& op, TExprContext& ctx, TPositionHandle pos)
{
const auto left = ConvertComparisonNode(op.Left(), ctx, pos);
const auto left = ConvertComparisonNode(TExprBase(op.Ref().HeadPtr()), ctx, pos);
if (left.size() != 1U) {
return std::nullopt;
}

const auto right = ConvertComparisonNode(op.Right(), ctx, pos);
const auto right = ConvertComparisonNode(TExprBase(op.Ref().TailPtr()), ctx, pos);
if (right.size() != 1U) {
return std::nullopt;
}

return std::make_pair(left.front(), right.front());
}

TVector<std::pair<TExprBase, TExprBase>> ExtractComparisonParameters(const TCoCompare& predicate, TExprContext& ctx, TPositionHandle pos)
std::vector<std::pair<TExprBase, TExprBase>> ExtractComparisonParameters(const TCoCompare& predicate, TExprContext& ctx, TPositionHandle pos)
{
TVector<std::pair<TExprBase, TExprBase>> out;
std::vector<std::pair<TExprBase, TExprBase>> out;
auto left = ConvertComparisonNode(predicate.Left(), ctx, pos);

if (left.empty()) {
Expand Down Expand Up @@ -280,13 +307,13 @@ TExprBase BuildOneElementComparison(const std::pair<TExprBase, TExprBase>& param
YQL_ENSURE(!compareOperator.empty(), "Unsupported comparison node: " << predicate.Ptr()->Content());

return Build<TKqpOlapFilterBinaryOp>(ctx, pos)
.Operator(ctx.NewAtom(pos, compareOperator, TNodeFlags::Default))
.Operator().Value(compareOperator, TNodeFlags::Default).Build()
.Left(parameter.first)
.Right(parameter.second)
.Done();
}

TMaybeNode<TExprBase> ComparisonPushdown(const TVector<std::pair<TExprBase, TExprBase>>& parameters, const TCoCompare& predicate,
TMaybeNode<TExprBase> ComparisonPushdown(const std::vector<std::pair<TExprBase, TExprBase>>& parameters, const TCoCompare& predicate,
TExprContext& ctx, TPositionHandle pos)
{
ui32 conditionsCount = parameters.size();
Expand Down Expand Up @@ -342,7 +369,7 @@ TMaybeNode<TExprBase> ComparisonPushdown(const TVector<std::pair<TExprBase, TExp

for (ui32 j = 0; j < i; ++j) {
andConditions.emplace_back(Build<TKqpOlapFilterBinaryOp>(ctx, pos)
.Operator(ctx.NewAtom(pos, "eq"))
.Operator().Value("eq", TNodeFlags::Default).Build()
.Left(parameters[j].first)
.Right(parameters[j].second)
.Done());
Expand All @@ -360,6 +387,16 @@ TMaybeNode<TExprBase> ComparisonPushdown(const TVector<std::pair<TExprBase, TExp
.Done();
}

TMaybeNode<TExprBase> SimplePredicatePushdown(const TCoCompare& predicate, TExprContext& ctx, TPositionHandle pos)
{
auto parameters = ExtractComparisonParameters(predicate, ctx, pos);
if (parameters.empty()) {
return NullNode;
}

return ComparisonPushdown(parameters, predicate, ctx, pos);
}

// TODO: Check how to reduce columns if they are not needed. Unfortunately columnshard need columns list
// for every column present in program even if it is not used in result set.
//#define ENABLE_COLUMNS_PRUNING
Expand Down Expand Up @@ -419,8 +456,6 @@ TMaybeNode<TExprBase> SafeCastPredicatePushdown(const TCoFlatMap& inputFlatmap,
* FlatMap (Member(), FlatMap(SafeCast(), Just(Comparison))
* FlatMap (SafeCast(), FlatMap(SafeCast(), Just(Comparison))
*/
TVector<std::pair<TExprBase, TExprBase>> out;

auto left = ConvertComparisonNode(inputFlatmap.Input(), ctx, pos);
if (left.empty()) {
return NullNode;
Expand All @@ -434,7 +469,7 @@ TMaybeNode<TExprBase> SafeCastPredicatePushdown(const TCoFlatMap& inputFlatmap,

auto predicate = flatmap.Lambda().Body().Cast<TCoJust>().Input().Cast<TCoCompare>();

TVector<std::pair<TExprBase, TExprBase>> parameters;
std::vector<std::pair<TExprBase, TExprBase>> parameters;
if (left.size() != right.size()) {
return NullNode;
}
Expand All @@ -446,18 +481,16 @@ TMaybeNode<TExprBase> SafeCastPredicatePushdown(const TCoFlatMap& inputFlatmap,
return ComparisonPushdown(parameters, predicate, ctx, pos);
}

TMaybeNode<TExprBase> SimplePredicatePushdown(const TCoCompare& predicate, TExprContext& ctx, TPositionHandle pos)
TMaybeNode<TExprBase> CoalescePushdown(const TCoCoalesce& coalesce, TExprContext& ctx, TPositionHandle pos)
{
auto parameters = ExtractComparisonParameters(predicate, ctx, pos);
if (parameters.empty()) {
return NullNode;
if constexpr (NSsa::RuntimeVersion >= 4U) {
if (!FindNode(coalesce.Ptr(), [](const TExprNode::TPtr& node) { return TCoJsonValue::Match(node.Get()); })) {
if (const auto node = YqlCoalescePushdown(coalesce, ctx)) {
return node;
}
}
}

return ComparisonPushdown(parameters, predicate, ctx, pos);
}

TMaybeNode<TExprBase> CoalescePushdown(const TCoCoalesce& coalesce, TExprContext& ctx, TPositionHandle pos)
{
auto predicate = coalesce.Predicate();
if (auto maybeFlatmap = predicate.Maybe<TCoFlatMap>()) {
return SafeCastPredicatePushdown(maybeFlatmap.Cast(), ctx, pos);
Expand Down
Loading