Handle partition and add support for metadata listing

plugin/trino-hive/src/main/java/io/trino/plugin/hive/BackgroundHiveSplitLoader.java

@@ -28,7 +28,6 @@
 import io.trino.plugin.hive.HiveSplit.BucketConversion;
 import io.trino.plugin.hive.HiveSplit.BucketValidation;
 import io.trino.plugin.hive.acid.AcidTransaction;
-import io.trino.plugin.hive.metastore.Column;
 import io.trino.plugin.hive.metastore.Partition;
 import io.trino.plugin.hive.metastore.Table;
 import io.trino.plugin.hive.util.HiveBucketing.BucketingVersion;
@@ -100,8 +99,6 @@
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_BAD_DATA;
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR;
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_BUCKET_FILES;
-import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA;
-import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_PARTITION_VALUE;
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_UNKNOWN_ERROR;
 import static io.trino.plugin.hive.HivePartitionManager.partitionMatches;
 import static io.trino.plugin.hive.HiveSessionProperties.getMaxInitialSplitSize;
@@ -114,11 +111,11 @@
 import static io.trino.plugin.hive.util.HiveFileIterator.NestedDirectoryPolicy.FAIL;
 import static io.trino.plugin.hive.util.HiveFileIterator.NestedDirectoryPolicy.IGNORED;
 import static io.trino.plugin.hive.util.HiveFileIterator.NestedDirectoryPolicy.RECURSE;
-import static io.trino.plugin.hive.util.HiveUtil.checkCondition;
 import static io.trino.plugin.hive.util.HiveUtil.getFooterCount;
 import static io.trino.plugin.hive.util.HiveUtil.getHeaderCount;
 import static io.trino.plugin.hive.util.HiveUtil.getInputFormat;
 import static io.trino.plugin.hive.util.HiveUtil.getPartitionKeyColumnHandles;
+import static io.trino.plugin.hive.util.HiveUtil.getPartitionKeys;
 import static io.trino.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.lang.Integer.parseInt;
 import static java.lang.Math.max;
@@ -942,28 +939,6 @@ private static List<Path> getTargetPathsFromSymlink(FileSystem fileSystem, Path
  }
  }
 
