feat: add param rows_per_range for range-based btree index built (#439)

# Background

Currently the partition(range) number is configured by spark parameter
like: `spark.sql.adaptive.coalescePartitions.initialPartitionNum` or
`spark.sql.shuffle.partitions` when building btree index using
range-mode.

The current approach cannot dynamically adjust the number of ranges
based on changes in the total row count of the Dataset. This becomes
quite inconvenient when the total row count of the Dataset continues to
grow.

# Design

So, we add a new parameter `rows_per_range` for range-mode. This param
specifies the row number for each range. The spark partition number is
calculated by `Dataset.total_rows/rows_per_range`. This method can
dynamically adjust the number of Ranges based on the Dataset's row
count.

@hamersaw @puchengy Could you please take a look and see if this makes
sense? Thank you.

Co-authored-by: fangbo.0511 <fangbo.0511@bytedance.com>

Author: fangbo

docs/src/operations/ddl/create-index.md

@@ -35,10 +35,12 @@ The `CREATE INDEX` command supports options via the `WITH` clause to control ind
 
 For the `btree` method, the following options are supported:
 
-| Option      | Type   | Description                                  |
-|-------------|--------|----------------------------------------------|
-| `zone_size` | Long   | The number of rows per zone in the B-tree index. |
-| `build_mode`| String | Index building mode: 'fragment' builds indexes in parallel by fragment; 'range' sorts data by indexed columns first, then partitions and builds indexes in parallel by partition. Default is 'fragment'.|
+| Option           | Type   | Description                                                                                                                                                                                              |
+|------------------|--------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `zone_size`      | Long   | The number of rows per zone in the B-tree index.                                                                                                                                                         |
+| `build_mode`     | String | Index building mode: 'fragment' builds indexes in parallel by fragment; 'range' sorts data by indexed columns first, then partitions and builds indexes in parallel by partition. Default is 'fragment'. |
+| `rows_per_range` | Long   | The number of rows per range when built using range mode. Default is 1000000.                                                                                                                            |
+
 
 ### FTS Options
 

lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala

@@ -85,10 +85,11 @@ case class AddIndexExec(
     val uuid = UUID.randomUUID()
     val indexType = IndexUtils.buildIndexType(method)
 
+    val dataset = Utils.openDatasetBuilder(readOptions).build()
+
     val indexBuildResult =
-      createIndexJob(lanceDataset, readOptions, uuid.toString, fragmentIds).run()
+      createIndexJob(dataset, lanceDataset, readOptions, uuid.toString, fragmentIds).run()
 
-    val dataset = Utils.openDatasetBuilder(readOptions).build()
     try {
       // Merge index metadata after all fragments are indexed
       dataset.mergeIndexMetadata(uuid.toString, indexType, Optional.empty())
@@ -148,6 +149,7 @@ case class AddIndexExec(
   }
 
   private def createIndexJob(
+      dataset: Dataset,
       lanceDataset: LanceDataset,
       readOptions: LanceSparkReadOptions,
       uuid: String,
@@ -179,7 +181,8 @@ case class AddIndexExec(
               nsImpl,
               nsProps,
               tableId,
-              initialStorageOpts)
+              initialStorageOpts,
+              dataset.getVersion.getManifestSummary.getTotalRows)
 
           case Some("fragment") | None =>
             new FragmentBasedIndexJob(
@@ -362,6 +365,7 @@ case class FragmentIndexTask(
  * @param nsProps            Optional namespace properties for credential vending
  * @param tableId            Optional table identifier for credential vending
  * @param initialStorageOpts Optional initial storage options for the dataset
+ * @param totalRows          Total number of rows in the dataset
  */
 class RangeBasedBTreeIndexJob(
     addIndexExec: AddIndexExec,
@@ -370,9 +374,11 @@ class RangeBasedBTreeIndexJob(
     nsImpl: Option[String],
     nsProps: Option[Map[String, String]],
     tableId: Option[List[String]],
-    initialStorageOpts: Option[Map[String, String]]) extends IndexJob {
+    initialStorageOpts: Option[Map[String, String]],
+    totalRows: Long) extends IndexJob {
 
   private val VALUE_COLUMN_NAME = "value"
+  private val DEFAULT_ROWS_PER_RANGE = 1000000L
 
   override def run(): IndexBuildResult = {
     if (addIndexExec.columns.size != 1) {
@@ -402,11 +408,15 @@ class RangeBasedBTreeIndexJob(
       df.select(df.col(columns.head).as(VALUE_COLUMN_NAME), df.col(LanceDataset.ROW_ID_COLUMN.name))
 
     // Repartition the data to numRanges and sort by indexed column
+    val rowsPerRange = addIndexExec.args.find(_.name == "rows_per_range").map(
+      _.value.asInstanceOf[Long]).getOrElse(DEFAULT_ROWS_PER_RANGE)
+    val numRange = Math.max(1L, totalRows / rowsPerRange.longValue())
+
     val rangeDf = selectDf
       .repartitionByRange(
-        session.sessionState.conf.numShufflePartitions,
+        numRange.intValue(),
         selectDf.col(VALUE_COLUMN_NAME).asc)
-      .sortWithinPartitions(VALUE_COLUMN_NAME)
+      .sortWithinPartitions(selectDf.col(VALUE_COLUMN_NAME).asc)
 
     val indexBuilder = RangeBTreeIndexBuilder(
       encode(readOptions),

lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java

@@ -229,6 +229,32 @@ public void testCreateBTreeIndexWithRangeMode() {
     Assertions.assertEquals("text_15", r.getString(1));
   }
 
+  @Test
+  public void testCreateBTreeIndexWithRowsPerRange() {
+    prepareDataset();
+    Dataset<Row> result =
+        spark.sql(
+            String.format(
+                "alter table %s create index test_index_btree_param using btree (id) "
+                    + "with (zone_size=2048, build_mode='range', rows_per_range=2)",
+                fullTable));
+    Assertions.assertEquals(
+        "StructType(StructField(fragments_indexed,LongType,true),StructField(index_name,StringType,true))",
+        result.schema().toString());
+    Row row = result.collectAsList().get(0);
+    long fragmentsIndexed = row.getLong(0);
+    String indexName = row.getString(1);
+    Assertions.assertTrue(fragmentsIndexed >= 2, "Expected at least 2 fragments to be indexed");
+    Assertions.assertEquals("test_index_btree_param", indexName);
+    checkIndex("test_index_btree_param");
+    // Verify query using the indexed field with zone_size parameter
+    Dataset<Row> query = spark.sql(String.format("select * from %s where id=15", fullTable));
+    Assertions.assertEquals(1L, query.count());
+    Row r = query.collectAsList().get(0);
+    Assertions.assertEquals(15, r.getInt(0));
+    Assertions.assertEquals("text_15", r.getString(1));
+  }
+
   @Test
   public void testCreateBTreeIndexWithFragmentMode() {
     prepareDataset();