YQ-2892 support emulate YT in KqpRun (#2562)

Author: GrigoriyPA

ydb/core/fq/libs/common/cache_ut.cpp

@@ -2,7 +2,7 @@
 
 #include <ydb/services/ydb/ydb_common_ut.h>
 
-using namespace NFq;
+namespace NFq {
 
 using TCache = TTtlCache<int,int>;
 
@@ -66,3 +66,5 @@ Y_UNIT_TEST_SUITE(Cache) {
  UNIT_ASSERT_VALUES_EQUAL(cache.Size(), 0);
  }
 }
+
+} // namespace NFq

ydb/core/grpc_services/ya.make

@@ -106,6 +106,7 @@ PEERDIR(
  ydb/core/io_formats/ydb_dump
  ydb/core/kesus/tablet
  ydb/core/kqp/common
+ ydb/core/kqp/federated_query
  ydb/core/protos
  ydb/core/scheme
  ydb/core/sys_view

ydb/core/kqp/compute_actor/kqp_compute_actor.cpp

@@ -12,7 +12,6 @@
 #include <ydb/library/yql/providers/generic/actors/yql_generic_source_factory.h>
 #include <ydb/library/yql/providers/s3/actors/yql_s3_sink_factory.h>
 #include <ydb/library/yql/providers/s3/actors/yql_s3_source_factory.h>
-#include <ydb/library/yql/providers/yt/comp_nodes/dq/dq_yt_factory.h>
 #include <ydb/core/protos/ssa.pb.h>
 #include <ydb/library/yql/dq/proto/dq_tasks.pb.h>
 
@@ -24,22 +23,20 @@ using TCallableActorBuilderFunc = std::function<
  IComputationNode*(
  TCallable& callable, const TComputationNodeFactoryContext& ctx, TKqpScanComputeContext& computeCtx)>;
 
-TComputationNodeFactory GetKqpActorComputeFactory(TKqpScanComputeContext* computeCtx) {
+TComputationNodeFactory GetKqpActorComputeFactory(TKqpScanComputeContext* computeCtx, const std::optional<NKqp::TKqpFederatedQuerySetup>& federatedQuerySetup) {
  MKQL_ENSURE_S(computeCtx);
 
- auto computeFactory = GetKqpBaseComputeFactory(computeCtx);
- auto ytComputeFactory = NYql::GetDqYtFactory();
+ auto computeFactory = NKqp::MakeKqpFederatedQueryComputeFactory(
+ GetKqpBaseComputeFactory(computeCtx),
+ federatedQuerySetup
+ );
 
- return [computeFactory, ytComputeFactory, computeCtx]
+ return [computeFactory, computeCtx]
  (TCallable& callable, const TComputationNodeFactoryContext& ctx) -> IComputationNode* {
  if (auto compute = computeFactory(callable, ctx)) {
  return compute;
  }
 
- if (auto ytCompute = ytComputeFactory(callable, ctx)) {
- return ytCompute;
- }
-
  auto name = callable.GetType()->GetName();
 
  if (name == "KqpWideReadTable"sv) {

ydb/core/kqp/compute_actor/kqp_compute_actor.h

@@ -13,7 +13,7 @@ namespace NMiniKQL {
 
 class TKqpScanComputeContext;
 
-TComputationNodeFactory GetKqpActorComputeFactory(TKqpScanComputeContext* computeCtx);
+TComputationNodeFactory GetKqpActorComputeFactory(TKqpScanComputeContext* computeCtx, const std::optional<NKqp::TKqpFederatedQuerySetup>& federatedQuerySetup);
 
 } // namespace NMiniKQL
 
@@ -47,7 +47,8 @@ IActor* CreateKqpComputeActor(const TActorId& executerId, ui64 txId, NYql::NDqPr
  const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry,
  const NYql::NDq::TComputeRuntimeSettings& settings, const NYql::NDq::TComputeMemoryLimits& memoryLimits,
  NWilson::TTraceId traceId,
- TIntrusivePtr<NActors::TProtoArenaHolder> arena);
+ TIntrusivePtr<NActors::TProtoArenaHolder> arena,
+ const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup);
 
 IActor* CreateKqpScanComputeActor(const TActorId& executerId, ui64 txId,
  NYql::NDqProto::TDqTask* task, NYql::NDq::IDqAsyncIoFactory::TPtr asyncIoFactory,

ydb/core/kqp/compute_actor/kqp_pure_compute_actor.cpp

@@ -14,9 +14,11 @@ TKqpComputeActor::TKqpComputeActor(const TActorId& executerId, ui64 txId, NDqPro
  IDqAsyncIoFactory::TPtr asyncIoFactory,
  const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry,
  const TComputeRuntimeSettings& settings, const TComputeMemoryLimits& memoryLimits,
- NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena)
+ NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena,
+ const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup)
  : TBase(executerId, txId, task, std::move(asyncIoFactory), functionRegistry, settings, memoryLimits, /* ownMemoryQuota = */ true, /* passExceptions = */ true, /*taskCounters = */ nullptr, std::move(traceId), std::move(arena))
  , ComputeCtx(settings.StatsMode)
+ , FederatedQuerySetup(federatedQuerySetup)
 {
  if (GetTask().GetMeta().Is<NKikimrTxDataShard::TKqpTransaction::TScanTaskMeta>()) {
  Meta.ConstructInPlace();
@@ -44,7 +46,7 @@ void TKqpComputeActor::DoBootstrap() {
  execCtx.RandomProvider = TAppData::RandomProvider.Get();
  execCtx.TimeProvider = TAppData::TimeProvider.Get();
  execCtx.ComputeCtx = &ComputeCtx;
- execCtx.ComputationFactory = NMiniKQL::GetKqpActorComputeFactory(&ComputeCtx);
+ execCtx.ComputationFactory = NMiniKQL::GetKqpActorComputeFactory(&ComputeCtx, FederatedQuerySetup);
  execCtx.ApplyCtx = nullptr;
  execCtx.TypeEnv = nullptr;
  execCtx.PatternCache = GetKqpResourceManager()->GetPatternCache();
@@ -278,10 +280,11 @@ IActor* CreateKqpComputeActor(const TActorId& executerId, ui64 txId, NDqProto::T
  IDqAsyncIoFactory::TPtr asyncIoFactory,
  const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry,
  const TComputeRuntimeSettings& settings, const TComputeMemoryLimits& memoryLimits,
- NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena)
+ NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena,
+ const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup)
 {
  return new TKqpComputeActor(executerId, txId, task, std::move(asyncIoFactory),
- functionRegistry, settings, memoryLimits, std::move(traceId), std::move(arena));
+ functionRegistry, settings, memoryLimits, std::move(traceId), std::move(arena), federatedQuerySetup);
 }
 
 } // namespace NKqp

ydb/core/kqp/compute_actor/kqp_pure_compute_actor.h

@@ -29,7 +29,8 @@ class TKqpComputeActor : public TDqSyncComputeActorBase<TKqpComputeActor> {
  IDqAsyncIoFactory::TPtr asyncIoFactory,
  const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry,
  const TComputeRuntimeSettings& settings, const TComputeMemoryLimits& memoryLimits,
- NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena);
+ NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena,
+ const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup);
 
  void DoBootstrap();
 
@@ -61,13 +62,15 @@ class TKqpComputeActor : public TDqSyncComputeActorBase<TKqpComputeActor> {
  NMiniKQL::TKqpScanComputeContext::TScanData* ScanData = nullptr;
  TActorId SysViewActorId;
  const TDqTaskRunnerParameterProvider ParameterProvider;
+ const std::optional<TKqpFederatedQuerySetup> FederatedQuerySetup;
 };
 
 IActor* CreateKqpComputeActor(const TActorId& executerId, ui64 txId, NDqProto::TDqTask* task,
  IDqAsyncIoFactory::TPtr asyncIoFactory,
  const NKikimr::NMiniKQL::IFunctionRegistry* functionRegistry,
  const TComputeRuntimeSettings& settings, const TComputeMemoryLimits& memoryLimits,
- NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena);
+ NWilson::TTraceId traceId, TIntrusivePtr<NActors::TProtoArenaHolder> arena,
+ const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup);
 
 } // namespace NKqp
 } // namespace NKikimr

