fix: Info and infoUnion should always be synced

Author: shangm2

presto-main-base/src/main/java/com/facebook/presto/event/QueryMonitor.java

@@ -41,7 +41,7 @@
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.metadata.SessionPropertyManager;
 import com.facebook.presto.operator.OperatorInfo;
-import com.facebook.presto.operator.OperatorStats;
+import com.facebook.presto.operator.OperatorInfoUnion;
 import com.facebook.presto.operator.TableFinishInfo;
 import com.facebook.presto.operator.TaskStats;
 import com.facebook.presto.server.BasicQueryInfo;
@@ -82,6 +82,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.stream.Collectors;
 
@@ -601,10 +602,22 @@ private static QueryIOMetadata getQueryIOMetadata(QueryInfo queryInfo)
 
         Optional<QueryOutputMetadata> output = Optional.empty();
         if (queryInfo.getOutput().isPresent()) {
+            // Check both info (JSON) and infoUnion (Thrift) fields for TableFinishInfo
             Optional<TableFinishInfo> tableFinishInfo = queryInfo.getQueryStats().getOperatorSummaries().stream()
-                    .map(OperatorStats::getInfo)
-                    .filter(TableFinishInfo.class::isInstance)
-                    .map(TableFinishInfo.class::cast)
+                    .map(operatorStats -> {
+                        // First try the info field (JSON serialization)
+                        OperatorInfo info = operatorStats.getInfo();
+                        if (info instanceof TableFinishInfo) {
+                            return (TableFinishInfo) info;
+                        }
+                        // Fall back to infoUnion field (Thrift serialization)
+                        OperatorInfoUnion infoUnion = operatorStats.getInfoUnion();
+                        if (infoUnion != null) {
+                            return infoUnion.getTableFinishInfo();
+                        }
+                        return null;
+                    })
+                    .filter(Objects::nonNull)
                     .findFirst();
 
             Optional<List<OutputColumnMetadata>> outputColumnsMetadata = queryInfo.getOutput().get().getColumns()

presto-main-base/src/main/java/com/facebook/presto/operator/OperatorInfoUnion.java

@@ -257,7 +257,7 @@ else if (infoUnion.getTableWriterMergeInfo() != null) {
             return infoUnion.getTableWriterMergeInfo();
         }
         else {
-            throw new IllegalArgumentException("OperatorInfoUnion is of an unknown type: " + infoUnion.getClass().getName());
+            return null;
         }
     }
 }

presto-main-base/src/main/java/com/facebook/presto/operator/OperatorStats.java

@@ -243,7 +243,7 @@ public OperatorStats(
         this.blockedReason = blockedReason;
 
         this.info = info;
-        this.infoUnion = null;
+        this.infoUnion = (info != null) ? OperatorInfoUnion.convertToOperatorInfoUnion(info) : null;
         this.nullJoinBuildKeyCount = nullJoinBuildKeyCount;
         this.joinBuildKeyCount = joinBuildKeyCount;
         this.nullJoinProbeKeyCount = nullJoinProbeKeyCount;
@@ -390,7 +390,7 @@ public OperatorStats(
         this.blockedReason = blockedReason;
 
         this.infoUnion = infoUnion;
-        this.info = null;
+        this.info = (infoUnion != null) ? OperatorInfoUnion.convertToOperatorInfo(infoUnion) : null;
         this.nullJoinBuildKeyCount = nullJoinBuildKeyCount;
         this.joinBuildKeyCount = joinBuildKeyCount;
         this.nullJoinProbeKeyCount = nullJoinProbeKeyCount;
@@ -668,13 +668,19 @@ public Optional<BlockedReason> getBlockedReason()
     @JsonProperty
     public OperatorInfo getInfo()
     {
+        if (info == null && infoUnion != null) {
+            return OperatorInfoUnion.convertToOperatorInfo(infoUnion);
+        }
         return info;
     }
 
     @Nullable
     @ThriftField(39)
     public OperatorInfoUnion getInfoUnion()
     {
+        if (infoUnion == null && info != null) {
+            return OperatorInfoUnion.convertToOperatorInfoUnion(info);
+        }
         return infoUnion;
     }