fix: handling of struct fields in Spark (#5453)

Fixes some issues I encountered trying to get some files with nested
structs/lists to load.

I'm still hitting a very strange issue deep in C Arrow Data interface
with the 3.5GB github archives dataset, but this fixes some other
issues.

---------

Signed-off-by: Andrew Duffy <[email protected]>
Signed-off-by: Robert Kruszewski <[email protected]>
Co-authored-by: Robert Kruszewski <[email protected]>

Author: a10y

java/settings.gradle.kts

@@ -20,4 +20,3 @@ rootProject.name = "vortex-root"
 // API bindings
 include("vortex-jni")
 include("vortex-spark")
-

java/vortex-spark/build.gradle.kts

@@ -91,7 +91,6 @@ tasks.withType<ShadowJar> {
     }
 }
 
-
 tasks.withType<Test>().all {
     classpath +=
         project(":vortex-jni")

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

@@ -3,7 +3,6 @@
 
 package dev.vortex.spark;
 
-import com.google.common.collect.Streams;
 import dev.vortex.api.DType;
 import java.util.Optional;
 import org.apache.spark.sql.connector.catalog.Column;
@@ -107,13 +106,20 @@ public static DataType toDataType(DType dType) {
                 return DataTypes.BinaryType;
             case STRUCT:
                 // For each of the inner struct fields, we capture them together here.
-                var struct = new StructType();
+                var fieldNames = dType.getFieldNames();
+                var fieldTypes = dType.getFieldTypes();
 
-                Streams.forEachPair(
-                        dType.getFieldNames().stream(),
-                        dType.getFieldTypes().stream(),
-                        (name, type) -> struct.add(name, toDataType(type)));
-                return struct;
+                // NOTE: it's very important we do this with a for loop. Using the streams API can easily
+                //  lead to StackOverflowError being thrown.
+                var fields = new StructField[fieldNames.size()];
+                for (int i = 0; i < fieldNames.size(); i++) {
+                    var name = fieldNames.get(i);
+                    try (var type = fieldTypes.get(i)) {
+                        fields[i] = new StructField(name, toDataType(type), dType.isNullable(), Metadata.empty());
+                    }
+                }
+
+                return DataTypes.createStructType(fields);
             case LIST:
                 return DataTypes.createArrayType(toDataType(dType.getElementType()), dType.isNullable());
             case EXTENSION:
@@ -151,10 +157,17 @@ public static DataType toDataType(DType dType) {
      * Convert a STRUCT Vortex type to a Spark {@link Column}.
      */
     public static Column[] toColumns(DType dType) {
-        return Streams.zip(dType.getFieldNames().stream(), dType.getFieldTypes().stream(), (name, fieldType) -> {
-                    var dataType = toDataType(fieldType);
-                    return Column.create(name, dataType, fieldType.isNullable());
-                })
-                .toArray(Column[]::new);
+        var fieldNames = dType.getFieldNames();
+        var fieldTypes = dType.getFieldTypes();
+        var columns = new Column[fieldNames.size()];
+
+        for (int i = 0; i < columns.length; i++) {
+            var name = fieldNames.get(i);
+            try (var type = fieldTypes.get(i)) {
+                columns[i] = Column.create(name, toDataType(type), type.isNullable());
+            }
+        }
+
+        return columns;
     }
 }

java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceBasicTest.java

@@ -5,6 +5,7 @@
 
 import static org.junit.jupiter.api.Assertions.*;
 
+import dev.vortex.relocated.org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.spark.sql.types.DataTypes;
 import org.apache.spark.sql.types.StructField;
 import org.apache.spark.sql.types.StructType;
@@ -48,6 +49,48 @@ public void testSparkToArrowSchemaConversion() {
         assertEquals("active", arrowSchema.getFields().get(3).getName());
     }
 
+    @Test
+    @DisplayName("SparkToArrowSchema should convert nested types")
+    public void testNestedSparkToArrowSchemaConversion() {
+        // Create a more complex spark schema
+        StructType sparkSchema = DataTypes.createStructType(new StructField[] {
+            DataTypes.createStructField(
+                    "inner",
+                    DataTypes.createStructType(new StructField[] {
+                        DataTypes.createStructField("id", DataTypes.IntegerType, false),
+                        DataTypes.createStructField("name", DataTypes.StringType, true),
+                        DataTypes.createStructField("value", DataTypes.DoubleType, false),
+                        DataTypes.createStructField("active", DataTypes.BooleanType, true)
+                    }),
+                    false)
+        });
+
+        // Convert to Arrow schema
+        var arrowSchema = dev.vortex.spark.write.SparkToArrowSchema.convert(sparkSchema);
+
+        // Verify conversion
+        assertNotNull(arrowSchema, "Arrow schema should not be null");
+        assertEquals(1, arrowSchema.getFields().size(), "Arrow schema should have same number of fields");
+
+        // Should contain the right inner fields
+        var nestedFields = arrowSchema.getFields().get(0).getChildren();
+
+        // Verify field types are preserved
+        assertInstanceOf(ArrowType.Struct.class, arrowSchema.getFields().get(0).getType());
+
+        assertEquals("id", nestedFields.get(0).getName());
+        assertInstanceOf(ArrowType.Int.class, nestedFields.get(0).getType());
+
+        assertEquals("name", nestedFields.get(1).getName());
+        assertInstanceOf(ArrowType.Utf8.class, nestedFields.get(1).getType());
+
+        assertEquals("value", nestedFields.get(2).getName());
+        assertInstanceOf(ArrowType.FloatingPoint.class, nestedFields.get(2).getType());
+
+        assertEquals("active", nestedFields.get(3).getName());
+        assertInstanceOf(ArrowType.Bool.class, nestedFields.get(3).getType());
+    }
+
     @Test
     @DisplayName("VortexWriterCommitMessage should store metadata correctly")
     public void testWriterCommitMessage() {

vortex-dtype/src/field_names.rs

@@ -322,15 +322,27 @@ impl From<&[&'static str]> for FieldNames {
     }
 }
 
+impl<const N: usize> From<[&str; N]> for FieldNames {
+    fn from(value: [&str; N]) -> Self {
+        Self(value.iter().cloned().map(FieldName::from).collect())
+    }
+}
+
+impl From<Vec<&str>> for FieldNames {
+    fn from(value: Vec<&str>) -> Self {
+        Self(value.into_iter().map(FieldName::from).collect())
+    }
+}
+
 impl From<&[FieldName]> for FieldNames {
     fn from(value: &[FieldName]) -> Self {
         Self(Arc::from(value))
     }
 }
 
-impl<const N: usize> From<[&'static str; N]> for FieldNames {
-    fn from(value: [&'static str; N]) -> Self {
-        Self(value.into_iter().map(FieldName::from).collect())
+impl<const N: usize> From<&[&str; N]> for FieldNames {
+    fn from(value: &[&str; N]) -> Self {
+        Self(value.iter().cloned().map(FieldName::from).collect())
     }
 }
 

vortex-jni/src/array.rs

@@ -120,7 +120,6 @@ fn data_type_no_views(data_type: DataType) -> DataType {
             DataType::LargeList(FieldRef::new(new_inner))
         }
         DataType::Struct(fields) => {
-            // Things
             let viewless_fields: Vec<FieldRef> = fields
                 .iter()
                 .map(|field_ref| {

vortex-jni/src/dtype.rs

@@ -128,7 +128,9 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDTypeMethods_getFieldTypes(
     let dtype = unsafe { &*(dtype_ptr as *const DType) };
 
     try_or_throw(&mut env, |env| {
-        let array_list = env.new_object("java/util/ArrayList", "()V", &[])?;
+        let array_list = env
+            .new_object("java/util/ArrayList", "()V", &[])
+            .map_err(|e| JNIError::Vortex(vortex_err!("failure constructing ArrayList: {e}")))?;
         let field_types = env.get_list(&array_list)?;
         let Some(struct_dtype) = dtype.as_struct_fields_opt() else {
             throw_runtime!("DType should be STRUCT, was {dtype}");

vortex-jni/src/errors.rs

@@ -76,7 +76,7 @@ impl JNIDefault for jobject {
     }
 }
 
-/// Run the provided function inside of the JNIEnv context. Throws an exception if the function returns an error.
+/// Run the provided function inside the JNIEnv context. Throws an exception if the function returns an error.
 #[allow(clippy::expect_used)]
 #[inline]
 pub fn try_or_throw<'a, F, T>(env: &mut JNIEnv<'a>, function: F) -> T
@@ -87,12 +87,21 @@ where
     match function(env) {
         Ok(result) => result,
         Err(error) => {
+            // Propagate the exception instead of throwing our own.
+            if env
+                .exception_check()
+                .expect("checking exception should succeed")
+            {
+                return T::jni_default();
+            }
+
             let msg = error.to_string();
-            env.throw((RUNTIME_EXC_CLASS, msg))
-                .expect("throwing exception back to Java failed, everything is bad");
+            match env.throw(msg) {
+                Ok(()) => {}
+                Err(err) => log::warn!("Failed throwing exception back up to Java: {err}"),
+            }
+
             T::jni_default()
         }
     }
 }
-
-pub static RUNTIME_EXC_CLASS: &str = "java/lang/RuntimeException";