feat: add num_segments option for configurable zonemap segment count

Allow users to control the number of index segments created during
distributed zonemap builds via the num_segments DDL option.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: beinan

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

@@ -39,6 +39,7 @@ For the `zonemap` method, the following options are supported:
 | Option          | Type | Description                                  |
 |-----------------|------|----------------------------------------------|
 | `rows_per_zone` | Long | The approximate number of rows per zonemap zone. |
+| `num_segments` | Integer | Number of index segments to create. Each segment covers a batch of fragments. Defaults to `min(fragment_count, spark.default.parallelism)`. |
 
 ### BTree Options
 

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

@@ -92,6 +92,12 @@ case class AddIndexExec(
     val btreeBuildMode = IndexUtils.btreeBuildMode(indexType, args)
     val useLogicalSegmentCommit = IndexUtils.useLogicalSegmentCommit(indexType)
 
+    val numSegmentsOpt = args.find(_.name == "num_segments")
+    if (numSegmentsOpt.isDefined && !useLogicalSegmentCommit) {
+      throw new IllegalArgumentException(
+        "num_segments option is only supported for index types that use segmented builds (e.g., zonemap)")
+    }
+
     // Create distributed index job and run it
     val createdSegments = createIndexJob(
       lanceDataset,
@@ -317,8 +323,9 @@ class FragmentBasedIndexJob(
     val encodedReadOptions = encode(readOptions)
     val columns = addIndexExec.columns.toList
     val argsJson = IndexUtils.toJson(addIndexExec.args)
+    val numSegments = addIndexExec.args.find(_.name == "num_segments").map(_.value.asInstanceOf[Number].intValue())
     val fragmentBatches = if (groupFragmentsIntoSegments) {
-      batchFragments(fragmentIds)
+      batchFragments(fragmentIds, numSegments)
     } else {
       fragmentIds.map(fid => List(fid))
     }
@@ -346,9 +353,11 @@ class FragmentBasedIndexJob(
       .toSeq
   }
 
-  private def batchFragments(fragmentIds: List[Integer]): Seq[List[Integer]] = {
-    val targetTasks =
-      math.max(1, math.min(fragmentIds.size, addIndexExec.session.sparkContext.defaultParallelism))
+  private def batchFragments(fragmentIds: List[Integer], numSegments: Option[Int] = None): Seq[List[Integer]] = {
+    val targetTasks = numSegments match {
+      case Some(n) => math.max(1, math.min(fragmentIds.size, n))
+      case None => math.max(1, math.min(fragmentIds.size, addIndexExec.session.sparkContext.defaultParallelism))
+    }
     val batchSize = math.ceil(fragmentIds.size.toDouble / targetTasks.toDouble).toInt
     fragmentIds.grouped(batchSize).map(_.toList).toSeq
   }

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

@@ -335,6 +335,53 @@ public void testRepeatedCreateZonemapIndexReplacesExistingSegments() {
     }
   }
 
+  @Test
+  public void testCreateZonemapIndexWithNumSegments() {
+    prepareDataset();
+
+    Dataset<Row> result =
+        spark.sql(
+            String.format(
+                "alter table %s create index test_index_zonemap_segments using zonemap (id) with (num_segments = 3)",
+                fullTable));
+
+    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_zonemap_segments", indexName);
+
+    // Verify the number of segments matches the requested num_segments
+    org.lance.Dataset lanceDataset = org.lance.Dataset.open().uri(tableDir).build();
+    try {
+      int fragmentCount = lanceDataset.getFragments().size();
+      int expectedSegmentCount = Math.min(fragmentCount, 3);
+      List<Index> segments =
+          lanceDataset.getIndexes().stream()
+              .filter(index -> "test_index_zonemap_segments".equals(index.name()))
+              .collect(Collectors.toList());
+
+      Assertions.assertEquals(
+          expectedSegmentCount,
+          segments.size(),
+          "Expected num_segments=3 to produce exactly 3 segments (or fewer if fragment count < 3)");
+
+      int coveredFragments =
+          segments.stream()
+              .map(index -> index.fragments().orElse(Collections.emptyList()).size())
+              .mapToInt(Integer::intValue)
+              .sum();
+      Assertions.assertEquals(
+          fragmentCount,
+          coveredFragments,
+          "Expected committed segments to cover all fragments exactly once");
+    } finally {
+      lanceDataset.close();
+    }
+  }
+
+
   @Test
   public void testCreateBTreeIndexWithRangeMode() {
     prepareDataset();