feat(java): add Dataset.sample() API (#6500)

## Summary

- Expose the Rust core `Dataset::sample()` method through the Java JNI
interface
- Allows Java users to randomly sample `n` rows from a dataset, with
optional fragment ID filtering
- Returns results in row-id order for efficient underlying take
operations

## Changes

- **`java/lance-jni/src/blocking_dataset.rs`**: Add `nativeSample` JNI
function that projects schema from column names, calls
`dataset.sample()`, and serializes the result as Arrow IPC bytes
- **`java/src/main/java/org/lance/Dataset.java`**: Add `sample(n,
columns)` and `sample(n, columns, fragmentIds)` public methods with
input validation
- **`java/src/test/java/org/lance/DatasetTest.java`**: Add `testSample`
covering basic sampling, fragment-filtered sampling, and over-sampling
(n > total rows)

## Test plan

- [ ] `testSample` — verifies correct row counts for basic sample,
fragment-filtered sample, and over-sample cases
- [ ] Rust JNI compiles cleanly with `cargo check`

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: beinan

java/lance-jni/src/blocking_dataset.rs

@@ -1978,6 +1978,64 @@ fn inner_take(
     Ok(**byte_array)
 }
 
+#[unsafe(no_mangle)]
+pub extern "system" fn Java_org_lance_Dataset_nativeSample(
+    mut env: JNIEnv,
+    java_dataset: JObject,
+    n: jlong,
+    columns_obj: JObject,      // List<String>
+    fragment_ids_obj: JObject, // Optional<List<Integer>>
+) -> jbyteArray {
+    match inner_sample(&mut env, java_dataset, n, columns_obj, fragment_ids_obj) {
+        Ok(byte_array) => byte_array,
+        Err(e) => {
+            let _ = env.throw_new("java/lang/RuntimeException", format!("{:?}", e));
+            std::ptr::null_mut()
+        }
+    }
+}
+
+fn inner_sample(
+    env: &mut JNIEnv,
+    java_dataset: JObject,
+    n: jlong,
+    columns_obj: JObject,      // List<String>
+    fragment_ids_obj: JObject, // Optional<List<Integer>>
+) -> Result<jbyteArray> {
+    let columns: Vec<String> = env.get_strings(&columns_obj)?;
+    let fragment_ids: Option<Vec<i32>> = env.get_ints_opt(&fragment_ids_obj)?;
+    let fragment_ids_u32: Option<Vec<u32>> =
+        fragment_ids.map(|ids| ids.iter().map(|&id| id as u32).collect());
+
+    let result = {
+        let dataset_guard =
+            unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
+        let dataset = &dataset_guard.inner;
+
+        let projection = dataset
+            .schema()
+            .project_preserve_system_columns(&columns)
+            .map_err(|e| Error::runtime_error(e.to_string()))?;
+
+        match RT.block_on(dataset.sample(n as usize, &projection, fragment_ids_u32.as_deref())) {
+            Ok(res) => res,
+            Err(e) => {
+                return Err(e.into());
+            }
+        }
+    };
+
+    let mut buffer = Vec::new();
+    {
+        let mut writer = StreamWriter::try_new(&mut buffer, &result.schema())?;
+        writer.write(&result)?;
+        writer.finish()?;
+    }
+
+    let byte_array = env.byte_array_from_slice(&buffer)?;
+    Ok(**byte_array)
+}
+
 #[unsafe(no_mangle)]
 pub extern "system" fn Java_org_lance_Dataset_nativeDelete(
     mut env: JNIEnv,

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

@@ -756,6 +756,57 @@ public void close() throws IOException {
 
   private native byte[] nativeTake(List<Long> indices, List<String> columns);
 
+  /**
+   * Randomly sample n rows from the dataset.
+   *
+   * <p>The returned rows are in row-id order (not random order), which allows the underlying take
+   * operation to use an efficient sorted code path.
+   *
+   * @param n the number of rows to sample
+   * @param columns the columns to include in the result
+   * @return an ArrowReader containing the sampled rows
+   * @throws IOException if an I/O error occurs
+   */
+  public ArrowReader sample(long n, List<String> columns) throws IOException {
+    return sample(n, columns, Optional.empty());
+  }
+
+  /**
+   * Randomly sample n rows from specific fragments of the dataset.
+   *
+   * <p>The returned rows are in row-id order (not random order), which allows the underlying take
+   * operation to use an efficient sorted code path.
+   *
+   * @param n the number of rows to sample
+   * @param columns the columns to include in the result
+   * @param fragmentIds optional list of fragment IDs to restrict sampling to
+   * @return an ArrowReader containing the sampled rows
+   * @throws IOException if an I/O error occurs
+   */
+  public ArrowReader sample(long n, List<String> columns, Optional<List<Integer>> fragmentIds)
+      throws IOException {
+    Preconditions.checkArgument(n > 0, "n must be greater than 0");
+    Preconditions.checkNotNull(columns, "columns cannot be null");
+    Preconditions.checkArgument(!columns.isEmpty(), "columns cannot be empty");
+    Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
+    try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
+      byte[] arrowData = nativeSample(n, columns, fragmentIds);
+      ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(arrowData);
+      ReadableByteChannel readChannel = Channels.newChannel(byteArrayInputStream);
+      return new ArrowStreamReader(readChannel, allocator) {
+        @Override
+        public void close() throws IOException {
+          super.close();
+          readChannel.close();
+          byteArrayInputStream.close();
+        }
+      };
+    }
+  }
+
+  private native byte[] nativeSample(
+      long n, List<String> columns, Optional<List<Integer>> fragmentIds);
+
   /**
    * Delete rows of data by predicate.
    *

java/src/test/java/org/lance/DatasetTest.java

@@ -839,6 +839,49 @@ void testTake(@TempDir Path tempDir) throws IOException, ClosedChannelException
     }
   }
 
+  @Test
+  void testSample(@TempDir Path tempDir) throws IOException {
+    String testMethodName = new Object() {}.getClass().getEnclosingMethod().getName();
+    String datasetPath = tempDir.resolve(testMethodName).toString();
+    try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+      TestUtils.SimpleTestDataset testDataset =
+          new TestUtils.SimpleTestDataset(allocator, datasetPath);
+      dataset = testDataset.createEmptyDataset();
+
+      try (Dataset dataset2 = testDataset.write(1, 20)) {
+        // Sample without fragment filter
+        List<String> columns = Arrays.asList("id", "name");
+        try (ArrowReader reader = dataset2.sample(5, columns)) {
+          assertTrue(reader.loadNextBatch());
+          VectorSchemaRoot result = reader.getVectorSchemaRoot();
+          assertNotNull(result);
+          assertEquals(5, result.getRowCount());
+          assertEquals(2, result.getSchema().getFields().size());
+        }
+
+        // Sample with fragment filter
+        List<Fragment> fragments = dataset2.getFragments();
+        assertFalse(fragments.isEmpty());
+        List<Integer> fragmentIds =
+            fragments.stream().map(f -> f.getId()).collect(Collectors.toList());
+        try (ArrowReader reader = dataset2.sample(3, columns, Optional.of(fragmentIds))) {
+          assertTrue(reader.loadNextBatch());
+          VectorSchemaRoot result = reader.getVectorSchemaRoot();
+          assertNotNull(result);
+          assertEquals(3, result.getRowCount());
+        }
+
+        // Sample more than available rows returns all rows
+        try (ArrowReader reader = dataset2.sample(100, columns)) {
+          assertTrue(reader.loadNextBatch());
+          VectorSchemaRoot result = reader.getVectorSchemaRoot();
+          assertNotNull(result);
+          assertEquals(20, result.getRowCount());
+        }
+      }
+    }
+  }
+
   @Test
   void testCountRows(@TempDir Path tempDir) {
     String testMethodName = new Object() {}.getClass().getEnclosingMethod().getName();