Maintain partition schema even after altering table schema

Author: NeerajUnnikrishnan

presto-hive/src/main/java/io/prestosql/plugin/hive/HivePageSource.java

@@ -126,7 +126,12 @@ public HivePageSource(
             }
 
             if (columnMapping.getKind() == PREFILLED) {
-                prefilledValues[columnIndex] = typedPartitionKey(columnMapping.getPrefilledValue(), type, name);
+                if (columnMapping.getPrefilledValue() == null) {
+                    prefilledValues[columnIndex] = null;
+                }
+                else {
+                    prefilledValues[columnIndex] = typedPartitionKey(columnMapping.getPrefilledValue(), type, name);
+                }
             }
         }
         this.coercers = coercers.build();

presto-hive/src/main/java/io/prestosql/plugin/hive/HivePageSourceProvider.java

@@ -50,12 +50,15 @@
 import io.prestosql.spi.type.TypeManager;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hive.serde.serdeConstants;
+import org.apache.hadoop.hive.serde2.SerDeUtils;
 import org.eclipse.jetty.util.URIUtil;
 
 import javax.inject.Inject;
 
 import java.net.URI;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -75,8 +78,10 @@
 import static io.prestosql.plugin.hive.HivePageSourceProvider.ColumnMapping.toColumnHandles;
 import static io.prestosql.plugin.hive.HiveUtil.isPartitionFiltered;
 import static io.prestosql.plugin.hive.coercions.HiveCoercer.createCoercer;
+import static io.prestosql.plugin.hive.metastore.MetastoreUtil.META_PARTITION_COLUMNS;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toList;
+import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.META_TABLE_COLUMNS;
 
 public class HivePageSourceProvider
         implements ConnectorPageSourceProvider
@@ -170,6 +175,24 @@ private ConnectorPageSource createPageSourceInternal(ConnectorSession session,
         Configuration configuration = hdfsEnvironment.getConfiguration(
                 new HdfsEnvironment.HdfsContext(session, hiveSplit.getDatabase(), hiveSplit.getTable()), path);
 
+        Properties schema = hiveSplit.getSchema();
+        String columnNameDelimiter = schema.containsKey(serdeConstants.COLUMN_NAME_DELIMITER) ? schema
+                .getProperty(serdeConstants.COLUMN_NAME_DELIMITER) : String.valueOf(SerDeUtils.COMMA);
+        List<String> partitionColumnNames;
+        if (schema.containsKey(META_PARTITION_COLUMNS)) {
+            partitionColumnNames = Arrays.asList(schema.getProperty(META_PARTITION_COLUMNS).split(columnNameDelimiter));
+        }
+        else if (schema.containsKey(META_TABLE_COLUMNS)) {
+            partitionColumnNames = Arrays.asList(schema.getProperty(META_TABLE_COLUMNS).split(columnNameDelimiter));
+        }
+        else {
+            partitionColumnNames = new ArrayList<>();
+        }
+
+        List<String> tableColumns = hiveColumns.stream().map(cols -> cols.getName()).collect(toList());
+
+        List<String> missingColumns = tableColumns.stream().filter(cols -> !partitionColumnNames.contains(cols)).collect(toList());
+
         List<IndexMetadata> indexes = new ArrayList<>();
         if (indexCache != null && session.isHeuristicIndexFilterEnabled()) {
             indexes.addAll(this.indexCache.getIndices(session
@@ -221,7 +244,7 @@ session, hiveSplit, assignUniqueIndicesToPartitionColumns(hiveColumns), typeMana
                     hiveTable.getDisjunctCompactEffectivePredicate(),
                     hiveSplit.getBucketConversion(),
                     hiveSplit.getBucketNumber(),
-                    hiveSplit.getLastModifiedTime());
+                    hiveSplit.getLastModifiedTime(), missingColumns);
         }
 
         Optional<ConnectorPageSource> pageSource = createHivePageSource(
@@ -249,7 +272,7 @@ session, hiveSplit, assignUniqueIndicesToPartitionColumns(hiveColumns), typeMana
                 splitMetadata,
                 hiveSplit.isCacheable(),
                 hiveSplit.getLastModifiedTime(),
-                hiveSplit.getCustomSplitInfo());
+                hiveSplit.getCustomSplitInfo(), missingColumns);
         if (pageSource.isPresent()) {
             return pageSource.get();
         }
