fix: row_range vs row_indices pushdown in JNI (#3514)

Author: a10y

java/vortex-jni/src/main/java/dev/vortex/api/ScanOptions.java

@@ -34,6 +34,12 @@ public interface ScanOptions {
      */
     Optional<Expression> predicate();
 
+    /**
+     * Optional start (inclusive) and end (exclusive) row indices to select a range of rows
+     * in the scan.
+     */
+    Optional<long[]> rowRange();
+
     /**
      * Optional row indices to select specific rows.
      * These must be sorted in ascending order.

java/vortex-jni/src/main/java/dev/vortex/jni/JNIFile.java

@@ -47,9 +47,10 @@ public ArrayIterator newScan(ScanOptions options) {
         }
 
         long[] rowIndices = options.rowIndices().orElse(null);
+        long[] rowRange = options.rowRange().orElse(null);
 
         return new JNIArrayIterator(
-                NativeFileMethods.scan(pointer.getAsLong(), options.columns(), predicateProto, rowIndices));
+                NativeFileMethods.scan(pointer.getAsLong(), options.columns(), predicateProto, rowRange, rowIndices));
     }
 
     @Override

java/vortex-jni/src/main/java/dev/vortex/jni/NativeFileMethods.java

@@ -49,5 +49,10 @@ private NativeFileMethods() {}
      */
     public static native void close(long pointer);
 
-    public static native long scan(long pointer, List<String> columns, byte[] predicateProto, long[] rowIndices);
+    /**
+     * Build a new native scan operator that will materialize Arrays from the file, pushing down the optional
+     * predicate, row range or row indices to perform data skipping.
+     */
+    public static native long scan(
+            long pointer, List<String> columns, byte[] predicateProto, long[] rowRange, long[] rowIndices);
 }

java/vortex-spark/src/main/java/dev/vortex/spark/ArrowUtils.java

@@ -20,16 +20,14 @@
 import dev.vortex.relocated.org.apache.arrow.vector.types.TimeUnit;
 import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType;
 import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.Field;
+import java.util.stream.Collectors;
 import org.apache.spark.sql.types.DataType;
 import org.apache.spark.sql.types.DataTypes;
 import org.apache.spark.sql.types.Metadata;
 import org.apache.spark.sql.types.StructField;
 
-import java.util.stream.Collectors;
-
 public final class ArrowUtils {
-    private ArrowUtils() {
-    }
+    private ArrowUtils() {}
 
     public static DataType fromArrowField(Field field) {
         switch (field.getType().getTypeID()) {

vortex-jni/src/file.rs

@@ -121,6 +121,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFileMethods_scan(
     pointer: jlong,
     project_cols: JObject,
     predicate: JByteArray,
+    row_range: JLongArray,
     row_indices: JLongArray,
 ) -> jlong {
     // Return a new pointer to some native memory for the scan.
@@ -166,6 +167,15 @@ pub extern "system" fn Java_dev_vortex_jni_NativeFileMethods_scan(
             scan_builder = scan_builder.with_row_indices(indices_buffer);
         }
 
+        if !row_range.is_null() {
+            let indices = unsafe { env.get_array_elements(&row_range, ReleaseMode::NoCopyBack) }?;
+            let start_idx =
+                u64::try_from(indices[0]).map_err(|_| vortex_err!("i64 row_index overflow"))?;
+            let end_idx =
+                u64::try_from(indices[1]).map_err(|_| vortex_err!("i64 row_index overflow"))?;
+            scan_builder = scan_builder.with_row_range(start_idx..end_idx);
+        }
+
         Ok(NativeArrayIterator::new(Box::new(scan_builder.into_array_iter()?)).into_raw())
     })
 }

vortex-layout/src/scan/mod.rs

@@ -76,11 +76,6 @@ impl<A: 'static + Send + Sync> ScanBuilder<A> {
         self
     }
 
-    pub fn with_some_row_range(mut self, row_range: Option<Range<u64>>) -> Self {
-        self.row_range = row_range;
-        self
-    }
-
     pub fn with_selection(mut self, selection: Selection) -> Self {
         self.selection = selection;
         self