- private static List<HivePartitionKey> getPartitionKeys(Table table, Optional<Partition> partition)
- {
- if (partition.isEmpty()) {
- return ImmutableList.of();
- }
- ImmutableList.Builder<HivePartitionKey> partitionKeys = ImmutableList.builder();
- List<Column> keys = table.getPartitionColumns();
- List<String> values = partition.get().getValues();
- checkCondition(keys.size() == values.size(), HIVE_INVALID_METADATA, "Expected %s partition key values, but got %s", keys.size(), values.size());
- for (int i = 0; i < keys.size(); i++) {
- String name = keys.get(i).getName();
- HiveType hiveType = keys.get(i).getType();
- if (!hiveType.isSupportedType(table.getStorage().getStorageFormat())) {
- throw new TrinoException(NOT_SUPPORTED, format("Unsupported Hive type %s found in partition keys of table %s.%s", hiveType, table.getDatabaseName(), table.getTableName()));
- }
- String value = values.get(i);
- checkCondition(value != null, HIVE_INVALID_PARTITION_VALUE, "partition key value cannot be null for field: %s", name);
- partitionKeys.add(new HivePartitionKey(name, value));
- }
- return partitionKeys.build();
- }
-
  private static Properties getPartitionSchema(Table table, Optional<Partition> partition)
  {
  if (partition.isEmpty()) {

plugin/trino-hive/src/main/java/io/trino/plugin/hive/util/HiveUtil.java

@@ -32,6 +32,7 @@
 import io.trino.plugin.hive.HiveType;
 import io.trino.plugin.hive.avro.TrinoAvroSerDe;
 import io.trino.plugin.hive.metastore.Column;
+import io.trino.plugin.hive.metastore.Partition;
 import io.trino.plugin.hive.metastore.SortingColumn;
 import io.trino.plugin.hive.metastore.Table;
 import io.trino.spi.ErrorCodeSupplier;
@@ -1143,4 +1144,26 @@ public static boolean isIcebergTable(Table table)
  {
  return ICEBERG_TABLE_TYPE_VALUE.equalsIgnoreCase(table.getParameters().get(ICEBERG_TABLE_TYPE_NAME));
  }
+
+ public static List<HivePartitionKey> getPartitionKeys(Table table, Optional<Partition> partition)
+ {
+ if (partition.isEmpty()) {
+ return ImmutableList.of();
+ }
+ ImmutableList.Builder<HivePartitionKey> partitionKeys = ImmutableList.builder();
+ List<Column> keys = table.getPartitionColumns();
+ List<String> values = partition.get().getValues();
+ checkCondition(keys.size() == values.size(), HIVE_INVALID_METADATA, "Expected %s partition key values, but got %s", keys.size(), values.size());
+ for (int i = 0; i < keys.size(); i++) {
+ String name = keys.get(i).getName();
+ HiveType hiveType = keys.get(i).getType();
+ if (!hiveType.isSupportedType(table.getStorage().getStorageFormat())) {
+ throw new TrinoException(NOT_SUPPORTED, format("Unsupported Hive type %s found in partition keys of table %s.%s", hiveType, table.getDatabaseName(), table.getTableName()));
+ }
+ String value = values.get(i);
+ checkCondition(value != null, HIVE_INVALID_PARTITION_VALUE, "partition key value cannot be null for field: %s", name);
+ partitionKeys.add(new HivePartitionKey(name, value));
+ }
+ return partitionKeys.build();
+ }
 }

plugin/trino-hudi/pom.xml

@@ -17,7 +17,7 @@
 
  <properties>
  <air.main.basedir>${project.parent.basedir}</air.main.basedir>
- <dep.hudi.version>0.9.0</dep.hudi.version>
+ <dep.hudi.version>0.10.0</dep.hudi.version>
  </properties>
 
  <dependencies>
@@ -124,6 +124,45 @@
  </exclusion>
  </exclusions>
  </dependency>
+ <dependency>
+ <groupId>org.apache.hudi</groupId>
+ <artifactId>hudi-hadoop-mr</artifactId>
+ <version>${dep.hudi.version}</version>
+ <exclusions>
+ <exclusion>
+ <groupId>org.apache.hbase</groupId>
+ <artifactId>hbase-server</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>org.apache.orc</groupId>
+ <artifactId>orc-core</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-annotations</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-databind</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>org.apache.httpcomponents</groupId>
+ <artifactId>httpclient</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>org.apache.httpcomponents</groupId>
+ <artifactId>fluent-hc</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>org.rocksdb</groupId>
+ <artifactId>rocksdbjni</artifactId>
+ </exclusion>
+ <exclusion>
+ <groupId>com.esotericsoftware</groupId>
+ <artifactId>kryo-shaded</artifactId>
+ </exclusion>
+ </exclusions>
+ </dependency>
  <dependency>
  <groupId>org.weakref</groupId>
  <artifactId>jmxutils</artifactId>

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

@@ -15,6 +15,7 @@
 package io.trino.plugin.hudi;
 
 import com.google.common.collect.ImmutableList;
+import io.airlift.log.Logger;
 import io.trino.parquet.Field;
 import io.trino.parquet.ParquetCorruptionException;
 import io.trino.parquet.ParquetDataSource;
@@ -67,7 +68,6 @@
 import java.util.Optional;
 import java.util.TimeZone;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static io.trino.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext;
 import static io.trino.parquet.ParquetTypeUtils.getColumnIO;
 import static io.trino.parquet.ParquetTypeUtils.getDescriptors;
@@ -94,6 +94,8 @@
 public class HudiPageSourceProvider
  implements ConnectorPageSourceProvider
 {
+ private static final Logger log = Logger.get(HudiPageSourceProvider.class);
+
  private final HdfsEnvironment hdfsEnvironment;
  private final FileFormatDataSourceStats fileFormatDataSourceStats;
  private final ParquetReaderOptions parquetReaderOptions;
@@ -165,6 +167,8 @@ private static ConnectorPageSource createParquetPageSource(
  ConnectorIdentity identity,
  List<HivePartitionKey> partitionKeys)
  {
+ log.debug(">>> Creating Parquet Page Source with partition keys: " + partitionKeys);
+
  ParquetDataSource dataSource = null;
  // TODO: Reuse some elements of ParquetPageSourceFactory and extract the try block to a new HudiParquetReader class.
  try {
@@ -231,9 +235,10 @@ && predicateMatches(parquetPredicate, block, dataSource, descriptorsByPath, parq
  .collect(toUnmodifiableList()))
  .orElse(columns);
 
- for (HiveColumnHandle column : baseColumns) {
+ // TODO: add a check for patition column type
+ /*for (HiveColumnHandle column : baseColumns) {
  checkArgument(column == PARQUET_ROW_INDEX_COLUMN || column.getColumnType() == REGULAR, "column type must be REGULAR: %s", column);
- }
+ }*/
 
  ImmutableList.Builder<Type> trinoTypes = ImmutableList.builder();
  ImmutableList.Builder<Optional<Field>> internalFields = ImmutableList.builder();

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

@@ -70,6 +70,7 @@ public boolean isRemotelyAccessible()
  return true;
  }
 
+ @JsonProperty
  @Override
  public List<HostAddress> getAddresses()
  {

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

@@ -14,11 +14,17 @@
 
 package io.trino.plugin.hudi;
 
+import com.google.common.collect.ImmutableList;
+import io.airlift.log.Logger;
 import io.trino.plugin.base.classloader.ClassLoaderSafeConnectorSplitSource;
 import io.trino.plugin.hive.HdfsEnvironment;
+import io.trino.plugin.hive.HivePartitionKey;
 import io.trino.plugin.hive.authentication.HiveIdentity;
+import io.trino.plugin.hive.metastore.Column;
 import io.trino.plugin.hive.metastore.HiveMetastore;
+import io.trino.plugin.hive.metastore.Partition;
 import io.trino.plugin.hive.metastore.Table;
+import io.trino.plugin.hive.util.HiveUtil;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorSplitManager;
 import io.trino.spi.connector.ConnectorSplitSource;
@@ -28,16 +34,38 @@
 import io.trino.spi.connector.DynamicFilter;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.connector.TableNotFoundException;
+import io.trino.spi.predicate.TupleDomain;
 import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.mapred.InputFormat;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.hadoop.HoodieParquetInputFormat;
 
 import javax.inject.Inject;
 
+import java.io.IOException;
+import java.util.List;
+import java.util.Optional;
+import java.util.Properties;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static io.trino.plugin.hive.util.ConfigurationUtils.toJobConf;
+import static io.trino.plugin.hive.util.HiveUtil.getPartitionKeys;
+import static io.trino.plugin.hudi.HudiSessionProperties.isHudiMetadataEnabled;
+import static io.trino.plugin.hudi.HudiUtil.getMetaClient;
+import static io.trino.plugin.hudi.HudiUtil.getPartitionSchema;
+import static io.trino.plugin.hudi.HudiUtil.isHudiParquetInputFormat;
 import static java.util.Objects.requireNonNull;
+import static org.apache.hadoop.hive.metastore.api.hive_metastoreConstants.META_TABLE_LOCATION;
+import static org.apache.hudi.common.table.timeline.TimelineUtils.getPartitionsWritten;
 
 public class HudiSplitManager
  implements ConnectorSplitManager
 {
+ private static final Logger log = Logger.get(HudiSplitManager.class);
+
  private final HudiTransactionManager transactionManager;
  private final HdfsEnvironment hdfsEnvironment;
 
@@ -57,13 +85,78 @@ public ConnectorSplitSource getSplits(
  DynamicFilter dynamicFilter,
  Constraint constraint)
  {
+ log.debug(" >>>> Getting Splits <<<< ");
+ HiveIdentity identity = new HiveIdentity(session);
  HudiTableHandle hudiTable = (HudiTableHandle) tableHandle;
  SchemaTableName tableName = hudiTable.getSchemaTableName();
  HiveMetastore metastore = transactionManager.get(transaction).getMetastore();
- Table table = metastore.getTable(new HiveIdentity(session), tableName.getSchemaName(), tableName.getTableName())
+ Table table = metastore.getTable(identity, tableName.getSchemaName(), tableName.getTableName())
  .orElseThrow(() -> new TableNotFoundException(tableName));
- Configuration conf = hdfsEnvironment.getConfiguration(new HdfsEnvironment.HdfsContext(session), new Path(table.getStorage().getLocation()));
- HudiSplitSource splitSource = new HudiSplitSource(hudiTable, conf);
+ HdfsEnvironment.HdfsContext context = new HdfsEnvironment.HdfsContext(session);
+ FileSystem fs = null;
+ try {
+ fs = hdfsEnvironment.getFileSystem(context, new Path(table.getStorage().getLocation()));
+ }
+ catch (IOException e) {
+ e.printStackTrace();
+ }
+ Configuration conf = hdfsEnvironment.getConfiguration(context, new Path(table.getStorage().getLocation()));
+ HoodieTableMetaClient metaClient = hudiTable.getMetaClient().orElseGet(() -> getMetaClient(conf, hudiTable.getBasePath()));
+ List<String> partitionValues = getPartitionsWritten(metaClient.getActiveTimeline());
+ log.debug("Fetched partitions from Hudi: " + partitionValues);
+ hudiTable.getPartitions().ifPresent(p -> p.forEach(p1 -> log.debug("Partitions from TableHandle: " + p1)));
+
+ List<String> columnNames = table.getPartitionColumns().stream()
+ .map(Column::getName)
+ .collect(toImmutableList());
+ log.debug("Column Names: " + columnNames);
+ HudiSplitSource splitSource;
+ String tablePath = table.getStorage().getLocation();
+ Optional<FileStatus[]> fileStatuses = Optional.empty();
+ if (!columnNames.isEmpty()) {
+ List<List<String>> partitionNames = metastore.getPartitionNamesByFilter(identity, tableName.getSchemaName(), tableName.getTableName(), columnNames, TupleDomain.all())
+ .orElseThrow(() -> new TableNotFoundException(hudiTable.getSchemaTableName()))
+ .stream()
+ .map(HiveUtil::toPartitionValues)
+ .collect(toImmutableList());
+ log.debug("Partition Names: " + partitionNames);
+
+ Optional<Partition> partition = metastore.getPartition(identity, table, partitionNames.get(0));
+
+ log.debug("Fetched partitions from Metastore: " + partition.get());
+ Properties schema = getPartitionSchema(table, partition);
+ String dataDir = schema.getProperty(META_TABLE_LOCATION);
+ log.debug("Partition schema: " + schema);
+
+ List<HivePartitionKey> partitionKeys = getPartitionKeys(table, partition);
+ partitionKeys.forEach(p -> log.warn("Fetched partitions from HiveUtil: " + p));
+
+ InputFormat inputFormat = HiveUtil.getInputFormat(conf, schema, false);
+ log.debug(">>> Check for inputFormat: " + isHudiParquetInputFormat(inputFormat));
+
+ try {
+ if (isHudiParquetInputFormat(inputFormat)) {
+ fileStatuses = Optional.of(((HoodieParquetInputFormat) inputFormat).listStatus(toJobConf(conf)));
+ }
+ if (fileStatuses.isPresent()) {
+ log.debug(">>> Total Files: " + fileStatuses.get().length);
+ if (fileStatuses.get().length == 0 && fs != null) {
+ fileStatuses = Optional.of(fs.listStatus(new Path(dataDir)));
+ log.debug(">>> Total Files: " + fileStatuses.get().length);
+ }
+ }
+ log.debug(">>> Total Splits: " + inputFormat.getSplits(toJobConf(conf), 0).length);
+ }
+ catch (IOException e) {
+ e.printStackTrace();
+ }
+ splitSource = new HudiSplitSource(hudiTable, conf, partitionKeys, isHudiMetadataEnabled(session), fileStatuses, tablePath, dataDir);
+ }
+ else {
+ // no partitions, so data dir is same as table path
+ splitSource = new HudiSplitSource(hudiTable, conf, ImmutableList.of(), isHudiMetadataEnabled(session), fileStatuses, tablePath, tablePath);
+ }
+
  return new ClassLoaderSafeConnectorSplitSource(splitSource, Thread.currentThread().getContextClassLoader());
  }
 }