Merge stream lookup changes (#12292)

Co-authored-by: Mikhail Surin <[email protected]>
Co-authored-by: ulya-sidorina <[email protected]>
Co-authored-by: Pavel Velikhov <[email protected]>

Author: 4 people

ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp

@@ -415,6 +415,7 @@ void BuildStreamLookupChannels(TKqpTasksGraph& graph, const TStageInfo& stageInf
     }
 
     settings->SetLookupStrategy(streamLookup.GetLookupStrategy());
+    settings->SetKeepRowsOrder(streamLookup.GetKeepRowsOrder());
 
     TTransform streamLookupTransform;
     streamLookupTransform.Type = "StreamLookupInputTransformer";

ydb/core/kqp/opt/physical/kqp_opt_phy_source.cpp

@@ -110,6 +110,11 @@ TExprBase KqpRewriteReadTable(TExprBase node, TExprContext& ctx, const TKqpOptim
         matched->Settings = settings.BuildNode(ctx, matched->Settings.Pos());
     }
 
+    if (kqpCtx.Config->HasMaxSequentialReadsInFlight()) {
+        settings.SequentialInFlight = *kqpCtx.Config->MaxSequentialReadsInFlight.Get();
+        matched->Settings = settings.BuildNode(ctx, matched->Settings.Pos());
+    }
+
     TVector<TExprBase> inputs;
     TVector<TCoArgument> args;
     TNodeOnNodeOwnedMap argReplaces;

ydb/core/kqp/provider/yql_kikimr_settings.cpp

@@ -87,6 +87,7 @@ TKikimrConfiguration::TKikimrConfiguration() {
     REGISTER_SETTING(*this, OptimizerHints).Parser([](const TString& v) { return NYql::TOptimizerHints::Parse(v); });
     REGISTER_SETTING(*this, OverridePlanner);
     REGISTER_SETTING(*this, UseGraceJoinCoreForMap);
+    REGISTER_SETTING(*this, EnableOrderPreservingLookupJoin);
 
     REGISTER_SETTING(*this, OptUseFinalizeByKey);
     REGISTER_SETTING(*this, CostBasedOptimizationLevel);
@@ -96,6 +97,7 @@ TKikimrConfiguration::TKikimrConfiguration() {
     REGISTER_SETTING(*this, MaxDPccpDPTableSize);
 
     REGISTER_SETTING(*this, MaxTasksPerStage);
+    REGISTER_SETTING(*this, MaxSequentialReadsInFlight);
 
     /* Runtime */
     REGISTER_SETTING(*this, ScanQuery);
@@ -120,6 +122,10 @@ bool TKikimrSettings::SpillingEnabled() const {
     return GetFlagValue(_KqpEnableSpilling.Get());
 }
 
+bool TKikimrSettings::OrderPreservingLookupJoinEnabled() const {
+    return GetFlagValue(EnableOrderPreservingLookupJoin.Get());
+}
+
 bool TKikimrSettings::DisableLlvmForUdfStages() const {
     return GetFlagValue(_KqpDisableLlvmForUdfStages.Get());
 }
@@ -148,6 +154,10 @@ bool TKikimrSettings::HasOptUseFinalizeByKey() const {
     return GetFlagValue(OptUseFinalizeByKey.Get().GetOrElse(true)) != EOptionalFlag::Disabled;
 }
 
+bool TKikimrSettings::HasMaxSequentialReadsInFlight() const {
+    return !MaxSequentialReadsInFlight.Get().Empty();
+}
+
 EOptionalFlag TKikimrSettings::GetOptPredicateExtract() const {
     return GetOptionalFlagValue(OptEnablePredicateExtract.Get());
 }

ydb/core/kqp/provider/yql_kikimr_settings.h

@@ -54,6 +54,7 @@ struct TKikimrSettings {
     NCommon::TConfSetting<ui64, false> EnableSpillingNodes;
     NCommon::TConfSetting<TString, false> OverridePlanner;
     NCommon::TConfSetting<bool, false> UseGraceJoinCoreForMap;
+    NCommon::TConfSetting<bool, false> EnableOrderPreservingLookupJoin;
 
     NCommon::TConfSetting<TString, false> OptOverrideStatistics;
     NCommon::TConfSetting<NYql::TOptimizerHints, false> OptimizerHints;
@@ -72,6 +73,7 @@ struct TKikimrSettings {
 
 
     NCommon::TConfSetting<ui32, false> MaxTasksPerStage;
+    NCommon::TConfSetting<ui32, false> MaxSequentialReadsInFlight;
 
     /* Runtime */
     NCommon::TConfSetting<bool, true> ScanQuery;
@@ -88,7 +90,8 @@ struct TKikimrSettings {
     bool HasOptEnableOlapPushdown() const;
     bool HasOptEnableOlapProvideComputeSharding() const;
     bool HasOptUseFinalizeByKey() const;
-
+    bool HasMaxSequentialReadsInFlight() const;
+    bool OrderPreservingLookupJoinEnabled() const;
     EOptionalFlag GetOptPredicateExtract() const;
     EOptionalFlag GetUseLlvm() const;
     NDq::EHashJoinMode GetHashJoinMode() const;

ydb/core/kqp/query_compiler/kqp_query_compiler.cpp

@@ -1277,6 +1277,7 @@ class TKqpQueryCompiler : public IKqpQueryCompiler {
             YQL_ENSURE(streamLookup.LookupStrategy().Maybe<TCoAtom>());
             TString lookupStrategy = streamLookup.LookupStrategy().Maybe<TCoAtom>().Cast().StringValue();
             streamLookupProto.SetLookupStrategy(GetStreamLookupStrategy(lookupStrategy));
+            streamLookupProto.SetKeepRowsOrder(Config->OrderPreservingLookupJoinEnabled());
 
             switch (streamLookupProto.GetLookupStrategy()) {
                 case NKqpProto::EStreamLookupStrategy::LOOKUP: {

ydb/core/kqp/runtime/kqp_stream_lookup_actor.cpp

@@ -268,6 +268,8 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
                     RuntimeError(TStringBuilder() << "Unexpected event: " << ev->GetTypeRewrite(),
                         NYql::NDqProto::StatusIds::INTERNAL_ERROR);
             }
+        } catch (const NKikimr::TMemoryLimitExceededException& e) {
+            RuntimeError("Memory limit exceeded at stream lookup", NYql::NDqProto::StatusIds::PRECONDITION_FAILED);
         } catch (const yexception& e) {
             RuntimeError(e.what(), NYql::NDqProto::StatusIds::INTERNAL_ERROR);
         }
@@ -400,6 +402,7 @@ class TKqpStreamLookupActor : public NActors::TActorBootstrapped<TKqpStreamLooku
             }
         }
 
+        auto guard = BindAllocator();
         StreamLookupWorker->AddResult(TKqpStreamLookupWorker::TShardReadResult{
             read.ShardId, THolder<TEventHandle<TEvDataShard::TEvReadResult>>(ev.Release())
         });