Refactor HudiSplitSource (trinodb#16)

* Refactor HudiSplitSource

* Address some code style issues

* Remove metaClient and partitions from HudiTableHandle

* Use intersect for HudiPredicates

* Refactor HudiUtil.getSplits to HudiSplitFactory

* Fix error-prone-checks failure

* Extract HudiTableInfo

Author: 7c00

plugin/trino-hudi/pom.xml

@@ -49,6 +49,10 @@
             <groupId>io.airlift</groupId>
             <artifactId>bootstrap</artifactId>
         </dependency>
+        <dependency>
+            <groupId>io.airlift</groupId>
+            <artifactId>concurrent</artifactId>
+        </dependency>
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>configuration</artifactId>
@@ -82,6 +86,10 @@
             <groupId>com.google.inject</groupId>
             <artifactId>guice</artifactId>
         </dependency>
+        <dependency>
+            <groupId>javax.annotation</groupId>
+            <artifactId>javax.annotation-api</artifactId>
+        </dependency>
         <dependency>
             <groupId>javax.inject</groupId>
             <artifactId>javax.inject</artifactId>

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/ForHudiSplitManager.java

@@ -0,0 +1,31 @@
+/*
+ * Licensed 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 io.trino.plugin.hudi;
+
+import javax.inject.Qualifier;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target({FIELD, PARAMETER, METHOD})
+@Qualifier
+public @interface ForHudiSplitManager
+{
+}

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiConfig.java

@@ -39,6 +39,8 @@ public class HudiConfig
     private boolean sizeBasedSplitWeightsEnabled = true;
     private DataSize standardSplitWeightSize = DataSize.of(128, MEGABYTE);
     private double minimumAssignedSplitWeight = 0.05;
+    private int maxSplitsPerSecond = Integer.MAX_VALUE;
+    private int maxOutstandingSplits = 1000;
 
     @NotNull
     public String getBaseFileFormat()
@@ -194,4 +196,30 @@ public double getMinimumAssignedSplitWeight()
     {
         return minimumAssignedSplitWeight;
     }
+
+    @Min(1)
+    public int getMaxSplitsPerSecond()
+    {
+        return maxSplitsPerSecond;
+    }
+
+    @Config("hudi.max-splits-per-second")
+    public HudiConfig setMaxSplitsPerSecond(int maxSplitsPerSecond)
+    {
+        this.maxSplitsPerSecond = maxSplitsPerSecond;
+        return this;
+    }
+
+    @Min(1)
+    public int getMaxOutstandingSplits()
+    {
+        return maxOutstandingSplits;
+    }
+
+    @Config("hudi.max-outstanding-splits")
+    public HudiConfig setMaxOutstandingSplits(int maxOutstandingSplits)
+    {
+        this.maxOutstandingSplits = maxOutstandingSplits;
+        return this;
+    }
 }

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiConnector.java

@@ -50,7 +50,6 @@ public class HudiConnector
 
     private final LifeCycleManager lifeCycleManager;
     private final HudiTransactionManager transactionManager;
-    private final HudiMetadataFactory metadataFactory;
     private final ConnectorSplitManager splitManager;
     private final ConnectorPageSourceProvider pageSourceProvider;
     private final ConnectorNodePartitioningProvider nodePartitioningProvider;
@@ -62,7 +61,6 @@ public class HudiConnector
     public HudiConnector(
             LifeCycleManager lifeCycleManager,
             HudiTransactionManager transactionManager,
-            HudiMetadataFactory metadataFactory,
             ConnectorSplitManager splitManager,
             ConnectorPageSourceProvider pageSourceProvider,
             ConnectorNodePartitioningProvider nodePartitioningProvider,
@@ -73,7 +71,6 @@ public HudiConnector(
     {
         this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null");
         this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
-        this.metadataFactory = requireNonNull(metadataFactory, "metadataFactory is null");
         this.splitManager = requireNonNull(splitManager, "splitManager is null");
         this.pageSourceProvider = requireNonNull(pageSourceProvider, "pageSourceProvider is null");
         this.nodePartitioningProvider = requireNonNull(nodePartitioningProvider, "nodePartitioningProvider is null");

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiErrorCode.java

@@ -33,7 +33,7 @@ public enum HudiErrorCode
 
     HudiErrorCode(int code, ErrorType type)
     {
-        errorCode = new ErrorCode(code + 0x0100_0000, name(), type);
+        errorCode = new ErrorCode(code + 0x0507_0000, name(), type);
     }
 
     @Override

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java

@@ -20,7 +20,6 @@
 import io.airlift.log.Logger;
 import io.trino.plugin.hive.HdfsEnvironment;
 import io.trino.plugin.hive.HiveColumnHandle;
-import io.trino.plugin.hive.HivePartition;
 import io.trino.plugin.hive.acid.AcidSchema;
 import io.trino.plugin.hive.metastore.Column;
 import io.trino.plugin.hive.metastore.HiveMetastore;
@@ -32,23 +31,23 @@
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorTableHandle;
 import io.trino.spi.connector.ConnectorTableMetadata;
-import io.trino.spi.connector.ConnectorTableProperties;
 import io.trino.spi.connector.Constraint;
 import io.trino.spi.connector.ConstraintApplicationResult;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.connector.SchemaTablePrefix;
+import io.trino.spi.connector.TableColumnsMetadata;
 import io.trino.spi.connector.TableNotFoundException;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.type.TypeManager;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
 import org.apache.hudi.common.model.HoodieTableType;
-import org.apache.hudi.common.table.HoodieTableMetaClient;
 
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.stream.Stream;
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
@@ -67,7 +66,6 @@
 import static io.trino.plugin.hive.util.HiveUtil.hiveColumnHandles;
 import static io.trino.plugin.hive.util.HiveUtil.isHiveSystemSchema;
 import static io.trino.plugin.hudi.HudiErrorCode.HUDI_UNKNOWN_TABLE_TYPE;
-import static io.trino.plugin.hudi.HudiUtil.splitPredicate;
 import static io.trino.spi.connector.SchemaTableName.schemaTableName;
 import static java.lang.String.format;
 import static java.util.Collections.singletonList;
@@ -120,8 +118,7 @@ public HudiTableHandle getTableHandle(ConnectorSession session, SchemaTableName
                 table.get().getStorage().getLocation(),
                 HoodieTableType.COPY_ON_WRITE,
                 TupleDomain.all(),
-                TupleDomain.all(),
-                Optional.of(getTableMetaClient(session, table.get())));
+                TupleDomain.all());
     }
 
     @Override
@@ -135,12 +132,14 @@ public ConnectorTableMetadata getTableMetadata(ConnectorSession session, Connect
     public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(ConnectorSession session, ConnectorTableHandle tableHandle, Constraint constraint)
     {
         HudiTableHandle handle = (HudiTableHandle) tableHandle;
-        HudiPredicates predicates = splitPredicate(constraint.getSummary());
-        HudiTableHandle newHudiTableHandle = handle.withPredicates(predicates);
+        HudiPredicates predicates = HudiPredicates.from(constraint.getSummary());
+        HudiTableHandle newHudiTableHandle = handle.withPredicates(
+                predicates.getPartitionColumnPredicates(),
+                predicates.getRegularColumnPredicates());
 
         if (handle.getPartitionPredicates().equals(newHudiTableHandle.getPartitionPredicates())
                 && handle.getRegularPredicates().equals(newHudiTableHandle.getRegularPredicates())) {
-            log.info("No new predicates to apply");
+            log.debug("No new predicates to apply");
             return Optional.empty();
         }
 
@@ -150,12 +149,6 @@ public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(C
                 false));
     }
 
-    @Override
-    public ConnectorTableProperties getTableProperties(ConnectorSession session, ConnectorTableHandle tableHandle)
-    {
-        return new ConnectorTableProperties();
-    }
-
     @Override
     public Map<String, ColumnHandle> getColumnHandles(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
@@ -175,11 +168,8 @@ public ColumnMetadata getColumnMetadata(ConnectorSession session, ConnectorTable
     @Override
     public Optional<Object> getInfo(ConnectorTableHandle table)
     {
-        return ((HudiTableHandle) table).getPartitions()
-                .map(partitions -> new HudiInputInfo(
-                        partitions.stream()
-                                .map(HivePartition::getPartitionId)
-                                .collect(toImmutableList())));
+        HudiTableHandle hudiTable = (HudiTableHandle) table;
+        return Optional.of(HudiTableInfo.from(hudiTable));
     }
 
     @Override
@@ -191,40 +181,26 @@ public List<SchemaTableName> listTables(ConnectorSession session, Optional<Strin
                 tableNames.add(schemaTableName(schemaName, tableName));
             }
         }
-
-        tableNames.addAll(listMaterializedViews(session, optionalSchemaName));
         return tableNames.build();
     }
 
     @Override
-    public Map<SchemaTableName, List<ColumnMetadata>> listTableColumns(ConnectorSession session, SchemaTablePrefix prefix)
+    public Stream<TableColumnsMetadata> streamTableColumns(ConnectorSession session, SchemaTablePrefix prefix)
     {
         List<SchemaTableName> tables = prefix.getTable()
                 .map(ignored -> singletonList(prefix.toSchemaTableName()))
                 .orElseGet(() -> listTables(session, prefix.getSchema()));
-
-        ImmutableMap.Builder<SchemaTableName, List<ColumnMetadata>> columns = ImmutableMap.builder();
-        for (SchemaTableName table : tables) {
-            try {
-                columns.put(table, getTableMetadata(table).getColumns());
-            }
-            catch (TableNotFoundException e) {
-                // table disappeared during listing operation
-            }
-        }
-        return columns.buildOrThrow();
+        return tables.stream().map(table -> {
+            List<ColumnMetadata> columns = getTableMetadata(table).getColumns();
+            return TableColumnsMetadata.forTable(table, columns);
+        });
     }
 
     HiveMetastore getMetastore()
     {
         return metastore;
     }
 
-    void rollback()
-    {
-        // TODO: cleanup open transaction when write will be supported
-    }
-
     private static Function<HiveColumnHandle, ColumnMetadata> columnMetadataGetter(Table table)
     {
         ImmutableList.Builder<String> columnNames = ImmutableList.builder();
@@ -288,13 +264,6 @@ private boolean isHudiTable(ConnectorSession session, Table table)
         return true;
     }
 
-    private HoodieTableMetaClient getTableMetaClient(ConnectorSession session, Table table)
-    {
-        String basePath = table.getStorage().getLocation();
-        Configuration conf = hdfsEnvironment.getConfiguration(new HdfsEnvironment.HdfsContext(session), new Path(basePath));
-        return HoodieTableMetaClient.builder().setConf(conf).setBasePath(basePath).build();
-    }
-
     private ConnectorTableMetadata getTableMetadata(SchemaTableName tableName)
     {
         Table table = metastore.getTable(tableName.getSchemaName(), tableName.getTableName())

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiModule.java

@@ -16,6 +16,7 @@
 
 import com.google.inject.Binder;
 import com.google.inject.Module;
+import com.google.inject.Provides;
 import com.google.inject.Scopes;
 import io.trino.plugin.base.security.AllowAllAccessControl;
 import io.trino.plugin.base.session.SessionPropertiesProvider;
@@ -32,9 +33,15 @@
 import io.trino.spi.connector.ConnectorPageSourceProvider;
 import io.trino.spi.connector.ConnectorSplitManager;
 
+import javax.inject.Singleton;
+
+import java.util.concurrent.ExecutorService;
+
 import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
+import static io.airlift.concurrent.Threads.daemonThreadsNamed;
 import static io.airlift.configuration.ConfigBinder.configBinder;
+import static java.util.concurrent.Executors.newCachedThreadPool;
 import static org.weakref.jmx.guice.ExportBinder.newExporter;
 
 public class HudiModule
@@ -67,4 +74,12 @@ public void configure(Binder binder)
 
         newOptionalBinder(binder, ConnectorAccessControl.class).setDefault().to(AllowAllAccessControl.class).in(Scopes.SINGLETON);
     }
+
+    @ForHudiSplitManager
+    @Singleton
+    @Provides
+    public ExecutorService createExecutorService()
+    {
+        return newCachedThreadPool(daemonThreadsNamed("hudi-split-manager-%d"));
+    }
 }

plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiPageSource.java

@@ -25,6 +25,7 @@
 import java.util.List;
 import java.util.Map;
 
+import static io.trino.plugin.base.util.Closables.closeAllSuppress;
 import static io.trino.plugin.hudi.HudiErrorCode.HUDI_BAD_DATA;
 import static java.util.Objects.requireNonNull;
 
@@ -89,11 +90,11 @@ public Page getNextPage()
             return new Page(positionCount, blocksWithPartitionColumns);
         }
         catch (TrinoException e) {
-            closeWithSuppression(e);
+            closeAllSuppress(e, this);
             throw e;
         }
         catch (RuntimeException e) {
-            closeWithSuppression(e);
+            closeAllSuppress(e, this);
             throw new TrinoException(HUDI_BAD_DATA, e);
         }
     }
@@ -110,18 +111,4 @@ public void close()
     {
         dataPageSource.close();
     }
-
-    private void closeWithSuppression(Throwable throwable)
-    {
-        requireNonNull(throwable, "throwable is null");
-        try {
-            close();
-        }
-        catch (Exception e) {
-            // Self-suppression not permitted
-            if (e != throwable) {
-                throwable.addSuppressed(e);
-            }
-        }
-    }
 }