diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 67acc7170aa..9d93afc95ba 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -22,7 +22,6 @@ use arrow::ffi_stream::ArrowArrayStreamReader; use arrow::ffi_stream::FFI_ArrowArrayStream; use arrow::ipc::writer::StreamWriter; use arrow::record_batch::RecordBatchIterator; -use arrow_schema::DataType; use arrow_schema::Schema as ArrowSchema; use chrono::{DateTime, Utc}; use jni::objects::{JMap, JString, JValue}; @@ -62,7 +61,6 @@ use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; use std::collections::HashMap; use std::future::IntoFuture; use std::iter::empty; -use std::str::FromStr; use std::sync::Arc; use std::time::{Duration, UNIX_EPOCH}; use uuid::Uuid; @@ -2314,17 +2312,23 @@ pub extern "system" fn Java_org_lance_Dataset_nativeAlterColumns( mut env: JNIEnv, java_dataset: JObject, column_alterations_obj: JObject, // List + cast_schema_addr: jlong, ) { ok_or_throw_without_return!( env, - inner_alter_columns(&mut env, java_dataset, column_alterations_obj) + inner_alter_columns( + &mut env, + java_dataset, + column_alterations_obj, + cast_schema_addr + ) ) } fn create_column_alteration( env: &mut JNIEnv, column_alteration_jobj: JObject, // ColumnAlteration -) -> Result { +) -> Result<(ColumnAlteration, bool)> { let path_obj = env .get_field(&column_alteration_jobj, "path", "Ljava/lang/String;")? .l()?; @@ -2363,48 +2367,55 @@ fn create_column_alteration( None }; + // The cast target type (if any) is not read here: it is transferred separately through the + // Arrow C Data Interface (see inner_alter_columns), because ArrowType#toString() does not + // round-trip through DataType::from_str for parameterized types. This flag records whether a + // cast was requested so the caller can attach the imported type in order. let data_type_obj = env .get_field(&column_alteration_jobj, "dataType", "Ljava/util/Optional;")? .l()?; - let data_type = if env + let wants_cast = env .call_method(&data_type_obj, "isPresent", "()Z", &[])? - .z()? - { - let j_data_type: JObject = env - .call_method(data_type_obj, "get", "()Ljava/lang/Object;", &[])? - .l()?; - let jstring: JString = env - .call_method(j_data_type, "toString", "()Ljava/lang/String;", &[])? - .l()? - .into(); - let data_type_str: String = env.get_string(&jstring)?.into(); // Intermediate variable - DataType::from_str(&data_type_str) - .map_err(|e| Error::input_error(e.to_string())) - .ok() - } else { - None - }; + .z()?; - Ok(ColumnAlteration { + let alteration = ColumnAlteration { path, rename, nullable, - data_type, - }) + data_type: None, + }; + Ok((alteration, wants_cast)) } fn inner_alter_columns( env: &mut JNIEnv, java_dataset: JObject, column_alterations_obj: JObject, // List + cast_schema_addr: jlong, ) -> Result<()> { let list = env.get_list(&column_alterations_obj)?; let mut iter = list.iter(env)?; let mut column_alterations = Vec::new(); + let mut cast_flags = Vec::new(); while let Some(elem) = iter.next(env)? { - let alteration = create_column_alteration(env, elem)?; + let (alteration, wants_cast) = create_column_alteration(env, elem)?; column_alterations.push(alteration); + cast_flags.push(wants_cast); + } + + // Cast target types arrive as one Arrow schema field per requested cast, in the same order + // as the alterations that requested one. + let cast_schema = unsafe { FFI_ArrowSchema::from_raw(cast_schema_addr as *mut _) }; + let cast_schema = ArrowSchema::try_from(&cast_schema) + .map_err(|_| Error::input_error("ArrowSchema conversion error".to_string()))?; + let mut cast_types = cast_schema.fields.iter().map(|f| f.data_type().clone()); + for (alteration, wants_cast) in column_alterations.iter_mut().zip(cast_flags) { + if wants_cast { + alteration.data_type = Some(cast_types.next().ok_or_else(|| { + Error::input_error("Missing cast type for column alteration".to_string()) + })?); + } } let mut dataset_guard = diff --git a/java/src/main/java/org/lance/Dataset.java b/java/src/main/java/org/lance/Dataset.java index 6ee3e13488d..02b7c81ce39 100644 --- a/java/src/main/java/org/lance/Dataset.java +++ b/java/src/main/java/org/lance/Dataset.java @@ -55,6 +55,7 @@ import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.ipc.ArrowStreamReader; 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 java.io.ByteArrayInputStream; @@ -63,6 +64,7 @@ import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.nio.channels.ReadableByteChannel; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -733,11 +735,31 @@ public void dropColumns(List columns) { public void alterColumns(List columnAlterations) { try (LockManager.WriteLock writeLock = lockManager.acquireWriteLock()) { Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed"); - nativeAlterColumns(columnAlterations); + // Cast target types are carried across the FFI boundary through the Arrow C Data + // Interface rather than ArrowType#toString(), which does not round-trip reliably on + // the native side (parameterized types such as Int(64, true) fail to parse and the + // cast would otherwise be silently dropped). One field is exported per alteration that + // requests a type change, in the same order as {@code columnAlterations}. + List castFields = new ArrayList<>(); + int castIndex = 0; + for (ColumnAlteration alteration : columnAlterations) { + if (alteration.getDataType().isPresent()) { + castFields.add(new Field("f" + castIndex++, castFieldType(alteration), null)); + } + } + try (ArrowSchema castSchema = ArrowSchema.allocateNew(allocator)) { + Data.exportSchema(allocator, new Schema(castFields), null, castSchema); + nativeAlterColumns(columnAlterations, castSchema.memoryAddress()); + } } } - private native void nativeAlterColumns(List columnAlterations); + private static FieldType castFieldType(ColumnAlteration alteration) { + boolean nullable = alteration.getNullable().orElse(true); + return new FieldType(nullable, alteration.getDataType().get(), null); + } + + private native void nativeAlterColumns(List columnAlterations, long castAddr); /** * Create a new Dataset Scanner. diff --git a/java/src/test/java/org/lance/DatasetTest.java b/java/src/test/java/org/lance/DatasetTest.java index c24f1722ad1..25610a5c96c 100644 --- a/java/src/test/java/org/lance/DatasetTest.java +++ b/java/src/test/java/org/lance/DatasetTest.java @@ -660,6 +660,40 @@ void testAlterColumns(@TempDir Path tempDir) { } } + @Test + void testAlterColumnsCastType(@TempDir Path tempDir) { + String testMethodName = new Object() {}.getClass().getEnclosingMethod().getName(); + String datasetPath = tempDir.resolve(testMethodName).toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + dataset = testDataset.createEmptyDataset(); + + // Widen "id" from Int32 to Int64. The cast target type is a parameterized ArrowType, which + // must survive the trip to the native side; regression test for a dropped cast that left the + // stored type unchanged. + ColumnAlteration widenId = + new ColumnAlteration.Builder("id").castTo(new ArrowType.Int(64, true)).build(); + dataset.alterColumns(Collections.singletonList(widenId)); + + assertEquals(new ArrowType.Int(64, true), dataset.getSchema().findField("id").getType()); + + // A cast combined with rename must apply both. + ColumnAlteration renameAndWiden = + new ColumnAlteration.Builder("id") + .rename("id_long") + .castTo(new ArrowType.Int(64, true)) + .build(); + dataset.alterColumns(Collections.singletonList(renameAndWiden)); + + List fieldNames = + dataset.getSchema().getFields().stream().map(Field::getName).collect(Collectors.toList()); + assertFalse(fieldNames.contains("id")); + assertTrue(fieldNames.contains("id_long")); + assertEquals(new ArrowType.Int(64, true), dataset.getSchema().findField("id_long").getType()); + } + } + @Test void testAddColumnBySqlExpressions(@TempDir Path tempDir) { String testMethodName = new Object() {}.getClass().getEnclosingMethod().getName();