This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-java.git
The following commit(s) were added to refs/heads/main by this push:
new 3698b9a feat: configure SessionContext and RuntimeEnv via builder
(#28)
3698b9a is described below
commit 3698b9a1af9a741c5935a32de6aaa8e4a2947b69
Author: Andy Grove <[email protected]>
AuthorDate: Wed May 13 10:44:15 2026 -0600
feat: configure SessionContext and RuntimeEnv via builder (#28)
---
native/Cargo.lock | 100 +++++++++++++++++
native/Cargo.toml | 4 +
native/build.rs | 24 ++++
native/src/lib.rs | 50 ++++++++-
native/src/proto.rs | 40 ++++---
pom.xml | 9 ++
proto/session_options.proto | 39 +++++++
.../java/org/apache/datafusion/SessionContext.java | 14 +++
.../apache/datafusion/SessionContextBuilder.java | 124 +++++++++++++++++++++
.../datafusion/SessionContextBuilderTest.java | 116 +++++++++++++++++++
10 files changed, 502 insertions(+), 18 deletions(-)
diff --git a/native/Cargo.lock b/native/Cargo.lock
index 004fa98..495cc60 100644
--- a/native/Cargo.lock
+++ b/native/Cargo.lock
@@ -1143,6 +1143,8 @@ dependencies = [
"datafusion-proto",
"jni",
"prost",
+ "prost-build",
+ "protoc-bin-vendored",
"tokio",
]
@@ -2084,6 +2086,12 @@ dependencies = [
"simd-adler32",
]
+[[package]]
+name = "multimap"
+version = "0.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084"
+
[[package]]
name = "num-bigint"
version = "0.4.6"
@@ -2332,6 +2340,25 @@ dependencies = [
"prost-derive",
]
+[[package]]
+name = "prost-build"
+version = "0.14.3"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7"
+dependencies = [
+ "heck",
+ "itertools",
+ "log",
+ "multimap",
+ "petgraph",
+ "prettyplease",
+ "prost",
+ "prost-types",
+ "regex",
+ "syn",
+ "tempfile",
+]
+
[[package]]
name = "prost-derive"
version = "0.14.3"
@@ -2345,6 +2372,79 @@ dependencies = [
"syn",
]
+[[package]]
+name = "prost-types"
+version = "0.14.3"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7"
+dependencies = [
+ "prost",
+]
+
+[[package]]
+name = "protoc-bin-vendored"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "d1c381df33c98266b5f08186583660090a4ffa0889e76c7e9a5e175f645a67fa"
+dependencies = [
+ "protoc-bin-vendored-linux-aarch_64",
+ "protoc-bin-vendored-linux-ppcle_64",
+ "protoc-bin-vendored-linux-s390_64",
+ "protoc-bin-vendored-linux-x86_32",
+ "protoc-bin-vendored-linux-x86_64",
+ "protoc-bin-vendored-macos-aarch_64",
+ "protoc-bin-vendored-macos-x86_64",
+ "protoc-bin-vendored-win32",
+]
+
+[[package]]
+name = "protoc-bin-vendored-linux-aarch_64"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "c350df4d49b5b9e3ca79f7e646fde2377b199e13cfa87320308397e1f37e1a4c"
+
+[[package]]
+name = "protoc-bin-vendored-linux-ppcle_64"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "a55a63e6c7244f19b5c6393f025017eb5d793fd5467823a099740a7a4222440c"
+
+[[package]]
+name = "protoc-bin-vendored-linux-s390_64"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "1dba5565db4288e935d5330a07c264a4ee8e4a5b4a4e6f4e83fad824cc32f3b0"
+
+[[package]]
+name = "protoc-bin-vendored-linux-x86_32"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "8854774b24ee28b7868cd71dccaae8e02a2365e67a4a87a6cd11ee6cdbdf9cf5"
+
+[[package]]
+name = "protoc-bin-vendored-linux-x86_64"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "b38b07546580df720fa464ce124c4b03630a6fb83e05c336fea2a241df7e5d78"
+
+[[package]]
+name = "protoc-bin-vendored-macos-aarch_64"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "89278a9926ce312e51f1d999fee8825d324d603213344a9a706daa009f1d8092"
+
+[[package]]
+name = "protoc-bin-vendored-macos-x86_64"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "81745feda7ccfb9471d7a4de888f0652e806d5795b61480605d4943176299756"
+
+[[package]]
+name = "protoc-bin-vendored-win32"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index";
+checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3"
+
[[package]]
name = "psm"
version = "0.1.31"
diff --git a/native/Cargo.toml b/native/Cargo.toml
index f2a6818..01dd002 100644
--- a/native/Cargo.toml
+++ b/native/Cargo.toml
@@ -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"
diff --git a/native/build.rs b/native/build.rs
new file mode 100644
index 0000000..495e147
--- /dev/null
+++ b/native/build.rs
@@ -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");
+}
diff --git a/native/src/lib.rs b/native/src/lib.rs
index 4ea1a61..677aaf9 100644
--- a/native/src/lib.rs
+++ b/native/src/lib.rs
@@ -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};
@@ -29,13 +34,16 @@ use datafusion::config::TableParquetOptions;
use datafusion::dataframe::DataFrame;
use datafusion::dataframe::DataFrameWriteOptions;
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();
@@ -53,6 +61,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>,
diff --git a/native/src/proto.rs b/native/src/proto.rs
index 429a91e..d90ed2c 100644
--- a/native/src/proto.rs
+++ b/native/src/proto.rs
@@ -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,
@@ -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())
+ },
+ )
}
diff --git a/pom.xml b/pom.xml
index 9039d8b..47d5b4e 100644
--- a/pom.xml
+++ b/pom.xml
@@ -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>
diff --git a/proto/session_options.proto b/proto/session_options.proto
new file mode 100644
index 0000000..96342a6
--- /dev/null
+++ b/proto/session_options.proto
@@ -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;
+}
diff --git a/src/main/java/org/apache/datafusion/SessionContext.java
b/src/main/java/org/apache/datafusion/SessionContext.java
index 1aec343..a496883 100644
--- a/src/main/java/org/apache/datafusion/SessionContext.java
+++ b/src/main/java/org/apache/datafusion/SessionContext.java
@@ -54,6 +54,18 @@ public final class SessionContext implements AutoCloseable {
}
}
+ 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.
@@ -251,6 +263,8 @@ public final class SessionContext implements AutoCloseable {
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);
diff --git a/src/main/java/org/apache/datafusion/SessionContextBuilder.java
b/src/main/java/org/apache/datafusion/SessionContextBuilder.java
new file mode 100644
index 0000000..bb0cac9
--- /dev/null
+++ b/src/main/java/org/apache/datafusion/SessionContextBuilder.java
@@ -0,0 +1,124 @@
+/*
+ * 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.
+ */
+
+package org.apache.datafusion;
+
+import org.apache.datafusion.protobuf.MemoryLimit;
+import org.apache.datafusion.protobuf.SessionOptions;
+
+/**
+ * Builder for a configured {@link SessionContext}. Each setter is optional;
unset fields leave the
+ * DataFusion default in place.
+ */
+public final class SessionContextBuilder {
+ private Integer batchSize;
+ private Integer targetPartitions;
+ private Boolean collectStatistics;
+ private Boolean informationSchema;
+ private Long memoryLimitBytes;
+ private Double memoryLimitFraction;
+ private String tempDirectory;
+
+ SessionContextBuilder() {}
+
+ public SessionContextBuilder batchSize(int batchSize) {
+ if (batchSize <= 0) {
+ throw new IllegalArgumentException("batchSize must be positive, got " +
batchSize);
+ }
+ this.batchSize = batchSize;
+ return this;
+ }
+
+ public SessionContextBuilder targetPartitions(int targetPartitions) {
+ if (targetPartitions <= 0) {
+ throw new IllegalArgumentException(
+ "targetPartitions must be positive, got " + targetPartitions);
+ }
+ this.targetPartitions = targetPartitions;
+ return this;
+ }
+
+ public SessionContextBuilder collectStatistics(boolean collectStatistics) {
+ this.collectStatistics = collectStatistics;
+ return this;
+ }
+
+ public SessionContextBuilder informationSchema(boolean informationSchema) {
+ this.informationSchema = informationSchema;
+ return this;
+ }
+
+ /**
+ * Cap the memory pool at {@code maxMemoryBytes}, reserving {@code fraction}
of it for queries.
+ */
+ public SessionContextBuilder memoryLimit(long maxMemoryBytes, double
fraction) {
+ if (maxMemoryBytes <= 0) {
+ throw new IllegalArgumentException(
+ "maxMemoryBytes must be positive, got " + maxMemoryBytes);
+ }
+ if (fraction <= 0.0 || fraction > 1.0) {
+ throw new IllegalArgumentException("fraction must be in (0, 1], got " +
fraction);
+ }
+ this.memoryLimitBytes = maxMemoryBytes;
+ this.memoryLimitFraction = fraction;
+ return this;
+ }
+
+ /** Directory the DiskManager uses for spill files. */
+ public SessionContextBuilder tempDirectory(String path) {
+ this.tempDirectory = path;
+ return this;
+ }
+
+ /**
+ * Construct a {@link SessionContext} with the configured options.
+ *
+ * @throws RuntimeException if the native side fails to construct the
context.
+ */
+ public SessionContext build() {
+ return new SessionContext(toBytes());
+ }
+
+ byte[] toBytes() {
+ SessionOptions.Builder b = SessionOptions.newBuilder();
+ if (batchSize != null) {
+ b.setBatchSize(batchSize);
+ }
+ if (targetPartitions != null) {
+ b.setTargetPartitions(targetPartitions);
+ }
+ if (collectStatistics != null) {
+ b.setCollectStatistics(collectStatistics);
+ }
+ if (informationSchema != null) {
+ b.setInformationSchema(informationSchema);
+ }
+ if (memoryLimitBytes != null && memoryLimitFraction != null) {
+ b.setMemoryLimit(
+ MemoryLimit.newBuilder()
+ .setMaxMemoryBytes(memoryLimitBytes)
+ .setMemoryFraction(memoryLimitFraction)
+ .build());
+ }
+ if (tempDirectory != null) {
+ b.setTempDirectory(tempDirectory);
+ }
+ return b.build().toByteArray();
+ }
+}
diff --git a/src/test/java/org/apache/datafusion/SessionContextBuilderTest.java
b/src/test/java/org/apache/datafusion/SessionContextBuilderTest.java
new file mode 100644
index 0000000..ed6cff2
--- /dev/null
+++ b/src/test/java/org/apache/datafusion/SessionContextBuilderTest.java
@@ -0,0 +1,116 @@
+/*
+ * 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.
+ */
+
+package org.apache.datafusion;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Path;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.ipc.ArrowReader;
+import org.apache.datafusion.protobuf.SessionOptions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+class SessionContextBuilderTest {
+
+ @Test
+ void protoRoundTripPreservesAllFields() throws Exception {
+ byte[] bytes =
+ SessionContext.builder()
+ .batchSize(8192)
+ .targetPartitions(4)
+ .collectStatistics(true)
+ .informationSchema(true)
+ .memoryLimit(1L << 30, 0.8)
+ .tempDirectory("/tmp/df")
+ .toBytes();
+
+ SessionOptions parsed = SessionOptions.parseFrom(bytes);
+ assertTrue(parsed.hasBatchSize());
+ assertEquals(8192L, parsed.getBatchSize());
+ assertTrue(parsed.hasTargetPartitions());
+ assertEquals(4L, parsed.getTargetPartitions());
+ assertTrue(parsed.hasCollectStatistics());
+ assertTrue(parsed.getCollectStatistics());
+ assertTrue(parsed.hasInformationSchema());
+ assertTrue(parsed.getInformationSchema());
+ assertTrue(parsed.hasMemoryLimit());
+ assertEquals(1L << 30, parsed.getMemoryLimit().getMaxMemoryBytes());
+ assertEquals(0.8, parsed.getMemoryLimit().getMemoryFraction(), 1e-9);
+ assertTrue(parsed.hasTempDirectory());
+ assertEquals("/tmp/df", parsed.getTempDirectory());
+ }
+
+ @Test
+ void unsetFieldsAreAbsentInProto() throws Exception {
+ byte[] bytes = SessionContext.builder().batchSize(8192).toBytes();
+ SessionOptions parsed = SessionOptions.parseFrom(bytes);
+ assertTrue(parsed.hasBatchSize());
+ assertFalse(parsed.hasTargetPartitions());
+ assertFalse(parsed.hasCollectStatistics());
+ assertFalse(parsed.hasInformationSchema());
+ assertFalse(parsed.hasMemoryLimit());
+ assertFalse(parsed.hasTempDirectory());
+ }
+
+ @Test
+ void informationSchemaEnabledMakesMetaQueryRun() throws Exception {
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx =
SessionContext.builder().informationSchema(true).build();
+ DataFrame df = ctx.sql("SELECT table_name FROM
information_schema.tables");
+ ArrowReader reader = df.collect(allocator)) {
+ // information_schema.tables always has at least one row.
+ assertTrue(reader.loadNextBatch());
+ }
+ }
+
+ @Test
+ void informationSchemaDisabledByDefaultThrows() {
+ try (SessionContext ctx = SessionContext.builder().build()) {
+ assertThrows(
+ RuntimeException.class,
+ () -> ctx.sql("SELECT table_name FROM information_schema.tables"));
+ }
+ }
+
+ @Test
+ void buildWithEveryKnobSetCanExecuteSelectOne(@TempDir Path tempDir) throws
Exception {
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx =
+ SessionContext.builder()
+ .batchSize(8192)
+ .targetPartitions(4)
+ .collectStatistics(true)
+ .informationSchema(true)
+ .memoryLimit(1L << 30, 0.8)
+ .tempDirectory(tempDir.toAbsolutePath().toString())
+ .build();
+ DataFrame df = ctx.sql("SELECT 1");
+ ArrowReader reader = df.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ assertEquals(1, reader.getVectorSchemaRoot().getRowCount());
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
