feat(java): count indexed rows on explicit scalar segments

java/lance-jni/src/blocking_dataset.rs

@@ -4018,73 +4018,105 @@ pub extern "system" fn Java_org_lance_Dataset_nativeCountIndexedRows(
 fn inner_count_indexed_rows(
     env: &mut JNIEnv,
     java_dataset: JObject,
-    _jindex_name: JString,
+    jindex_name: JString,
     jfilter: JString,
     jfragment_ids: JObject, // Optional<List<Integer>>
 ) -> Result<i64> {
+    let index_name: String = jindex_name.extract(env)?;
     let filter: String = jfilter.extract(env)?;
+    let fragment_ids = extract_count_fragment_ids(env, &jfragment_ids)?;
 
-    // Extract optional fragment IDs
-    let fragment_ids: Option<Vec<u32>> = if env
-        .call_method(&jfragment_ids, "isPresent", "()Z", &[])?
-        .z()?
-    {
-        let list_obj = env
-            .call_method(&jfragment_ids, "get", "()Ljava/lang/Object;", &[])?
-            .l()?;
-        let list = env.get_list(&list_obj)?;
-        let mut ids = Vec::new();
-        let mut iter = list.iter(env)?;
-        while let Some(elem) = iter.next(env)? {
-            let int_val = env.call_method(&elem, "intValue", "()I", &[])?.i()?;
-            ids.push(int_val as u32);
-        }
-        Some(ids)
-    } else {
-        None
-    };
-
-    let count = {
-        let dataset_guard =
-            unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
-
-        // Use a scanner with fragment filtering to count rows
-        // This ensures we only count rows in the specified fragments
-        let inner = dataset_guard.inner.clone();
+    count_indexed_rows(env, java_dataset, index_name, filter, None, fragment_ids)
+}
 
-        block_on(async {
-            let mut scanner = inner.scan();
+#[unsafe(no_mangle)]
+pub extern "system" fn Java_org_lance_Dataset_nativeCountIndexedRowsWithSegments(
+    mut env: JNIEnv,
+    java_dataset: JObject,
+    jindex_name: JString,
+    jfilter: JString,
+    jsegment_uuids: JObject, // List<UUID>
+    jfragment_ids: JObject,  // Optional<List<Integer>>
+) -> jlong {
+    ok_or_throw_with_return!(
+        env,
+        inner_count_indexed_rows_with_segments(
+            &mut env,
+            java_dataset,
+            jindex_name,
+            jfilter,
+            jsegment_uuids,
+            jfragment_ids,
+        ),
+        -1
+    )
+}
 
-            // Apply filter
-            if !filter.is_empty() {
-                scanner.filter(&filter)?;
-            }
+fn inner_count_indexed_rows_with_segments(
+    env: &mut JNIEnv,
+    java_dataset: JObject,
+    jindex_name: JString,
+    jfilter: JString,
+    jsegment_uuids: JObject,
+    jfragment_ids: JObject,
+) -> Result<i64> {
+    let index_name: String = jindex_name.extract(env)?;
+    let filter: String = jfilter.extract(env)?;
+    let segment_uuids =
+        import_vec_to_rust(env, &jsegment_uuids, |env, obj| obj.extract_object(env))?;
+    let fragment_ids = extract_count_fragment_ids(env, &jfragment_ids)?;
 
-            // Empty projection and enable row_id for count_rows to work
-            // count_rows() requires metadata-only projection
-            scanner.project::<String>(&[])?;
-            scanner.with_row_id();
-
-            // Apply fragment filter if specified
-            if let Some(frag_ids) = fragment_ids {
-                // Convert FileFragment to Fragment by extracting metadata
-                let filtered_fragments: Vec<_> = inner
-                    .get_fragments()
-                    .into_iter()
-                    .filter(|f| frag_ids.contains(&(f.id() as u32)))
-                    .map(|f| f.metadata().clone())
-                    .collect();
-                scanner.with_fragments(filtered_fragments);
-            }
+    count_indexed_rows(
+        env,
+        java_dataset,
+        index_name,
+        filter,
+        Some(segment_uuids),
+        fragment_ids,
+    )
+}
 
-            // Use the scanner's count_rows method
-            let count = scanner.count_rows().await?;
+fn extract_count_fragment_ids(
+    env: &mut JNIEnv,
+    jfragment_ids: &JObject,
+) -> Result<Option<Vec<u32>>> {
+    env.get_ints_opt(jfragment_ids)?
+        .map(|fragment_ids| {
+            fragment_ids
+                .into_iter()
+                .map(|fragment_id| {
+                    u32::try_from(fragment_id).map_err(|_| {
+                        Error::input_error(format!(
+                            "fragment ID must be non-negative, got {fragment_id}"
+                        ))
+                    })
+                })
+                .collect()
+        })
+        .transpose()
+}
 
-            Ok::<i64, lance::Error>(count as i64)
-        })?
+fn count_indexed_rows(
+    env: &mut JNIEnv,
+    java_dataset: JObject,
+    index_name: String,
+    filter: String,
+    segment_uuids: Option<Vec<Uuid>>,
+    fragment_ids: Option<Vec<u32>>,
+) -> Result<i64> {
+    let count = {
+        let dataset_guard =
+            unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
+        block_on(dataset_guard.inner.count_indexed_rows(
+            &index_name,
+            &filter,
+            segment_uuids.as_deref(),
+            fragment_ids.as_deref(),
+        ))?
     };
 
-    Ok(count)
+    i64::try_from(count)
+        .map_err(|_| Error::input_error(format!("indexed row count {count} exceeds Java long")))
 }
 
 //////////////////////////////

java/src/main/java/org/lance/Dataset.java

@@ -71,6 +71,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 
@@ -1482,6 +1483,10 @@ public Session session() {
    * counts matching row addresses, which is more efficient than scanning when the index covers the
    * filter column.
    *
+   * <p>Planning is pinned to {@code indexName}. A filter that cannot be answered by that scalar
+   * index is rejected instead of scanning the table or selecting another index. Deleted rows are
+   * excluded.
+   *
    * @param indexName the name of the scalar index to use
    * @param filter the filter expression (e.g., "column = 5")
    * @param fragmentIds optional list of fragment IDs to restrict the count to
@@ -1499,9 +1504,58 @@ public long countIndexedRows(
     }
   }
 
+  /**
+   * Count rows matching a filter using explicit physical segments of a scalar index.
+   *
+   * <p>Only {@code segmentUuids} are opened. Their current fragment coverage defines the count
+   * scope: matching deleted rows inside that scope are excluded, and rows outside it are not
+   * counted. When {@code fragmentIds} is omitted, the scope is derived from that coverage. When it
+   * is present, its set must equal the coverage; order does not matter. A mismatch is rejected with
+   * an error that reports both sets.
+   *
+   * <p>The selection is accepted only when it includes every segment that contributes to that
+   * coverage. After fragment reuse, one source segment can advertise every destination fragment
+   * while still depending on its siblings. An incomplete selection is rejected, and the error names
+   * the missing segment UUIDs. A segment whose coverage does not overlap the selection can be
+   * queried alone.
+   *
+   * <p>An empty segment list, duplicate segment UUIDs, an unknown UUID, a UUID from a different
+   * index, or a segment without fragment coverage is rejected. The existing three-argument method
+   * remains available and does not take a segment list.
+   *
+   * @param indexName the logical scalar index name that every selected segment must belong to
+   * @param filter the filter expression (e.g., "column = 5")
+   * @param segmentUuids physical segment UUIDs to open; must be non-empty and contain no duplicates
+   * @param fragmentIds optional fragment IDs that must match the selected segments' current
+   *     coverage
+   * @return count of matching rows in the selected segment scope
+   */
+  public long countIndexedRows(
+      String indexName,
+      String filter,
+      List<UUID> segmentUuids,
+      Optional<List<Integer>> fragmentIds) {
+    try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
+      Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
+      Preconditions.checkArgument(
+          indexName != null && !indexName.isEmpty(), "indexName cannot be null or empty");
+      Preconditions.checkArgument(
+          filter != null && !filter.isEmpty(), "filter cannot be null or empty");
+      Preconditions.checkNotNull(segmentUuids, "segmentUuids cannot be null");
+      Preconditions.checkArgument(!segmentUuids.isEmpty(), "segmentUuids cannot be empty");
+      return nativeCountIndexedRowsWithSegments(indexName, filter, segmentUuids, fragmentIds);
+    }
+  }
+
   private native long nativeCountIndexedRows(
       String indexName, String filter, Optional<List<Integer>> fragmentIds);
 
+  private native long nativeCountIndexedRowsWithSegments(
+      String indexName,
+      String filter,
+      List<UUID> segmentUuids,
+      Optional<List<Integer>> fragmentIds);
+
   /**
    * Calculate the size of the dataset.
    *