@@ -291,6 +314,7 @@ private static List<HiveColumnHandle> assignUniqueIndicesToPartitionColumns(List
      * @param predicateColumns Map of all columns handles being part of predicate
      * @param additionPredicates Predicates related to OR clause.
      * Remaining columns are same as for createHivePageSource.
+     * @param missingColumns
      * @return
      */
     private static ConnectorPageSource createSelectivePageSource(
@@ -310,7 +334,7 @@ private static ConnectorPageSource createSelectivePageSource(
             Optional<List<TupleDomain<HiveColumnHandle>>> additionPredicates,
             Optional<HiveSplit.BucketConversion> bucketConversion,
             OptionalInt bucketNumber,
-            long dataSourceLastModifiedTime)
+            long dataSourceLastModifiedTime, List<String> missingColumns)
     {
         Set<HiveColumnHandle> interimColumns = ImmutableSet.<HiveColumnHandle>builder()
                 .addAll(predicateColumns.values())
@@ -325,7 +349,7 @@ private static ConnectorPageSource createSelectivePageSource(
                 split.getColumnCoercions(),
                 path,
                 bucketNumber,
-                true);
+                true, missingColumns);
 
         List<ColumnMapping> regularAndInterimColumnMappings = ColumnMapping.extractRegularAndInterimColumnMappings(
                 columnMappings);
@@ -411,7 +435,7 @@ public static Optional<ConnectorPageSource> createHivePageSource(
             SplitMetadata splitMetadata,
             boolean splitCacheable,
             long dataSourceLastModifiedTime,
-            Map<String, String> customSplitInfo)
+            Map<String, String> customSplitInfo, List<String> missingColumns)
     {
         List<ColumnMapping> columnMappings = ColumnMapping.buildColumnMappings(
                 partitionKeys,
@@ -420,7 +444,7 @@ public static Optional<ConnectorPageSource> createHivePageSource(
                 columnCoercions,
                 path,
                 bucketNumber,
-                true);
+                true, missingColumns);
         List<ColumnMapping> regularAndInterimColumnMappings = ColumnMapping.extractRegularAndInterimColumnMappings(
                 columnMappings);
 
@@ -603,7 +627,7 @@ public ColumnMappingKind getKind()
         public String getPrefilledValue()
         {
             checkState(kind == ColumnMappingKind.PREFILLED);
-            return prefilledValue.get();
+            return prefilledValue.isPresent() ? prefilledValue.get() : HIVE_DEFAULT_PARTITION_VALUE;
         }
 
         public HiveColumnHandle getHiveColumnHandle()
@@ -628,6 +652,7 @@ public Optional<HiveType> getCoercionFrom()
          * @param requiredInterimColumns columns that are needed for processing, but shouldn't be returned to engine (may overlaps with columns)
          * @param columnCoercions map from hive column index to hive type
          * @param bucketNumber empty if table is not bucketed, a number within [0, # bucket in table) otherwise
+         * @param missingColumns
          */
         public static List<ColumnMapping> buildColumnMappings(
                 List<HivePartitionKey> partitionKeys,
@@ -636,7 +661,7 @@ public static List<ColumnMapping> buildColumnMappings(
                 Map<Integer, HiveType> columnCoercions,
                 Path path,
                 OptionalInt bucketNumber,
-                boolean filterPushDown)
+                boolean filterPushDown, List<String> missingColumns)
         {
             Map<String, HivePartitionKey> partitionKeysByName = uniqueIndex(partitionKeys, HivePartitionKey::getName);
             int regularIndex = 0;
@@ -645,6 +670,11 @@ public static List<ColumnMapping> buildColumnMappings(
             for (HiveColumnHandle column : columns) {
                 Optional<HiveType> coercionFrom = Optional.ofNullable(columnCoercions.get(column.getHiveColumnIndex()));
                 if (column.getColumnType() == REGULAR) {
+                    if (missingColumns.contains(column.getColumnName())) {
+                        columnMappings.add(new ColumnMapping(ColumnMappingKind.PREFILLED, column, Optional.empty(),
+                                OptionalInt.empty(), coercionFrom));
+                        continue;
+                    }
                     checkArgument(regularColumnIndices.add(column.getHiveColumnIndex()), "duplicate hiveColumnIndex in columns list");
 
                     columnMappings.add(regular(column, regularIndex, coercionFrom));

presto-hive/src/main/java/io/prestosql/plugin/hive/metastore/MetastoreUtil.java

@@ -29,6 +29,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.isNullOrEmpty;
@@ -53,6 +54,9 @@
 
 public class MetastoreUtil
 {
+    public static final String META_PARTITION_COLUMNS = "partition_metadata";
+    public static final String COLUMN_COMMENTS = "columns.comments";
+
     private MetastoreUtil()
     {
     }
@@ -132,10 +136,14 @@ private static Properties getHiveSchema(
             first = false;
         }
         String columnNames = columnNameBuilder.toString();
+        String partitionColumnNames = dataColumns.stream()
+                .map(Column::getName)
+                .collect(Collectors.joining(","));
         String columnTypes = columnTypeBuilder.toString();
         schema.setProperty(META_TABLE_COLUMNS, columnNames);
+        schema.setProperty(META_PARTITION_COLUMNS, partitionColumnNames);
         schema.setProperty(META_TABLE_COLUMN_TYPES, columnTypes);
-        schema.setProperty("columns.comments", columnCommentBuilder.toString());
+        schema.setProperty(COLUMN_COMMENTS, columnCommentBuilder.toString());
 
         schema.setProperty(SERIALIZATION_DDL, toThriftDdl(tableName, dataColumns));
 

presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveFileFormats.java

@@ -605,7 +605,7 @@ private void testCursorProvider(HiveRecordCursorProvider cursorProvider,
                 null,
                 false,
                 -1L,
-                ImmutableMap.of());
+                ImmutableMap.of(), ImmutableList.of());
 
         RecordCursor cursor = ((RecordPageSource) pageSource.get()).getCursor();
 
@@ -658,7 +658,7 @@ private void testPageSourceFactory(HivePageSourceFactory sourceFactory,
                 null,
                 false,
                 -1L,
-                ImmutableMap.of());
+                ImmutableMap.of(), ImmutableList.of());
 
         assertTrue(pageSource.isPresent());
 

presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveIntegrationSmokeTest.java

@@ -3150,10 +3150,10 @@ public void testRenameColumn()
 
         assertUpdate(createTable, "SELECT count(*) FROM orders");
         assertUpdate("ALTER TABLE test_rename_column RENAME COLUMN orderkey TO new_orderkey");
-        assertQuery("SELECT new_orderkey, orderstatus FROM test_rename_column", "SELECT orderkey, orderstatus FROM orders");
+        assertQuery("SELECT new_orderkey, orderstatus FROM test_rename_column", "SELECT NULL, orderstatus FROM orders where orderstatus != 'dfd'");
         assertQueryFails("ALTER TABLE test_rename_column RENAME COLUMN \"$path\" TO test", ".* Cannot rename hidden column");
         assertQueryFails("ALTER TABLE test_rename_column RENAME COLUMN orderstatus TO new_orderstatus", "Renaming partition columns is not supported");
-        assertQuery("SELECT new_orderkey, orderstatus FROM test_rename_column", "SELECT orderkey, orderstatus FROM orders");
+        assertQuery("SELECT new_orderkey, orderstatus FROM test_rename_column", "SELECT NULL, orderstatus FROM orders");
         assertUpdate("DROP TABLE test_rename_column");
     }
 

presto-hive/src/test/java/io/prestosql/plugin/hive/TestOrcPageSourceMemoryTracking.java

@@ -517,7 +517,7 @@ public ConnectorPageSource newPageSource(FileFormatDataSourceStats stats, Connec
                     null,
                     false,
                     -1L,
-                    ImmutableMap.of())
+                    ImmutableMap.of(), ImmutableList.of())
                     .get();
         }
 

presto-hive/src/test/java/io/prestosql/plugin/hive/metastore/TestMetastoreUtil.java

@@ -29,6 +29,7 @@
 import java.util.List;
 import java.util.Properties;
 
+import static io.prestosql.plugin.hive.metastore.MetastoreUtil.META_PARTITION_COLUMNS;
 import static org.apache.hadoop.hive.serde.serdeConstants.COLUMN_NAME_DELIMITER;
 import static org.testng.Assert.assertEquals;
 
@@ -142,6 +143,7 @@ public void testHiveSchemaTable()
         Properties expected = MetaStoreUtils.getTableMetadata(TEST_TABLE_WITH_UNSUPPORTED_FIELDS);
         expected.remove(COLUMN_NAME_DELIMITER);
         Properties actual = MetastoreUtil.getHiveSchema(ThriftMetastoreUtil.fromMetastoreApiTable(TEST_TABLE_WITH_UNSUPPORTED_FIELDS, TEST_SCHEMA));
+        actual.remove(META_PARTITION_COLUMNS);
         assertEquals(actual, expected);
     }
 
@@ -151,6 +153,7 @@ public void testHiveSchemaPartition()
         Properties expected = MetaStoreUtils.getPartitionMetadata(TEST_PARTITION_WITH_UNSUPPORTED_FIELDS, TEST_TABLE_WITH_UNSUPPORTED_FIELDS);
         expected.remove(COLUMN_NAME_DELIMITER);
         Properties actual = MetastoreUtil.getHiveSchema(ThriftMetastoreUtil.fromMetastoreApiPartition(TEST_PARTITION_WITH_UNSUPPORTED_FIELDS), ThriftMetastoreUtil.fromMetastoreApiTable(TEST_TABLE_WITH_UNSUPPORTED_FIELDS, TEST_SCHEMA));
+        actual.remove(META_PARTITION_COLUMNS);
         assertEquals(actual, expected);
     }