Skip to content

Handle Write in DQ recapture #8212

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 23, 2024
Merged

Handle Write in DQ recapture #8212

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7862086
Handle Write in DQ recapture
rvu1024 Aug 23, 2024
24a3845
Use another method for write recapture
rvu1024 Aug 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ydb/library/yql/dq/integration/yql_dq_integration.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ class IDqIntegration {
virtual TMaybe<ui64> EstimateReadSize(ui64 dataSizePerJob, ui32 maxTasksPerStage, const TVector<const TExprNode*>& nodes, TExprContext& ctx) = 0;
virtual TExprNode::TPtr WrapRead(const TDqSettings& config, const TExprNode::TPtr& read, TExprContext& ctx) = 0;
virtual TMaybe<TOptimizerStatistics> ReadStatistics(const TExprNode::TPtr& readWrap, TExprContext& ctx) = 0;
virtual TExprNode::TPtr RecaptureWrite(const TExprNode::TPtr& write, TExprContext& ctx) = 0;

// Nothing if callable is not for writing,
// false if callable is for writing and there are some errors (they are added to ctx),
Expand Down
12 changes: 11 additions & 1 deletion ydb/library/yql/dq/opt/dq_opt_log.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -349,7 +349,7 @@ NNodes::TExprBase DqReplicateFieldSubset(NNodes::TExprBase node, TExprContext& c
return node;
}

IGraphTransformer::TStatus DqWrapRead(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx, TTypeAnnotationContext& typesCtx, const TDqSettings& config) {
IGraphTransformer::TStatus DqWrapIO(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx, TTypeAnnotationContext& typesCtx, const TDqSettings& config) {
TOptimizeExprSettings settings{&typesCtx};
auto status = OptimizeExpr(input, output, [&](const TExprNode::TPtr& node, TExprContext& ctx) {
if (auto maybeRead = TMaybeNode<TCoRight>(node).Input()) {
Expand All @@ -364,6 +364,16 @@ IGraphTransformer::TStatus DqWrapRead(const TExprNode::TPtr& input, TExprNode::T
}
}
}
} else if (node->GetTypeAnn()->GetKind() == ETypeAnnotationKind::World
&& !TCoCommit::Match(node.Get())
&& node->ChildrenSize() > 1
&& TCoDataSink::Match(node->Child(1))) {
auto dataSinkName = node->Child(1)->Child(0)->Content();
auto dataSink = typesCtx.DataSinkMap.FindPtr(dataSinkName);
YQL_ENSURE(dataSink);
if (auto dqIntegration = (*dataSink)->GetDqIntegration()) {
return dqIntegration->RecaptureWrite(node, ctx);
}
}

return node;
Expand Down
2 changes: 1 addition & 1 deletion ydb/library/yql/dq/opt/dq_opt_log.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ NNodes::TExprBase DqSqlInDropCompact(NNodes::TExprBase node, TExprContext& ctx);

NNodes::TExprBase DqReplicateFieldSubset(NNodes::TExprBase node, TExprContext& ctx, const TParentsMap& parents);

IGraphTransformer::TStatus DqWrapRead(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx, TTypeAnnotationContext& typesCtx, const TDqSettings& config);
IGraphTransformer::TStatus DqWrapIO(const TExprNode::TPtr& input, TExprNode::TPtr& output, TExprContext& ctx, TTypeAnnotationContext& typesCtx, const TDqSettings& config);

NNodes::TExprBase DqExpandMatchRecognize(NNodes::TExprBase node, TExprContext& ctx, TTypeAnnotationContext& typeAnnCtx);

Expand Down
5 changes: 5 additions & 0 deletions ydb/library/yql/providers/common/dq/yql_dq_integration_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,11 @@ TMaybe<TOptimizerStatistics> TDqIntegrationBase::ReadStatistics(const TExprNode:
return Nothing();
}

TExprNode::TPtr TDqIntegrationBase::RecaptureWrite(const TExprNode::TPtr& write, TExprContext& ctx) {
Y_UNUSED(ctx);
return write;
}

TMaybe<bool> TDqIntegrationBase::CanWrite(const TExprNode&, TExprContext&) {
return Nothing();
}
Expand Down
1 change: 1 addition & 0 deletions ydb/library/yql/providers/common/dq/yql_dq_integration_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ class TDqIntegrationBase: public IDqIntegration {
TMaybe<ui64> EstimateReadSize(ui64 dataSizePerJob, ui32 maxTasksPerStage, const TVector<const TExprNode*>& nodes, TExprContext& ctx) override;
TExprNode::TPtr WrapRead(const TDqSettings& config, const TExprNode::TPtr& read, TExprContext& ctx) override;
TMaybe<TOptimizerStatistics> ReadStatistics(const TExprNode::TPtr& readWrap, TExprContext& ctx) override;
TExprNode::TPtr RecaptureWrite(const TExprNode::TPtr& write, TExprContext& ctx) override;
void RegisterMkqlCompiler(NCommon::TMkqlCallableCompilerBase& compiler) override;
TMaybe<bool> CanWrite(const TExprNode& write, TExprContext& ctx) override;
bool CanBlockRead(const NNodes::TExprBase& node, TExprContext& ctx, TTypeAnnotationContext& typesCtx) override;
Expand Down
2 changes: 1 addition & 1 deletion ydb/library/yql/providers/dq/provider/yql_dq_recapture.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ class TDqsRecaptureTransformer : public TSyncTransformerBase {

State_->TypeCtx->DqFallbackPolicy = State_->Settings->FallbackPolicy.Get().GetOrElse(EFallbackPolicy::Default);

IGraphTransformer::TStatus status = NDq::DqWrapRead(input, output, ctx, *State_->TypeCtx, *State_->Settings);
IGraphTransformer::TStatus status = NDq::DqWrapIO(input, output, ctx, *State_->TypeCtx, *State_->Settings);
if (input != output) {
YQL_CLOG(INFO, ProviderDq) << "DqsRecapture";
// TODO: Add before/after recapture transformers
Expand Down
Loading