[incubator-gobblin] branch master updated: [GOBBLIN-1142] Hive Distcp support filter on partitioned or snapshot tables

[hu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

hutran pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-gobblin.git


The following commit(s) were added to refs/heads/master by this push:
     new 4d76f45  [GOBBLIN-1142] Hive Distcp support filter on partitioned or snapshot tables
4d76f45 is described below

commit 4d76f4544e936af1ec1871143c95d86e25c239b0
Author: zhchen <zh...@linkedin.com>
AuthorDate: Thu May 7 14:35:15 2020 -0700

    [GOBBLIN-1142] Hive Distcp support filter on partitioned or snapshot tables
    
    Closes #2979 from zxcware/hive-distcp
---
 .../management/copy/hive/HiveDatasetFinder.java    | 15 +++++
 .../copy/predicates/TableTypeFilter.java           | 65 ++++++++++++++++++++++
 .../copy/hive/HiveDatasetFinderTest.java           | 48 ++++++++++++++++
 3 files changed, 128 insertions(+)

diff --git a/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinder.java b/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinder.java
index 98134b8..e04783f 100644
--- a/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinder.java
+++ b/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinder.java
@@ -62,6 +62,7 @@ import org.apache.gobblin.metrics.event.EventSubmitter;
 import org.apache.gobblin.metrics.event.sla.SlaEventSubmitter;
 import org.apache.gobblin.util.AutoReturnableObject;
 import org.apache.gobblin.util.ConfigUtils;
+import org.apache.gobblin.util.reflection.GobblinConstructorUtils;
 
 
 /**
@@ -76,6 +77,7 @@ public class HiveDatasetFinder implements IterableDatasetFinder<HiveDataset> {
   public static final String DB_KEY = HIVE_DATASET_PREFIX + ".database";
   public static final String TABLE_PATTERN_KEY = HIVE_DATASET_PREFIX + ".table.pattern";
   public static final String DEFAULT_TABLE_PATTERN = "*";
+  public static final String TABLE_FILTER = HIVE_DATASET_PREFIX + ".tableFilter";
 
   /*
    * By setting the prefix, only config keys with this prefix will be used to build a HiveDataset.
@@ -112,6 +114,7 @@ public class HiveDatasetFinder implements IterableDatasetFinder<HiveDataset> {
 
   protected Optional<String> configStoreUri;
   protected final Function<Table, String> configStoreDatasetUriBuilder;
+  protected final Optional<Predicate<Table>> tableFilter;
 
   protected final String datasetConfigPrefix;
   protected final ConfigClient configClient;
@@ -181,6 +184,14 @@ public class HiveDatasetFinder implements IterableDatasetFinder<HiveDataset> {
       throw new RuntimeException(e);
     }
     this.jobConfig = ConfigUtils.propertiesToConfig(properties);
+
+    String tableFilterPredicate = properties.getProperty(TABLE_FILTER);
+    if (StringUtils.isNotEmpty(tableFilterPredicate)) {
+      this.tableFilter = Optional.of((Predicate<Table>)GobblinConstructorUtils.invokeConstructor(
+          Predicate.class, tableFilterPredicate, properties));
+    } else {
+      this.tableFilter = Optional.absent();
+    }
   }
 
   protected static HiveMetastoreClientPool createClientPool(Properties properties) throws IOException {
@@ -249,6 +260,10 @@ public class HiveDatasetFinder implements IterableDatasetFinder<HiveDataset> {
 
           try (AutoReturnableObject<IMetaStoreClient> client = HiveDatasetFinder.this.clientPool.getClient()) {
             Table table = client.get().getTable(dbAndTable.getDb(), dbAndTable.getTable());
+            if (tableFilter.isPresent() && !tableFilter.get().apply(table)) {
+              continue;
+            }
+
             Config datasetConfig = getDatasetConfig(table);
             if (ConfigUtils.getBoolean(datasetConfig, HIVE_DATASET_IS_BLACKLISTED_KEY, DEFAULT_HIVE_DATASET_IS_BLACKLISTED_KEY)) {
               continue;
diff --git a/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/predicates/TableTypeFilter.java b/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/predicates/TableTypeFilter.java
new file mode 100644
index 0000000..b73610c
--- /dev/null
+++ b/gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/predicates/TableTypeFilter.java
@@ -0,0 +1,65 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.gobblin.data.management.copy.predicates;
+
+import java.util.Properties;
+
+import org.apache.hadoop.hive.metastore.api.Table;
+
+import com.google.common.base.Predicate;
+
+import javax.annotation.Nullable;
+
+
+/**
+ * A predicate to check if a hive {@link Table} is of a certain type in {@link TABLE_TYPE}
+ *
+ * <p> Example usage: {@link org.apache.gobblin.data.management.copy.hive.HiveDatasetFinder#TABLE_FILTER}
+ */
+public class TableTypeFilter implements Predicate<Table> {
+
+  public static final String FILTER_TYPE = "tableTypeFilter.type";
+
+  private enum TABLE_TYPE {
+    SNAPSHOT,
+    PARTITIONED
+  }
+
+  private final TABLE_TYPE tableType;
+
+  public TableTypeFilter(Properties props) {
+    tableType = TABLE_TYPE.valueOf(
+        props.getProperty(FILTER_TYPE, TABLE_TYPE.SNAPSHOT.name()).toUpperCase());
+  }
+
+  @Override
+  public boolean apply(@Nullable Table input) {
+    if (input == null) {
+      return false;
+    }
+
+    switch (tableType) {
+      case SNAPSHOT:
+        return input.getPartitionKeys() == null || input.getPartitionKeys().size() == 0;
+      case PARTITIONED:
+        return input.getPartitionKeys() != null && input.getPartitionKeys().size() > 0;
+      default:
+        throw new UnsupportedOperationException("Invalid type: " + tableType);
+    }
+  }
+}
diff --git a/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinderTest.java b/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinderTest.java
index 70ef3d1..0794520 100644
--- a/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinderTest.java
+++ b/gobblin-data-management/src/test/java/org/apache/gobblin/data/management/copy/hive/HiveDatasetFinderTest.java
@@ -24,6 +24,7 @@ import java.util.Properties;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.hive.metastore.IMetaStoreClient;
+import org.apache.hadoop.hive.metastore.api.FieldSchema;
 import org.apache.hadoop.hive.metastore.api.StorageDescriptor;
 import org.apache.hadoop.hive.metastore.api.Table;
 import org.mockito.Mockito;
@@ -36,6 +37,7 @@ import com.google.common.collect.SetMultimap;
 import com.google.common.collect.Sets;
 import com.typesafe.config.Config;
 
+import org.apache.gobblin.data.management.copy.predicates.TableTypeFilter;
 import org.apache.gobblin.hive.HiveMetastoreClientPool;
 import org.apache.gobblin.util.AutoReturnableObject;
 
@@ -61,6 +63,36 @@ public class HiveDatasetFinderTest {
   }
 
   @Test
+  public void testTableFilter() throws Exception {
+    List<HiveDatasetFinder.DbAndTable> dbAndTables = Lists.newArrayList();
+    dbAndTables.add(new HiveDatasetFinder.DbAndTable("db1", "table1"));
+    dbAndTables.add(new HiveDatasetFinder.DbAndTable("db1", "table2"));
+    dbAndTables.add(new DatePartitionTable("db1", "table3"));
+    dbAndTables.add(new DatePartitionTable("db1", "table4"));
+    HiveMetastoreClientPool pool = getTestPool(dbAndTables);
+
+    Properties properties = new Properties();
+    properties.put(HiveDatasetFinder.HIVE_DATASET_PREFIX + "." + WhitelistBlacklist.WHITELIST, "");
+
+    // Test snapshot only
+    properties.put(HiveDatasetFinder.TABLE_FILTER, TableTypeFilter.class.getName());
+    properties.put(TableTypeFilter.FILTER_TYPE, "snapshot");
+    HiveDatasetFinder finder = new TestHiveDatasetFinder(FileSystem.getLocal(new Configuration()), properties, pool);
+    List<HiveDataset> datasets = Lists.newArrayList(finder.getDatasetsIterator());
+    Assert.assertEquals(datasets.size(), 2);
+    Assert.assertTrue(datasets.stream().anyMatch(dataset -> dataset.getDbAndTable().toString().equals("db1.table1")));
+    Assert.assertTrue(datasets.stream().anyMatch(dataset -> dataset.getDbAndTable().toString().equals("db1.table2")));
+
+    // Test partitioned table only
+    properties.put(TableTypeFilter.FILTER_TYPE, "partitioned");
+    finder = new TestHiveDatasetFinder(FileSystem.getLocal(new Configuration()), properties, pool);
+    datasets = Lists.newArrayList(finder.getDatasetsIterator());
+    Assert.assertEquals(datasets.size(), 2);
+    Assert.assertTrue(datasets.stream().anyMatch(dataset -> dataset.getDbAndTable().toString().equals("db1.table3")));
+    Assert.assertTrue(datasets.stream().anyMatch(dataset -> dataset.getDbAndTable().toString().equals("db1.table4")));
+  }
+
+  @Test
   public void testException() throws Exception {
 
     List<HiveDatasetFinder.DbAndTable> dbAndTables = Lists.newArrayList();
@@ -201,6 +233,9 @@ public class HiveDatasetFinderTest {
       Table table = new Table();
       table.setDbName(dbAndTable.getDb());
       table.setTableName(dbAndTable.getTable());
+      if (dbAndTable instanceof DatePartitionTable) {
+        table.setPartitionKeys(((DatePartitionTable) dbAndTable).getPartitionKeys());
+      }
       StorageDescriptor sd = new StorageDescriptor();
       sd.setLocation("/tmp/test");
       table.setSd(sd);
@@ -235,4 +270,17 @@ public class HiveDatasetFinderTest {
     }
   }
 
+  private static class DatePartitionTable extends HiveDatasetFinder.DbAndTable {
+
+    static final FieldSchema DATE_PARTITION_KEY = new FieldSchema("datepartition", "String", "");
+
+    public DatePartitionTable(String db, String table) {
+      super(db, table);
+    }
+
+    public List<FieldSchema> getPartitionKeys() {
+      return Lists.newArrayList(DATE_PARTITION_KEY);
+    }
+  }
+
 }