ydb/core/kqp/compute_actor/kqp_scan_compute_actor.cpp

@@ -181,7 +181,7 @@ void TKqpScanComputeActor::DoBootstrap() {
  NDq::TDqTaskRunnerContext execCtx;
  execCtx.FuncRegistry = AppData()->FunctionRegistry;
  execCtx.ComputeCtx = &ComputeCtx;
- execCtx.ComputationFactory = NMiniKQL::GetKqpActorComputeFactory(&ComputeCtx);
+ execCtx.ComputationFactory = NMiniKQL::GetKqpActorComputeFactory(&ComputeCtx, std::nullopt);
  execCtx.RandomProvider = TAppData::RandomProvider.Get();
  execCtx.TimeProvider = TAppData::TimeProvider.Get();
  execCtx.ApplyCtx = nullptr;

ydb/core/kqp/compute_actor/ya.make

@@ -24,7 +24,6 @@ PEERDIR(
  ydb/library/yql/dq/actors/compute
  ydb/library/yql/providers/generic/actors
  ydb/library/yql/providers/s3/actors
- ydb/library/yql/providers/yt/comp_nodes/dq
  ydb/library/yql/public/issue
 )
 

ydb/core/kqp/executer_actor/kqp_data_executer.cpp

@@ -128,12 +128,13 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
  const NKikimrConfig::TTableServiceConfig::EChannelTransportVersion chanTransportVersion,
  const NKikimrConfig::TTableServiceConfig::TAggregationConfig& aggregation,
  const TActorId& creator, TDuration maximalSecretsSnapshotWaitTime, const TIntrusivePtr<TUserRequestContext>& userRequestContext,
- const bool enableOlapSink, ui32 statementResultIndex)
+ const bool enableOlapSink, ui32 statementResultIndex, const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup)
  : TBase(std::move(request), database, userToken, counters, executerRetriesConfig, chanTransportVersion, aggregation,
  maximalSecretsSnapshotWaitTime, userRequestContext, statementResultIndex, TWilsonKqp::DataExecuter, "DataExecuter", streamResult
  )
  , AsyncIoFactory(std::move(asyncIoFactory))
  , EnableOlapSink(enableOlapSink)
