You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hudi.apache.org by zh...@apache.org on 2025/03/21 08:04:23 UTC

(hudi) branch HUDI-8990 created (now aa9ed6ce12c)

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

zhangyue19921010 pushed a change to branch HUDI-8990
in repository https://gitbox.apache.org/repos/asf/hudi.git


      at aa9ed6ce12c finish Spark related conding and test

This branch includes the following new commits:

     new 475b4f108e0 finish hashing_config initial && related UT
     new 60de489d9a9 finish hashing_config initial && related UT
     new 989172d58a5 finsih PartitionBucketIndexCalculator part
     new 97bf409a5ad finish Flink related coding and UT && Spark partitioner and index without test
     new 184231de7b0 finish Spark call update hashing config need more test
     new aa9ed6ce12c finish Spark related conding and test

The 6 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.



(hudi) 03/06: finsih PartitionBucketIndexCalculator part

Posted by zh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zhangyue19921010 pushed a commit to branch HUDI-8990
in repository https://gitbox.apache.org/repos/asf/hudi.git

commit 989172d58a5fc49a9a1823951652e7ba8b584cca
Author: YueZhang <zh...@163.com>
AuthorDate: Tue Mar 18 21:42:42 2025 +0800

    finsih PartitionBucketIndexCalculator part
---
 .../org/apache/hudi/config/HoodieIndexConfig.java  |  13 +-
 .../bucket/PartitionBucketIndexCalculator.java     | 302 +++++++++++++++++++++
 ...tionRule.java => PartitionBucketIndexRule.java} |   8 +-
 .../index/bucket/PartitionBucketIndexUtils.java    | 148 ++++++++++
 .../hudi/index/bucket/SimpleBucketIndexUtils.java  |  81 ------
 .../bucket/TestPartitionBucketIndexCalculator.java | 144 ++++++++++
 .../apache/spark/sql/BucketPartitionUtils.scala    |   6 +-
 .../hudi/common/util/hash/BucketIndexUtil.java     |   4 +-
 .../hudi/common/util/hash/TestBucketIndexUtil.java |  16 +-
 .../apache/hudi/configuration/FlinkOptions.java    |  12 +
 .../hudi/configuration/OptionsInference.java       |  17 ++
 .../sink/bucket/BucketBulkInsertWriterHelper.java  |  19 +-
 .../sink/bucket/BucketStreamWriteFunction.java     |  29 +-
 .../sink/partitioner/BucketIndexPartitioner.java   |  31 ++-
 .../java/org/apache/hudi/sink/utils/Pipelines.java |  12 +-
 .../org/apache/hudi/table/HoodieTableSink.java     |   2 +
 .../hudi/table/catalog/HoodieCatalogUtil.java      |   4 +-
 .../hudi/sink/utils/BulkInsertFunctionWrapper.java |   5 +-
 .../hudi/table/catalog/TestHoodieCatalog.java      |   8 +-
 style/checkstyle-suppressions.xml                  |   1 +
 20 files changed, 733 insertions(+), 129 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
index bfad563a9e7..99a42bd41d0 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
@@ -28,7 +28,7 @@ import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.exception.HoodieIndexException;
 import org.apache.hudi.exception.HoodieNotSupportedException;
 import org.apache.hudi.index.HoodieIndex;
-import org.apache.hudi.index.bucket.HoodieSimpleBucketIndexPartitionRule;
+import org.apache.hudi.index.bucket.PartitionBucketIndexRule;
 import org.apache.hudi.keygen.constant.KeyGeneratorOptions;
 
 import org.slf4j.Logger;
