feat: implement java bindings #1928

beinan · 2024-02-08T01:23:19Z

This pull request is still under development and addresses the interconnection between Java and Rust using JNI for both write and read paths. Here's a summary of the progress:

Write Path:

Code currently implements sending data from Java to Rust through JNI.
arrow_c in java and arrow ffi in rust are used for efficient data serialization and transfer.
Resource management ensures proper release on both sides, preventing memory leaks.

Read Path:

Development for the read path (data flow from Rust to Java) is ongoing within this PR.
The same approach with Arrow and JNI will be utilized for consistency and performance.

Next Steps:

Continuous development and testing on both write and read paths.
Open to feedback and suggestions on design and implementation for all aspects.

Highlights:

This PR establishes a single source for building a complete Java-Rust communication channel using JNI, which will have a better compatibility for integration with Java projects still running on java 8 such as Presto, Spark and hive/hadoop .
Arrow integration ensures optimized data handling and minimizes data copies.
Careful resource management prevents memory leaks and improves overall stability.

eddyxu · 2024-02-10T00:56:02Z

java/lance-jni/src/blocking_dataset.rs

+}
+
+impl BlockingDataset {
+    pub fn write(reader: ArrowArrayStreamReader, uri: &str) -> Result<Self> {


Is async calls not encouraged cross JNI?

I'm just thinking async call in jni might introduce complexities. Also I prefer to rely on the thread management in java compute engine (I'm still targeting to implement a lance connector for Trino/Presto). But anyway, async api will be more efficient in many places, I think we could implement the async jni later, can we?

hey @beinan Any plans to add a Trino connector for this? If so, any sense of ETA? Thanks for building this :)

eddyxu · 2024-02-10T00:57:57Z

java/lance-jni/src/blocking_dataset.rs

@@ -0,0 +1,36 @@
+use arrow::ffi_stream::ArrowArrayStreamReader;


Lets add apache license to headers :)

cargo clippy will pick this rule up. Lets make github action to run on this PR?

Good call, I will add the github actions for both java and rust-jni projects

eddyxu · 2024-02-10T00:58:28Z

java/lance-jni/src/blocking_dataset.rs

+        )?;
+        Ok(BlockingDataset { inner, rt })
+    }
+    pub fn open(uri: &str) -> Result<Self> {


it'd be great that we can pass Read / Write options from JNI.

Yeah, I'm also thinking about that, I will add the options very soon

eddyxu · 2024-02-10T01:01:31Z

java/lance-jni/src/blocking_dataset.rs

+    pub fn count_rows(&self) -> Result<usize> {
+        self.rt.block_on(self.inner.count_rows())
+    }
+    pub fn close(&self) {}


So input and output are FFI_ArrowArrayStream iiuc
?

I think the input of the write path can be either a FFI_ArrowArrayStream or FFI_ArrowArray (I have implement the stream api for write, I'm thinking to implement append api for both structure). For the output of the read path, I think it will be something very similar but I didn't get a chance to try to implement it.

beinan · 2024-02-13T08:24:04Z

Hi @eddyxu,

I integrated a mvn plugin to build rust project in maven, you can run mvn clean package to generate the jar with a native library embedded.

Then you can find jar atlance/java/target/lance-0.1-SNAPSHOT.jar

The structure of the jar is simple

jar tf target/lance-0.1-SNAPSHOT.jar
META-INF/MANIFEST.MF
META-INF/
nativelib/
nativelib/darwin-aarch64/
com/
com/lancedb/
com/lancedb/lance/
META-INF/maven/
META-INF/maven/com.lancedb/
META-INF/maven/com.lancedb/lance/
nativelib/darwin-aarch64/liblance_jni.dylib
com/lancedb/lance/Dataset.class
META-INF/maven/com.lancedb/lance/pom.xml
META-INF/maven/com.lancedb/lance/pom.properties

I only added one target platform "darwin-aarch64" for macos on arm. (I will figure out how to add more target platform later)

It will also run the test in both Rust and Java project.

By the way, I also removed quite a few redundant code.

eddyxu · 2024-02-16T23:27:15Z

.github/workflows/java.yml

+      matrix:
+        include:
+          - os: ubuntu-22.04
+            java-version: 8


I vote to not support java8 :) Even ok to not support java11 if it is up to me :)

eddyxu · 2024-02-20T00:08:16Z

.github/workflows/java.yml

+          distribution: temurin
+          java-version: 17
+          cache: 'maven'
+      - run: echo "JAVA_17=$JAVA_HOME" >> $GITHUB_ENV


Should we use a matrix to split it into two jobs?

