From 0adbc574a9b815f398aa65509d162b55c0eedfff Mon Sep 17 00:00:00 2001
From: Feilong Liu <feilongliu@meta.com>
Date: Mon, 2 Dec 2024 16:35:49 -0800
Subject: [PATCH] Add delete node processing in subfield prune optimizer

---
 .../sql/planner/optimizations/PushdownSubfields.java  | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PushdownSubfields.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PushdownSubfields.java
index 3c04454458e95..31f15d9a694ec 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PushdownSubfields.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PushdownSubfields.java
@@ -35,6 +35,7 @@
 import com.facebook.presto.spi.function.StandardFunctionResolution;
 import com.facebook.presto.spi.plan.AggregationNode;
 import com.facebook.presto.spi.plan.CteProducerNode;
+import com.facebook.presto.spi.plan.DeleteNode;
 import com.facebook.presto.spi.plan.DistinctLimitNode;
 import com.facebook.presto.spi.plan.EquiJoinClause;
 import com.facebook.presto.spi.plan.FilterNode;
@@ -408,6 +409,16 @@ public PlanNode visitTableWriter(TableWriterNode node, RewriteContext<Context> c
             return context.defaultRewrite(node, context.get());
         }
 
+        @Override
+        public PlanNode visitDelete(DeleteNode node, RewriteContext<Context> context)
+        {
+            if (node.getInputDistribution().isPresent()) {
+                context.get().variables.addAll(node.getInputDistribution().get().getInputVariables());
+            }
+            context.get().variables.add(node.getRowId());
+            return context.defaultRewrite(node, context.get());
+        }
+
         @Override
         public PlanNode visitTopN(TopNNode node, RewriteContext<Context> context)
         {