With blaze, kyuubi will use only one task for small file merging

[ziyangRen]

Describe the bug
With blaze, kyuubi will use only one task for small file merging：

DAG for this stage is as follow：

To Reproduce
I use spark3.3.1 with blaze v4.0.0 and kyuubi-extension-spark-3-3_2.12-1.8.1.jar
and related configrations are as follow:
--conf spark.sql.extensions=org.apache.spark.sql.blaze.BlazeSparkSessionExtension,org.apache.kyuubi.sql.KyuubiSparkSQLExtension
--conf spark.sql.optimizer.insertRepartitionBeforeWriteIfNoShuffle.enabled=true
--conf spark.blaze.enable=true
--conf spark.shuffle.service.enabled=false
--conf spark.shuffle.manager=org.apache.spark.sql.execution.blaze.shuffle.celeborn.BlazeCelebornShuffleManager

[zuston]

#663 Maybe you can cherry-pick this patch to check weather is fixed

[ziyangRen]

@zuston I will add a minimal reproduction example here, using the store_sales table from the tpcds dataset. Here's my step:

1. First, create a target partition table：

CREATE TABLE test.store_sales_with_batch (
  ss_sold_date_sk         INTEGER,
  ss_sold_time_sk         INTEGER,
  ss_item_sk              INTEGER,
  ss_customer_sk          INTEGER,
  ss_cdemo_sk             INTEGER,
  ss_hdemo_sk             INTEGER,
  ss_addr_sk              INTEGER,
  ss_store_sk             INTEGER,
  ss_promo_sk             INTEGER,
  ss_ticket_number        BIGINT,
  ss_quantity             INTEGER,
  ss_wholesale_cost       DECIMAL(7,2),
  ss_list_price           DECIMAL(7,2),
  ss_sales_price          DECIMAL(7,2),
  ss_ext_discount_amt     DECIMAL(7,2),
  ss_ext_sales_price      DECIMAL(7,2),
  ss_ext_wholesale_cost   DECIMAL(7,2),
  ss_ext_list_price       DECIMAL(7,2),
  ss_ext_tax              DECIMAL(7,2),
  ss_coupon_amt           DECIMAL(7,2),
  ss_net_paid             DECIMAL(7,2),
  ss_net_paid_inc_tax     DECIMAL(7,2),
  ss_net_profit           DECIMAL(7,2)
)
PARTITIONED BY (
   batch       INTEGER
);

2. inserted into the target partition table with dynamic partitioning syntax:
   note that static partition syntax is unable to reproduce the problem

    INSERT OVERWRITE TABLE test.store_sales_with_batch
        PARTITION (batch)
        SELECT
          ss_sold_date_sk,
          ss_sold_time_sk,
          ss_item_sk,
          ss_customer_sk,
          ss_cdemo_sk,
          ss_hdemo_sk,
          ss_addr_sk,
          ss_store_sk,
          ss_promo_sk,
          ss_ticket_number,
          ss_quantity,
          ss_wholesale_cost,
          ss_list_price,
          ss_sales_price,
          ss_ext_discount_amt,
          ss_ext_sales_price,
          ss_ext_wholesale_cost,
          ss_ext_list_price,
          ss_ext_tax,
          ss_coupon_amt,
          ss_net_paid,
          ss_net_paid_inc_tax,
          ss_net_profit,
          20241128 AS batch 
        FROM tpc_test.store_sales;

4. running job with kyuubi:
   --jars kyuubi-extension-spark-3-3_2.12-1.10.0.jar --conf spark.sql.extensions=org.apache.kyuubi.sql.KyuubiSparkSQLExtension --conf spark.sql.optimizer.insertRepartitionBeforeWriteIfNoShuffle.enabled=true
   and you can get jar by:
   wget https://repo1.maven.org/maven2/org/apache/kyuubi/kyuubi-extension-spark-3-3_2.12/1.10.0/kyuubi-extension-spark-3-3_2.12-1.10.0.jar

[ziyangRen]

#663 Maybe you can cherry-pick this patch to check weather is fixed
fixed!! thank you