perf(java): reduce fragment statistics allocations

java/lance-jni/src/blocking_dataset.rs

@@ -1652,7 +1652,7 @@ pub extern "system" fn Java_org_lance_Dataset_nativeGetFragmentStatistics<'a>(
     ok_or_throw!(env, inner_get_fragment_statistics(&mut env, jdataset))
 }
 
-/// Returns per-fragment statistics flattened as [id0, rowCount0, dataFileNum0, id1, ...].
+/// Returns per-fragment statistics in their final Java primitive arrays.
 ///
 /// Row count semantics match Java `FragmentMetadata.getNumRows()`:
 /// physical rows minus deleted rows, with absent values treated as 0.
@@ -1661,28 +1661,62 @@ fn inner_get_fragment_statistics<'local>(
     env: &mut JNIEnv<'local>,
     jdataset: JObject,
 ) -> Result<JObject<'local>> {
-    let stats: Vec<i64> = {
+    // Three 4096-entry typed buffers use 64 KiB while keeping JNI calls amortized.
+    const CHUNK_SIZE: usize = 4096;
+
+    let fragments = {
         let dataset =
             unsafe { env.get_rust_field::<_, _, BlockingDataset>(jdataset, NATIVE_DATASET) }?;
-        let fragments = dataset.inner.get_fragments();
-        let mut stats = Vec::with_capacity(fragments.len() * 3);
-        for f in fragments.iter() {
-            let meta = f.metadata();
-            let physical_rows = meta.physical_rows.unwrap_or(0) as i64;
-            let deleted_rows = meta
+        dataset.inner.fragments().clone()
+    };
+    let fragment_count = i32::try_from(fragments.len()).map_err(|_| {
+        Error::runtime_error(format!(
+            "Fragment statistics contain {} fragments, exceeding the Java array limit of {}",
+            fragments.len(),
+            i32::MAX
+        ))
+    })?;
+    let ids = env.new_int_array(fragment_count)?;
+    let row_counts = env.new_long_array(fragment_count)?;
+    let data_file_nums = env.new_int_array(fragment_count)?;
+
+    let chunk_capacity = fragments.len().min(CHUNK_SIZE);
+    let mut id_chunk = Vec::with_capacity(chunk_capacity);
+    let mut row_count_chunk = Vec::with_capacity(chunk_capacity);
+    let mut data_file_num_chunk = Vec::with_capacity(chunk_capacity);
+
+    for (chunk_index, fragment_chunk) in fragments.chunks(CHUNK_SIZE).enumerate() {
+        id_chunk.clear();
+        row_count_chunk.clear();
+        data_file_num_chunk.clear();
+
+        for fragment in fragment_chunk {
+            let physical_rows = fragment.physical_rows.unwrap_or(0) as i64;
+            let deleted_rows = fragment
                 .deletion_file
                 .as_ref()
-                .and_then(|d| d.num_deleted_rows)
+                .and_then(|deletion_file| deletion_file.num_deleted_rows)
                 .unwrap_or(0) as i64;
-            stats.push(f.id() as i64);
-            stats.push(physical_rows - deleted_rows);
-            stats.push(meta.files.len() as i64);
+            id_chunk.push(fragment.id as i32);
+            row_count_chunk.push(physical_rows - deleted_rows);
+            data_file_num_chunk.push(fragment.files.len() as i32);
         }
-        stats
-    };
-    let jarray = env.new_long_array(stats.len() as i32)?;
-    env.set_long_array_region(&jarray, 0, &stats)?;
-    Ok(jarray.into())
+
+        let offset = (chunk_index * CHUNK_SIZE) as i32;
+        env.set_int_array_region(&ids, offset, &id_chunk)?;
+        env.set_long_array_region(&row_counts, offset, &row_count_chunk)?;
+        env.set_int_array_region(&data_file_nums, offset, &data_file_num_chunk)?;
+    }
+
+    Ok(env.new_object(
+        "org/lance/FragmentStatistics",
+        "([I[J[I)V",
+        &[
+            JValue::Object(&ids),
+            JValue::Object(&row_counts),
+            JValue::Object(&data_file_nums),
+        ],
+    )?)
 }
 
 #[unsafe(no_mangle)]

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

