refactor: pass csv and parquet read options via protobuf (#29)
diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md
index c6818f6..0b42a2e 100644
--- a/docs/source/contributor-guide/development.md
+++ b/docs/source/contributor-guide/development.md
@@ -74,4 +74,40 @@
- `mvnw`, `mvnw.cmd` — bundled Maven wrapper.
- `src/` — Java sources and tests.
- `native/` — Rust crate (JNI + Arrow C Data Interface).
+- `proto/` — Protobuf definitions shared between Java and Rust.
- `docs/` — Sphinx documentation source and build scripts.
+
+## Passing structured options across the JNI boundary
+
+When a JNI call needs to carry more than a handful of scalar arguments —
+for example, a struct of nullable knobs like `CsvReadOptions` or
+`SessionOptions` — encode the call's configuration as a protobuf message
+rather than expanding the JNI signature with one parameter per field.
+
+Add a `.proto` file under `proto/`, declare `package datafusion_java;`,
+and follow the conventions already in use:
+
+- One message per logical bundle of options.
+- Use `optional` for fields whose unset-ness must survive the boundary
+ (so the Rust side can leave a DataFusion default in place).
+- Use proto enums for closed sets of choices instead of strings.
+ Prefix every value with the enum's name (e.g.
+ `FILE_COMPRESSION_TYPE_GZIP`, not bare `GZIP`) because proto3 enum
+ values are scoped at the package level, and the zero value must be a
+ `_UNSPECIFIED` sentinel — the Rust side should reject `UNSPECIFIED`
+ rather than silently default it.
+- Keep large opaque payloads (Arrow IPC schemas, plan nodes) as separate
+ `byte[]` JNI arguments next to the options proto, not inside it.
+- Suffix the message name with `Proto` if a sibling Java class would
+ otherwise shadow it (e.g. `CsvReadOptionsProto` vs the public
+ `CsvReadOptions`).
+
+The proto is compiled by both `prost-build` (Rust, via `native/build.rs`)
+and the Maven `protobuf-maven-plugin` (Java). The Java side builds the
+message, serializes to bytes, and passes the byte array through JNI; the
+Rust side decodes once and folds the fields into the corresponding
+DataFusion struct.
+
+This pattern keeps JNI signatures short, makes nullable and enum fields
+explicit in a single typed schema, and lets new fields be added without
+touching the signature on either side.
diff --git a/native/build.rs b/native/build.rs
index 495e147..5a27cb0 100644
--- a/native/build.rs
+++ b/native/build.rs
@@ -16,9 +16,15 @@
// under the License.
fn main() {
- println!("cargo:rerun-if-changed=../proto/session_options.proto");
+ const PROTOS: &[&str] = &[
+ "../proto/session_options.proto",
+ "../proto/csv_read_options.proto",
+ "../proto/parquet_read_options.proto",
+ ];
+ for p in PROTOS {
+ println!("cargo:rerun-if-changed={p}");
+ }
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");
+ prost_build::compile_protos(PROTOS, &["../proto"]).expect("failed to compile protos");
}
diff --git a/native/src/csv.rs b/native/src/csv.rs
index a1b6fdd..3201951 100644
--- a/native/src/csv.rs
+++ b/native/src/csv.rs
@@ -15,73 +15,63 @@
// specific language governing permissions and limitations
// under the License.
-use std::str::FromStr;
-
-use datafusion::arrow::datatypes::Schema;
-use datafusion::arrow::ipc::reader::StreamReader;
use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
use datafusion::error::DataFusionError;
use datafusion::prelude::{CsvReadOptions, SessionContext};
use jni::objects::{JByteArray, JClass, JString};
-use jni::sys::{jboolean, jbyte, jlong};
+use jni::sys::jlong;
use jni::JNIEnv;
+use prost::Message;
use crate::errors::{try_unwrap_or_throw, JniResult};
+use crate::proto_gen::{CsvReadOptionsProto, FileCompressionType as ProtoFileCompressionType};
use crate::runtime;
+use crate::schema::decode_optional_schema;
-#[allow(clippy::too_many_arguments)]
fn with_csv_options<R>(
env: &mut JNIEnv,
- has_header: jboolean,
- delimiter: jbyte,
- quote: jbyte,
- terminator_set: jboolean,
- terminator_value: jbyte,
- escape_set: jboolean,
- escape_value: jbyte,
- comment_set: jboolean,
- comment_value: jbyte,
- newlines_in_values_set: jboolean,
- newlines_in_values_value: jboolean,
- schema_infer_max_records: jlong,
- file_extension: JString,
- file_compression_type: JString,
+ options_bytes: JByteArray,
schema_ipc_bytes: JByteArray,
f: impl FnOnce(CsvReadOptions) -> JniResult<R>,
) -> JniResult<R> {
- let file_ext: String = env.get_string(&file_extension)?.into();
- let compression: String = env.get_string(&file_compression_type)?.into();
- let compression = FileCompressionType::from_str(&compression)?;
+ let bytes: Vec<u8> = env.convert_byte_array(&options_bytes)?;
+ let p = CsvReadOptionsProto::decode(bytes.as_slice())?;
- let schema: Option<Schema> = if !schema_ipc_bytes.is_null() {
- let bytes: Vec<u8> = env.convert_byte_array(&schema_ipc_bytes)?;
- let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)?;
- Some((*reader.schema()).clone())
- } else {
- None
+ let schema = decode_optional_schema(env, schema_ipc_bytes)?;
+
+ let compression = match p.file_compression_type() {
+ ProtoFileCompressionType::Unspecified => {
+ return Err("CsvReadOptionsProto.file_compression_type is UNSPECIFIED".into());
+ }
+ ProtoFileCompressionType::Uncompressed => FileCompressionType::UNCOMPRESSED,
+ ProtoFileCompressionType::Gzip => FileCompressionType::GZIP,
+ ProtoFileCompressionType::Bzip2 => FileCompressionType::BZIP2,
+ ProtoFileCompressionType::Xz => FileCompressionType::XZ,
+ ProtoFileCompressionType::Zstd => FileCompressionType::ZSTD,
};
+ let file_ext = p.file_extension;
let mut opts = CsvReadOptions::new()
- .has_header(has_header != 0)
- .delimiter(delimiter as u8)
- .quote(quote as u8)
+ .has_header(p.has_header)
+ .delimiter(p.delimiter as u8)
+ .quote(p.quote as u8)
.file_extension(&file_ext)
.file_compression_type(compression);
- if terminator_set != 0 {
- opts = opts.terminator(Some(terminator_value as u8));
+ if let Some(t) = p.terminator {
+ opts = opts.terminator(Some(t as u8));
}
- if escape_set != 0 {
- opts = opts.escape(escape_value as u8);
+ if let Some(e) = p.escape {
+ opts = opts.escape(e as u8);
}
- if comment_set != 0 {
- opts = opts.comment(comment_value as u8);
+ if let Some(c) = p.comment {
+ opts = opts.comment(c as u8);
}
- if newlines_in_values_set != 0 {
- opts = opts.newlines_in_values(newlines_in_values_value != 0);
+ if let Some(n) = p.newlines_in_values {
+ opts = opts.newlines_in_values(n);
}
- if schema_infer_max_records >= 0 {
- opts = opts.schema_infer_max_records(schema_infer_max_records as usize);
+ if let Some(n) = p.schema_infer_max_records {
+ opts = opts.schema_infer_max_records(n as usize);
}
if let Some(ref s) = schema {
opts = opts.schema(s);
@@ -90,7 +80,6 @@
f(opts)
}
-#[allow(clippy::too_many_arguments)]
#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_SessionContext_registerCsvWithOptions<'local>(
mut env: JNIEnv<'local>,
@@ -98,20 +87,7 @@
handle: jlong,
name: JString<'local>,
path: JString<'local>,
- has_header: jboolean,
- delimiter: jbyte,
- quote: jbyte,
- terminator_set: jboolean,
- terminator_value: jbyte,
- escape_set: jboolean,
- escape_value: jbyte,
- comment_set: jboolean,
- comment_value: jbyte,
- newlines_in_values_set: jboolean,
- newlines_in_values_value: jboolean,
- schema_infer_max_records: jlong,
- file_extension: JString<'local>,
- file_compression_type: JString<'local>,
+ options_bytes: JByteArray<'local>,
schema_ipc_bytes: JByteArray<'local>,
) {
try_unwrap_or_throw(&mut env, (), |env| -> JniResult<()> {
@@ -121,55 +97,23 @@
let ctx = unsafe { &*(handle as *const SessionContext) };
let name: String = env.get_string(&name)?.into();
let path: String = env.get_string(&path)?.into();
- with_csv_options(
- env,
- has_header,
- delimiter,
- quote,
- terminator_set,
- terminator_value,
- escape_set,
- escape_value,
- comment_set,
- comment_value,
- newlines_in_values_set,
- newlines_in_values_value,
- schema_infer_max_records,
- file_extension,
- file_compression_type,
- schema_ipc_bytes,
- |opts| {
- runtime().block_on(async {
- ctx.register_csv(&name, &path, opts).await?;
- Ok::<(), DataFusionError>(())
- })?;
- Ok(())
- },
- )
+ with_csv_options(env, options_bytes, schema_ipc_bytes, |opts| {
+ runtime().block_on(async {
+ ctx.register_csv(&name, &path, opts).await?;
+ Ok::<(), DataFusionError>(())
+ })?;
+ Ok(())
+ })
})
}
-#[allow(clippy::too_many_arguments)]
#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_SessionContext_readCsvWithOptions<'local>(
mut env: JNIEnv<'local>,
_class: JClass<'local>,
handle: jlong,
path: JString<'local>,
- has_header: jboolean,
- delimiter: jbyte,
- quote: jbyte,
- terminator_set: jboolean,
- terminator_value: jbyte,
- escape_set: jboolean,
- escape_value: jbyte,
- comment_set: jboolean,
- comment_value: jbyte,
- newlines_in_values_set: jboolean,
- newlines_in_values_value: jboolean,
- schema_infer_max_records: jlong,
- file_extension: JString<'local>,
- file_compression_type: JString<'local>,
+ options_bytes: JByteArray<'local>,
schema_ipc_bytes: JByteArray<'local>,
) -> jlong {
try_unwrap_or_throw(&mut env, 0, |env| -> JniResult<jlong> {
@@ -178,27 +122,9 @@
}
let ctx = unsafe { &*(handle as *const SessionContext) };
let path: String = env.get_string(&path)?.into();
- with_csv_options(
- env,
- has_header,
- delimiter,
- quote,
- terminator_set,
- terminator_value,
- escape_set,
- escape_value,
- comment_set,
- comment_value,
- newlines_in_values_set,
- newlines_in_values_value,
- schema_infer_max_records,
- file_extension,
- file_compression_type,
- schema_ipc_bytes,
- |opts| {
- let df = runtime().block_on(ctx.read_csv(path, opts))?;
- Ok(Box::into_raw(Box::new(df)) as jlong)
- },
- )
+ with_csv_options(env, options_bytes, schema_ipc_bytes, |opts| {
+ let df = runtime().block_on(ctx.read_csv(path, opts))?;
+ Ok(Box::into_raw(Box::new(df)) as jlong)
+ })
})
}
diff --git a/native/src/lib.rs b/native/src/lib.rs
index 677aaf9..d6e304c 100644
--- a/native/src/lib.rs
+++ b/native/src/lib.rs
@@ -18,17 +18,17 @@
mod csv;
mod errors;
mod proto;
+mod schema;
-pub(crate) mod session_options {
+pub(crate) mod proto_gen {
include!(concat!(env!("OUT_DIR"), "/datafusion_java.rs"));
}
use std::path::PathBuf;
use std::sync::{Arc, OnceLock};
-use datafusion::arrow::datatypes::{Schema, SchemaRef};
+use datafusion::arrow::datatypes::SchemaRef;
use datafusion::arrow::ffi_stream::FFI_ArrowArrayStream;
-use datafusion::arrow::ipc::reader::StreamReader;
use datafusion::arrow::record_batch::RecordBatchIterator;
use datafusion::config::TableParquetOptions;
use datafusion::dataframe::DataFrame;
@@ -43,7 +43,9 @@
use tokio::runtime::Runtime;
use crate::errors::{try_unwrap_or_throw, JniResult};
-use crate::session_options::SessionOptions;
+use crate::proto_gen::ParquetReadOptionsProto;
+use crate::proto_gen::SessionOptions;
+use crate::schema::decode_optional_schema;
pub(crate) fn runtime() -> &'static Runtime {
static RT: OnceLock<Runtime> = OnceLock::new();
@@ -313,37 +315,27 @@
})
}
-#[allow(clippy::too_many_arguments)]
fn with_parquet_options<R>(
env: &mut JNIEnv,
- file_extension: JString,
- parquet_pruning_set: jboolean,
- parquet_pruning_value: jboolean,
- skip_metadata_set: jboolean,
- skip_metadata_value: jboolean,
- metadata_size_hint: jlong,
+ options_bytes: JByteArray,
schema_ipc_bytes: JByteArray,
f: impl FnOnce(ParquetReadOptions) -> JniResult<R>,
) -> JniResult<R> {
- let file_ext: String = env.get_string(&file_extension)?.into();
+ let bytes: Vec<u8> = env.convert_byte_array(&options_bytes)?;
+ let p = ParquetReadOptionsProto::decode(bytes.as_slice())?;
- let schema: Option<Schema> = if !schema_ipc_bytes.is_null() {
- let bytes: Vec<u8> = env.convert_byte_array(&schema_ipc_bytes)?;
- let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)?;
- Some((*reader.schema()).clone())
- } else {
- None
- };
+ let schema = decode_optional_schema(env, schema_ipc_bytes)?;
+ let file_ext = p.file_extension;
let mut opts = ParquetReadOptions::default().file_extension(&file_ext);
- if parquet_pruning_set != 0 {
- opts = opts.parquet_pruning(parquet_pruning_value != 0);
+ if let Some(v) = p.parquet_pruning {
+ opts = opts.parquet_pruning(v);
}
- if skip_metadata_set != 0 {
- opts = opts.skip_metadata(skip_metadata_value != 0);
+ if let Some(v) = p.skip_metadata {
+ opts = opts.skip_metadata(v);
}
- if metadata_size_hint >= 0 {
- opts = opts.metadata_size_hint(Some(metadata_size_hint as usize));
+ if let Some(v) = p.metadata_size_hint {
+ opts = opts.metadata_size_hint(Some(v as usize));
}
if let Some(ref s) = schema {
opts = opts.schema(s);
@@ -361,12 +353,7 @@
handle: jlong,
name: JString<'local>,
path: JString<'local>,
- file_extension: JString<'local>,
- parquet_pruning_set: jboolean,
- parquet_pruning_value: jboolean,
- skip_metadata_set: jboolean,
- skip_metadata_value: jboolean,
- metadata_size_hint: jlong,
+ options_bytes: JByteArray<'local>,
schema_ipc_bytes: JByteArray<'local>,
) {
try_unwrap_or_throw(&mut env, (), |env| -> JniResult<()> {
@@ -376,23 +363,13 @@
let ctx = unsafe { &*(handle as *const SessionContext) };
let name: String = env.get_string(&name)?.into();
let path: String = env.get_string(&path)?.into();
- with_parquet_options(
- env,
- file_extension,
- parquet_pruning_set,
- parquet_pruning_value,
- skip_metadata_set,
- skip_metadata_value,
- metadata_size_hint,
- schema_ipc_bytes,
- |opts| {
- runtime().block_on(async {
- ctx.register_parquet(&name, &path, opts).await?;
- Ok::<(), DataFusionError>(())
- })?;
- Ok(())
- },
- )
+ with_parquet_options(env, options_bytes, schema_ipc_bytes, |opts| {
+ runtime().block_on(async {
+ ctx.register_parquet(&name, &path, opts).await?;
+ Ok::<(), DataFusionError>(())
+ })?;
+ Ok(())
+ })
})
}
@@ -402,12 +379,7 @@
_class: JClass<'local>,
handle: jlong,
path: JString<'local>,
- file_extension: JString<'local>,
- parquet_pruning_set: jboolean,
- parquet_pruning_value: jboolean,
- skip_metadata_set: jboolean,
- skip_metadata_value: jboolean,
- metadata_size_hint: jlong,
+ options_bytes: JByteArray<'local>,
schema_ipc_bytes: JByteArray<'local>,
) -> jlong {
try_unwrap_or_throw(&mut env, 0, |env| -> JniResult<jlong> {
@@ -416,19 +388,9 @@
}
let ctx = unsafe { &*(handle as *const SessionContext) };
let path: String = env.get_string(&path)?.into();
- with_parquet_options(
- env,
- file_extension,
- parquet_pruning_set,
- parquet_pruning_value,
- skip_metadata_set,
- skip_metadata_value,
- metadata_size_hint,
- schema_ipc_bytes,
- |opts| {
- let df = runtime().block_on(ctx.read_parquet(path, opts))?;
- Ok(Box::into_raw(Box::new(df)) as jlong)
- },
- )
+ with_parquet_options(env, options_bytes, schema_ipc_bytes, |opts| {
+ let df = runtime().block_on(ctx.read_parquet(path, opts))?;
+ Ok(Box::into_raw(Box::new(df)) as jlong)
+ })
})
}
diff --git a/native/src/schema.rs b/native/src/schema.rs
new file mode 100644
index 0000000..968a73a
--- /dev/null
+++ b/native/src/schema.rs
@@ -0,0 +1,37 @@
+// 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.
+
+use datafusion::arrow::datatypes::Schema;
+use datafusion::arrow::ipc::reader::StreamReader;
+use jni::objects::JByteArray;
+use jni::JNIEnv;
+
+use crate::errors::JniResult;
+
+/// Decode an optional Arrow-IPC schema byte array passed in from Java.
+/// Returns `None` if the byte-array reference is null.
+pub(crate) fn decode_optional_schema(
+ env: &mut JNIEnv,
+ bytes: JByteArray,
+) -> JniResult<Option<Schema>> {
+ if bytes.is_null() {
+ return Ok(None);
+ }
+ let buf: Vec<u8> = env.convert_byte_array(&bytes)?;
+ let reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
+ Ok(Some((*reader.schema()).clone()))
+}
diff --git a/proto/csv_read_options.proto b/proto/csv_read_options.proto
new file mode 100644
index 0000000..840867d
--- /dev/null
+++ b/proto/csv_read_options.proto
@@ -0,0 +1,48 @@
+// 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 read CSV files. Fields with non-null Java defaults are
+// always sent; fields marked `optional` preserve unset-ness so the Rust
+// side can leave a DataFusion default in place.
+message CsvReadOptionsProto {
+ bool has_header = 1;
+ uint32 delimiter = 2; // single byte, sent as uint32
+ uint32 quote = 3; // single byte
+ optional uint32 terminator = 4;
+ optional uint32 escape = 5;
+ optional uint32 comment = 6;
+ optional bool newlines_in_values = 7;
+ optional uint64 schema_infer_max_records = 8;
+ string file_extension = 9;
+ FileCompressionType file_compression_type = 10;
+}
+
+enum FileCompressionType {
+ FILE_COMPRESSION_TYPE_UNSPECIFIED = 0;
+ FILE_COMPRESSION_TYPE_UNCOMPRESSED = 1;
+ FILE_COMPRESSION_TYPE_GZIP = 2;
+ FILE_COMPRESSION_TYPE_BZIP2 = 3;
+ FILE_COMPRESSION_TYPE_XZ = 4;
+ FILE_COMPRESSION_TYPE_ZSTD = 5;
+}
diff --git a/proto/parquet_read_options.proto b/proto/parquet_read_options.proto
new file mode 100644
index 0000000..3355dd7
--- /dev/null
+++ b/proto/parquet_read_options.proto
@@ -0,0 +1,33 @@
+// 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 read Parquet files. `file_extension` has a non-null
+// Java default and is always sent; the rest preserve unset-ness so the
+// Rust side can leave a DataFusion default in place.
+message ParquetReadOptionsProto {
+ string file_extension = 1;
+ optional bool parquet_pruning = 2;
+ optional bool skip_metadata = 3;
+ optional uint64 metadata_size_hint = 4;
+}
diff --git a/src/main/java/org/apache/datafusion/CsvReadOptions.java b/src/main/java/org/apache/datafusion/CsvReadOptions.java
index 4b35b95..f441820 100644
--- a/src/main/java/org/apache/datafusion/CsvReadOptions.java
+++ b/src/main/java/org/apache/datafusion/CsvReadOptions.java
@@ -109,47 +109,51 @@
return this;
}
- boolean hasHeader() {
- return hasHeader;
- }
-
- byte delimiter() {
- return delimiter;
- }
-
- byte quote() {
- return quote;
- }
-
- Byte terminator() {
- return terminator;
- }
-
- Byte escape() {
- return escape;
- }
-
- Byte comment() {
- return comment;
- }
-
- Boolean newlinesInValues() {
- return newlinesInValues;
- }
-
- Long schemaInferMaxRecords() {
- return schemaInferMaxRecords;
- }
-
- String fileExtension() {
- return fileExtension;
- }
-
- FileCompressionType fileCompressionType() {
- return fileCompressionType;
+ byte[] toBytes() {
+ org.apache.datafusion.protobuf.CsvReadOptionsProto.Builder b =
+ org.apache.datafusion.protobuf.CsvReadOptionsProto.newBuilder()
+ .setHasHeader(hasHeader)
+ .setDelimiter(delimiter & 0xFF)
+ .setQuote(quote & 0xFF)
+ .setFileExtension(fileExtension)
+ .setFileCompressionType(toProto(fileCompressionType));
+ if (terminator != null) {
+ b.setTerminator(terminator & 0xFF);
+ }
+ if (escape != null) {
+ b.setEscape(escape & 0xFF);
+ }
+ if (comment != null) {
+ b.setComment(comment & 0xFF);
+ }
+ if (newlinesInValues != null) {
+ b.setNewlinesInValues(newlinesInValues);
+ }
+ if (schemaInferMaxRecords != null) {
+ b.setSchemaInferMaxRecords(schemaInferMaxRecords);
+ }
+ return b.build().toByteArray();
}
Schema schema() {
return schema;
}
+
+ private static org.apache.datafusion.protobuf.FileCompressionType toProto(FileCompressionType t) {
+ switch (t) {
+ case UNCOMPRESSED:
+ return org.apache.datafusion.protobuf.FileCompressionType
+ .FILE_COMPRESSION_TYPE_UNCOMPRESSED;
+ case GZIP:
+ return org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_GZIP;
+ case BZIP2:
+ return org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_BZIP2;
+ case XZ:
+ return org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_XZ;
+ case ZSTD:
+ return org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_ZSTD;
+ default:
+ throw new IllegalArgumentException("unhandled FileCompressionType: " + t);
+ }
+ }
}
diff --git a/src/main/java/org/apache/datafusion/ParquetReadOptions.java b/src/main/java/org/apache/datafusion/ParquetReadOptions.java
index b29d762..6b47738 100644
--- a/src/main/java/org/apache/datafusion/ParquetReadOptions.java
+++ b/src/main/java/org/apache/datafusion/ParquetReadOptions.java
@@ -62,20 +62,20 @@
return this;
}
- String fileExtension() {
- return fileExtension;
- }
-
- Boolean parquetPruning() {
- return parquetPruning;
- }
-
- Boolean skipMetadata() {
- return skipMetadata;
- }
-
- Long metadataSizeHint() {
- return metadataSizeHint;
+ byte[] toBytes() {
+ org.apache.datafusion.protobuf.ParquetReadOptionsProto.Builder b =
+ org.apache.datafusion.protobuf.ParquetReadOptionsProto.newBuilder()
+ .setFileExtension(fileExtension);
+ if (parquetPruning != null) {
+ b.setParquetPruning(parquetPruning);
+ }
+ if (skipMetadata != null) {
+ b.setSkipMetadata(skipMetadata);
+ }
+ if (metadataSizeHint != null) {
+ b.setMetadataSizeHint(metadataSizeHint);
+ }
+ return b.build().toByteArray();
}
Schema schema() {
diff --git a/src/main/java/org/apache/datafusion/SessionContext.java b/src/main/java/org/apache/datafusion/SessionContext.java
index a496883..3cea058 100644
--- a/src/main/java/org/apache/datafusion/SessionContext.java
+++ b/src/main/java/org/apache/datafusion/SessionContext.java
@@ -133,20 +133,7 @@
nativeHandle,
name,
path,
- options.hasHeader(),
- options.delimiter(),
- options.quote(),
- options.terminator() != null,
- options.terminator() != null ? options.terminator() : 0,
- options.escape() != null,
- options.escape() != null ? options.escape() : 0,
- options.comment() != null,
- options.comment() != null ? options.comment() : 0,
- options.newlinesInValues() != null,
- options.newlinesInValues() != null && options.newlinesInValues(),
- options.schemaInferMaxRecords() != null ? options.schemaInferMaxRecords() : -1L,
- options.fileExtension(),
- options.fileCompressionType().name(),
+ options.toBytes(),
options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
}
@@ -168,20 +155,7 @@
readCsvWithOptions(
nativeHandle,
path,
- options.hasHeader(),
- options.delimiter(),
- options.quote(),
- options.terminator() != null,
- options.terminator() != null ? options.terminator() : 0,
- options.escape() != null,
- options.escape() != null ? options.escape() : 0,
- options.comment() != null,
- options.comment() != null ? options.comment() : 0,
- options.newlinesInValues() != null,
- options.newlinesInValues() != null && options.newlinesInValues(),
- options.schemaInferMaxRecords() != null ? options.schemaInferMaxRecords() : -1L,
- options.fileExtension(),
- options.fileCompressionType().name(),
+ options.toBytes(),
options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
return new DataFrame(dfHandle);
}
@@ -203,12 +177,7 @@
nativeHandle,
name,
path,
- options.fileExtension(),
- options.parquetPruning() != null,
- options.parquetPruning() != null && options.parquetPruning(),
- options.skipMetadata() != null,
- options.skipMetadata() != null && options.skipMetadata(),
- options.metadataSizeHint() != null ? options.metadataSizeHint() : -1L,
+ options.toBytes(),
options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
}
@@ -230,12 +199,7 @@
readParquetWithOptions(
nativeHandle,
path,
- options.fileExtension(),
- options.parquetPruning() != null,
- options.parquetPruning() != null && options.parquetPruning(),
- options.skipMetadata() != null,
- options.skipMetadata() != null && options.skipMetadata(),
- options.metadataSizeHint() != null ? options.metadataSizeHint() : -1L,
+ options.toBytes(),
options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
return new DataFrame(dfHandle);
}
@@ -272,66 +236,16 @@
private static native byte[] tableSchemaIpc(long handle, String tableName);
private static native void registerParquetWithOptions(
- long handle,
- String name,
- String path,
- String fileExtension,
- boolean parquetPruningSet,
- boolean parquetPruningValue,
- boolean skipMetadataSet,
- boolean skipMetadataValue,
- long metadataSizeHint,
- byte[] schemaIpcBytes);
+ long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);
private static native long readParquetWithOptions(
- long handle,
- String path,
- String fileExtension,
- boolean parquetPruningSet,
- boolean parquetPruningValue,
- boolean skipMetadataSet,
- boolean skipMetadataValue,
- long metadataSizeHint,
- byte[] schemaIpcBytes);
+ long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);
private static native void registerCsvWithOptions(
- long handle,
- String name,
- String path,
- boolean hasHeader,
- byte delimiter,
- byte quote,
- boolean terminatorSet,
- byte terminatorValue,
- boolean escapeSet,
- byte escapeValue,
- boolean commentSet,
- byte commentValue,
- boolean newlinesInValuesSet,
- boolean newlinesInValuesValue,
- long schemaInferMaxRecords,
- String fileExtension,
- String fileCompressionType,
- byte[] schemaIpcBytes);
+ long handle, String name, String path, byte[] optionsBytes, byte[] schemaIpcBytes);
private static native long readCsvWithOptions(
- long handle,
- String path,
- boolean hasHeader,
- byte delimiter,
- byte quote,
- boolean terminatorSet,
- byte terminatorValue,
- boolean escapeSet,
- byte escapeValue,
- boolean commentSet,
- byte commentValue,
- boolean newlinesInValuesSet,
- boolean newlinesInValuesValue,
- long schemaInferMaxRecords,
- String fileExtension,
- String fileCompressionType,
- byte[] schemaIpcBytes);
+ long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);
private static native void closeSessionContext(long handle);
}
diff --git a/src/main/java/org/apache/datafusion/SessionContextBuilder.java b/src/main/java/org/apache/datafusion/SessionContextBuilder.java
index bb0cac9..beabb5f 100644
--- a/src/main/java/org/apache/datafusion/SessionContextBuilder.java
+++ b/src/main/java/org/apache/datafusion/SessionContextBuilder.java
@@ -69,8 +69,7 @@
*/
public SessionContextBuilder memoryLimit(long maxMemoryBytes, double fraction) {
if (maxMemoryBytes <= 0) {
- throw new IllegalArgumentException(
- "maxMemoryBytes must be positive, got " + maxMemoryBytes);
+ 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);
diff --git a/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java b/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
index 6f49403..72f8e20 100644
--- a/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
+++ b/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
@@ -20,7 +20,7 @@
package org.apache.datafusion;
import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -30,31 +30,34 @@
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 org.apache.datafusion.protobuf.CsvReadOptionsProto;
+import org.apache.datafusion.protobuf.FileCompressionType;
import org.junit.jupiter.api.Test;
+import com.google.protobuf.InvalidProtocolBufferException;
+
class CsvReadOptionsTest {
@Test
- void defaultsMatchDataFusion() {
- CsvReadOptions opts = new CsvReadOptions();
- assertTrue(opts.hasHeader());
- assertEquals((byte) ',', opts.delimiter());
- assertEquals((byte) '"', opts.quote());
- assertNull(opts.terminator());
- assertNull(opts.escape());
- assertNull(opts.comment());
- assertNull(opts.newlinesInValues());
- assertNull(opts.schemaInferMaxRecords());
- assertEquals(".csv", opts.fileExtension());
- assertEquals(CsvReadOptions.FileCompressionType.UNCOMPRESSED, opts.fileCompressionType());
- assertNull(opts.schema());
+ void defaultsRoundTripThroughProto() throws InvalidProtocolBufferException {
+ CsvReadOptionsProto p = CsvReadOptionsProto.parseFrom(new CsvReadOptions().toBytes());
+
+ assertTrue(p.getHasHeader());
+ assertEquals((int) ',', p.getDelimiter());
+ assertEquals((int) '"', p.getQuote());
+ assertEquals(".csv", p.getFileExtension());
+ assertEquals(
+ FileCompressionType.FILE_COMPRESSION_TYPE_UNCOMPRESSED, p.getFileCompressionType());
+
+ assertFalse(p.hasTerminator());
+ assertFalse(p.hasEscape());
+ assertFalse(p.hasComment());
+ assertFalse(p.hasNewlinesInValues());
+ assertFalse(p.hasSchemaInferMaxRecords());
}
@Test
- void fluentSettersChainAndMutate() {
- Schema schema =
- new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(32, true)), null)));
-
+ void fullyConfiguredRoundTripsThroughProto() throws InvalidProtocolBufferException {
CsvReadOptions opts =
new CsvReadOptions()
.hasHeader(false)
@@ -66,19 +69,40 @@
.newlinesInValues(true)
.schemaInferMaxRecords(10L)
.fileExtension(".tsv")
- .fileCompressionType(CsvReadOptions.FileCompressionType.GZIP)
- .schema(schema);
+ .fileCompressionType(CsvReadOptions.FileCompressionType.GZIP);
- assertEquals(false, opts.hasHeader());
- assertEquals((byte) '|', opts.delimiter());
- assertEquals((byte) '\'', opts.quote());
- assertEquals(Byte.valueOf((byte) '\n'), opts.terminator());
- assertEquals(Byte.valueOf((byte) '\\'), opts.escape());
- assertEquals(Byte.valueOf((byte) '#'), opts.comment());
- assertEquals(Boolean.TRUE, opts.newlinesInValues());
- assertEquals(Long.valueOf(10L), opts.schemaInferMaxRecords());
- assertEquals(".tsv", opts.fileExtension());
- assertEquals(CsvReadOptions.FileCompressionType.GZIP, opts.fileCompressionType());
+ CsvReadOptionsProto p = CsvReadOptionsProto.parseFrom(opts.toBytes());
+
+ assertFalse(p.getHasHeader());
+ assertEquals((int) '|', p.getDelimiter());
+ assertEquals((int) '\'', p.getQuote());
+ assertEquals((int) '\n', p.getTerminator());
+ assertEquals((int) '\\', p.getEscape());
+ assertEquals((int) '#', p.getComment());
+ assertTrue(p.getNewlinesInValues());
+ assertEquals(10L, p.getSchemaInferMaxRecords());
+ assertEquals(".tsv", p.getFileExtension());
+ assertEquals(FileCompressionType.FILE_COMPRESSION_TYPE_GZIP, p.getFileCompressionType());
+ }
+
+ @Test
+ void schemaIsHeldByReferenceAndNotInProto() {
+ Schema schema =
+ new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(32, true)), null)));
+ CsvReadOptions opts = new CsvReadOptions().schema(schema);
+
assertSame(schema, opts.schema());
}
+
+ @Test
+ void allCompressionTypesMapThroughProto() throws InvalidProtocolBufferException {
+ for (CsvReadOptions.FileCompressionType t : CsvReadOptions.FileCompressionType.values()) {
+ CsvReadOptionsProto p =
+ CsvReadOptionsProto.parseFrom(new CsvReadOptions().fileCompressionType(t).toBytes());
+ assertEquals(
+ "FILE_COMPRESSION_TYPE_" + t.name(),
+ p.getFileCompressionType().name(),
+ "mismatch for " + t);
+ }
+ }
}
diff --git a/src/test/java/org/apache/datafusion/ParquetReadOptionsTest.java b/src/test/java/org/apache/datafusion/ParquetReadOptionsTest.java
index 3eca311..634a953 100644
--- a/src/test/java/org/apache/datafusion/ParquetReadOptionsTest.java
+++ b/src/test/java/org/apache/datafusion/ParquetReadOptionsTest.java
@@ -20,7 +20,7 @@
package org.apache.datafusion;
import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertSame;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -30,44 +30,47 @@
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 org.apache.datafusion.protobuf.ParquetReadOptionsProto;
import org.junit.jupiter.api.Test;
+import com.google.protobuf.InvalidProtocolBufferException;
+
class ParquetReadOptionsTest {
@Test
- void defaultsMatchDataFusion() {
- ParquetReadOptions opts = new ParquetReadOptions();
- assertEquals(".parquet", opts.fileExtension());
- assertNull(opts.parquetPruning());
- assertNull(opts.skipMetadata());
- assertNull(opts.metadataSizeHint());
- assertNull(opts.schema());
+ void defaultsRoundTripThroughProto() throws InvalidProtocolBufferException {
+ ParquetReadOptionsProto p =
+ ParquetReadOptionsProto.parseFrom(new ParquetReadOptions().toBytes());
+
+ assertEquals(".parquet", p.getFileExtension());
+ assertFalse(p.hasParquetPruning());
+ assertFalse(p.hasSkipMetadata());
+ assertFalse(p.hasMetadataSizeHint());
}
@Test
- void fluentSettersChainAndMutate() {
- Schema schema =
- new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(32, true)), null)));
-
+ void fullyConfiguredRoundTripsThroughProto() throws InvalidProtocolBufferException {
ParquetReadOptions opts =
new ParquetReadOptions()
- .fileExtension(".parq")
+ .fileExtension(".par")
.parquetPruning(true)
.skipMetadata(false)
- .metadataSizeHint(1_048_576L)
- .schema(schema);
+ .metadataSizeHint(4096L);
- assertEquals(".parq", opts.fileExtension());
- assertEquals(Boolean.TRUE, opts.parquetPruning());
- assertEquals(Boolean.FALSE, opts.skipMetadata());
- assertEquals(Long.valueOf(1_048_576L), opts.metadataSizeHint());
- assertTrue(opts.schema() == schema);
+ ParquetReadOptionsProto p = ParquetReadOptionsProto.parseFrom(opts.toBytes());
+
+ assertEquals(".par", p.getFileExtension());
+ assertTrue(p.getParquetPruning());
+ assertFalse(p.getSkipMetadata());
+ assertEquals(4096L, p.getMetadataSizeHint());
}
@Test
- void schemaSetterRetainsReferenceIdentity() {
- Schema schema = new Schema(List.of());
+ void schemaIsHeldByReferenceAndNotInProto() {
+ Schema schema =
+ new Schema(List.of(new Field("x", FieldType.nullable(new ArrowType.Int(32, true)), null)));
ParquetReadOptions opts = new ParquetReadOptions().schema(schema);
+
assertSame(schema, opts.schema());
}
}