Merge 464b6ca into 580503a

Author: GrigoriyPA

ydb/core/kqp/common/kqp_tx.cpp

@@ -392,6 +392,7 @@ bool HasUncommittedChangesRead(THashSet<NKikimr::TTableId>& modifiedTables, cons
                 case NKqpProto::TKqpPhyConnection::kResult:
                 case NKqpProto::TKqpPhyConnection::kValue:
                 case NKqpProto::TKqpPhyConnection::kMerge:
+                case NKqpProto::TKqpPhyConnection::kDqSourceStreamLookup:
                 case NKqpProto::TKqpPhyConnection::TYPE_NOT_SET:
                     break;
                 }

ydb/core/kqp/common/kqp_user_request_context.h

@@ -19,6 +19,7 @@ namespace NKikimr::NKqp {
         NActors::TActorId RunScriptActorId;
         TString PoolId;
         std::optional<NResourcePool::TPoolSettings> PoolConfig;
+        bool IsStreamingQuery = false;
 
         TUserRequestContext() = default;
 

ydb/core/kqp/compile_service/kqp_compile_actor.cpp

@@ -96,6 +96,10 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
             } else {
                 Config->_ResultRowsLimit.Clear();
             }
+
+            if (UserRequestContext && UserRequestContext->IsStreamingQuery) {
+                Config->HashJoinMode = NYql::NDq::EHashJoinMode::Map;
+            }
         }
         PerStatementResult = perStatementResult && Config->EnablePerStatementQueryExecution;
 

ydb/core/kqp/compute_actor/kqp_compute_actor.cpp

@@ -6,19 +6,20 @@
 #include <ydb/core/base/appdata.h>
 #include <ydb/core/kqp/runtime/kqp_compute.h>
 #include <ydb/core/kqp/runtime/kqp_read_actor.h>
-#include <ydb/core/kqp/runtime/kqp_write_actor.h>
 #include <ydb/core/kqp/runtime/kqp_read_table.h>
 #include <ydb/core/kqp/runtime/kqp_sequencer_factory.h>
 #include <ydb/core/kqp/runtime/kqp_stream_lookup_factory.h>
 #include <ydb/core/kqp/runtime/kqp_vector_actor.h>
-#include <ydb/library/yql/providers/generic/actors/yql_generic_provider_factories.h>
+#include <ydb/core/kqp/runtime/kqp_write_actor.h>
 #include <ydb/library/formats/arrow/protos/ssa.pb.h>
+#include <ydb/library/yql/dq/actors/input_transforms/dq_input_transform_lookup_factory.h>
+#include <ydb/library/yql/dq/comp_nodes/dq_block_hash_join.h>
+#include <ydb/library/yql/dq/comp_nodes/dq_hash_combine.h>
 #include <ydb/library/yql/dq/proto/dq_tasks.pb.h>
-#include <ydb/library/yql/providers/solomon/actors/dq_solomon_read_actor.h>
+#include <ydb/library/yql/providers/generic/actors/yql_generic_provider_factories.h>
 #include <ydb/library/yql/providers/pq/async_io/dq_pq_read_actor.h>
 #include <ydb/library/yql/providers/pq/async_io/dq_pq_write_actor.h>
