vortex-jni includes vortex-geo, jni arrow import/export uses session methods

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

@@ -120,7 +120,7 @@ public static DataSource open(Session session, List<NativeReadable> readables, i
     /** Arrow schema of the data source (and of scans produced from it). */
     public Schema arrowSchema(BufferAllocator allocator) {
         try (ArrowSchema schema = ArrowSchema.allocateNew(allocator)) {
-            NativeDataSource.arrowSchema(pointer, schema.memoryAddress());
+            NativeDataSource.arrowSchema(session.nativePointer(), pointer, schema.memoryAddress());
             return Data.importSchema(allocator, schema, null);
         }
     }

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

@@ -42,7 +42,7 @@ static Scan fromPointer(Session session, long pointer) {
      */
     public Schema arrowSchema(BufferAllocator allocator) {
         try (ArrowSchema schema = ArrowSchema.allocateNew(allocator)) {
-            NativeScan.arrowSchema(pointer, schema.memoryAddress());
+            NativeScan.arrowSchema(session.nativePointer(), pointer, schema.memoryAddress());
             return Data.importSchema(allocator, schema, null);
         }
     }

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

@@ -40,8 +40,13 @@ public static native long openFiles(
     /** Free a data source pointer. */
     public static native void free(long pointer);
 
-    /** Export the data source's schema into the Arrow C Data Interface struct at {@code schemaAddress}. */
-    public static native void arrowSchema(long pointer, long schemaAddress);
+    /**
+     * Export the data source's schema into the Arrow C Data Interface struct at {@code schemaAddress}. Extension dtypes
+     * are dispatched through the session's registered Arrow export plugins.
+     *
+     * @param sessionPointer pointer from {@link NativeSession#newSession()}
+     */
+    public static native void arrowSchema(long sessionPointer, long pointer, long schemaAddress);
 
     /**
      * Populate {@code out} with {@code [rows, cardinality]}. Cardinality is one of {@code 0=unknown},

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

@@ -43,8 +43,13 @@ public static native long create(
     /** Free a scan pointer. */
     public static native void free(long pointer);
 
-    /** Export the scan's schema into the Arrow C Data Interface struct at {@code schemaAddress}. */
-    public static native void arrowSchema(long pointer, long schemaAddress);
+    /**
+     * Export the scan's schema into the Arrow C Data Interface struct at {@code schemaAddress}. Extension dtypes are
+     * dispatched through the session's registered Arrow export plugins.
+     *
+     * @param sessionPointer pointer from {@link NativeSession#newSession()}
+     */
+    public static native void arrowSchema(long sessionPointer, long pointer, long schemaAddress);
 
     /** Fill {@code out} with {@code [count, cardinality]}. */
     public static native void partitionCount(long pointer, long[] out);

java/vortex-jni/src/test/java/dev/vortex/api/GeoTypesTest.java

@@ -0,0 +1,164 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+package dev.vortex.api;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import dev.vortex.arrow.ArrowAllocation;
+import dev.vortex.jni.NativeLoader;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.arrow.c.ArrowArray;
+import org.apache.arrow.c.ArrowSchema;
+import org.apache.arrow.c.Data;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.ipc.ArrowReader;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+/**
+ * Round-trips a Vortex geo extension column ({@code vortex.geo.wkb}) through the JNI boundary. Geo columns cross the
+ * boundary as Arrow fields tagged with the GeoArrow extension name ({@code geoarrow.wkb}) and JSON metadata carrying
+ * the CRS.
+ */
+public final class GeoTypesTest {
+    private static final String EXTENSION_NAME_KEY = "ARROW:extension:name";
+    private static final String EXTENSION_METADATA_KEY = "ARROW:extension:metadata";
+    private static final String GEOARROW_WKB = "geoarrow.wkb";
+    private static final String CRS_METADATA = "{\"crs\":\"OGC:CRS84\"}";
+
+    @TempDir
+    static Path tempDir;
+
+    static String writePath;
+
+    private static final List<byte[]> WKB_POINTS = new ArrayList<>();
+
+    @BeforeAll
+    public static void loadLibrary() {
+        NativeLoader.loadJni();
+    }
+
+    @BeforeAll
+    static void setup() throws IOException {
+        writePath = tempDir.resolve("geo.vortex").toAbsolutePath().toUri().toString();
+
+        WKB_POINTS.add(wkbPoint(1.0, 2.0));
+        WKB_POINTS.add(wkbPoint(-111.7610, 34.8697));
+        WKB_POINTS.add(null);
+        WKB_POINTS.add(wkbPoint(0.0, 0.0));
+
+        BufferAllocator allocator = ArrowAllocation.rootAllocator();
+        Schema schema = new Schema(List.of(wkbField("geom")));
+        Session session = Session.create();
+        try (VortexWriter writer = VortexWriter.create(session, writePath, schema, new HashMap<>(), allocator);
+                VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+            VarBinaryVector geomVec = (VarBinaryVector) root.getVector("geom");
+            geomVec.allocateNew(WKB_POINTS.size());
+            for (int i = 0; i < WKB_POINTS.size(); i++) {
+                byte[] wkb = WKB_POINTS.get(i);
+                if (wkb == null) {
+                    geomVec.setNull(i);
+                } else {
+                    geomVec.setSafe(i, wkb);
+                }
+            }
+            root.setRowCount(WKB_POINTS.size());
+
+            try (ArrowArray arrowArray = ArrowArray.allocateNew(allocator);
+                    ArrowSchema arrowSchemaFfi = ArrowSchema.allocateNew(allocator)) {
+                Data.exportVectorSchemaRoot(allocator, root, null, arrowArray, arrowSchemaFfi);
+                writer.writeBatch(arrowArray.memoryAddress(), arrowSchemaFfi.memoryAddress());
+            }
+        }
+    }
+
+    @Test
+    public void testSchemaCarriesGeoArrowExtension() {
+        BufferAllocator allocator = ArrowAllocation.rootAllocator();
+        Session session = Session.create();
+        DataSource ds = DataSource.open(session, writePath);
+
+        Schema schema = ds.arrowSchema(allocator);
+        Field geom = schema.findField("geom");
+        assertEquals(GEOARROW_WKB, geom.getMetadata().get(EXTENSION_NAME_KEY));
+        assertEquals(CRS_METADATA, geom.getMetadata().get(EXTENSION_METADATA_KEY));
+        assertTrue(geom.isNullable());
+    }
+
+    @Test
+    public void testScanSchemaCarriesGeoArrowExtension() {
+        BufferAllocator allocator = ArrowAllocation.rootAllocator();
+        Session session = Session.create();
+        DataSource ds = DataSource.open(session, writePath);
+
+        Schema schema = ds.scan(ScanOptions.of()).arrowSchema(allocator);
+        Field geom = schema.findField("geom");
+        assertEquals(GEOARROW_WKB, geom.getMetadata().get(EXTENSION_NAME_KEY));
+        assertEquals(CRS_METADATA, geom.getMetadata().get(EXTENSION_METADATA_KEY));
+    }
+
+    @Test
+    public void testWkbValuesRoundTrip() throws Exception {
+        BufferAllocator allocator = ArrowAllocation.rootAllocator();
+        Session session = Session.create();
+        DataSource ds = DataSource.open(session, writePath);
+
+        List<byte[]> values = new ArrayList<>();
+        Scan scan = ds.scan(ScanOptions.of());
+        while (scan.hasNext()) {
+            Partition partition = scan.next();
+            try (ArrowReader arrowReader = partition.scanArrow(allocator)) {
+                while (arrowReader.loadNextBatch()) {
+                    VectorSchemaRoot root = arrowReader.getVectorSchemaRoot();
+                    var geomVec = root.getVector("geom");
+                    for (int i = 0; i < root.getRowCount(); i++) {
+                        values.add(geomVec.isNull(i) ? null : (byte[]) geomVec.getObject(i));
+                    }
+                }
+            }
+        }
+
+        assertEquals(WKB_POINTS.size(), values.size());
+        for (int i = 0; i < WKB_POINTS.size(); i++) {
+            byte[] expected = WKB_POINTS.get(i);
+            if (expected == null) {
+                assertNull(values.get(i));
+            } else {
+                assertArrayEquals(expected, values.get(i));
+            }
+        }
+    }
+
+    private static Field wkbField(String name) {
+        Map<String, String> metadata = Map.of(EXTENSION_NAME_KEY, GEOARROW_WKB, EXTENSION_METADATA_KEY, CRS_METADATA);
+        return new Field(name, new FieldType(true, ArrowType.Binary.INSTANCE, null, metadata), null);
+    }
+
+    /** Little-endian WKB encoding of {@code POINT(x y)}. */
+    private static byte[] wkbPoint(double x, double y) {
+        ByteBuffer buffer = ByteBuffer.allocate(21).order(ByteOrder.LITTLE_ENDIAN);
+        buffer.put((byte) 1); // little-endian marker
+        buffer.putInt(1); // geometry type: point
+        buffer.putDouble(x);
+        buffer.putDouble(y);
+        return buffer.array();
+    }
+}

java/vortex-jni/src/test/java/dev/vortex/jni/JNIWriterTest.java

@@ -31,6 +31,7 @@
 import org.apache.arrow.vector.VarBinaryVector;
 import org.apache.arrow.vector.VarCharVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.ViewVarBinaryVector;
 import org.apache.arrow.vector.ViewVarCharVector;
 import org.apache.arrow.vector.complex.StructVector;
 import org.apache.arrow.vector.ipc.ArrowReader;
@@ -261,8 +262,9 @@ public void testParquetVariantRoundTrip() throws IOException {
                 assertTrue(reader.loadNextBatch());
                 VectorSchemaRoot resultRoot = reader.getVectorSchemaRoot();
                 StructVector variant = (StructVector) resultRoot.getVector("variant");
-                VarBinaryVector metadata = variant.getChild("metadata", VarBinaryVector.class);
-                VarBinaryVector value = variant.getChild("value", VarBinaryVector.class);
+                // Binary columns cross the boundary as their native view types.
+                ViewVarBinaryVector metadata = variant.getChild("metadata", ViewVarBinaryVector.class);
+                ViewVarBinaryVector value = variant.getChild("value", ViewVarBinaryVector.class);
 
                 assertArrayEquals(VARIANT_METADATA, metadata.get(0));
                 assertArrayEquals(VARIANT_INT8_42, value.get(0));

vortex-jni/Cargo.toml

@@ -33,6 +33,7 @@ tracing-subscriber = { workspace = true, features = ["env-filter"] }
 url = { workspace = true }
 vortex = { workspace = true, features = ["object_store", "files"] }
 vortex-arrow = { workspace = true }
+vortex-geo = { workspace = true }
 vortex-parquet-variant = { workspace = true }
 
 [dev-dependencies]

vortex-jni/src/data_source.rs

@@ -8,8 +8,10 @@
 //! and bare file paths are both accepted. Filesystems are cached per base URL so repeated
 //! globs against the same bucket share a single client.
 
+use std::ptr;
 use std::sync::Arc;
 
+use arrow_array::ffi::FFI_ArrowSchema;
 use jni::EnvUnowned;
 use jni::objects::JClass;
 use jni::objects::JLongArray;
@@ -29,9 +31,9 @@ use vortex::io::runtime::BlockingRuntime;
 use vortex::io::session::RuntimeSessionExt;
 use vortex::scan::DataSourceRef;
 use vortex::utils::aliases::hash_map::HashMap;
+use vortex_arrow::ArrowSessionExt;
 
 use crate::RUNTIME;
-use crate::dtype::export_dtype_to_arrow;
 use crate::errors::try_or_throw;
 use crate::file::extract_properties;
 use crate::io::JavaFileSystem;
@@ -208,20 +210,25 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_free(
 }
 
 /// Export the data source's schema into the Arrow C Data Interface schema struct at
-/// `schema_addr`.
 #[unsafe(no_mangle)]
 pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_arrowSchema(
     mut env: EnvUnowned,
     _class: JClass,
+    session_ptr: jlong,
     pointer: jlong,
     schema_addr: jlong,
 ) {
     try_or_throw(&mut env, |_| {
         if schema_addr == 0 {
             throw_runtime!("null arrow schema address");
         }
+        let session = unsafe { session_ref(session_ptr) };
         let ds = unsafe { NativeDataSource::from_ptr(pointer) };
-        export_dtype_to_arrow(ds.inner.dtype(), schema_addr)?;
+        let arrow_schema = session.arrow().to_arrow_schema(ds.inner.dtype())?;
+        let ffi_schema = FFI_ArrowSchema::try_from(&arrow_schema)?;
+        unsafe {
+            ptr::write(schema_addr as *mut FFI_ArrowSchema, ffi_schema);
+        }
         Ok(())
     });
 }

vortex-jni/src/lib.rs

@@ -19,7 +19,6 @@ macro_rules! throw_runtime {
 }
 
 mod data_source;
-mod dtype;
 mod errors;
 mod expression;
 mod file;