support writing schema metadata through java LanceFileWriter API.

Author: 喆宇

java/lance-jni/Cargo.lock

@@ -149,6 +149,37 @@ pub extern "system" fn Java_com_lancedb_lance_file_LanceFileWriter_closeNative<'
     JObject::null()
 }
 
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_file_LanceFileWriter_nativeWriteSchemaMetadata<
+    'local,
+>(
+    mut env: JNIEnv<'local>,
+    writer: JObject,
+    schema_metadata: JObject, // Map<String, String>
+) -> JObject<'local> {
+    if let Err(e) = inner_write_schema_metadata(&mut env, writer, schema_metadata) {
+        e.throw(&mut env);
+        return JObject::null();
+    }
+    JObject::null()
+}
+
+fn inner_write_schema_metadata(
+    env: &mut JNIEnv<'_>,
+    writer: JObject,
+    schema_metadata: JObject, // Map<String, String>
+) -> Result<()> {
+    let metadata_map = JMap::from_env(env, &schema_metadata)?;
+    let metadata = to_rust_map(env, &metadata_map)?;
+    let writer_guard =
+        unsafe { env.get_rust_field::<_, _, BlockingFileWriter>(writer, NATIVE_WRITER) }?;
+    let mut writer = writer_guard.inner.lock().unwrap();
+    metadata.into_iter().for_each(|(k, v)| {
+        writer.add_schema_metadata(k, v);
+    });
+    Ok(())
+}
+
 #[no_mangle]
 pub extern "system" fn Java_com_lancedb_lance_file_LanceFileWriter_writeNative<'local>(
     mut env: JNIEnv<'local>,

java/lance-jni/src/file_writer.rs

@@ -120,6 +120,19 @@ public void write(VectorSchemaRoot batch) throws IOException {
     }
   }
 
+  /**
+   * Write a schema metadata map to underlying file. User can retrieve those values from {@link
+   * LanceFileReader#schema() reader schema}.
+   *
+   * @param metadata metadata
+   * @throws IOException IOException
+   */
+  public void writeSchemaMetadata(Map<String, String> metadata) throws IOException {
+    nativeWriteSchemaMetadata(metadata);
+  }
+
+  private native void nativeWriteSchemaMetadata(Map<String, String> metadata) throws IOException;
+
   /**
    * Close the LanceFileWriter
    *

java/src/main/java/com/lancedb/lance/file/LanceFileWriter.java

@@ -27,6 +27,7 @@
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.arrow.vector.util.Text;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
@@ -258,4 +259,47 @@ void testInvalidPath() {
     } catch (IOException e) {
     }
   }
+
+  @Test
+  void testWriteSchemaMetadata(@TempDir Path tempDir) throws Exception {
+    String filePath = tempDir.resolve("write_schema_metadata.lance").toString();
+    BufferAllocator allocator = new RootAllocator();
+    try (LanceFileWriter writer = LanceFileWriter.open(filePath, allocator, null)) {
+      try (VectorSchemaRoot batch = createBatch(allocator)) {
+        writer.write(batch);
+        writer.writeSchemaMetadata(Collections.singletonMap("testKey", "testValue"));
+        writer.write(batch);
+        // repeatedly write
+        writer.writeSchemaMetadata(Collections.singletonMap("testKey1", "testValue1"));
+      }
+    }
+
+    try (LanceFileReader reader = LanceFileReader.open(filePath, allocator)) {
+      Schema fileSchema = reader.schema();
+      Map<String, String> metadata = fileSchema.getCustomMetadata();
+
+      Assertions.assertTrue(metadata.containsKey("testKey"));
+      Assertions.assertEquals("testValue", metadata.get("testKey"));
+
+      Assertions.assertTrue(metadata.containsKey("testKey1"));
+      Assertions.assertEquals("testValue1", metadata.get("testKey1"));
+    }
+  }
+
+  @Test
+  void testWriteNullSchemaMetadata(@TempDir Path tempDir) throws Exception {
+    String filePath = tempDir.resolve("write_null_schema_metadata.lance").toString();
+    BufferAllocator allocator = new RootAllocator();
+    try (LanceFileWriter writer = LanceFileWriter.open(filePath, allocator, null)) {
+      try (VectorSchemaRoot batch = createBatch(allocator)) {
+        writer.write(batch);
+        Assertions.assertThrows(
+            Exception.class,
+            () -> writer.writeSchemaMetadata(Collections.singletonMap("someKey", null)));
+        Assertions.assertThrows(
+            Exception.class,
+            () -> writer.writeSchemaMetadata(Collections.singletonMap(null, "someValue")));
+      }
+    }
+  }
 }