HIVE-29413: Avoid code duplication by updating getPartCols method for iceberg tables

Author: ramitg254

ql/src/java/org/apache/hadoop/hive/ql/ddl/table/column/show/ShowColumnsOperation.java

@@ -66,10 +66,7 @@ private List<FieldSchema> getColumnsByPattern() throws HiveException {
 
   private List<FieldSchema> getCols() throws HiveException {
     Table table = context.getDb().getTable(desc.getTableName());
-    List<FieldSchema> allColumns = new ArrayList<>();
-    allColumns.addAll(table.getCols());
-    allColumns.addAll(table.getPartCols());
-    return allColumns;
+    return new ArrayList<>(table.getAllCols());
   }
 
   private Matcher getMatcher() {

ql/src/java/org/apache/hadoop/hive/ql/ddl/table/info/desc/DescTableOperation.java

@@ -130,7 +130,7 @@ private Deserializer getDeserializer(Table table) throws SQLException {
   private void getColumnsNoColumnPath(Table table, Partition partition, List<FieldSchema> cols) throws HiveException {
     cols.addAll(partition == null || table.getTableType() == TableType.VIRTUAL_VIEW ?
         table.getCols() : partition.getCols());
-    if (!desc.isFormatted()) {
+    if (!desc.isFormatted() && !table.hasNonNativePartitionSupport()) {
       cols.addAll(table.getPartCols());
     }
 

ql/src/java/org/apache/hadoop/hive/ql/ddl/table/info/desc/formatter/TextDescTableFormatter.java

@@ -174,9 +174,7 @@ private void addPartitionData(DataOutputStream out, HiveConf conf, String column
       List<FieldSchema> partitionColumns = null;
       // TODO (HIVE-29413): Refactor to a generic getPartCols() implementation
       if (table.isPartitioned()) {
-        partitionColumns = table.hasNonNativePartitionSupport() ?
-            table.getStorageHandler().getPartitionKeys(table) :
-            table.getPartCols();
+        partitionColumns = table.getPartCols();
       }
       if (CollectionUtils.isNotEmpty(partitionColumns) &&
           conf.getBoolVar(ConfVars.HIVE_DISPLAY_PARTITION_COLUMNS_SEPARATELY)) {

ql/src/java/org/apache/hadoop/hive/ql/ddl/table/partition/PartitionUtils.java

@@ -23,7 +23,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.Map.Entry;
 
 import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.conf.HiveConf.ConfVars;

ql/src/java/org/apache/hadoop/hive/ql/metadata/DummyPartition.java

@@ -91,9 +91,7 @@ public List<String> getValues() {
     values = new ArrayList<>();
 
     // TODO (HIVE-29413): Refactor to a generic getPartCols() implementation
-    for (FieldSchema fs : table.hasNonNativePartitionSupport()
-        ? table.getStorageHandler().getPartitionKeys(table)
-        : table.getPartCols()) {
+    for (FieldSchema fs : table.getPartCols()) {
       String val = partSpec.get(fs.getName());
       values.add(val);
     }

ql/src/java/org/apache/hadoop/hive/ql/metadata/Partition.java

@@ -173,7 +173,8 @@ protected void initialize(Table table,
           // set default if location is not set and this is a physical
           // table partition (not a view partition)
           if (table.getDataLocation() != null) {
-            Path partPath = new Path(table.getDataLocation(), Warehouse.makePartName(table.getPartCols(), tPartition.getValues()));
+            Path partPath = new Path(table.getDataLocation(),
+                Warehouse.makePartName(table.getPartCols(), tPartition.getValues()));
             tPartition.getSd().setLocation(partPath.toString());
           }
         }

ql/src/java/org/apache/hadoop/hive/ql/metadata/Table.java

@@ -57,6 +57,7 @@
 import org.apache.hadoop.hive.metastore.api.SerDeInfo;
 import org.apache.hadoop.hive.metastore.api.SkewedInfo;
 import org.apache.hadoop.hive.metastore.api.StorageDescriptor;
+import org.apache.hadoop.hive.metastore.HiveMetaHook;
 import org.apache.hadoop.hive.metastore.api.hive_metastoreConstants;
 import org.apache.hadoop.hive.metastore.utils.MetaStoreServerUtils;
 import org.apache.hadoop.hive.metastore.utils.MetaStoreUtils;
@@ -118,6 +119,8 @@ public class Table implements Serializable {
   private transient StorageHandlerInfo storageHandlerInfo;
   private transient MaterializedViewMetadata materializedViewMetadata;
 
+  private List<FieldSchema> cachedPartCols;
+
   private TableSpec tableSpec;
 
   private boolean materializedTable;
@@ -193,6 +196,9 @@ public Table makeCopy() {
 
     newTab.setMetaTable(this.getMetaTable());
     newTab.setSnapshotRef(this.getSnapshotRef());
+    if (this.cachedPartCols != null) {
+      newTab.cachedPartCols = new ArrayList<>(this.cachedPartCols);
+    }
     return newTab;
   }
 
@@ -214,6 +220,7 @@ public org.apache.hadoop.hive.metastore.api.Table getTTable() {
    */
   public void setTTable(org.apache.hadoop.hive.metastore.api.Table tTable) {
     this.tTable = tTable;
+    clearCachedPartCols();
   }
 
   /**
@@ -248,7 +255,7 @@ public void setTTable(org.apache.hadoop.hive.metastore.api.Table tTable) {
     org.apache.hadoop.hive.metastore.api.Table t = new org.apache.hadoop.hive.metastore.api.Table();
     {
       t.setSd(sd);
-      t.setPartitionKeys(new ArrayList<FieldSchema>());
+      t.setPartitionKeys(null);
       t.setParameters(new HashMap<String, String>());
       t.setTableType(TableType.MANAGED_TABLE.toString());
       t.setDbName(databaseName);
@@ -595,24 +602,43 @@ public boolean equals(Object obj) {
   }
 
   public List<FieldSchema> getPartCols() {
-    List<FieldSchema> partKeys = tTable.getPartitionKeys();
-    if (partKeys == null) {
-      partKeys = new ArrayList<>();
-      tTable.setPartitionKeys(partKeys);
+    if (cachedPartCols != null) {
+      return cachedPartCols;
+    }
+    List<FieldSchema> partKeys;
+    if (isTableTypeSet() && hasNonNativePartitionSupport()) {
+      partKeys = getStorageHandler().getPartitionKeys(this);
+    } else {
+      partKeys = tTable.getPartitionKeys();
+      if (partKeys == null) {
+        partKeys = new ArrayList<>();
+        tTable.setPartitionKeys(partKeys);
+      }
     }
+    cachedPartCols = partKeys;
     return partKeys;
   }
 
+  private void clearCachedPartCols() {
+    cachedPartCols = null;
+  }
+
+  private boolean isTableTypeSet() {
+    if (tTable.getParameters() == null) {
+      return false;
+    }
+    String tableType = tTable.getParameters().get(HiveMetaHook.TABLE_TYPE);
+    return tableType != null;
+  }
+
   public FieldSchema getPartColByName(String colName) {
     return getPartCols().stream()
       .filter(key -> key.getName().toLowerCase().equals(colName))
       .findFirst().orElse(null);
   }
 
   public List<String> getPartColNames() {
-    List<FieldSchema> partCols = hasNonNativePartitionSupport() ?
-        getStorageHandler().getPartitionKeys(this) : getPartCols();
-    return partCols.stream().map(FieldSchema::getName)
+    return getPartCols().stream().map(FieldSchema::getName)
       .collect(Collectors.toList());
   }
 
@@ -761,14 +787,22 @@ private List<FieldSchema> getColsInternal(boolean forMs) {
    * @return List&lt;FieldSchema&gt;
    */
   public List<FieldSchema> getAllCols() {
-    ArrayList<FieldSchema> f_list = new ArrayList<FieldSchema>();
-    f_list.addAll(getCols());
-    f_list.addAll(getPartCols());
-    return f_list;
+    List<FieldSchema> allCols = new ArrayList<>(getCols());
+    Set<String> colNames = new HashSet<>();
+    for (FieldSchema col : allCols) {
+      colNames.add(col.getName());
+    }
+    for (FieldSchema col : getPartCols()) {
+      if (!colNames.contains(col.getName())) {
+        allCols.add(col);
+      }
+    }
+    return allCols;
   }
 
   public void setPartCols(List<FieldSchema> partCols) {
     tTable.setPartitionKeys(partCols);
+    clearCachedPartCols();
   }
 
   public String getCatName() {
@@ -812,7 +846,7 @@ public void setOutputFormatClass(String name) throws HiveException {
   }
 
   public boolean isPartitioned() {
-    return hasNonNativePartitionSupport() ? getStorageHandler().isPartitioned(this) : 
+    return hasNonNativePartitionSupport() ? getStorageHandler().isPartitioned(this) :
         CollectionUtils.isNotEmpty(getPartCols());
   }
 
@@ -1153,7 +1187,7 @@ public static void validateColumns(List<FieldSchema> columns, List<FieldSchema>
       }
       colNames.add(colName);
     }
-    if (partCols != null) {
+    if (partCols != null && !icebergTable) {
       // there is no overlap between columns and partitioning columns
       for (FieldSchema partCol: partCols) {
         String colName = normalize(partCol.getName());

ql/src/java/org/apache/hadoop/hive/ql/optimizer/ColumnPrunerProcFactory.java

@@ -807,8 +807,7 @@ public Object process(Node nd, Stack<Node> stack, NodeProcessorCtx ctx,
         for (FieldNode col : cols) {
           int index = originalOutputColumnNames.indexOf(col.getFieldName());
           Table tab = cppCtx.getParseContext().getViewProjectToTableSchema().get(op);
-          List<FieldSchema> fullFieldList = new ArrayList<FieldSchema>(tab.getCols());
-          fullFieldList.addAll(tab.getPartCols());
+          List<FieldSchema> fullFieldList = new ArrayList<>(tab.getAllCols());
           cppCtx.getParseContext().getColumnAccessInfo()
               .add(tab.getCompleteName(), fullFieldList.get(index).getName());
         }

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

@@ -175,7 +175,8 @@ private void analyzeAcidExport(ASTNode ast, Table exportTable, ASTNode tokRefOrN
     //now generate insert statement
     //insert into newTableName select * from ts <where partition spec>
     StringBuilder rewrittenQueryStr = generateExportQuery(
-            newTable.getPartCols(), tokRefOrNameExportTable, (ASTNode) tokRefOrNameExportTable.parent, newTableName);
+        newTable.getPartCols(),
+        tokRefOrNameExportTable, (ASTNode) tokRefOrNameExportTable.parent, newTableName);
     ReparseResult rr = ParseUtils.parseRewrittenQuery(ctx, rewrittenQueryStr);
     Context rewrittenCtx = rr.rewrittenCtx;
     rewrittenCtx.setIsUpdateDeleteMerge(false); //it's set in parseRewrittenQuery()

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

@@ -83,7 +83,8 @@ public ColumnStatsAutoGatherContext(SemanticAnalyzer sa, HiveConf conf,
     this.isInsertInto = isInsertInto;
     this.origCtx = ctx;
     columns = tbl.getCols();
-    partitionColumns = tbl.getPartCols();
+    // current behaviour intact until we have getCols() giving only non-partition columns for non native tables as well
+    partitionColumns = tbl.hasNonNativePartitionSupport() ? new ArrayList<>() : tbl.getPartCols();
   }
 
   public List<LoadFileDesc> getLoadFileWork() {