Merge 6315ad7 into db1f75f

yumkam · web-flow · commit a7b411f4c28b · 2025-01-17T10:05:39.000Z
diff --git a/ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup.cpp b/ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup.cpp
@@ -522,6 +522,16 @@ std::pair<
             } else {
                 result[i] = { EOutputRowItemSource::LookupOther, lookupPayloadColumns.at(name) };
             }
+        } else if (leftLabel.empty()) {
+            const auto name = prefixedName;
+            if (auto j = leftJoinColumns.FindPtr(name)) {
+                result[i] = { EOutputRowItemSource::InputKey, lookupKeyColumns.at(rightNames[*j]) };
+            } else if (auto k = inputColumns.FindPtr(name)) {
+                result[i] = { EOutputRowItemSource::InputOther, otherInputIndexes.size() };
+                otherInputIndexes.push_back(*k);
+            } else {
+                Y_ABORT();
+            }
         } else {
             Y_ABORT();
         }
diff --git a/ydb/library/yql/dq/type_ann/dq_type_ann.cpp b/ydb/library/yql/dq/type_ann/dq_type_ann.cpp
@@ -625,8 +625,21 @@ TStatus AnnotateDqCnStreamLookup(const TExprNode::TPtr& input, TExprContext& ctx
         ctx.AddError(TIssue(ctx.GetPosition(joinType.Pos()), "Streamlookup supports only LEFT JOIN ... ANY"));
         return TStatus::Error;
     }
+    auto rightInput = cnStreamLookup.RightInput();
+    if (!rightInput.Raw()->IsCallable("TDqLookupSourceWrap")) {
+        ctx.AddError(TIssue(ctx.GetPosition(rightInput.Pos()), "Unknown DqCnStreamLookup RightInput type (must be TDqLookupSourceWrap)"));
+        return TStatus::Error;
+    }
     const auto leftRowType = GetSeqItemType(leftInputType);