-#include <ydb/library/yql/dq/comp_nodes/dq_block_hash_join.h>
-#include <ydb/library/yql/dq/comp_nodes/dq_hash_combine.h>
+#include <ydb/library/yql/providers/solomon/actors/dq_solomon_read_actor.h>
 
 namespace NKikimr {
 namespace NMiniKQL {
@@ -90,6 +91,7 @@ NYql::NDq::IDqAsyncIoFactory::TPtr CreateKqpAsyncIoFactory(
     RegisterKqpWriteActor(*factory, counters);
     RegisterSequencerActorFactory(*factory, counters);
     RegisterKqpVectorResolveActor(*factory, counters);
+    NYql::NDq::RegisterDqInputTransformLookupActorFactory(*factory);
 
     if (federatedQuerySetup) {
         auto s3HttpRetryPolicy = NYql::GetHTTPDefaultRetryPolicy(NYql::THttpRetryPolicyOptions{.RetriedCurlCodes = NYql::FqRetriedCurlCodes()});

ydb/core/kqp/compute_actor/ya.make

@@ -25,12 +25,13 @@ PEERDIR(
     ydb/library/formats/arrow/protos
     ydb/library/formats/arrow/common
     ydb/library/yql/dq/actors/compute
+    ydb/library/yql/dq/actors/input_transforms
+    ydb/library/yql/dq/comp_nodes
     ydb/library/yql/providers/generic/actors
     ydb/library/yql/providers/pq/async_io
     ydb/library/yql/providers/s3/actors_factory
     ydb/library/yql/providers/solomon/actors
     yql/essentials/public/issue
-    ydb/library/yql/dq/comp_nodes
 )
 
 GENERATE_ENUM_SERIALIZATION(kqp_compute_state.h)

ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp

@@ -517,6 +517,47 @@ void TKqpTasksGraph::BuildVectorResolveChannels(const TStageInfo& stageInfo, ui3
         inputStageInfo, outputIndex, enableSpilling, logFunc);
 }
 
+void TKqpTasksGraph::BuildDqSourceStreamLookupChannels(const TStageInfo& stageInfo, ui32 inputIndex, const TStageInfo& inputStageInfo,
+    ui32 outputIndex, const NKqpProto::TKqpPhyCnDqSourceStreamLookup& dqSourceStreamLookup, const TChannelLogFunc& logFunc) {
+    YQL_ENSURE(stageInfo.Tasks.size() == 1);
+
+    auto* settings = GetMeta().Allocate<NDqProto::TDqInputTransformLookupSettings>();
+    settings->SetLeftLabel(dqSourceStreamLookup.GetLeftLabel());
+    settings->SetRightLabel(dqSourceStreamLookup.GetRightLabel());
+    settings->SetJoinType(dqSourceStreamLookup.GetJoinType());
+    settings->SetNarrowInputRowType(dqSourceStreamLookup.GetConnectionInputRowType());
+    settings->SetNarrowOutputRowType(dqSourceStreamLookup.GetConnectionOutputRowType());
+    settings->SetCacheLimit(dqSourceStreamLookup.GetCacheLimit());
+    settings->SetCacheTtlSeconds(dqSourceStreamLookup.GetCacheTtlSeconds());
+    settings->SetMaxDelayedRows(dqSourceStreamLookup.GetMaxDelayedRows());
+    settings->SetIsMultiget(dqSourceStreamLookup.GetIsMultiGet());
+
+    const auto& leftJointKeys = dqSourceStreamLookup.GetLeftJoinKeyNames();
+    settings->MutableLeftJoinKeyNames()->Assign(leftJointKeys.begin(), leftJointKeys.end());
+
+    const auto& rightJointKeys = dqSourceStreamLookup.GetRightJoinKeyNames();
+    settings->MutableRightJoinKeyNames()->Assign(rightJointKeys.begin(), rightJointKeys.end());
+
+    auto& streamLookupSource = *settings->MutableRightSource();
+    streamLookupSource.SetSerializedRowType(dqSourceStreamLookup.GetLookupRowType());
+    const auto& compiledSource = dqSourceStreamLookup.GetLookupSource();
+    streamLookupSource.SetProviderName(compiledSource.GetType());
+    *streamLookupSource.MutableLookupSource() = compiledSource.GetSettings();
+
+    TTransform dqSourceStreamLookupTransform = {
+        .Type = "StreamLookupInputTransform",
+        .InputType = dqSourceStreamLookup.GetInputStageRowType(),
+        .OutputType = dqSourceStreamLookup.GetOutputStageRowType(),
+    };
+    YQL_ENSURE(dqSourceStreamLookupTransform.Settings.PackFrom(*settings));
+
+    for (const auto taskId : stageInfo.Tasks) {
+        GetTask(taskId).Inputs[inputIndex].Transform = dqSourceStreamLookupTransform;
+    }
+
+    BuildUnionAllChannels(*this, stageInfo, inputIndex, inputStageInfo, outputIndex, /* enableSpilling */ false, logFunc);
+}
+
 void TKqpTasksGraph::BuildKqpStageChannels(TStageInfo& stageInfo, ui64 txId, bool enableSpilling, bool enableShuffleElimination) {
     auto& stage = stageInfo.Meta.GetStage(stageInfo.Id);
 
@@ -709,6 +750,12 @@ void TKqpTasksGraph::BuildKqpStageChannels(TStageInfo& stageInfo, ui64 txId, boo
                 break;
             }
 
+            case NKqpProto::TKqpPhyConnection::kDqSourceStreamLookup: {
+                BuildDqSourceStreamLookupChannels(stageInfo, inputIdx, inputStageInfo, outputIdx,
+                    input.GetDqSourceStreamLookup(), log);
+                break;
+            }
+
             default:
                 YQL_ENSURE(false, "Unexpected stage input type: " << (ui32)input.GetTypeCase());
         }
@@ -1369,6 +1416,8 @@ void TKqpTasksGraph::FillInputDesc(NYql::NDqProto::TTaskInput& inputDesc, const
             }
 
             transformProto->MutableSettings()->PackFrom(*input.Meta.VectorResolveSettings);
+        } else {
+            *transformProto->MutableSettings() = input.Transform->Settings;
         }
     }
 }
