Cache and reuse ValidTxnWriteIdList

Author: deniskuzZ

ql/src/java/org/apache/hadoop/hive/ql/parse/CalcitePlanner.java

@@ -30,6 +30,7 @@
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.regex.Pattern;
 import org.antlr.runtime.ClassicToken;
 import org.antlr.runtime.CommonToken;
@@ -1700,6 +1701,11 @@ public RelNode apply(RelOptCluster cluster, RelOptSchema relOptSchema, SchemaPlu
           materializationValidator.getAutomaticRewritingValidationResult());
       perfLogger.perfLogEnd(this.getClass().getName(), PerfLogger.VALIDATE_QUERY_MATERIALIZATION);
 
+      // this step primes the cache containing the validTxnWriteIdList. It will fetch
+      // all the tables into the MetaStore Client cache with one HMS call. These are
+      // uses within the partition prune rule.
+      callAndCacheValidTxnWriteIdList(calcitePlan);
+
       // 2. Apply pre-join order optimizations
       perfLogger.perfLogBegin(this.getClass().getName(), PerfLogger.PREJOIN_ORDERING);
       calcitePlan = applyPreJoinOrderingTransforms(calcitePlan, mdProvider.getMetadataProvider(), executorProvider);
@@ -2454,15 +2460,20 @@ private RelNode applyPostJoinOrderingTransform(RelNode basePlan, RelMetadataProv
     }
 
     protected Set<TableName> getTablesUsed(RelNode plan) {
+      Predicate<Table> tableFilter = table -> AcidUtils.isTransactionalTable(table) ||
+          table.isNonNative() && table.getStorageHandler().areSnapshotsSupported();
+      return getTablesUsed(plan, tableFilter);
+    }
+
+    private Set<TableName> getTablesUsed(RelNode plan, Predicate<Table> tableFilter) {
       Set<TableName> tablesUsed = new HashSet<>();
       new RelVisitor() {
         @Override
         public void visit(RelNode node, int ordinal, RelNode parent) {
           if (node instanceof TableScan) {
             TableScan ts = (TableScan) node;
             Table table = ((RelOptHiveTable) ts.getTable()).getHiveTableMD();
-            if (AcidUtils.isTransactionalTable(table) ||
-                  table.isNonNative() && table.getStorageHandler().areSnapshotsSupported()) {
+            if (tableFilter.test(table)) {
               tablesUsed.add(table.getFullTableName());
             }
           }
@@ -5112,9 +5123,21 @@ private ImmutableMap<String, Integer> buildHiveColNameToInputPosMap(
       return hiveColNameToInputPosMapBuilder.build();
     }
 
-    private QBParseInfo getQBParseInfo(QB qb) throws CalciteSemanticException {
+    private QBParseInfo getQBParseInfo(QB qb) {
       return qb.getParseInfo();
     }
+
+    private void callAndCacheValidTxnWriteIdList(RelNode relNode) {
+      var transactionalTables = getTablesUsed(relNode, AcidUtils::isTransactionalTable)
+          .stream()
+          .map(TableName::getNotEmptyDbTable)
+          .toList();
+      try {
+        getValidTxnWriteIdList(transactionalTables);
+      } catch (SemanticException e) {
+        throw new RuntimeException(e.getMessage(), e);
+      }
+    }
   }
 
   @Override

ql/src/java/org/apache/hadoop/hive/ql/parse/SemanticAnalyzer.java

@@ -15166,30 +15166,29 @@ private String getQueryStringForCache(ASTNode ast) {
 
   private ValidTxnWriteIdList getQueryValidTxnWriteIdList() throws SemanticException {
     // TODO: Once HIVE-18948 is in, should be able to retrieve writeIdList from the conf.
-    //cachedWriteIdList = AcidUtils.getValidTxnWriteIdList(conf);
-    //
-    List<String> transactionalTables = tablesFromReadEntities(inputs)
-      .stream()
-      .filter(AcidUtils::isTransactionalTable)
-      .map(Table::getFullyQualifiedName)
-      .collect(Collectors.toList());
-    
-    if (transactionalTables.size() > 0) {
-      String txnString = queryState.getValidTxnList();
-      if (txnString == null) {
-        return null;
-      }
-      try {
-        return getTxnMgr().getValidWriteIds(transactionalTables, txnString);
-      } catch (Exception err) {
-        String msg = "Error while getting the txnWriteIdList for tables " + transactionalTables
-                + " and validTxnList " + conf.get(ValidTxnList.VALID_TXNS_KEY);
-        throw new SemanticException(msg, err);
-      }
-    }
+    // cachedWriteIdList = AcidUtils.getValidTxnWriteIdList(conf);
+    var transactionalTables = tablesFromReadEntities(inputs)
+        .stream()
+        .filter(AcidUtils::isTransactionalTable)
+        .map(Table::getFullyQualifiedName)
+        .toList();
 
-    // No transactional tables.
-    return null;
+    return getValidTxnWriteIdList(transactionalTables);
+  }
+
+  protected ValidTxnWriteIdList getValidTxnWriteIdList(List<String> transactionalTables)
+      throws SemanticException {
+    if (transactionalTables.isEmpty()) {
+      return null;
+    }
+    String txnString = queryState.getValidTxnList();
+    try {
+      return getTxnMgr().getValidWriteIds(transactionalTables, txnString);
+    } catch (Exception err) {
+      String msg = "Error while getting the txnWriteIdList for tables " + transactionalTables
+          + " and validTxnList " + txnString;
+      throw new SemanticException(msg, err);
+    }
   }
 
   private QueryResultsCache.LookupInfo createLookupInfoForQuery(ASTNode astNode) throws SemanticException {