-    const auto rightRowType = GetSeqItemType(cnStreamLookup.RightInput().Raw()->GetTypeAnn());
+    if (!leftRowType) {
+        ctx.AddError(TIssue(ctx.GetPosition(rightInput.Pos()), "DqCnStreamLookup: Failed to annotate left row type"));
+        return TStatus::Error;
+    }
+    const auto rightRowType = GetSeqItemType(rightInput.Raw()->GetTypeAnn());
+    if (!rightRowType) {
+        ctx.AddError(TIssue(ctx.GetPosition(rightInput.Pos()), "DqCnStreamLookup: Failed to annotate right row type"));
+        return TStatus::Error;
+    }
     const auto outputRowType = GetDqJoinResultType<true>(
         input->Pos(),
         *leftRowType->Cast<TStructExprType>(),
@@ -637,6 +650,10 @@ TStatus AnnotateDqCnStreamLookup(const TExprNode::TPtr& input, TExprContext& ctx
         cnStreamLookup.JoinKeys(),
         ctx
     );
+    if (!outputRowType) {
+        ctx.AddError(TIssue(ctx.GetPosition(rightInput.Pos()), "DqCnStreamLookup: Failed to annotate output row type"));
+        return TStatus::Error;
+    }
     if (!EnsureConvertibleTo<ui64>(cnStreamLookup.MaxCachedRows().Ref(), "MaxCachedRows", ctx) ||
         !EnsureConvertibleTo<ui64>(cnStreamLookup.TTL().Ref(), "TTL", ctx) ||
         !EnsureConvertibleTo<ui64>(cnStreamLookup.MaxDelayedRows().Ref(), "MaxDelayedRows", ctx)) {
diff --git a/ydb/library/yql/providers/dq/opt/logical_optimize.cpp b/ydb/library/yql/providers/dq/opt/logical_optimize.cpp
@@ -225,36 +225,64 @@ class TDqsLogicalOptProposalTransformer : public TOptimizeTransformerBase {
         return TDqLookupSourceWrap(lookupSourceWrap);
     }
 
-    TMaybeNode<TExprBase> RewriteStreamEquiJoinWithLookup(TExprBase node, TExprContext& ctx) {
-        Y_UNUSED(ctx);
-        const auto equiJoin = node.Cast<TCoEquiJoin>();
-        if (equiJoin.ArgCount() != 4) { // 2 parties join
-            return node;
+    // Recursively walk join tree and replace right-side of StreamLookupJoin
+    ui32 RewriteStreamJoinTuple(ui32 idx, const TCoEquiJoin& equiJoin, const TCoEquiJoinTuple& joinTuple, std::vector<TExprNode::TPtr>& args, TExprContext& ctx, bool& changed) {
+        // recursion depth O(args.size())
+        Y_ENSURE(idx < args.size());
+        // handle left side
+        if (!joinTuple.LeftScope().Maybe<TCoAtom>()) {
+            idx = RewriteStreamJoinTuple(idx, equiJoin, joinTuple.LeftScope().Cast<TCoEquiJoinTuple>(), args, ctx, changed);
+        } else {
+            ++idx;
+        }
+        // handle right side
+        if (!joinTuple.RightScope().Maybe<TCoAtom>()) {
+            return RewriteStreamJoinTuple(idx, equiJoin, joinTuple.RightScope().Cast<TCoEquiJoinTuple>(), args, ctx, changed);
         }
-        const auto left = equiJoin.Arg(0).Cast<TCoEquiJoinInput>().List();
-        const auto right = equiJoin.Arg(1).Cast<TCoEquiJoinInput>().List();
-        const auto joinTuple = equiJoin.Arg(equiJoin.ArgCount() - 2).Cast<TCoEquiJoinTuple>();
+        Y_ENSURE(idx < args.size());
         if (!IsStreamLookup(joinTuple)) {
-            return node;
+            return idx + 1;
         }
-        if (!right.Maybe<TDqSourceWrap>() && !right.Maybe<TDqReadWrap>()) {
-            return node;
+        auto right = equiJoin.Arg(idx).Cast<TCoEquiJoinInput>();
+        auto rightList = right.List();
+        if (auto maybeExtractMembers = rightList.Maybe<TCoExtractMembers>()) {
+            rightList = maybeExtractMembers.Cast().Input();
         }
+        TExprNode::TPtr lookupSourceWrap;
+        if (auto maybeSource = rightList.Maybe<TDqSourceWrap>()) {
+            lookupSourceWrap = LookupSourceFromSource(maybeSource.Cast(), ctx).Ptr();
+        } else if (auto maybeRead = rightList.Maybe<TDqReadWrap>()) {
+            lookupSourceWrap = LookupSourceFromRead(maybeRead.Cast(), ctx).Ptr();
+        } else {
+            return idx + 1;
+        }
+        changed = true;
+        args[idx] =
+            Build<TCoEquiJoinInput>(ctx, joinTuple.Pos())
+                .List(lookupSourceWrap)
+                .Scope(right.Scope())
+            .Done().Ptr();
+        return idx + 1;
+    }
 
-        TDqLookupSourceWrap lookupSourceWrap =  right.Maybe<TDqSourceWrap>()
-            ? LookupSourceFromSource(right.Cast<TDqSourceWrap>(), ctx)
-            : LookupSourceFromRead(right.Cast<TDqReadWrap>(), ctx)
-        ;
-
-        return Build<TCoEquiJoin>(ctx, node.Pos())
-                .Add(equiJoin.Arg(0))
-                .Add<TCoEquiJoinInput>()
-                        .List(lookupSourceWrap)
-                        .Scope(equiJoin.Arg(1).Cast<TCoEquiJoinInput>().Scope())
-                .Build()
-                .Add(equiJoin.Arg(2))
-                .Add(equiJoin.Arg(3))
-            .Done();
+    TMaybeNode<TExprBase> RewriteStreamEquiJoinWithLookup(TExprBase node, TExprContext& ctx) {
+        const auto equiJoin = node.Cast<TCoEquiJoin>();
+        auto argCount = equiJoin.ArgCount();
+        const auto joinTuple = equiJoin.Arg(argCount - 2).Cast<TCoEquiJoinTuple>();
+        std::vector<TExprNode::TPtr> args(argCount);
+        bool changed = false;
+        auto rightIdx = RewriteStreamJoinTuple(0u, equiJoin, joinTuple, args, ctx, changed);
+        Y_ENSURE(rightIdx + 2 == argCount);
+        if (!changed) {
+            return node;
+        }
+        // fill copies of remaining args
+        for (ui32 i = 0; i < argCount; ++i) {
+            if (!args[i]) {
+                args[i] = equiJoin.Arg(i).Ptr();
+            }
+        }
+        return Build<TCoEquiJoin>(ctx, node.Pos()).Add(std::move(args)).Done();
     }
 
     TMaybeNode<TExprBase> OptimizeEquiJoinWithCosts(TExprBase node, TExprContext& ctx) {
diff --git a/ydb/library/yql/providers/dq/opt/physical_optimize.cpp b/ydb/library/yql/providers/dq/opt/physical_optimize.cpp
@@ -312,10 +312,16 @@ class TDqsPhysicalOptProposalTransformer : public TOptimizeTransformerBase {
         if (!maxDelayedRows) {
             maxDelayedRows = ctx.NewAtom(pos, 1'000'000);
         }
+        auto rightInput = join.RightInput().Ptr();
+        if (auto maybe = TExprBase(rightInput).Maybe<TCoExtractMembers>()) {
+            rightInput = maybe.Cast().Input().Ptr();
+        }
+        auto leftLabel = join.LeftLabel().Maybe<NNodes::TCoAtom>() ? join.LeftLabel().Cast<NNodes::TCoAtom>().Ptr() : ctx.NewAtom(pos, "");
+        Y_ENSURE(join.RightLabel().Maybe<NNodes::TCoAtom>());
         auto cn = Build<TDqCnStreamLookup>(ctx, pos)
             .Output(left.Output().Cast())
-            .LeftLabel(join.LeftLabel().Cast<NNodes::TCoAtom>())
-            .RightInput(join.RightInput())
+            .LeftLabel(leftLabel)
+            .RightInput(rightInput)
             .RightLabel(join.RightLabel().Cast<NNodes::TCoAtom>())
             .JoinKeys(join.JoinKeys())
             .JoinType(join.JoinType())
diff --git a/ydb/library/yql/providers/dq/planner/execution_planner.cpp b/ydb/library/yql/providers/dq/planner/execution_planner.cpp
@@ -608,7 +608,7 @@ namespace NYql::NDqs {
         settings.SetRightLabel(streamLookup.RightLabel().StringValue());
         settings.SetJoinType(streamLookup.JoinType().StringValue());
         for (const auto& k: streamLookup.LeftJoinKeyNames()) {
-            *settings.AddLeftJoinKeyNames() = RemoveAliases(k.StringValue());
+            *settings.AddLeftJoinKeyNames() = streamLookup.LeftLabel().StringValue().empty() ? k.StringValue() : RemoveAliases(k.StringValue());
         }
         for (const auto& k: streamLookup.RightJoinKeyNames()) {
             *settings.AddRightJoinKeyNames() = RemoveAliases(k.StringValue());
diff --git a/ydb/tests/fq/generic/streaming/test_join.py b/ydb/tests/fq/generic/streaming/test_join.py
@@ -412,6 +412,115 @@ def freeze(json):
             ]
         ),
     ),
+    # 8
+    (
+        R'''
+            $input = SELECT * FROM myyds.`{input_topic}`
+                    WITH (
+                        FORMAT=json_each_row,
+                        SCHEMA (
+                            za Int32,
+                            yb STRING,
+                            yc Int32,
+                            zd Int32,
+                        )
+                    )            ;
+
+            $enriched1 = select a, b, c, d, e, f, za, yb, yc, zd
+                from
+                    $input as e
+                left join {streamlookup} any ydb_conn_{table_name}.db as u
+                on(e.za = u.a AND e.yb = u.b)
+            ;
+
+            $enriched2 = SELECT e.a AS a, e.b AS b, e.c AS c, e.d AS d, e.e AS e, e.f AS f, za, yb, yc, zd, u.c AS c2, u.d AS d2
+                from
+                    $enriched1 as e
+                left join {streamlookup} any ydb_conn_{table_name}.db as u
+                on(e.za = u.a AND e.yb = u.b)
+            ;
+
+            $enriched = select a, b, c, d, e, f, za, yb, yc, zd, (c2 IS NOT DISTINCT FROM c) as eq1, (d2 IS NOT DISTINCT FROM d) as eq2
+                from
+                    $enriched2 as e
+            ;
+
+            insert into myyds.`{output_topic}`
+            select Unwrap(Yson::SerializeJson(Yson::From(TableRow()))) from $enriched;
+            ''',
+        ResequenceId(
+            [
+                (
+                    '{"id":1,"za":1,"yb":"2","yc":100,"zd":101}',
+                    '{"a":1,"b":"2","c":3,"d":4,"e":5,"f":6,"za":1,"yb":"2","yc":100,"zd":101,"eq1":true,"eq2":true}',
+                ),
+                (
+                    '{"id":2,"za":7,"yb":"8","yc":106,"zd":107}',
+                    '{"a":7,"b":"8","c":9,"d":10,"e":11,"f":12,"za":7,"yb":"8","yc":106,"zd":107,"eq1":true,"eq2":true}',
+                ),
+                (
+                    '{"id":3,"za":2,"yb":"1","yc":114,"zd":115}',
+                    '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":2,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}',
+                ),
+                (
+                    '{"id":3,"za":null,"yb":"1","yc":114,"zd":115}',
+                    '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":null,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}',
+                ),
+            ]
+        ),
+    ),
+    # 9
+    (
+        R'''
+            $input = SELECT * FROM myyds.`{input_topic}`
+                    WITH (
+                        FORMAT=json_each_row,
+                        SCHEMA (
+                            za Int32,
+                            yb STRING,
+                            yc Int32,
+                            zd Int32,
+                        )
+                    )            ;
+
+            $enriched12 = select u.a as a, u.b as b, u.c as c, u.d as d, u.e as e, u.f as f, za, yb, yc, zd, u2.c as c2, u2.d as d2
+                from
+                    $input as e
+                left join {streamlookup} any ydb_conn_{table_name}.db as u
+                on(e.za = u.a AND e.yb = u.b)
+                left join {streamlookup} any ydb_conn_{table_name}.db as u2
+                on(e.yb = u2.b AND e.za = u2.a)
+            ;
+
+            $enriched = select a, b, c, d, e, f, za, yb, yc, zd, (c2 IS NOT DISTINCT FROM c) as eq1, (d2 IS NOT DISTINCT FROM d) as eq2
+                from
+                    $enriched12 as e
+            ;
+
+            insert into myyds.`{output_topic}`
+            select Unwrap(Yson::SerializeJson(Yson::From(TableRow()))) from $enriched;
+            ''',
+        ResequenceId(
+            [
+                (
+                    '{"id":1,"za":1,"yb":"2","yc":100,"zd":101}',
+                    '{"a":1,"b":"2","c":3,"d":4,"e":5,"f":6,"za":1,"yb":"2","yc":100,"zd":101,"eq1":true,"eq2":true}',
+                ),
+                (
+                    '{"id":2,"za":7,"yb":"8","yc":106,"zd":107}',
+                    '{"a":7,"b":"8","c":9,"d":10,"e":11,"f":12,"za":7,"yb":"8","yc":106,"zd":107,"eq1":true,"eq2":true}',
+                ),
+                (
+                    '{"id":3,"za":2,"yb":"1","yc":114,"zd":115}',
+                    '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":2,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}',
+                ),
+                (
+                    '{"id":3,"za":null,"yb":"1","yc":114,"zd":115}',
+                    '{"a":null,"b":null,"c":null,"d":null,"e":null,"f":null,"za":null,"yb":"1","yc":114,"zd":115,"eq1":true,"eq2":true}',
+                ),
+            ]
+        ),
+    ),
 ]
 
 

Original file line number	Diff line number	Diff line change
`@@ -608,7 +608,7 @@ namespace NYql::NDqs {`
`608`	`608`	`settings.SetRightLabel(streamLookup.RightLabel().StringValue());`
`609`	`609`	`settings.SetJoinType(streamLookup.JoinType().StringValue());`
`610`	`610`	`for (const auto& k: streamLookup.LeftJoinKeyNames()) {`
`611`		`- *settings.AddLeftJoinKeyNames() = RemoveAliases(k.StringValue());`
	`611`	`+ *settings.AddLeftJoinKeyNames() = streamLookup.LeftLabel().StringValue().empty() ? k.StringValue() : RemoveAliases(k.StringValue());`
`612`	`612`	`}`
`613`	`613`	`for (const auto& k: streamLookup.RightJoinKeyNames()) {`
`614`	`614`	`*settings.AddRightJoinKeyNames() = RemoveAliases(k.StringValue());`