44#include < ydb/library/yql/core/yql_opt_utils.h>
55#include < ydb/library/yql/core/yql_expr_type_annotation.h>
66#include < ydb/library/yql/core/yql_expr_optimize.h>
7+ #include < ydb/library/yql/core/yql_type_helpers.h>
78
89#include < ydb/library/yql/utils/log/log.h>
910
@@ -130,8 +131,158 @@ TExprNode::TListType OriginalJoinOutputMembers(const TDqPhyMapJoin& mapJoin, TEx
130131 }
131132 return structMembers;
132133}
134+
135+ TExprNode::TPtr ExpandJoinInput (const TStructExprType& type, TExprNode::TPtr&& arg, TExprContext& ctx) {
136+ return ctx.Builder (arg->Pos ())
137+ .Callable (" ExpandMap"
138+ .Add (0 , std::move (arg))
139+ .Lambda (1 )
140+ .Param (" item"
141+ .Do ([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
142+ auto i = 0U ;
143+ for (const auto & item : type.GetItems ()) {
144+ parent.Callable (i++, " Member"
145+ .Arg (0 , " item"
146+ .Atom (1 , item->GetName ())
147+ .Seal ();
148+ }
149+ return parent;
150+ })
151+ .Seal ()
152+ .Seal ().Build ();
153+ }
154+
133155} // anonymous namespace end
134156
157+ TExprBase DqPeepholeRewriteMapJoinWithGraceCore (const TExprBase& node, TExprContext& ctx) {
158+ if (!node.Maybe <TDqPhyGraceJoin>()) {
159+ return node;
160+ }
161+ const auto graceJoin = node.Cast <TDqPhyGraceJoin>();
162+ const auto pos = graceJoin.Pos ();
163+
164+ const TString leftTableLabel (GetTableLabel (graceJoin.LeftLabel ()));
165+ const TString rightTableLabel (GetTableLabel (graceJoin.RightLabel ()));
166+
167+ auto [leftKeyColumnNodes, rightKeyColumnNodes] = JoinKeysToAtoms (ctx, graceJoin, leftTableLabel, rightTableLabel);
168+ const auto keyWidth = leftKeyColumnNodes.size ();
169+
170+ ui32 outputIndex = 0 ;
171+ const auto makeRenames = [&ctx, &outputIndex, pos](TStringBuf, const TStructExprType& type) {
172+ TExprNode::TListType renames;
173+ for (auto i = 0u ; i < type.GetSize (); i++) {
174+ renames.emplace_back (ctx.NewAtom (pos, ctx.GetIndexAsString (i)));
175+ renames.emplace_back (ctx.NewAtom (pos, ctx.GetIndexAsString (outputIndex++)));
176+ }
177+ return renames;
178+ };
179+
180+ const auto itemTypeLeft = GetSequenceItemType (graceJoin.LeftInput (), false , ctx)->Cast <TStructExprType>();
181+ const auto itemTypeRight = GetSequenceItemType (graceJoin.RightInput (), false , ctx)->Cast <TStructExprType>();
182+
183+ std::vector<TString> fullColNames;
184+
185+ for (auto i = 0u ; i < itemTypeLeft->GetSize (); i++) {
186+ TString name (itemTypeLeft->GetItems ()[i]->GetName ());
187+ if (leftTableLabel) {
188+ name = leftTableLabel + " ."
189+ }
190+ fullColNames.push_back (name);
191+ }
192+ for (auto i = 0u ; i < itemTypeRight->GetSize (); i++) {
193+ TString name (itemTypeRight->GetItems ()[i]->GetName ());
194+ if (rightTableLabel) {
195+ name = rightTableLabel + " ."
196+ }
197+ fullColNames.push_back (name);
198+ }
199+
200+ TExprNode::TListType leftRenames = makeRenames (leftTableLabel, *itemTypeLeft);
201+ TExprNode::TListType rightRenames, rightPayloads;
202+ const bool withRightSide = graceJoin.JoinType ().Value () != " LeftOnly" JoinType ().Value () != " LeftSemi"
203+ if (withRightSide) {
204+ rightRenames = makeRenames (rightTableLabel, *itemTypeRight);
205+ rightPayloads.reserve (rightRenames.size () >> 1U );
206+ for (auto it = rightRenames.cbegin (); rightRenames.cend () != it; ++++it)
207+ rightPayloads.emplace_back (*it);
208+ }
209+
210+ TTypeAnnotationNode::TListType keyTypesLeft (keyWidth);
211+ TTypeAnnotationNode::TListType keyTypesRight (keyWidth);
212+ TTypeAnnotationNode::TListType keyTypes (keyWidth);
213+ for (auto i = 0U ; i < keyTypes.size (); ++i) {
214+ const auto keyTypeLeft = itemTypeLeft->FindItemType (leftKeyColumnNodes[i]->Content ());
215+ const auto keyTypeRight = itemTypeRight->FindItemType (rightKeyColumnNodes[i]->Content ());
216+ bool optKey = false ;
217+ keyTypes[i] = JoinDryKeyType (keyTypeLeft, keyTypeRight, optKey, ctx);
218+ if (!keyTypes[i]) {
219+ keyTypes.clear ();
220+ keyTypesLeft.clear ();
221+ keyTypesRight.clear ();
222+ break ;
223+ }
224+ keyTypesLeft[i] = optKey ? ctx.MakeType <TOptionalExprType>(keyTypes[i]) : keyTypes[i];
225+ keyTypesRight[i] = optKey ? ctx.MakeType <TOptionalExprType>(keyTypes[i]) : keyTypes[i];
226+ }
227+
228+ auto leftInput = ExpandJoinInput (*itemTypeLeft, ctx.NewCallable (graceJoin.LeftInput ().Pos (), " ToFlow" LeftInput ().Ptr ()}), ctx);
229+ auto rightInput = ExpandJoinInput (*itemTypeRight, ctx.NewCallable (graceJoin.RightInput ().Pos (), " ToFlow" RightInput ().Ptr ()}), ctx);
230+ YQL_ENSURE (!keyTypes.empty ());
231+
232+ for (auto i = 0U ; i < leftKeyColumnNodes.size (); i++) {
233+ const auto origName = TString (leftKeyColumnNodes[i]->Content ());
234+ auto index = itemTypeLeft->FindItem (origName);
235+ YQL_ENSURE (index);
236+ leftKeyColumnNodes[i] = ctx.NewAtom (leftKeyColumnNodes[i]->Pos (), ctx.GetIndexAsString (*index));
237+ }
238+ for (auto i = 0U ; i < rightKeyColumnNodes.size (); i++) {
239+ const auto origName = TString (rightKeyColumnNodes[i]->Content ());
240+ auto index = itemTypeRight->FindItem (origName);
241+ YQL_ENSURE (index);
242+ rightKeyColumnNodes[i] = ctx.NewAtom (rightKeyColumnNodes[i]->Pos (), ctx.GetIndexAsString (*index));
243+ }
244+
245+ auto [leftKeyColumnNodesCopy, rightKeyColumnNodesCopy] = JoinKeysToAtoms (ctx, graceJoin, leftTableLabel, rightTableLabel);
246+
247+ auto graceJoinCore = Build<TCoGraceJoinCore>(ctx, pos)
248+ .LeftInput (std::move (leftInput))
249+ .RightInput (std::move (rightInput))
250+ .JoinKind (graceJoin.JoinType ())
251+ .LeftKeysColumns (ctx.NewList (pos, std::move (leftKeyColumnNodes)))
252+ .RightKeysColumns (ctx.NewList (pos, std::move (rightKeyColumnNodes)))
253+ .LeftRenames (ctx.NewList (pos, std::move (leftRenames)))
254+ .RightRenames (ctx.NewList (pos, std::move (rightRenames)))
255+ .LeftKeysColumnNames (ctx.NewList (pos, std::move (leftKeyColumnNodesCopy)))
256+ .RightKeysColumnNames (ctx.NewList (pos, std::move (rightKeyColumnNodesCopy)))
257+ .Flags ()
258+ .Build ()
259+ .Done ();
260+
261+ auto graceNode = ctx.Builder (pos)
262+ .Callable (" NarrowMap"
263+ .Add (0 , graceJoinCore.Ptr ())
264+ .Lambda (1 )
265+ .Params (" output" size ())
266+ .Callable (" AsStruct"
267+ .Do ([&](TExprNodeBuilder& parent) -> TExprNodeBuilder& {
268+ ui32 i = 0U ;
269+ for (const auto & colName : fullColNames) {
270+ parent.List (i)
271+ .Atom (0 , colName)
272+ .Arg (1 , " output"
273+ .Seal ();
274+ i++;
275+ }
276+ return parent;
277+ })
278+ .Seal ()
279+ .Seal ()
280+ .Seal ()
281+ .Build ();
282+
283+ return TExprBase (graceNode);
284+ }
285+
135286/* *
136287 * Rewrites a `KqpMapJoin` to the `MapJoinCore`.
137288 *
@@ -142,10 +293,11 @@ TExprNode::TListType OriginalJoinOutputMembers(const TDqPhyMapJoin& mapJoin, TEx
142293 * (rely on the fact that there will be only one element in the `FlatMap`-stream)
143294 * - Align key types using `StrictCast`, use internal columns to store converted left keys
144295 */
145- TExprBase DqPeepholeRewriteMapJoin (const TExprBase& node, TExprContext& ctx) {
296+ TExprBase DqPeepholeRewriteMapJoinWithMapCore (const TExprBase& node, TExprContext& ctx) {
146297 if (!node.Maybe <TDqPhyMapJoin>()) {
147298 return node;
148299 }
300+
149301 const auto mapJoin = node.Cast <TDqPhyMapJoin>();
150302 const auto pos = mapJoin.Pos ();
151303
0 commit comments