+ , FederatedQuerySetup(federatedQuerySetup)
  {
  Target = creator;
 
@@ -2206,7 +2207,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
 
  Planner = CreateKqpPlanner(TasksGraph, TxId, SelfId(), GetSnapshot(),
  Database, UserToken, Deadline.GetOrElse(TInstant::Zero()), Request.StatsMode, false, Nothing(),
- ExecuterSpan, std::move(ResourceSnapshot), ExecuterRetriesConfig, useDataQueryPool, localComputeTasks, Request.MkqlMemoryLimit, AsyncIoFactory, singlePartitionOptAllowed, GetUserRequestContext());
+ ExecuterSpan, std::move(ResourceSnapshot), ExecuterRetriesConfig, useDataQueryPool, localComputeTasks, Request.MkqlMemoryLimit, AsyncIoFactory, singlePartitionOptAllowed, GetUserRequestContext(), FederatedQuerySetup);
 
  auto err = Planner->PlanExecution();
  if (err) {
@@ -2404,6 +2405,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
 private:
  NYql::NDq::IDqAsyncIoFactory::TPtr AsyncIoFactory;
  bool EnableOlapSink = false;
+ const std::optional<TKqpFederatedQuerySetup> FederatedQuerySetup;
 
  bool HasExternalSources = false;
  bool SecretSnapshotRequired = false;
@@ -2445,10 +2447,11 @@ IActor* CreateKqpDataExecuter(IKqpGateway::TExecPhysicalRequest&& request, const
  TKqpRequestCounters::TPtr counters, bool streamResult, const NKikimrConfig::TTableServiceConfig::TAggregationConfig& aggregation,
  const NKikimrConfig::TTableServiceConfig::TExecuterRetriesConfig& executerRetriesConfig,
  NYql::NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, const NKikimrConfig::TTableServiceConfig::EChannelTransportVersion chanTransportVersion, const TActorId& creator,
- TDuration maximalSecretsSnapshotWaitTime, const TIntrusivePtr<TUserRequestContext>& userRequestContext, const bool enableOlapSink, ui32 statementResultIndex)
+ TDuration maximalSecretsSnapshotWaitTime, const TIntrusivePtr<TUserRequestContext>& userRequestContext, const bool enableOlapSink, ui32 statementResultIndex,
+ const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup)
 {
  return new TKqpDataExecuter(std::move(request), database, userToken, counters, streamResult, executerRetriesConfig,
- std::move(asyncIoFactory), chanTransportVersion, aggregation, creator, maximalSecretsSnapshotWaitTime, userRequestContext, enableOlapSink, statementResultIndex);
+ std::move(asyncIoFactory), chanTransportVersion, aggregation, creator, maximalSecretsSnapshotWaitTime, userRequestContext, enableOlapSink, statementResultIndex, federatedQuerySetup);
 }
 
 } // namespace NKqp

ydb/core/kqp/executer_actor/kqp_executer.h

@@ -6,6 +6,7 @@
 #include <ydb/core/kqp/query_data/kqp_query_data.h>
 #include <ydb/core/kqp/gateway/kqp_gateway.h>
 #include <ydb/core/kqp/counters/kqp_counters.h>
+#include <ydb/core/kqp/federated_query/kqp_federated_query_helpers.h>
 #include <ydb/core/tx/long_tx_service/public/lock_handle.h>
 #include <ydb/library/yql/dq/actors/compute/dq_compute_actor_async_io_factory.h>
 #include <ydb/core/protos/table_service_config.pb.h>
@@ -93,7 +94,7 @@ IActor* CreateKqpExecuter(IKqpGateway::TExecPhysicalRequest&& request, const TSt
  NYql::NDq::IDqAsyncIoFactory::TPtr asyncIoFactory, TPreparedQueryHolder::TConstPtr preparedQuery,
  const NKikimrConfig::TTableServiceConfig::EChannelTransportVersion chanTransportVersion, const TActorId& creator,
  TDuration maximalSecretsSnapshotWaitTime, const TIntrusivePtr<TUserRequestContext>& userRequestContext,
- const bool enableOlapSink, ui32 statementResultIndex);
+ const bool enableOlapSink, ui32 statementResultIndex, const std::optional<TKqpFederatedQuerySetup>& federatedQuerySetup);
 
 IActor* CreateKqpSchemeExecuter(
  TKqpPhyTxHolder::TConstPtr phyTx, NKikimrKqp::EQueryType queryType, const TActorId& target,