feat: wire row limit through the scan ABI

Apply ScanRequest.limit end to end so all consumers agree: - scan core: ScanRequest gains `limit: Option<usize>`, applied as df.limit(0, Some(n)) after filters. - C ABI: df_scan_create gains an `int64 limit` parameter (negative means none); header updated. - JNI shim: build_request maps the proto's optional limit through. Tests: Rust limit_caps_row_count (demo provider, cap 2 across two partitions) and Java limitCapsRows (limit 2 over a 3-row CSV). Existing df_scan_create call sites pass -1. Rust suite 11, Java scan suite 5; clippy/fmt clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
timsaucer · Jun 12, 2026 · Jun 12, 2026 · Jun 12, 2026 · Jun 12, 2026 · Jun 12, 2026
commit d5eefee5dd327e7a815bf1c4c2eeaedcd75dae9d
diff --git a/core/src/test/java/org/apache/datafusion/scan/DatafusionScanTest.java b/core/src/test/java/org/apache/datafusion/scan/DatafusionScanTest.java
@@ -164,6 +164,25 @@ void filterPushdownSelectsRows() throws Exception {
     }
   }
 
+  @Test
+  void limitCapsRows() throws Exception {
+    byte[] config = csvConfig(writeCsv());
+    // Pushed limit of 2 over the 3-row CSV.
+    byte[] request = ScanRequest.newBuilder().setLimit(2).build().toByteArray();
+
+    try (BufferAllocator allocator = new RootAllocator();
+        DatafusionScan scan = DatafusionScan.create(PROVIDER, config, request)) {
+      int rows = 0;
+      try (ArrowReader reader = scan.execute(allocator)) {
+        VectorSchemaRoot root = reader.getVectorSchemaRoot();
+        while (reader.loadNextBatch()) {
+          rows += root.getRowCount();
+        }
+      }
+      assertEquals(2, rows, "limit should cap the scan at 2 rows");
+    }
+  }
+
   /** Serialize the LogicalExprNode for {@code id >= value}, as the engine's filter pushdown would. */
   private static byte[] idAtLeast(long value) {
     LogicalExprNode column =

diff --git a/native-ffi/include/datafusion_scan.h b/native-ffi/include/datafusion_scan.h
@@ -84,9 +84,10 @@ int32_t df_scan_schema(DfStr provider, DfBytes options, DfBytes partition,
 // Plan a scan. On success writes an owned handle to *out_handle (release with
 // df_scan_close). projection is an array of column-name DfStr (empty = all);
 // filters is an array of serialized datafusion.LogicalExprNode DfBytes;
-// target_partitions / batch_size <= 0 keep DataFusion defaults.
+// target_partitions / batch_size <= 0 keep DataFusion defaults; limit < 0 means
+// no row limit.
 int32_t df_scan_create(DfStr provider, DfBytes options, DfBytes partition,
-                       int32_t target_partitions, int32_t batch_size,
+                       int32_t target_partitions, int32_t batch_size, int64_t limit,
                        const DfKeyValue* config_overrides, size_t config_overrides_len,
                        const DfStr* projection, size_t projection_len,
                        const DfBytes* filters, size_t filters_len,

diff --git a/native-ffi/src/abi.rs b/native-ffi/src/abi.rs
@@ -113,10 +113,10 @@ pub unsafe extern "C" fn df_scan_schema(
 /// Plan a scan. On success writes an owned [`DfScanHandle`] pointer to
 /// `*out_handle`; the caller must release it with [`df_scan_close`].
 ///
-/// `config_keys`/`config_values` ... here folded into a single
-/// `config_overrides` array of [`DfKeyValue`]. `projection` is an array of
-/// column-name [`DfStr`]s (empty selects all). `filters` is an array of
-/// serialized `datafusion.LogicalExprNode` [`DfBytes`].
+/// Session config overrides are a single `config_overrides` array of
+/// [`DfKeyValue`]. `projection` is an array of column-name [`DfStr`]s (empty
+/// selects all). `filters` is an array of serialized `datafusion.LogicalExprNode`
+/// [`DfBytes`]. `limit` is the pushed row limit; a negative value means none.
 ///
 /// # Safety
 /// Array args follow the `(ptr, len)` borrow contract; `out_handle` must be a
@@ -129,6 +129,7 @@ pub unsafe extern "C" fn df_scan_create(
     partition: DfBytes,
     target_partitions: c_int,
     batch_size: c_int,
+    limit: i64,
     config_overrides: *const DfKeyValue,
     config_overrides_len: usize,
     projection: *const DfStr,
@@ -163,6 +164,7 @@ pub unsafe extern "C" fn df_scan_create(
             partition: partition.as_slice(),
             target_partitions,
             batch_size,
+            limit: if limit < 0 { None } else { Some(limit as usize) },
             config_overrides: overrides,
             projection: cols,
             filters: filter_bytes,

diff --git a/native-ffi/src/scan.rs b/native-ffi/src/scan.rs
@@ -63,6 +63,8 @@ pub struct ScanRequest<'a> {
     pub projection: Vec<String>,
     /// Each entry is a serialized `datafusion.LogicalExprNode`.
     pub filters: Vec<Vec<u8>>,
+    /// Optional row limit pushed into the scan. `None` means no limit.
+    pub limit: Option<usize>,
 }
 
 /// A planned scan. Holds the context alive for the plan's lifetime.
@@ -117,6 +119,9 @@ pub fn create(req: ScanRequest<'_>) -> ScanResult<ScanHandle> {
             .map_err(|e| ScanError::new(DfStatus::Planning, e.to_string()))?;
         df = df.filter(expr)?;
     }
+    if let Some(fetch) = req.limit {
+        df = df.limit(0, Some(fetch))?;
+    }
 
     // task_ctx() borrows df; capture before create_physical_plan consumes it.
     let task_ctx = Arc::new(df.task_ctx());

diff --git a/native-ffi/tests/listing.rs b/native-ffi/tests/listing.rs
@@ -130,6 +130,7 @@ fn listing_csv_scans_rows() {
             EMPTY,
             0,
             0,
+            -1,
             ptr::null(),
             0,
             ptr::null(),

diff --git a/native-ffi/tests/roundtrip.rs b/native-ffi/tests/roundtrip.rs
@@ -29,8 +29,8 @@ use datafusion::arrow::ffi::FFI_ArrowSchema;
 use datafusion::arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
 
 use datafusion_scan_ffi::abi::{
-    df_error_free, df_scan_abi_version, df_scan_close, df_scan_create, df_scan_execute_partition,
-    df_scan_partition_count, df_scan_schema, DfScanHandle,
+    df_error_free, df_scan_abi_version, df_scan_close, df_scan_create, df_scan_execute,
+    df_scan_execute_partition, df_scan_partition_count, df_scan_schema, DfScanHandle,
 };
 use datafusion_scan_ffi::ffi_types::{DfBytes, DfStr};
 use datafusion_scan_ffi::{demo, ABI_VERSION};
@@ -138,6 +138,48 @@ fn execute_partition_roundtrips_arrow_c_stream() {
     unsafe { df_scan_close(handle) };
 }
 
+#[test]
+fn limit_caps_row_count() {
+    demo::register();
+    // demo provider has 5 rows across two partitions; cap at 2.
+    let mut handle: *mut DfScanHandle = ptr::null_mut();
+    let mut err: *mut c_char = ptr::null_mut();
+    let status = unsafe {
+        df_scan_create(
+            provider(),
+            EMPTY_BYTES,
+            EMPTY_BYTES,
+            0,
+            0,
+            2, // limit
+            ptr::null(),
+            0,
+            ptr::null(),
+            0,
+            ptr::null(),
+            0,
+            &mut handle,
+            &mut err,
+        )
+    };
+    assert_eq!(status, 0, "err: {:?}", unsafe { take_err(err) });
+
+    // Read the whole plan; the limit must hold across partitions.
+    let mut stream = FFI_ArrowArrayStream::empty();
+    let mut err2: *mut c_char = ptr::null_mut();
+    assert_eq!(
+        unsafe { df_scan_execute(handle, &mut stream, &mut err2) },
+        0,
+        "err: {:?}",
+        unsafe { take_err(err2) }
+    );
+    let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream) }.expect("import");
+    let rows: usize = reader.map(|b| b.expect("batch").num_rows()).sum();
+    assert_eq!(rows, 2, "limit should cap the scan at 2 rows");
+
+    unsafe { df_scan_close(handle) };
+}
+
 #[test]
 fn close_is_null_safe() {
     unsafe { df_scan_close(ptr::null_mut()) };
@@ -154,6 +196,7 @@ fn create_full_scan() -> *mut DfScanHandle {
             EMPTY_BYTES,
             0,
             0,
+            -1,
             ptr::null(),
             0,
             ptr::null(),

diff --git a/native-jni/src/lib.rs b/native-jni/src/lib.rs
@@ -96,14 +96,13 @@ fn build_request<'a>(
         ProtoScanRequest::decode(scan_request)
             .map_err(|e| format!("failed to decode ScanRequest: {e}"))?
     };
-    // NOTE: `req.limit` is carried in the proto but not yet applied by the scan
-    // core or the C ABI; wire it through in a follow-up so both consumers agree.
     Ok(ScanRequest {
         provider,
         options: config,
         partition: &[],
         target_partitions: req.target_partitions,
         batch_size: req.batch_size,
+        limit: req.limit.map(|l| l as usize),
         config_overrides: req.config_overrides.into_iter().collect(),
         projection: req.projection,
         filters: req.filters,
-Original file line number
+Diff line change
@@ Expand Up / @@ -130,6 +130,7 @@ fn listing_csv_scans_rows() { @@
                 EMPTY,
 ,
 ,
+                -1,
                 ptr::null(),
 ,
                 ptr::null(),
@@ Expand Down @@