From 66fe5236b3db45a55b4db1a140960aeec0b8eb51 Mon Sep 17 00:00:00 2001
From: NeerajUnnikrishnan <neeraju14@gmail.com>
Date: Fri, 16 Jul 2021 11:30:17 -0400
Subject: [PATCH] Fix for bucketing column rename

---
 .../prestosql/plugin/hive/HiveBucketing.java  |  7 ++-----
 .../io/prestosql/plugin/hive/HiveUtil.java    |  3 ++-
 .../hive/TestHiveIntegrationSmokeTest.java    | 20 +++++++++++++++++++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveBucketing.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveBucketing.java
index 66f815c68..6325591ae 100644
--- a/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveBucketing.java
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveBucketing.java
@@ -49,7 +49,6 @@ import java.util.stream.Collectors;
 import static io.prestosql.plugin.hive.HiveBucketing.BucketingVersion.BUCKETING_V1;
 import static io.prestosql.plugin.hive.HiveBucketing.BucketingVersion.BUCKETING_V2;
 import static io.prestosql.plugin.hive.HiveColumnHandle.BUCKET_COLUMN_NAME;
-import static io.prestosql.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA;
 import static io.prestosql.plugin.hive.HiveUtil.getRegularColumnHandles;
 import static java.lang.String.format;
 import static java.util.Map.Entry;
@@ -152,9 +151,7 @@ public final class HiveBucketing
         for (String bucketColumnName : hiveBucketProperty.get().getBucketedBy()) {
             HiveColumnHandle bucketColumnHandle = map.get(bucketColumnName);
             if (bucketColumnHandle == null) {
-                throw new PrestoException(
-                        HIVE_INVALID_METADATA,
-                        format("Table '%s.%s' is bucketed on non-existent column '%s'", table.getDatabaseName(), table.getTableName(), bucketColumnName));
+                return Optional.empty();
             }
             bucketColumns.add(bucketColumnHandle);
         }
@@ -166,7 +163,7 @@ public final class HiveBucketing
 
     public static Optional<HiveBucketFilter> getHiveBucketFilter(Table table, TupleDomain<ColumnHandle> effectivePredicate)
     {
-        if (!table.getStorage().getBucketProperty().isPresent()) {
+        if (!getHiveBucketHandle(table).isPresent()) {
             return Optional.empty();
         }
 
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveUtil.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveUtil.java
index 1e115b926..744b8432a 100644
--- a/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveUtil.java
+++ b/presto-hive/src/main/java/io/prestosql/plugin/hive/HiveUtil.java
@@ -123,6 +123,7 @@ import static com.google.common.collect.Iterables.filter;
 import static com.google.common.collect.Lists.newArrayList;
 import static com.google.common.collect.Lists.transform;
 import static io.prestosql.plugin.hive.HiveBucketing.bucketedOnTimestamp;
+import static io.prestosql.plugin.hive.HiveBucketing.getHiveBucketHandle;
 import static io.prestosql.plugin.hive.HiveColumnHandle.bucketColumnHandle;
 import static io.prestosql.plugin.hive.util.CustomSplitConversionUtils.recreateSplitWithCustomInfo;
 import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
@@ -878,7 +879,7 @@ public final class HiveUtil
 
         // add hidden columns
         columns.add(HiveColumnHandle.pathColumnHandle());
-        if (table.getStorage().getBucketProperty().isPresent()) {
+        if (getHiveBucketHandle(table).isPresent()) {
             if (!bucketedOnTimestamp(table.getStorage().getBucketProperty().get(), table)) {
                 columns.add(bucketColumnHandle());
             }
diff --git a/presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveIntegrationSmokeTest.java b/presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveIntegrationSmokeTest.java
index 9695fff76..4e54ab294 100644
--- a/presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveIntegrationSmokeTest.java
+++ b/presto-hive/src/test/java/io/prestosql/plugin/hive/TestHiveIntegrationSmokeTest.java
@@ -3200,6 +3200,26 @@ public class TestHiveIntegrationSmokeTest
         assertUpdate("DROP TABLE test_drop_bucketing_column");
     }
 
+    @Test
+    private void testRenameBucketingColumn()
+    {
+        @Language("SQL") String createTable = "" +
+                "CREATE TABLE test_rename_bucketing_column\n" +
+                "WITH (\n" +
+                "  bucket_count = 5, bucketed_by = ARRAY ['orderstatus']\n" +
+                ")\n" +
+                "AS\n" +
+                "SELECT custkey, orderkey, orderstatus FROM orders";
+
+        assertUpdate(createTable, "SELECT count(*) FROM orders");
+        assertQuery("SELECT orderkey, orderstatus FROM test_rename_bucketing_column", "SELECT orderkey, orderstatus FROM orders");
+
+        assertUpdate("ALTER TABLE test_rename_bucketing_column RENAME COLUMN orderstatus TO orderstatus1");
+        assertQuery("SELECT orderkey, orderstatus1  FROM test_rename_bucketing_column", "SELECT orderkey, orderstatus  FROM orders");
+
+        assertUpdate("DROP TABLE test_rename_bucketing_column");
+    }
+
     @Test
     public void testAvroTypeValidation()
     {