FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

feat: configure SessionContext and RuntimeEnv via builder by andygrove · Pull Request #28 · apache/datafusion-java · GitHub

Repository navigation

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .java  (3) .lock  (1) .proto  (1) .rs  (3) .toml  (1) .xml  (1) All 6 file types selected
Only manifest files
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
100 changes: 100 additions & 0 deletions native/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions native/Cargo.toml
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,7 @@ datafusion-proto = "53.1.0"
jni = "0.21"
prost = "0.14"
tokio = { version = "1", features = ["rt-multi-thread"] }

[build-dependencies]
prost-build = "0.14"
protoc-bin-vendored = "3"
24 changes: 24 additions & 0 deletions native/build.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

fn main() {
println!("cargo:rerun-if-changed=../proto/session_options.proto");
let protoc = protoc_bin_vendored::protoc_bin_path().expect("vendored protoc not available");
std::env::set_var("PROTOC", protoc);
prost_build::compile_protos(&["../proto/session_options.proto"], &["../proto"])
.expect("failed to compile session_options.proto");
}
50 changes: 49 additions & 1 deletion native/src/lib.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,11 @@ mod csv;
mod errors;
mod proto;

pub(crate) mod session_options {
include!(concat!(env!("OUT_DIR"), "/datafusion_java.rs"));
}

use std::path::PathBuf;
use std::sync::{Arc, OnceLock};

use datafusion::arrow::datatypes::{Schema, SchemaRef};
Expand All @@ -27,13 +32,16 @@ use datafusion::arrow::ipc::reader::StreamReader;
use datafusion::arrow::record_batch::RecordBatchIterator;
use datafusion::dataframe::DataFrame;
use datafusion::error::DataFusionError;
use datafusion::prelude::{ParquetReadOptions, SessionContext};
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
use jni::objects::{JByteArray, JClass, JObjectArray, JString};
use jni::sys::{jboolean, jint, jlong};
use jni::JNIEnv;
use prost::Message;
use tokio::runtime::Runtime;

use crate::errors::{try_unwrap_or_throw, JniResult};
use crate::session_options::SessionOptions;

pub(crate) fn runtime() -> &'static Runtime {
static RT: OnceLock<Runtime> = OnceLock::new();
Expand All @@ -51,6 +59,46 @@ pub extern "system" fn Java_org_apache_datafusion_SessionContext_createSessionCo
})
}

#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_SessionContext_createSessionContextWithOptions<
'local,
>(
mut env: JNIEnv<'local>,
_class: JClass<'local>,
options_bytes: JByteArray<'local>,
) -> jlong {
try_unwrap_or_throw(&mut env, 0, |env| -> JniResult<jlong> {
let bytes: Vec<u8> = env.convert_byte_array(&options_bytes)?;
let opts = SessionOptions::decode(bytes.as_slice())?;

let mut config = SessionConfig::new();
if let Some(v) = opts.batch_size {
config = config.with_batch_size(v as usize);
}
if let Some(v) = opts.target_partitions {
config = config.with_target_partitions(v as usize);
}
if let Some(v) = opts.collect_statistics {
config = config.with_collect_statistics(v);
}
if let Some(v) = opts.information_schema {
config = config.with_information_schema(v);
}

let mut runtime = RuntimeEnvBuilder::new();
if let Some(mem) = opts.memory_limit {
runtime = runtime.with_memory_limit(mem.max_memory_bytes as usize, mem.memory_fraction);
}
if let Some(dir) = opts.temp_directory {
runtime = runtime.with_temp_file_path(PathBuf::from(dir));
}

let runtime_env = runtime.build()?;
let ctx = SessionContext::new_with_config_rt(config, Arc::new(runtime_env));
Ok(Box::into_raw(Box::new(ctx)) as jlong)
})
}

#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_SessionContext_createDataFrame<'local>(
mut env: JNIEnv<'local>,
Expand Down
40 changes: 23 additions & 17 deletions native/src/proto.rs
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ use crate::errors::{try_unwrap_or_throw, JniResult};
use crate::runtime;