If you look at the CI job, testing another Java version only took 6s. If we split to another job, it would have to recompile the Rust (and also do the installation steps). I think we should keep as-is so that our CI usage is kept minimal.

Thank you for the comments @wjones127 ! Yes, that's why I merge these two together as recompiling the rust part really takes a lot of time.

eddyxu · 2024-02-20T00:10:46Z

.github/workflows/java.yml

+  pull_request:
+    paths:
+      - java/**
+      - rust/**


If we do expect to run this everytime with rust change, could we change this package as a workspace package with the rest of rust code, and share the same workspace Cargo.toml

It sounds great to me!

wjones127

Hi @beinan. Thanks for working on this.

The things I would like to see:

Make the Tokio runtime a lazy static, so we aren't recreating runtimes.
Make this crate part of the workspace, if possible.
Minimize scope of your unsafe calls.

I think additional features like supporting ReadParams can be left for a future PR. More important that we get the foundation of FFI and async runtime working.

wjones127 · 2024-02-20T18:23:40Z

java/lance-jni/src/lib.rs

+    unsafe {
+        match ArrowArrayStreamReader::from_raw(arrow_array_stream_addr as *mut FFI_ArrowArrayStream)
+        {


Let's make sure when we use unsafe blocks, we only put the minimal unsafe parts there and leave a comment why the calls are safe.

let stream_ptr = arrow_array_stream_addr as *mut FFI_ArrowArrayStream; // SAFETY: the pointer is recieved directly from Java's // ArrowArrayStream.memory_address(), which guarantees to return a non-null // valid pointer. let stream = unsafe { ArrowArrayStreamReader::from_raw(stream_ptr) }; match stream { .... }

@eddyxu This is still outstanding

wjones127 · 2024-02-20T18:35:36Z

java/lance-jni/src/blocking_dataset.rs

+        let rt = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .build()?;


Could you move this to a lazy static variable? We don't want to create a new runtime every time we write to a dataset.

Note, however, in the future we should probably port the executor we did in Python #1172. This allows things like writing to a Lance dataset from a scan of another Lance dataset, and support for KeyboardInterrupt (#1438).

java/lance-jni/src/lib.rs

wjones127 · 2024-02-20T18:43:06Z

.github/workflows/java.yml

+          distribution: temurin
+          java-version: 17
+          cache: 'maven'
+      - run: echo "JAVA_17=$JAVA_HOME" >> $GITHUB_ENV


If you look at the CI job, testing another Java version only took 6s. If we split to another job, it would have to recompile the Rust (and also do the installation steps). I think we should keep as-is so that our CI usage is kept minimal.

eddyxu · 2024-02-26T23:28:36Z

@beinan @LuQQiu thanks for working on this. Please let us know what else we can help to move this forward. We are excited to have a Java SDK!

chebbyChefNEQ · 2024-02-27T02:48:11Z

.github/workflows/java.yml

+        uses: actions/setup-java@v4
+        with:
+          distribution: temurin
+          java-version: 11


this should be 17 instead?

chebbyChefNEQ · 2024-02-27T02:49:42Z

.github/workflows/rust.yml

@@ -35,9 +35,6 @@ jobs:
        toolchain:
          - stable
          - nightly
-    defaults:


I think https://github.com/lancedb/lance/blob/main/.github/workflows/bump-version/action.yml also needs to be updated

chebbyChefNEQ · 2024-02-27T02:50:07Z

java/lance-jni/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "lance-jni"
+version = "0.1.0"


should this version stay consistent with the rust crate?

Sure i can use the rust version. Still need to find a way to bump maven version i guess.

chebbyChefNEQ · 2024-02-27T02:52:33Z

java/lance-jni/src/jni_helpers.rs

+use jni::JNIEnv;
+
+pub fn throw_java_exception(env: &mut JNIEnv, err_msg: &str) {
+    env.throw_new("java/lang/RuntimeException", err_msg)


nit: should we make a com.lancedb.lance.LanceException class? a bare RuntimeException might make catch statements unnecessarily wide

ok, will do as follow up.

Co-authored-by: Will Jones <willjones127@gmail.com>

java/lance-jni/src/blocking_dataset.rs

java/lance-jni/src/lib.rs

Co-authored-by: Will Jones <willjones127@gmail.com>

#1928 moved the workspace `Cargo.toml` to the root directory

eddyxu reviewed Feb 10, 2024

View reviewed changes

beinan force-pushed the java_jni branch 2 times, most recently from 378ed98 to f809e86 Compare February 13, 2024 07:15

beinan force-pushed the java_jni branch 4 times, most recently from b735c42 to 1dfe53a Compare February 16, 2024 23:18

eddyxu reviewed Feb 16, 2024

View reviewed changes

beinan force-pushed the java_jni branch 7 times, most recently from 3c4badc to 824181c Compare February 17, 2024 05:51

LuQQiu force-pushed the java_jni branch 3 times, most recently from 29bbe7c to 3447b43 Compare February 19, 2024 05:53

eddyxu reviewed Feb 20, 2024

View reviewed changes

wjones127 requested changes Feb 20, 2024

View reviewed changes

beinan force-pushed the java_jni branch 3 times, most recently from 392454e to 43d6959 Compare February 21, 2024 01:19

eddyxu force-pushed the java_jni branch from a6be864 to 15b8c07 Compare February 27, 2024 00:50

chebbyChefNEQ reviewed Feb 27, 2024

View reviewed changes

beinan and others added 15 commits February 27, 2024 07:52

Java jni code skeleton

08853d1

Add plugin to build rust project in java mvn build

da9e820

Fix build of jni project with new lance and rust version

1e9e42f

Add github workflow for Java JNI project

d03978f

Add version to avoid compile time error

71c676a

Add WriteParams

4bac79c

Add cargo fmt and run cargo fmt & clippy

b891c13

Improve WriteOptions: no java defaults, use map

19575e3

Add code style check for both java and rust

32c3d16

Update java/lance-jni/src/lib.rs

fed49df

Co-authored-by: Will Jones <willjones127@gmail.com>

lazy static

5a63331

ca

e4ca66d

exclude python package

fd1dd55

fix ci

05ec38e

install java17 first

2a1bf4b

eddyxu force-pushed the java_jni branch from 7c715e2 to 2a1bf4b Compare February 27, 2024 15:54

sync versions

13f1800

wjones127 reviewed Feb 27, 2024

View reviewed changes

java/lance-jni/src/blocking_dataset.rs Outdated Show resolved Hide resolved

address remaining comments

285b723

wjones127 reviewed Feb 27, 2024

View reviewed changes

java/lance-jni/src/lib.rs Outdated Show resolved Hide resolved

eddyxu and others added 2 commits February 27, 2024 09:23

Update java/lance-jni/src/lib.rs

5518d1e

Co-authored-by: Will Jones <willjones127@gmail.com>

fmt

4cff35a

wjones127 approved these changes Feb 27, 2024

View reviewed changes

eddyxu merged commit e79db4f into lancedb:main Feb 27, 2024
17 checks passed

chebbyChefNEQ mentioned this pull request Feb 28, 2024

fix: crate publish #2008

Merged

chebbyChefNEQ added a commit that referenced this pull request Feb 28, 2024

fix: crate publish (#2008)

e9cd804

#1928 moved the workspace `Cargo.toml` to the root directory

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: implement java bindings #1928

feat: implement java bindings #1928

beinan commented Feb 8, 2024

eddyxu Feb 10, 2024

beinan Feb 13, 2024

bazooka720 Feb 28, 2024

eddyxu Feb 10, 2024

beinan Feb 13, 2024

eddyxu Feb 10, 2024

beinan Feb 13, 2024

eddyxu Feb 10, 2024 •

edited

beinan Feb 13, 2024

beinan commented Feb 13, 2024

eddyxu Feb 16, 2024 •

edited

eddyxu Feb 20, 2024

wjones127 Feb 20, 2024

beinan Feb 27, 2024

eddyxu Feb 20, 2024

beinan Feb 27, 2024

wjones127 left a comment

wjones127 Feb 20, 2024

chebbyChefNEQ Feb 27, 2024

wjones127 Feb 27, 2024

eddyxu Feb 27, 2024

wjones127 Feb 20, 2024

wjones127 Feb 20, 2024

eddyxu commented Feb 26, 2024

chebbyChefNEQ Feb 27, 2024

eddyxu Feb 27, 2024

chebbyChefNEQ Feb 27, 2024

eddyxu Feb 27, 2024

chebbyChefNEQ Feb 27, 2024

eddyxu Feb 27, 2024

chebbyChefNEQ Feb 27, 2024

eddyxu Feb 27, 2024

		@@ -0,0 +1,36 @@
		use arrow::ffi_stream::ArrowArrayStreamReader;

feat: implement java bindings #1928

feat: implement java bindings #1928

Conversation

beinan commented Feb 8, 2024

Write Path:

Read Path:

Next Steps:

Highlights:

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eddyxu Feb 10, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

beinan commented Feb 13, 2024

eddyxu Feb 16, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wjones127 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eddyxu commented Feb 26, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eddyxu Feb 10, 2024 •

edited

eddyxu Feb 16, 2024 •

edited