@@ -1411,31 +1411,20 @@ public List<Fragment> getFragments() {
    * Get per-fragment statistics for all fragments in this dataset version.
    *
    * <p>Unlike {@link #getFragments()}, this is a metadata-only bulk operation: no per-fragment Java
-   * objects are materialized, making it suitable for planning over datasets with a very large
-   * number of fragments. Row counts match {@link FragmentMetadata#getNumRows()} (physical rows
-   * minus deleted rows).
+   * objects are materialized, and native code fills the returned primitive arrays directly. This
+   * makes it suitable for planning over datasets with a very large number of fragments. Row counts
+   * match {@link FragmentMetadata#getNumRows()} (physical rows minus deleted rows).
    *
    * @return per-fragment statistics as parallel arrays, in manifest order
    */
   public FragmentStatistics getFragmentStatistics() {
     try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
       Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
-      // Flattened as [id0, rowCount0, dataFileNum0, id1, ...] to keep the JNI surface primitive
-      long[] flat = nativeGetFragmentStatistics();
-      int count = flat.length / 3;
-      int[] ids = new int[count];
-      long[] rowCounts = new long[count];
-      int[] dataFileNums = new int[count];
-      for (int i = 0; i < count; i++) {
-        ids[i] = (int) flat[3 * i];
-        rowCounts[i] = flat[3 * i + 1];
-        dataFileNums[i] = (int) flat[3 * i + 2];
-      }
-      return new FragmentStatistics(ids, rowCounts, dataFileNums);
+      return nativeGetFragmentStatistics();
     }
   }
 
-  private native long[] nativeGetFragmentStatistics();
+  private native FragmentStatistics nativeGetFragmentStatistics();
 
   /**
    * Gets the arrow schema of the dataset.

java/src/test/java/org/lance/FragmentTest.java

@@ -499,6 +499,9 @@ void testFragmentStatistics(@TempDir Path tempDir) {
             stats.getDataFileNums());
 
         assertEquals(30, Arrays.stream(stats.getRowCounts()).sum());
+
+        dataset.delete("id < 5");
+        assertArrayEquals(new long[] {16, 4}, dataset.getFragmentStatistics().getRowCounts());
       }
     }
   }
@@ -514,4 +517,56 @@ void testFragmentStatisticsOnEmptyDataset(@TempDir Path tempDir) {
       }
     }
   }
+
+  @Test
+  void testFragmentStatisticsAcrossNativeChunks(@TempDir Path tempDir) {
+    String datasetPath = tempDir.resolve("fragment_statistics_chunks").toString();
+    try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+      TestUtils.SimpleTestDataset testDataset =
+          new TestUtils.SimpleTestDataset(allocator, datasetPath);
+      testDataset.createEmptyDataset().close();
+
+      FragmentMetadata template = testDataset.createNewFragment(1);
+      int fragmentCount = 4097;
+      List<FragmentMetadata> fragments = new ArrayList<>(fragmentCount);
+      for (int id = 0; id < fragmentCount; id++) {
+        fragments.add(
+            new FragmentMetadata(
+                id,
+                template.getFiles(),
+                template.getPhysicalRows(),
+                template.getDeletionFile(),
+                template.getRowIdMeta()));
+      }
+
+      FragmentOperation.Append appendOp = new FragmentOperation.Append(fragments);
+      try (Dataset dataset = Dataset.commit(allocator, datasetPath, appendOp, Optional.of(1L))) {
+        FragmentStatistics stats = dataset.getFragmentStatistics();
+        int lastIndex = fragmentCount - 1;
+        assertEquals(fragmentCount, stats.size());
+        assertEquals(0, stats.getIds()[0]);
+        assertEquals(lastIndex, stats.getIds()[lastIndex]);
+        assertEquals(1, stats.getRowCounts()[0]);
+        assertEquals(1, stats.getRowCounts()[lastIndex]);
+        assertEquals(1, stats.getDataFileNums()[0]);
+        assertEquals(1, stats.getDataFileNums()[lastIndex]);
+      }
+    }
+  }
+
+  @Test
+  void testFragmentStatisticsPreservesLegacyMissingRowCount() {
+    String historicalPath =
+        Path.of("..", "test_data", "v0.7.5", "with_deletions")
+            .toAbsolutePath()
+            .normalize()
+            .toString();
+    try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+        Dataset dataset = Dataset.open(historicalPath, allocator)) {
+      FragmentStatistics stats = dataset.getFragmentStatistics();
+      assertArrayEquals(new int[] {0}, stats.getIds());
+      assertArrayEquals(new long[] {0}, stats.getRowCounts());
+      assertArrayEquals(new int[] {1}, stats.getDataFileNums());
+    }
+  }
 }

rust/lance/Cargo.toml

@@ -205,6 +205,10 @@ harness = false
 name = "count_pushdown"
 harness = false
 
+[[bench]]
+name = "fragment_statistics"
+harness = false
+
 [[bench]]
 name = "vector_index"
 harness = false

rust/lance/benches/fragment_statistics.rs

@@ -0,0 +1,152 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+//! Compares the old and new per-fragment JNI payload preparation paths.
+//!
+//! The flattened baseline mirrors the old Java `Dataset.getFragmentStatistics()` path: it creates
+//! `FileFragment` wrappers and allocates three `i64` values per fragment before JNI copies and
+//! Java-side array splitting. The typed-chunk path mirrors the current implementation's native
+//! preparation before it copies directly into the three final Java arrays.
+//!
+//! At 100,000 fragments, the old flattened path allocates a 2.4 MB native vector, followed by a
+//! 2.4 MB Java `long[]` JNI copy and 1.6 MB across the three final Java primitive arrays. The
+//! typed-chunk path bounds native staging memory to 64 KiB and creates only the 1.6 MB final Java
+//! arrays.
+//!
+//! ```text
+//! cargo bench -p lance --bench fragment_statistics
+//! ```
+
+use std::hint::black_box;
+
+use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
+use criterion::{Criterion, Throughput, criterion_group, criterion_main};
+use lance::Dataset;
+use lance::dataset::transaction::Operation;
+use lance_core::utils::tempfile::TempStrDir;
+use lance_table::format::Fragment;
+
+const NUM_FRAGMENTS: usize = 100_000;
+const STATISTICS_CHUNK_SIZE: usize = 4096;
+
+struct Fixture {
+    _data_dir: TempStrDir,
+    dataset: Dataset,
+}
+
+impl Fixture {
+    async fn open() -> Self {
+        let data_dir = TempStrDir::default();
+        let schema =
+            lance_core::datatypes::Schema::try_from(&ArrowSchema::new(vec![ArrowField::new(
+                "value",
+                DataType::Int64,
+                false,
+            )]))
+            .unwrap();
+        let fragments = (0..NUM_FRAGMENTS)
+            .map(|id| {
+                let mut fragment = Fragment::new(id as u64);
+                fragment.physical_rows = Some(1_000 + id % 100);
+                fragment
+            })
+            .collect();
+        let operation = Operation::Overwrite {
+            fragments,
+            schema,
+            config_upsert_values: None,
+            initial_bases: None,
+        };
+        let dataset = Dataset::commit(
+            data_dir.as_str(),
+            operation,
+            None,
+            None,
+            None,
+            Default::default(),
+            false,
+        )
+        .await
+        .unwrap();
+
+        Self {
+            _data_dir: data_dir,
+            dataset,
+        }
+    }
+}
+
+fn legacy_flattened_fragment_statistics(dataset: &Dataset) -> Vec<i64> {
+    let fragments = dataset.get_fragments();
+    let mut statistics = Vec::with_capacity(fragments.len() * 3);
+    for fragment in fragments.iter() {
+        let metadata = fragment.metadata();
+        let physical_rows = metadata.physical_rows.unwrap_or(0) as i64;
+        let deleted_rows = metadata
+            .deletion_file
+            .as_ref()
+            .and_then(|deletion_file| deletion_file.num_deleted_rows)
+            .unwrap_or(0) as i64;
+        statistics.push(metadata.id as i64);
+        statistics.push(physical_rows - deleted_rows);
+        statistics.push(metadata.files.len() as i64);
+    }
+    statistics
+}
+
+fn prepare_typed_fragment_statistics_chunks(dataset: &Dataset) -> usize {
+    let fragments = dataset.fragments();
+    let chunk_capacity = fragments.len().min(STATISTICS_CHUNK_SIZE);
+    let mut ids = Vec::with_capacity(chunk_capacity);
+    let mut row_counts = Vec::with_capacity(chunk_capacity);
+    let mut data_file_nums = Vec::with_capacity(chunk_capacity);
+    let mut value_count = 0;
+
+    for fragments in fragments.chunks(STATISTICS_CHUNK_SIZE) {
+        ids.clear();
+        row_counts.clear();
+        data_file_nums.clear();
+        for fragment in fragments {
+            let physical_rows = fragment.physical_rows.unwrap_or(0) as i64;
+            let deleted_rows = fragment
+                .deletion_file
+                .as_ref()
+                .and_then(|deletion_file| deletion_file.num_deleted_rows)
+                .unwrap_or(0) as i64;
+            ids.push(fragment.id as i32);
+            row_counts.push(physical_rows - deleted_rows);
+            data_file_nums.push(fragment.files.len() as i32);
+        }
+        value_count += ids.len() + row_counts.len() + data_file_nums.len();
+        black_box((&ids, &row_counts, &data_file_nums));
+    }
+
+    value_count
+}
+
+fn bench_fragment_statistics(c: &mut Criterion) {
+    let runtime = tokio::runtime::Runtime::new().unwrap();
+    let fixture = runtime.block_on(Fixture::open());
+
+    assert_eq!(
+        legacy_flattened_fragment_statistics(&fixture.dataset).len(),
+        NUM_FRAGMENTS * 3
+    );
+    assert_eq!(
+        prepare_typed_fragment_statistics_chunks(&fixture.dataset),
+        NUM_FRAGMENTS * 3
+    );
+
+    let mut group = c.benchmark_group("fragment_statistics/100k_fragments");
+    group.throughput(Throughput::Elements(NUM_FRAGMENTS as u64));
+    group.bench_function("legacy_materialize_flattened_statistics", |b| {
+        b.iter(|| black_box(legacy_flattened_fragment_statistics(&fixture.dataset)))
+    });
+    group.bench_function("prepare_typed_statistics_chunks", |b| {
+        b.iter(|| black_box(prepare_typed_fragment_statistics_chunks(&fixture.dataset)))
+    });
+    group.finish();
+}
+
+criterion_group!(benches, bench_fragment_statistics);
+criterion_main!(benches);