#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_SessionContext_createDataFrameFromProto<'local>(
pub extern "system" fn Java_org_apache_datafusion_SessionContext_createDataFrameFromProto<
'local,
>(
mut env: JNIEnv<'local>,
_class: JClass<'local>,
handle: jlong,
Expand Down Expand Up @@ -60,22 +62,26 @@ pub extern "system" fn Java_org_apache_datafusion_SessionContext_tableSchemaIpc<
handle: jlong,
name: JString<'local>,
) -> jbyteArray {
try_unwrap_or_throw(&mut env, std::ptr::null_mut(), |env| -> JniResult<jbyteArray> {
if handle == 0 {
return Err("SessionContext handle is null".into());
}
let ctx = unsafe { &*(handle as *const SessionContext) };
let name: String = env.get_string(&name)?.into();
try_unwrap_or_throw(
&mut env,
std::ptr::null_mut(),
|env| -> JniResult<jbyteArray> {
if handle == 0 {
return Err("SessionContext handle is null".into());
}
let ctx = unsafe { &*(handle as *const SessionContext) };
let name: String = env.get_string(&name)?.into();

let df = runtime().block_on(ctx.table(name.as_str()))?;
let schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());
let df = runtime().block_on(ctx.table(name.as_str()))?;
let schema: SchemaRef = Arc::new(df.schema().as_arrow().clone());

let mut buf: Vec<u8> = Vec::new();
{
let mut writer = StreamWriter::try_new(&mut buf, schema.as_ref())?;
writer.finish()?;
}
let arr = env.byte_array_from_slice(&buf)?;
Ok(arr.into_raw())
})
let mut buf: Vec<u8> = Vec::new();
{
let mut writer = StreamWriter::try_new(&mut buf, schema.as_ref())?;
writer.finish()?;
}
let arr = env.byte_array_from_slice(&buf)?;
Ok(arr.into_raw())
},
)
}
9 changes: 9 additions & 0 deletions pom.xml
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -192,8 +192,17 @@ under the License.
</configuration>
<executions>
<execution>
<id>compile-upstream</id>
<goals><goal>compile</goal></goals>
</execution>
<execution>
<id>compile-local</id>
<goals><goal>compile</goal></goals>
<configuration>
<protoSourceRoot>${project.basedir}/proto</protoSourceRoot>
<clearOutputDirectory>false</clearOutputDirectory>
</configuration>
</execution>
</executions>
</plugin>
</plugins>
Expand Down
39 changes: 39 additions & 0 deletions proto/session_options.proto
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

syntax = "proto3";

package datafusion_java;

option java_package = "org.apache.datafusion.protobuf";
option java_multiple_files = true;

// Options used to construct a SessionContext. All fields are optional —
// unset fields leave DataFusion's default behavior in place.
message SessionOptions {
optional uint64 batch_size = 1;
optional uint64 target_partitions = 2;
optional bool collect_statistics = 3;
optional bool information_schema = 4;
optional MemoryLimit memory_limit = 5;
optional string temp_directory = 6;
}

message MemoryLimit {
uint64 max_memory_bytes = 1;
double memory_fraction = 2;
}
14 changes: 14 additions & 0 deletions src/main/java/org/apache/datafusion/SessionContext.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,18 @@ public SessionContext() {
}
}

SessionContext(byte[] optionsBytes) {
this.nativeHandle = createSessionContextWithOptions(optionsBytes);
if (this.nativeHandle == 0) {
throw new RuntimeException("Failed to create native SessionContext");
}
}

/** Start configuring a {@link SessionContext}. */
public static SessionContextBuilder builder() {
return new SessionContextBuilder();
}

/**
* Parse and plan {@code query}, returning a lazy {@link DataFrame}. The query is not executed
* until {@link DataFrame#collect} is called.
Expand Down Expand Up @@ -251,6 +263,8 @@ public void close() {

private static native long createSessionContext();

private static native long createSessionContextWithOptions(byte[] optionsBytes);

private static native long createDataFrame(long handle, String sql);

private static native long createDataFrameFromProto(long handle, byte[] planBytes);
Expand Down
Loading
Loading

Back | FazBrowse Home | New Git URL