fix: Fix update error with pushdownsubfields optimization

[zation99]

Description

• Fix PushdownSubfields optimizer to handle UpdateNode in the plan tree, preventing verify failures when pushdown-subfields-enabled=true is set and UPDATE operations are executed
• Add visitUpdate override that registers all source output variables in the optimizer context, following the same pattern as visitExplainAnalyze

Motivation and Context

Update with subfields pushdown enabled reports error.

Test Plan

added test for Iceberg connector for update

== NO RELEASE NOTE ==

Summary by Sourcery

Handle UPDATE plans in the PushdownSubfields optimizer and add coverage for DML with subfields pushdown enabled for the Iceberg connector.

Bug Fixes:

• Register source output variables for UpdateNode in the PushdownSubfields optimizer to prevent verification failures when subfields pushdown is enabled for UPDATE statements.

Tests:

• Add Iceberg logical planner test covering UPDATE and DELETE operations with subfields pushdown enabled.

[sourcery-ai]

Reviewer's Guide

Handle UPDATE plan nodes in the PushdownSubfields optimizer to avoid verification failures when pushdown-subfields is enabled, and add Iceberg logical planner coverage for DML (UPDATE/DELETE) with subfield pushdown enabled.

Sequence diagram for UPDATE planning with PushdownSubfields optimization

sequenceDiagram
    actor User
    participant Coordinator
    participant Planner
    participant PushdownSubfields
    participant Rewriter
    participant UpdateNode

    User->>Coordinator: submit UPDATE statement
    Coordinator->>Planner: plan UPDATE
    Planner->>Planner: build logical plan with UpdateNode
    Planner->>PushdownSubfields: optimize(plan, session)
    PushdownSubfields->>Rewriter: rewrite(plan, rootContext)

    Rewriter->>UpdateNode: visitUpdate(node, context)
    UpdateNode-->>Rewriter: getSource()
    UpdateNode-->>Rewriter: getSource().getOutputVariables()
    Rewriter->>Rewriter: context.variables.addAll(sourceOutputVariables)
    Rewriter->>Rewriter: context.defaultRewrite(node, context)

    Rewriter-->>PushdownSubfields: optimizedPlan
    PushdownSubfields-->>Planner: return optimizedPlan
    Planner-->>Coordinator: executable plan
    Coordinator-->>User: execute UPDATE successfully

Loading

Updated class diagram for PushdownSubfields optimizer handling UpdateNode

classDiagram
    class PushdownSubfields {
        +PlanNode optimize(PlanNode plan, Session session, Map planSymbolAllocator)
    }

    class Context {
        +Set~VariableReferenceExpression~ variables
    }

    class Rewriter {
        +PlanNode visitDelete(DeleteNode node, RewriteContext~Context~ context)
        +PlanNode visitUpdate(UpdateNode node, RewriteContext~Context~ context)
        +PlanNode visitTopN(TopNNode node, RewriteContext~Context~ context)
        +PlanNode visitExplainAnalyze(ExplainAnalyzeNode node, RewriteContext~Context~ context)
        +PlanNode defaultRewrite(PlanNode node, Context context)
    }

    class PlanNode {
        <<abstract>>
        +List~VariableReferenceExpression~ getOutputVariables()
    }

    class DeleteNode {
        +PlanNode getSource()
    }

    class UpdateNode {
        +PlanNode getSource()
        +List~VariableReferenceExpression~ getOutputVariables()
    }

    class TopNNode {
        +PlanNode getSource()
    }

    class ExplainAnalyzeNode {
        +PlanNode getSource()
    }

    class RewriteContext~Context~ {
        +Context get()
        +PlanNode defaultRewrite(PlanNode node, Context context)
    }

    PushdownSubfields *-- Rewriter
    Rewriter --> RewriteContext~Context~
    RewriteContext~Context~ --> Context

    PlanNode <|-- DeleteNode
    PlanNode <|-- UpdateNode
    PlanNode <|-- TopNNode
    PlanNode <|-- ExplainAnalyzeNode

    Rewriter ..> DeleteNode : visitDelete
    Rewriter ..> UpdateNode : visitUpdate
    Rewriter ..> TopNNode : visitTopN
    Rewriter ..> ExplainAnalyzeNode : visitExplainAnalyze

    Rewriter ..> Context : uses
    Rewriter ..> PlanNode : rewrites
    UpdateNode ..> PlanNode : source
    DeleteNode ..> PlanNode : source
    TopNNode ..> PlanNode : source
    ExplainAnalyzeNode ..> PlanNode : source

Loading

File-Level Changes

Change	Details	Files
Extend PushdownSubfields optimizer to support UpdateNode by registering its source output variables in the optimization context.	Add visitUpdate override mirroring visitDelete/visitExplainAnalyze behavior Populate optimizer context variable set with all variables from the UPDATE source node Delegate to defaultRewrite to continue recursive optimization after registering variables	presto-main-base/src/main/java/com/facebook/presto/sql/planner/optimizations/PushdownSubfields.java
Add an Iceberg logical planner test that exercises UPDATE and DELETE statements with pushdown subfields enabled to prevent regressions.	Create testPushdownSubfieldsWithDml that creates a test table and inserts sample rows Verify UPDATE changes rows when the WHERE predicate matches and is a no-op when it does not match Verify DELETE removes the expected row and final table contents are as expected Ensure cleanup of the test table in a finally block	presto-iceberg/src/test/java/com/facebook/presto/iceberg/TestIcebergLogicalPlanner.java

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've left some high level feedback:

• The new testPushdownSubfieldsWithDml relies on the environment’s default for pushdown-subfields-enabled; consider explicitly enabling this session property in the test so it clearly exercises the intended optimizer behavior and doesn’t become a no-op if defaults change.
• In visitUpdate, you only register node.getSource().getOutputVariables(); if the optimizer relies on seeing all referenced variables, consider whether any update-specific variables (e.g., assignment targets) should also be tracked for consistency with other plan node handlers.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- The new `testPushdownSubfieldsWithDml` relies on the environment’s default for `pushdown-subfields-enabled`; consider explicitly enabling this session property in the test so it clearly exercises the intended optimizer behavior and doesn’t become a no-op if defaults change.
- In `visitUpdate`, you only register `node.getSource().getOutputVariables()`; if the optimizer relies on seeing all referenced variables, consider whether any update-specific variables (e.g., assignment targets) should also be tracked for consistency with other plan node handlers.

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[hantangwangd]

Thanks @zation99, lgtm!