Repository: apache/hudi
Pull request #13060: [HUDI-8990] Partition bucket index supports query pruning based on bucket id
Author: login: zhangyue19921010
Labels: name: size:L
description: PR with lines of changes in (300, 1000]
Created: 2025-03-31T13:31:55Z | Merged: 2025-04-05T05:44:09Z

### Change Logs

Followed PR https://github.com/apache/hudi/pull/13017
Support Flink && Spark query adopt partition bucket index based buckId pruning

### Impact

Flink && Spark query adopt partition bucket index based buckId pruning

### Risk level (write none, low medium or high below)

low

### Documentation Update

_Describe any necessary documentation update if there is any new feature, config, or user-facing change. If not, put "none"._

- _The config description must be updated if new configs are added or the default value of the configs are changed_
- _Any new feature or user-facing change requires updating the Hudi website. Please create a Jira ticket, attach the
  ticket number here and follow the [instruction](https://hudi.apache.org/contribute/developer-setup#website) to make
  changes to the website._

### Contributor's checklist

- [ ] Read through [contributor's guide](https://hudi.apache.org/contribute/how-to-contribute)
- [ ] Change Logs and Impact were stated clearly
- [ ] Adequate tests were added if applicable
- [ ] CI passed



Diff:
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/partition/NumBucketsFunction.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/partition/NumBucketsFunction.java
index 0587304d76027..f9deed6fe7236 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/partition/NumBucketsFunction.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/partition/NumBucketsFunction.java
@@ -18,6 +18,8 @@
 
 package org.apache.hudi.index.bucket.partition;
 
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
 import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.config.HoodieWriteConfig;
 
@@ -38,7 +40,7 @@ public class NumBucketsFunction implements Serializable {
   /**
    * The default number of buckets to use when partition-specific buckets are not configured.
    */
-  private final int defaultBucketNumber;
+  private final int defaultBucketsNum;
 
   /**
    * Flag indicating whether partition-level bucket index is enabled.
@@ -49,30 +51,34 @@ public class NumBucketsFunction implements Serializable {
    * Calculator for partition-specific bucket numbers.
    */
   private final PartitionBucketIndexCalculator calculator;
-  private final String expressions;
-  private final String ruleType;
 
   /**
    * Creates a NumBucketsFunction with the given configuration.
    *
-   * @param config The Flink configuration containing bucket index settings.
+   * @param expressions       The expressions declared as the rules with {@code ruleType}
+   * @param ruleType          The rule type
+   * @param defaultBucketsNum The default buckets number
    */
-  public NumBucketsFunction(String expressions, String ruleType, int defaultBucketNumber) {
-    this.defaultBucketNumber = defaultBucketNumber;
-    this.expressions = expressions;
-    this.ruleType = ruleType;
+  public NumBucketsFunction(String expressions, String ruleType, int defaultBucketsNum) {
+    this.defaultBucketsNum = defaultBucketsNum;
     this.isPartitionLevelBucketIndexEnabled = StringUtils.nonEmpty(expressions);
     if (isPartitionLevelBucketIndexEnabled) {
       this.calculator = PartitionBucketIndexCalculator.getInstance(
-          expressions, ruleType, defaultBucketNumber);
+          expressions, ruleType, defaultBucketsNum);
       LOG.info("Initialized partition-level bucket index with expressions: {}, rule: {}, default bucket number: {}",
-          expressions, ruleType, defaultBucketNumber);
+          expressions, ruleType, defaultBucketsNum);
     } else {
       this.calculator = null;
-      LOG.info("Using fixed bucket number: {}", defaultBucketNumber);
+      LOG.info("Using fixed bucket number: {}", defaultBucketsNum);
     }
   }
 
+  public NumBucketsFunction(int defaultBucketsNum) {
+    this.defaultBucketsNum = defaultBucketsNum;
+    this.isPartitionLevelBucketIndexEnabled = false;
+    this.calculator = null;
+  }
+
   public static NumBucketsFunction fromWriteConfig(HoodieWriteConfig writeConfig) {
     String expression = writeConfig.getBucketIndexPartitionExpression();
     String ruleType = writeConfig.getBucketIndexPartitionRuleType();
@@ -80,6 +86,15 @@ public static NumBucketsFunction fromWriteConfig(HoodieWriteConfig writeConfig)
     return new NumBucketsFunction(expression, ruleType, numBuckets);
   }
 
+  public static NumBucketsFunction fromMetaClient(HoodieTableMetaClient metaClient, int defaultBucketNumber) {
+    if (PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(metaClient.getStorageConf(), metaClient.getBasePath().toString())) {
+      PartitionBucketIndexHashingConfig hashingConfig = PartitionBucketIndexHashingConfig.loadingLatestHashingConfig(metaClient);
+      return new NumBucketsFunction(hashingConfig.getExpressions(), hashingConfig.getRule(), hashingConfig.getDefaultBucketNumber());
+    } else {
+      return new NumBucketsFunction(defaultBucketNumber);
+    }
+  }
+
   /**
    * Gets the number of buckets for the given partition path.
    *
@@ -90,7 +105,7 @@ public int getNumBuckets(String partitionPath) {
     if (isPartitionLevelBucketIndexEnabled && calculator != null) {
       return calculator.computeNumBuckets(partitionPath);
     }
-    return defaultBucketNumber;
+    return defaultBucketsNum;
   }
 
   /**
@@ -99,7 +114,7 @@ public int getNumBuckets(String partitionPath) {
    * @return The default bucket number.
    */
   public int getDefaultBucketNumber() {
-    return defaultBucketNumber;
+    return defaultBucketsNum;
   }
 
   /**
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java b/hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
index c186e4f049e3e..4467607ef8b58 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
@@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
-      return allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1);
+      if (instant.isPresent()) {
+        Option<String> res = getHashingConfigInstantToLoadBeforeOrOn(allCommittedHashingConfig, instant.get());
+        // fall back to look up archived hashing config instant before return empty
+        return res.isPresent() ? res : getHashingConfigInstantToLoadBeforeOrOn(getArchiveHashingConfigInstants(metaClient), instant.get());
+      } else {
+        return Option.of(allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1));
+      }
     } catch (Exception e) {
       throw new HoodieException("Failed to get hashing config instant to load.", e);
     }
   }
 
+  private static Option<String> getHashingConfigInstantToLoadBeforeOrOn(List<String> hashingConfigInstants, String instant) {
+    List<String> res = hashingConfigInstants.stream().filter(hashingConfigInstant -> {
+      return hashingConfigInstant.compareTo(instant) <= 0;
+    }).collect(Collectors.toList());
+    return res.isEmpty() ? Option.empty() : Option.of(res.get(res.size() - 1));
+  }
+
   public static PartitionBucketIndexHashingConfig loadingLatestHashingConfig(HoodieTableMetaClient metaClient) {
-    String instantToLoad = getLatestHashingConfigInstantToLoad(metaClient);
-    Option<PartitionBucketIndexHashingConfig> latestHashingConfig = loadHashingConfig(metaClient.getStorage(), getHashingConfigPath(metaClient.getBasePath().toString(), instantToLoad));
+    Option<String> instantToLoad = getHashingConfigInstantToLoad(metaClient, Option.empty());
+    ValidationUtils.checkArgument(instantToLoad.isPresent(), "Can not load latest hashing config " + instantToLoad);
+    Option<PartitionBucketIndexHashingConfig> latestHashingConfig = loadHashingConfig(metaClient.getStorage(), getHashingConfigPath(metaClient.getBasePath().toString(), instantToLoad.get()));
     ValidationUtils.checkArgument(latestHashingConfig.isPresent(), "Can not load latest hashing config " + instantToLoad);
 
     return latestHashingConfig.get();
   }
 
+  public static Option<PartitionBucketIndexHashingConfig> loadingLatestHashingConfigBeforeOrOn(HoodieTableMetaClient metaClient, String instant) {
+    Option<String> hashingConfigInstantToLoad = getHashingConfigInstantToLoad(metaClient, Option.of(instant));
+    if (hashingConfigInstantToLoad.isPresent()) {
+      Option<PartitionBucketIndexHashingConfig> latestHashingConfig = loadHashingConfig(metaClient.getStorage(),
+          getHashingConfigPath(metaClient.getBasePath().toString(), hashingConfigInstantToLoad.get()));
+      ValidationUtils.checkArgument(latestHashingConfig.isPresent(), "Can not load hashing config " + hashingConfigInstantToLoad + " based on " + instant);
+      return latestHashingConfig;
+    } else {
+      return Option.empty();
+    }
+  }
+
   /**
    * Archive hashing config.
    */
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
index a8ff13fe62121..351137c3d51b7 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
@@ -27,7 +27,6 @@
 import org.apache.hudi.index.bucket.BucketIdentifier;
 import org.apache.hudi.source.prune.ColumnStatsProbe;
 import org.apache.hudi.source.prune.PartitionPruners;
-import org.apache.hudi.source.prune.PrimaryKeyPruners;
 import org.apache.hudi.source.stats.FileStatsIndex;
 import org.apache.hudi.storage.StoragePath;
 import org.apache.hudi.storage.StoragePathInfo;
@@ -50,6 +49,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -66,11 +66,11 @@ public class FileIndex implements Serializable {
   private final boolean tableExists;
   private final HoodieMetadataConfig metadataConfig;
   private final org.apache.hadoop.conf.Configuration hadoopConf;
-  private final PartitionPruners.PartitionPruner partitionPruner; // for partition pruning
-  private final ColumnStatsProbe colStatsProbe;                   // for probing column stats
-  private final int dataBucket;                                   // for bucket pruning
-  private List<String> partitionPaths;                            // cache of partition paths
-  private final FileStatsIndex fileStatsIndex;                    // for data skipping
+  private final PartitionPruners.PartitionPruner partitionPruner;  // for partition pruning
+  private final ColumnStatsProbe colStatsProbe;                    // for probing column stats
+  private final Function<String, Integer> partitionBucketIdFunc;   // for bucket pruning
+  private List<String> partitionPaths;                             // cache of partition paths
+  private final FileStatsIndex fileStatsIndex;                     // for data skipping
 
   private FileIndex(
       StoragePath path,
@@ -78,15 +78,15 @@ private FileIndex(
       RowType rowType,
       ColumnStatsProbe colStatsProbe,
       PartitionPruners.PartitionPruner partitionPruner,
-      int dataBucket) {
+      Function<String, Integer> partitionBucketIdFunc) {
     this.path = path;
     this.hadoopConf = HadoopConfigurations.getHadoopConf(conf);
     this.tableExists = StreamerUtil.tableExists(path.toString(), hadoopConf);
     this.metadataConfig = StreamerUtil.metadataConfig(conf);
     this.colStatsProbe = isDataSkippingFeasible(conf.get(FlinkOptions.READ_DATA_SKIPPING_ENABLED)) ? colStatsProbe : null;
     this.partitionPruner = partitionPruner;
-    this.dataBucket = dataBucket;
     this.fileStatsIndex = new FileStatsIndex(path.toString(), rowType, metadataConfig);
+    this.partitionBucketIdFunc = partitionBucketIdFunc;
   }
 
   /**
@@ -157,26 +157,30 @@ public List<StoragePathInfo> getFilesInPartitions() {
     if (partitions.length < 1) {
       return Collections.emptyList();
     }
-    List<StoragePathInfo> allFiles = FSUtils.getFilesInPartitions(
-            new HoodieFlinkEngineContext(hadoopConf),
-            new HoodieHadoopStorage(path, HadoopFSUtils.getStorageConf(hadoopConf)), metadataConfig, path.toString(), partitions)
-        .values().stream()
-        .flatMap(Collection::stream)
-        .collect(Collectors.toList());
-
-    if (allFiles.isEmpty()) {
+    Map<String, List<StoragePathInfo>> filesInPartitions = FSUtils.getFilesInPartitions(
+        new HoodieFlinkEngineContext(hadoopConf),
+        new HoodieHadoopStorage(path, HadoopFSUtils.getStorageConf(hadoopConf)), metadataConfig, path.toString(), partitions);
+    int totalFilesNum = filesInPartitions.values().stream().mapToInt(List::size).sum();
+    if (totalFilesNum < 1) {
       // returns early for empty table.
-      return allFiles;
+      return Collections.emptyList();
     }
 
+    List<StoragePathInfo> allFiles;
     // bucket pruning
-    if (this.dataBucket >= 0) {
-      String bucketIdStr = BucketIdentifier.bucketIdStr(this.dataBucket);
-      List<StoragePathInfo> filesAfterBucketPruning = allFiles.stream()
-          .filter(fileInfo -> fileInfo.getPath().getName().contains(bucketIdStr))
-          .collect(Collectors.toList());
-      logPruningMsg(allFiles.size(), filesAfterBucketPruning.size(), "bucket pruning");
-      allFiles = filesAfterBucketPruning;
+    if (this.partitionBucketIdFunc != null) {
+      allFiles = filesInPartitions.entrySet().stream().flatMap(entry -> {
+        String bucketIdStr = BucketIdentifier.bucketIdStr(partitionBucketIdFunc.apply(entry.getKey()));
+        return entry.getValue().stream().filter(fileInfo -> fileInfo.getPath().getName().contains(bucketIdStr));
+      }).collect(Collectors.toList());
+      logPruningMsg(totalFilesNum, allFiles.size(), "bucket pruning");
+    } else {
+      allFiles = filesInPartitions.values().stream().flatMap(Collection::stream).collect(Collectors.toList());
+    }
+
+    if (allFiles.isEmpty()) {
+      // returns early for empty table.
+      return allFiles;
     }
 
     // data skipping
@@ -286,7 +290,7 @@ public static class Builder {
     private RowType rowType;
     private ColumnStatsProbe columnStatsProbe;
     private PartitionPruners.PartitionPruner partitionPruner;
-    private int dataBucket = PrimaryKeyPruners.BUCKET_ID_NO_PRUNING;
+    private Function<String, Integer> partitionBucketIdFunc;
 
     private Builder() {
     }
@@ -316,14 +320,14 @@ public Builder partitionPruner(PartitionPruners.PartitionPruner partitionPruner)
       return this;
     }
 
-    public Builder dataBucket(int dataBucket) {
-      this.dataBucket = dataBucket;
+    public Builder partitionBucketIdFunc(Function<String, Integer> partitionBucketIdFunc) {
+      this.partitionBucketIdFunc = partitionBucketIdFunc;
       return this;
     }
 
     public FileIndex build() {
       return new FileIndex(Objects.requireNonNull(path), Objects.requireNonNull(conf), Objects.requireNonNull(rowType),
-          columnStatsProbe, partitionPruner, dataBucket);
+          columnStatsProbe, partitionPruner, partitionBucketIdFunc);
     }
   }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PartitionBucketIdFunc.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PartitionBucketIdFunc.java
new file mode 100644
index 0000000000000..26fa6122bf9cf
--- /dev/null
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PartitionBucketIdFunc.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.hudi.source.prune;
+
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.index.bucket.partition.NumBucketsFunction;
+
+import javax.annotation.Nullable;
+
+import java.io.Serializable;
+import java.util.function.Function;
+
+/**
+ * A function to fetch the bucket id for buckets pruning.
+ */
+public class PartitionBucketIdFunc implements Function<String, Integer>, Serializable {
+  private final NumBucketsFunction numBucketsFunction;   // partition -> numBuckets
+  private final Function<Integer, Integer> bucketIdFunc; // numBuckets -> bucketId
+
+  private PartitionBucketIdFunc(NumBucketsFunction numBucketsFunction, Function<Integer, Integer> bucketIdFunc) {
+    this.numBucketsFunction = numBucketsFunction;
+    this.bucketIdFunc = bucketIdFunc;
+  }
+
+  /**
+   * Creates the {@code PartitionBucketIdFunc} instance, will return null if the bucket pruning does not take effect.
+   *
+   * @param bucketIdFuncOpt   The bucket id function option
+   * @param metaClient        The meta client
+   * @param defaultBucketsNum The default num buckets
+   */
+  @Nullable
+  public static PartitionBucketIdFunc create(Option<Function<Integer, Integer>> bucketIdFuncOpt, HoodieTableMetaClient metaClient, int defaultBucketsNum) {
+    if (bucketIdFuncOpt.isPresent()) {
+      return new PartitionBucketIdFunc(NumBucketsFunction.fromMetaClient(metaClient, defaultBucketsNum), bucketIdFuncOpt.get());
+    } else {
+      return null;
+    }
+  }
+
+  /**
+   * Returns the bucket id for the given partition path.
+   */
+  @Override
+  public Integer apply(String partitionPath) {
+    return this.bucketIdFunc.apply(this.numBucketsFunction.getNumBuckets(partitionPath));
+  }
+}
\ No newline at end of file
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
index 05b69bdcf7656..d8bd17e0ecfbf 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
@@ -35,6 +35,7 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -43,9 +44,7 @@
 public class PrimaryKeyPruners {
   private static final Logger LOG = LoggerFactory.getLogger(PrimaryKeyPruners.class);
 
-  public static final int BUCKET_ID_NO_PRUNING = -1;
-
-  public static int getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
+  public static Function<Integer, Integer> getBucketIdFunc(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
     List<String> pkFields = Arrays.asList(conf.getString(FlinkOptions.RECORD_KEY_FIELD).split(","));
     // step1: resolve the hash key values
     final boolean logicalTimestamp = OptionsResolver.isConsistentLogicalTimestampEnabled(conf);
@@ -60,7 +59,7 @@ public static int getBucketId(List<ResolvedExpression> hashKeyFilters, Configura
         .map(Pair::getValue)
         .collect(Collectors.toList());
     // step2: generate bucket id
-    return BucketIdentifier.getBucketId(values, conf.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS));
+    return (numBuckets) -> BucketIdentifier.getBucketId(values, numBuckets);
   }
 
   private static Pair<FieldReferenceExpression, ValueLiteralExpression> castChildAs(List<Expression> children) {
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
index 5c9cde00ba90e..deb12a3d18640 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
@@ -47,6 +47,7 @@
 import org.apache.hudi.source.StreamReadMonitoringFunction;
 import org.apache.hudi.source.StreamReadOperator;
 import org.apache.hudi.source.prune.ColumnStatsProbe;
+import org.apache.hudi.source.prune.PartitionBucketIdFunc;
 import org.apache.hudi.source.prune.PartitionPruners;
 import org.apache.hudi.source.prune.PrimaryKeyPruners;
 import org.apache.hudi.source.rebalance.partitioner.StreamReadAppendPartitioner;
@@ -116,6 +117,7 @@
 import java.util.Set;
 import java.util.StringJoiner;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -156,7 +158,7 @@ public class HoodieTableSource implements
   private List<Predicate> predicates;
   private ColumnStatsProbe columnStatsProbe;
   private PartitionPruners.PartitionPruner partitionPruner;
-  private int dataBucket;
+  private Option<Function<Integer, Integer>> dataBucketFunc; // numBuckets -> bucketId
   private transient FileIndex fileIndex;
 
   public HoodieTableSource(
@@ -165,7 +167,8 @@ public HoodieTableSource(
       List<String> partitionKeys,
       String defaultPartName,
       Configuration conf) {
-    this(schema, path, partitionKeys, defaultPartName, conf, null, null, null, PrimaryKeyPruners.BUCKET_ID_NO_PRUNING, null, null, null, null);
+    this(schema, path, partitionKeys, defaultPartName, conf, null, null, null,
+        Option.empty(), null, null, null, null);
   }
 
   public HoodieTableSource(
@@ -177,7 +180,7 @@ public HoodieTableSource(
       @Nullable List<Predicate> predicates,
       @Nullable ColumnStatsProbe columnStatsProbe,
       @Nullable PartitionPruners.PartitionPruner partitionPruner,
-      int dataBucket,
+      Option<Function<Integer, Integer>> dataBucketFunc,
       @Nullable int[] requiredPos,
       @Nullable Long limit,
       @Nullable HoodieTableMetaClient metaClient,
@@ -191,7 +194,7 @@ public HoodieTableSource(
     this.predicates = Optional.ofNullable(predicates).orElse(Collections.emptyList());
     this.columnStatsProbe = columnStatsProbe;
     this.partitionPruner = partitionPruner;
-    this.dataBucket = dataBucket;
+    this.dataBucketFunc = dataBucketFunc;
     this.requiredPos = Optional.ofNullable(requiredPos).orElseGet(() -> IntStream.range(0, this.tableRowType.getFieldCount()).toArray());
     this.limit = Optional.ofNullable(limit).orElse(NO_LIMIT_CONSTANT);
     this.hadoopConf = new HadoopStorageConfiguration(HadoopConfigurations.getHadoopConf(conf));
@@ -266,7 +269,7 @@ public ChangelogMode getChangelogMode() {
   @Override
   public DynamicTableSource copy() {
     return new HoodieTableSource(schema, path, partitionKeys, defaultPartName,
-        conf, predicates, columnStatsProbe, partitionPruner, dataBucket, requiredPos, limit, metaClient, internalSchemaManager);
+        conf, predicates, columnStatsProbe, partitionPruner, dataBucketFunc, requiredPos, limit, metaClient, internalSchemaManager);
   }
 
   @Override
@@ -281,7 +284,7 @@ public Result applyFilters(List<ResolvedExpression> filters) {
     this.predicates = ExpressionPredicates.fromExpression(splitFilters.f0);
     this.columnStatsProbe = ColumnStatsProbe.newInstance(splitFilters.f0);
     this.partitionPruner = createPartitionPruner(splitFilters.f1, columnStatsProbe);
-    this.dataBucket = getDataBucket(splitFilters.f0);
+    this.dataBucketFunc = getDataBucketFunc(splitFilters.f0);
     // refuse all the filters now
     return SupportsFilterPushDown.Result.of(new ArrayList<>(splitFilters.f1), new ArrayList<>(filters));
   }
@@ -365,16 +368,16 @@ private PartitionPruners.PartitionPruner createPartitionPruner(List<ResolvedExpr
         .build();
   }
 
-  private int getDataBucket(List<ResolvedExpression> dataFilters) {
+  private Option<Function<Integer, Integer>> getDataBucketFunc(List<ResolvedExpression> dataFilters) {
     if (!OptionsResolver.isBucketIndexType(conf) || dataFilters.isEmpty()) {
-      return PrimaryKeyPruners.BUCKET_ID_NO_PRUNING;
+      return Option.empty();
     }
     Set<String> indexKeyFields = Arrays.stream(OptionsResolver.getIndexKeys(conf)).collect(Collectors.toSet());
     List<ResolvedExpression> indexKeyFilters = dataFilters.stream().filter(expr -> ExpressionUtils.isEqualsLitExpr(expr, indexKeyFields)).collect(Collectors.toList());
     if (!ExpressionUtils.isFilteringByAllFields(indexKeyFilters, indexKeyFields)) {
-      return PrimaryKeyPruners.BUCKET_ID_NO_PRUNING;
+      return Option.empty();
     }
-    return PrimaryKeyPruners.getBucketId(indexKeyFilters, conf);
+    return Option.of(PrimaryKeyPruners.getBucketIdFunc(indexKeyFilters, conf));
   }
 
   private List<MergeOnReadInputSplit> buildInputSplits() {
@@ -611,7 +614,8 @@ private FileIndex getOrBuildFileIndex() {
           .rowType(this.tableRowType)
           .columnStatsProbe(this.columnStatsProbe)
           .partitionPruner(this.partitionPruner)
-          .dataBucket(this.dataBucket)
+          .partitionBucketIdFunc(PartitionBucketIdFunc.create(this.dataBucketFunc,
+              this.metaClient, conf.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS)))
           .build();
     }
     return this.fileIndex;
@@ -695,7 +699,7 @@ public ColumnStatsProbe getColumnStatsProbe() {
   }
 
   @VisibleForTesting
-  public int getDataBucket() {
-    return dataBucket;
+  public Option<Function<Integer, Integer>> getDataBucketFunc() {
+    return dataBucketFunc;
   }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/ITTestHoodieDataSource.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/ITTestHoodieDataSource.java
index 522da0e039aa3..560987711adcf 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/ITTestHoodieDataSource.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/ITTestHoodieDataSource.java
@@ -1321,6 +1321,48 @@ void testBulkInsertWithPartitionBucketIndex(String operationType, String tableTy
     assertEquals(expected.stream().sorted().collect(Collectors.toList()), actual.stream().sorted().collect(Collectors.toList()));
   }
 
+  @Test
+  void tesQueryWithPartitionBucketIndexPruning() {
+    String operationType = "upsert";
+    String tableType = "MERGE_ON_READ";
+    TableEnvironment tableEnv = batchTableEnv;
+    // csv source
+    String csvSourceDDL = TestConfigurations.getCsvSourceDDL("csv_source", "test_source_5.data");
+    tableEnv.executeSql(csvSourceDDL);
+    String catalogName = "hudi_" + operationType;
+    String hudiCatalogDDL = catalog(catalogName)
+        .catalogPath(tempFile.getAbsolutePath())
+        .end();
+
+    tableEnv.executeSql(hudiCatalogDDL);
+    String dbName = "hudi";
+    tableEnv.executeSql("create database " + catalogName + "." + dbName);
+    String basePath = tempFile.getAbsolutePath() + "/hudi/hoodie_sink";
+
+    String hoodieTableDDL = sql(catalogName + ".hudi.hoodie_sink")
+        .option(FlinkOptions.PATH, basePath)
+        .option(FlinkOptions.TABLE_TYPE, tableType)
+        .option(FlinkOptions.OPERATION, operationType)
+        .option(FlinkOptions.WRITE_BULK_INSERT_SHUFFLE_INPUT, true)
+        .option(FlinkOptions.INDEX_TYPE, "BUCKET")
+        .option(FlinkOptions.HIVE_STYLE_PARTITIONING, "true")
+        .option(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, "1")
+        .option(FlinkOptions.BUCKET_INDEX_PARTITION_RULE, "regex")
+        .option(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS, "partition=(par1|par2),2")
+        .end();
+    tableEnv.executeSql(hoodieTableDDL);
+
+    String insertInto = "insert into " + catalogName + ".hudi.hoodie_sink select * from csv_source";
+    execInsertSql(tableEnv, insertInto);
+
+    List<Row> result1 = execSelectSql(tableEnv, "select * from " + catalogName + ".hudi.hoodie_sink");
+    assertRowsEquals(result1, TestData.DATA_SET_SOURCE_INSERT);
+    // apply filters which will prune based on partition level bucket index
+    List<Row> result2 = CollectionUtil.iterableToList(
+        () -> tableEnv.sqlQuery("select * from " + catalogName + ".hudi.hoodie_sink where uuid = 'id5'").execute().collect());
+    assertRowsEquals(result2, "[+I[id5, Sophia, 18, 1970-01-01T00:00:05, par3]]");
+  }
+
   @Test
   void testBulkInsertWithSortByRecordKey() {
     TableEnvironment tableEnv = batchTableEnv;
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/TestHoodieTableSource.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/TestHoodieTableSource.java
index ff0c849e0c146..931a576612e0f 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/TestHoodieTableSource.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/TestHoodieTableSource.java
@@ -19,15 +19,16 @@
 package org.apache.hudi.table;
 
 import org.apache.hudi.common.config.HoodieMetadataConfig;
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
 import org.apache.hudi.common.table.HoodieTableMetaClient;
 import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.source.ExpressionPredicates;
 import org.apache.hudi.source.prune.ColumnStatsProbe;
-import org.apache.hudi.source.prune.PrimaryKeyPruners;
 import org.apache.hudi.storage.StoragePath;
 import org.apache.hudi.storage.StoragePathInfo;
 import org.apache.hudi.table.format.mor.MergeOnReadInputFormat;
 import org.apache.hudi.util.SerializableSchema;
+import org.apache.hudi.util.StreamerUtil;
 import org.apache.hudi.utils.TestConfigurations;
 import org.apache.hudi.utils.TestData;
 
@@ -76,6 +77,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /**
  * Test cases for HoodieTableSource.
@@ -199,7 +201,9 @@ void testBucketPruning(boolean hiveStylePartitioning) throws Exception {
     tableSource1.applyFilters(Collections.singletonList(
         createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1")));
 
-    assertThat(tableSource1.getDataBucket(), is(1));
+    int numBuckets = (int)FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.defaultValue();
+
+    assertThat(tableSource1.getDataBucketFunc().get().apply(numBuckets), is(1));
     List<StoragePathInfo> fileList = tableSource1.getReadFiles();
     assertThat("Files should be pruned by bucket id 1", fileList.size(), CoreMatchers.is(2));
 
@@ -214,7 +218,7 @@ void testBucketPruning(boolean hiveStylePartitioning) throws Exception {
     tableSource2.applyFilters(Arrays.asList(
         createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1"),
         createLitEquivalenceExpr("name", 1, DataTypes.STRING().notNull(), "Danny")));
-    assertThat(tableSource2.getDataBucket(), is(3));
+    assertThat(tableSource2.getDataBucketFunc().get().apply(numBuckets), is(3));
     List<StoragePathInfo> fileList2 = tableSource2.getReadFiles();
     assertThat("Files should be pruned by bucket id 3", fileList2.size(), CoreMatchers.is(3));
 
@@ -223,7 +227,7 @@ void testBucketPruning(boolean hiveStylePartitioning) throws Exception {
     tableSource2.applyFilters(Arrays.asList(
         createLitEquivalenceExpr("name", 1, DataTypes.STRING().notNull(), "Danny"),
         createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1")));
-    assertThat(tableSource2.getDataBucket(), is(3));
+    assertThat(tableSource2.getDataBucketFunc().get().apply(numBuckets), is(3));
     assertThat("Files should be pruned by bucket id 3", tableSource2.getReadFiles().size(),
         CoreMatchers.is(3));
 
@@ -238,7 +242,7 @@ void testBucketPruning(boolean hiveStylePartitioning) throws Exception {
     tableSource3.applyFilters(Collections.singletonList(
         createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1")));
 
-    assertThat(tableSource3.getDataBucket(), is(PrimaryKeyPruners.BUCKET_ID_NO_PRUNING));
+    assertTrue(tableSource3.getDataBucketFunc().isEmpty());
     List<StoragePathInfo> fileList3 = tableSource3.getReadFiles();
     assertThat("Partial pk filtering does not prune any files", fileList3.size(),
         CoreMatchers.is(7));
@@ -253,7 +257,7 @@ void testBucketPruning(boolean hiveStylePartitioning) throws Exception {
         createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1"),
         createLitEquivalenceExpr("name", 1, DataTypes.STRING().notNull(), "Danny")));
 
-    assertThat(tableSource4.getDataBucket(), is(1));
+    assertThat(tableSource4.getDataBucketFunc().get().apply(numBuckets), is(1));
     List<StoragePathInfo> fileList4 = tableSource4.getReadFiles();
     assertThat("Files should be pruned by bucket id 1", fileList4.size(), CoreMatchers.is(2));
   }
@@ -269,6 +273,7 @@ void testBucketPruningSpecialKeyDataType(boolean logicalTimestamp) throws Except
     conf1.setString(FlinkOptions.PRECOMBINE_FIELD, f1);
     conf1.removeConfig(FlinkOptions.PARTITION_PATH_FIELD);
     conf1.setBoolean(KEYGENERATOR_CONSISTENT_LOGICAL_TIMESTAMP_ENABLED.key(), logicalTimestamp);
+    int numBuckets = (int)FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.defaultValue();
 
     // test timestamp filtering
     TestData.writeDataAsBatch(TestData.DATA_SET_INSERT_HOODIE_KEY_SPECIAL_DATA_TYPE, conf1);
@@ -277,7 +282,7 @@ void testBucketPruningSpecialKeyDataType(boolean logicalTimestamp) throws Except
         createLitEquivalenceExpr(f1, 0, DataTypes.TIMESTAMP(3).notNull(),
             LocalDateTime.ofInstant(Instant.ofEpochMilli(1), ZoneId.of("UTC")))));
 
-    assertThat(tableSource1.getDataBucket(), is(logicalTimestamp ? 1 : 0));
+    assertThat(tableSource1.getDataBucketFunc().get().apply(numBuckets), is(logicalTimestamp ? 1 : 0));
     List<StoragePathInfo> fileList = tableSource1.getReadFiles();
     assertThat("Files should be pruned", fileList.size(), CoreMatchers.is(1));
 
@@ -293,7 +298,7 @@ void testBucketPruningSpecialKeyDataType(boolean logicalTimestamp) throws Except
     tableSource2.applyFilters(Collections.singletonList(
         createLitEquivalenceExpr(f2, 1, DataTypes.DATE().notNull(), LocalDate.ofEpochDay(1))));
 
-    assertThat(tableSource2.getDataBucket(), is(1));
+    assertThat(tableSource2.getDataBucketFunc().get().apply(numBuckets), is(1));
     List<StoragePathInfo> fileList2 = tableSource2.getReadFiles();
     assertThat("Files should be pruned", fileList2.size(), CoreMatchers.is(1));
 
@@ -310,7 +315,7 @@ void testBucketPruningSpecialKeyDataType(boolean logicalTimestamp) throws Except
         createLitEquivalenceExpr(f3, 1, DataTypes.DECIMAL(3, 2).notNull(),
             new BigDecimal("1.11"))));
 
-    assertThat(tableSource3.getDataBucket(), is(0));
+    assertThat(tableSource3.getDataBucketFunc().get().apply(numBuckets), is(0));
     List<StoragePathInfo> fileList3 = tableSource3.getReadFiles();
     assertThat("Files should be pruned", fileList3.size(), CoreMatchers.is(1));
   }
@@ -343,6 +348,126 @@ void testFilterPushDownWithParquetPredicates() {
     assertEquals(ExpressionPredicates.fromExpression(expectedFilters).toString(), actualPredicates);
   }
 
+  /**
+   * test single primary key filtering
+   * @throws Exception
+   */
+  @Test
+  void testPartitionBucketPruningWithSinglePK() throws Exception {
+    String tablePath1 = new Path(tempFile.getAbsolutePath(), "tbl1").toString();
+    int bucketNumber = 10000;
+    String expression = "par1|par2|par3|par4,4";
+    String rule = "regex";
+    Configuration conf1 = TestConfigurations.getDefaultConf(tablePath1);
+    conf1.setString(FlinkOptions.INDEX_TYPE, "BUCKET");
+    conf1.set(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, bucketNumber);
+    conf1.set(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS, expression);
+    conf1.set(FlinkOptions.BUCKET_INDEX_PARTITION_RULE, rule);
+
+    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf1);
+    PartitionBucketIndexHashingConfig.saveHashingConfig(metaClient, expression, rule, bucketNumber, null);
+
+    // test single primary key filtering
+    TestData.writeDataAsBatch(TestData.DATA_SET_INSERT, conf1);
+    HoodieTableSource tableSource1 = createHoodieTableSource(conf1);
+    tableSource1.applyFilters(Collections.singletonList(
+        createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1")));
+
+    assertThat(tableSource1.getDataBucketFunc().get().apply(4), is(1));
+    List<StoragePathInfo> fileList = tableSource1.getReadFiles();
+    assertThat("Files should be pruned by bucket id 1", fileList.size(), CoreMatchers.is(2));
+  }
+
+  /**
+   * test multiple primary keys filtering
+   * @throws Exception
+   */
+  @Test
+  void testPartitionBucketPruningWithMultiPK() throws Exception {
+    String tablePath = new Path(tempFile.getAbsolutePath(), "tbl1").toString();
+    int bucketNumber = 10000;
+    String expression = "par1|par2|par3|par4,4";
+    String rule = "regex";
+    Configuration conf = TestConfigurations.getDefaultConf(tablePath);
+    conf.setString(FlinkOptions.INDEX_TYPE, "BUCKET");
+    conf.set(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, bucketNumber);
+    conf.set(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS, expression);
+    conf.set(FlinkOptions.BUCKET_INDEX_PARTITION_RULE, rule);
+    conf.setString(FlinkOptions.RECORD_KEY_FIELD, "uuid,name");
+    conf.setString(FlinkOptions.KEYGEN_TYPE, "COMPLEX");
+
+    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
+    PartitionBucketIndexHashingConfig.saveHashingConfig(metaClient, expression, rule, bucketNumber, null);
+    TestData.writeDataAsBatch(TestData.DATA_SET_INSERT, conf);
+    HoodieTableSource tableSource = createHoodieTableSource(conf);
+    tableSource.applyFilters(Arrays.asList(
+        createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1"),
+        createLitEquivalenceExpr("name", 1, DataTypes.STRING().notNull(), "Danny")));
+    assertThat(tableSource.getDataBucketFunc().get().apply(4), is(3));
+    List<StoragePathInfo> fileList = tableSource.getReadFiles();
+    assertThat("Files should be pruned by bucket id 3", fileList.size(), CoreMatchers.is(3));
+  }
+
+  /**
+   * test partial primary keys filtering
+   * @throws Exception
+   */
+  @Test
+  void testPartialPartitionBucketPruningWithMultiPK() throws Exception {
+    String tablePath = new Path(tempFile.getAbsolutePath(), "tbl1").toString();
+    int bucketNumber = 10000;
+    String expression = "par1|par2|par3|par4,4";
+    String rule = "regex";
+    Configuration conf = TestConfigurations.getDefaultConf(tablePath);
+    conf.setString(FlinkOptions.INDEX_TYPE, "BUCKET");
+    conf.set(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, bucketNumber);
+    conf.set(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS, expression);
+    conf.set(FlinkOptions.BUCKET_INDEX_PARTITION_RULE, rule);
+    conf.setString(FlinkOptions.RECORD_KEY_FIELD, "uuid,name");
+    conf.setString(FlinkOptions.KEYGEN_TYPE, "COMPLEX");
+
+    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
+    PartitionBucketIndexHashingConfig.saveHashingConfig(metaClient, expression, rule, bucketNumber, null);
+    TestData.writeDataAsBatch(TestData.DATA_SET_INSERT, conf);
+    HoodieTableSource tableSource = createHoodieTableSource(conf);
+    tableSource.applyFilters(Collections.singletonList(
+        createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1")));
+
+    assertTrue(tableSource.getDataBucketFunc().isEmpty());
+    List<StoragePathInfo> fileList = tableSource.getReadFiles();
+    assertThat("Partial pk filtering does not prune any files", fileList.size(),
+        CoreMatchers.is(7));
+  }
+
+  /**
+   * test single primary keys filtering together with non-primary key predicate
+   * @throws Exception
+   */
+  @Test
+  void testPartitionBucketPruningWithMixedFilter() throws Exception {
+    String tablePath = new Path(tempFile.getAbsolutePath(), "tbl1").toString();
+    int bucketNumber = 10000;
+    String expression = "par1|par2|par3|par4,4";
+    String rule = "regex";
+    Configuration conf = TestConfigurations.getDefaultConf(tablePath);
+    conf.setString(FlinkOptions.INDEX_TYPE, "BUCKET");
+    conf.set(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, bucketNumber);
+    conf.set(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS, expression);
+    conf.set(FlinkOptions.BUCKET_INDEX_PARTITION_RULE, rule);
+
+    HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
+    PartitionBucketIndexHashingConfig.saveHashingConfig(metaClient, expression, rule, bucketNumber, null);
+    TestData.writeDataAsBatch(TestData.DATA_SET_INSERT, conf);
+    HoodieTableSource tableSource = createHoodieTableSource(conf);
+    tableSource.applyFilters(Arrays.asList(
+        createLitEquivalenceExpr("uuid", 0, DataTypes.STRING().notNull(), "id1"),
+        createLitEquivalenceExpr("name", 1, DataTypes.STRING().notNull(), "Danny")));
+
+    assertThat(tableSource.getDataBucketFunc().get().apply(4), is(1));
+    List<StoragePathInfo> fileList = tableSource.getReadFiles();
+    assertThat("Files should be pruned by bucket id 1", fileList.size(), CoreMatchers.is(2));
+  }
+
   // -------------------------------------------------------------------------
   //  Utilities
   // -------------------------------------------------------------------------
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/BucketIndexSupport.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/BucketIndexSupport.scala
index f91fa8b0a42c8..068ec48214232 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/BucketIndexSupport.scala
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/BucketIndexSupport.scala
@@ -42,7 +42,7 @@ import scala.collection.{mutable, JavaConverters}
 class BucketIndexSupport(spark: SparkSession,
                          metadataConfig: HoodieMetadataConfig,
                          metaClient: HoodieTableMetaClient)
-  extends SparkBaseIndexSupport (spark, metadataConfig, metaClient){
+  extends SparkBaseIndexSupport (spark, metadataConfig, metaClient) {
 
   private val log = LoggerFactory.getLogger(getClass)
 
@@ -101,8 +101,8 @@ class BucketIndexSupport(spark: SparkSession,
     candidateFiles.toSet
   }
 
-  def filterQueriesWithBucketHashField(queryFilters: Seq[Expression]): Option[BitSet] = {
-    val bucketNumber = metadataConfig.getIntOrDefault(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS)
+  def filterQueriesWithBucketHashField(queryFilters: Seq[Expression],
+                                       bucketNumber: Int = metadataConfig.getIntOrDefault(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS)): Option[BitSet] = {
     if (indexBucketHashFieldsOpt.isEmpty || queryFilters.isEmpty) {
       None
     } else {
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieFileIndex.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieFileIndex.scala
index a8c0dd87ec034..2b0779d70ac33 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieFileIndex.scala
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieFileIndex.scala
@@ -27,6 +27,7 @@ import org.apache.hudi.common.model.{FileSlice, HoodieBaseFile, HoodieLogFile}
 import org.apache.hudi.common.table.{HoodieTableConfig, HoodieTableMetaClient}
 import org.apache.hudi.common.util.StringUtils
 import org.apache.hudi.exception.HoodieException
+import org.apache.hudi.index.bucket.partition.PartitionBucketIndexUtils
 import org.apache.hudi.keygen.{TimestampBasedAvroKeyGenerator, TimestampBasedKeyGenerator}
 import org.apache.hudi.storage.{StoragePath, StoragePathInfo}
 import org.apache.hudi.util.JFunction
@@ -106,13 +107,18 @@ case class HoodieFileIndex(spark: SparkSession,
 
   /**
    * NOTE: [[indicesSupport]] is a transient state, since it's only relevant while logical plan
-   *       is handled by the Spark's driver
-   *       The order of elements is important as in this order indices will be applied
-   *       during `lookupCandidateFilesInMetadataTable`
+   * is handled by the Spark's driver
+   * The order of elements is important as in this order indices will be applied
+   * during `lookupCandidateFilesInMetadataTable`
    */
   @transient private lazy val indicesSupport: List[SparkBaseIndexSupport] = List(
     new RecordLevelIndexSupport(spark, metadataConfig, metaClient),
-    new BucketIndexSupport(spark, metadataConfig, metaClient),
+    if (PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(metaClient.getStorageConf, metaClient.getBasePath.toString)) {
+      new PartitionBucketIndexSupport(spark, metadataConfig, metaClient,
+        options.get(DataSourceReadOptions.TIME_TRAVEL_AS_OF_INSTANT.key).map(HoodieSqlCommonUtils.formatQueryInstant))
+    } else {
+      new BucketIndexSupport(spark, metadataConfig, metaClient)
+    },
     new SecondaryIndexSupport(spark, metadataConfig, metaClient),
     new ExpressionIndexSupport(spark, schema, metadataConfig, metaClient),
     new BloomFiltersIndexSupport(spark, metadataConfig, metaClient),
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/PartitionBucketIndexSupport.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/PartitionBucketIndexSupport.scala
new file mode 100644
index 0000000000000..fd8a008bc36a4
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/PartitionBucketIndexSupport.scala
@@ -0,0 +1,86 @@
+/*
+ * 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.hudi
+
+import org.apache.hudi.BaseHoodieTableFileIndex.PartitionPath
+import org.apache.hudi.common.config.HoodieMetadataConfig
+import org.apache.hudi.common.model.{FileSlice, PartitionBucketIndexHashingConfig}
+import org.apache.hudi.common.table.HoodieTableMetaClient
+import org.apache.hudi.index.bucket.partition.PartitionBucketIndexCalculator
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.expressions.Expression
+import org.slf4j.LoggerFactory
+
+class PartitionBucketIndexSupport(spark: SparkSession,
+                         metadataConfig: HoodieMetadataConfig,
+                         metaClient: HoodieTableMetaClient,
+                         specifiedQueryInstant: Option[String] = Option.empty)
+  extends BucketIndexSupport(spark, metadataConfig, metaClient) {
+
+  private val log = LoggerFactory.getLogger(getClass)
+  private def initCalc(): Option[PartitionBucketIndexCalculator] = {
+    if (specifiedQueryInstant.isDefined) {
+      val hashingConfigOption = PartitionBucketIndexHashingConfig.loadingLatestHashingConfigBeforeOrOn(metaClient, specifiedQueryInstant.get)
+      if (hashingConfigOption.isPresent) {
+        Option.apply(PartitionBucketIndexCalculator.getInstance(hashingConfigOption.get().getExpressions,
+          hashingConfigOption.get().getRule,
+          hashingConfigOption.get().getDefaultBucketNumber))
+      } else {
+        // could be empty, if user upgrade to partition level bucket index after specifiedQueryInstant
+        Option.empty
+      }
+    } else {
+      // specifiedQueryInstant is null or empty, so load latest hashing config directly
+      val hashingConfig = PartitionBucketIndexHashingConfig.loadingLatestHashingConfig(metaClient)
+      Option.apply(PartitionBucketIndexCalculator.getInstance(hashingConfig.getExpressions,
+        hashingConfig.getRule, hashingConfig.getDefaultBucketNumber))
+    }
+  }
+
+
+  private val calc: Option[PartitionBucketIndexCalculator] = initCalc()
+
+  override def computeCandidateFileNames(fileIndex: HoodieFileIndex,
+                                         queryFilters: Seq[Expression],
+                                         queryReferencedColumns: Seq[String],
+                                         prunedPartitionsAndFileSlices: Seq[(Option[BaseHoodieTableFileIndex.PartitionPath], Seq[FileSlice])],
+                                         shouldPushDownFilesFilter: Boolean): Option[Set[String]] = {
+    if (calc.isEmpty) {
+      // fall back to BucketIndexSupport
+      super.computeCandidateFileNames(fileIndex, queryFilters, queryReferencedColumns, prunedPartitionsAndFileSlices, shouldPushDownFilesFilter)
+    } else {
+      Option.apply(prunedPartitionsAndFileSlices.flatMap(v => {
+        val partitionPathOption = v._1
+        val fileSlices = v._2
+        val numBuckets = calc.get.computeNumBuckets(partitionPathOption.getOrElse(new PartitionPath("", Array())).path)
+        val bucketIdsBitMapByFilter = filterQueriesWithBucketHashField(queryFilters, numBuckets)
+        if (bucketIdsBitMapByFilter.isDefined && bucketIdsBitMapByFilter.get.cardinality() > 0) {
+          val allFilesName = getPrunedPartitionsAndFileNames(fileIndex, Seq((partitionPathOption, fileSlices)))._2
+          getCandidateFiles(allFilesName, bucketIdsBitMapByFilter.get)
+        } else {
+          Seq()
+        }}).toSet)
+    }
+  }
+}
+
+object PartitionBucketIndexSupport {
+  val INDEX_NAME = "BUCKET"
+}
+
diff --git a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/functional/TestPartitionBucketIndexSupport.scala b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/functional/TestPartitionBucketIndexSupport.scala
new file mode 100644
index 0000000000000..029e887d601e1
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/hudi/functional/TestPartitionBucketIndexSupport.scala
@@ -0,0 +1,201 @@
+/*
+ * 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.hudi.functional
+
+import org.apache.hudi.{BaseHoodieTableFileIndex, HoodieFileIndex, PartitionBucketIndexSupport}
+import org.apache.hudi.common.config.{HoodieMetadataConfig, TypedProperties}
+import org.apache.hudi.common.fs.FSUtils
+import org.apache.hudi.common.model.{FileSlice, HoodieBaseFile, PartitionBucketIndexHashingConfig}
+import org.apache.hudi.common.table.HoodieTableConfig
+import org.apache.hudi.config.HoodieIndexConfig
+import org.apache.hudi.index.bucket.BucketIdentifier
+import org.apache.hudi.keygen.NonpartitionedKeyGenerator
+import org.apache.hudi.keygen.constant.KeyGeneratorOptions
+import org.apache.hudi.storage.{StoragePath, StoragePathInfo}
+
+import org.apache.avro.generic.GenericData
+import org.apache.spark.sql.HoodieCatalystExpressionUtils
+import org.apache.spark.sql.catalyst.encoders.DummyExpressionHolder
+import org.junit.jupiter.api.{BeforeEach, Tag, Test}
+import org.mockito.Mockito
+
+@Tag("functional")
+class TestPartitionBucketIndexSupport extends TestBucketIndexSupport {
+
+  private val DEFAULT_RULE = "regex"
+  private val EXPRESSION_BUCKET_NUMBER = 19
+  private val DEFAULT_EXPRESSIONS = "\\d{4}\\-(06\\-(01|17|18)|11\\-(01|10|11))," + EXPRESSION_BUCKET_NUMBER
+  private val DEFAULT_BUCKET_NUMBER = 10
+  private val DEFAULT_PARTITION_PATH = Array("2025-06-17", "2025-06-18")
+  private var fileIndex: HoodieFileIndex = null
+  @BeforeEach
+  override def setUp(): Unit = {
+    super.setUp()
+    PartitionBucketIndexHashingConfig.saveHashingConfig(metaClient, DEFAULT_EXPRESSIONS, DEFAULT_RULE, DEFAULT_BUCKET_NUMBER, null)
+    fileIndex = Mockito.mock(classOf[HoodieFileIndex])
+  }
+
+  @Test
+  override def testSingleHashFieldsExpression: Unit = {
+    val configProperties = new TypedProperties()
+    configProperties.setProperty(HoodieIndexConfig.BUCKET_INDEX_HASH_FIELD.key, "A")
+    configProperties.setProperty(HoodieTableConfig.RECORDKEY_FIELDS.key, "A")
+    configProperties.setProperty(KeyGeneratorOptions.RECORDKEY_FIELD_NAME.key, "A")
+
+    configProperties.setProperty(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS.key, String.valueOf(DEFAULT_BUCKET_NUMBER))
+    metaClient.getTableConfig.setValue(HoodieTableConfig.CREATE_SCHEMA.key(), avroSchemaStr)
+    val metadataConfig = HoodieMetadataConfig.newBuilder
+      .fromProperties(configProperties)
+      .enable(configProperties.getBoolean(HoodieMetadataConfig.ENABLE.key, true)).build()
+    val bucketIndexSupport: PartitionBucketIndexSupport = new PartitionBucketIndexSupport(spark, metadataConfig, metaClient)
+
+
+    // init
+    val testKeyGenerator = new NonpartitionedKeyGenerator(configProperties)
+    var record = new GenericData.Record(avroSchema)
+    record.put("A", "1")
+    val bucket1Id4 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "2")
+    val bucket2Id5 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "3")
+    val bucket3Id6 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "4")
+    val bucket4Id7 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "5")
+    val bucket5Id8 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A", EXPRESSION_BUCKET_NUMBER)
+    assert(bucket1Id4 == 4 && bucket2Id5 == 5 && bucket3Id6 == 6 && bucket4Id7 == 7 && bucket5Id8 == 8)
+
+    // fileIdStr
+    val token = FSUtils.makeWriteToken(1, 0, 1)
+    val bucket1Id4FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket1Id4) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket2Id5FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket2Id5) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket3Id6FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket3Id6) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket4Id7FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket4Id7) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket5Id8FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket5Id8) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+
+
+    val allFileNames = Set.apply(bucket1Id4FileName, bucket2Id5FileName, bucket3Id6FileName, bucket4Id7FileName, bucket5Id8FileName)
+    var equalTo = "A = 3"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket3Id6FileName), allFileNames)
+    equalTo = "A = 3 And A = 4 and B = '6'"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.empty, allFileNames)
+    equalTo = "A = 5 And B = 'abc'"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket5Id8FileName), allFileNames)
+    equalTo = "A = C and A = 1"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket1Id4FileName), allFileNames)
+    equalTo = "A = 5 Or A = 2"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket5Id8FileName, bucket2Id5FileName), allFileNames)
+    equalTo = "A = 5 Or A = 2 Or A = 8"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket5Id8FileName, bucket2Id5FileName), allFileNames)
+    equalTo = "A = 5 Or (A = 2 and B = 'abc')"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket5Id8FileName, bucket2Id5FileName), allFileNames)
+    equalTo = "A = 5 And (A = 2 Or B = 'abc')"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket5Id8FileName), allFileNames)
+    equalTo = "A = 5 And (A = 2 Or B = 'abc')"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket5Id8FileName), allFileNames)
+  }
+
+  @Test
+  override def testMultipleHashFieldsExpress(): Unit = {
+    val configProperties = new TypedProperties()
+    configProperties.setProperty(HoodieIndexConfig.BUCKET_INDEX_HASH_FIELD.key, "A,B")
+    configProperties.setProperty(HoodieTableConfig.RECORDKEY_FIELDS.key, "A,B")
+    configProperties.setProperty(KeyGeneratorOptions.RECORDKEY_FIELD_NAME.key, "A,B")
+    configProperties.setProperty(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS.key, String.valueOf(DEFAULT_BUCKET_NUMBER))
+    metaClient.getTableConfig.setValue(HoodieTableConfig.CREATE_SCHEMA.key(), avroSchemaStr)
+    val metadataConfig = HoodieMetadataConfig.newBuilder
+      .fromProperties(configProperties)
+      .enable(configProperties.getBoolean(HoodieMetadataConfig.ENABLE.key, true)).build()
+    val bucketIndexSupport = new PartitionBucketIndexSupport(spark, metadataConfig, metaClient)
+    val keyGenerator = bucketIndexSupport.getKeyGenerator
+    assert(keyGenerator.isInstanceOf[NonpartitionedKeyGenerator])
+
+    // init
+    val testKeyGenerator = new NonpartitionedKeyGenerator(configProperties)
+    var record = new GenericData.Record(avroSchema)
+    record.put("A", "1")
+    record.put("B", "2")
+    val bucket1Id4 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A,B", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "2")
+    record.put("B", "3")
+    val bucket2Id5 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A,B", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "3")
+    record.put("B", "4")
+    val bucket3Id6 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A,B", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "4")
+    record.put("B", "5")
+    val bucket4Id7 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A,B", EXPRESSION_BUCKET_NUMBER)
+    record = new GenericData.Record(avroSchema)
+    record.put("A", "5")
+    record.put("B", "6")
+    val bucket5Id8 = BucketIdentifier.getBucketId(testKeyGenerator.getKey(record).getRecordKey, "A,B", EXPRESSION_BUCKET_NUMBER)
+    assert(bucket1Id4 == 3 && bucket2Id5 == 16 && bucket3Id6 == 10 && bucket4Id7 == 4 && bucket5Id8 == 17)
+
+    // fileIdStr
+    val token = FSUtils.makeWriteToken(1, 0, 1)
+    val bucket1Id4FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket1Id4) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket2Id5FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket2Id5) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket3Id6FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket3Id6) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket4Id7FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket4Id7) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+    val bucket5Id8FileName = FSUtils.makeBaseFileName("00000000000000000", token, BucketIdentifier.newBucketFileIdPrefix(bucket5Id8) + "-0", HoodieTableConfig.BASE_FILE_FORMAT.defaultValue.getFileExtension)
+
+    val allFileNames = Set.apply(bucket1Id4FileName, bucket2Id5FileName, bucket3Id6FileName, bucket4Id7FileName, bucket5Id8FileName)
+
+    var equalTo = "A = 2 and B = '3'"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket2Id5FileName), allFileNames)
+
+    equalTo = "A = 4 and B = '5'"
+    exprFilePathAnswerCheck(bucketIndexSupport, equalTo, Set.apply(bucket4Id7FileName), allFileNames)
+  }
+
+  def exprFilePathAnswerCheck(bucketIndexSupport: PartitionBucketIndexSupport, exprRaw: String, expectResult: Set[String],
+                              allFileStatus: Set[String]): Unit = {
+    val resolveExpr = HoodieCatalystExpressionUtils.resolveExpr(spark, exprRaw, structSchema)
+    val optimizerPlan = spark.sessionState.optimizer.execute(DummyExpressionHolder(Seq(resolveExpr)))
+    val optimizerExpr = optimizerPlan.asInstanceOf[DummyExpressionHolder].exprs.head
+
+    // split input files into different partitions
+    val partitionPath1 = DEFAULT_PARTITION_PATH(0)
+    val allFileSlices1: Seq[FileSlice] = allFileStatus.slice(0, 3).map(fileName => {
+      val slice = new FileSlice(partitionPath1, "00000000000000000", FSUtils.getFileId(fileName))
+      slice.setBaseFile(new HoodieBaseFile(new StoragePathInfo(new StoragePath(fileName), 0L, false, 0, 0, 0)))
+      slice
+    }).toSeq
+
+    val partitionPath2 = DEFAULT_PARTITION_PATH(1)
+    val allFileSlices2: Seq[FileSlice] = allFileStatus.slice(3, 5).map(fileName => {
+      val slice = new FileSlice(partitionPath1, "00000000000000000", FSUtils.getFileId(fileName))
+      slice.setBaseFile(new HoodieBaseFile(new StoragePathInfo(new StoragePath(fileName), 0L, false, 0, 0, 0)))
+      slice
+    }).toSeq
+
+    val input = Seq((Option.apply(new BaseHoodieTableFileIndex.PartitionPath(partitionPath1, Array())), allFileSlices1),
+      (Option.apply(new BaseHoodieTableFileIndex.PartitionPath(partitionPath2, Array())), allFileSlices2))
+    val candidate = bucketIndexSupport.computeCandidateFileNames(fileIndex, splitConjunctivePredicates(optimizerExpr),
+      Seq(), input, false)
+
+    assert(candidate.get.equals(expectResult))
+  }
+}
diff --git a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTableWithPartitionBucketIndex.scala b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTableWithPartitionBucketIndex.scala
index 4b05dd456fb4a..bca3d04461731 100644
--- a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTableWithPartitionBucketIndex.scala
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTableWithPartitionBucketIndex.scala
@@ -313,6 +313,101 @@ class TestInsertTableWithPartitionBucketIndex extends HoodieSparkSqlTestBase {
     }
   }
 
+  test("Test BucketID Pruning With Partition Bucket Index") {
+    withSQLConf(
+      "hoodie.datasource.write.operation" -> "upsert") {
+      withTempDir { tmp =>
+        withTable(generateTableName) { tableName =>
+          val tablePath = tmp.getCanonicalPath + "/" + tableName
+          val defaultBucketNumber = 2
+          // Create a partitioned table
+          spark.sql(
+            s"""
+               |create table $tableName (
+               |  id int,
+               |  dt string,
+               |  name string,
+               |  price double,
+               |  ts long
+               |) using hudi
+               | tblproperties (
+               | primaryKey = 'id',
+               | type = 'mor',
+               | preCombineField = 'ts',
+               | hoodie.index.type = 'BUCKET',
+               | hoodie.bucket.index.hash.field = 'id',
+               | hoodie.bucket.index.num.buckets = $defaultBucketNumber)
+               | partitioned by (dt)
+               | location '$tablePath'
+               | """.stripMargin)
+
+          // Note: Do not write the field alias, the partition field must be placed last.
+          // do commit 1
+          spark.sql(
+            s"""
+               | insert into $tableName values
+               | (1, 'a1,1', 1.0, 1, "2021-01-05"),
+               | (11, 'a11,11', 11.0, 11, "2021-01-05"),
+               | (111, 'a111', 111.0, 111, "2021-01-05"),
+               | (1111, 'a1111', 1111.0, 1111, "2021-01-05")
+               | """.stripMargin)
+
+          // upgrade to partition level bucket index and rescale dt=2021-01-05 from 1 to 2
+          // do bucket rescale commit 2
+          val expressions = "dt=(2021\\-01\\-05|2021\\-01\\-07),4"
+          val rule = "regex"
+          spark.sql(s"call partition_bucket_index_manager(table => '$tableName', overwrite => '$expressions', rule => '$rule', bucketNumber => $defaultBucketNumber)")
+
+          // do commit 3 update id = 1111
+          spark.sql(
+            s"""
+               | insert into $tableName values
+               | (1111, 'a2222,2222', 2222.0, 2222, "2021-01-05")
+               | """.stripMargin)
+
+          // do bucket rescale commit 4
+          val expressions2 = "dt=(2021\\-01\\-05|2021\\-01\\-07),3"
+          val rule2 = "regex"
+          spark.sql(s"call partition_bucket_index_manager(table => '$tableName', overwrite => '$expressions2', rule => '$rule2', bucketNumber => $defaultBucketNumber)")
+
+          // do commit 5 update id = 1111
+          spark.sql(
+            s"""
+               | insert into $tableName values
+               | (1111, 'a3333,3333', 3333.0, 3333, "2021-01-05")
+               | """.stripMargin)
+          val metaClient = createMetaClient(spark, tablePath)
+          val committedInstant = metaClient.getActiveTimeline.filterCompletedInstants().getInstants
+          val size = committedInstant.size()
+
+          val writeCommit5 = committedInstant.get(size - 1)
+          val bucketRescaleCommit4 = committedInstant.get(size - 2)
+          val writeCommit3 = committedInstant.get(size - 3)
+          val bucketRescaleCommit2 = committedInstant.get(size - 4)
+          val writeCommit1 = committedInstant.get(size - 5)
+
+          // normal query with bucket id pruning at latest writeCommit5
+          checkAnswer(s"select id, price, ts, dt from $tableName where id = '1111'")(
+            Seq(1111, 3333.0, 3333, "2021-01-05")
+          )
+          // time travel to writeCommit3
+          checkAnswer(s"select id, price, ts, dt from $tableName timestamp as of '${writeCommit3.requestedTime()}' where id = '1111'")(
+            Seq(1111, 2222.0, 2222, "2021-01-05")
+          )
+          // time travel to writeCommit1
+          checkAnswer(s"select id, price, ts, dt from $tableName timestamp as of '${writeCommit1.requestedTime()}' where id = '1111'")(
+            Seq(1111, 1111.0, 1111, "2021-01-05")
+          )
+
+          // incremental query with bucket id pruning
+          checkAnswer(s"select id, price, ts, dt from hudi_table_changes('$tableName', 'latest_state', '${bucketRescaleCommit2.getCompletionTime}', '${writeCommit3.getCompletionTime}') where id = '1111'")(
+            Seq(1111, 2222.0, 2222, "2021-01-05")
+          )
+        }
+      }
+    }
+  }
+
   test("Test Rollback Hashing Config") {
     withSQLConf(
       "hoodie.datasource.write.operation" -> "upsert",


Reviews:
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T01:11:35Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T01:11:59Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T01:14:29Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T01:16:09Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T01:16:58Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-01T01:18:50Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T15:04:10Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T15:11:07Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T15:12:38Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T01:38:33Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T06:18:49Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T06:19:05Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T06:19:37Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T06:21:36Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:42:39Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:43:39Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:44:03Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:46:02Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:46:49Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:47:02Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:47:15Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T08:47:46Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T15:36:05Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T15:36:15Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T15:36:21Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T15:36:25Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T15:37:47Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T15:38:03Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-03T15:38:37Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T00:33:01Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T13:59:44Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T14:00:34Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T14:03:39Z
user: login: danny0405
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-05T01:23:23Z
user: login: zhangyue19921010
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-05T01:49:35Z
user: login: danny0405
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-05T05:42:15Z

Review comments:
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
-      return allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1);
+      if (instant.isPresent()) {
+        Option<String> res = getHashingConfigInstantToLoadBeforeOrOn(allCommittedHashingConfig, instant.get());
+        // fall back to look up archived hashing config instant before return empty
+        return res.isPresent() ? res : getHashingConfigInstantToLoadBeforeOrOn(getArchiveHashingConfigInstants(metaClient), instant.get());
+      } else {
+        return Option.of(allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1));
+      }
     } catch (Exception e) {
       throw new HoodieException("Failed to get hashing config instant to load.", e);
     }
   }
 
+  private static Option<String> getHashingConfigInstantToLoadBeforeOrOn(List<String> hashingConfigInstants, String instant) {
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: danny0405
body: We should always load the latest bucket config to comply with the latest bucket id mappings. Otherwise the query would fail. We do not ensure snapshot isolation here because the reader read being affected by the writer.
created_at: 2025-04-01T01:11:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,25 +69,30 @@ public class FileIndex implements Serializable {
   private final org.apache.hadoop.conf.Configuration hadoopConf;
   private final PartitionPruners.PartitionPruner partitionPruner; // for partition pruning
   private final ColumnStatsProbe colStatsProbe;                   // for probing column stats
-  private final int dataBucket;                                   // for bucket pruning
+  private final int dataBucketHashing;                                   // for bucket pruning
   private List<String> partitionPaths;                            // cache of partition paths
   private final FileStatsIndex fileStatsIndex;                    // for data skipping
+  private final NumBucketsFunction numBucketsFunction;
 
   private FileIndex(
       StoragePath path,
       Configuration conf,
       RowType rowType,
       ColumnStatsProbe colStatsProbe,
       PartitionPruners.PartitionPruner partitionPruner,
-      int dataBucket) {
+      int dataBucketHashing,
+      NumBucketsFunction numBucketsFunction) {
     this.path = path;
     this.hadoopConf = HadoopConfigurations.getHadoopConf(conf);
     this.tableExists = StreamerUtil.tableExists(path.toString(), hadoopConf);
     this.metadataConfig = StreamerUtil.metadataConfig(conf);
     this.colStatsProbe = isDataSkippingFeasible(conf.get(FlinkOptions.READ_DATA_SKIPPING_ENABLED)) ? colStatsProbe : null;
     this.partitionPruner = partitionPruner;
-    this.dataBucket = dataBucket;
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
user: login: danny0405
body: Can we still use the bucket id here?
created_at: 2025-04-01T01:11:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -45,7 +45,7 @@ public class PrimaryKeyPruners {
 
   public static final int BUCKET_ID_NO_PRUNING = -1;
 
-  public static int getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
+  public static int getBucketFieldHashing(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
user: login: danny0405
body: Can we still return the bucket id? We can add a new param for the function:

```java

  // the input of bucketIdFunc is the partition path
  getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf, Function<String, Integer> bucketIdFunc)
```
created_at: 2025-04-01T01:14:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,196 @@
+/*
+ * 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.hudi.table;
+
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.index.bucket.BucketIdentifier;
+import org.apache.hudi.source.prune.PrimaryKeyPruners;
+import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.storage.StoragePathInfo;
+import org.apache.hudi.util.SerializableSchema;
+import org.apache.hudi.util.StreamerUtil;
+import org.apache.hudi.utils.TestConfigurations;
+import org.apache.hudi.utils.TestData;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.expressions.CallExpression;
+import org.apache.flink.table.expressions.FieldReferenceExpression;
+import org.apache.flink.table.expressions.ResolvedExpression;
+import org.apache.flink.table.expressions.ValueLiteralExpression;
+import org.apache.flink.table.functions.BuiltInFunctionDefinitions;
+import org.apache.flink.table.types.DataType;
+import org.apache.hadoop.fs.Path;
+import org.hamcrest.CoreMatchers;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.core.Is.is;
+
+public class TestPartitionBucketPruning {
path: hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/TestPartitionBucketPruning.java
user: login: danny0405
body: Should we move the tests into `TestHoodieTableSource`, can we at least add a IT test for the query result validation.
created_at: 2025-04-01T01:16:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -103,6 +104,15 @@ case class HoodieFileIndex(spark: SparkSession,
     endCompletionTime = options.get(DataSourceReadOptions.END_COMMIT.key)) with FileIndex {
 
   @transient protected var hasPushedDownPartitionPredicates: Boolean = false
+  private val isPartitionSimpleBucketIndex = PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(spark.sparkContext.hadoopConfiguration,
+    metaClient.getBasePath.toString)
+
+  @transient private lazy val bucketIndexSupport = if (isPartitionSimpleBucketIndex) {
+    val specifiedQueryInstant = options.get(DataSourceReadOptions.TIME_TRAVEL_AS_OF_INSTANT.key).map(HoodieSqlCommonUtils.formatQueryInstant)
+    new PartitionBucketIndexSupport(spark, metadataConfig, metaClient, specifiedQueryInstant)
path: hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieFileIndex.scala
user: login: danny0405
body: always query from the latest hash config because there is no SI for reader/writers.
created_at: 2025-04-01T01:16:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: danny0405
body: Can we assume there is always config files once the partition bucket index is enabled?
created_at: 2025-04-01T01:18:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -45,7 +45,7 @@ public class PrimaryKeyPruners {
 
   public static final int BUCKET_ID_NO_PRUNING = -1;
 
-  public static int getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
+  public static int getBucketFieldHashing(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
user: login: zhangyue19921010
body: Sorry Danny, I didn't get this. Is that possible to get full partition path and use it in this new bucketIdFunc In the original code's call site？
```
  @Override
  public Result applyFilters(List<ResolvedExpression> filters) {
    List<ResolvedExpression> simpleFilters = filterSimpleCallExpression(filters);
    Tuple2<List<ResolvedExpression>, List<ResolvedExpression>> splitFilters = splitExprByPartitionCall(simpleFilters, this.partitionKeys, this.tableRowType);
    this.predicates = ExpressionPredicates.fromExpression(splitFilters.f0);
    this.columnStatsProbe = ColumnStatsProbe.newInstance(splitFilters.f0);
    this.partitionPruner = createPartitionPruner(splitFilters.f1, columnStatsProbe);
    this.dataBucket = getDataBucket(splitFilters.f0);
    // refuse all the filters now
    return SupportsFilterPushDown.Result.of(new ArrayList<>(splitFilters.f1), new ArrayList<>(filters));
  }
```

What is PR did is get and pass hashing value to `getFilesInPartitions`, then compute numBuckets , finally compute the final bucket id `hashing value % numBuckets`

created_at: 2025-04-02T15:04:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: zhangyue19921010
body: Yes we can.
Currently, we have two methods to use partition-level bucket indexing:

1. Enabling during table creation via DDL: When enabled during table creation, this method initializes a 0000000000000.hashing_config file through the catalog. 

2. Upgrading existing table-level bucket indexes via the CALL command: For tables that already use table-level bucket indexing, invoking a CALL command triggers an upgrade process. This generates a replace-commit instant  and initializes a corresponding <replace-commit>.hashing_config file
created_at: 2025-04-02T15:11:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
-      return allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1);
+      if (instant.isPresent()) {
+        Option<String> res = getHashingConfigInstantToLoadBeforeOrOn(allCommittedHashingConfig, instant.get());
+        // fall back to look up archived hashing config instant before return empty
+        return res.isPresent() ? res : getHashingConfigInstantToLoadBeforeOrOn(getArchiveHashingConfigInstants(metaClient), instant.get());
+      } else {
+        return Option.of(allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1));
+      }
     } catch (Exception e) {
       throw new HoodieException("Failed to get hashing config instant to load.", e);
     }
   }
 
+  private static Option<String> getHashingConfigInstantToLoadBeforeOrOn(List<String> hashingConfigInstants, String instant) {
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: zhangyue19921010
body: The method `getHashingConfigInstantToLoadBeforeOrOn` is designed to handle Time Travel scenarios. For example, suppose we have a sequence of commits and replace-commits:

C1, C2, R3 (a bucket rescale operation), C4, R5 (another bucket rescale).

If we perform a Time Travel query targeting commit C4 (i.e., specifiedQueryInstant = C4), Maybe we should load the hashing configuration associated with R3 (the latest bucket rescale operation before or equal to C4).
created_at: 2025-04-02T15:12:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -45,7 +45,7 @@ public class PrimaryKeyPruners {
 
   public static final int BUCKET_ID_NO_PRUNING = -1;
 
-  public static int getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
+  public static int getBucketFieldHashing(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
user: login: danny0405
body: Maybe we evolve this dataBucket to a function `(num_buckets_per_partition) -> (int)bucketId` to make it somehow more flexible.
created_at: 2025-04-03T01:38:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,25 +69,30 @@ public class FileIndex implements Serializable {
   private final org.apache.hadoop.conf.Configuration hadoopConf;
   private final PartitionPruners.PartitionPruner partitionPruner; // for partition pruning
   private final ColumnStatsProbe colStatsProbe;                   // for probing column stats
-  private final int dataBucket;                                   // for bucket pruning
+  private final int dataBucketHashing;                                   // for bucket pruning
   private List<String> partitionPaths;                            // cache of partition paths
   private final FileStatsIndex fileStatsIndex;                    // for data skipping
+  private final NumBucketsFunction numBucketsFunction;
 
   private FileIndex(
       StoragePath path,
       Configuration conf,
       RowType rowType,
       ColumnStatsProbe colStatsProbe,
       PartitionPruners.PartitionPruner partitionPruner,
-      int dataBucket) {
+      int dataBucketHashing,
+      NumBucketsFunction numBucketsFunction) {
     this.path = path;
     this.hadoopConf = HadoopConfigurations.getHadoopConf(conf);
     this.tableExists = StreamerUtil.tableExists(path.toString(), hadoopConf);
     this.metadataConfig = StreamerUtil.metadataConfig(conf);
     this.colStatsProbe = isDataSkippingFeasible(conf.get(FlinkOptions.READ_DATA_SKIPPING_ENABLED)) ? colStatsProbe : null;
     this.partitionPruner = partitionPruner;
-    this.dataBucket = dataBucket;
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
user: login: zhangyue19921010
body: evolve this dataBucket to a function
created_at: 2025-04-03T06:18:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -45,7 +45,7 @@ public class PrimaryKeyPruners {
 
   public static final int BUCKET_ID_NO_PRUNING = -1;
 
-  public static int getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
+  public static int getBucketFieldHashing(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
user: login: zhangyue19921010
body: Sure, changed
created_at: 2025-04-03T06:19:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,196 @@
+/*
+ * 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.hudi.table;
+
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.index.bucket.BucketIdentifier;
+import org.apache.hudi.source.prune.PrimaryKeyPruners;
+import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.storage.StoragePathInfo;
+import org.apache.hudi.util.SerializableSchema;
+import org.apache.hudi.util.StreamerUtil;
+import org.apache.hudi.utils.TestConfigurations;
+import org.apache.hudi.utils.TestData;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.expressions.CallExpression;
+import org.apache.flink.table.expressions.FieldReferenceExpression;
+import org.apache.flink.table.expressions.ResolvedExpression;
+import org.apache.flink.table.expressions.ValueLiteralExpression;
+import org.apache.flink.table.functions.BuiltInFunctionDefinitions;
+import org.apache.flink.table.types.DataType;
+import org.apache.hadoop.fs.Path;
+import org.hamcrest.CoreMatchers;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.File;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.core.Is.is;
+
+public class TestPartitionBucketPruning {
path: hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/TestPartitionBucketPruning.java
user: login: zhangyue19921010
body: change also add a IT named `tesQueryWithPartitionBucketIndexPruning` to do query result validation.
created_at: 2025-04-03T06:19:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -103,6 +104,15 @@ case class HoodieFileIndex(spark: SparkSession,
     endCompletionTime = options.get(DataSourceReadOptions.END_COMMIT.key)) with FileIndex {
 
   @transient protected var hasPushedDownPartitionPredicates: Boolean = false
+  private val isPartitionSimpleBucketIndex = PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(spark.sparkContext.hadoopConfiguration,
+    metaClient.getBasePath.toString)
+
+  @transient private lazy val bucketIndexSupport = if (isPartitionSimpleBucketIndex) {
+    val specifiedQueryInstant = options.get(DataSourceReadOptions.TIME_TRAVEL_AS_OF_INSTANT.key).map(HoodieSqlCommonUtils.formatQueryInstant)
+    new PartitionBucketIndexSupport(spark, metadataConfig, metaClient, specifiedQueryInstant)
path: hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieFileIndex.scala
user: login: zhangyue19921010
body: Same as `TimeTravel` scenarios mentioned above.
created_at: 2025-04-03T06:21:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1321,6 +1321,49 @@ void testBulkInsertWithPartitionBucketIndex(String operationType, String tableTy
     assertEquals(expected.stream().sorted().collect(Collectors.toList()), actual.stream().sorted().collect(Collectors.toList()));
   }
 
+  @Test
+  void tesQueryWithPartitionBucketIndexPruning() {
+    String operationType = "upsert";
+    String tableType = "MERGE_ON_READ";
+    TableEnvironment tableEnv = batchTableEnv;
+    // csv source
+    String csvSourceDDL = TestConfigurations.getCsvSourceDDL("csv_source", "test_source_5.data");
+    tableEnv.executeSql(csvSourceDDL);
+    String catalogName = "hudi_" + operationType;
+    String hudiCatalogDDL = catalog(catalogName)
+        .catalogPath(tempFile.getAbsolutePath())
+        .end();
+
+    tableEnv.executeSql(hudiCatalogDDL);
+    String dbName = "hudi";
+    tableEnv.executeSql("create database " + catalogName + "." + dbName);
+    String basePath = tempFile.getAbsolutePath() + "/hudi/hoodie_sink";
+
+    String hoodieTableDDL = sql(catalogName + ".hudi.hoodie_sink")
+        .option(FlinkOptions.PATH, basePath)
+        .option(FlinkOptions.TABLE_TYPE, tableType)
+        .option(FlinkOptions.OPERATION, operationType)
+        .option(FlinkOptions.WRITE_BULK_INSERT_SHUFFLE_INPUT, true)
+        .option(FlinkOptions.INDEX_TYPE, "BUCKET")
+        .option(FlinkOptions.HIVE_STYLE_PARTITIONING, "true")
+        .option(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, "1")
+        .option(FlinkOptions.BUCKET_INDEX_PARTITION_RULE, "regex")
+        .option(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS, "partition=(par1|par2),2")
+        .end();
+    tableEnv.executeSql(hoodieTableDDL);
+
+    String insertInto = "insert into " + catalogName + ".hudi.hoodie_sink select * from csv_source";
+    execInsertSql(tableEnv, insertInto);
+
+    List<Row> result1 = CollectionUtil.iterableToList(
path: hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/ITTestHoodieDataSource.java
user: login: danny0405
body: use `execSelectSql(TableEnvironment tEnv, String select)`
created_at: 2025-04-03T08:42:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -69,7 +70,7 @@ public class FileIndex implements Serializable {
   private final org.apache.hadoop.conf.Configuration hadoopConf;
   private final PartitionPruners.PartitionPruner partitionPruner; // for partition pruning
   private final ColumnStatsProbe colStatsProbe;                   // for probing column stats
-  private final int dataBucketHashing;                                   // for bucket pruning
+  private final Option<Functions.Function1<Integer, Integer>> dataBucket;                                   // for bucket pruning
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
user: login: danny0405
body: Can we use Java function, can we align the indentation of the comments:

```java
  // for partition pruning
  // for probing column stats
  // for bucket pruning
```
created_at: 2025-04-03T08:43:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -69,7 +70,7 @@ public class FileIndex implements Serializable {
   private final org.apache.hadoop.conf.Configuration hadoopConf;
   private final PartitionPruners.PartitionPruner partitionPruner; // for partition pruning
   private final ColumnStatsProbe colStatsProbe;                   // for probing column stats
-  private final int dataBucketHashing;                                   // for bucket pruning
+  private final Option<Functions.Function1<Integer, Integer>> dataBucket;                                   // for bucket pruning
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
user: login: danny0405
body: dataBucket -> dataBucketFunc
created_at: 2025-04-03T08:44:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -43,9 +45,7 @@
 public class PrimaryKeyPruners {
   private static final Logger LOG = LoggerFactory.getLogger(PrimaryKeyPruners.class);
 
-  public static final int BUCKET_ID_NO_PRUNING = -1;
-
-  public static int getBucketFieldHashing(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
+  public static Option<Functions.Function1<Integer, Integer>> getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
user: login: danny0405
body: getBucketId -> getBucketIdFunc
created_at: 2025-04-03T08:46:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -43,9 +45,7 @@
 public class PrimaryKeyPruners {
   private static final Logger LOG = LoggerFactory.getLogger(PrimaryKeyPruners.class);
 
-  public static final int BUCKET_ID_NO_PRUNING = -1;
-
-  public static int getBucketFieldHashing(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
+  public static Option<Functions.Function1<Integer, Integer>> getBucketId(List<ResolvedExpression> hashKeyFilters, Configuration conf) {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/prune/PrimaryKeyPruners.java
user: login: danny0405
body: It looks like the return value is always non-empty, can we just return the function instead of the option instead?
created_at: 2025-04-03T08:46:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -158,7 +158,7 @@ public class HoodieTableSource implements
   private List<Predicate> predicates;
   private ColumnStatsProbe columnStatsProbe;
   private PartitionPruners.PartitionPruner partitionPruner;
-  private int dataBucketHashing;
+  private Option<Functions.Function1<Integer, Integer>> dataBucket;
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: danny0405
body: dataBucket -> dataBucketFunc
created_at: 2025-04-03T08:47:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -179,7 +180,7 @@ public HoodieTableSource(
       @Nullable List<Predicate> predicates,
       @Nullable ColumnStatsProbe columnStatsProbe,
       @Nullable PartitionPruners.PartitionPruner partitionPruner,
-      int dataBucketHashing,
+      Option<Functions.Function1<Integer, Integer>> dataBucket,
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: danny0405
body: dataBucket -> dataBucketFunc
created_at: 2025-04-03T08:47:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -698,12 +699,7 @@ public ColumnStatsProbe getColumnStatsProbe() {
   }
 
   @VisibleForTesting
-  public int getDataBucket() {
-    return BucketIdentifier.getBucketId(this.dataBucketHashing, conf.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS));
-  }
-
-  @VisibleForTesting
-  public int getDataBucketHashing() {
-    return dataBucketHashing;
+  public Option<Functions.Function1<Integer, Integer>> getDataBucket() {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: danny0405
body: getDataBucketFunc
created_at: 2025-04-03T08:47:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -69,7 +70,7 @@ public class FileIndex implements Serializable {
   private final org.apache.hadoop.conf.Configuration hadoopConf;
   private final PartitionPruners.PartitionPruner partitionPruner; // for partition pruning
   private final ColumnStatsProbe colStatsProbe;                   // for probing column stats
-  private final int dataBucketHashing;                                   // for bucket pruning
+  private final Option<Functions.Function1<Integer, Integer>> dataBucket;                                   // for bucket pruning
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
user: login: zhangyue19921010
body: chenged
created_at: 2025-04-03T15:36:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -158,7 +158,7 @@ public class HoodieTableSource implements
   private List<Predicate> predicates;
   private ColumnStatsProbe columnStatsProbe;
   private PartitionPruners.PartitionPruner partitionPruner;
-  private int dataBucketHashing;
+  private Option<Functions.Function1<Integer, Integer>> dataBucket;
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: zhangyue19921010
body: changed
created_at: 2025-04-03T15:36:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -179,7 +180,7 @@ public HoodieTableSource(
       @Nullable List<Predicate> predicates,
       @Nullable ColumnStatsProbe columnStatsProbe,
       @Nullable PartitionPruners.PartitionPruner partitionPruner,
-      int dataBucketHashing,
+      Option<Functions.Function1<Integer, Integer>> dataBucket,
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: zhangyue19921010
body: all changed
created_at: 2025-04-03T15:36:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -698,12 +699,7 @@ public ColumnStatsProbe getColumnStatsProbe() {
   }
 
   @VisibleForTesting
-  public int getDataBucket() {
-    return BucketIdentifier.getBucketId(this.dataBucketHashing, conf.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS));
-  }
-
-  @VisibleForTesting
-  public int getDataBucketHashing() {
-    return dataBucketHashing;
+  public Option<Functions.Function1<Integer, Integer>> getDataBucket() {
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: zhangyue19921010
body: changed
created_at: 2025-04-03T15:36:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -365,16 +368,16 @@ private PartitionPruners.PartitionPruner createPartitionPruner(List<ResolvedExpr
         .build();
   }
 
-  private int getDataBucket(List<ResolvedExpression> dataFilters) {
+  private Option<Function<Integer, Integer>> getDataBucketFunc(List<ResolvedExpression> dataFilters) {
     if (!OptionsResolver.isBucketIndexType(conf) || dataFilters.isEmpty()) {
-      return PrimaryKeyPruners.BUCKET_ID_NO_PRUNING;
+      return Option.empty();
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: zhangyue19921010
body: > It looks like the return value is always non-empty, can we just return the function instead of the option instead?

Here may return `Option.empty()` when `getDataBucketFunc`
created_at: 2025-04-03T15:37:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -365,16 +368,16 @@ private PartitionPruners.PartitionPruner createPartitionPruner(List<ResolvedExpr
         .build();
   }
 
-  private int getDataBucket(List<ResolvedExpression> dataFilters) {
+  private Option<Function<Integer, Integer>> getDataBucketFunc(List<ResolvedExpression> dataFilters) {
     if (!OptionsResolver.isBucketIndexType(conf) || dataFilters.isEmpty()) {
-      return PrimaryKeyPruners.BUCKET_ID_NO_PRUNING;
+      return Option.empty();
     }
     Set<String> indexKeyFields = Arrays.stream(OptionsResolver.getIndexKeys(conf)).collect(Collectors.toSet());
     List<ResolvedExpression> indexKeyFilters = dataFilters.stream().filter(expr -> ExpressionUtils.isEqualsLitExpr(expr, indexKeyFields)).collect(Collectors.toList());
     if (!ExpressionUtils.isFilteringByAllFields(indexKeyFilters, indexKeyFields)) {
-      return PrimaryKeyPruners.BUCKET_ID_NO_PRUNING;
+      return Option.empty();
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSource.java
user: login: zhangyue19921010
body: > It looks like the return value is always non-empty, can we just return the function instead of the option instead?

Here also may return `Option.empty()` when `getDataBucketFunc`
created_at: 2025-04-03T15:38:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1321,6 +1321,49 @@ void testBulkInsertWithPartitionBucketIndex(String operationType, String tableTy
     assertEquals(expected.stream().sorted().collect(Collectors.toList()), actual.stream().sorted().collect(Collectors.toList()));
   }
 
+  @Test
+  void tesQueryWithPartitionBucketIndexPruning() {
+    String operationType = "upsert";
+    String tableType = "MERGE_ON_READ";
+    TableEnvironment tableEnv = batchTableEnv;
+    // csv source
+    String csvSourceDDL = TestConfigurations.getCsvSourceDDL("csv_source", "test_source_5.data");
+    tableEnv.executeSql(csvSourceDDL);
+    String catalogName = "hudi_" + operationType;
+    String hudiCatalogDDL = catalog(catalogName)
+        .catalogPath(tempFile.getAbsolutePath())
+        .end();
+
+    tableEnv.executeSql(hudiCatalogDDL);
+    String dbName = "hudi";
+    tableEnv.executeSql("create database " + catalogName + "." + dbName);
+    String basePath = tempFile.getAbsolutePath() + "/hudi/hoodie_sink";
+
+    String hoodieTableDDL = sql(catalogName + ".hudi.hoodie_sink")
+        .option(FlinkOptions.PATH, basePath)
+        .option(FlinkOptions.TABLE_TYPE, tableType)
+        .option(FlinkOptions.OPERATION, operationType)
+        .option(FlinkOptions.WRITE_BULK_INSERT_SHUFFLE_INPUT, true)
+        .option(FlinkOptions.INDEX_TYPE, "BUCKET")
+        .option(FlinkOptions.HIVE_STYLE_PARTITIONING, "true")
+        .option(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS, "1")
+        .option(FlinkOptions.BUCKET_INDEX_PARTITION_RULE, "regex")
+        .option(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS, "partition=(par1|par2),2")
+        .end();
+    tableEnv.executeSql(hoodieTableDDL);
+
+    String insertInto = "insert into " + catalogName + ".hudi.hoodie_sink select * from csv_source";
+    execInsertSql(tableEnv, insertInto);
+
+    List<Row> result1 = CollectionUtil.iterableToList(
path: hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/ITTestHoodieDataSource.java
user: login: zhangyue19921010
body: done
created_at: 2025-04-03T15:38:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
-      return allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1);
+      if (instant.isPresent()) {
+        Option<String> res = getHashingConfigInstantToLoadBeforeOrOn(allCommittedHashingConfig, instant.get());
+        // fall back to look up archived hashing config instant before return empty
+        return res.isPresent() ? res : getHashingConfigInstantToLoadBeforeOrOn(getArchiveHashingConfigInstants(metaClient), instant.get());
+      } else {
+        return Option.of(allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1));
+      }
     } catch (Exception e) {
       throw new HoodieException("Failed to get hashing config instant to load.", e);
     }
   }
 
+  private static Option<String> getHashingConfigInstantToLoadBeforeOrOn(List<String> hashingConfigInstants, String instant) {
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: danny0405
body: > Maybe we should load the hashing configuration associated with R3 (the latest bucket rescale operation before or equal to C4).

We can not do that because the data file layouts had been changed.
created_at: 2025-04-04T00:33:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
-      return allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1);
+      if (instant.isPresent()) {
+        Option<String> res = getHashingConfigInstantToLoadBeforeOrOn(allCommittedHashingConfig, instant.get());
+        // fall back to look up archived hashing config instant before return empty
+        return res.isPresent() ? res : getHashingConfigInstantToLoadBeforeOrOn(getArchiveHashingConfigInstants(metaClient), instant.get());
+      } else {
+        return Option.of(allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1));
+      }
     } catch (Exception e) {
       throw new HoodieException("Failed to get hashing config instant to load.", e);
     }
   }
 
+  private static Option<String> getHashingConfigInstantToLoadBeforeOrOn(List<String> hashingConfigInstants, String instant) {
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: zhangyue19921010
body: Yes, Danny. After Bucket Rescale is completed, the data layout will change. Therefore, for Spark's Time Travel，something like travel to specific time point snapshot view(https://hudi.apache.org/docs/sql_queries#time-travel-query), 

Such as executing a query like `SELECT * FROM <table_name> TIMESTAMP AS OF <instant1> WHERE <filter_conditions>`, the Hudi would init specifiedQueryTimestamp through HoodieBaseRelation.
```
  protected lazy val specifiedQueryTimestamp: Option[String] =
    optParams.get(DataSourceReadOptions.TIME_TRAVEL_AS_OF_INSTANT.key)
      .map(HoodieSqlCommonUtils.formatQueryInstant)
```
Then get `Schema` and build `fsView` based on `specifiedQueryTimestamp`

For constructing the FsView, Hudi will call `getLatestMergedFileSlicesBeforeOrOn(String partitionStr, String maxInstantTime)`, travel fs view to the specified version. At this point, it is also necessary to load the corresponding hashing_config that was valid at that specific timestamp to ensure the historical data layout
```
  protected def listLatestFileSlices(globPaths: Seq[StoragePath], partitionFilters: Seq[Expression], dataFilters: Seq[Expression]): Seq[FileSlice] = {
    queryTimestamp match {
      case Some(ts) =>
        specifiedQueryTimestamp.foreach(t => validateTimestampAsOf(metaClient, t))

        val partitionDirs = if (globPaths.isEmpty) {
          fileIndex.listFiles(partitionFilters, dataFilters)
        } else {
          val inMemoryFileIndex = HoodieInMemoryFileIndex.create(sparkSession, globPaths)
          inMemoryFileIndex.listFiles(partitionFilters, dataFilters)
        }

        val fsView = new HoodieTableFileSystemView(
          metaClient, timeline, sparkAdapter.getSparkPartitionedFileUtils.toFileStatuses(partitionDirs)
            .map(fileStatus => HadoopFSUtils.convertToStoragePathInfo(fileStatus))
            .asJava)

        fsView.getPartitionPaths.asScala.flatMap { partitionPath =>
          val relativePath = getRelativePartitionPath(convertToStoragePath(basePath), partitionPath)
          fsView.getLatestMergedFileSlicesBeforeOrOn(relativePath, ts).iterator().asScala
        }.toSeq

      case _ => Seq()
    }
  }
```
created_at: 2025-04-04T13:59:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
-      return allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1);
+      if (instant.isPresent()) {
+        Option<String> res = getHashingConfigInstantToLoadBeforeOrOn(allCommittedHashingConfig, instant.get());
+        // fall back to look up archived hashing config instant before return empty
+        return res.isPresent() ? res : getHashingConfigInstantToLoadBeforeOrOn(getArchiveHashingConfigInstants(metaClient), instant.get());
+      } else {
+        return Option.of(allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1));
+      }
     } catch (Exception e) {
       throw new HoodieException("Failed to get hashing config instant to load.", e);
     }
   }
 
+  private static Option<String> getHashingConfigInstantToLoadBeforeOrOn(List<String> hashingConfigInstants, String instant) {
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: zhangyue19921010
body: For example 

DeltaCommit1 ==> Write `C1_File1, C1_File2`
Bucket-Rescale Commit2 ==> Write `C2_File1, C2_File2, C2_File3(Replaced  C1_File1, C1_File2)`
DeltaCommit 3 ==> Write `C3_File1_Log1`
Bucket-Rescale Commit4 ==> Write `C4_File1(Replaced C2_File1, C2_File2, C2_File3 and C3_File1_Log1)`



For Sql `SELECT * FROM hudi_table TIMESTAMP AS OF <DeltaCommit 3>`, we need to load Bucket-Rescale Commit2 instead of load latest hashing config Bucket-Rescale Commit4
created_at: 2025-04-04T14:00:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -196,24 +196,51 @@ public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(Hoodie
 
   /**
    * Get Latest committed hashing config instant to load.
+   * If instant is empty, then return latest hashing config instant
    */
-  public static String getLatestHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+  public static Option<String> getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient, Option<String> instant) {
     try {
       List<String> allCommittedHashingConfig = getCommittedHashingConfigInstants(metaClient);
-      return allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1);
+      if (instant.isPresent()) {
+        Option<String> res = getHashingConfigInstantToLoadBeforeOrOn(allCommittedHashingConfig, instant.get());
+        // fall back to look up archived hashing config instant before return empty
+        return res.isPresent() ? res : getHashingConfigInstantToLoadBeforeOrOn(getArchiveHashingConfigInstants(metaClient), instant.get());
+      } else {
+        return Option.of(allCommittedHashingConfig.get(allCommittedHashingConfig.size() - 1));
+      }
     } catch (Exception e) {
       throw new HoodieException("Failed to get hashing config instant to load.", e);
     }
   }
 
+  private static Option<String> getHashingConfigInstantToLoadBeforeOrOn(List<String> hashingConfigInstants, String instant) {
path: hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
user: login: zhangyue19921010
body: Add a new Spark UT `test("Test BucketID Pruning With Partition Bucket Index")`
Without This PR will throw Exception 
```
Expected Array([1111,3333.0,3333,2021-01-05]), but got Array()
ScalaTestFailureLocation: org.apache.spark.sql.hudi.common.HoodieSparkSqlTestBase at (HoodieSparkSqlTestBase.scala:135)
org.scalatest.exceptions.TestFailedException: Expected Array([1111,3333.0,3333,2021-01-05]), but got Array()

```

With PR in `Always load latest hashing config` logic, will throw exception
```
Expected Array([1111,2222.0,2222,2021-01-05]), but got Array()
ScalaTestFailureLocation: org.apache.spark.sql.hudi.common.HoodieSparkSqlTestBase at (HoodieSparkSqlTestBase.scala:135)
org.scalatest.exceptions.TestFailedException: Expected Array([1111,2222.0,2222,2021-01-05]), but got Array()
```
created_at: 2025-04-04T14:03:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -157,28 +157,34 @@ public List<StoragePathInfo> getFilesInPartitions() {
     if (partitions.length < 1) {
       return Collections.emptyList();
     }
-    List<StoragePathInfo> allFiles = FSUtils.getFilesInPartitions(
-            new HoodieFlinkEngineContext(hadoopConf),
-            new HoodieHadoopStorage(path, HadoopFSUtils.getStorageConf(hadoopConf)), metadataConfig, path.toString(), partitions)
-        .values().stream()
-        .flatMap(Collection::stream)
-        .collect(Collectors.toList());
+    Map<String, List<StoragePathInfo>> partition2Files = FSUtils.getFilesInPartitions(
+        new HoodieFlinkEngineContext(hadoopConf),
+        new HoodieHadoopStorage(path, HadoopFSUtils.getStorageConf(hadoopConf)), metadataConfig, path.toString(), partitions);
+
+    List<StoragePathInfo> allFiles;
+    // bucket pruning
+    if (this.partitionBucketIdFunc != null) {
+      allFiles = partition2Files.entrySet().stream().flatMap(entry -> {
+        String partitionPath = entry.getKey();
+        int bucketId = partitionBucketIdFunc.apply(partitionPath);
+        String bucketIdStr = BucketIdentifier.bucketIdStr(bucketId);
+        List<StoragePathInfo> filesInPartition = entry.getValue();
+        return filesInPartition.stream().filter(fileInfo -> fileInfo.getPath().getName().contains(bucketIdStr));
+      }).collect(Collectors.toList());
+    } else {
+      allFiles = FSUtils.getFilesInPartitions(
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
user: login: danny0405
body: should use `partition2Files` to avoid redundant files fetching.
created_at: 2025-04-05T01:23:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -157,28 +157,34 @@ public List<StoragePathInfo> getFilesInPartitions() {
     if (partitions.length < 1) {
       return Collections.emptyList();
     }
-    List<StoragePathInfo> allFiles = FSUtils.getFilesInPartitions(
-            new HoodieFlinkEngineContext(hadoopConf),
-            new HoodieHadoopStorage(path, HadoopFSUtils.getStorageConf(hadoopConf)), metadataConfig, path.toString(), partitions)
-        .values().stream()
-        .flatMap(Collection::stream)
-        .collect(Collectors.toList());
+    Map<String, List<StoragePathInfo>> partition2Files = FSUtils.getFilesInPartitions(
+        new HoodieFlinkEngineContext(hadoopConf),
+        new HoodieHadoopStorage(path, HadoopFSUtils.getStorageConf(hadoopConf)), metadataConfig, path.toString(), partitions);
+
+    List<StoragePathInfo> allFiles;
+    // bucket pruning
+    if (this.partitionBucketIdFunc != null) {
+      allFiles = partition2Files.entrySet().stream().flatMap(entry -> {
+        String partitionPath = entry.getKey();
+        int bucketId = partitionBucketIdFunc.apply(partitionPath);
+        String bucketIdStr = BucketIdentifier.bucketIdStr(bucketId);
+        List<StoragePathInfo> filesInPartition = entry.getValue();
+        return filesInPartition.stream().filter(fileInfo -> fileInfo.getPath().getName().contains(bucketIdStr));
+      }).collect(Collectors.toList());
+    } else {
+      allFiles = FSUtils.getFilesInPartitions(
path: hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/source/FileIndex.java
user: login: zhangyue19921010
body: Opsss. changed this.
created_at: 2025-04-05T01:49:35Z
author_association: CONTRIBUTOR

PR comments:
user: login: zhangyue19921010
created_at: 2025-03-31T16:10:52Z
body: @hudi-bot run azure
author_association: CONTRIBUTOR
user: login: danny0405
created_at: 2025-04-04T10:25:44Z
body: Thanks for the contribution, I have made some refactoring with the given patch:

[13060.patch.zip](https://github.com/user-attachments/files/19602219/13060.patch.zip)

author_association: CONTRIBUTOR
user: login: zhangyue19921010
created_at: 2025-04-04T14:05:38Z
body: > Thanks for the contribution, I have made some refactoring with the given patch:
> 
> [13060.patch.zip](https://github.com/user-attachments/files/19602219/13060.patch.zip)

Thanks for your help Danny, All changed.
author_association: CONTRIBUTOR
user: login: danny0405
created_at: 2025-04-05T01:53:44Z
body: I have made another patch to fix the bucket pruning logging: 
[fix_the_bucket_pruning_logging.patch.zip](https://github.com/user-attachments/files/19612894/fix_the_bucket_pruning_logging.patch.zip)


author_association: CONTRIBUTOR
user: login: hudi-bot
created_at: 2025-04-05T04:32:33Z
body: <!--
Meta data
{
  "version" : 1,
  "metaDataEntries" : [ {
    "hash" : "2a2886f4340ef90adb653849a38307c5b7f9a5c9",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4551",
    "triggerID" : "2a2886f4340ef90adb653849a38307c5b7f9a5c9",
    "triggerType" : "PUSH"
  }, {
    "hash" : "cb8f72d1483384de0f72270bdb3bba4001380f8f",
    "status" : "UNKNOWN",
    "url" : "TBD",
    "triggerID" : "cb8f72d1483384de0f72270bdb3bba4001380f8f",
    "triggerType" : "PUSH"
  }, {
    "hash" : "6f0939dfbbcddd07213c638d01933c25fe486941",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4552",
    "triggerID" : "6f0939dfbbcddd07213c638d01933c25fe486941",
    "triggerType" : "PUSH"
  }, {
    "hash" : "6f0939dfbbcddd07213c638d01933c25fe486941",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4553",
    "triggerID" : "2766718497",
    "triggerType" : "MANUAL"
  }, {
    "hash" : "2bafa0c28b3f912a5cdac376387eb1e8f6edaa83",
    "status" : "UNKNOWN",
    "url" : "TBD",
    "triggerID" : "2bafa0c28b3f912a5cdac376387eb1e8f6edaa83",
    "triggerType" : "PUSH"
  }, {
    "hash" : "cbb666d497e8b4ae98e342197ac951f01bbf0a2c",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4602",
    "triggerID" : "cbb666d497e8b4ae98e342197ac951f01bbf0a2c",
    "triggerType" : "PUSH"
  }, {
    "hash" : "95c0d03c1707ea4e7d119f5751d042ad6a742664",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4603",
    "triggerID" : "95c0d03c1707ea4e7d119f5751d042ad6a742664",
    "triggerType" : "PUSH"
  }, {
    "hash" : "6b38a5b0831e624603ada818aa19f7368336dee1",
    "status" : "UNKNOWN",
    "url" : "TBD",
    "triggerID" : "6b38a5b0831e624603ada818aa19f7368336dee1",
    "triggerType" : "PUSH"
  }, {
    "hash" : "1610e8ca29bc433cda1a20d7d31d98e2d1bd7bd5",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4619",
    "triggerID" : "1610e8ca29bc433cda1a20d7d31d98e2d1bd7bd5",
    "triggerType" : "PUSH"
  }, {
    "hash" : "b53f457aa1ef63f665add37a75d4f8e9a614780c",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4651",
    "triggerID" : "b53f457aa1ef63f665add37a75d4f8e9a614780c",
    "triggerType" : "PUSH"
  }, {
    "hash" : "2d79e40b46ab9af5b518314c9031bfc46434841d",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4654",
    "triggerID" : "2d79e40b46ab9af5b518314c9031bfc46434841d",
    "triggerType" : "PUSH"
  }, {
    "hash" : "b3a978b204ecd08df87cf4116647ce95088eedc0",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4656",
    "triggerID" : "b3a978b204ecd08df87cf4116647ce95088eedc0",
    "triggerType" : "PUSH"
  }, {
    "hash" : "de0276c24c4b560b6564e5320ea4fb6ab9f8cdbe",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4659",
    "triggerID" : "de0276c24c4b560b6564e5320ea4fb6ab9f8cdbe",
    "triggerType" : "PUSH"
  }, {
    "hash" : "69bfd3369acf7ebabb3b58055ab583c1f2e2ff23",
    "status" : "DELETED",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4670",
    "triggerID" : "69bfd3369acf7ebabb3b58055ab583c1f2e2ff23",
    "triggerType" : "PUSH"
  }, {
    "hash" : "ecd6bb61d253d969a26fd97b7c52d5d79c5c02aa",
    "status" : "SUCCESS",
    "url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4672",
    "triggerID" : "ecd6bb61d253d969a26fd97b7c52d5d79c5c02aa",
    "triggerType" : "PUSH"
  } ]
}-->
## CI report:

* cb8f72d1483384de0f72270bdb3bba4001380f8f UNKNOWN
* 2bafa0c28b3f912a5cdac376387eb1e8f6edaa83 UNKNOWN
* 6b38a5b0831e624603ada818aa19f7368336dee1 UNKNOWN
* ecd6bb61d253d969a26fd97b7c52d5d79c5c02aa Azure: [SUCCESS](https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=4672) 

<details>
<summary>Bot commands</summary>
  @hudi-bot supports the following commands:

 - `@hudi-bot run azure` re-run the last Azure build
</details>
author_association: COLLABORATOR
user: login: zhangyue19921010
created_at: 2025-04-05T04:40:43Z
body: > I have made another patch to fix the bucket pruning logging: [fix_the_bucket_pruning_logging.patch.zip](https://github.com/user-attachments/files/19612894/fix_the_bucket_pruning_logging.patch.zip)

Done. Also CI passed.
author_association: CONTRIBUTOR