@@ -289,10 +289,19 @@ public class HoodieIndexConfig extends HoodieConfig {
 
   public static final ConfigProperty<String> BUCKET_INDEX_PARTITION_RULE_TYPE = ConfigProperty
       .key("hoodie.bucket.index.partition.rule.type")
-      .defaultValue(HoodieSimpleBucketIndexPartitionRule.REGEX.value)
+      .defaultValue(PartitionBucketIndexRule.REGEX.name)
       .markAdvanced()
       .withDocumentation("Rule parser for expressions when using partition level bucket index, default regex.");
 
+  public static final ConfigProperty<Boolean> BUCKET_INDEX_PARTITION_LEVEL_ENABLE = ConfigProperty
+      .key("hoodie.bucket.index.partition.level.enable")
+      .defaultValue(false)
+      .withDocumentation("Inner config to determine whether to enable partition level bucket index");
+  public static final ConfigProperty<String> BUCKET_INDEX_PARTITION_LOAD_INSTANT = ConfigProperty
+      .key("hoodie.bucket.index.partition.load.instant")
+      .defaultValue("")
+      .withDocumentation("Inner config to determine which hashing_config to load when using partition level bucket index");
+
   public static final ConfigProperty<String> BUCKET_INDEX_PARTITION_EXPRESSIONS = ConfigProperty
       .key("hoodie.bucket.index.partition.expressions")
       .noDefaultValue()
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
new file mode 100644
index 00000000000..ff8af59a73d
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
@@ -0,0 +1,302 @@
+/*
+ * 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.index.bucket;
+
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.ValidationUtils;
+import org.apache.hudi.exception.HoodieException;
+import org.apache.hudi.exception.HoodieIOException;
+import org.apache.hudi.hadoop.fs.HadoopFSUtils;
+import org.apache.hudi.storage.HoodieStorage;
+import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.storage.hadoop.HoodieHadoopStorage;
+
+import org.apache.commons.collections.map.LRUMap;
+import org.apache.hadoop.conf.Configuration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * A singleton implementation of PartitionBucketIndexCalculator that ensures only one instance
+ * exists for each unique hashingInstantToLoad value.
+ */
+public class PartitionBucketIndexCalculator implements Serializable {
+  private static final Logger LOG = LoggerFactory.getLogger(PartitionBucketIndexCalculator.class);
+  // Map to store singleton instances for each instantToLoad + configuration hash combination
+  private static final ConcurrentMap<String, PartitionBucketIndexCalculator> INSTANCES = new ConcurrentHashMap<>();
+  private static final int CACHE_SIZE = 100_000;
+  private PartitionBucketIndexHashingConfig hashingConfig;
+  private int defaultBucketNumber;
+  private final String instantToLoad;
+  // Cache for partition to bucket number mapping
+  @SuppressWarnings("unchecked")
+  private final Map<String, Integer> partitionToBucketCache = new LRUMap(CACHE_SIZE);
+  private RuleEngine ruleEngine;
+
+  /**
+   * Private constructor to prevent direct instantiation
+   *
+   * @param instantToLoad The instant to load
+   * @param hadoopConf The Hadoop configuration
+   */
+  private PartitionBucketIndexCalculator(String instantToLoad, Configuration hadoopConf, String basePath) {
+    this.instantToLoad = instantToLoad;
+    StoragePath metaPath = new StoragePath(basePath, HoodieTableMetaClient.METAFOLDER_NAME);
+    StoragePath hashingBase = new StoragePath(metaPath, HoodieTableMetaClient.BUCKET_INDEX_METAFOLDER_CONFIG_FOLDER);
+    StoragePath hashingConfigPath =
+        new StoragePath(hashingBase,
+            instantToLoad + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
+
+    try (HoodieHadoopStorage storage = new HoodieHadoopStorage(hashingConfigPath, HadoopFSUtils.getStorageConf(hadoopConf))) {
+      init(storage, hashingConfigPath);
+    } catch (IOException e) {
+      throw new HoodieIOException("Failed to initialize PartitionBucketIndexCalculator ", e);
+    }
+  }
+
+  private PartitionBucketIndexCalculator(String instantToLoad, HoodieTableMetaClient client) {
+    this.instantToLoad = instantToLoad;
+    String metaPath = client.getHashingMetadataConfigPath();
+    StoragePath hashingConfigPath = new StoragePath(metaPath, instantToLoad + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
+
+    try (HoodieStorage storage = client.getStorage()) {
+      init(storage, hashingConfigPath);
+    } catch (IOException e) {
+      throw new HoodieIOException("Failed to initialize PartitionBucketIndexCalculator ", e);
+    }
+  }
+
+  private void init(HoodieStorage storage, StoragePath hashingConfigPath) {
+    Option<PartitionBucketIndexHashingConfig> config = PartitionBucketIndexUtils.loadHashingConfig(storage, hashingConfigPath);
+    ValidationUtils.checkArgument(config.isPresent());
+    this.hashingConfig = config.get();
+    this.defaultBucketNumber = config.get().getDefaultBucketNumber();
+    String expressions = config.get().getExpressions();
+    String ruleType = config.get().getRule();
+    this.ruleEngine = createRuleEngine(ruleType, expressions);
+  }
+
+  /**
+   * Gets the singleton instance for the specified instantToLoad and configuration
+   *
+   * @param instantToLoad The instant to load
+   * @param hadoopConf The Hadoop configuration
+   * @return The singleton instance
+   */
+  public static PartitionBucketIndexCalculator getInstance(String instantToLoad, Configuration hadoopConf, String basePath) {
+    // Using instantToLoad as the key for the cache
+    return INSTANCES.computeIfAbsent(instantToLoad,
+        key -> {
+          LOG.info("Creating new PartitionBucketIndexCalculator instance for instantToLoad: {}", key);
+          return new PartitionBucketIndexCalculator(key, hadoopConf, basePath);
+        });
+  }
+
+  public static PartitionBucketIndexCalculator getInstance(String instantToLoad, HoodieTableMetaClient client) {
+    // Using instantToLoad as the key for the cache
+    return INSTANCES.computeIfAbsent(instantToLoad,
+        key -> {
+          LOG.info("Creating new PartitionBucketIndexCalculator instance for instantToLoad: {}", key);
+          return new PartitionBucketIndexCalculator(key, client);
+        });
+  }
+
+  /**
+   * Computes the bucket number for a given partition path
+   *
+   * @param partitionPath The partition path
+   * @return The computed bucket number
+   */
+  public int computeNumBuckets(String partitionPath) {
+    // Check cache first
+    Integer cachedBucketNumber = partitionToBucketCache.get(partitionPath);
+    if (cachedBucketNumber != null) {
+      return cachedBucketNumber;
+    }
+
+    // Calculate bucket number using the rule engine
+    int bucketNumber = ruleEngine.calculateBucketNumber(partitionPath);
+
+    // If no rule matched, use default bucket number
+    if (bucketNumber == -1) {
+      bucketNumber = defaultBucketNumber;
+      LOG.debug("No rule matched for partition: {}. Using default bucket number: {}",
+          partitionPath, defaultBucketNumber);
+    }
+
+    // Update cache
+    partitionToBucketCache.put(partitionPath, bucketNumber);
+
+    return bucketNumber;
+  }
+
+  /**
+   * Gets the instant to load
+   *
+   * @return The instant to load
+   */
+  public String getInstantToLoad() {
+    return instantToLoad;
+  }
+
+  /**
+   * Gets the hashing configuration
+   *
+   * @return The hashing configuration
+   */
+  public PartitionBucketIndexHashingConfig getHashingConfig() {
+    return hashingConfig;
+  }
+
+  /**
+   * Clears the instance cache (useful for testing or memory management)
+   */
+  public static void clearInstanceCache() {
+    INSTANCES.clear();
+  }
+
+  public int getCacheSize() {
+    return partitionToBucketCache.size();
+  }
+
+  public void clearCache() {
+    partitionToBucketCache.clear();
+    LOG.info("Cleared partition to bucket number cache");
+  }
+
+  // -------------------------------------------------------------------------------------------------------
+
+  /**
+   * Interface for rule engines that calculate bucket numbers
+   */
+  private interface RuleEngine extends Serializable {
+    /**
+     * Calculate bucket number for a partition path
+     * @param partitionPath The partition path
+     * @return The calculated bucket number, or -1 if no rule matches
+     */
+    int calculateBucketNumber(String partitionPath);
+  }
+
+  /**
+   * Factory method to create the appropriate rule engine based on rule type
+   */
+  private static PartitionBucketIndexCalculator.RuleEngine createRuleEngine(String ruleType, String expressions) {
+    switch (PartitionBucketIndexRule.valueOf(ruleType.toUpperCase())) {
+      case REGEX:
+        return new PartitionBucketIndexCalculator.RegexRuleEngine(expressions);
+      default:
+        LOG.error("Unsupported rule type: {}.", ruleType);
+        throw new UnsupportedOperationException("Unsupported rule type " + ruleType);
+    }
+  }
+
+  /**
+   * Regex-based rule engine implementation
+   */
+  private static class RegexRuleEngine implements PartitionBucketIndexCalculator.RuleEngine {
+    private static final long serialVersionUID = 1L;
+    private final List<RegexRule> rules = new ArrayList<>();
+
+    /**
+     * Represents a single regex rule with its pattern and bucket number
+     */
+    private static class RegexRule implements Serializable {
+      private static final long serialVersionUID = 1L;
+      private final Pattern pattern;
+      private final int bucketNumber;
+
+      public RegexRule(String regex, int bucketNumber) {
+        this.pattern = Pattern.compile(regex);
+        this.bucketNumber = bucketNumber;
+      }
+
+      public boolean matches(String input) {
+        Matcher matcher = pattern.matcher(input);
+        return matcher.find();
+      }
+
+      @Override
+      public String toString() {
+        return pattern.pattern() + " -> " + bucketNumber;
+      }
+    }
+
+    /**
+     * Initialize the regex rule engine with expressions
+     * @param expressions Format: "expression1,bucketNumber1;expression2,bucketNumber2;..."
+     */
+    public RegexRuleEngine(String expressions) {
+      parseExpressions(expressions);
+    }
+
+    /**
+     * Parse the expressions string and create regex rules
+     */
+    private void parseExpressions(String expressions) {
+      String[] ruleExpressions = expressions.split(";");
+
+      for (String ruleExpression : ruleExpressions) {
+        String[] parts = ruleExpression.trim().split(",");
+        if (parts.length != 2) {
+          throw new HoodieException("Invalid regex expression format. Expected 'pattern,bucketNumber' but got: " + ruleExpression);
+        }
+
+        String regex = parts[0].trim();
+        int bucketNumber;
+        try {
+          bucketNumber = Integer.parseInt(parts[1].trim());
+        } catch (NumberFormatException e) {
+          throw new HoodieException("Invalid bucket number in expression: " + ruleExpression, e);
+        }
+
+        rules.add(new PartitionBucketIndexCalculator.RegexRuleEngine.RegexRule(regex, bucketNumber));
+        LOG.info("Added regex rule: {} with bucket number: {}", regex, bucketNumber);
+      }
+
+      LOG.info("Initialized {} regex rules", rules.size());
+    }
+
+    @Override
+    public int calculateBucketNumber(String partitionPath) {
+      // Check each rule in order (priority by position)
+      for (PartitionBucketIndexCalculator.RegexRuleEngine.RegexRule rule : rules) {
+        if (rule.matches(partitionPath)) {
+          LOG.debug("Partition '{}' matched regex rule: {}", partitionPath, rule);
+          return rule.bucketNumber;
+        }
+      }
+
+      // No rule matched
+      return -1;
+    }
+  }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndexPartitionRule.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexRule.java
similarity index 84%
rename from hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndexPartitionRule.java
rename to hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexRule.java
index 8d86ed8080b..8b4a73e251d 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndexPartitionRule.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexRule.java
@@ -18,12 +18,12 @@
 
 package org.apache.hudi.index.bucket;
 
-public enum HoodieSimpleBucketIndexPartitionRule {
+public enum PartitionBucketIndexRule {
   REGEX("regex");
 
-  public final String value;
+  public final String name;
 
-  HoodieSimpleBucketIndexPartitionRule(String rule) {
-    this.value = rule;
+  PartitionBucketIndexRule(String name) {
+    this.name = name;
   }
 }
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
new file mode 100644
index 00000000000..82b61acd7c7
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
@@ -0,0 +1,148 @@
+/*
+ * 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.index.bucket;
+
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.common.table.timeline.HoodieInstant;
+import org.apache.hudi.common.table.timeline.HoodieTimeline;
+import org.apache.hudi.common.util.FileIOUtils;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.exception.HoodieException;
+import org.apache.hudi.exception.HoodieIOException;
+import org.apache.hudi.hadoop.fs.HadoopFSUtils;
+import org.apache.hudi.storage.HoodieInstantWriter;
+import org.apache.hudi.storage.HoodieStorage;
+import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.storage.StoragePathInfo;
+import org.apache.hudi.storage.hadoop.HoodieHadoopStorage;
+
+import org.apache.hadoop.conf.Configuration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class PartitionBucketIndexUtils {
+  public static final String INITIAL_HASHING_CONFIG_INSTANT = HoodieTimeline.INIT_INSTANT_TS;
+
+  private static final Logger LOG = LoggerFactory.getLogger(PartitionBucketIndexUtils.class);
+
+  public static boolean isPartitionSimpleBucketIndex(Configuration conf) {
+    StoragePath storagePath = new StoragePath(HoodieTableMetaClient.PARTITION_BUCKET_INDEX_HASHING_FOLDER);
+    try (HoodieHadoopStorage storage = new HoodieHadoopStorage(storagePath, HadoopFSUtils.getStorageConf(conf))) {
+      return storage.exists(storagePath);
+    } catch (IOException e) {
+      throw new HoodieIOException("Failed to list PARTITION_BUCKET_INDEX_HASHING_FOLDER folder ", e);
+    }
+  }
+
+  public static boolean initHashingConfig(HoodieTableMetaClient metaClient,
+                                          String expressions,
+                                          String rule,
+                                          int defaultBucketNumber,
+                                          String instant) {
+    if (StringUtils.isNullOrEmpty(expressions)) {
+      return false;
+    }
+    String hashingInstant = StringUtils.isNullOrEmpty(instant) ? INITIAL_HASHING_CONFIG_INSTANT : instant;
+    HoodieStorage storage = metaClient.getStorage();
+    PartitionBucketIndexHashingConfig hashingConfig =
+        new PartitionBucketIndexHashingConfig(expressions, defaultBucketNumber, rule, PartitionBucketIndexHashingConfig.CURRENT_VERSION, hashingInstant);
+    StoragePath hashingConfigPath = new StoragePath(metaClient.getHashingMetadataConfigPath(), hashingConfig.getFilename());
+
+    try {
+      Option<byte []> content = Option.of(hashingConfig.toJsonString().getBytes(StandardCharsets.UTF_8));
+      storage.createImmutableFileInPath(hashingConfigPath, content.map(HoodieInstantWriter::convertByteArrayToWriter));
+    } catch (IOException ioe) {
+      throw new HoodieIOException("Failed to initHashingConfig ", ioe);
+    }
+    return true;
+  }
+
+  public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(HoodieStorage storage, StoragePathInfo hashingConfig) {
+    return loadHashingConfig(storage, hashingConfig.getPath());
+  }
+
+  public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(HoodieStorage storage, StoragePath hashingConfig) {
+    if (hashingConfig == null) {
+      return Option.empty();
+    }
+    try (InputStream is = storage.open(hashingConfig)) {
+      byte[] content = FileIOUtils.readAsByteArray(is);
+      return Option.of(PartitionBucketIndexHashingConfig.fromBytes(content));
+    } catch (IOException e) {
+      LOG.error("Error when loading hashing config, for path: " + hashingConfig.getName(), e);
+      throw new HoodieIOException("Error while loading hashing config", e);
+    }
+  }
+
+  public static String getHashingConfigInstantToLoad(HoodieTableMetaClient metaClient) {
+    List<String> instants = metaClient.getActiveTimeline().getCompletedReplaceTimeline()
+        .getInstants().stream().map(HoodieInstant::requestedTime).collect(Collectors.toList());
+    Option<HoodieInstant> earliestInstant = metaClient.getActiveTimeline().getCommitsTimeline().filterCompletedInstants().firstInstant();
+    String instantToLoad = "";
+    try {
+      // 按时间降序排列
+      List<String> hashingConfigInstants = metaClient.getStorage()
+          .listDirectEntries(new StoragePath(metaClient.getHashingMetadataConfigPath())).stream().map(info -> {
+            String instant = getHashingConfigInstant(info.getPath().getName());
+            if (StringUtils.isNullOrEmpty(instant)) {
+              throw new HoodieException("Failed to get hashing config instant to load.");
+            }
+            return instant;
+          }).sorted().collect(Collectors.toList());
+
+      for (String instant : hashingConfigInstants) {
+        if (instants.contains(instant)) {
+          instantToLoad = instant;
+          break;
+        } else if (earliestInstant.isPresent() && instant.compareTo(earliestInstant.get().requestedTime()) < 0){
+          instantToLoad = instant;
+        }
+      }
+
+      if (StringUtils.isNullOrEmpty(instantToLoad)) {
+        throw new HoodieException("Failed to get hashing config instant to load.");
+      }
+
+      return instantToLoad;
+    } catch (IOException e) {
+      throw new HoodieIOException("Failed to get hashing config instant to load.", e);
+    }
+  }
+
+  public static String getHashingConfigInstant(String hashingConfig) {
+    int lastIndex = hashingConfig.lastIndexOf('/');
+    if (lastIndex == -1) {
+      return null;
+    }
+    String fileName = hashingConfig.substring(lastIndex + 1);
+    int dotIndex = fileName.indexOf('.');
+    if (dotIndex == -1) {
+      return null;
+    }
+    return fileName.substring(0, dotIndex);
+  }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java
deleted file mode 100644
index f3a20c6a002..00000000000
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java
+++ /dev/null
@@ -1,81 +0,0 @@
-/*
- * 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.index.bucket;
-
-import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
-import org.apache.hudi.common.table.HoodieTableMetaClient;
-import org.apache.hudi.common.table.timeline.HoodieTimeline;
-import org.apache.hudi.common.util.FileIOUtils;
-import org.apache.hudi.common.util.Option;
-import org.apache.hudi.common.util.StringUtils;
-import org.apache.hudi.exception.HoodieIOException;
-import org.apache.hudi.storage.HoodieInstantWriter;
-import org.apache.hudi.storage.HoodieStorage;
-import org.apache.hudi.storage.StoragePath;
-import org.apache.hudi.storage.StoragePathInfo;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.io.InputStream;
-import java.nio.charset.StandardCharsets;
-
-public class SimpleBucketIndexUtils {
-  public static final String INITIAL_HASHING_CONFIG_INSTANT = HoodieTimeline.INIT_INSTANT_TS;
-
-  private static final Logger LOG = LoggerFactory.getLogger(SimpleBucketIndexUtils.class);
-
-  public static boolean initHashingConfig(HoodieTableMetaClient metaClient,
-                                              String expressions,
-                                              String rule,
-                                              int defaultBucketNumber,
-                                              String instant) {
-    if (StringUtils.isNullOrEmpty(expressions)) {
-      return false;
-    }
-    String hashingInstant = StringUtils.isNullOrEmpty(instant) ? INITIAL_HASHING_CONFIG_INSTANT : instant;
-    HoodieStorage storage = metaClient.getStorage();
-    PartitionBucketIndexHashingConfig hashingConfig =
-        new PartitionBucketIndexHashingConfig(expressions, defaultBucketNumber, rule, PartitionBucketIndexHashingConfig.CURRENT_VERSION, hashingInstant);
-    StoragePath hashingConfigPath = new StoragePath(metaClient.getHashingMetadataConfigPath(), hashingConfig.getFilename());
-
-    try {
-      Option<byte []> content = Option.of(hashingConfig.toJsonString().getBytes(StandardCharsets.UTF_8));
-      storage.createImmutableFileInPath(hashingConfigPath, content.map(HoodieInstantWriter::convertByteArrayToWriter));
-    } catch (IOException ioe) {
-      throw new HoodieIOException("Failed to initHashingConfig ", ioe);
-    }
-    return true;
-  }
-
-  public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(HoodieStorage storage,
-                                                                            StoragePathInfo hashingConfig) {
-    if (hashingConfig == null) {
-      return Option.empty();
-    }
-    try (InputStream is = storage.open(hashingConfig.getPath())) {
-      byte[] content = FileIOUtils.readAsByteArray(is);
-      return Option.of(PartitionBucketIndexHashingConfig.fromBytes(content));
-    } catch (IOException e) {
-      LOG.error("Error when loading hashing config, for path: " + hashingConfig.getPath().getName(), e);
-      throw new HoodieIOException("Error while loading hashing config", e);
-    }
-  }
-}
diff --git a/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/index/bucket/TestPartitionBucketIndexCalculator.java b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/index/bucket/TestPartitionBucketIndexCalculator.java
new file mode 100644
index 00000000000..820952ce0e5
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/test/java/org/apache/hudi/index/bucket/TestPartitionBucketIndexCalculator.java
@@ -0,0 +1,144 @@
+/*
+ * 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.index.bucket;
+
+import org.apache.hudi.common.testutils.HoodieCommonTestHarness;
+import org.apache.hudi.exception.HoodieException;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.Arguments;
+
+import java.io.IOException;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class TestPartitionBucketIndexCalculator extends HoodieCommonTestHarness {
+
+  private static final String DEFAULT_RULE = "regex";
+  private static final String DEFAULT_EXPRESSIONS = "\\d{4}-(06-(01|17|18)|11-(01|10|11)),256";
+  private static final int DEFAULT_BUCKET_NUMBER = 10;
+  private PartitionBucketIndexCalculator calc;
+
+  void setUp(String expression, String rule, int defaultBucketNumber) throws IOException {
+    initMetaClient();
+    PartitionBucketIndexUtils.initHashingConfig(metaClient, expression, rule, defaultBucketNumber, null);
+    this.calc = PartitionBucketIndexCalculator.getInstance(PartitionBucketIndexUtils.INITIAL_HASHING_CONFIG_INSTANT, metaClient);
+  }
+
+  /**
+   * Test basic regex rule matching
+   */
+  @ParameterizedTest(name = "Partition {0} should map to bucket {1}")
+  @MethodSource("providePartitionsAndExpectedBuckets")
+  void testBasicRegexRuleMatching(String partition, int expectedBucket) throws IOException {
+    setUp(DEFAULT_EXPRESSIONS, DEFAULT_RULE, DEFAULT_BUCKET_NUMBER);
+    assertEquals(expectedBucket, calc.computeNumBuckets(partition));
+  }
+
+  /**
+   * Provide test data for parameterized tests
+   */
+  private static Stream<Arguments> providePartitionsAndExpectedBuckets() {
+    return Stream.of(
+        // Matching the regex pattern - should get bucket 256
+        Arguments.of("2023-06-01", 256),
+        Arguments.of("2023-06-17", 256),
+        Arguments.of("2023-06-18", 256),
+        Arguments.of("2023-11-01", 256),
+        Arguments.of("2023-11-10", 256),
+        Arguments.of("2023-11-11", 256),
+        Arguments.of("2022-06-01", 256),
+        Arguments.of("2021-11-10", 256),
+
+        // Not matching the regex pattern - should get default bucket 10
+        Arguments.of("2023-06-02", 10),
+        Arguments.of("2023-06-19", 10),
+        Arguments.of("2023-07-01", 10),
+        Arguments.of("2023-11-12", 10),
+        Arguments.of("2023-11-02", 10),
+        Arguments.of("2023-05-01", 10),
+        Arguments.of("not-a-date", 10),
+        Arguments.of("", 10)
+    );
+  }
+
+  @Test
+  void testCaching() throws IOException {
+    setUp(DEFAULT_EXPRESSIONS, DEFAULT_RULE, DEFAULT_BUCKET_NUMBER);
+    // Calculate bucket number for the same partition multiple times
+    String partition = "2023-06-01";
+
+    // First call should calculate and cache
+    int bucket1 = calc.computeNumBuckets(partition);
+    assertEquals(256, bucket1);
+    assertEquals(1, calc.getCacheSize());
+
+    // Second call should use the cache
+    int bucket2 = calc.computeNumBuckets(partition);
+    assertEquals(256, bucket2);
+    assertEquals(1, calc.getCacheSize());
+
+    // Different partition should increase cache size
+    calc.computeNumBuckets("2023-11-10");
+    assertEquals(2, calc.getCacheSize());
+
+    // Clear cache should reset size
+    calc.clearCache();
+    assertEquals(0, calc.getCacheSize());
+  }
+
+  /**
+   * Test multiple regex rules with priority
+   */
+  @Test
+  void testMultipleRegexRulesWithPriority() throws IOException {
+    String expressions = "\\d{4}-06-01,256;\\d{4}-06-\\d{2},128;\\d{4}-11-\\d{2},64";
+    setUp(expressions, DEFAULT_RULE, DEFAULT_BUCKET_NUMBER);
+    PartitionBucketIndexCalculator calc = PartitionBucketIndexCalculator.getInstance(PartitionBucketIndexUtils.INITIAL_HASHING_CONFIG_INSTANT, metaClient);
+    // Setup configuration with multiple rules
+
+    // First rule should take priority
+    assertEquals(256, calc.computeNumBuckets("2023-06-01"));
+
+    // Second rule should match when first doesn't
+    assertEquals(128, calc.computeNumBuckets("2023-06-02"));
+    assertEquals(128, calc.computeNumBuckets("2023-06-30"));
+
+    // Third rule should match for November dates
+    assertEquals(64, calc.computeNumBuckets("2023-11-01"));
+    assertEquals(64, calc.computeNumBuckets("2023-11-30"));
+
+    // Default for non-matching partitions
+    assertEquals(10, calc.computeNumBuckets("2023-07-01"));
+    assertEquals(10, calc.computeNumBuckets("2023-12-25"));
+  }
+
+  /**
+   * Test invalid regex expressions
+   */
+  @Test
+  void testInvalidExpressions() {
+    assertThrows(HoodieException.class, () -> setUp("\\d{4-06-01256", DEFAULT_RULE, DEFAULT_BUCKET_NUMBER));
+    assertThrows(HoodieException.class, () -> setUp("\\d{4}-06-01,a", DEFAULT_RULE, DEFAULT_BUCKET_NUMBER));
+  }
+}
diff --git a/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala b/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
index 7b197727eb9..740e6799348 100644
--- a/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
+++ b/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
@@ -41,14 +41,14 @@ object BucketPartitionUtils {
     val getPartitionKey = getPartitionKeyExtractor()
     val partitioner = new Partitioner {
 
-      private val partitionIndexFunc: Functions.Function2[String, Integer, Integer] =
-        BucketIndexUtil.getPartitionIndexFunc(bucketNum, partitionNum)
+      private val partitionIndexFunc: Functions.Function3[Integer, String, Integer, Integer] =
+        BucketIndexUtil.getPartitionIndexFunc(partitionNum)
 
       override def numPartitions: Int = partitionNum
 
       override def getPartition(value: Any): Int = {
         val partitionKeyPair = value.asInstanceOf[(String, Int)]
-        partitionIndexFunc.apply(partitionKeyPair._1, partitionKeyPair._2)
+        partitionIndexFunc.apply(bucketNum, partitionKeyPair._1, partitionKeyPair._2)
       }
     }
     // use internalRow to avoid extra convert.
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/util/hash/BucketIndexUtil.java b/hudi-common/src/main/java/org/apache/hudi/common/util/hash/BucketIndexUtil.java
index ea3f6a2a12c..8081e65871b 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/util/hash/BucketIndexUtil.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/util/hash/BucketIndexUtil.java
@@ -35,8 +35,8 @@ public class BucketIndexUtil {
    * @param parallelism Parallelism of the task
    * @return The partition index of this bucket.
    */
-  public static Functions.Function2<String, Integer, Integer> getPartitionIndexFunc(int bucketNum, int parallelism) {
-    return (partition, curBucket) -> {
+  public static Functions.Function3<Integer, String, Integer, Integer> getPartitionIndexFunc(int parallelism) {
+    return (bucketNum, partition, curBucket) -> {
       int partitionIndex = (partition.hashCode() & Integer.MAX_VALUE) % parallelism * bucketNum;
       int globalIndex = partitionIndex + curBucket;
       return globalIndex % parallelism;
diff --git a/hudi-common/src/test/java/org/apache/hudi/common/util/hash/TestBucketIndexUtil.java b/hudi-common/src/test/java/org/apache/hudi/common/util/hash/TestBucketIndexUtil.java
index 91c0da003f4..fc2c1fff13a 100644
--- a/hudi-common/src/test/java/org/apache/hudi/common/util/hash/TestBucketIndexUtil.java
+++ b/hudi-common/src/test/java/org/apache/hudi/common/util/hash/TestBucketIndexUtil.java
@@ -73,8 +73,8 @@ public class TestBucketIndexUtil {
   @MethodSource("partitionParams")
   void testPartition(int parallelism, int bucketNumber, boolean partitioned) {
     Map<Integer, Integer> parallelism2TaskCount = new HashMap<>();
-    final Functions.Function2<String, Integer, Integer> partitionIndexFunc =
-        BucketIndexUtil.getPartitionIndexFunc(bucketNumber, parallelism);
+    final Functions.Function3<Integer, String, Integer, Integer> partitionIndexFunc =
+        BucketIndexUtil.getPartitionIndexFunc(parallelism);
     initPartitionData(parallelism2TaskCount, bucketNumber, partitionIndexFunc);
     checkResult(parallelism2TaskCount, parallelism, bucketNumber, partitioned);
   }
@@ -83,8 +83,8 @@ public class TestBucketIndexUtil {
   @MethodSource("noPartitionParams")
   void testNoPartition(int parallelism, int bucketNumber, boolean partitioned) {
     Map<Integer, Integer> parallelism2TaskCount = new HashMap<>();
-    final Functions.Function2<String, Integer, Integer> partitionIndexFunc =
-        BucketIndexUtil.getPartitionIndexFunc(bucketNumber, parallelism);
+    final Functions.Function3<Integer, String, Integer, Integer> partitionIndexFunc =
+        BucketIndexUtil.getPartitionIndexFunc(parallelism);
     initNoPartitionData(parallelism2TaskCount, bucketNumber, partitionIndexFunc);
     checkResult(parallelism2TaskCount, parallelism, bucketNumber, partitioned);
   }
@@ -132,21 +132,21 @@ public class TestBucketIndexUtil {
   }
 
   private void initPartitionData(Map<Integer, Integer> parallelism2TaskCount, int bucketNumber,
-                                 Functions.Function2<String, Integer, Integer> partitionIndexFunc) {
+                                 Functions.Function3<Integer, String, Integer, Integer> partitionIndexFunc) {
     parallelism2TaskCount.clear();
     Arrays.asList("year=2021/month=01/day=01", "year=2021/month=01/day=02", "year=2021/month=01/day=03", "year=2021/month=01/day=04",
         "year=2021/month=01/day=05", "year=2021/month=01/day=06", "year=2021/month=01/day=07", "year=2021/month=01/day=08").forEach(partition -> {
           for (int bucketIndex = 0; bucketIndex < bucketNumber; bucketIndex++) {
-            putIndexCount(parallelism2TaskCount, partitionIndexFunc.apply(partition, bucketIndex));
+            putIndexCount(parallelism2TaskCount, partitionIndexFunc.apply(bucketNumber, partition, bucketIndex));
           }
         });
   }
 
   private void initNoPartitionData(Map<Integer, Integer> parallelism2TaskCount, int bucketNumber,
-                                   Functions.Function2<String, Integer, Integer> partitionIndexFunc) {
+                                   Functions.Function3<Integer, String, Integer, Integer> partitionIndexFunc) {
     parallelism2TaskCount.clear();
     for (int bucketIndex = 0; bucketIndex < bucketNumber; bucketIndex++) {
-      putIndexCount(parallelism2TaskCount, partitionIndexFunc.apply("", bucketIndex));
+      putIndexCount(parallelism2TaskCount, partitionIndexFunc.apply(bucketNumber, "", bucketIndex));
     }
   }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java
index e1c85c123a7..c256cf630ef 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java
@@ -490,6 +490,18 @@ public class FlinkOptions extends HoodieConfig {
           + "numbers (separated by commas).Multiple rules are separated by semicolons like "
           + "hoodie.bucket.index.partition.expressions=expression1,bucket-number1;expression2,bucket-number2");
 
+  public static final ConfigOption<Boolean> BUCKET_INDEX_PARTITION_LEVEL_ENABLE = ConfigOptions
+      .key(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LEVEL_ENABLE.key())
+      .booleanType()
+      .defaultValue(false)
+      .withDescription("Inner config to determine whether to enable partition level bucket index");
+
+  public static final ConfigOption<String> BUCKET_INDEX_PARTITION_LOAD_INSTANT = ConfigOptions
+      .key(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.key())
+      .stringType()
+      .defaultValue(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.defaultValue())
+      .withDescription("Inner config to determine which hashing_config to load when using partition level bucket index");
+
   public static final ConfigOption<String> PARTITION_PATH_FIELD = ConfigOptions
       .key(KeyGeneratorOptions.PARTITIONPATH_FIELD_NAME.key())
       .stringType()
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java
index 388503b5032..66d2e08417e 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java
@@ -18,7 +18,12 @@
 
 package org.apache.hudi.configuration;
 
+import org.apache.hudi.client.HoodieFlinkWriteClient;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.index.HoodieIndex;
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils;
 import org.apache.hudi.util.ClientIds;
+import org.apache.hudi.util.FlinkWriteClients;
 
 import org.apache.flink.configuration.Configuration;
 import org.slf4j.Logger;
@@ -90,4 +95,16 @@ public class OptionsInference {
       }
     }
   }
+
+  public static void setupIndexConfigs(Configuration conf) {
+    HoodieIndex.BucketIndexEngineType engineType = OptionsResolver.getBucketEngineType(conf);
+    if (engineType.equals(HoodieIndex.BucketIndexEngineType.SIMPLE) &&
+        PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(HadoopConfigurations.getHadoopConf(conf))) {
+      try (HoodieFlinkWriteClient writeClient = FlinkWriteClients.createWriteClientV2(conf)) {
+        HoodieTableMetaClient metaClient = writeClient.getHoodieTable().getMetaClient();
+        String hashingConfigToLoad = PartitionBucketIndexUtils.getHashingConfigInstantToLoad(metaClient);
+        conf.set(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT, hashingConfigToLoad);
+      }
+    }
+  }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java
index 7685bffb468..8369a0b768e 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java
@@ -18,13 +18,17 @@
 
 package org.apache.hudi.sink.bucket;
 
+import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.index.bucket.BucketIdentifier;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
 import org.apache.hudi.io.storage.row.HoodieRowDataCreateHandle;
 import org.apache.hudi.sink.bulk.BulkInsertWriterHelper;
 import org.apache.hudi.sink.bulk.RowDataKeyGen;
 import org.apache.hudi.sink.bulk.sort.SortOperatorGen;
 import org.apache.hudi.table.HoodieTable;
+import org.apache.hudi.util.Lazy;
 
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.table.api.DataTypes;
@@ -92,16 +96,25 @@ public class BucketBulkInsertWriterHelper extends BulkInsertWriterHelper {
     return new SortOperatorGen(rowType, new String[] {FILE_GROUP_META_FIELD});
   }
 
-  private static String getFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, String indexKeys, int numBuckets, boolean needFixedFileIdSuffix) {
+  private static String getFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, String indexKeys, Configuration conf, boolean needFixedFileIdSuffix,
+                                  Lazy<org.apache.hadoop.conf.Configuration> hadoopConf) {
     String recordKey = keyGen.getRecordKey(record);
     String partition = keyGen.getPartitionPath(record);
+    final int numBuckets;
+    if (!StringUtils.isNullOrEmpty(conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT))) {
+      numBuckets = PartitionBucketIndexCalculator.getInstance(conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT), hadoopConf.get(), conf.get(FlinkOptions.PATH))
+          .computeNumBuckets(partition);
+    } else {
+      numBuckets = conf.get(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
+    }
     final int bucketNum = BucketIdentifier.getBucketId(recordKey, indexKeys, numBuckets);
     String bucketId = partition + bucketNum;
     return bucketIdToFileId.computeIfAbsent(bucketId, k -> needFixedFileIdSuffix ? BucketIdentifier.newBucketFileIdForNBCC(bucketNum) : BucketIdentifier.newBucketFileIdPrefix(bucketNum));
   }
 
-  public static RowData rowWithFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, String indexKeys, int numBuckets, boolean needFixedFileIdSuffix) {
-    final String fileId = getFileId(bucketIdToFileId, keyGen, record, indexKeys, numBuckets, needFixedFileIdSuffix);
+  public static RowData rowWithFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, String indexKeys, Configuration conf, boolean needFixedFileIdSuffix,
+                                      Lazy<org.apache.hadoop.conf.Configuration> hadoopConf) {
+    final String fileId = getFileId(bucketIdToFileId, keyGen, record, indexKeys, conf, needFixedFileIdSuffix, hadoopConf);
     return GenericRowData.of(StringData.fromString(fileId), record);
   }
 
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java
index cf13ebfc0b4..2a3e1e154dd 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java
@@ -21,10 +21,12 @@ package org.apache.hudi.sink.bucket;
 import org.apache.hudi.client.model.HoodieFlinkInternalRow;
 import org.apache.hudi.common.model.HoodieRecordLocation;
 import org.apache.hudi.common.util.Functions;
+import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.common.util.hash.BucketIndexUtil;
 import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.configuration.OptionsResolver;
 import org.apache.hudi.index.bucket.BucketIdentifier;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
 import org.apache.hudi.sink.StreamWriteFunction;
 
 import org.apache.flink.configuration.Configuration;
@@ -54,8 +56,6 @@ public class BucketStreamWriteFunction extends StreamWriteFunction {
 
   private int parallelism;
 
-  private int bucketNum;
-
   private String indexKeyFields;
 
   private boolean isNonBlockingConcurrencyControl;
@@ -76,13 +76,16 @@ public class BucketStreamWriteFunction extends StreamWriteFunction {
   /**
    * Functions for calculating the task partition to dispatch.
    */
-  private Functions.Function2<String, Integer, Integer> partitionIndexFunc;
+  private Functions.Function3<Integer, String, Integer, Integer> partitionIndexFunc;
 
   /**
    * To prevent strings compare for each record, define this only during open()
    */
   private boolean isInsertOverwrite;
 
+  private PartitionBucketIndexCalculator calc;
+  private boolean isPartitionLevelBucketIndexEnable = false;
+
   /**
    * Constructs a BucketStreamWriteFunction.
    *
@@ -95,15 +98,19 @@ public class BucketStreamWriteFunction extends StreamWriteFunction {
   @Override
   public void open(Configuration parameters) throws IOException {
     super.open(parameters);
-    this.bucketNum = config.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
     this.indexKeyFields = OptionsResolver.getIndexKeyField(config);
     this.isNonBlockingConcurrencyControl = OptionsResolver.isNonBlockingConcurrencyControl(config);
     this.taskID = getRuntimeContext().getIndexOfThisSubtask();
     this.parallelism = getRuntimeContext().getNumberOfParallelSubtasks();
     this.bucketIndex = new HashMap<>();
     this.incBucketIndex = new HashSet<>();
-    this.partitionIndexFunc = BucketIndexUtil.getPartitionIndexFunc(bucketNum, parallelism);
+    this.partitionIndexFunc = BucketIndexUtil.getPartitionIndexFunc(parallelism);
     this.isInsertOverwrite = OptionsResolver.isInsertOverwrite(config);
+    String hashingInstant = config.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT);
+    this.isPartitionLevelBucketIndexEnable = !StringUtils.isNullOrEmpty(hashingInstant);
+    if (isPartitionLevelBucketIndexEnable) {
+      calc = PartitionBucketIndexCalculator.getInstance(hashingInstant, metaClient);
+    }
   }
 
   @Override
@@ -132,7 +139,10 @@ public class BucketStreamWriteFunction extends StreamWriteFunction {
       bootstrapIndexIfNeed(partition);
     }
     Map<Integer, String> bucketToFileId = bucketIndex.computeIfAbsent(partition, p -> new HashMap<>());
-    final int bucketNum = BucketIdentifier.getBucketId(record.getRecordKey(), indexKeyFields, this.bucketNum);
+
+    int numBuckets = getOrComputeNumBuckets(partition);
+
+    final int bucketNum = BucketIdentifier.getBucketId(record.getRecordKey(), indexKeyFields, numBuckets);
     final String bucketId = partition + "/" + bucketNum;
 
     final HoodieRecordLocation location;
@@ -156,7 +166,12 @@ public class BucketStreamWriteFunction extends StreamWriteFunction {
    * partitionIndex == this taskID belongs to this task.
    */
   public boolean isBucketToLoad(int bucketNumber, String partition) {
-    return this.partitionIndexFunc.apply(partition, bucketNumber) == taskID;
+    int numBuckets = getOrComputeNumBuckets(partition);
+    return this.partitionIndexFunc.apply(numBuckets, partition, bucketNumber) == taskID;
+  }
+
+  private int getOrComputeNumBuckets(String partition) {
+    return isPartitionLevelBucketIndexEnable ? calc.computeNumBuckets(partition) : config.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
   }
 
   /**
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java
index 53ec4def01e..ecf4656d7ee 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java
@@ -20,10 +20,15 @@ package org.apache.hudi.sink.partitioner;
 
 import org.apache.hudi.common.model.HoodieKey;
 import org.apache.hudi.common.util.Functions;
+import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.common.util.hash.BucketIndexUtil;
+import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.index.bucket.BucketIdentifier;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
+import org.apache.hudi.util.Lazy;
 
 import org.apache.flink.api.common.functions.Partitioner;
+import org.apache.flink.configuration.Configuration;
 
 /**
  * Bucket index input partitioner.
@@ -33,22 +38,36 @@ import org.apache.flink.api.common.functions.Partitioner;
  */
 public class BucketIndexPartitioner<T extends HoodieKey> implements Partitioner<T> {
 
-  private final int bucketNum;
+  private final Configuration conf;
   private final String indexKeyFields;
+  private final String hashingInstantToLoad;
+  private org.apache.hadoop.conf.Configuration hadoopConf;
 
-  private Functions.Function2<String, Integer, Integer> partitionIndexFunc;
+  private Functions.Function3<Integer, String, Integer, Integer> partitionIndexFunc;
 
-  public BucketIndexPartitioner(int bucketNum, String indexKeyFields) {
-    this.bucketNum = bucketNum;
+  public BucketIndexPartitioner(Configuration conf, String indexKeyFields, Lazy<org.apache.hadoop.conf.Configuration> hadoopConf) {
+    this.conf = conf;
     this.indexKeyFields = indexKeyFields;
+    this.hashingInstantToLoad = conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT);
+    if (!StringUtils.isNullOrEmpty(hashingInstantToLoad)) {
+      this.hadoopConf = hadoopConf.get();
+    }
   }
 
   @Override
   public int partition(HoodieKey key, int numPartitions) {
     if (this.partitionIndexFunc == null) {
-      this.partitionIndexFunc = BucketIndexUtil.getPartitionIndexFunc(bucketNum, numPartitions);
+      this.partitionIndexFunc = BucketIndexUtil.getPartitionIndexFunc(numPartitions);
+    }
+
+    int bucketNum;
+    if (!StringUtils.isNullOrEmpty(hashingInstantToLoad)) {
+      PartitionBucketIndexCalculator calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, hadoopConf, conf.get(FlinkOptions.PATH));
+      bucketNum = calc.computeNumBuckets(key.getPartitionPath());
+    } else {
+      bucketNum = conf.get(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
     }
     int curBucket = BucketIdentifier.getBucketId(key.getRecordKey(), indexKeyFields, bucketNum);
-    return this.partitionIndexFunc.apply(key.getPartitionPath(), curBucket);
+    return this.partitionIndexFunc.apply(bucketNum, key.getPartitionPath(), curBucket);
   }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java
index 66269af4b96..d604567246d 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java
@@ -23,6 +23,7 @@ import org.apache.hudi.client.model.HoodieFlinkInternalRowTypeInfo;
 import org.apache.hudi.common.model.HoodieKey;
 import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.configuration.HadoopConfigurations;
 import org.apache.hudi.configuration.OptionsResolver;
 import org.apache.hudi.exception.HoodieException;
 import org.apache.hudi.exception.HoodieNotSupportedException;
@@ -53,6 +54,7 @@ import org.apache.hudi.sink.partitioner.BucketAssignFunction;
 import org.apache.hudi.sink.partitioner.BucketIndexPartitioner;
 import org.apache.hudi.sink.transform.RowDataToHoodieFunctions;
 import org.apache.hudi.table.format.FilePathUtils;
+import org.apache.hudi.util.Lazy;
 
 import org.apache.flink.api.common.functions.Partitioner;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
@@ -120,9 +122,8 @@ public class Pipelines {
             "Consistent hashing bucket index does not work with bulk insert using FLINK engine. Use simple bucket index or Spark engine.");
       }
       String indexKeys = OptionsResolver.getIndexKeyField(conf);
-      int numBuckets = conf.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
-
-      BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(numBuckets, indexKeys);
+      Lazy<org.apache.hadoop.conf.Configuration> hadoopConf = Lazy.lazily(() -> HadoopConfigurations.getHadoopConf(conf));
+      BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(conf, indexKeys, hadoopConf);
       RowDataKeyGen keyGen = RowDataKeyGen.instance(conf, rowType);
       RowType rowTypeWithFileId = BucketBulkInsertWriterHelper.rowTypeWithFileId(rowType);
       InternalTypeInfo<RowData> typeInfo = InternalTypeInfo.of(rowTypeWithFileId);
@@ -130,7 +131,7 @@ public class Pipelines {
 
       Map<String, String> bucketIdToFileId = new HashMap<>();
       dataStream = dataStream.partitionCustom(partitioner, keyGen::getHoodieKey)
-          .map(record -> BucketBulkInsertWriterHelper.rowWithFileId(bucketIdToFileId, keyGen, record, indexKeys, numBuckets, needFixedFileIdSuffix), typeInfo)
+          .map(record -> BucketBulkInsertWriterHelper.rowWithFileId(bucketIdToFileId, keyGen, record, indexKeys, conf, needFixedFileIdSuffix, hadoopConf), typeInfo)
           .setParallelism(PARALLELISM_VALUE);
       if (conf.getBoolean(FlinkOptions.WRITE_BULK_INSERT_SORT_INPUT)) {
         SortOperatorGen sortOperatorGen = BucketBulkInsertWriterHelper.getFileIdSorterGen(rowTypeWithFileId);
@@ -337,11 +338,10 @@ public class Pipelines {
       HoodieIndex.BucketIndexEngineType bucketIndexEngineType = OptionsResolver.getBucketEngineType(conf);
       switch (bucketIndexEngineType) {
         case SIMPLE:
-          int bucketNum = conf.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
           String indexKeyFields = OptionsResolver.getIndexKeyField(conf);
           // [HUDI-9036] BucketIndexPartitioner is also used in bulk insert mode,
           // keep use of HoodieKey here in partitionCustom for now
-          BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(bucketNum, indexKeyFields);
+          BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(conf, indexKeyFields, Lazy.lazily(() -> HadoopConfigurations.getHadoopConf(conf)));
           return dataStream
               .partitionCustom(
                   partitioner,
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSink.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSink.java
index 51ceec2eefb..4f96515a64b 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSink.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/HoodieTableSink.java
@@ -82,6 +82,8 @@ public class HoodieTableSink implements
       OptionsInference.setupSinkTasks(conf, dataStream.getExecutionConfig().getParallelism());
       // set up client id
       OptionsInference.setupClientId(conf);
+      // set up index related configs
+      OptionsInference.setupIndexConfigs(conf);
 
       RowType rowType = (RowType) schema.toSinkRowDataType().notNull().getLogicalType();
 
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java
index db1d8364ad3..32f7ebfa0e0 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java
@@ -27,7 +27,7 @@ import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.configuration.HadoopConfigurations;
 import org.apache.hudi.exception.HoodieCatalogException;
-import org.apache.hudi.index.bucket.SimpleBucketIndexUtils;
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils;
 import org.apache.hudi.internal.schema.InternalSchema;
 import org.apache.hudi.internal.schema.Type;
 import org.apache.hudi.internal.schema.convert.AvroInternalSchemaConverter;
@@ -310,6 +310,6 @@ public class HoodieCatalogUtil {
     int bucketNumber = options.containsKey(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.key()) ?
         Integer.valueOf(options.get(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.key())) : FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.defaultValue();
 
-    return SimpleBucketIndexUtils.initHashingConfig(metaClient, expressions, rule, bucketNumber, null);
+    return PartitionBucketIndexUtils.initHashingConfig(metaClient, expressions, rule, bucketNumber, null);
   }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java
index 1dec40ceba0..a4da9049dce 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java
@@ -21,6 +21,7 @@ package org.apache.hudi.sink.utils;
 import org.apache.hudi.adapter.CollectOutputAdapter;
 import org.apache.hudi.adapter.TestStreamConfigs;
 import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.configuration.HadoopConfigurations;
 import org.apache.hudi.configuration.OptionsResolver;
 import org.apache.hudi.exception.HoodieException;
 import org.apache.hudi.sink.StreamWriteOperatorCoordinator;
@@ -32,6 +33,7 @@ import org.apache.hudi.sink.bulk.sort.SortOperatorGen;
 import org.apache.hudi.sink.common.AbstractWriteFunction;
 import org.apache.hudi.sink.event.WriteMetadataEvent;
 import org.apache.hudi.util.AvroSchemaConverter;
+import org.apache.hudi.util.Lazy;
 import org.apache.hudi.util.StreamerUtil;
 
 import org.apache.flink.api.common.ExecutionConfig;
@@ -218,7 +220,8 @@ public class BulkInsertFunctionWrapper<I> implements TestFunctionWrapper<I> {
     int numBuckets = conf.getInteger(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
     boolean needFixedFileIdSuffix = OptionsResolver.isNonBlockingConcurrencyControl(conf);
     this.bucketIdToFileId = new HashMap<>();
-    this.mapFunction = r -> BucketBulkInsertWriterHelper.rowWithFileId(bucketIdToFileId, keyGen, r, indexKeys, numBuckets, needFixedFileIdSuffix);
+    this.mapFunction = r -> BucketBulkInsertWriterHelper.rowWithFileId(bucketIdToFileId, keyGen, r, indexKeys, conf, needFixedFileIdSuffix,
+        Lazy.lazily(() -> HadoopConfigurations.getHadoopConf(conf)));
   }
 
   private void setupSortOperator() throws Exception {
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
index 84e982df8f3..f5bbc420b9a 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
@@ -31,7 +31,7 @@ import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.configuration.HadoopConfigurations;
 import org.apache.hudi.exception.HoodieIOException;
 import org.apache.hudi.exception.HoodieValidationException;
-import org.apache.hudi.index.bucket.SimpleBucketIndexUtils;
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils;
 import org.apache.hudi.keygen.ComplexAvroKeyGenerator;
 import org.apache.hudi.keygen.NonpartitionedAvroKeyGenerator;
 import org.apache.hudi.keygen.SimpleAvroKeyGenerator;
@@ -340,7 +340,7 @@ public class TestHoodieCatalog {
   }
 
   @Test
-  void testCreateTableWithPartitionSimpleBucketIndex() throws TableAlreadyExistException, DatabaseNotExistException, IOException {
+  void testCreateTableWithPartitionBucketIndex() throws TableAlreadyExistException, DatabaseNotExistException, IOException {
     String rule = "regex";
     String expressions = "\\d{4}-(06-(01|17|18)|11-(01|10|11)),256";
     String defaultBucketNumber = "20";
@@ -361,9 +361,9 @@ public class TestHoodieCatalog {
             new HadoopStorageConfiguration(HadoopConfigurations.getHadoopConf(flinkConf)), tablePathStr);
     HoodieStorage storage = metaClient.getStorage();
     StoragePath initialHashingConfig =
-        new StoragePath(metaClient.getHashingMetadataConfigPath(), SimpleBucketIndexUtils.INITIAL_HASHING_CONFIG_INSTANT + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
+        new StoragePath(metaClient.getHashingMetadataConfigPath(), PartitionBucketIndexUtils.INITIAL_HASHING_CONFIG_INSTANT + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
     StoragePathInfo info = storage.getPathInfo(initialHashingConfig);
-    Option<PartitionBucketIndexHashingConfig> hashingConfig = SimpleBucketIndexUtils.loadHashingConfig(storage, info);
+    Option<PartitionBucketIndexHashingConfig> hashingConfig = PartitionBucketIndexUtils.loadHashingConfig(storage, info);
     assertTrue(hashingConfig.isPresent());
     assertEquals(hashingConfig.get().getDefaultBucketNumber(), Integer.parseInt(defaultBucketNumber));
     assertEquals(hashingConfig.get().getRule(), rule);
diff --git a/style/checkstyle-suppressions.xml b/style/checkstyle-suppressions.xml
index 8f6697b4083..c349443bccc 100644
--- a/style/checkstyle-suppressions.xml
+++ b/style/checkstyle-suppressions.xml
@@ -26,6 +26,7 @@
   <!-- Member Names expected to start with "_"  -->
   <suppress checks="naming" files="TestRecord.java" lines="1-9999"/>
   <suppress checks="IllegalImport" files="Option.java" />
+  <suppress checks="IllegalImport" files="PartitionBucketIndexCalculator.java" />
   <!-- java.util.Optional part of DataSource V2 API signature -->
   <suppress checks="IllegalImport" files="DefaultSource.java" />
   <suppress checks="IllegalImport" files="HoodieTableSource.java" />


(hudi) 05/06: finish Spark call update hashing config need more test

Posted by zh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zhangyue19921010 pushed a commit to branch HUDI-8990
in repository https://gitbox.apache.org/repos/asf/hudi.git

commit 184231de7b02e33baad0d6ee5b08f2075f0e7f9c
Author: YueZhang <zh...@163.com>
AuthorDate: Thu Mar 20 16:58:36 2025 +0800

    finish Spark call update hashing config need more test
---
 .../org/apache/hudi/config/HoodieWriteConfig.java  |   8 +
 .../bucket/PartitionBucketIndexCalculator.java     |  24 +-
 .../index/bucket/PartitionBucketIndexUtils.java    |  15 +-
 .../BucketIndexBulkInsertPartitionerWithRows.java  |   9 +-
 .../RDDSimpleBucketBulkInsertPartitioner.java      |   2 +-
 .../BucketBulkInsertDataInternalWriterHelper.java  |   2 +-
 .../action/commit/SparkBucketIndexPartitioner.java |   2 +-
 .../model/PartitionBucketIndexHashingConfig.java   |   4 +
 .../hudi/common/model/WriteOperationType.java      |   1 +
 .../DatasetBucketRescaleCommitActionExecutor.java  |  77 ++++++
 .../org/apache/hudi/HoodieSparkSqlWriter.scala     |  12 +-
 .../spark/sql/hudi/ProvidesHoodieConfig.scala      |  24 +-
 .../procedures/PartitionBucketIndexManager.scala   | 289 +++++++++++++++++++++
 .../sql/hudi/common/HoodieSparkSqlTestBase.scala   |   1 +
 .../spark/sql/hudi/dml/TestInsertTable.scala       |  88 +++++++
 .../TestPartitionBucketIndexManager.scala          |  31 +++
 16 files changed, 575 insertions(+), 14 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java
index 557f8672e19..623f02e623c 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java
@@ -1959,6 +1959,14 @@ public class HoodieWriteConfig extends HoodieConfig {
     return getString(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT);
   }
 
+  public String getBucketIndexPartitionExpression() {
+    return getString(HoodieIndexConfig.BUCKET_INDEX_PARTITION_EXPRESSIONS);
+  }
+
+  public String getBucketIndexPartitionRuleType() {
+    return getString(HoodieIndexConfig.BUCKET_INDEX_PARTITION_RULE_TYPE);
+  }
+
   public String getIndexClass() {
     return getString(HoodieIndexConfig.INDEX_CLASS_NAME);
   }
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
index 65a08363402..acd48bb9961 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
@@ -37,6 +37,7 @@ import org.slf4j.LoggerFactory;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -55,7 +56,7 @@ public class PartitionBucketIndexCalculator implements Serializable {
   private static final int CACHE_SIZE = 100_000;
   private PartitionBucketIndexHashingConfig hashingConfig;
   private int defaultBucketNumber;
-  private final String instantToLoad;
+  private String instantToLoad;
   // Cache for partition to bucket number mapping
   @SuppressWarnings("unchecked")
   private final Map<String, Integer> partitionToBucketCache = new LRUMap(CACHE_SIZE);
@@ -89,6 +90,14 @@ public class PartitionBucketIndexCalculator implements Serializable {
     }
   }
 
+  private PartitionBucketIndexCalculator(String instantToLoad, PartitionBucketIndexHashingConfig config) {
+    this.hashingConfig = config;
+    this.defaultBucketNumber = config.getDefaultBucketNumber();
+    String expressions = config.getExpressions();
+    String ruleType = config.getRule();
+    this.ruleEngine = createRuleEngine(ruleType, expressions);
+  }
+
   private void init(HoodieStorage storage, StoragePath hashingConfigPath) {
     Option<PartitionBucketIndexHashingConfig> config = PartitionBucketIndexUtils.loadHashingConfig(storage, hashingConfigPath);
     ValidationUtils.checkArgument(config.isPresent());
@@ -124,6 +133,15 @@ public class PartitionBucketIndexCalculator implements Serializable {
         });
   }
 
+  public static PartitionBucketIndexCalculator getInstance(String instantToLoad, PartitionBucketIndexHashingConfig config) {
+    // Using instantToLoad as the key for the cache
+    return INSTANCES.computeIfAbsent(instantToLoad,
+        key -> {
+          LOG.info("Creating new PartitionBucketIndexCalculator instance for instantToLoad: {}", key);
+          return new PartitionBucketIndexCalculator(key, config);
+        });
+  }
+
   /**
    * Computes the bucket number for a given partition path
    *
@@ -182,6 +200,10 @@ public class PartitionBucketIndexCalculator implements Serializable {
     return partitionToBucketCache.size();
   }
 
+  public Map<String, Integer> getPartitionToBucket() {
+    return new HashMap<>(partitionToBucketCache);
+  }
+
   public void clearCache() {
     partitionToBucketCache.clear();
     LOG.info("Cleared partition to bucket number cache");
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
index 69e29bd92e3..ad1e1d0b3bf 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
@@ -42,6 +42,7 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 
 public class PartitionBucketIndexUtils {
@@ -77,11 +78,14 @@ public class PartitionBucketIndexUtils {
       return false;
     }
     String hashingInstant = StringUtils.isNullOrEmpty(instant) ? INITIAL_HASHING_CONFIG_INSTANT : instant;
-    HoodieStorage storage = metaClient.getStorage();
     PartitionBucketIndexHashingConfig hashingConfig =
         new PartitionBucketIndexHashingConfig(expressions, defaultBucketNumber, rule, PartitionBucketIndexHashingConfig.CURRENT_VERSION, hashingInstant);
-    StoragePath hashingConfigPath = new StoragePath(metaClient.getHashingMetadataConfigPath(), hashingConfig.getFilename());
+    return saveHashingConfig(hashingConfig, metaClient);
+  }
 
+  public static boolean saveHashingConfig(PartitionBucketIndexHashingConfig hashingConfig, HoodieTableMetaClient metaClient) {
+    StoragePath hashingConfigPath = new StoragePath(metaClient.getHashingMetadataConfigPath(), hashingConfig.getFilename());
+    HoodieStorage storage = metaClient.getStorage();
     try {
       Option<byte []> content = Option.of(hashingConfig.toJsonString().getBytes(StandardCharsets.UTF_8));
       storage.createImmutableFileInPath(hashingConfigPath, content.map(HoodieInstantWriter::convertByteArrayToWriter));
@@ -153,4 +157,11 @@ public class PartitionBucketIndexUtils {
     }
     return hashingConfigName.substring(0, dotIndex);
   }
+
+  public static Map<String, Integer> getAllBucketNumbers(PartitionBucketIndexCalculator calc, List<String> partitions) {
+    for (String partition : partitions) {
+      calc.computeNumBuckets(partition);
+    }
+    return calc.getPartitionToBucket();
+  }
 }
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
index e80c8bfac0e..3a94fa78cb3 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
@@ -41,12 +41,19 @@ public class BucketIndexBulkInsertPartitionerWithRows implements BulkInsertParti
     this.indexKeyFields = indexKeyFields;
     this.bucketNum = bucketNum;
     String hashingInstantToLoad = table.getConfig().getHashingConfigInstantToLoad();
-    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    this.isPartitionBucketIndexEnable = StringUtils.nonEmpty(hashingInstantToLoad);
     if (isPartitionBucketIndexEnable) {
       calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, table.getMetaClient());
     }
   }
 
+  public BucketIndexBulkInsertPartitionerWithRows(String indexKeyFields, int bucketNum, PartitionBucketIndexCalculator calc) {
+    this.indexKeyFields = indexKeyFields;
+    this.bucketNum = bucketNum;
+    this.isPartitionBucketIndexEnable = true;
+    this.calc = calc;
+  }
+
   @Override
   public Dataset<Row> repartitionRecords(Dataset<Row> rows, int outputPartitions) {
     return BucketPartitionUtils$.MODULE$.createDataFrame(rows, indexKeyFields, bucketNum, outputPartitions, calc, isPartitionBucketIndexEnable);
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
index d651342c368..393f0054398 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
@@ -49,7 +49,7 @@ public class RDDSimpleBucketBulkInsertPartitioner<T extends HoodieRecordPayload>
     ValidationUtils.checkArgument(table.getIndex() instanceof HoodieSimpleBucketIndex);
     this.isNonBlockingConcurrencyControl = table.getConfig().isNonBlockingConcurrencyControl();
     String hashingInstantToLoad = table.getConfig().getHashingConfigInstantToLoad();
-    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    this.isPartitionBucketIndexEnable = StringUtils.nonEmpty(hashingInstantToLoad);
     if (isPartitionBucketIndexEnable) {
       calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, table.getMetaClient());
     }
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java
index d2b6a04c2e1..7f55c02b851 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java
@@ -71,7 +71,7 @@ public class BucketBulkInsertDataInternalWriterHelper extends BulkInsertDataInte
     this.handles = new HashMap<>();
     this.isNonBlockingConcurrencyControl = writeConfig.isNonBlockingConcurrencyControl();
     String hashingInstantToLoad = writeConfig.getHashingConfigInstantToLoad();
-    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    this.isPartitionBucketIndexEnable = StringUtils.nonEmpty(hashingInstantToLoad);
     if (isPartitionBucketIndexEnable) {
       calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, hoodieTable.getMetaClient());
     }
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java
index c760d2aabdd..b099392e9e1 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java
@@ -99,7 +99,7 @@ public class SparkBucketIndexPartitioner<T> extends SparkHoodiePartitioner<T> {
               + table.getIndex().getClass().getSimpleName());
     }
     String hashingInstantToLoad = table.getConfig().getHashingConfigInstantToLoad();
-    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    this.isPartitionBucketIndexEnable = StringUtils.nonEmpty(hashingInstantToLoad);
     if (isPartitionBucketIndexEnable) {
       calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, table.getMetaClient());
     }
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 46b38896925..a658003f851 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
@@ -62,6 +62,10 @@ public class PartitionBucketIndexHashingConfig implements Serializable {
     return JsonUtils.getObjectMapper().writerWithDefaultPrettyPrinter().writeValueAsString(this);
   }
 
+  public String getInstant() {
+    return this.instant;
+  }
+
   public static <T> T fromJsonString(String jsonStr, Class<T> clazz) throws Exception {
     if (jsonStr == null || jsonStr.isEmpty()) {
       // For empty commit file (no data or somethings bad happen).
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java b/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java
index 1036e8fca44..bc82a04023c 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java
@@ -41,6 +41,7 @@ public enum WriteOperationType {
   BOOTSTRAP("bootstrap"),
   // insert overwrite with static partitioning
   INSERT_OVERWRITE("insert_overwrite"),
+  BUCKET_RESCALE("bucket_rescale"),
   // cluster
   CLUSTER("cluster"),
   // delete partition
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/DatasetBucketRescaleCommitActionExecutor.java b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/DatasetBucketRescaleCommitActionExecutor.java
new file mode 100644
index 00000000000..1c6d90eb6f7
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/DatasetBucketRescaleCommitActionExecutor.java
@@ -0,0 +1,77 @@
+/*
+ * 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.commit;
+
+import org.apache.hudi.client.SparkRDDWriteClient;
+import org.apache.hudi.client.WriteStatus;
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
+import org.apache.hudi.common.util.ValidationUtils;
+import org.apache.hudi.config.HoodieWriteConfig;
+import org.apache.hudi.execution.bulkinsert.BucketIndexBulkInsertPartitionerWithRows;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils;
+import org.apache.hudi.table.BulkInsertPartitioner;
+import org.apache.hudi.table.action.HoodieWriteMetadata;
+
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+
+public class DatasetBucketRescaleCommitActionExecutor extends DatasetBulkInsertOverwriteCommitActionExecutor {
+
+  private final PartitionBucketIndexCalculator calc;
+
+  public DatasetBucketRescaleCommitActionExecutor(HoodieWriteConfig config,
+                                                  SparkRDDWriteClient writeClient,
+                                                  String instantTime) {
+    super(config, writeClient, instantTime);
+    String expression = config.getBucketIndexPartitionExpression();
+    String instant = config.getHashingConfigInstantToLoad();
+    String rule = config.getBucketIndexPartitionRuleType();
+    int bucketNumber = config.getBucketIndexNumBuckets();
+    PartitionBucketIndexHashingConfig hashingConfig = new PartitionBucketIndexHashingConfig(expression,
+        bucketNumber, rule, PartitionBucketIndexHashingConfig.CURRENT_VERSION, instant);
+    this.calc = PartitionBucketIndexCalculator.getInstance(instantTime, hashingConfig);
+  }
+
+  /**
+   * Create BulkInsertPartitioner with prepared PartitionBucketIndexCalculator
+   * @param populateMetaFields
+   * @param isTablePartitioned
+   * @return BulkInsertPartitioner
+   */
+  @Override
+  protected BulkInsertPartitioner<Dataset<Row>> getPartitioner(boolean populateMetaFields, boolean isTablePartitioned) {
+    return new BucketIndexBulkInsertPartitionerWithRows(writeConfig.getBucketIndexHashFieldWithDefault(),
+        writeConfig.getBucketIndexNumBuckets(), calc);
+  }
+
+  /**
+   * create new hashing_config during afterExecute and before commit finished
+   * @param result
+   */
+  @Override
+  protected void afterExecute(HoodieWriteMetadata<JavaRDD<WriteStatus>> result) {
+    super.afterExecute(result);
+
+    PartitionBucketIndexHashingConfig hashingConfig = calc.getHashingConfig();
+    boolean res = PartitionBucketIndexUtils.saveHashingConfig(hashingConfig, table.getMetaClient());
+    ValidationUtils.checkArgument(res);
+  }
+}
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala
index 58cfebf596d..7c59e433c8a 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala
@@ -30,7 +30,7 @@ import org.apache.hudi.avro.AvroSchemaUtils.resolveNullableSchema
 import org.apache.hudi.avro.HoodieAvroUtils
 import org.apache.hudi.client.{HoodieWriteResult, SparkRDDWriteClient}
 import org.apache.hudi.client.common.HoodieSparkEngineContext
-import org.apache.hudi.commit.{DatasetBulkInsertCommitActionExecutor, DatasetBulkInsertOverwriteCommitActionExecutor, DatasetBulkInsertOverwriteTableCommitActionExecutor}
+import org.apache.hudi.commit.{DatasetBucketRescaleCommitActionExecutor, DatasetBulkInsertCommitActionExecutor, DatasetBulkInsertOverwriteCommitActionExecutor, DatasetBulkInsertOverwriteTableCommitActionExecutor}
 import org.apache.hudi.common.config._
 import org.apache.hudi.common.engine.HoodieEngineContext
 import org.apache.hudi.common.fs.FSUtils
@@ -40,7 +40,7 @@ import org.apache.hudi.common.model.HoodieTableType.{COPY_ON_WRITE, MERGE_ON_REA
 import org.apache.hudi.common.table.{HoodieTableConfig, HoodieTableMetaClient, HoodieTableVersion, TableSchemaResolver}
 import org.apache.hudi.common.table.log.block.HoodieLogBlock.HoodieLogBlockType
 import org.apache.hudi.common.table.timeline.HoodieInstantTimeGenerator
-import org.apache.hudi.common.util.{CommitUtils, Option => HOption, StringUtils}
+import org.apache.hudi.common.util.{CommitUtils, StringUtils, Option => HOption}
 import org.apache.hudi.common.util.ConfigUtils.getAllConfigKeys
 import org.apache.hudi.config.{HoodieCompactionConfig, HoodieInternalConfig, HoodieWriteConfig}
 import org.apache.hudi.config.HoodieBootstrapConfig.{BASE_PATH, INDEX_CLASS_NAME}
@@ -62,7 +62,6 @@ import org.apache.hudi.sync.common.util.SyncUtilHelpers
 import org.apache.hudi.sync.common.util.SyncUtilHelpers.getHoodieMetaSyncException
 import org.apache.hudi.util.{SparkConfigUtils, SparkKeyGenUtils}
 import org.apache.hudi.util.SparkConfigUtils.getStringWithAltKeys
-
 import org.apache.avro.Schema
 import org.apache.avro.generic.GenericData
 import org.apache.hadoop.conf.Configuration
@@ -81,7 +80,6 @@ import org.apache.spark.sql.types.StructType
 import org.slf4j.LoggerFactory
 
 import java.util.function.BiConsumer
-
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 import scala.util.{Failure, Success, Try}
@@ -815,7 +813,7 @@ class HoodieSparkSqlWriterInternal {
     val overwriteOperationType = Option(hoodieConfig.getString(HoodieInternalConfig.BULKINSERT_OVERWRITE_OPERATION_TYPE))
       .map(WriteOperationType.fromValue)
       .orNull
-    val instantTime = writeClient.createNewInstantTime()
+    var instantTime = writeClient.createNewInstantTime()
     val executor = mode match {
       case _ if overwriteOperationType == null =>
         // Don't need to overwrite
@@ -823,6 +821,10 @@ class HoodieSparkSqlWriterInternal {
       case SaveMode.Append if overwriteOperationType == WriteOperationType.INSERT_OVERWRITE =>
         // INSERT OVERWRITE PARTITION uses Append mode
         new DatasetBulkInsertOverwriteCommitActionExecutor(writeConfig, writeClient, instantTime)
+      case SaveMode.Append if overwriteOperationType == WriteOperationType.BUCKET_RESCALE => {
+        instantTime = writeConfig.getHashingConfigInstantToLoad()
+        new DatasetBucketRescaleCommitActionExecutor(writeConfig, writeClient, instantTime)
+    }
       case SaveMode.Overwrite if overwriteOperationType == WriteOperationType.INSERT_OVERWRITE_TABLE =>
         new DatasetBulkInsertOverwriteTableCommitActionExecutor(writeConfig, writeClient, instantTime)
       case _ =>
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala
index ff3261bbcd6..6cf5053bdce 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala
@@ -33,7 +33,6 @@ import org.apache.hudi.hive.ddl.HiveSyncMode
 import org.apache.hudi.keygen.{ComplexKeyGenerator, CustomAvroKeyGenerator, CustomKeyGenerator}
 import org.apache.hudi.sql.InsertMode
 import org.apache.hudi.sync.common.HoodieSyncConfig
-
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.{SaveMode, SparkSession}
 import org.apache.spark.sql.catalyst.catalog.HoodieCatalogTable
@@ -50,8 +49,8 @@ import org.apache.spark.sql.types.StructType
 import org.slf4j.LoggerFactory
 
 import java.util.Locale
-
 import scala.collection.JavaConverters._
+import scala.collection.mutable
 
 trait ProvidesHoodieConfig extends Logging {
 
@@ -85,6 +84,27 @@ trait ProvidesHoodieConfig extends Logging {
       defaultOpts = defaultOpts, overridingOpts = overridingOpts)
   }
 
+  def buildBucketRescaleHoodieConfig(hoodieCatalogTable: HoodieCatalogTable): Map[String, String] = {
+    val sparkSession: SparkSession = hoodieCatalogTable.spark
+    val tableConfig = hoodieCatalogTable.tableConfig
+    val preCombineField = Option(tableConfig.getPreCombineField).getOrElse("")
+    val hiveSyncConfig = buildHiveSyncConfig(sparkSession, hoodieCatalogTable, tableConfig)
+
+    val defaultOpts = Map[String, String](
+      OPERATION.key -> BULK_INSERT_OPERATION_OPT_VAL,
+      HoodieInternalConfig.BULKINSERT_OVERWRITE_OPERATION_TYPE.key -> WriteOperationType.BUCKET_RESCALE.value(),
+      KEYGENERATOR_CLASS_NAME.key -> classOf[SqlKeyGenerator].getCanonicalName,
+      SqlKeyGenerator.ORIGINAL_KEYGEN_CLASS_NAME -> tableConfig.getKeyGeneratorClassName,
+      SqlKeyGenerator.PARTITION_SCHEMA -> hoodieCatalogTable.partitionSchema.toDDL,
+      HoodieSyncConfig.META_SYNC_ENABLED.key -> hiveSyncConfig.getString(HoodieSyncConfig.META_SYNC_ENABLED.key),
+      HiveSyncConfigHolder.HIVE_SYNC_ENABLED.key -> "false"
+    )
+
+    val overridingOpts = buildOverridingOpts(hoodieCatalogTable, preCombineField)
+    combineOptions(hoodieCatalogTable, tableConfig, sparkSession.sqlContext.conf,
+      defaultOpts = defaultOpts, overridingOpts = overridingOpts)
+  }
+
   /**
    * Deduce the sql write operation for INSERT_INTO
    */
diff --git a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/PartitionBucketIndexManager.scala b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/PartitionBucketIndexManager.scala
new file mode 100644
index 00000000000..4aaeec80a1a
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/PartitionBucketIndexManager.scala
@@ -0,0 +1,289 @@
+/*
+ * 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.spark.sql.hudi.command.procedures
+
+import org.apache.avro.Schema
+import org.apache.hadoop.conf.Configuration
+import org.apache.hudi.{AvroConversionUtils, HoodieCLIUtils, HoodieSparkSqlWriter}
+import org.apache.hudi.avro.HoodieAvroUtils
+import org.apache.hudi.client.SparkRDDWriteClient
+import org.apache.hudi.client.common.HoodieSparkEngineContext
+import org.apache.hudi.commit.DatasetBucketRescaleCommitActionExecutor
+import org.apache.hudi.common.config.{HoodieMetadataConfig, HoodieReaderConfig, SerializableSchema}
+import org.apache.hudi.common.engine.HoodieEngineContext
+import org.apache.hudi.common.fs.FSUtils
+import org.apache.hudi.common.model.{FileSlice, PartitionBucketIndexHashingConfig, WriteOperationType}
+import org.apache.hudi.common.table.read.HoodieFileGroupReader
+import org.apache.hudi.common.table.{HoodieTableMetaClient, TableSchemaResolver}
+import org.apache.hudi.common.table.view.HoodieTableFileSystemView
+import org.apache.hudi.common.util.{Option, StringUtils}
+import org.apache.hudi.config.{HoodieIndexConfig, HoodieWriteConfig}
+import org.apache.hudi.exception.HoodieException
+import org.apache.hudi.index.bucket.{PartitionBucketIndexCalculator, PartitionBucketIndexUtils}
+import org.apache.hudi.internal.schema.InternalSchema
+import org.apache.hudi.internal.schema.utils.SerDeHelper
+import org.apache.hudi.storage.hadoop.HadoopStorageConfiguration
+import org.apache.hudi.storage.{StoragePath, StoragePathInfo}
+import org.apache.hudi.table.SparkBroadcastManager
+import org.apache.spark.internal.Logging
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.{HoodieUnsafeUtils, Row, SaveMode}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.catalog.HoodieCatalogTable
+import org.apache.spark.sql.catalyst.expressions.PredicateHelper
+import org.apache.spark.sql.hudi.ProvidesHoodieConfig
+import org.apache.spark.sql.types.{DataTypes, Metadata, StructField, StructType}
+
+import java.util
+import java.util.function.Supplier
+import java.util.stream.Collectors
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+class PartitionBucketIndexManager() extends BaseProcedure
+  with ProcedureBuilder
+  with PredicateHelper
+  with ProvidesHoodieConfig
+  with Logging {
+
+  private val PARAMETERS = Array[ProcedureParameter](
+    ProcedureParameter.required(0, "table", DataTypes.StringType),
+    ProcedureParameter.optional(1, "overwrite", DataTypes.StringType),
+    ProcedureParameter.optional(2, "bucket-number", DataTypes.IntegerType, -1),
+    ProcedureParameter.optional(3, "add", DataTypes.StringType),
+    ProcedureParameter.optional(4, "dry-run", DataTypes.BooleanType, true),
+    ProcedureParameter.optional(5, "rollback", DataTypes.StringType),
+    ProcedureParameter.optional(6, "show-config", DataTypes.BooleanType, false),
+    ProcedureParameter.optional(7, "rule", DataTypes.StringType, "regex")
+  )
+
+  private val OUTPUT_TYPE = new StructType(Array[StructField](
+    StructField("result", DataTypes.StringType, nullable = true, Metadata.empty),
+    StructField("operation", DataTypes.StringType, nullable = true, Metadata.empty),
+    StructField("details", DataTypes.StringType, nullable = true, Metadata.empty)
+  ))
+
+  def parameters: Array[ProcedureParameter] = PARAMETERS
+
+  def outputType: StructType = OUTPUT_TYPE
+
+  override def call(args: ProcedureArgs): Seq[Row] = {
+    super.checkArgs(PARAMETERS, args)
+
+    val table = getArgValueOrDefault(args, PARAMETERS(0)).get.asInstanceOf[String]
+    val overwrite = getArgValueOrDefault(args, PARAMETERS(1)).orNull.asInstanceOf[String]
+    val bucketNumber = getArgValueOrDefault(args, PARAMETERS(2)).get.asInstanceOf[Int]
+    val add = getArgValueOrDefault(args, PARAMETERS(3)).orNull.asInstanceOf[String]
+    val dryRun = getArgValueOrDefault(args, PARAMETERS(4)).get.asInstanceOf[Boolean]
+    val rollback = getArgValueOrDefault(args, PARAMETERS(5)).orNull.asInstanceOf[String]
+    val showConfig = getArgValueOrDefault(args, PARAMETERS(6)).get.asInstanceOf[Boolean]
+    val rule = getArgValueOrDefault(args, PARAMETERS(7)).orNull.asInstanceOf[String]
+
+    try {
+      // Get table metadata
+      val hoodieCatalogTable = HoodieCLIUtils.getHoodieCatalogTable(sparkSession, table)
+      val basePath = hoodieCatalogTable.tableLocation
+      val metaClient = createMetaClient(jsc, basePath)
+      val writeClient = HoodieCLIUtils.createHoodieWriteClient(sparkSession, basePath, Map.empty,
+        scala.Option.apply(table))
+      val context = writeClient.getEngineContext
+
+      // Determine which operation to perform
+      if (showConfig) {
+        handleShowConfig(metaClient)
+      } else if (rollback != null) {
+        handleRollback(metaClient, rollback)
+      } else if (overwrite != null) {
+        handleOverwrite(hoodieCatalogTable, writeClient, context, metaClient, overwrite, bucketNumber, rule, dryRun)
+      } else if (add != null) {
+        handleAdd(metaClient, add, dryRun)
+      } else {
+        Seq(Row("ERROR", "INVALID_OPERATION", "No valid operation specified"))
+      }
+    } catch {
+      case e: Exception => Seq(Row("ERROR", "EXCEPTION", e.getMessage))
+    }
+  }
+
+  /**
+   * Handle the overwrite operation.
+   */
+  private def handleOverwrite(catalog: HoodieCatalogTable,
+                              writeClient: SparkRDDWriteClient[_],
+                              context : HoodieEngineContext,
+                              metaClient: HoodieTableMetaClient,
+                              expression: String,
+                              bucketNumber: Int,
+                              rule: String,
+                              dryRun: Boolean): Seq[Row] = {
+    val config = buildBucketRescaleHoodieConfig(catalog).asJava
+    val basePath = metaClient.getBasePath
+    val mdtEnable = metaClient.getStorage().exists(new StoragePath(metaClient.getBasePath, HoodieTableMetaClient.METADATA_TABLE_FOLDER_PATH))
+
+    // get all partition paths
+    val allPartitions = FSUtils.getAllPartitionPaths(context, metaClient.getStorage, metaClient.getBasePath, mdtEnable)
+
+    // get Map<partitionPath, bucketNumber> based on latest hashing_config
+    val latestHashingConfigInstant = PartitionBucketIndexUtils.getHashingConfigInstantToLoad(metaClient)
+    val calcWithLatestInstant = PartitionBucketIndexCalculator.getInstance(latestHashingConfigInstant, metaClient)
+    val partition2BucketWithLatestHashingConfig = PartitionBucketIndexUtils.getAllBucketNumbers(calcWithLatestInstant, allPartitions)
+
+    // get Map<partitionPath, bucketNumber> based on new given expression
+    val defaultBucketNumber = if (bucketNumber != -1) {
+      bucketNumber
+    } else {
+      // reuse latest default bucket number
+      calcWithLatestInstant.getHashingConfig.getDefaultBucketNumber
+    }
+    val instantTime = writeClient.createNewInstantTime()
+    config.put(HoodieIndexConfig.BUCKET_INDEX_PARTITION_EXPRESSIONS.key(), expression)
+    config.put(HoodieIndexConfig.BUCKET_INDEX_PARTITION_RULE_TYPE.key(), rule)
+    config.put(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.key(), instantTime)
+    config.put(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS.key(), bucketNumber.toString)
+
+    val newConfig = new PartitionBucketIndexHashingConfig(expression, defaultBucketNumber, rule, PartitionBucketIndexHashingConfig.CURRENT_VERSION, instantTime)
+    val calcWithNewExpression = PartitionBucketIndexCalculator.getInstance(instantTime, newConfig)
+    val partition2BucketWithNewHashingConfig = PartitionBucketIndexUtils.getAllBucketNumbers(calcWithNewExpression, allPartitions)
+
+    // get partitions need to be rescaled
+    val rescalePartitionsMap = getDifferentPartitions(partition2BucketWithNewHashingConfig.asScala, partition2BucketWithLatestHashingConfig.asScala)
+    val partitionsToRescale = rescalePartitionsMap.keys
+
+    // get all fileSlices need to read
+    val allFilesMap = FSUtils.getFilesInPartitions(context, metaClient.getStorage(), HoodieMetadataConfig.newBuilder.enable(mdtEnable).build,
+      metaClient.getBasePath.toString, partitionsToRescale.toArray)
+    val files = allFilesMap.values().asScala.flatMap(x => x.asScala).toList
+    val view = new HoodieTableFileSystemView(metaClient, metaClient.getActiveTimeline, files.asJava)
+    val allFileSlice = partitionsToRescale.flatMap(partitionPath => {
+      view.getLatestFileSlices(partitionPath).iterator().asScala
+    }).toList
+
+    // read all fileSlice para and get DF
+    var tableSchemaWithMetaFields: Schema = null
+    try tableSchemaWithMetaFields = HoodieAvroUtils.addMetadataFields(new TableSchemaResolver(metaClient).getTableAvroSchema(false), false)
+    catch {
+      case e: Exception =>
+        throw new HoodieException("Failed to get table schema during clustering", e)
+    }
+
+    // broadcast reader context.
+    val broadcastManager = new SparkBroadcastManager(context, metaClient)
+    broadcastManager.prepareAndBroadcast()
+    val sparkSchemaWithMetaFields = AvroConversionUtils.convertAvroSchemaToStructType(tableSchemaWithMetaFields)
+    val serializableTableSchemaWithMetaFields = new SerializableSchema(tableSchemaWithMetaFields)
+    val latestInstantTime = metaClient.getActiveTimeline.getCommitsTimeline.filterCompletedInstants().lastInstant().get()
+
+    val res: RDD[InternalRow] = spark.sparkContext.parallelize(allFileSlice, allFileSlice.size).flatMap(fileSlice => {
+      // instantiate other supporting cast
+      val readerSchema = serializableTableSchemaWithMetaFields.get
+      val readerContextOpt = broadcastManager.retrieveFileGroupReaderContext(basePath)
+      val internalSchemaOption: Option[InternalSchema] = Option.empty()
+      // instantiate FG reader
+      val fileGroupReader = new HoodieFileGroupReader(readerContextOpt.get(),
+        metaClient.getStorage,
+        basePath.toString,
+        latestInstantTime.requestedTime(),
+        fileSlice,
+        readerSchema,
+        readerSchema,
+        internalSchemaOption, // not support evolution of schema for now
+        metaClient,
+        metaClient.getTableConfig.getProps,
+        0,
+        java.lang.Long.MAX_VALUE,
+        HoodieReaderConfig.MERGE_USE_RECORD_POSITIONS.defaultValue())
+      fileGroupReader.initRecordIterators()
+      val iterator = fileGroupReader.getClosableIterator.asInstanceOf[HoodieFileGroupReader.HoodieFileGroupReaderIterator[InternalRow]]
+      iterator.asScala
+    })
+    val dataFrame = HoodieUnsafeUtils.createDataFrameFromRDD(sparkSession, res, sparkSchemaWithMetaFields)
+
+    val (success, _, _, _, _, _) = HoodieSparkSqlWriter.write(
+      sparkSession.sqlContext,
+      SaveMode.Overwrite,
+      config.asScala.toMap,
+      dataFrame)
+
+    val details = s"Expression: $expression, Bucket Number: $bucketNumber, Dry Run: $dryRun"
+    Seq(Row("SUCCESS", "OVERWRITE", success))
+  }
+
+  /**
+   * Compares two maps of partition paths to bucket numbers and returns a map of partition paths
+   * that have different bucket numbers in the two maps, using the new bucket numbers.
+   *
+   * @param newMap The new map of partition paths to bucket numbers
+   * @param oldMap The old map of partition paths to bucket numbers
+   * @return A map containing only the partition paths that have different bucket numbers, with values from newMap
+   */
+  def getDifferentPartitions(newMap: mutable.Map[String, Integer], oldMap: mutable.Map[String, Integer]): mutable.Map[String, Integer] = {
+    newMap.filter {
+      case (partitionPath, bucketNumber) =>
+        // Include partition path if it's not in old map or has a different bucket number
+        !oldMap.contains(partitionPath) || oldMap(partitionPath) != bucketNumber
+    }
+  }
+
+  /**
+   * Handle the add operation.
+   */
+  private def handleAdd(metaClient: HoodieTableMetaClient, expression: String, dryRun: Boolean): Seq[Row] = {
+    // In a real implementation, this would call PartitionBucketIndexManager
+    // For now, just return a placeholder result
+    val details = s"Expression: $expression, Dry Run: $dryRun"
+
+    // Here would be the actual call to PartitionBucketIndexManager
+    // PartitionBucketIndexManager.addExpression(metaClient, expression, dryRun)
+
+    Seq(Row("SUCCESS", "ADD", details))
+  }
+
+  /**
+   * Handle the rollback operation.
+   */
+  private def handleRollback(metaClient: HoodieTableMetaClient, instantTime: String): Seq[Row] = {
+    // In a real implementation, this would call PartitionBucketIndexManager
+    // For now, just return a placeholder result
+    val details = s"Rolled back bucket rescale action: $instantTime"
+
+    // Here would be the actual call to PartitionBucketIndexManager
+    // PartitionBucketIndexManager.rollback(metaClient, instantTime)
+
+    Seq(Row("SUCCESS", "ROLLBACK", details))
+  }
+
+  /**
+   * Handle the show-config operation.
+   */
+  private def handleShowConfig(metaClient: HoodieTableMetaClient): Seq[Row] = {
+
+
+
+    Seq(Row("SUCCESS", "SHOW_CONFIG", null))
+  }
+
+  override def build: Procedure = new PartitionBucketIndexManager()
+}
+
+object PartitionBucketIndexManager {
+  val NAME = "PartitionBucketIndexManager"
+
+  def builder: Supplier[ProcedureBuilder] = new Supplier[ProcedureBuilder] {
+    override def get() = new PartitionBucketIndexManager()
+  }
+}
\ No newline at end of file
diff --git a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/common/HoodieSparkSqlTestBase.scala b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/common/HoodieSparkSqlTestBase.scala
index e0e60e3f6b4..5fd8fd4ad19 100644
--- a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/common/HoodieSparkSqlTestBase.scala
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/common/HoodieSparkSqlTestBase.scala
@@ -74,6 +74,7 @@ class HoodieSparkSqlTestBase extends FunSuite with BeforeAndAfterAll {
   TimeZone.setDefault(DateTimeUtils.getTimeZone("UTC"))
   protected lazy val spark: SparkSession = SparkSession.builder()
     .config("spark.sql.warehouse.dir", sparkWareHouse.getCanonicalPath)
+    .config("spark.driver.bindAddress", "127.0.0.1")
     .config("spark.sql.session.timeZone", "UTC")
     .config("hoodie.insert.shuffle.parallelism", "4")
     .config("hoodie.upsert.shuffle.parallelism", "4")
diff --git a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala
index 7c1bc0aa2a2..532a000f0f6 100644
--- a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala
@@ -1847,6 +1847,94 @@ class TestInsertTable extends HoodieSparkSqlTestBase {
     }
   }
 
+  test("Test Bulk Insert Into Partition Bucket Index Table") {
+    withSQLConf(
+      "hoodie.datasource.write.operation" -> "bulk_insert",
+      "hoodie.bulkinsert.shuffle.parallelism" -> "1") {
+      Seq("mor", "cow").foreach { tableType =>
+        Seq("true", "false").foreach { bulkInsertAsRow =>
+          withTempDir { tmp =>
+            val tableName = generateTableName
+            // 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,name',
+                 | type = '$tableType',
+                 | preCombineField = 'ts',
+                 | hoodie.index.type = 'BUCKET',
+                 | hoodie.bucket.index.hash.field = 'id,name',
+                 | hoodie.datasource.write.row.writer.enable = '$bulkInsertAsRow')
+                 | partitioned by (dt)
+                 | location '${tmp.getCanonicalPath}'
+                 """.stripMargin)
+
+            // Note: Do not write the field alias, the partition field must be placed last.
+            spark.sql(
+              s"""
+                 | insert into $tableName values
+                 | (1, 'a1,1', 10, 1000, "2021-01-05"),
+                 | (2, 'a2', 20, 2000, "2021-01-06"),
+                 | (3, 'a3,3', 30, 3000, "2021-01-07")
+                 """.stripMargin)
+
+            checkAnswer(s"select id, name, price, ts, dt from $tableName")(
+              Seq(1, "a1,1", 10.0, 1000, "2021-01-05"),
+              Seq(2, "a2", 20.0, 2000, "2021-01-06"),
+              Seq(3, "a3,3", 30.0, 3000, "2021-01-07")
+            )
+
+            // for COW with disabled Spark native row writer, multiple bulk inserts are restricted
+            if (tableType != "cow" && bulkInsertAsRow != "false") {
+              spark.sql(
+                s"""
+                   | insert into $tableName values
+                   | (1, 'a1', 10, 1000, "2021-01-05"),
+                   | (3, "a3", 30, 3000, "2021-01-07")
+                 """.stripMargin)
+
+              checkAnswer(s"select id, name, price, ts, dt from $tableName")(
+                Seq(1, "a1,1", 10.0, 1000, "2021-01-05"),
+                Seq(1, "a1", 10.0, 1000, "2021-01-05"),
+                Seq(2, "a2", 20.0, 2000, "2021-01-06"),
+                Seq(3, "a3,3", 30.0, 3000, "2021-01-07"),
+                Seq(3, "a3", 30.0, 3000, "2021-01-07")
+              )
+
+              // there are two files in partition(dt = '2021-01-05')
+              checkAnswer(s"select count(distinct _hoodie_file_name) from $tableName where dt = '2021-01-05'")(
+                Seq(2)
+              )
+
+              // would generate 6 other files in partition(dt = '2021-01-05')
+              spark.sql(
+                s"""
+                   | insert into $tableName values
+                   | (4, 'a1,1', 10, 1000, "2021-01-05"),
+                   | (5, 'a1,1', 10, 1000, "2021-01-05"),
+                   | (6, 'a1,1', 10, 1000, "2021-01-05"),
+                   | (7, 'a1,1', 10, 1000, "2021-01-05"),
+                   | (8, 'a1,1', 10, 1000, "2021-01-05"),
+                   | (10, 'a3,3', 30, 3000, "2021-01-05")
+                 """.stripMargin)
+
+              checkAnswer(s"select count(distinct _hoodie_file_name) from $tableName where dt = '2021-01-05'")(
+                Seq(8)
+              )
+            }
+          }
+        }
+      }
+    }
+  }
+
   test("Test not supported multiple BULK INSERTs into COW with SIMPLE BUCKET and disabled Spark native row writer") {
     withSQLConf("hoodie.datasource.write.operation" -> "bulk_insert",
       "hoodie.bulkinsert.shuffle.parallelism" -> "1") {
diff --git a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/procedure/TestPartitionBucketIndexManager.scala b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/procedure/TestPartitionBucketIndexManager.scala
new file mode 100644
index 00000000000..93d2a557afe
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/procedure/TestPartitionBucketIndexManager.scala
@@ -0,0 +1,31 @@
+package org.apache.spark.sql.hudi.procedure
+
+import org.apache.hudi.common.table.HoodieTableMetaClient
+import org.apache.hudi.hadoop.fs.HadoopFSUtils
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils
+
+class TestPartitionBucketIndexManager extends HoodieSparkProcedureTestBase {
+
+  test("Case1: Test Call drop_partition Procedure For Multiple Partitions: '*' stands for all partitions in leaf partition") {
+    withTempDir { tmp =>
+      val tableName = generateTableName
+      val tablePath = tmp.getCanonicalPath + "/" + tableName
+      val metaClient = getTableMetaClient(tablePath)
+      val expressions = "\\d{4}-(06-(01|17|18)|11-(01|10|11)),256"
+      val rule ="regex"
+      val defaultBucketNumber = 10
+      PartitionBucketIndexUtils.initHashingConfig(metaClient, expressions, rule, defaultBucketNumber, null)
+
+      spark.sql(s"""call PartitionBucketIndexManager(table => '$tableName', show-config => 'true')""").collect()
+
+    }
+  }
+
+  private def getTableMetaClient(tablePath: String): HoodieTableMetaClient = {
+    HoodieTableMetaClient.builder()
+      .setBasePath(tablePath)
+      .setConf(HadoopFSUtils.getStorageConf(spark.sparkContext.hadoopConfiguration))
+      .build()
+  }
+
+}


(hudi) 01/06: finish hashing_config initial && related UT

Posted by zh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zhangyue19921010 pushed a commit to branch HUDI-8990
in repository https://gitbox.apache.org/repos/asf/hudi.git

commit 475b4f108e01e137dcac0fde994ac63ac42e1f15
Author: YueZhang <zh...@163.com>
AuthorDate: Mon Mar 17 20:02:24 2025 +0800

    finish hashing_config initial && related UT
---
 .../org/apache/hudi/config/HoodieIndexConfig.java  | 15 ++++
 .../HoodieSimpleBucketIndexPartitionRule.java      | 29 +++++++
 .../hudi/index/bucket/SimpleBucketIndexUtils.java  | 81 ++++++++++++++++++
 .../model/PartitionBucketIndexHashingConfig.java   | 98 ++++++++++++++++++++++
 .../hudi/common/table/HoodieTableMetaClient.java   | 17 +++-
 .../apache/hudi/configuration/FlinkOptions.java    | 16 ++++
 .../apache/hudi/table/catalog/HoodieCatalog.java   |  3 +-
 .../hudi/table/catalog/HoodieCatalogUtil.java      | 11 +++
 .../hudi/table/catalog/HoodieHiveCatalog.java      |  7 +-
 .../hudi/table/catalog/TestHoodieCatalog.java      | 36 ++++++++
 10 files changed, 308 insertions(+), 5 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
index 48b30dbe6d0..626e894fd1a 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
@@ -28,6 +28,7 @@ import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.exception.HoodieIndexException;
 import org.apache.hudi.exception.HoodieNotSupportedException;
 import org.apache.hudi.index.HoodieIndex;
+import org.apache.hudi.index.bucket.HoodieSimpleBucketIndexPartitionRule;
 import org.apache.hudi.keygen.constant.KeyGeneratorOptions;
 
 import org.slf4j.Logger;
@@ -286,6 +287,20 @@ public class HoodieIndexConfig extends HoodieConfig {
       .withDocumentation("Only applies if index type is BUCKET. Determine the number of buckets in the hudi table, "
           + "and each partition is divided to N buckets.");
 
+  public static final ConfigProperty<String> BUCKET_INDEX_PARTITION_RULE_TYPE = ConfigProperty
+      .key("hoodie.bucket.index.partition.rule.type")
+      .defaultValue(HoodieSimpleBucketIndexPartitionRule.REGEX.value)
+      .markAdvanced()
+      .withDocumentation("Rule parser for expressions when using partition level bucket index, default regex.");
+
+  public static final ConfigProperty<String> BUCKET_INDEX_PARTITION_EXPRESSIONS = ConfigProperty
+      .key("hoodie.bucket.index.partition.expressions")
+      .noDefaultValue()
+      .markAdvanced()
+      .withDocumentation("Users can use this parameter to specify expression and the corresponding bucket " +
+          "numbers (separated by commas).Multiple rules are separated by semicolons like " +
+          "hoodie.bucket.index.partition.expressions=expression1,bucket-number1;expression2,bucket-number2");
+
   public static final ConfigProperty<String> BUCKET_INDEX_MAX_NUM_BUCKETS = ConfigProperty
       .key("hoodie.bucket.index.max.num.buckets")
       .noDefaultValue()
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndexPartitionRule.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndexPartitionRule.java
new file mode 100644
index 00000000000..8d86ed8080b
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndexPartitionRule.java
@@ -0,0 +1,29 @@
+/*
+ * 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.index.bucket;
+
+public enum HoodieSimpleBucketIndexPartitionRule {
+  REGEX("regex");
+
+  public final String value;
+
+  HoodieSimpleBucketIndexPartitionRule(String rule) {
+    this.value = rule;
+  }
+}
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java
new file mode 100644
index 00000000000..74c8022c9aa
--- /dev/null
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java
@@ -0,0 +1,81 @@
+/*
+ * 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.index.bucket;
+
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
+import org.apache.hudi.common.table.timeline.HoodieTimeline;
+import org.apache.hudi.common.util.FileIOUtils;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.exception.HoodieIOException;
+import org.apache.hudi.storage.HoodieInstantWriter;
+import org.apache.hudi.storage.HoodieStorage;
+import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.storage.StoragePathInfo;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+
+public class SimpleBucketIndexUtils {
+  public static final String INITIAL_HASHING_CONFIG_INSTANT = HoodieTimeline.INIT_INSTANT_TS;
+
+  private static final Logger LOG = LoggerFactory.getLogger(SimpleBucketIndexUtils.class);
+
+  public static boolean initHashingConfig(HoodieTableMetaClient metaClient,
+                                              String expressions,
+                                              String rule,
+                                              int defaultBucketNumber,
+                                              String instant){
+    if (StringUtils.isNullOrEmpty(expressions)) {
+      return false;
+    }
+    String hashingInstant = StringUtils.isNullOrEmpty(instant) ? INITIAL_HASHING_CONFIG_INSTANT : instant;
+    HoodieStorage storage = metaClient.getStorage();
+    PartitionBucketIndexHashingConfig hashingConfig =
+        new PartitionBucketIndexHashingConfig(expressions, defaultBucketNumber, rule, PartitionBucketIndexHashingConfig.CURRENT_VERSION, hashingInstant);
+    StoragePath hashingConfigPath = new StoragePath(metaClient.getHashingMetadataConfigPath(), hashingConfig.getFilename());
+
+    try {
+      Option<byte []> content = Option.of(hashingConfig.toJsonString().getBytes(StandardCharsets.UTF_8));
+      storage.createImmutableFileInPath(hashingConfigPath, content.map(HoodieInstantWriter::convertByteArrayToWriter));
+    } catch (IOException ioe) {
+      throw new HoodieIOException("Failed to initHashingConfig ", ioe);
+    }
+    return true;
+  }
+
+  public static Option<PartitionBucketIndexHashingConfig> loadHashingConfig(HoodieStorage storage,
+                                                                            StoragePathInfo hashingConfig) {
+    if (hashingConfig == null) {
+      return Option.empty();
+    }
+    try (InputStream is = storage.open(hashingConfig.getPath())) {
+      byte[] content = FileIOUtils.readAsByteArray(is);
+      return Option.of(PartitionBucketIndexHashingConfig.fromBytes(content));
+    } catch (IOException e) {
+      LOG.error("Error when loading hashing config, for path: " + hashingConfig.getPath().getName(), e);
+      throw new HoodieIOException("Error while loading hashing config", e);
+    }
+  }
+}
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
new file mode 100644
index 00000000000..225629b52ff
--- /dev/null
+++ b/hudi-common/src/main/java/org/apache/hudi/common/model/PartitionBucketIndexHashingConfig.java
@@ -0,0 +1,98 @@
+/*
+ * 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.common.model;
+
+import org.apache.hudi.common.util.JsonUtils;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.nio.charset.StandardCharsets;
+
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class PartitionBucketIndexHashingConfig implements Serializable {
+  private static final Logger LOG = LoggerFactory.getLogger(PartitionBucketIndexHashingConfig.class);
+  public static final String HASHING_CONFIG_FILE_SUFFIX = ".hashing_config";
+  public static final Integer CURRENT_VERSION = 1;
+  private final String expressions;
+  private final int defaultBucketNumber;
+  private final String rule;
+
+  private final int version;
+  private final String instant;
+
+  @JsonCreator
+  public PartitionBucketIndexHashingConfig(@JsonProperty("expressions") String expressions,
+                                           @JsonProperty("defaultBucketNumber") int defaultBucketNumber,
+                                           @JsonProperty("rule") String rule,
+                                           @JsonProperty("version") int version,
+                                           @JsonProperty("instant") String instant) {
+    this.expressions = expressions;
+    this.defaultBucketNumber = defaultBucketNumber;
+    this.rule = rule;
+    this.version = version;
+    this.instant = instant;
+  }
+
+  public String getFilename() {
+    return instant + HASHING_CONFIG_FILE_SUFFIX;
+  }
+
+  public String toJsonString() throws IOException {
+    return JsonUtils.getObjectMapper().writerWithDefaultPrettyPrinter().writeValueAsString(this);
+  }
+
+  public static <T> T fromJsonString(String jsonStr, Class<T> clazz) throws Exception {
+    if (jsonStr == null || jsonStr.isEmpty()) {
+      // For empty commit file (no data or somethings bad happen).
+      return clazz.newInstance();
+    }
+    return JsonUtils.getObjectMapper().readValue(jsonStr, clazz);
+  }
+
+
+  public static PartitionBucketIndexHashingConfig fromBytes(byte[] bytes) throws IOException {
+    try {
+      return fromJsonString(new String(bytes, StandardCharsets.UTF_8), PartitionBucketIndexHashingConfig.class);
+    } catch (Exception e) {
+      throw new IOException("unable to load hashing config", e);
+    }
+  }
+
+  public int getVersion() {
+    return version;
+  }
+
+  public String getRule() {
+    return rule;
+  }
+
+  public int getDefaultBucketNumber() {
+    return defaultBucketNumber;
+  }
+
+  public String getExpressions() {
+    return expressions;
+  }
+}
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java b/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java
index 31b7df91d0b..d50e9f4b7cd 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/table/HoodieTableMetaClient.java
@@ -111,6 +111,7 @@ public class HoodieTableMetaClient implements Serializable {
   public static final String METADATA_STR = "metadata";
   public static final String METAFOLDER_NAME = ".hoodie";
   public static final String TIMELINEFOLDER_NAME = "timeline";
+  public static final String BUCKET_INDEX_METAFOLDER_NAME = ".bucket_index";
   public static final String TEMPFOLDER_NAME = METAFOLDER_NAME + StoragePath.SEPARATOR + ".temp";
   public static final String AUXILIARYFOLDER_NAME = METAFOLDER_NAME + StoragePath.SEPARATOR + ".aux";
   public static final String BOOTSTRAP_INDEX_ROOT_FOLDER_PATH = AUXILIARYFOLDER_NAME + StoragePath.SEPARATOR + ".bootstrap";
@@ -118,7 +119,13 @@ public class HoodieTableMetaClient implements Serializable {
   public static final String HEARTBEAT_FOLDER_NAME = METAFOLDER_NAME + StoragePath.SEPARATOR + ".heartbeat";
   public static final String METADATA_TABLE_FOLDER_PATH = METAFOLDER_NAME + StoragePath.SEPARATOR + METADATA_STR;
   public static final String HASHING_METADATA_FOLDER_NAME =
-      ".bucket_index" + StoragePath.SEPARATOR + "consistent_hashing_metadata";
+      BUCKET_INDEX_METAFOLDER_NAME + StoragePath.SEPARATOR + "consistent_hashing_metadata";
+  public static final String PARTITION_BUCKET_INDEX_HASHING_FOLDER =
+      BUCKET_INDEX_METAFOLDER_NAME + StoragePath.SEPARATOR + "partition_bucket_index_meta";
+  public static final String BUCKET_INDEX_METAFOLDER_CONFIG_FOLDER =
+      PARTITION_BUCKET_INDEX_HASHING_FOLDER + StoragePath.SEPARATOR + "configs";
+  public static final String BUCKET_INDEX_METAFOLDER_CONFIG_ARCHIVE_FOLDER =
+      PARTITION_BUCKET_INDEX_HASHING_FOLDER + StoragePath.SEPARATOR + "archive";
   public static final String BOOTSTRAP_INDEX_BY_PARTITION_FOLDER_PATH = BOOTSTRAP_INDEX_ROOT_FOLDER_PATH
       + StoragePath.SEPARATOR + ".partitions";
   public static final String BOOTSTRAP_INDEX_BY_FILE_ID_FOLDER_PATH =
@@ -361,6 +368,14 @@ public class HoodieTableMetaClient implements Serializable {
     return new StoragePath(metaPath, HASHING_METADATA_FOLDER_NAME).toString();
   }
 
+  /**
+   * Used for partition level bucket index to save hashing_config.
+   * @return
+   */
+  public String getHashingMetadataConfigPath() {
+    return new StoragePath(metaPath, BUCKET_INDEX_METAFOLDER_CONFIG_FOLDER).toString();
+  }
+
   /**
    * @return Temp Folder path
    */
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java
index d65c80f89e9..e1c85c123a7 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/FlinkOptions.java
@@ -474,6 +474,22 @@ public class FlinkOptions extends HoodieConfig {
       .defaultValue(4) // default 4 buckets per partition
       .withDescription("Hudi bucket number per partition. Only affected if using Hudi bucket index.");
 
+  @AdvancedConfig
+  public static final ConfigOption<String> BUCKET_INDEX_PARTITION_RULE = ConfigOptions
+      .key(HoodieIndexConfig.BUCKET_INDEX_PARTITION_RULE_TYPE.key())
+      .stringType()
+      .defaultValue(HoodieIndexConfig.BUCKET_INDEX_PARTITION_RULE_TYPE.defaultValue())
+      .withDescription("Rule parser for expressions when using partition level bucket index, default regex.");
+
+  @AdvancedConfig
+  public static final ConfigOption<String> BUCKET_INDEX_PARTITION_EXPRESSIONS = ConfigOptions
+      .key(HoodieIndexConfig.BUCKET_INDEX_PARTITION_EXPRESSIONS.key())
+      .stringType()
+      .noDefaultValue()
+      .withDescription("Users can use this parameter to specify expression and the corresponding bucket "
+          + "numbers (separated by commas).Multiple rules are separated by semicolons like "
+          + "hoodie.bucket.index.partition.expressions=expression1,bucket-number1;expression2,bucket-number2");
+
   public static final ConfigOption<String> PARTITION_PATH_FIELD = ConfigOptions
       .key(KeyGeneratorOptions.PARTITIONPATH_FIELD_NAME.key())
       .stringType()
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalog.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalog.java
index f9088b4096c..cb62dedc4c1 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalog.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalog.java
@@ -352,12 +352,13 @@ public class HoodieCatalog extends AbstractCatalog {
     }
     conf.setString(FlinkOptions.TABLE_NAME, tablePath.getObjectName());
     try {
-      StreamerUtil.initTableIfNotExists(conf);
+      HoodieTableMetaClient metaClient = StreamerUtil.initTableIfNotExists(conf);
       // prepare the non-table-options properties
       if (!StringUtils.isNullOrEmpty(resolvedTable.getComment())) {
         options.put(TableOptionProperties.COMMENT, resolvedTable.getComment());
       }
       TableOptionProperties.createProperties(tablePathStr, hadoopConf, options);
+      HoodieCatalogUtil.initPartitionBucketIndexMeta(metaClient, catalogTable);
     } catch (IOException e) {
       throw new CatalogException(String.format("Initialize table path %s exception.", tablePathStr), e);
     }
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java
index 99a23ded5e8..db1d8364ad3 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieCatalogUtil.java
@@ -27,6 +27,7 @@ import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.configuration.HadoopConfigurations;
 import org.apache.hudi.exception.HoodieCatalogException;
+import org.apache.hudi.index.bucket.SimpleBucketIndexUtils;
 import org.apache.hudi.internal.schema.InternalSchema;
 import org.apache.hudi.internal.schema.Type;
 import org.apache.hudi.internal.schema.convert.AvroInternalSchemaConverter;
@@ -301,4 +302,14 @@ public class HoodieCatalogUtil {
     }
     return true;
   }
+
+  public static boolean initPartitionBucketIndexMeta(HoodieTableMetaClient metaClient, CatalogBaseTable catalogTable) {
+    Map<String, String> options = catalogTable.getOptions();
+    String expressions = options.getOrDefault(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS.key(), "");
+    String rule = options.getOrDefault(FlinkOptions.BUCKET_INDEX_PARTITION_RULE.key(), FlinkOptions.BUCKET_INDEX_PARTITION_RULE.defaultValue());
+    int bucketNumber = options.containsKey(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.key()) ?
+        Integer.valueOf(options.get(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.key())) : FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.defaultValue();
+
+    return SimpleBucketIndexUtils.initHashingConfig(metaClient, expressions, rule, bucketNumber, null);
+  }
 }
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieHiveCatalog.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieHiveCatalog.java
index 17d59b873f9..e66a6745003 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieHiveCatalog.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/table/catalog/HoodieHiveCatalog.java
@@ -474,7 +474,8 @@ public class HoodieHiveCatalog extends AbstractCatalog {
       //create hive table
       client.createTable(hiveTable);
       //init hoodie metaClient
-      initTableIfNotExists(tablePath, (CatalogTable) table);
+      HoodieTableMetaClient metaClient = initTableIfNotExists(tablePath, (CatalogTable) table);
+      HoodieCatalogUtil.initPartitionBucketIndexMeta(metaClient, table);
     } catch (AlreadyExistsException e) {
       if (!ignoreIfExists) {
         throw new TableAlreadyExistException(getName(), tablePath, e);
@@ -485,7 +486,7 @@ public class HoodieHiveCatalog extends AbstractCatalog {
     }
   }
 
-  private void initTableIfNotExists(ObjectPath tablePath, CatalogTable catalogTable) {
+  private HoodieTableMetaClient initTableIfNotExists(ObjectPath tablePath, CatalogTable catalogTable) {
     Configuration flinkConf = Configuration.fromMap(catalogTable.getOptions());
     final String avroSchema = AvroSchemaConverter.convertToSchema(
         catalogTable.getSchema().toPersistedRowDataType().getLogicalType(),
@@ -527,7 +528,7 @@ public class HoodieHiveCatalog extends AbstractCatalog {
     StreamerUtil.checkPreCombineKey(flinkConf, fields);
 
     try {
-      StreamerUtil.initTableIfNotExists(flinkConf, hiveConf);
+      return StreamerUtil.initTableIfNotExists(flinkConf, hiveConf);
     } catch (IOException e) {
       throw new HoodieCatalogException("Initialize table exception.", e);
     }
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
index f6737128698..1e518c5ef45 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
@@ -21,6 +21,7 @@ package org.apache.hudi.table.catalog;
 import org.apache.hudi.common.model.DefaultHoodieRecordPayload;
 import org.apache.hudi.common.model.HoodieCommitMetadata;
 import org.apache.hudi.common.model.HoodieReplaceCommitMetadata;
+import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
 import org.apache.hudi.common.table.HoodieTableConfig;
 import org.apache.hudi.common.table.HoodieTableMetaClient;
 import org.apache.hudi.common.table.timeline.HoodieInstant;
@@ -30,10 +31,14 @@ import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.configuration.HadoopConfigurations;
 import org.apache.hudi.exception.HoodieIOException;
 import org.apache.hudi.exception.HoodieValidationException;
+import org.apache.hudi.index.bucket.SimpleBucketIndexUtils;
 import org.apache.hudi.keygen.ComplexAvroKeyGenerator;
 import org.apache.hudi.keygen.NonpartitionedAvroKeyGenerator;
 import org.apache.hudi.keygen.SimpleAvroKeyGenerator;
 import org.apache.hudi.sink.partitioner.profile.WriteProfiles;
+import org.apache.hudi.storage.HoodieStorage;
+import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.storage.StoragePathInfo;
 import org.apache.hudi.storage.hadoop.HadoopStorageConfiguration;
 import org.apache.hudi.util.StreamerUtil;
 import org.apache.hudi.utils.TestConfigurations;
@@ -334,6 +339,37 @@ public class TestHoodieCatalog {
     assertEquals(keyGeneratorClassName, NonpartitionedAvroKeyGenerator.class.getName());
   }
 
+  @Test
+  void testCreateTableWithPartitionSimpleBucketIndex() throws TableAlreadyExistException, DatabaseNotExistException, IOException {
+    String rule = "regex";
+    String expressions = "\\d{4}-(06-(01|17|18)|11-(01|10|11)),256";
+    String defaultBucketNumber = "20";
+    ObjectPath tablePath = new ObjectPath(TEST_DEFAULT_DATABASE, "tb1");
+    EXPECTED_CATALOG_TABLE.getOptions().put(FlinkOptions.BUCKET_INDEX_PARTITION_RULE.key(), rule);
+    EXPECTED_CATALOG_TABLE.getOptions().put(FlinkOptions.BUCKET_INDEX_PARTITION_EXPRESSIONS.key(), expressions);
+    EXPECTED_CATALOG_TABLE.getOptions().put(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS.key(), defaultBucketNumber);
+    // test create table
+    catalog.createTable(tablePath, EXPECTED_CATALOG_TABLE, true);
+
+    // test table exist
+    assertTrue(catalog.tableExists(tablePath));
+
+    String tablePathStr = catalog.inferTablePath(catalogPathStr, tablePath);
+    Configuration flinkConf = TestConfigurations.getDefaultConf(tablePathStr);
+    HoodieTableMetaClient metaClient = HoodieTestUtils
+        .createMetaClient(
+            new HadoopStorageConfiguration(HadoopConfigurations.getHadoopConf(flinkConf)), tablePathStr);
+    HoodieStorage storage = metaClient.getStorage();
+    StoragePath initialHashing_config =
+        new StoragePath(metaClient.getHashingMetadataConfigPath(), SimpleBucketIndexUtils.INITIAL_HASHING_CONFIG_INSTANT + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
+    StoragePathInfo info = storage.getPathInfo(initialHashing_config);
+    Option<PartitionBucketIndexHashingConfig> hashing_config = SimpleBucketIndexUtils.loadHashingConfig(storage, info);
+    assertTrue(hashing_config.isPresent());
+    assertEquals(hashing_config.get().getDefaultBucketNumber(), Integer.parseInt(defaultBucketNumber));
+    assertEquals(hashing_config.get().getRule(), rule);
+    assertEquals(hashing_config.get().getExpressions(), expressions);
+  }
+
   @Test
   void testCreateTableWithoutPreCombineKey() {
     Map<String, String> options = getDefaultCatalogOption();


(hudi) 04/06: finish Flink related coding and UT && Spark partitioner and index without test

Posted by zh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zhangyue19921010 pushed a commit to branch HUDI-8990
in repository https://gitbox.apache.org/repos/asf/hudi.git

commit 97bf409a5ad2cb939bc941f2b94432960c17a945
Author: YueZhang <zh...@163.com>
AuthorDate: Wed Mar 19 20:28:17 2025 +0800

    finish Flink related coding and UT && Spark partitioner and index without test
---
 .../org/apache/hudi/config/HoodieWriteConfig.java  |  4 ++
 .../hudi/index/bucket/HoodieSimpleBucketIndex.java |  4 ++
 .../bucket/PartitionBucketIndexCalculator.java     |  7 +-
 .../index/bucket/PartitionBucketIndexUtils.java    | 34 +++++----
 .../BucketIndexBulkInsertPartitionerWithRows.java  | 14 +++-
 .../RDDSimpleBucketBulkInsertPartitioner.java      | 14 +++-
 .../BucketBulkInsertDataInternalWriterHelper.java  | 11 ++-
 .../action/commit/SparkBucketIndexPartitioner.java | 82 +++++++++++++++++++---
 .../apache/spark/sql/BucketPartitionUtils.scala    | 21 ++++--
 .../hudi/configuration/OptionsInference.java       |  2 +-
 .../sink/bucket/BucketBulkInsertWriterHelper.java  | 11 +--
 .../sink/bucket/BucketStreamWriteFunction.java     |  2 -
 .../sink/partitioner/BucketIndexPartitioner.java   | 12 ++--
 .../java/org/apache/hudi/sink/utils/Pipelines.java | 16 +++--
 .../hudi/sink/utils/BulkInsertFunctionWrapper.java |  3 +-
 .../apache/hudi/table/ITTestHoodieDataSource.java  | 78 ++++++++++++++++++++
 .../org/apache/hudi/utils/TestConfigurations.java  |  3 +
 .../BaseDatasetBulkInsertCommitActionExecutor.java |  2 +-
 18 files changed, 259 insertions(+), 61 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java
index 959c94242b1..557f8672e19 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieWriteConfig.java
@@ -1955,6 +1955,10 @@ public class HoodieWriteConfig extends HoodieConfig {
     return HoodieIndex.IndexType.valueOf(getString(HoodieIndexConfig.INDEX_TYPE));
   }
 
+  public String getHashingConfigInstantToLoad() {
+    return getString(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT);
+  }
+
   public String getIndexClass() {
     return getString(HoodieIndexConfig.INDEX_CLASS_NAME);
   }
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java
index 9901ea0b257..e19e80b82a4 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java
@@ -140,6 +140,10 @@ public class HoodieSimpleBucketIndex extends HoodieBucketIndex {
     return BucketIdentifier.getBucketId(key.getRecordKey(), indexKeyFields, numBuckets);
   }
 
+  public int getBucketID(HoodieKey key, int numBuckets) {
+    return BucketIdentifier.getBucketId(key.getRecordKey(), indexKeyFields, numBuckets);
+  }
+
   @Override
   public boolean canIndexLogFiles() {
     return false;
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
index ff8af59a73d..65a08363402 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexCalculator.java
@@ -69,12 +69,7 @@ public class PartitionBucketIndexCalculator implements Serializable {
    */
   private PartitionBucketIndexCalculator(String instantToLoad, Configuration hadoopConf, String basePath) {
     this.instantToLoad = instantToLoad;
-    StoragePath metaPath = new StoragePath(basePath, HoodieTableMetaClient.METAFOLDER_NAME);
-    StoragePath hashingBase = new StoragePath(metaPath, HoodieTableMetaClient.BUCKET_INDEX_METAFOLDER_CONFIG_FOLDER);
-    StoragePath hashingConfigPath =
-        new StoragePath(hashingBase,
-            instantToLoad + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
-
+    StoragePath hashingConfigPath = PartitionBucketIndexUtils.getHashingConfigPath(basePath, instantToLoad);
     try (HoodieHadoopStorage storage = new HoodieHadoopStorage(hashingConfigPath, HadoopFSUtils.getStorageConf(hadoopConf))) {
       init(storage, hashingConfigPath);
     } catch (IOException e) {
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
index 82b61acd7c7..69e29bd92e3 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
@@ -49,8 +49,8 @@ public class PartitionBucketIndexUtils {
 
   private static final Logger LOG = LoggerFactory.getLogger(PartitionBucketIndexUtils.class);
 
-  public static boolean isPartitionSimpleBucketIndex(Configuration conf) {
-    StoragePath storagePath = new StoragePath(HoodieTableMetaClient.PARTITION_BUCKET_INDEX_HASHING_FOLDER);
+  public static boolean isPartitionSimpleBucketIndex(Configuration conf, String basePath) {
+    StoragePath storagePath = getHashingConfigStorageFolder(basePath);
     try (HoodieHadoopStorage storage = new HoodieHadoopStorage(storagePath, HadoopFSUtils.getStorageConf(conf))) {
       return storage.exists(storagePath);
     } catch (IOException e) {
@@ -58,6 +58,16 @@ public class PartitionBucketIndexUtils {
     }
   }
 
+  public static StoragePath getHashingConfigStorageFolder(String basePath) {
+    StoragePath metaPath = new StoragePath(basePath, HoodieTableMetaClient.METAFOLDER_NAME);
+    return new StoragePath(metaPath, HoodieTableMetaClient.BUCKET_INDEX_METAFOLDER_CONFIG_FOLDER);
+  }
+
+  public static StoragePath getHashingConfigPath(String basePath, String instantToLoad) {
+    StoragePath hashingBase = getHashingConfigStorageFolder(basePath);
+    return new StoragePath(hashingBase, instantToLoad + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
+  }
+
   public static boolean initHashingConfig(HoodieTableMetaClient metaClient,
                                           String expressions,
                                           String rule,
@@ -104,7 +114,6 @@ public class PartitionBucketIndexUtils {
     Option<HoodieInstant> earliestInstant = metaClient.getActiveTimeline().getCommitsTimeline().filterCompletedInstants().firstInstant();
     String instantToLoad = "";
     try {
-      // 按时间降序排列
       List<String> hashingConfigInstants = metaClient.getStorage()
           .listDirectEntries(new StoragePath(metaClient.getHashingMetadataConfigPath())).stream().map(info -> {
             String instant = getHashingConfigInstant(info.getPath().getName());
@@ -115,11 +124,15 @@ public class PartitionBucketIndexUtils {
           }).sorted().collect(Collectors.toList());
 
       for (String instant : hashingConfigInstants) {
-        if (instants.contains(instant)) {
+        if (!earliestInstant.isPresent()) {
           instantToLoad = instant;
           break;
-        } else if (earliestInstant.isPresent() && instant.compareTo(earliestInstant.get().requestedTime()) < 0){
+        } else if (instants.contains(instant)) {
           instantToLoad = instant;
+          break;
+        } else if (instant.compareTo(earliestInstant.get().requestedTime()) < 0){
+          instantToLoad = instant;
+          break;
         }
       }
 
@@ -133,16 +146,11 @@ public class PartitionBucketIndexUtils {
     }
   }
 
-  public static String getHashingConfigInstant(String hashingConfig) {
-    int lastIndex = hashingConfig.lastIndexOf('/');
-    if (lastIndex == -1) {
-      return null;
-    }
-    String fileName = hashingConfig.substring(lastIndex + 1);
-    int dotIndex = fileName.indexOf('.');
+  public static String getHashingConfigInstant(String hashingConfigName) {
+    int dotIndex = hashingConfigName.indexOf('.');
     if (dotIndex == -1) {
       return null;
     }
-    return fileName.substring(0, dotIndex);
+    return hashingConfigName.substring(0, dotIndex);
   }
 }
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
index ef32d24deb8..e80c8bfac0e 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/BucketIndexBulkInsertPartitionerWithRows.java
@@ -18,7 +18,10 @@
 
 package org.apache.hudi.execution.bulkinsert;
 
+import org.apache.hudi.common.util.StringUtils;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
 import org.apache.hudi.table.BulkInsertPartitioner;
+import org.apache.hudi.table.HoodieTable;
 
 import org.apache.spark.sql.BucketPartitionUtils$;
 import org.apache.spark.sql.Dataset;
@@ -31,15 +34,22 @@ public class BucketIndexBulkInsertPartitionerWithRows implements BulkInsertParti
 
   private final String indexKeyFields;
   private final int bucketNum;
+  private boolean isPartitionBucketIndexEnable = false;
+  private PartitionBucketIndexCalculator calc;
 
-  public BucketIndexBulkInsertPartitionerWithRows(String indexKeyFields, int bucketNum) {
+  public BucketIndexBulkInsertPartitionerWithRows(String indexKeyFields, int bucketNum, HoodieTable table) {
     this.indexKeyFields = indexKeyFields;
     this.bucketNum = bucketNum;
+    String hashingInstantToLoad = table.getConfig().getHashingConfigInstantToLoad();
+    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    if (isPartitionBucketIndexEnable) {
+      calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, table.getMetaClient());
+    }
   }
 
   @Override
   public Dataset<Row> repartitionRecords(Dataset<Row> rows, int outputPartitions) {
-    return BucketPartitionUtils$.MODULE$.createDataFrame(rows, indexKeyFields, bucketNum, outputPartitions);
+    return BucketPartitionUtils$.MODULE$.createDataFrame(rows, indexKeyFields, bucketNum, outputPartitions, calc, isPartitionBucketIndexEnable);
   }
 
   @Override
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
index 8304c1031cb..d651342c368 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
@@ -23,9 +23,11 @@ import org.apache.hudi.common.model.HoodieKey;
 import org.apache.hudi.common.model.HoodieRecord;
 import org.apache.hudi.common.model.HoodieRecordLocation;
 import org.apache.hudi.common.model.HoodieRecordPayload;
+import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.common.util.ValidationUtils;
 import org.apache.hudi.index.bucket.BucketIdentifier;
 import org.apache.hudi.index.bucket.HoodieSimpleBucketIndex;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
 import org.apache.hudi.table.HoodieTable;
 
 import org.apache.spark.Partitioner;
@@ -39,11 +41,18 @@ import java.util.stream.Collectors;
 public class RDDSimpleBucketBulkInsertPartitioner<T extends HoodieRecordPayload> extends RDDBucketIndexPartitioner<T> {
 
   private final boolean isNonBlockingConcurrencyControl;
+  private boolean isPartitionBucketIndexEnable = false;
+  private PartitionBucketIndexCalculator calc;
 
   public RDDSimpleBucketBulkInsertPartitioner(HoodieTable table) {
     super(table, null, false);
     ValidationUtils.checkArgument(table.getIndex() instanceof HoodieSimpleBucketIndex);
     this.isNonBlockingConcurrencyControl = table.getConfig().isNonBlockingConcurrencyControl();
+    String hashingInstantToLoad = table.getConfig().getHashingConfigInstantToLoad();
+    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    if (isPartitionBucketIndexEnable) {
+      calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, table.getMetaClient());
+    }
   }
 
   @Override
@@ -64,7 +73,7 @@ public class RDDSimpleBucketBulkInsertPartitioner<T extends HoodieRecordPayload>
       public int getPartition(Object key) {
         HoodieKey hoodieKey = (HoodieKey) key;
         String partitionPath = hoodieKey.getPartitionPath();
-        int bucketID = index.getBucketID(hoodieKey);
+        int bucketID = isPartitionBucketIndexEnable ? index.getBucketID(hoodieKey, calc.computeNumBuckets(partitionPath)) : index.getBucketID(hoodieKey);
         String fileID = partitionMapper.get(partitionPath).get(bucketID);
         return fileIdPrefixToBucketIndex.get(fileID);
       }
@@ -75,11 +84,11 @@ public class RDDSimpleBucketBulkInsertPartitioner<T extends HoodieRecordPayload>
                                                        Map<String, Integer> fileIdPrefixToBucketIndex) {
 
     HoodieSimpleBucketIndex index = (HoodieSimpleBucketIndex) table.getIndex();
-    int numBuckets = index.getNumBuckets();
     return records
         .map(HoodieRecord::getPartitionPath)
         .distinct().collect().stream()
         .collect(Collectors.toMap(p -> p, p -> {
+          // p ==> partition path
           Map<Integer, HoodieRecordLocation> locationMap = index.loadBucketIdToFileIdMappingForPartition(table, p);
           Map<Integer, String> bucketIdToFileIdPrefixMap = new HashMap<>();
           HashSet<Integer> existsBucketID = new HashSet<>();
@@ -94,6 +103,7 @@ public class RDDSimpleBucketBulkInsertPartitioner<T extends HoodieRecordPayload>
             doAppend.add(true);
           });
 
+          int numBuckets = isPartitionBucketIndexEnable ? calc.computeNumBuckets(p) : index.getNumBuckets();
           // Generate a file that does not exist
           for (int i = 0; i < numBuckets; i++) {
             if (!existsBucketID.contains(i)) {
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java
index 68123c02b79..d2b6a04c2e1 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/BucketBulkInsertDataInternalWriterHelper.java
@@ -19,10 +19,12 @@
 package org.apache.hudi.table.action.commit;
 
 import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.common.util.collection.Pair;
 import org.apache.hudi.config.HoodieIndexConfig;
 import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.index.bucket.BucketIdentifier;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
 import org.apache.hudi.io.storage.row.HoodieRowCreateHandle;
 import org.apache.hudi.keygen.constant.KeyGeneratorOptions;
 import org.apache.hudi.table.HoodieTable;
@@ -44,6 +46,8 @@ import java.util.Objects;
 public class BucketBulkInsertDataInternalWriterHelper extends BulkInsertDataInternalWriterHelper {
 
   private static final Logger LOG = LoggerFactory.getLogger(BucketBulkInsertDataInternalWriterHelper.class);
+  private final boolean isPartitionBucketIndexEnable;
+  private PartitionBucketIndexCalculator calc;
 
   private Pair<UTF8String, Integer> lastFileId; // for efficient code path
   // p -> (fileId -> handle)
@@ -66,13 +70,18 @@ public class BucketBulkInsertDataInternalWriterHelper extends BulkInsertDataInte
     this.bucketNum = writeConfig.getInt(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS);
     this.handles = new HashMap<>();
     this.isNonBlockingConcurrencyControl = writeConfig.isNonBlockingConcurrencyControl();
+    String hashingInstantToLoad = writeConfig.getHashingConfigInstantToLoad();
+    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    if (isPartitionBucketIndexEnable) {
+      calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, hoodieTable.getMetaClient());
+    }
   }
 
   public void write(InternalRow row) throws IOException {
     try {
       UTF8String partitionPath = extractPartitionPath(row);
       UTF8String recordKey = extractRecordKey(row);
-      int bucketId = BucketIdentifier.getBucketId(String.valueOf(recordKey), indexKeyFields, bucketNum);
+      int bucketId = BucketIdentifier.getBucketId(String.valueOf(recordKey), indexKeyFields, isPartitionBucketIndexEnable ? calc.computeNumBuckets(partitionPath.toString()) : bucketNum);
       if (lastFileId == null || !Objects.equals(lastFileId.getKey(), partitionPath) || !Objects.equals(lastFileId.getValue(), bucketId)) {
         // NOTE: It's crucial to make a copy here, since [[UTF8String]] could be pointing into
         //       a mutable underlying buffer
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java
index b7cdef72ab1..c760d2aabdd 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/table/action/commit/SparkBucketIndexPartitioner.java
@@ -23,12 +23,14 @@ import org.apache.hudi.common.model.HoodieKey;
 import org.apache.hudi.common.model.HoodieRecordLocation;
 import org.apache.hudi.common.model.WriteOperationType;
 import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.common.util.ValidationUtils;
 import org.apache.hudi.common.util.collection.Pair;
 import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.exception.HoodieException;
 import org.apache.hudi.index.bucket.BucketIdentifier;
 import org.apache.hudi.index.bucket.HoodieBucketIndex;
+import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
 import org.apache.hudi.table.HoodieTable;
 import org.apache.hudi.table.WorkloadProfile;
 import org.apache.hudi.table.WorkloadStat;
@@ -49,11 +51,17 @@ import static org.apache.hudi.common.model.WriteOperationType.INSERT_OVERWRITE_T
 
 /**
  * Packs incoming records to be inserted into buckets (1 bucket = 1 RDD partition).
+ *
+ * For Partition Level bucket index
+ * Creates direct lookup arrays during initialization, Provides O(1) lookup time for partition path and bucket ID,
+ * The trade-off is a small amount of additional memory usage for the lookup arrays.
  */
-public class SparkBucketIndexPartitioner<T> extends
-    SparkHoodiePartitioner<T> {
+public class SparkBucketIndexPartitioner<T> extends SparkHoodiePartitioner<T> {
 
-  private final int numBuckets;
+  private final int totalPartitions;
+  private boolean isPartitionBucketIndexEnable = false;
+  private PartitionBucketIndexCalculator calc;
+  private int numBuckets;
   private final String indexKeyField;
   private final int totalPartitionPaths;
   private final List<String> partitionPaths;
@@ -70,6 +78,15 @@ public class SparkBucketIndexPartitioner<T> extends
   private Map<String, Set<String>> updatePartitionPathFileIds;
 
   private final boolean isNonBlockingConcurrencyControl;
+  /**
+   * Direct mapping from partition number to partition path
+   */
+  private String[] partitionNumberToPath;
+
+  /**
+   * Direct mapping from partition number to local bucket ID
+   */
+  private Integer[] partitionNumberToLocalBucketId;
 
   public SparkBucketIndexPartitioner(WorkloadProfile profile,
                                      HoodieEngineContext context,
@@ -81,6 +98,11 @@ public class SparkBucketIndexPartitioner<T> extends
           " Bucket index partitioner should only be used by BucketIndex other than "
               + table.getIndex().getClass().getSimpleName());
     }
+    String hashingInstantToLoad = table.getConfig().getHashingConfigInstantToLoad();
+    this.isPartitionBucketIndexEnable = StringUtils.isNullOrEmpty(hashingInstantToLoad);
+    if (isPartitionBucketIndexEnable) {
+      calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, table.getMetaClient());
+    }
     this.numBuckets = ((HoodieBucketIndex) table.getIndex()).getNumBuckets();
     this.indexKeyField = config.getBucketIndexHashField();
     this.totalPartitionPaths = profile.getPartitionPaths().size();
@@ -89,12 +111,33 @@ public class SparkBucketIndexPartitioner<T> extends
     int i = 0;
     for (Object partitionPath : profile.getPartitionPaths()) {
       partitionPathOffset.put(partitionPath.toString(), i);
-      i += numBuckets;
+      if (isPartitionBucketIndexEnable) {
+        i += calc.computeNumBuckets(partitionPath.toString());
+      } else {
+        i += numBuckets;
+      }
     }
+    this.totalPartitions = i;
     assignUpdates(profile);
     WriteOperationType operationType = profile.getOperationType();
     this.isOverwrite = INSERT_OVERWRITE.equals(operationType) || INSERT_OVERWRITE_TABLE.equals(operationType);
     this.isNonBlockingConcurrencyControl = config.isNonBlockingConcurrencyControl();
+
+    if (isPartitionBucketIndexEnable) {
+      this.partitionNumberToPath = new String[totalPartitions];
+      this.partitionNumberToLocalBucketId = new Integer[totalPartitions];
+
+      for (String partitionPath : partitionPaths) {
+        int offset = partitionPathOffset.get(partitionPath);
+        int numBuckets = calc.computeNumBuckets(partitionPath);
+
+        for (int j = 0; j < numBuckets; j++) {
+          int partitionNumber = offset + j;
+          partitionNumberToPath[partitionNumber] = partitionPath;
+          partitionNumberToLocalBucketId[partitionNumber] = j;
+        }
+      }
+    }
   }
 
   private void assignUpdates(WorkloadProfile profile) {
@@ -115,33 +158,49 @@ public class SparkBucketIndexPartitioner<T> extends
 
   @Override
   public BucketInfo getBucketInfo(int bucketNumber) {
-    String partitionPath = partitionPaths.get(bucketNumber / numBuckets);
+    Pair<Integer, String> res = computeBucketAndPartitionPath(bucketNumber);
+    int bucket = res.getLeft();
+    String partitionPath = res.getRight();
     // Insert overwrite always generates new bucket file id
     if (isOverwrite) {
       ValidationUtils.checkArgument(!isNonBlockingConcurrencyControl,
           "Insert overwrite is not supported with non-blocking concurrency control");
-      return new BucketInfo(BucketType.INSERT, BucketIdentifier.newBucketFileIdPrefix(bucketNumber % numBuckets), partitionPath);
+      return new BucketInfo(BucketType.INSERT, BucketIdentifier.newBucketFileIdPrefix(bucket), partitionPath);
     }
     Option<String> fileIdOption = Option.fromJavaOptional(updatePartitionPathFileIds
         .getOrDefault(partitionPath, Collections.emptySet()).stream()
-        .filter(e -> e.startsWith(BucketIdentifier.bucketIdStr(bucketNumber % numBuckets)))
+        .filter(e -> e.startsWith(BucketIdentifier.bucketIdStr(bucket)))
         .findFirst());
     if (fileIdOption.isPresent()) {
       return new BucketInfo(BucketType.UPDATE, fileIdOption.get(), partitionPath);
     } else {
       // Always write into log file instead of base file if using NB-CC
       if (isNonBlockingConcurrencyControl) {
-        String fileId = BucketIdentifier.newBucketFileIdForNBCC(bucketNumber % numBuckets);
+        String fileId = BucketIdentifier.newBucketFileIdForNBCC(bucket);
         return new BucketInfo(BucketType.UPDATE, fileId, partitionPath);
       }
-      String fileIdPrefix = BucketIdentifier.newBucketFileIdPrefix(bucketNumber % numBuckets);
+      String fileIdPrefix = BucketIdentifier.newBucketFileIdPrefix(bucket);
       return new BucketInfo(BucketType.INSERT, fileIdPrefix, partitionPath);
     }
   }
 
+  private Pair<Integer, String> computeBucketAndPartitionPath(int bucketNumber) {
+    Integer bucket;
+    String partitionPath;
+    if (isPartitionBucketIndexEnable) {
+      bucket = partitionNumberToLocalBucketId[bucketNumber];
+      partitionPath = partitionNumberToPath[bucketNumber];
+      ValidationUtils.checkArgument(bucket != null && partitionPath != null);
+    } else {
+      bucket = bucketNumber % numBuckets;
+      partitionPath = partitionPaths.get(bucketNumber / numBuckets);
+    }
+    return Pair.of(bucket, partitionPath);
+  }
+
   @Override
   public int numPartitions() {
-    return totalPartitionPaths * numBuckets;
+    return totalPartitions;
   }
 
   @Override
@@ -151,7 +210,8 @@ public class SparkBucketIndexPartitioner<T> extends
     Option<HoodieRecordLocation> location = keyLocation._2;
     int bucketId = location.isPresent()
         ? BucketIdentifier.bucketIdFromFileId(location.get().getFileId())
-        : BucketIdentifier.getBucketId(keyLocation._1.getRecordKey(), indexKeyField, numBuckets);
+        : BucketIdentifier.getBucketId(keyLocation._1.getRecordKey(), indexKeyField,
+        isPartitionBucketIndexEnable ? calc.computeNumBuckets(partitionPath) : numBuckets);
     return partitionPathOffset.get(partitionPath) + bucketId;
   }
 }
diff --git a/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala b/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
index 740e6799348..464438671c1 100644
--- a/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
+++ b/hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala
@@ -21,16 +21,27 @@ package org.apache.spark.sql
 import org.apache.hudi.common.model.HoodieRecord
 import org.apache.hudi.common.util.Functions
 import org.apache.hudi.common.util.hash.BucketIndexUtil
-import org.apache.hudi.index.bucket.BucketIdentifier
+import org.apache.hudi.index.bucket.{BucketIdentifier, PartitionBucketIndexCalculator}
 import org.apache.spark.Partitioner
 import org.apache.spark.sql.catalyst.InternalRow
 
 object BucketPartitionUtils {
-  def createDataFrame(df: DataFrame, indexKeyFields: String, bucketNum: Int, partitionNum: Int): DataFrame = {
+  def createDataFrame(df: DataFrame, indexKeyFields: String, bucketNum: Int, partitionNum: Int,
+                      calc: PartitionBucketIndexCalculator, isPartitionBucketIndexEnable: Boolean): DataFrame = {
+
+    def computeBucketNumber(): String => Integer = partition => {
+      val bucketNumber = if (isPartitionBucketIndexEnable) {
+        calc.computeNumBuckets(partition)
+      } else {
+        bucketNum
+      }
+      bucketNumber
+    }
+
     def getPartitionKeyExtractor(): InternalRow => (String, Int) = row => {
-      val kb = BucketIdentifier
-        .getBucketId(row.getString(HoodieRecord.RECORD_KEY_META_FIELD_ORD), indexKeyFields, bucketNum)
       val partition = row.getString(HoodieRecord.PARTITION_PATH_META_FIELD_ORD)
+      val kb = BucketIdentifier
+        .getBucketId(row.getString(HoodieRecord.RECORD_KEY_META_FIELD_ORD), indexKeyFields, computeBucketNumber().apply(partition))
       if (partition == null || partition.trim.isEmpty) {
         ("", kb)
       } else {
@@ -48,7 +59,7 @@ object BucketPartitionUtils {
 
       override def getPartition(value: Any): Int = {
         val partitionKeyPair = value.asInstanceOf[(String, Int)]
-        partitionIndexFunc.apply(bucketNum, partitionKeyPair._1, partitionKeyPair._2)
+        partitionIndexFunc.apply(computeBucketNumber().apply(partitionKeyPair._1), partitionKeyPair._1, partitionKeyPair._2)
       }
     }
     // use internalRow to avoid extra convert.
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java
index 66d2e08417e..e2e2e6c2d55 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/configuration/OptionsInference.java
@@ -99,7 +99,7 @@ public class OptionsInference {
   public static void setupIndexConfigs(Configuration conf) {
     HoodieIndex.BucketIndexEngineType engineType = OptionsResolver.getBucketEngineType(conf);
     if (engineType.equals(HoodieIndex.BucketIndexEngineType.SIMPLE) &&
-        PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(HadoopConfigurations.getHadoopConf(conf))) {
+        PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(HadoopConfigurations.getHadoopConf(conf), conf.get(FlinkOptions.PATH))) {
       try (HoodieFlinkWriteClient writeClient = FlinkWriteClients.createWriteClientV2(conf)) {
         HoodieTableMetaClient metaClient = writeClient.getHoodieTable().getMetaClient();
         String hashingConfigToLoad = PartitionBucketIndexUtils.getHashingConfigInstantToLoad(metaClient);
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java
index 8369a0b768e..c3fc9a2e1af 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketBulkInsertWriterHelper.java
@@ -27,6 +27,7 @@ import org.apache.hudi.io.storage.row.HoodieRowDataCreateHandle;
 import org.apache.hudi.sink.bulk.BulkInsertWriterHelper;
 import org.apache.hudi.sink.bulk.RowDataKeyGen;
 import org.apache.hudi.sink.bulk.sort.SortOperatorGen;
+import org.apache.hudi.storage.StorageConfiguration;
 import org.apache.hudi.table.HoodieTable;
 import org.apache.hudi.util.Lazy;
 
@@ -97,13 +98,13 @@ public class BucketBulkInsertWriterHelper extends BulkInsertWriterHelper {
   }
 
   private static String getFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, String indexKeys, Configuration conf, boolean needFixedFileIdSuffix,
-                                  Lazy<org.apache.hadoop.conf.Configuration> hadoopConf) {
+                                  StorageConfiguration<org.apache.hadoop.conf.Configuration> storageConf) {
     String recordKey = keyGen.getRecordKey(record);
     String partition = keyGen.getPartitionPath(record);
     final int numBuckets;
     if (!StringUtils.isNullOrEmpty(conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT))) {
-      numBuckets = PartitionBucketIndexCalculator.getInstance(conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT), hadoopConf.get(), conf.get(FlinkOptions.PATH))
-          .computeNumBuckets(partition);
+      numBuckets = PartitionBucketIndexCalculator.getInstance(conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT),
+              storageConf.unwrapAs(org.apache.hadoop.conf.Configuration.class), conf.get(FlinkOptions.PATH)).computeNumBuckets(partition);
     } else {
       numBuckets = conf.get(FlinkOptions.BUCKET_INDEX_NUM_BUCKETS);
     }
@@ -113,8 +114,8 @@ public class BucketBulkInsertWriterHelper extends BulkInsertWriterHelper {
   }
 
   public static RowData rowWithFileId(Map<String, String> bucketIdToFileId, RowDataKeyGen keyGen, RowData record, String indexKeys, Configuration conf, boolean needFixedFileIdSuffix,
-                                      Lazy<org.apache.hadoop.conf.Configuration> hadoopConf) {
-    final String fileId = getFileId(bucketIdToFileId, keyGen, record, indexKeys, conf, needFixedFileIdSuffix, hadoopConf);
+                                      StorageConfiguration<org.apache.hadoop.conf.Configuration> storageConf) {
+    final String fileId = getFileId(bucketIdToFileId, keyGen, record, indexKeys, conf, needFixedFileIdSuffix, storageConf);
     return GenericRowData.of(StringData.fromString(fileId), record);
   }
 
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java
index 2a3e1e154dd..9cc23e1d0e0 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/bucket/BucketStreamWriteFunction.java
@@ -139,9 +139,7 @@ public class BucketStreamWriteFunction extends StreamWriteFunction {
       bootstrapIndexIfNeed(partition);
     }
     Map<Integer, String> bucketToFileId = bucketIndex.computeIfAbsent(partition, p -> new HashMap<>());
-
     int numBuckets = getOrComputeNumBuckets(partition);
-
     final int bucketNum = BucketIdentifier.getBucketId(record.getRecordKey(), indexKeyFields, numBuckets);
     final String bucketId = partition + "/" + bucketNum;
 
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java
index ecf4656d7ee..ca7cf4023c5 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/partitioner/BucketIndexPartitioner.java
@@ -25,7 +25,7 @@ import org.apache.hudi.common.util.hash.BucketIndexUtil;
 import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.index.bucket.BucketIdentifier;
 import org.apache.hudi.index.bucket.PartitionBucketIndexCalculator;
-import org.apache.hudi.util.Lazy;
+import org.apache.hudi.storage.StorageConfiguration;
 
 import org.apache.flink.api.common.functions.Partitioner;
 import org.apache.flink.configuration.Configuration;
@@ -41,17 +41,16 @@ public class BucketIndexPartitioner<T extends HoodieKey> implements Partitioner<
   private final Configuration conf;
   private final String indexKeyFields;
   private final String hashingInstantToLoad;
-  private org.apache.hadoop.conf.Configuration hadoopConf;
+  private StorageConfiguration<org.apache.hadoop.conf.Configuration> storageConf;
 
   private Functions.Function3<Integer, String, Integer, Integer> partitionIndexFunc;
 
-  public BucketIndexPartitioner(Configuration conf, String indexKeyFields, Lazy<org.apache.hadoop.conf.Configuration> hadoopConf) {
+  public BucketIndexPartitioner(Configuration conf, String indexKeyFields,
+                                StorageConfiguration<org.apache.hadoop.conf.Configuration> storageConf) {
     this.conf = conf;
     this.indexKeyFields = indexKeyFields;
     this.hashingInstantToLoad = conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT);
-    if (!StringUtils.isNullOrEmpty(hashingInstantToLoad)) {
-      this.hadoopConf = hadoopConf.get();
-    }
+    this.storageConf = storageConf;
   }
 
   @Override
@@ -62,6 +61,7 @@ public class BucketIndexPartitioner<T extends HoodieKey> implements Partitioner<
 
     int bucketNum;
     if (!StringUtils.isNullOrEmpty(hashingInstantToLoad)) {
+      org.apache.hadoop.conf.Configuration hadoopConf = storageConf.unwrapAs(org.apache.hadoop.conf.Configuration.class);
       PartitionBucketIndexCalculator calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, hadoopConf, conf.get(FlinkOptions.PATH));
       bucketNum = calc.computeNumBuckets(key.getPartitionPath());
     } else {
diff --git a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java
index d604567246d..d278d807064 100644
--- a/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java
+++ b/hudi-flink-datasource/hudi-flink/src/main/java/org/apache/hudi/sink/utils/Pipelines.java
@@ -53,8 +53,9 @@ import org.apache.hudi.sink.compact.CompactionPlanOperator;
 import org.apache.hudi.sink.partitioner.BucketAssignFunction;
 import org.apache.hudi.sink.partitioner.BucketIndexPartitioner;
 import org.apache.hudi.sink.transform.RowDataToHoodieFunctions;
+import org.apache.hudi.storage.StorageConfiguration;
+import org.apache.hudi.storage.hadoop.HadoopStorageConfiguration;
 import org.apache.hudi.table.format.FilePathUtils;
-import org.apache.hudi.util.Lazy;
 
 import org.apache.flink.api.common.functions.Partitioner;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
@@ -122,8 +123,9 @@ public class Pipelines {
             "Consistent hashing bucket index does not work with bulk insert using FLINK engine. Use simple bucket index or Spark engine.");
       }
       String indexKeys = OptionsResolver.getIndexKeyField(conf);
-      Lazy<org.apache.hadoop.conf.Configuration> hadoopConf = Lazy.lazily(() -> HadoopConfigurations.getHadoopConf(conf));
-      BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(conf, indexKeys, hadoopConf);
+      StorageConfiguration<org.apache.hadoop.conf.Configuration> storageConf = StringUtils.isNullOrEmpty(conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT)) ?
+          null : new HadoopStorageConfiguration(HadoopConfigurations.getHadoopConf(conf));
+      BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(conf, indexKeys, storageConf);
       RowDataKeyGen keyGen = RowDataKeyGen.instance(conf, rowType);
       RowType rowTypeWithFileId = BucketBulkInsertWriterHelper.rowTypeWithFileId(rowType);
       InternalTypeInfo<RowData> typeInfo = InternalTypeInfo.of(rowTypeWithFileId);
@@ -131,7 +133,7 @@ public class Pipelines {
 
       Map<String, String> bucketIdToFileId = new HashMap<>();
       dataStream = dataStream.partitionCustom(partitioner, keyGen::getHoodieKey)
-          .map(record -> BucketBulkInsertWriterHelper.rowWithFileId(bucketIdToFileId, keyGen, record, indexKeys, conf, needFixedFileIdSuffix, hadoopConf), typeInfo)
+          .map(record -> BucketBulkInsertWriterHelper.rowWithFileId(bucketIdToFileId, keyGen, record, indexKeys, conf, needFixedFileIdSuffix, storageConf), typeInfo)
           .setParallelism(PARALLELISM_VALUE);
       if (conf.getBoolean(FlinkOptions.WRITE_BULK_INSERT_SORT_INPUT)) {
         SortOperatorGen sortOperatorGen = BucketBulkInsertWriterHelper.getFileIdSorterGen(rowTypeWithFileId);
@@ -341,7 +343,9 @@ public class Pipelines {
           String indexKeyFields = OptionsResolver.getIndexKeyField(conf);
           // [HUDI-9036] BucketIndexPartitioner is also used in bulk insert mode,
           // keep use of HoodieKey here in partitionCustom for now
-          BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(conf, indexKeyFields, Lazy.lazily(() -> HadoopConfigurations.getHadoopConf(conf)));
+          StorageConfiguration<org.apache.hadoop.conf.Configuration> storageConf = StringUtils.isNullOrEmpty(conf.get(FlinkOptions.BUCKET_INDEX_PARTITION_LOAD_INSTANT)) ?
+              null : new HadoopStorageConfiguration(HadoopConfigurations.getHadoopConf(conf));
+          BucketIndexPartitioner<HoodieKey> partitioner = new BucketIndexPartitioner<>(conf, indexKeyFields, storageConf);
           return dataStream
               .partitionCustom(
                   partitioner,
@@ -505,6 +509,8 @@ public class Pipelines {
         : databaseName + "." + conf.getString(FlinkOptions.TABLE_NAME);
   }
 
+
+
   /**
    * Dummy sink that does nothing.
    */
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java
index a4da9049dce..eb6a8e3c25d 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/sink/utils/BulkInsertFunctionWrapper.java
@@ -32,6 +32,7 @@ import org.apache.hudi.sink.bulk.sort.SortOperator;
 import org.apache.hudi.sink.bulk.sort.SortOperatorGen;
 import org.apache.hudi.sink.common.AbstractWriteFunction;
 import org.apache.hudi.sink.event.WriteMetadataEvent;
+import org.apache.hudi.storage.hadoop.HadoopStorageConfiguration;
 import org.apache.hudi.util.AvroSchemaConverter;
 import org.apache.hudi.util.Lazy;
 import org.apache.hudi.util.StreamerUtil;
@@ -221,7 +222,7 @@ public class BulkInsertFunctionWrapper<I> implements TestFunctionWrapper<I> {
     boolean needFixedFileIdSuffix = OptionsResolver.isNonBlockingConcurrencyControl(conf);
     this.bucketIdToFileId = new HashMap<>();
     this.mapFunction = r -> BucketBulkInsertWriterHelper.rowWithFileId(bucketIdToFileId, keyGen, r, indexKeys, conf, needFixedFileIdSuffix,
-        Lazy.lazily(() -> HadoopConfigurations.getHadoopConf(conf)));
+        new HadoopStorageConfiguration(HadoopConfigurations.getHadoopConf(conf)));
   }
 
   private void setupSortOperator() throws Exception {
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 a2a31b492e9..17ff21c3938 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
@@ -23,13 +23,16 @@ import org.apache.hudi.common.model.DefaultHoodieRecordPayload;
 import org.apache.hudi.common.model.HoodieTableType;
 import org.apache.hudi.common.model.WriteOperationType;
 import org.apache.hudi.common.table.HoodieTableConfig;
+import org.apache.hudi.common.table.HoodieTableMetaClient;
 import org.apache.hudi.common.table.cdc.HoodieCDCSupplementalLoggingMode;
 import org.apache.hudi.common.table.timeline.HoodieTimeline;
+import org.apache.hudi.common.table.view.HoodieTableFileSystemView;
 import org.apache.hudi.common.util.CollectionUtils;
 import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.configuration.FlinkOptions;
 import org.apache.hudi.metadata.HoodieTableMetadata;
 import org.apache.hudi.storage.StoragePath;
+import org.apache.hudi.storage.StoragePathInfo;
 import org.apache.hudi.table.catalog.HoodieCatalogTestUtils;
 import org.apache.hudi.table.catalog.HoodieHiveCatalog;
 import org.apache.hudi.util.StreamerUtil;
@@ -66,10 +69,12 @@ import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 
 import java.io.File;
+import java.io.IOException;
 import java.time.Instant;
 import java.time.LocalDateTime;
 import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
@@ -89,6 +94,7 @@ import static org.apache.hudi.utils.TestData.assertRowsEquals;
 import static org.apache.hudi.utils.TestData.assertRowsEqualsUnordered;
 import static org.apache.hudi.utils.TestData.map;
 import static org.apache.hudi.utils.TestData.row;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -1258,6 +1264,78 @@ public class ITTestHoodieDataSource {
         + "+I[id8, Han, 56, 1970-01-01T00:00:08, par4]]");
   }
 
+  @ParameterizedTest
+  @ValueSource(strings = {"bulk_insert", "upsert"})
+  void testBulkInsertWithPartitionBucketIndex(String operationType) throws IOException {
+    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.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(
+        () -> tableEnv.sqlQuery("select * from " + catalogName + ".hudi.hoodie_sink").execute().collect());
+    assertRowsEquals(result1, TestData.DATA_SET_SOURCE_INSERT);
+    // apply filters
+    List<Row> result2 = CollectionUtil.iterableToList(
+        () -> tableEnv.sqlQuery("select * from " + catalogName + ".hudi.hoodie_sink where uuid > 'id5'").execute().collect());
+    assertRowsEquals(result2, "["
+        + "+I[id6, Emma, 20, 1970-01-01T00:00:06, par3], "
+        + "+I[id7, Bob, 44, 1970-01-01T00:00:07, par4], "
+        + "+I[id8, Han, 56, 1970-01-01T00:00:08, par4]]");
+    HoodieTableMetaClient metaClient = StreamerUtil.createMetaClient(basePath, new org.apache.hadoop.conf.Configuration());
+
+    List<StoragePathInfo> allFiles = metaClient.getStorage().listDirectEntries(new StoragePath(basePath)).stream().flatMap(path -> {
+      try {
+        return metaClient.getStorage().listDirectEntries(path.getPath()).stream();
+      } catch (IOException e) {
+        return Stream.empty();
+      }
+    }).collect(Collectors.toList());
+
+    HoodieTableFileSystemView fsView = new HoodieTableFileSystemView(metaClient,
+        metaClient.getActiveTimeline().getCommitsTimeline().filterCompletedInstants(), allFiles);
+    List<String> actual = fsView.getAllFileGroups().map(group -> {
+      return group.getPartitionPath() + group.getFileGroupId().getFileId().split("-")[0];
+    }).collect(Collectors.toList());
+
+    // based on expression partition=(par1|par2),2 and default bucket number 1
+    // par1 and par2 have two buckets.
+    // par3 and par4 have one bucket.
+    ArrayList<String> expected = new ArrayList<>();
+    expected.add("partition=par1" + "00000000");
+    expected.add("partition=par1" + "00000001");
+    expected.add("partition=par2" + "00000000");
+    expected.add("partition=par2" + "00000001");
+    expected.add("partition=par3" + "00000000");
+    expected.add("partition=par4" + "00000000");
+
+    assertEquals(expected.stream().sorted().collect(Collectors.toList()), actual.stream().sorted().collect(Collectors.toList()));
+  }
+
   @Test
   void testBulkInsertWithSortByRecordKey() {
     TableEnvironment tableEnv = batchTableEnv;
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/utils/TestConfigurations.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/utils/TestConfigurations.java
index 92d44a8e7ef..ad07f95c3b6 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/utils/TestConfigurations.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/utils/TestConfigurations.java
@@ -400,6 +400,9 @@ public class TestConfigurations {
       return TestConfigurations.getCreateHoodieTableDDL(this.tableName, this.fields, options,
           this.withPartition, this.pkField, this.partitionField);
     }
+    public Map<String, String> getOptions() {
+      return this.options;
+    }
   }
 
   /**
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java
index 9d643211c92..3bb0c2d585a 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/java/org/apache/hudi/commit/BaseDatasetBulkInsertCommitActionExecutor.java
@@ -120,7 +120,7 @@ public abstract class BaseDatasetBulkInsertCommitActionExecutor implements Seria
       if (writeConfig.getIndexType() == HoodieIndex.IndexType.BUCKET) {
         if (writeConfig.getBucketIndexEngineType() == HoodieIndex.BucketIndexEngineType.SIMPLE) {
           return new BucketIndexBulkInsertPartitionerWithRows(writeConfig.getBucketIndexHashFieldWithDefault(),
-              writeConfig.getBucketIndexNumBuckets());
+              writeConfig.getBucketIndexNumBuckets(), table);
         } else {
           return new ConsistentBucketIndexBulkInsertPartitionerWithRows(table, Collections.emptyMap(), true);
         }


(hudi) 06/06: finish Spark related conding and test

Posted by zh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zhangyue19921010 pushed a commit to branch HUDI-8990
in repository https://gitbox.apache.org/repos/asf/hudi.git

commit aa9ed6ce12c49d9c79034453a9bba3bfaf16a9d2
Author: YueZhang <zh...@163.com>
AuthorDate: Fri Mar 21 16:03:50 2025 +0800

    finish Spark related conding and test
---
 .../hudi/index/bucket/HoodieSimpleBucketIndex.java |  10 +-
 .../index/bucket/PartitionBucketIndexUtils.java    |  29 ++-
 .../RDDSimpleBucketBulkInsertPartitioner.java      |  10 +-
 .../hudi/HoodieDatasetBulkInsertHelper.scala       |  15 +-
 .../hudi/common/model/WriteOperationType.java      |   2 +
 .../apache/hudi/table/ITTestHoodieDataSource.java  |  16 +-
 .../org/apache/hudi/HoodieSparkSqlWriter.scala     |  19 +-
 .../spark/sql/hudi/ProvidesHoodieConfig.scala      |   6 +-
 .../hudi/command/procedures/HoodieProcedures.scala |   1 +
 .../procedures/PartitionBucketIndexManager.scala   | 153 ++++++++------
 .../spark/sql/hudi/dml/TestInsertTable.scala       |  88 --------
 .../TestInsertTableWithPartitionBucketIndex.scala  | 223 +++++++++++++++++++++
 .../TestPartitionBucketIndexManager.scala          |   2 +-
 13 files changed, 400 insertions(+), 174 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java
index e19e80b82a4..e3bae439956 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/HoodieSimpleBucketIndex.java
@@ -26,6 +26,7 @@ import org.apache.hudi.common.table.HoodieTableMetaClient;
 import org.apache.hudi.common.table.timeline.HoodieActiveTimeline;
 import org.apache.hudi.common.table.timeline.HoodieInstant;
 import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.StringUtils;
 import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.exception.HoodieIOException;
 import org.apache.hudi.index.HoodieIndexUtils;
@@ -156,14 +157,21 @@ public class HoodieSimpleBucketIndex extends HoodieBucketIndex {
 
   private class SimpleBucketIndexLocationFunction implements Function<HoodieRecord, Option<HoodieRecordLocation>> {
     private final Map<Integer, HoodieRecordLocation> bucketIdToFileIdMapping;
+    private final boolean isPartitionBucketIndexEnable;
+    private PartitionBucketIndexCalculator calc;
 
     public SimpleBucketIndexLocationFunction(HoodieTable table, String partitionPath) {
       this.bucketIdToFileIdMapping = loadBucketIdToFileIdMappingForPartition(table, partitionPath);
+      String hashingInstantToLoad = table.getConfig().getHashingConfigInstantToLoad();
+      this.isPartitionBucketIndexEnable = StringUtils.nonEmpty(hashingInstantToLoad);
+      if (isPartitionBucketIndexEnable) {
+        calc = PartitionBucketIndexCalculator.getInstance(hashingInstantToLoad, table.getMetaClient());
+      }
     }
 
     @Override
     public Option<HoodieRecordLocation> apply(HoodieRecord record) {
-      int bucketId = getBucketID(record.getKey());
+      int bucketId = isPartitionBucketIndexEnable ? getBucketID(record.getKey(), calc.computeNumBuckets(record.getPartitionPath())) : getBucketID(record.getKey());
       return Option.ofNullable(bucketIdToFileIdMapping.get(bucketId));
     }
   }
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
index ad1e1d0b3bf..1fbd16c28da 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/PartitionBucketIndexUtils.java
@@ -22,6 +22,7 @@ import org.apache.hudi.common.model.PartitionBucketIndexHashingConfig;
 import org.apache.hudi.common.table.HoodieTableMetaClient;
 import org.apache.hudi.common.table.timeline.HoodieInstant;
 import org.apache.hudi.common.table.timeline.HoodieTimeline;
+import org.apache.hudi.common.table.view.HoodieTableFileSystemView;
 import org.apache.hudi.common.util.FileIOUtils;
 import org.apache.hudi.common.util.Option;
 import org.apache.hudi.common.util.StringUtils;
@@ -30,6 +31,7 @@ import org.apache.hudi.exception.HoodieIOException;
 import org.apache.hudi.hadoop.fs.HadoopFSUtils;
 import org.apache.hudi.storage.HoodieInstantWriter;
 import org.apache.hudi.storage.HoodieStorage;
+import org.apache.hudi.storage.StorageConfiguration;
 import org.apache.hudi.storage.StoragePath;
 import org.apache.hudi.storage.StoragePathInfo;
 import org.apache.hudi.storage.hadoop.HoodieHadoopStorage;
@@ -44,6 +46,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public class PartitionBucketIndexUtils {
   public static final String INITIAL_HASHING_CONFIG_INSTANT = HoodieTimeline.INIT_INSTANT_TS;
@@ -51,8 +54,12 @@ public class PartitionBucketIndexUtils {
   private static final Logger LOG = LoggerFactory.getLogger(PartitionBucketIndexUtils.class);
 
   public static boolean isPartitionSimpleBucketIndex(Configuration conf, String basePath) {
+    return isPartitionSimpleBucketIndex(HadoopFSUtils.getStorageConf(conf), basePath);
+  }
+
+  public static boolean isPartitionSimpleBucketIndex(StorageConfiguration conf, String basePath) {
     StoragePath storagePath = getHashingConfigStorageFolder(basePath);
-    try (HoodieHadoopStorage storage = new HoodieHadoopStorage(storagePath, HadoopFSUtils.getStorageConf(conf))) {
+    try (HoodieHadoopStorage storage = new HoodieHadoopStorage(storagePath, conf)) {
       return storage.exists(storagePath);
     } catch (IOException e) {
       throw new HoodieIOException("Failed to list PARTITION_BUCKET_INDEX_HASHING_FOLDER folder ", e);
@@ -164,4 +171,24 @@ public class PartitionBucketIndexUtils {
     }
     return calc.getPartitionToBucket();
   }
+
+  /**
+   * Used for test
+   * @return all File id in current table using `partitionPath + FileId` format
+   */
+  public static List<String> getAllFileIDWithPartition(HoodieTableMetaClient metaClient) throws IOException {
+    List<StoragePathInfo> allFiles = metaClient.getStorage().listDirectEntries(metaClient.getBasePath()).stream().flatMap(path -> {
+      try {
+        return metaClient.getStorage().listDirectEntries(path.getPath()).stream();
+      } catch (IOException e) {
+        return Stream.empty();
+      }
+    }).collect(Collectors.toList());
+
+    HoodieTableFileSystemView fsView = new HoodieTableFileSystemView(metaClient,
+        metaClient.getActiveTimeline().getCommitsTimeline().filterCompletedInstants(), allFiles);
+    return fsView.getAllFileGroups().map(group -> {
+      return group.getPartitionPath() + group.getFileGroupId().getFileId().split("-")[0];
+    }).collect(Collectors.toList());
+  }
 }
diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
index 393f0054398..638abe95ce9 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/execution/bulkinsert/RDDSimpleBucketBulkInsertPartitioner.java
@@ -66,7 +66,7 @@ public class RDDSimpleBucketBulkInsertPartitioner<T extends HoodieRecordPayload>
     return doPartition(records, new Partitioner() {
       @Override
       public int numPartitions() {
-        return index.getNumBuckets() * partitionMapper.size();
+        return computeNumPartitions();
       }
 
       @Override
@@ -77,6 +77,14 @@ public class RDDSimpleBucketBulkInsertPartitioner<T extends HoodieRecordPayload>
         String fileID = partitionMapper.get(partitionPath).get(bucketID);
         return fileIdPrefixToBucketIndex.get(fileID);
       }
+
+      private int computeNumPartitions () {
+        if (isPartitionBucketIndexEnable) {
+          return partitionMapper.values().stream().mapToInt(Map::size).sum();
+        } else {
+          return index.getNumBuckets() * partitionMapper.size();
+        }
+      }
     });
   }
 
diff --git a/hudi-client/hudi-spark-client/src/main/scala/org/apache/hudi/HoodieDatasetBulkInsertHelper.scala b/hudi-client/hudi-spark-client/src/main/scala/org/apache/hudi/HoodieDatasetBulkInsertHelper.scala
index f4bb5b81e85..ab5cd537cfc 100644
--- a/hudi-client/hudi-spark-client/src/main/scala/org/apache/hudi/HoodieDatasetBulkInsertHelper.scala
+++ b/hudi-client/hudi-spark-client/src/main/scala/org/apache/hudi/HoodieDatasetBulkInsertHelper.scala
@@ -33,7 +33,7 @@ import org.apache.hudi.index.HoodieIndex.BucketIndexEngineType
 import org.apache.hudi.index.{HoodieIndex, SparkHoodieIndexFactory}
 import org.apache.hudi.keygen.factory.HoodieSparkKeyGeneratorFactory
 import org.apache.hudi.keygen.{AutoRecordGenWrapperKeyGenerator, BuiltinKeyGenerator, KeyGenUtils}
-import org.apache.hudi.table.action.commit.{BulkInsertDataInternalWriterHelper, ConsistentBucketBulkInsertDataInternalWriterHelper, ParallelismHelper}
+import org.apache.hudi.table.action.commit.{BucketBulkInsertDataInternalWriterHelper, BulkInsertDataInternalWriterHelper, ConsistentBucketBulkInsertDataInternalWriterHelper, ParallelismHelper}
 import org.apache.hudi.table.{BulkInsertPartitioner, HoodieTable}
 import org.apache.hudi.util.JFunction.toJavaSerializableFunctionUnchecked
 import org.apache.spark.TaskContext
@@ -179,6 +179,19 @@ object HoodieDatasetBulkInsertHelper
               writeConfig.populateMetaFields,
               arePartitionRecordsSorted,
               shouldPreserveHoodieMetadata)
+          case HoodieIndex.IndexType.BUCKET if writeConfig.getBucketIndexEngineType
+            == BucketIndexEngineType.SIMPLE =>
+            new BucketBulkInsertDataInternalWriterHelper(
+              table,
+              writeConfig,
+              instantTime,
+              taskPartitionId,
+              taskId,
+              taskEpochId,
+              schema,
+              writeConfig.populateMetaFields,
+              arePartitionRecordsSorted,
+              shouldPreserveHoodieMetadata)
           case _ =>
             new BulkInsertDataInternalWriterHelper(
               table,
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java b/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java
index bc82a04023c..033edc822c0 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/model/WriteOperationType.java
@@ -101,6 +101,8 @@ public enum WriteOperationType {
         return INDEX;
       case "alter_schema":
         return ALTER_SCHEMA;
+      case "bucket_rescale":
+        return BUCKET_RESCALE;
       case "unknown":
         return UNKNOWN;
       default:
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 17ff21c3938..39fabe0cbde 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
@@ -30,6 +30,7 @@ import org.apache.hudi.common.table.view.HoodieTableFileSystemView;
 import org.apache.hudi.common.util.CollectionUtils;
 import org.apache.hudi.config.HoodieWriteConfig;
 import org.apache.hudi.configuration.FlinkOptions;
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils;
 import org.apache.hudi.metadata.HoodieTableMetadata;
 import org.apache.hudi.storage.StoragePath;
 import org.apache.hudi.storage.StoragePathInfo;
@@ -1307,20 +1308,7 @@ public class ITTestHoodieDataSource {
         + "+I[id7, Bob, 44, 1970-01-01T00:00:07, par4], "
         + "+I[id8, Han, 56, 1970-01-01T00:00:08, par4]]");
     HoodieTableMetaClient metaClient = StreamerUtil.createMetaClient(basePath, new org.apache.hadoop.conf.Configuration());
-
-    List<StoragePathInfo> allFiles = metaClient.getStorage().listDirectEntries(new StoragePath(basePath)).stream().flatMap(path -> {
-      try {
-        return metaClient.getStorage().listDirectEntries(path.getPath()).stream();
-      } catch (IOException e) {
-        return Stream.empty();
-      }
-    }).collect(Collectors.toList());
-
-    HoodieTableFileSystemView fsView = new HoodieTableFileSystemView(metaClient,
-        metaClient.getActiveTimeline().getCommitsTimeline().filterCompletedInstants(), allFiles);
-    List<String> actual = fsView.getAllFileGroups().map(group -> {
-      return group.getPartitionPath() + group.getFileGroupId().getFileId().split("-")[0];
-    }).collect(Collectors.toList());
+    List<String> actual = PartitionBucketIndexUtils.getAllFileIDWithPartition(metaClient);
 
     // based on expression partition=(par1|par2),2 and default bucket number 1
     // par1 and par2 have two buckets.
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala
index 7c59e433c8a..b6a72a73e7e 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/hudi/HoodieSparkSqlWriter.scala
@@ -42,7 +42,7 @@ import org.apache.hudi.common.table.log.block.HoodieLogBlock.HoodieLogBlockType
 import org.apache.hudi.common.table.timeline.HoodieInstantTimeGenerator
 import org.apache.hudi.common.util.{CommitUtils, StringUtils, Option => HOption}
 import org.apache.hudi.common.util.ConfigUtils.getAllConfigKeys
-import org.apache.hudi.config.{HoodieCompactionConfig, HoodieInternalConfig, HoodieWriteConfig}
+import org.apache.hudi.config.{HoodieCompactionConfig, HoodieIndexConfig, HoodieInternalConfig, HoodieWriteConfig}
 import org.apache.hudi.config.HoodieBootstrapConfig.{BASE_PATH, INDEX_CLASS_NAME}
 import org.apache.hudi.config.HoodieWriteConfig.{SPARK_SQL_MERGE_INTO_PREPPED_KEY, WRITE_TABLE_VERSION}
 import org.apache.hudi.exception.{HoodieException, HoodieRecordCreationException, HoodieWriteConflictException}
@@ -68,6 +68,8 @@ import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileSystem, Path}
 import org.apache.hadoop.hive.conf.HiveConf
 import org.apache.hadoop.hive.shims.ShimLoader
+import org.apache.hudi.index.HoodieIndex
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils
 import org.apache.spark.{SPARK_VERSION, SparkContext}
 import org.apache.spark.api.java.{JavaRDD, JavaSparkContext}
 import org.apache.spark.sql._
@@ -247,7 +249,7 @@ class HoodieSparkSqlWriterInternal {
     asyncCompactionTriggerFnDefined = streamingWritesParamsOpt.map(_.asyncCompactionTriggerFn.isDefined).orElse(Some(false)).get
     asyncClusteringTriggerFnDefined = streamingWritesParamsOpt.map(_.asyncClusteringTriggerFn.isDefined).orElse(Some(false)).get
     // re-use table configs and inject defaults.
-    val (parameters, hoodieConfig) = mergeParamsAndGetHoodieConfig(optParams, tableConfig, mode, streamingWritesParamsOpt.isDefined)
+    var (parameters, hoodieConfig) = mergeParamsAndGetHoodieConfig(optParams, tableConfig, mode, streamingWritesParamsOpt.isDefined)
     val databaseName = hoodieConfig.getStringOrDefault(HoodieTableConfig.DATABASE_NAME, "")
     val tblName = hoodieConfig.getStringOrThrow(HoodieWriteConfig.TBL_NAME,
       s"'${HoodieWriteConfig.TBL_NAME.key}' must be set.").trim
@@ -331,6 +333,19 @@ class HoodieSparkSqlWriterInternal {
           .initTable(HadoopFSUtils.getStorageConfWithCopy(sparkContext.hadoopConfiguration), path)
       }
 
+      // take care of partition level bucket index which is simple bucket index
+      if (parameters.contains(HoodieIndexConfig.INDEX_TYPE.key)
+        && parameters(HoodieIndexConfig.INDEX_TYPE.key) == HoodieIndex.IndexType.BUCKET.name
+        && (!parameters.contains(HoodieIndexConfig.BUCKET_INDEX_ENGINE_TYPE.key)
+        || (parameters.contains(HoodieIndexConfig.BUCKET_INDEX_ENGINE_TYPE.key)
+        && parameters(HoodieIndexConfig.BUCKET_INDEX_ENGINE_TYPE.key) == HoodieIndex.IndexType.SIMPLE.name))
+        && !parameters.contains(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.key)
+        && PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(tableMetaClient.getStorageConf, basePath.toString)) {
+        val hashingConfigToLoad: String = PartitionBucketIndexUtils.getHashingConfigInstantToLoad(tableMetaClient)
+        hoodieConfig.setValue(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.key, hashingConfigToLoad)
+        parameters = parameters ++ Map(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.key -> hashingConfigToLoad)
+      }
+
       var instantTime: String = null
       tableConfig = tableMetaClient.getTableConfig
 
diff --git a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala
index 6cf5053bdce..3908e51a5db 100644
--- a/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala
+++ b/hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/ProvidesHoodieConfig.scala
@@ -30,6 +30,8 @@ import org.apache.hudi.config.{HoodieIndexConfig, HoodieInternalConfig, HoodieWr
 import org.apache.hudi.config.HoodieWriteConfig.TBL_NAME
 import org.apache.hudi.hive.{HiveSyncConfig, HiveSyncConfigHolder, MultiPartKeysValueExtractor}
 import org.apache.hudi.hive.ddl.HiveSyncMode
+import org.apache.hudi.index.HoodieIndex
+import org.apache.hudi.index.HoodieIndex.BucketIndexEngineType
 import org.apache.hudi.keygen.{ComplexKeyGenerator, CustomAvroKeyGenerator, CustomKeyGenerator}
 import org.apache.hudi.sql.InsertMode
 import org.apache.hudi.sync.common.HoodieSyncConfig
@@ -91,9 +93,9 @@ trait ProvidesHoodieConfig extends Logging {
     val hiveSyncConfig = buildHiveSyncConfig(sparkSession, hoodieCatalogTable, tableConfig)
 
     val defaultOpts = Map[String, String](
-      OPERATION.key -> BULK_INSERT_OPERATION_OPT_VAL,
-      HoodieInternalConfig.BULKINSERT_OVERWRITE_OPERATION_TYPE.key -> WriteOperationType.BUCKET_RESCALE.value(),
       KEYGENERATOR_CLASS_NAME.key -> classOf[SqlKeyGenerator].getCanonicalName,
+      HoodieIndexConfig.INDEX_TYPE.key() -> HoodieIndex.IndexType.BUCKET.name(),
+      HoodieIndexConfig.BUCKET_INDEX_ENGINE_TYPE.key() -> BucketIndexEngineType.SIMPLE.name(),
       SqlKeyGenerator.ORIGINAL_KEYGEN_CLASS_NAME -> tableConfig.getKeyGeneratorClassName,
       SqlKeyGenerator.PARTITION_SCHEMA -> hoodieCatalogTable.partitionSchema.toDDL,
       HoodieSyncConfig.META_SYNC_ENABLED.key -> hiveSyncConfig.getString(HoodieSyncConfig.META_SYNC_ENABLED.key),
diff --git a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala
index 0a6afd4e00a..7bea2d114a4 100644
--- a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala
+++ b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/HoodieProcedures.scala
@@ -96,6 +96,7 @@ object HoodieProcedures {
       ,(RunTTLProcedure.NAME, RunTTLProcedure.builder)
       ,(DropPartitionProcedure.NAME, DropPartitionProcedure.builder)
       ,(TruncateTableProcedure.NAME, TruncateTableProcedure.builder)
+      ,(PartitionBucketIndexManager.NAME, PartitionBucketIndexManager.builder)
     )
   }
 }
diff --git a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/PartitionBucketIndexManager.scala b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/PartitionBucketIndexManager.scala
index 4aaeec80a1a..d949bbcf9a0 100644
--- a/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/PartitionBucketIndexManager.scala
+++ b/hudi-spark-datasource/hudi-spark/src/main/scala/org/apache/spark/sql/hudi/command/procedures/PartitionBucketIndexManager.scala
@@ -18,44 +18,38 @@
 package org.apache.spark.sql.hudi.command.procedures
 
 import org.apache.avro.Schema
-import org.apache.hadoop.conf.Configuration
+import org.apache.hudi.DataSourceWriteOptions.{BULK_INSERT_OPERATION_OPT_VAL, ENABLE_ROW_WRITER, OPERATION}
 import org.apache.hudi.{AvroConversionUtils, HoodieCLIUtils, HoodieSparkSqlWriter}
 import org.apache.hudi.avro.HoodieAvroUtils
 import org.apache.hudi.client.SparkRDDWriteClient
-import org.apache.hudi.client.common.HoodieSparkEngineContext
-import org.apache.hudi.commit.DatasetBucketRescaleCommitActionExecutor
 import org.apache.hudi.common.config.{HoodieMetadataConfig, HoodieReaderConfig, SerializableSchema}
 import org.apache.hudi.common.engine.HoodieEngineContext
 import org.apache.hudi.common.fs.FSUtils
-import org.apache.hudi.common.model.{FileSlice, PartitionBucketIndexHashingConfig, WriteOperationType}
+import org.apache.hudi.common.model.{PartitionBucketIndexHashingConfig, WriteOperationType}
 import org.apache.hudi.common.table.read.HoodieFileGroupReader
 import org.apache.hudi.common.table.{HoodieTableMetaClient, TableSchemaResolver}
 import org.apache.hudi.common.table.view.HoodieTableFileSystemView
-import org.apache.hudi.common.util.{Option, StringUtils}
-import org.apache.hudi.config.{HoodieIndexConfig, HoodieWriteConfig}
+import org.apache.hudi.common.util.{Option, ValidationUtils}
+import org.apache.hudi.config.{HoodieIndexConfig, HoodieInternalConfig}
 import org.apache.hudi.exception.HoodieException
 import org.apache.hudi.index.bucket.{PartitionBucketIndexCalculator, PartitionBucketIndexUtils}
 import org.apache.hudi.internal.schema.InternalSchema
-import org.apache.hudi.internal.schema.utils.SerDeHelper
-import org.apache.hudi.storage.hadoop.HadoopStorageConfiguration
-import org.apache.hudi.storage.{StoragePath, StoragePathInfo}
+import org.apache.hudi.storage.StoragePath
 import org.apache.hudi.table.SparkBroadcastManager
 import org.apache.spark.internal.Logging
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.{HoodieUnsafeUtils, Row, SaveMode}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.catalog.HoodieCatalogTable
 import org.apache.spark.sql.catalyst.expressions.PredicateHelper
 import org.apache.spark.sql.hudi.ProvidesHoodieConfig
 import org.apache.spark.sql.types.{DataTypes, Metadata, StructField, StructType}
 
 import java.util
 import java.util.function.Supplier
-import java.util.stream.Collectors
 import scala.collection.JavaConverters._
 import scala.collection.mutable
 
-class PartitionBucketIndexManager() extends BaseProcedure
+class PartitionBucketIndexManager extends BaseProcedure
   with ProcedureBuilder
   with PredicateHelper
   with ProvidesHoodieConfig
@@ -64,12 +58,14 @@ class PartitionBucketIndexManager() extends BaseProcedure
   private val PARAMETERS = Array[ProcedureParameter](
     ProcedureParameter.required(0, "table", DataTypes.StringType),
     ProcedureParameter.optional(1, "overwrite", DataTypes.StringType),
-    ProcedureParameter.optional(2, "bucket-number", DataTypes.IntegerType, -1),
+    ProcedureParameter.optional(2, "bucketNumber", DataTypes.IntegerType, -1),
     ProcedureParameter.optional(3, "add", DataTypes.StringType),
     ProcedureParameter.optional(4, "dry-run", DataTypes.BooleanType, true),
     ProcedureParameter.optional(5, "rollback", DataTypes.StringType),
     ProcedureParameter.optional(6, "show-config", DataTypes.BooleanType, false),
-    ProcedureParameter.optional(7, "rule", DataTypes.StringType, "regex")
+    ProcedureParameter.optional(7, "rule", DataTypes.StringType, "regex"),
+    // params => key=value, key2=value2
+    ProcedureParameter.optional(8, "options", DataTypes.StringType)
   )
 
   private val OUTPUT_TYPE = new StructType(Array[StructField](
@@ -93,6 +89,7 @@ class PartitionBucketIndexManager() extends BaseProcedure
     val rollback = getArgValueOrDefault(args, PARAMETERS(5)).orNull.asInstanceOf[String]
     val showConfig = getArgValueOrDefault(args, PARAMETERS(6)).get.asInstanceOf[Boolean]
     val rule = getArgValueOrDefault(args, PARAMETERS(7)).orNull.asInstanceOf[String]
+    val options = getArgValueOrDefault(args, PARAMETERS(8))
 
     try {
       // Get table metadata
@@ -102,6 +99,18 @@ class PartitionBucketIndexManager() extends BaseProcedure
       val writeClient = HoodieCLIUtils.createHoodieWriteClient(sparkSession, basePath, Map.empty,
         scala.Option.apply(table))
       val context = writeClient.getEngineContext
+      var config = buildBucketRescaleHoodieConfig(hoodieCatalogTable)
+
+      options match {
+        case Some(p) =>
+          config = config ++ HoodieCLIUtils.extractOptions(p.asInstanceOf[String])
+        case _ =>
+          logInfo("No options")
+      }
+
+      config = config ++ Map(OPERATION.key -> BULK_INSERT_OPERATION_OPT_VAL,
+        HoodieInternalConfig.BULKINSERT_OVERWRITE_OPERATION_TYPE.key -> WriteOperationType.BUCKET_RESCALE.value(),
+        ENABLE_ROW_WRITER.key() -> "true")
 
       // Determine which operation to perform
       if (showConfig) {
@@ -109,21 +118,21 @@ class PartitionBucketIndexManager() extends BaseProcedure
       } else if (rollback != null) {
         handleRollback(metaClient, rollback)
       } else if (overwrite != null) {
-        handleOverwrite(hoodieCatalogTable, writeClient, context, metaClient, overwrite, bucketNumber, rule, dryRun)
+        handleOverwrite(config, writeClient, context, metaClient, overwrite, bucketNumber, rule, dryRun)
       } else if (add != null) {
         handleAdd(metaClient, add, dryRun)
       } else {
         Seq(Row("ERROR", "INVALID_OPERATION", "No valid operation specified"))
       }
     } catch {
-      case e: Exception => Seq(Row("ERROR", "EXCEPTION", e.getMessage))
+      case e: Exception => throw new HoodieException(e.getMessage, e)
     }
   }
 
   /**
    * Handle the overwrite operation.
    */
-  private def handleOverwrite(catalog: HoodieCatalogTable,
+  private def handleOverwrite(config: Map[String, String],
                               writeClient: SparkRDDWriteClient[_],
                               context : HoodieEngineContext,
                               metaClient: HoodieTableMetaClient,
@@ -131,30 +140,42 @@ class PartitionBucketIndexManager() extends BaseProcedure
                               bucketNumber: Int,
                               rule: String,
                               dryRun: Boolean): Seq[Row] = {
-    val config = buildBucketRescaleHoodieConfig(catalog).asJava
     val basePath = metaClient.getBasePath
     val mdtEnable = metaClient.getStorage().exists(new StoragePath(metaClient.getBasePath, HoodieTableMetaClient.METADATA_TABLE_FOLDER_PATH))
 
     // get all partition paths
     val allPartitions = FSUtils.getAllPartitionPaths(context, metaClient.getStorage, metaClient.getBasePath, mdtEnable)
+    val usePartitionBucketIndexBefore = PartitionBucketIndexUtils.isPartitionSimpleBucketIndex(context.getStorageConf, basePath.toString)
 
-    // get Map<partitionPath, bucketNumber> based on latest hashing_config
-    val latestHashingConfigInstant = PartitionBucketIndexUtils.getHashingConfigInstantToLoad(metaClient)
-    val calcWithLatestInstant = PartitionBucketIndexCalculator.getInstance(latestHashingConfigInstant, metaClient)
-    val partition2BucketWithLatestHashingConfig = PartitionBucketIndexUtils.getAllBucketNumbers(calcWithLatestInstant, allPartitions)
+    var partition2BucketWithLatestHashingConfig: util.Map[String, Integer] = null
+    var calcWithLatestInstant: PartitionBucketIndexCalculator = null
 
+    if (usePartitionBucketIndexBefore) {
+      // get Map<partitionPath, bucketNumber> based on latest hashing_config
+      val latestHashingConfigInstant = PartitionBucketIndexUtils.getHashingConfigInstantToLoad(metaClient)
+      calcWithLatestInstant = PartitionBucketIndexCalculator.getInstance(latestHashingConfigInstant, metaClient)
+      partition2BucketWithLatestHashingConfig = PartitionBucketIndexUtils.getAllBucketNumbers(calcWithLatestInstant, allPartitions)
+    } else {
+      ValidationUtils.checkArgument(bucketNumber != -1)
+      partition2BucketWithLatestHashingConfig = allPartitions.asScala.map(partition => (partition, new Integer(bucketNumber))).toMap.asJava
+    }
     // get Map<partitionPath, bucketNumber> based on new given expression
     val defaultBucketNumber = if (bucketNumber != -1) {
       bucketNumber
-    } else {
+    } else if (calcWithLatestInstant != null ) {
       // reuse latest default bucket number
       calcWithLatestInstant.getHashingConfig.getDefaultBucketNumber
+    } else {
+      throw new HoodieException("Please set original bucket number before upgrade to partition bucket level bucket index")
     }
+
     val instantTime = writeClient.createNewInstantTime()
-    config.put(HoodieIndexConfig.BUCKET_INDEX_PARTITION_EXPRESSIONS.key(), expression)
-    config.put(HoodieIndexConfig.BUCKET_INDEX_PARTITION_RULE_TYPE.key(), rule)
-    config.put(HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.key(), instantTime)
-    config.put(HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS.key(), bucketNumber.toString)
+
+    val finalConfig = config ++ Map(HoodieIndexConfig.BUCKET_INDEX_PARTITION_EXPRESSIONS.key()->  expression,
+      HoodieIndexConfig.BUCKET_INDEX_PARTITION_RULE_TYPE.key() -> rule,
+      HoodieIndexConfig.BUCKET_INDEX_PARTITION_LOAD_INSTANT.key() -> instantTime,
+      HoodieIndexConfig.BUCKET_INDEX_NUM_BUCKETS.key() -> bucketNumber.toString
+    )
 
     val newConfig = new PartitionBucketIndexHashingConfig(expression, defaultBucketNumber, rule, PartitionBucketIndexHashingConfig.CURRENT_VERSION, instantTime)
     val calcWithNewExpression = PartitionBucketIndexCalculator.getInstance(instantTime, newConfig)
@@ -166,7 +187,9 @@ class PartitionBucketIndexManager() extends BaseProcedure
 
     // get all fileSlices need to read
     val allFilesMap = FSUtils.getFilesInPartitions(context, metaClient.getStorage(), HoodieMetadataConfig.newBuilder.enable(mdtEnable).build,
-      metaClient.getBasePath.toString, partitionsToRescale.toArray)
+      metaClient.getBasePath.toString, partitionsToRescale.map(relative => {
+        new StoragePath(basePath, relative)
+      }).map(storagePath => storagePath.toString).toArray)
     val files = allFilesMap.values().asScala.flatMap(x => x.asScala).toList
     val view = new HoodieTableFileSystemView(metaClient, metaClient.getActiveTimeline, files.asJava)
     val allFileSlice = partitionsToRescale.flatMap(partitionPath => {
@@ -185,42 +208,46 @@ class PartitionBucketIndexManager() extends BaseProcedure
     val broadcastManager = new SparkBroadcastManager(context, metaClient)
     broadcastManager.prepareAndBroadcast()
     val sparkSchemaWithMetaFields = AvroConversionUtils.convertAvroSchemaToStructType(tableSchemaWithMetaFields)
-    val serializableTableSchemaWithMetaFields = new SerializableSchema(tableSchemaWithMetaFields)
-    val latestInstantTime = metaClient.getActiveTimeline.getCommitsTimeline.filterCompletedInstants().lastInstant().get()
-
-    val res: RDD[InternalRow] = spark.sparkContext.parallelize(allFileSlice, allFileSlice.size).flatMap(fileSlice => {
-      // instantiate other supporting cast
-      val readerSchema = serializableTableSchemaWithMetaFields.get
-      val readerContextOpt = broadcastManager.retrieveFileGroupReaderContext(basePath)
-      val internalSchemaOption: Option[InternalSchema] = Option.empty()
-      // instantiate FG reader
-      val fileGroupReader = new HoodieFileGroupReader(readerContextOpt.get(),
-        metaClient.getStorage,
-        basePath.toString,
-        latestInstantTime.requestedTime(),
-        fileSlice,
-        readerSchema,
-        readerSchema,
-        internalSchemaOption, // not support evolution of schema for now
-        metaClient,
-        metaClient.getTableConfig.getProps,
-        0,
-        java.lang.Long.MAX_VALUE,
-        HoodieReaderConfig.MERGE_USE_RECORD_POSITIONS.defaultValue())
-      fileGroupReader.initRecordIterators()
-      val iterator = fileGroupReader.getClosableIterator.asInstanceOf[HoodieFileGroupReader.HoodieFileGroupReaderIterator[InternalRow]]
-      iterator.asScala
-    })
-    val dataFrame = HoodieUnsafeUtils.createDataFrameFromRDD(sparkSession, res, sparkSchemaWithMetaFields)
 
+    val res: RDD[InternalRow] = if (allFileSlice.isEmpty) {
+      spark.sparkContext.emptyRDD
+    } else {
+      val serializableTableSchemaWithMetaFields = new SerializableSchema(tableSchemaWithMetaFields)
+      val latestInstantTime = metaClient.getActiveTimeline.getCommitsTimeline.filterCompletedInstants().lastInstant().get()
+
+      spark.sparkContext.parallelize(allFileSlice, allFileSlice.size).flatMap(fileSlice => {
+        // instantiate other supporting cast
+        val readerSchema = serializableTableSchemaWithMetaFields.get
+        val readerContextOpt = broadcastManager.retrieveFileGroupReaderContext(basePath)
+        val internalSchemaOption: Option[InternalSchema] = Option.empty()
+        // instantiate FG reader
+        val fileGroupReader = new HoodieFileGroupReader(readerContextOpt.get(),
+          metaClient.getStorage,
+          basePath.toString,
+          latestInstantTime.requestedTime(),
+          fileSlice,
+          readerSchema,
+          readerSchema,
+          internalSchemaOption, // not support evolution of schema for now
+          metaClient,
+          metaClient.getTableConfig.getProps,
+          0,
+          java.lang.Long.MAX_VALUE,
+          HoodieReaderConfig.MERGE_USE_RECORD_POSITIONS.defaultValue())
+        fileGroupReader.initRecordIterators()
+        val iterator = fileGroupReader.getClosableIterator.asInstanceOf[HoodieFileGroupReader.HoodieFileGroupReaderIterator[InternalRow]]
+        iterator.asScala
+      })
+    }
+    val dataFrame = HoodieUnsafeUtils.createDataFrameFromRDD(sparkSession, res, sparkSchemaWithMetaFields)
     val (success, _, _, _, _, _) = HoodieSparkSqlWriter.write(
       sparkSession.sqlContext,
-      SaveMode.Overwrite,
-      config.asScala.toMap,
+      SaveMode.Append,
+      finalConfig,
       dataFrame)
 
-    val details = s"Expression: $expression, Bucket Number: $bucketNumber, Dry Run: $dryRun"
-    Seq(Row("SUCCESS", "OVERWRITE", success))
+    val details = s"Expression: $expression, Bucket Number: $bucketNumber, Dry Run: $dryRun, Instant: $instantTime"
+    Seq(Row("SUCCESS", "OVERWRITE", details))
   }
 
   /**
@@ -277,13 +304,13 @@ class PartitionBucketIndexManager() extends BaseProcedure
     Seq(Row("SUCCESS", "SHOW_CONFIG", null))
   }
 
-  override def build: Procedure = new PartitionBucketIndexManager()
+  override def build: Procedure = new PartitionBucketIndexManager
 }
 
 object PartitionBucketIndexManager {
-  val NAME = "PartitionBucketIndexManager"
+  val NAME = "partition_bucket_index_manager"
 
-  def builder: Supplier[ProcedureBuilder] = new Supplier[ProcedureBuilder] {
-    override def get() = new PartitionBucketIndexManager()
+  def builder : Supplier[ProcedureBuilder] = new Supplier[ProcedureBuilder] {
+    override def get() = new PartitionBucketIndexManager
   }
 }
\ No newline at end of file
diff --git a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala
index 532a000f0f6..7c1bc0aa2a2 100644
--- a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTable.scala
@@ -1847,94 +1847,6 @@ class TestInsertTable extends HoodieSparkSqlTestBase {
     }
   }
 
-  test("Test Bulk Insert Into Partition Bucket Index Table") {
-    withSQLConf(
-      "hoodie.datasource.write.operation" -> "bulk_insert",
-      "hoodie.bulkinsert.shuffle.parallelism" -> "1") {
-      Seq("mor", "cow").foreach { tableType =>
-        Seq("true", "false").foreach { bulkInsertAsRow =>
-          withTempDir { tmp =>
-            val tableName = generateTableName
-            // 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,name',
-                 | type = '$tableType',
-                 | preCombineField = 'ts',
-                 | hoodie.index.type = 'BUCKET',
-                 | hoodie.bucket.index.hash.field = 'id,name',
-                 | hoodie.datasource.write.row.writer.enable = '$bulkInsertAsRow')
-                 | partitioned by (dt)
-                 | location '${tmp.getCanonicalPath}'
-                 """.stripMargin)
-
-            // Note: Do not write the field alias, the partition field must be placed last.
-            spark.sql(
-              s"""
-                 | insert into $tableName values
-                 | (1, 'a1,1', 10, 1000, "2021-01-05"),
-                 | (2, 'a2', 20, 2000, "2021-01-06"),
-                 | (3, 'a3,3', 30, 3000, "2021-01-07")
-                 """.stripMargin)
-
-            checkAnswer(s"select id, name, price, ts, dt from $tableName")(
-              Seq(1, "a1,1", 10.0, 1000, "2021-01-05"),
-              Seq(2, "a2", 20.0, 2000, "2021-01-06"),
-              Seq(3, "a3,3", 30.0, 3000, "2021-01-07")
-            )
-
-            // for COW with disabled Spark native row writer, multiple bulk inserts are restricted
-            if (tableType != "cow" && bulkInsertAsRow != "false") {
-              spark.sql(
-                s"""
-                   | insert into $tableName values
-                   | (1, 'a1', 10, 1000, "2021-01-05"),
-                   | (3, "a3", 30, 3000, "2021-01-07")
-                 """.stripMargin)
-
-              checkAnswer(s"select id, name, price, ts, dt from $tableName")(
-                Seq(1, "a1,1", 10.0, 1000, "2021-01-05"),
-                Seq(1, "a1", 10.0, 1000, "2021-01-05"),
-                Seq(2, "a2", 20.0, 2000, "2021-01-06"),
-                Seq(3, "a3,3", 30.0, 3000, "2021-01-07"),
-                Seq(3, "a3", 30.0, 3000, "2021-01-07")
-              )
-
-              // there are two files in partition(dt = '2021-01-05')
-              checkAnswer(s"select count(distinct _hoodie_file_name) from $tableName where dt = '2021-01-05'")(
-                Seq(2)
-              )
-
-              // would generate 6 other files in partition(dt = '2021-01-05')
-              spark.sql(
-                s"""
-                   | insert into $tableName values
-                   | (4, 'a1,1', 10, 1000, "2021-01-05"),
-                   | (5, 'a1,1', 10, 1000, "2021-01-05"),
-                   | (6, 'a1,1', 10, 1000, "2021-01-05"),
-                   | (7, 'a1,1', 10, 1000, "2021-01-05"),
-                   | (8, 'a1,1', 10, 1000, "2021-01-05"),
-                   | (10, 'a3,3', 30, 3000, "2021-01-05")
-                 """.stripMargin)
-
-              checkAnswer(s"select count(distinct _hoodie_file_name) from $tableName where dt = '2021-01-05'")(
-                Seq(8)
-              )
-            }
-          }
-        }
-      }
-    }
-  }
-
   test("Test not supported multiple BULK INSERTs into COW with SIMPLE BUCKET and disabled Spark native row writer") {
     withSQLConf("hoodie.datasource.write.operation" -> "bulk_insert",
       "hoodie.bulkinsert.shuffle.parallelism" -> "1") {
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
new file mode 100644
index 00000000000..d437f4c9679
--- /dev/null
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/dml/TestInsertTableWithPartitionBucketIndex.scala
@@ -0,0 +1,223 @@
+/*
+ * 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.spark.sql.hudi.dml
+
+import org.apache.hudi.index.bucket.PartitionBucketIndexUtils
+import org.apache.hudi.testutils.HoodieClientTestUtils.createMetaClient
+import org.apache.spark.sql.hudi.common.HoodieSparkSqlTestBase
+import scala.collection.JavaConverters._
+
+class TestInsertTableWithPartitionBucketIndex extends HoodieSparkSqlTestBase {
+
+  /**
+   * Write operation including :
+   * Case1: Mor + bulk_insert + asRow enable + simple bucket index
+   * Case2: Mor + bulk_insert + asRow disable + simple bucket index
+   * Case3: Cow + bulk_insert + asRow enable + simple bucket index
+   * Case4: Cow + bulk_insert + asRow disable + simple bucket index
+   * Case5: Mor + bulk_insert + asRow enable + partition level bucket index
+   * Case6: Mor + bulk_insert + asRow disable + partition level bucket index
+   * Case7: Cow + bulk_insert + asRow enable + partition level bucket index
+   * Case8: Cow + bulk_insert + asRow disable + partition level bucket index
+   *
+   * Total WorkFlow are as followed:
+   * Step1: Create a new table using simple bucket index
+   * Step2: Upgrade table to partition level bucket index using `call partition_bucket_index_manager`
+   * Step3: Trigger another write operation through partition level bucket index
+   *
+   */
+  test("Test Bulk Insert Into Partition Bucket Index Table") {
+    withSQLConf(
+      "hoodie.datasource.write.operation" -> "bulk_insert",
+      "hoodie.bulkinsert.shuffle.parallelism" -> "2") {
+      Seq("mor", "cow").foreach { tableType =>
+        Seq("true", "false").foreach { bulkInsertAsRow =>
+          withTempDir { tmp =>
+            val tableName = generateTableName
+            // 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,name',
+                 | type = '$tableType',
+                 | preCombineField = 'ts',
+                 | hoodie.index.type = 'BUCKET',
+                 | hoodie.bucket.index.hash.field = 'id,name',
+                 | hoodie.bucket.index.num.buckets = 1,
+                 | hoodie.datasource.write.row.writer.enable = '$bulkInsertAsRow')
+                 | partitioned by (dt)
+                 | location '${tmp.getCanonicalPath}'
+           """.stripMargin)
+
+            // Note: Do not write the field alias, the partition field must be placed last.
+            spark.sql(
+              s"""
+                 | insert into $tableName values
+                 | (1, 'a1,1', 10, 1000, "2021-01-05"),
+                 | (11, 'a1,1', 10, 1000, "2021-01-05"),
+                 | (2, 'a2', 20, 2000, "2021-01-06"),
+                 | (22, 'a2', 20, 2000, "2021-01-06")
+           """.stripMargin)
+
+            // upgrade to partition level bucket index and rescale dt=2021-01-05 from 1 to 2
+            val expressions = "dt=(2021\\-01\\-05|2021\\-01\\-07),2"
+            val rule = "regex"
+            val defaultBucketNumber = 1
+            spark.sql(s"call partition_bucket_index_manager(table => '$tableName', overwrite => '$expressions', rule => '$rule', bucketNumber => $defaultBucketNumber)")
+
+            checkAnswer(s"select id, name, price, ts, dt from $tableName")(
+              Seq(1, "a1,1", 10.0, 1000, "2021-01-05"),
+              Seq(11, "a1,1", 10.0, 1000, "2021-01-05"),
+              Seq(2, "a2", 20.0, 2000, "2021-01-06"),
+              Seq(22, "a2", 20.0, 2000, "2021-01-06")
+            )
+
+            // insert into new partition 2021-01-07 with partition level bucket index
+            spark.sql(
+              s"""
+                 | insert into $tableName values
+                 | (3, 'a3,3', 30, 3000, "2021-01-07"),
+                 | (33, 'a3,3', 30, 3000, "2021-01-07")
+           """.stripMargin)
+
+            checkAnswer(s"select id, name, price, ts, dt from $tableName") (
+              Seq(1, "a1,1", 10.0, 1000, "2021-01-05"),
+              Seq(11, "a1,1", 10.0, 1000, "2021-01-05"),
+              Seq(2, "a2", 20.0, 2000, "2021-01-06"),
+              Seq(22, "a2", 20.0, 2000, "2021-01-06"),
+              Seq(3, "a3,3", 30.0, 3000, "2021-01-07"),
+              Seq(33, "a3,3", 30.0, 3000, "2021-01-07")
+            )
+            val metaClient = createMetaClient(spark, tmp.getCanonicalPath)
+            val actual: List[String] = PartitionBucketIndexUtils.getAllFileIDWithPartition(metaClient).asScala.toList
+            val expected: List[String] = List("dt=2021-01-05" + "00000000",
+              "dt=2021-01-05" + "00000000",
+              "dt=2021-01-05" + "00000001",
+              "dt=2021-01-06" + "00000000",
+              "dt=2021-01-07" + "00000000",
+              "dt=2021-01-07" + "00000001")
+            assert(actual.sorted == expected.sorted)
+          }
+        }
+      }
+    }
+  }
+
+  /**
+   * Write operation including :
+   * Case1: Mor + upsert + simple bucket index
+   * Case2: Cow + upsert + simple bucket index
+   * Case3: Mor + upsert + partition level bucket index
+   * Case4: Cow + upsert + partition level bucket index
+   *
+   * Total WorkFlow are as followed:
+   * Step1: Create a new table using simple bucket index
+   * Step2: Upgrade table to partition level bucket index using `call partition_bucket_index_manager`
+   * Step3: Trigger another write operation through partition level bucket index
+   *
+   */
+  test("Test Upsert Into Partition Bucket Index Table") {
+    withSQLConf(
+      "hoodie.datasource.write.operation" -> "upsert") {
+      Seq("cow", "mor").foreach { tableType =>
+        withTempDir { tmp =>
+          val tableName = generateTableName
+          // 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,name',
+               | type = '$tableType',
+               | preCombineField = 'ts',
+               | hoodie.index.type = 'BUCKET',
+               | hoodie.bucket.index.hash.field = 'id,name',
+               | hoodie.bucket.index.num.buckets = 1)
+               | partitioned by (dt)
+               | location '${tmp.getCanonicalPath}'
+               | """.stripMargin)
+
+          // Note: Do not write the field alias, the partition field must be placed last.
+          spark.sql(
+            s"""
+               | insert into $tableName values
+               | (1, 'a1,1', 10, 1000, "2021-01-05"),
+               | (11, 'a1,1', 10, 1000, "2021-01-05"),
+               | (2, 'a2', 20, 2000, "2021-01-06"),
+               | (22, 'a2', 20, 2000, "2021-01-06")
+               | """.stripMargin)
+
+          // upgrade to partition level bucket index and rescale dt=2021-01-05 from 1 to 2
+          val expressions = "dt=(2021\\-01\\-05|2021\\-01\\-07),2"
+          val rule = "regex"
+          val defaultBucketNumber = 1
+          spark.sql(s"call partition_bucket_index_manager(table => '$tableName', overwrite => '$expressions', rule => '$rule', bucketNumber => $defaultBucketNumber)")
+
+          checkAnswer(s"select id, name, price, ts, dt from $tableName")(
+            Seq(1, "a1,1", 10.0, 1000, "2021-01-05"),
+            Seq(11, "a1,1", 10.0, 1000, "2021-01-05"),
+            Seq(2, "a2", 20.0, 2000, "2021-01-06"),
+            Seq(22, "a2", 20.0, 2000, "2021-01-06")
+          )
+
+          // insert into new partition 2021-01-07 with partition level bucket index
+          // update (1, 'a1,1', 10, 1000, "2021-01-05") to (1, 'a1,1', 100, 1000, "2021-01-05")
+          spark.sql(
+            s"""
+               | insert into $tableName values
+               | (1, 'a1,1', 100, 1000, "2021-01-05"),
+               | (3, 'a3,3', 30, 3000, "2021-01-07"),
+               | (33, 'a3,3', 30, 3000, "2021-01-07")
+               | """.stripMargin)
+
+          checkAnswer(s"select id, name, price, ts, dt from $tableName")(
+            Seq(1, "a1,1", 100.0, 1000, "2021-01-05"),
+            Seq(11, "a1,1", 10.0, 1000, "2021-01-05"),
+            Seq(2, "a2", 20.0, 2000, "2021-01-06"),
+            Seq(22, "a2", 20.0, 2000, "2021-01-06"),
+            Seq(3, "a3,3", 30.0, 3000, "2021-01-07"),
+            Seq(33, "a3,3", 30.0, 3000, "2021-01-07")
+          )
+          val metaClient = createMetaClient(spark, tmp.getCanonicalPath)
+          val actual: List[String] = PartitionBucketIndexUtils.getAllFileIDWithPartition(metaClient).asScala.toList
+          val expected: List[String] = List("dt=2021-01-05" + "00000000",
+            "dt=2021-01-05" + "00000000",
+            "dt=2021-01-05" + "00000001",
+            "dt=2021-01-06" + "00000000",
+            "dt=2021-01-07" + "00000000",
+            "dt=2021-01-07" + "00000001")
+          // compare file group as expected
+          assert(actual.sorted == expected.sorted)
+        }
+      }
+    }
+  }
+}
diff --git a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/procedure/TestPartitionBucketIndexManager.scala b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/procedure/TestPartitionBucketIndexManager.scala
index 93d2a557afe..0bf5aaaa3a4 100644
--- a/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/procedure/TestPartitionBucketIndexManager.scala
+++ b/hudi-spark-datasource/hudi-spark/src/test/scala/org/apache/spark/sql/hudi/procedure/TestPartitionBucketIndexManager.scala
@@ -16,7 +16,7 @@ class TestPartitionBucketIndexManager extends HoodieSparkProcedureTestBase {
       val defaultBucketNumber = 10
       PartitionBucketIndexUtils.initHashingConfig(metaClient, expressions, rule, defaultBucketNumber, null)
 
-      spark.sql(s"""call PartitionBucketIndexManager(table => '$tableName', show-config => 'true')""").collect()
+      spark.sql(s"call PartitionBucketIndexManager(table => '$tableName', show-config => 'true')").collect()
 
     }
   }


(hudi) 02/06: finish hashing_config initial && related UT

Posted by zh...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

zhangyue19921010 pushed a commit to branch HUDI-8990
in repository https://gitbox.apache.org/repos/asf/hudi.git

commit 60de489d9a925d4f681894fc310081ffc3f27555
Author: YueZhang <zh...@163.com>
AuthorDate: Mon Mar 17 20:04:40 2025 +0800

    finish hashing_config initial && related UT
---
 .../java/org/apache/hudi/config/HoodieIndexConfig.java     |  6 +++---
 .../apache/hudi/index/bucket/SimpleBucketIndexUtils.java   |  2 +-
 .../common/model/PartitionBucketIndexHashingConfig.java    |  2 --
 .../org/apache/hudi/table/catalog/TestHoodieCatalog.java   | 14 +++++++-------
 4 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
index 626e894fd1a..bfad563a9e7 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/config/HoodieIndexConfig.java
@@ -297,9 +297,9 @@ public class HoodieIndexConfig extends HoodieConfig {
       .key("hoodie.bucket.index.partition.expressions")
       .noDefaultValue()
       .markAdvanced()
-      .withDocumentation("Users can use this parameter to specify expression and the corresponding bucket " +
-          "numbers (separated by commas).Multiple rules are separated by semicolons like " +
-          "hoodie.bucket.index.partition.expressions=expression1,bucket-number1;expression2,bucket-number2");
+      .withDocumentation("Users can use this parameter to specify expression and the corresponding bucket "
+          + "numbers (separated by commas).Multiple rules are separated by semicolons like "
+          + "hoodie.bucket.index.partition.expressions=expression1,bucket-number1;expression2,bucket-number2");
 
   public static final ConfigProperty<String> BUCKET_INDEX_MAX_NUM_BUCKETS = ConfigProperty
       .key("hoodie.bucket.index.max.num.buckets")
diff --git a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java
index 74c8022c9aa..f3a20c6a002 100644
--- a/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java
+++ b/hudi-client/hudi-client-common/src/main/java/org/apache/hudi/index/bucket/SimpleBucketIndexUtils.java
@@ -46,7 +46,7 @@ public class SimpleBucketIndexUtils {
                                               String expressions,
                                               String rule,
                                               int defaultBucketNumber,
-                                              String instant){
+                                              String instant) {
     if (StringUtils.isNullOrEmpty(expressions)) {
       return false;
     }
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 225629b52ff..46b38896925 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
@@ -38,7 +38,6 @@ public class PartitionBucketIndexHashingConfig implements Serializable {
   private final String expressions;
   private final int defaultBucketNumber;
   private final String rule;
-
   private final int version;
   private final String instant;
 
@@ -71,7 +70,6 @@ public class PartitionBucketIndexHashingConfig implements Serializable {
     return JsonUtils.getObjectMapper().readValue(jsonStr, clazz);
   }
 
-
   public static PartitionBucketIndexHashingConfig fromBytes(byte[] bytes) throws IOException {
     try {
       return fromJsonString(new String(bytes, StandardCharsets.UTF_8), PartitionBucketIndexHashingConfig.class);
diff --git a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
index 1e518c5ef45..84e982df8f3 100644
--- a/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
+++ b/hudi-flink-datasource/hudi-flink/src/test/java/org/apache/hudi/table/catalog/TestHoodieCatalog.java
@@ -360,14 +360,14 @@ public class TestHoodieCatalog {
         .createMetaClient(
             new HadoopStorageConfiguration(HadoopConfigurations.getHadoopConf(flinkConf)), tablePathStr);
     HoodieStorage storage = metaClient.getStorage();
-    StoragePath initialHashing_config =
+    StoragePath initialHashingConfig =
         new StoragePath(metaClient.getHashingMetadataConfigPath(), SimpleBucketIndexUtils.INITIAL_HASHING_CONFIG_INSTANT + PartitionBucketIndexHashingConfig.HASHING_CONFIG_FILE_SUFFIX);
-    StoragePathInfo info = storage.getPathInfo(initialHashing_config);
-    Option<PartitionBucketIndexHashingConfig> hashing_config = SimpleBucketIndexUtils.loadHashingConfig(storage, info);
-    assertTrue(hashing_config.isPresent());
-    assertEquals(hashing_config.get().getDefaultBucketNumber(), Integer.parseInt(defaultBucketNumber));
-    assertEquals(hashing_config.get().getRule(), rule);
-    assertEquals(hashing_config.get().getExpressions(), expressions);
+    StoragePathInfo info = storage.getPathInfo(initialHashingConfig);
+    Option<PartitionBucketIndexHashingConfig> hashingConfig = SimpleBucketIndexUtils.loadHashingConfig(storage, info);
+    assertTrue(hashingConfig.isPresent());
+    assertEquals(hashingConfig.get().getDefaultBucketNumber(), Integer.parseInt(defaultBucketNumber));
+    assertEquals(hashingConfig.get().getRule(), rule);
+    assertEquals(hashingConfig.get().getExpressions(), expressions);
   }
 
   @Test