(hudi) branch master updated: [HUDI-8098] With custom keygen, only read timestamp partition columns from the file (#11895)

[yi...@apache.org]

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

yihua pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git


The following commit(s) were added to refs/heads/master by this push:
     new be080a87ef4 [HUDI-8098] With custom keygen, only read timestamp partition columns from the file (#11895)
be080a87ef4 is described below

commit be080a87ef4132af94c6ec5adbd3cfd0a1adcc60
Author: Jon Vexler <jb...@gmail.com>
AuthorDate: Sun Sep 8 15:15:34 2024 -0400

    [HUDI-8098] With custom keygen, only read timestamp partition columns from the file (#11895)
    
    Co-authored-by: Jonathan Vexler <=>
---
 .../org/apache/hudi/keygen/CustomAvroKeyGenerator.java | 12 ++++++++++++
 .../apache/hudi/HoodieHadoopFsRelationFactory.scala    | 18 ++++++++++++------
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/keygen/CustomAvroKeyGenerator.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/keygen/CustomAvroKeyGenerator.java
index 319f32f0116..1aaee836e7e 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/keygen/CustomAvroKeyGenerator.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/keygen/CustomAvroKeyGenerator.java
@@ -106,6 +106,18 @@ public class CustomAvroKeyGenerator extends BaseKeyGenerator {
     }
   }
 
+  public static List<String> getTimestampFields(List<String> partitionPathFields) {
+    if (partitionPathFields.size() == 1 && partitionPathFields.get(0).isEmpty()) {
+      return Collections.emptyList(); // Corresponds to no partition case
+    } else {
+      return partitionPathFields.stream()
+          .map(CustomAvroKeyGenerator::getPartitionFieldAndKeyType)
+          .filter(fieldAndKeyType -> fieldAndKeyType.getRight().equals(PartitionKeyType.TIMESTAMP))
+          .map(Pair::getLeft)
+          .collect(Collectors.toList());
+    }
+  }
+
   public static Pair<String, PartitionKeyType> getPartitionFieldAndKeyType(String field) {
     String[] fieldWithType = field.split(BaseKeyGenerator.CUSTOM_KEY_GENERATOR_SPLIT_REGEX);
     if (fieldWithType.length != 2) {
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieHadoopFsRelationFactory.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieHadoopFsRelationFactory.scala
index c134b23251f..e6197856b41 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieHadoopFsRelationFactory.scala
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieHadoopFsRelationFactory.scala
@@ -78,17 +78,23 @@ abstract class HoodieBaseHadoopFsRelationFactory(val sqlContext: SQLContext,
   protected lazy val basePath: StoragePath = metaClient.getBasePath
   protected lazy val partitionColumns: Array[String] = tableConfig.getPartitionFields.orElse(Array.empty)
 
-  private lazy val keygenTypeHasVariablePartitionCols = !isNullOrEmpty(tableConfig.getKeyGeneratorClassName) &&
+  private lazy val keygenTypeHasVariablePartitionCols = isTimestampKeygen || isCustomKeygen
+
+  private lazy val isTimestampKeygen = !isNullOrEmpty(tableConfig.getKeyGeneratorClassName) &&
     (tableConfig.getKeyGeneratorClassName.equals(classOf[TimestampBasedKeyGenerator].getName) ||
-      tableConfig.getKeyGeneratorClassName.equals(classOf[TimestampBasedAvroKeyGenerator].getName) ||
-      tableConfig.getKeyGeneratorClassName.equals(classOf[CustomKeyGenerator].getName) ||
-      tableConfig.getKeyGeneratorClassName.equals(classOf[CustomAvroKeyGenerator].getName))
+    tableConfig.getKeyGeneratorClassName.equals(classOf[TimestampBasedAvroKeyGenerator].getName))
+
+  private lazy val isCustomKeygen = !isNullOrEmpty(tableConfig.getKeyGeneratorClassName) &&
+    (tableConfig.getKeyGeneratorClassName.equals(classOf[CustomKeyGenerator].getName) ||
+    tableConfig.getKeyGeneratorClassName.equals(classOf[CustomAvroKeyGenerator].getName))
 
   protected lazy val partitionColumnsToRead: Seq[String] = if (shouldExtractPartitionValuesFromPartitionPath || !keygenTypeHasVariablePartitionCols) {
     Seq.empty
-  } else {
-    //TODO: [HUDI-8098] filter for timestamp keygen columns when using custom keygen
+  } else if (isTimestampKeygen) {
     tableConfig.getPartitionFields.orElse(Array.empty).toSeq
+  } else {
+    //it's custom keygen
+    CustomAvroKeyGenerator.getTimestampFields(HoodieTableConfig.getPartitionFieldsForKeyGenerator(tableConfig).orElse(java.util.Collections.emptyList[String]())).asScala.toSeq
   }
 
   protected lazy val (tableAvroSchema: Schema, internalSchemaOpt: Option[InternalSchema]) = {