@@ -1724,6 +1773,7 @@ bool TKqpTasksGraph::BuildComputeTasks(TStageInfo& stageInfo, const ui32 nodesCo
                 case NKqpProto::TKqpPhyConnection::kMap:
                 case NKqpProto::TKqpPhyConnection::kParallelUnionAll:
                 case NKqpProto::TKqpPhyConnection::kVectorResolve:
+                case NKqpProto::TKqpPhyConnection::kDqSourceStreamLookup:
                     break;
                 default:
                     YQL_ENSURE(false, "Unexpected connection type: " << (ui32)input.GetTypeCase() << Endl);

ydb/core/kqp/executer_actor/kqp_tasks_graph.h

@@ -422,6 +422,8 @@ class TKqpTasksGraph : public NYql::NDq::TDqTasksGraph<TGraphMeta, TStageInfoMet
     void BuildVectorResolveChannels(const TStageInfo& stageInfo, ui32 inputIndex,
         const TStageInfo& inputStageInfo, ui32 outputIndex,
         const NKqpProto::TKqpPhyCnVectorResolve& vectorResolve, bool enableSpilling, const NYql::NDq::TChannelLogFunc& logFunc);
+    void BuildDqSourceStreamLookupChannels(const TStageInfo& stageInfo, ui32 inputIndex, const TStageInfo& inputStageInfo,
+        ui32 outputIndex, const NKqpProto::TKqpPhyCnDqSourceStreamLookup& dqSourceStreamLookup, const NYql::NDq::TChannelLogFunc& logFunc);
 
     void FillOutputDesc(NYql::NDqProto::TTaskOutput& outputDesc, const TTaskOutput& output, ui32 outputIdx,
         bool enableSpilling, const TStageInfo& stageInfo) const;

ydb/core/kqp/gateway/behaviour/streaming_query/optimization.cpp

@@ -39,8 +39,8 @@ bool ExploreStreamingQueryNode(TExprNode::TPtr node, TStreamingExploreCtx& res)
     const auto providerArg = node->ChildPtr(1);
     if (const auto maybeDataSource = TMaybeNode<TCoDataSource>(providerArg)) {
         const auto dataSourceCategory = maybeDataSource.Cast().Category().Value();
-        if (dataSourceCategory == NYql::PqProviderName) {
-            ++res.StreamingReads;
+        if (IsIn({NYql::PqProviderName, NYql::S3ProviderName, NYql::GenericProviderName}, dataSourceCategory)) {
+            res.StreamingReads += dataSourceCategory == NYql::PqProviderName;
             return true;
         }
 

ydb/core/kqp/host/kqp_host.cpp

@@ -1773,12 +1773,18 @@ class TKqpHost : public IKqpHost {
         }
 
         auto state = MakeIntrusive<NYql::TS3State>();
+
+        auto& configuration = *state->Configuration;
+        if (const auto requestContext = SessionCtx->GetUserRequestContext(); requestContext && requestContext->IsStreamingQuery) {
+            configuration.DisablePragma(configuration.UseRuntimeListing, false, "Runtime listing is not supported for streaming queries, pragma value was ignored");
+        }
+        configuration.WriteThroughDqIntegration = true;
+        configuration.AllowAtomicUploadCommit = queryType == EKikimrQueryType::Script;
+        configuration.Init(FederatedQuerySetup->S3GatewayConfig, TypesCtx);
+
         state->Types = TypesCtx.Get();
         state->FunctionRegistry = FuncRegistry;
         state->CredentialsFactory = FederatedQuerySetup->CredentialsFactory;
-        state->Configuration->WriteThroughDqIntegration = true;
-        state->Configuration->AllowAtomicUploadCommit = queryType == EKikimrQueryType::Script;
-        state->Configuration->Init(FederatedQuerySetup->S3GatewayConfig, TypesCtx);
         state->Gateway = FederatedQuerySetup->HttpGateway;
         state->GatewayRetryPolicy = NYql::GetHTTPDefaultRetryPolicy(NYql::THttpRetryPolicyOptions{.RetriedCurlCodes = NYql::FqRetriedCurlCodes()});
         state->ExecutorPoolId = AppData()->UserPoolId;
@@ -1980,6 +1986,7 @@ class TKqpHost : public IKqpHost {
         auto configProvider = CreateConfigProvider(*TypesCtx, gatewaysConfig, {}, allowSettings);
         TypesCtx->AddDataSource(ConfigProviderName, configProvider);
         TypesCtx->MatchRecognize = QueryServiceConfig.GetEnableMatchRecognize();
+        TypesCtx->StreamLookupJoin = true;
 
         YQL_ENSURE(TypesCtx->Initialize(*ExprCtx));
 

ydb/core/kqp/opt/logical/kqp_opt_log.cpp

@@ -41,6 +41,7 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase {
         AddHandler(0, &TCoTake::Match, HNDL(RewriteTakeSortToTopSort));
         AddHandler(0, &TCoFlatMap::Match, HNDL(RewriteSqlInToEquiJoin));
         AddHandler(0, &TCoFlatMap::Match, HNDL(RewriteSqlInCompactToJoin));
+        AddHandler(0, &TCoEquiJoin::Match, HNDL(RewriteStreamEquiJoinWithLookup));
         AddHandler(0, &TCoEquiJoin::Match, HNDL(OptimizeEquiJoinWithCosts));
         AddHandler(0, &TCoEquiJoin::Match, HNDL(RewriteEquiJoin));
         AddHandler(0, &TDqJoin::Match, HNDL(JoinToIndexLookup));
@@ -167,6 +168,12 @@ class TKqpLogicalOptTransformer : public TOptimizeTransformerBase {
         return output;
     }
 
+    TMaybeNode<TExprBase> RewriteStreamEquiJoinWithLookup(TExprBase node, TExprContext& ctx) {
+        TExprBase output = DqRewriteStreamEquiJoinWithLookup(node, ctx, TypesCtx);
+        DumpAppliedRule("KqpRewriteStreamEquiJoinWithLookup", node.Ptr(), output.Ptr(), ctx);
+        return output;
+    }
+
     TMaybeNode<TExprBase> OptimizeEquiJoinWithCosts(TExprBase node, TExprContext& ctx) {
         auto maxDPhypDPTableSize = Config->MaxDPHypDPTableSize.Get().GetOrElse(TDqSettings::TDefault::MaxDPHypDPTableSize);
         auto optLevel = Config->CostBasedOptimizationLevel.Get().GetOrElse(Config->DefaultCostBasedOptimizationLevel);