Skip to content

Commit c631662

Browse files
igormunkinigormunkin
authored
Merge 568b6c8 into a1bcf08 
2 parents a1bcf08 + 568b6c8 commit c631662

File tree

17 files changed

+154
-56
lines changed

17 files changed

+154
-56
lines changed

‎ydb/library/yql/core/peephole_opt/yql_opt_peephole_physical.cpp

+24-1
Original file line numberDiff line numberDiff line change
@@ -5570,7 +5570,7 @@ bool CollectBlockRewrites(const TMultiExprType* multiInputType, bool keepInputCo
55705570
std::string_view arrowFunctionName;
55715571
const bool rewriteAsIs = node->IsCallable({"AssumeStrict", "AssumeNonStrict", "Likely"});
55725572
if (node->IsList() || rewriteAsIs ||
5573-
node->IsCallable({"And", "Or", "Xor", "Not", "Coalesce", "Exists", "If", "Just", "Member", "Nth", "ToPg", "FromPg", "PgResolvedCall", "PgResolvedOp"}))
5573+
node->IsCallable({"And", "Or", "Xor", "Not", "Coalesce", "Exists", "If", "Just", "AsStruct", "Member", "Nth", "ToPg", "FromPg", "PgResolvedCall", "PgResolvedOp"}))
55745574
{
55755575
if (node->IsCallable() && !IsSupportedAsBlockType(node->Pos(), *node->GetTypeAnn(), ctx, types)) {
55765576
return true;
@@ -5609,6 +5609,29 @@ bool CollectBlockRewrites(const TMultiExprType* multiInputType, bool keepInputCo
56095609
}
56105610
}
56115611

5612+
// <AsStruct> arguments (i.e. members of the resulting structure)
5613+
// are literal tuples, that don't propagate their child rewrites.
5614+
// Hence, process these rewrites the following way: wrap the
5615+
// complete expressions, supported by the block engine, with
5616+
// <AsScalar> callable or apply the rewrite of one is found.
5617+
// Otherwise, abort this <AsStruct> rewrite, since one of its
5618+
// arguments is neither block nor scalar.
5619+
if (node->IsCallable("AsStruct")) {
5620+
for (ui32 index = 0; index < node->ChildrenSize(); index++) {
5621+
auto member = funcArgs[index];
5622+
auto child = member->TailPtr();
5623+
TExprNodePtr rewrite;
5624+
if (child->IsComplete() && IsSupportedAsBlockType(child->Pos(), *child->GetTypeAnn(), ctx, types)) {
5625+
rewrite = ctx.NewCallable(child->Pos(), "AsScalar", { child });
5626+
} else if (auto rit = rewrites.find(child.Get()); rit != rewrites.end()) {
5627+
rewrite = rit->second;
5628+
} else {
5629+
return true;
5630+
}
5631+
funcArgs[index] = ctx.NewList(member->Pos(), {member->HeadPtr(), rewrite});
5632+
}
5633+
}
5634+
56125635
const TString blockFuncName = rewriteAsIs ? ToString(node->Content()) :
56135636
(TString("Block") + (node->IsList() ? "AsTuple" : node->Content()));
56145637
if (node->IsCallable({"And", "Or", "Xor"}) && funcArgs.size() > 2) {

‎ydb/library/yql/core/type_ann/type_ann_blocks.cpp

+36
Original file line numberDiff line numberDiff line change
@@ -400,6 +400,42 @@ IGraphTransformer::TStatus BlockJustWrapper(const TExprNode::TPtr& input, TExprN
400400
return IGraphTransformer::TStatus::Ok;
401401
}
402402

403+
IGraphTransformer::TStatus BlockAsStructWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
404+
Y_UNUSED(output);
405+
if (!EnsureMinArgsCount(*input, 1, ctx.Expr)) {
406+
return IGraphTransformer::TStatus::Error;
407+
}
408+
409+
TVector<const TItemExprType*> members;
410+
bool onlyScalars = true;
411+
for (auto& child : input->Children()) {
412+
auto nameNode = child->Child(0);
413+
if (!EnsureAtom(*nameNode, ctx.Expr)) {
414+
return IGraphTransformer::TStatus::Error;
415+
}
416+
auto valueNode = child->Child(1);
417+
if (!EnsureBlockOrScalarType(*valueNode, ctx.Expr)) {
418+
return IGraphTransformer::TStatus::Error;
419+
}
420+
421+
bool isScalar;
422+
const TTypeAnnotationNode* blockItemType = GetBlockItemType(*valueNode->GetTypeAnn(), isScalar);
423+
424+
onlyScalars = onlyScalars && isScalar;
425+
members.push_back(ctx.Expr.MakeType<TItemExprType>(nameNode->Content(), blockItemType));
426+
}
427+
428+
const TTypeAnnotationNode* resultType = ctx.Expr.MakeType<TStructExprType>(members);
429+
430+
if (onlyScalars) {
431+
resultType = ctx.Expr.MakeType<TScalarExprType>(resultType);
432+
} else {
433+
resultType = ctx.Expr.MakeType<TBlockExprType>(resultType);
434+
}
435+
input->SetTypeAnn(resultType);
436+
return IGraphTransformer::TStatus::Ok;
437+
}
438+
403439
IGraphTransformer::TStatus BlockAsTupleWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx) {
404440
Y_UNUSED(output);
405441
if (!EnsureMinArgsCount(*input, 1, ctx.Expr)) {

‎ydb/library/yql/core/type_ann/type_ann_blocks.h

+1
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ namespace NTypeAnnImpl {
1818
IGraphTransformer::TStatus BlockLogicalWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
1919
IGraphTransformer::TStatus BlockIfWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
2020
IGraphTransformer::TStatus BlockJustWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
21+
IGraphTransformer::TStatus BlockAsStructWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
2122
IGraphTransformer::TStatus BlockAsTupleWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
2223
IGraphTransformer::TStatus BlockNthWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);
2324
IGraphTransformer::TStatus BlockMemberWrapper(const TExprNode::TPtr& input, TExprNode::TPtr& output, TContext& ctx);

‎ydb/library/yql/core/type_ann/type_ann_core.cpp

+1
Original file line numberDiff line numberDiff line change
@@ -12286,6 +12286,7 @@ template <NKikimr::NUdf::EDataSlot DataSlot>
1228612286
Functions["BlockNot"] = &BlockLogicalWrapper;
1228712287
Functions["BlockIf"] = &BlockIfWrapper;
1228812288
Functions["BlockJust"] = &BlockJustWrapper;
12289+
Functions["BlockAsStruct"] = &BlockAsStructWrapper;
1228912290
Functions["BlockAsTuple"] = &BlockAsTupleWrapper;
1229012291
Functions["BlockMember"] = &BlockMemberWrapper;
1229112292
Functions["BlockNth"] = &BlockNthWrapper;

‎ydb/library/yql/minikql/comp_nodes/mkql_block_tuple.cpp renamed to ‎ydb/library/yql/minikql/comp_nodes/mkql_block_container.cpp

+9-9
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#include "mkql_block_tuple.h"
1+
#include "mkql_block_container.h"
22

33
#include <ydb/library/yql/minikql/computation/mkql_block_impl.h>
44

@@ -15,9 +15,9 @@ namespace NMiniKQL {
1515

1616
namespace {
1717

18-
class TBlockAsTupleExec {
18+
class TBlockAsContainerExec {
1919
public:
20-
TBlockAsTupleExec(const TVector<TType*>& argTypes, const std::shared_ptr<arrow::DataType>& returnArrowType)
20+
TBlockAsContainerExec(const TVector<TType*>& argTypes, const std::shared_ptr<arrow::DataType>& returnArrowType)
2121
: ArgTypes(argTypes)
2222
, ReturnArrowType(returnArrowType)
2323
{}
@@ -66,10 +66,10 @@ class TBlockAsTupleExec {
6666
const std::shared_ptr<arrow::DataType> ReturnArrowType;
6767
};
6868

69-
std::shared_ptr<arrow::compute::ScalarKernel> MakeBlockAsTupleKernel(const TVector<TType*>& argTypes, TType* resultType) {
69+
std::shared_ptr<arrow::compute::ScalarKernel> MakeBlockAsContainerKernel(const TVector<TType*>& argTypes, TType* resultType) {
7070
std::shared_ptr<arrow::DataType> returnArrowType;
7171
MKQL_ENSURE(ConvertArrowType(AS_TYPE(TBlockType, resultType)->GetItemType(), returnArrowType), "Unsupported arrow type");
72-
auto exec = std::make_shared<TBlockAsTupleExec>(argTypes, returnArrowType);
72+
auto exec = std::make_shared<TBlockAsContainerExec>(argTypes, returnArrowType);
7373
auto kernel = std::make_shared<arrow::compute::ScalarKernel>(ConvertToInputTypes(argTypes), ConvertToOutputType(resultType),
7474
[exec](arrow::compute::KernelContext* ctx, const arrow::compute::ExecBatch& batch, arrow::Datum* res) {
7575
return exec->Exec(ctx, batch, res);
@@ -81,17 +81,17 @@ std::shared_ptr<arrow::compute::ScalarKernel> MakeBlockAsTupleKernel(const TVect
8181

8282
} // namespace
8383

84-
IComputationNode* WrapBlockAsTuple(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
84+
IComputationNode* WrapBlockAsContainer(TCallable& callable, const TComputationNodeFactoryContext& ctx) {
8585
TComputationNodePtrVector argsNodes;
8686
TVector<TType*> argsTypes;
8787
for (ui32 i = 0; i < callable.GetInputsCount(); ++i) {
8888
argsNodes.push_back(LocateNode(ctx.NodeLocator, callable, i));
8989
argsTypes.push_back(callable.GetInput(i).GetStaticType());
9090
}
9191

92-
auto kernel = MakeBlockAsTupleKernel(argsTypes, callable.GetType()->GetReturnType());
92+
auto kernel = MakeBlockAsContainerKernel(argsTypes, callable.GetType()->GetReturnType());
9393
return new TBlockFuncNode(ctx.Mutables, callable.GetType()->GetName(), std::move(argsNodes), argsTypes, *kernel, kernel);
9494
}
9595

96-
}
97-
}
96+
} // namespace NMiniKQL
97+
} // namespace NKikimr

‎ydb/library/yql/minikql/comp_nodes/mkql_block_container.h

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
#pragma once
2+
#include <ydb/library/yql/minikql/computation/mkql_computation_node.h>
3+
4+
namespace NKikimr {
5+
namespace NMiniKQL {
6+
7+
IComputationNode* WrapBlockAsContainer(TCallable& callable, const TComputationNodeFactoryContext& ctx);
8+
9+
} // namespace NMiniKQL
10+
} // namespace NKikimr

‎ydb/library/yql/minikql/comp_nodes/mkql_block_tuple.h

-10
This file was deleted.

‎ydb/library/yql/minikql/comp_nodes/mkql_factory.cpp

+3-2
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include "mkql_blocks.h"
99
#include "mkql_block_agg.h"
1010
#include "mkql_block_coalesce.h"
11+
#include "mkql_block_container.h"
1112
#include "mkql_block_exists.h"
1213
#include "mkql_block_getelem.h"
1314
#include "mkql_block_if.h"
@@ -16,7 +17,6 @@
1617
#include "mkql_block_compress.h"
1718
#include "mkql_block_skiptake.h"
1819
#include "mkql_block_top.h"
19-
#include "mkql_block_tuple.h"
2020
#include "mkql_callable.h"
2121
#include "mkql_chain_map.h"
2222
#include "mkql_chain1_map.h"
@@ -297,7 +297,8 @@ struct TCallableComputationNodeBuilderFuncMapFiller {
297297
{"BlockNot", &WrapBlockNot},
298298
{"BlockJust", &WrapBlockJust},
299299
{"BlockCompress", &WrapBlockCompress},
300-
{"BlockAsTuple", &WrapBlockAsTuple},
300+
{"BlockAsTuple", &WrapBlockAsContainer},
301+
{"BlockAsStruct", &WrapBlockAsContainer},
301302
{"BlockMember", &WrapBlockMember},
302303
{"BlockNth", &WrapBlockNth},
303304
{"BlockExpandChunked", &WrapBlockExpandChunked},

‎ydb/library/yql/minikql/comp_nodes/ya.make.inc

+1-1
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ SET(ORIG_SOURCES
1515
mkql_block_agg_some.cpp
1616
mkql_block_agg_sum.cpp
1717
mkql_block_coalesce.cpp
18+
mkql_block_container.cpp
1819
mkql_block_exists.cpp
1920
mkql_block_getelem.cpp
2021
mkql_block_if.cpp
@@ -24,7 +25,6 @@ SET(ORIG_SOURCES
2425
mkql_block_func.cpp
2526
mkql_block_skiptake.cpp
2627
mkql_block_top.cpp
27-
mkql_block_tuple.cpp
2828
mkql_blocks.cpp
2929
mkql_callable.cpp
3030
mkql_chain_map.cpp

‎ydb/library/yql/minikql/mkql_program_builder.cpp

+22
Original file line numberDiff line numberDiff line change
@@ -1657,6 +1657,28 @@ TRuntimeNode TProgramBuilder::BlockNth(TRuntimeNode tuple, ui32 index) {
16571657
return TRuntimeNode(callableBuilder.Build(), false);
16581658
}
16591659

1660+
TRuntimeNode TProgramBuilder::BlockAsStruct(const TArrayRef<std::pair<std::string_view, TRuntimeNode>>& args) {
1661+
MKQL_ENSURE(!args.empty(), "Expected at least one argument");
1662+
1663+
TBlockType::EShape resultShape = TBlockType::EShape::Scalar;
1664+
TVector<std::pair<std::string_view, TType*>> members;
1665+
for (const auto& x : args) {
1666+
auto blockType = AS_TYPE(TBlockType, x.second.GetStaticType());
1667+
members.emplace_back(x.first, blockType->GetItemType());
1668+
if (blockType->GetShape() == TBlockType::EShape::Many) {
1669+
resultShape = TBlockType::EShape::Many;
1670+
}
1671+
}
1672+
1673+
auto returnType = NewBlockType(NewStructType(members), resultShape);
1674+
TCallableBuilder callableBuilder(Env, __func__, returnType);
1675+
for (const auto& x : args) {
1676+
callableBuilder.Add(x.second);
1677+
}
1678+
1679+
return TRuntimeNode(callableBuilder.Build(), false);
1680+
}
1681+
16601682
TRuntimeNode TProgramBuilder::BlockAsTuple(const TArrayRef<const TRuntimeNode>& args) {
16611683
MKQL_ENSURE(!args.empty(), "Expected at least one argument");
16621684

‎ydb/library/yql/minikql/mkql_program_builder.h

+1
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,7 @@ class TProgramBuilder : public TTypeBuilder {
239239
TRuntimeNode BlockExists(TRuntimeNode data);
240240
TRuntimeNode BlockMember(TRuntimeNode structure, const std::string_view& memberName);
241241
TRuntimeNode BlockNth(TRuntimeNode tuple, ui32 index);
242+
TRuntimeNode BlockAsStruct(const TArrayRef<std::pair<std::string_view, TRuntimeNode>>& args);
242243
TRuntimeNode BlockAsTuple(const TArrayRef<const TRuntimeNode>& args);
243244
TRuntimeNode BlockToPg(TRuntimeNode input, TType* returnType);
244245
TRuntimeNode BlockFromPg(TRuntimeNode input, TType* returnType);

‎ydb/library/yql/providers/common/mkql/yql_provider_mkql.cpp

+8
Original file line numberDiff line numberDiff line change
@@ -2737,6 +2737,14 @@ TMkqlCommonCallableCompiler::TShared::TShared() {
27372737
return ctx.ProgramBuilder.BlockNth(tupleObj, index);
27382738
});
27392739

2740+
AddCallable("BlockAsStruct", [](const TExprNode& node, TMkqlBuildContext& ctx) {
2741+
std::vector<std::pair<std::string_view, TRuntimeNode>> members;
2742+
for (const auto& x : node.Children()) {
2743+
members.emplace_back(x->Head().Content(), MkqlBuildExpr(x->Tail(), ctx));
2744+
}
2745+
return ctx.ProgramBuilder.BlockAsStruct(members);
2746+
});
2747+
27402748
AddCallable("BlockAsTuple", [](const TExprNode& node, TMkqlBuildContext& ctx) {
27412749
TVector<TRuntimeNode> args;
27422750
for (const auto& x : node.Children()) {

‎ydb/library/yql/tests/sql/dq_file/part11/canondata/result.json

+9-9
Original file line numberDiff line numberDiff line change
@@ -535,23 +535,23 @@
535535
"test.test[blocks-filter_direct_col--Results]": [],
536536
"test.test[blocks-member--Analyze]": [
537537
{
538-
"checksum": "b08274fd137c1878d90520c832f06fd3",
539-
"size": 3676,
540-
"uri": "https://{canondata_backend}/1889210/75a1d72834c0a9de8b328ec130be934f6cc6cea0/resource.tar.gz#test.test_blocks-member--Analyze_/plan.txt"
538+
"checksum": "4d80733bb5655340645b981057ba9910",
539+
"size": 3703,
540+
"uri": "https://{canondata_backend}/1942525/5635585a917e888e1628404d5ff137a2e18f23ca/resource.tar.gz#test.test_blocks-member--Analyze_/plan.txt"
541541
}
542542
],
543543
"test.test[blocks-member--Debug]": [
544544
{
545-
"checksum": "a30b76ba380ee4f694dc731aa2771fed",
546-
"size": 1467,
547-
"uri": "https://{canondata_backend}/1937027/96028d31f8e29253c9276a86f02284e2a71add76/resource.tar.gz#test.test_blocks-member--Debug_/opt.yql_patched"
545+
"checksum": "be5b35d7624905acc850940a1ff74780",
546+
"size": 1837,
547+
"uri": "https://{canondata_backend}/1775319/6624c18402d2e5473f3dcf5d9248a5e624496fd5/resource.tar.gz#test.test_blocks-member--Debug_/opt.yql_patched"
548548
}
549549
],
550550
"test.test[blocks-member--Plan]": [
551551
{
552-
"checksum": "b08274fd137c1878d90520c832f06fd3",
553-
"size": 3676,
554-
"uri": "https://{canondata_backend}/1889210/75a1d72834c0a9de8b328ec130be934f6cc6cea0/resource.tar.gz#test.test_blocks-member--Plan_/plan.txt"
552+
"checksum": "4d80733bb5655340645b981057ba9910",
553+
"size": 3703,
554+
"uri": "https://{canondata_backend}/1942525/5635585a917e888e1628404d5ff137a2e18f23ca/resource.tar.gz#test.test_blocks-member--Plan_/plan.txt"
555555
}
556556
],
557557
"test.test[blocks-member--Results]": [],

‎ydb/library/yql/tests/sql/hybrid_file/part10/canondata/result.json

+6-6
Original file line numberDiff line numberDiff line change
@@ -673,16 +673,16 @@
673673
],
674674
"test.test[blocks-member--Debug]": [
675675
{
676-
"checksum": "49012dc14c1d467d864bc5079aa4542f",
677-
"size": 1982,
678-
"uri": "https://{canondata_backend}/1597364/2a4f282ea286021ee8f9098fb8207324c7ebe684/resource.tar.gz#test.test_blocks-member--Debug_/opt.yql_patched"
676+
"checksum": "282430bc5d30ff486ec087f2577eadd7",
677+
"size": 2755,
678+
"uri": "https://{canondata_backend}/1689644/13287ae9d02a4e7590ca31d651f8da70af99729d/resource.tar.gz#test.test_blocks-member--Debug_/opt.yql_patched"
679679
}
680680
],
681681
"test.test[blocks-member--Plan]": [
682682
{
683-
"checksum": "794e5e7aaffc457f9e8f888953e1c89e",
684-
"size": 4045,
685-
"uri": "https://{canondata_backend}/1597364/07eb39555ae99a903261332d5998f32599b281fe/resource.tar.gz#test.test_blocks-member--Plan_/plan.txt"
683+
"checksum": "dd2dd3d0cedb748dd1909eb4bfb6cfc5",
684+
"size": 4072,
685+
"uri": "https://{canondata_backend}/1597364/81dc17780dff76af9a3a69947365d6afbefb3093/resource.tar.gz#test.test_blocks-member--Plan_/plan.txt"
686686
}
687687
],
688688
"test.test[blocks-minmax_strings--Debug]": [

‎ydb/library/yql/tests/sql/sql2yql/canondata/result.json

+6-6
Original file line numberDiff line numberDiff line change
@@ -3578,9 +3578,9 @@
35783578
],
35793579
"test_sql2yql.test[blocks-member]": [
35803580
{
3581-
"checksum": "88c9131d7ea3c80ecc268cb7f458fc86",
3582-
"size": 1230,
3583-
"uri": "https://{canondata_backend}/1937027/f015781f1ee8e2e10d049f80211c1f81b56abfb9/resource.tar.gz#test_sql2yql.test_blocks-member_/sql.yql"
3581+
"checksum": "801120991f062aeba5589bf360ea68b4",
3582+
"size": 1781,
3583+
"uri": "https://{canondata_backend}/1775319/914ef4c98185f02d24bc7021ac55e78c12bb3586/resource.tar.gz#test_sql2yql.test_blocks-member_/sql.yql"
35843584
}
35853585
],
35863586
"test_sql2yql.test[blocks-minmax_strings]": [
@@ -21904,9 +21904,9 @@
2190421904
],
2190521905
"test_sql_format.test[blocks-member]": [
2190621906
{
21907-
"checksum": "7fcb44569c4f84aee80ea32b5961e0ed",
21908-
"size": 146,
21909-
"uri": "https://{canondata_backend}/1937027/f015781f1ee8e2e10d049f80211c1f81b56abfb9/resource.tar.gz#test_sql_format.test_blocks-member_/formatted.sql"
21907+
"checksum": "a516f2fe71075d584f4e29663ede2562",
21908+
"size": 381,
21909+
"uri": "https://{canondata_backend}/1937424/630ba1f00b4d62293128d029101c2a4bb885b850/resource.tar.gz#test_sql_format.test_blocks-member_/formatted.sql"
2191021910
}
2191121911
],
2191221912
"test_sql_format.test[blocks-minmax_strings]": [

‎ydb/library/yql/tests/sql/suites/blocks/member.sql

+5
Original file line numberDiff line numberDiff line change
@@ -4,4 +4,9 @@ pragma UseBlocks;
44

55
SELECT
66
val.a as a,
7+
<|qq:key,qkrq:"QKRQ"|> as q,
8+
/* XXX: <AddMember> callable always expands to <AsStruct>. */
9+
AddMember(val, "k", key) as wik,
10+
/* XXX: <RemoveMember> callable always expands to <AsStruct>. */
11+
RemoveMember(val, "x") as wox,
712
FROM Input;

0 commit comments

Comments
 (0)