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 63a86a5 feat(json): expose NdJsonReadOptions via registerJson and
readJson (#47)
63a86a5 is described below
commit 63a86a588ab274a796784373d5d5519b9953b71c
Author: Lantao Jin <[email protected]>
AuthorDate: Mon May 18 05:00:28 2026 +0800
feat(json): expose NdJsonReadOptions via registerJson and readJson (#47)
---
.../java/org/apache/datafusion/CsvReadOptions.java | 32 +----
.../org/apache/datafusion/FileCompressionType.java | 33 +++++
.../apache/datafusion/FileCompressionTypes.java | 47 +++++++
.../org/apache/datafusion/NdJsonReadOptions.java | 78 +++++++++++
.../java/org/apache/datafusion/SessionContext.java | 67 ++++++++++
.../org/apache/datafusion/CsvReadOptionsTest.java | 23 +++-
...OptionsTest.java => NdJsonReadOptionsTest.java} | 75 +++++------
.../apache/datafusion/SessionContextJsonTest.java | 148 +++++++++++++++++++++
native/build.rs | 2 +
native/src/json.rs | 116 ++++++++++++++++
native/src/lib.rs | 1 +
proto/csv_read_options.proto | 11 +-
...d_options.proto => file_compression_type.proto} | 21 +--
native/build.rs => proto/json_read_options.proto | 31 +++--
14 files changed, 573 insertions(+), 112 deletions(-)
diff --git a/core/src/main/java/org/apache/datafusion/CsvReadOptions.java
b/core/src/main/java/org/apache/datafusion/CsvReadOptions.java
index f441820..ce62afe 100644
--- a/core/src/main/java/org/apache/datafusion/CsvReadOptions.java
+++ b/core/src/main/java/org/apache/datafusion/CsvReadOptions.java
@@ -33,15 +33,6 @@ import org.apache.arrow.vector.types.pojo.Schema;
*/
public final class CsvReadOptions {
- /** Compression of the file. Names match DataFusion's {@code
FileCompressionType} variants. */
- public enum FileCompressionType {
- UNCOMPRESSED,
- GZIP,
- BZIP2,
- XZ,
- ZSTD
- }
-
private boolean hasHeader = true;
private byte delimiter = (byte) ',';
private byte quote = (byte) '"';
@@ -90,6 +81,9 @@ public final class CsvReadOptions {
}
public CsvReadOptions schemaInferMaxRecords(long n) {
+ if (n < 0) {
+ throw new IllegalArgumentException("schemaInferMaxRecords must be
non-negative, got " + n);
+ }
this.schemaInferMaxRecords = n;
return this;
}
@@ -116,7 +110,7 @@ public final class CsvReadOptions {
.setDelimiter(delimiter & 0xFF)
.setQuote(quote & 0xFF)
.setFileExtension(fileExtension)
- .setFileCompressionType(toProto(fileCompressionType));
+
.setFileCompressionType(FileCompressionTypes.toProto(fileCompressionType));
if (terminator != null) {
b.setTerminator(terminator & 0xFF);
}
@@ -138,22 +132,4 @@ public final class CsvReadOptions {
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/core/src/main/java/org/apache/datafusion/FileCompressionType.java
b/core/src/main/java/org/apache/datafusion/FileCompressionType.java
new file mode 100644
index 0000000..cae1a9c
--- /dev/null
+++ b/core/src/main/java/org/apache/datafusion/FileCompressionType.java
@@ -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.
+ */
+
+package org.apache.datafusion;
+
+/**
+ * Compression of a file-format source. Shared by {@link CsvReadOptions} and
{@link
+ * NdJsonReadOptions} (and any future format that exposes the same set of
compressions). Variant
+ * names match DataFusion's Rust {@code FileCompressionType} so they
round-trip across JNI.
+ */
+public enum FileCompressionType {
+ UNCOMPRESSED,
+ GZIP,
+ BZIP2,
+ XZ,
+ ZSTD
+}
diff --git a/core/src/main/java/org/apache/datafusion/FileCompressionTypes.java
b/core/src/main/java/org/apache/datafusion/FileCompressionTypes.java
new file mode 100644
index 0000000..2e7ea56
--- /dev/null
+++ b/core/src/main/java/org/apache/datafusion/FileCompressionTypes.java
@@ -0,0 +1,47 @@
+/*
+ * 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;
+
+/**
+ * Internal helpers for translating between the user-facing {@link
FileCompressionType} enum and the
+ * protobuf-generated {@code
org.apache.datafusion.protobuf.FileCompressionType} enum used on the
+ * wire. Same variant set, different Java types.
+ */
+final class FileCompressionTypes {
+ private FileCompressionTypes() {}
+
+ 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/core/src/main/java/org/apache/datafusion/NdJsonReadOptions.java
b/core/src/main/java/org/apache/datafusion/NdJsonReadOptions.java
new file mode 100644
index 0000000..1fa1bf1
--- /dev/null
+++ b/core/src/main/java/org/apache/datafusion/NdJsonReadOptions.java
@@ -0,0 +1,78 @@
+/*
+ * 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.arrow.vector.types.pojo.Schema;
+
+/**
+ * Configuration knobs for newline-delimited JSON sources passed to {@link
+ * SessionContext#registerJson(String, String, NdJsonReadOptions)} and {@link
+ * SessionContext#readJson(String, NdJsonReadOptions)}.
+ *
+ * <p>Mirrors a subset of DataFusion's {@code NdJsonReadOptions}. All setters
return {@code this}
+ * for fluent chaining. Defaults match the Rust struct: {@code fileExtension =
".json"}, {@code
+ * fileCompressionType = UNCOMPRESSED}; {@code schemaInferMaxRecords} unset
(the DataFusion default
+ * is used).
+ */
+public final class NdJsonReadOptions {
+
+ private String fileExtension = ".json";
+ private FileCompressionType fileCompressionType =
FileCompressionType.UNCOMPRESSED;
+ private Long schemaInferMaxRecords;
+ private Schema schema;
+
+ public NdJsonReadOptions fileExtension(String ext) {
+ this.fileExtension = ext;
+ return this;
+ }
+
+ public NdJsonReadOptions fileCompressionType(FileCompressionType t) {
+ this.fileCompressionType = t;
+ return this;
+ }
+
+ public NdJsonReadOptions schemaInferMaxRecords(long n) {
+ if (n < 0) {
+ throw new IllegalArgumentException("schemaInferMaxRecords must be
non-negative, got " + n);
+ }
+ this.schemaInferMaxRecords = n;
+ return this;
+ }
+
+ public NdJsonReadOptions schema(Schema schema) {
+ this.schema = schema;
+ return this;
+ }
+
+ byte[] toBytes() {
+ org.apache.datafusion.protobuf.NdJsonReadOptionsProto.Builder b =
+ org.apache.datafusion.protobuf.NdJsonReadOptionsProto.newBuilder()
+ .setFileExtension(fileExtension)
+
.setFileCompressionType(FileCompressionTypes.toProto(fileCompressionType));
+ if (schemaInferMaxRecords != null) {
+ b.setSchemaInferMaxRecords(schemaInferMaxRecords);
+ }
+ return b.build().toByteArray();
+ }
+
+ Schema schema() {
+ return schema;
+ }
+}
diff --git a/core/src/main/java/org/apache/datafusion/SessionContext.java
b/core/src/main/java/org/apache/datafusion/SessionContext.java
index 1b2a075..51384a6 100644
--- a/core/src/main/java/org/apache/datafusion/SessionContext.java
+++ b/core/src/main/java/org/apache/datafusion/SessionContext.java
@@ -188,6 +188,67 @@ public final class SessionContext implements AutoCloseable
{
return new DataFrame(dfHandle);
}
+ public void registerJson(String name, String path) {
+ registerJson(name, path, new NdJsonReadOptions());
+ }
+
+ /**
+ * Register a newline-delimited JSON file (or directory of NDJSON files) as
a table with the
+ * supplied {@link NdJsonReadOptions}.
+ *
+ * @throws RuntimeException if registration fails (path not found, schema
inference error, etc.).
+ */
+ public void registerJson(String name, String path, NdJsonReadOptions
options) {
+ if (nativeHandle == 0) {
+ throw new IllegalStateException("SessionContext is closed");
+ }
+ if (name == null) {
+ throw new IllegalArgumentException("registerJson name must be non-null");
+ }
+ if (path == null) {
+ throw new IllegalArgumentException("registerJson path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("registerJson options must be
non-null");
+ }
+ registerJsonWithOptions(
+ nativeHandle,
+ name,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) :
null);
+ }
+
+ /** Read a newline-delimited JSON file as a {@link DataFrame} without
registering it. */
+ public DataFrame readJson(String path) {
+ return readJson(path, new NdJsonReadOptions());
+ }
+
+ /**
+ * Read a newline-delimited JSON file as a {@link DataFrame} with the
supplied {@link
+ * NdJsonReadOptions}.
+ *
+ * @throws RuntimeException if the read fails.
+ */
+ public DataFrame readJson(String path, NdJsonReadOptions options) {
+ if (nativeHandle == 0) {
+ throw new IllegalStateException("SessionContext is closed");
+ }
+ if (path == null) {
+ throw new IllegalArgumentException("readJson path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("readJson options must be non-null");
+ }
+ long dfHandle =
+ readJsonWithOptions(
+ nativeHandle,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) :
null);
+ return new DataFrame(dfHandle);
+ }
+
public void registerParquet(String name, String path) {
registerParquet(name, path, new ParquetReadOptions());
}
@@ -277,5 +338,11 @@ public final class SessionContext implements AutoCloseable
{
private static native long readCsvWithOptions(
long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);
+ private static native void registerJsonWithOptions(
+ long handle, String name, String path, byte[] optionsBytes, byte[]
schemaIpcBytes);
+
+ private static native long readJsonWithOptions(
+ long handle, String path, byte[] optionsBytes, byte[] schemaIpcBytes);
+
private static native void closeSessionContext(long handle);
}
diff --git a/core/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
b/core/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
index 72f8e20..2e35522 100644
--- a/core/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
+++ b/core/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
@@ -22,6 +22,7 @@ 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.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.util.List;
@@ -31,7 +32,6 @@ 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;
@@ -47,7 +47,8 @@ class CsvReadOptionsTest {
assertEquals((int) '"', p.getQuote());
assertEquals(".csv", p.getFileExtension());
assertEquals(
- FileCompressionType.FILE_COMPRESSION_TYPE_UNCOMPRESSED,
p.getFileCompressionType());
+
org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_UNCOMPRESSED,
+ p.getFileCompressionType());
assertFalse(p.hasTerminator());
assertFalse(p.hasEscape());
@@ -69,7 +70,7 @@ class CsvReadOptionsTest {
.newlinesInValues(true)
.schemaInferMaxRecords(10L)
.fileExtension(".tsv")
- .fileCompressionType(CsvReadOptions.FileCompressionType.GZIP);
+ .fileCompressionType(FileCompressionType.GZIP);
CsvReadOptionsProto p = CsvReadOptionsProto.parseFrom(opts.toBytes());
@@ -82,7 +83,9 @@ class CsvReadOptionsTest {
assertTrue(p.getNewlinesInValues());
assertEquals(10L, p.getSchemaInferMaxRecords());
assertEquals(".tsv", p.getFileExtension());
- assertEquals(FileCompressionType.FILE_COMPRESSION_TYPE_GZIP,
p.getFileCompressionType());
+ assertEquals(
+
org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_GZIP,
+ p.getFileCompressionType());
}
@Test
@@ -96,7 +99,7 @@ class CsvReadOptionsTest {
@Test
void allCompressionTypesMapThroughProto() throws
InvalidProtocolBufferException {
- for (CsvReadOptions.FileCompressionType t :
CsvReadOptions.FileCompressionType.values()) {
+ for (FileCompressionType t : FileCompressionType.values()) {
CsvReadOptionsProto p =
CsvReadOptionsProto.parseFrom(new
CsvReadOptions().fileCompressionType(t).toBytes());
assertEquals(
@@ -105,4 +108,14 @@ class CsvReadOptionsTest {
"mismatch for " + t);
}
}
+
+ @Test
+ void schemaInferMaxRecordsRejectsNegative() {
+ // The proto wire field is uint64, so a negative long would be
reinterpreted
+ // as a huge unsigned value on the Rust side and silently expand schema
+ // inference across the full dataset. Reject at the Java setter instead.
+ CsvReadOptions opts = new CsvReadOptions();
+ assertThrows(IllegalArgumentException.class, () ->
opts.schemaInferMaxRecords(-1L));
+ assertThrows(IllegalArgumentException.class, () ->
opts.schemaInferMaxRecords(Long.MIN_VALUE));
+ }
}
diff --git a/core/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
b/core/src/test/java/org/apache/datafusion/NdJsonReadOptionsTest.java
similarity index 50%
copy from core/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
copy to core/src/test/java/org/apache/datafusion/NdJsonReadOptionsTest.java
index 72f8e20..f55551e 100644
--- a/core/src/test/java/org/apache/datafusion/CsvReadOptionsTest.java
+++ b/core/src/test/java/org/apache/datafusion/NdJsonReadOptionsTest.java
@@ -22,7 +22,7 @@ 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.assertSame;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import java.util.List;
@@ -30,79 +30,70 @@ import org.apache.arrow.vector.types.pojo.ArrowType;
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.apache.datafusion.protobuf.NdJsonReadOptionsProto;
import org.junit.jupiter.api.Test;
import com.google.protobuf.InvalidProtocolBufferException;
-class CsvReadOptionsTest {
+class NdJsonReadOptionsTest {
@Test
void defaultsRoundTripThroughProto() throws InvalidProtocolBufferException {
- CsvReadOptionsProto p = CsvReadOptionsProto.parseFrom(new
CsvReadOptions().toBytes());
+ NdJsonReadOptionsProto p = NdJsonReadOptionsProto.parseFrom(new
NdJsonReadOptions().toBytes());
- assertTrue(p.getHasHeader());
- assertEquals((int) ',', p.getDelimiter());
- assertEquals((int) '"', p.getQuote());
- assertEquals(".csv", p.getFileExtension());
+ assertEquals(".json", p.getFileExtension());
assertEquals(
- FileCompressionType.FILE_COMPRESSION_TYPE_UNCOMPRESSED,
p.getFileCompressionType());
-
- assertFalse(p.hasTerminator());
- assertFalse(p.hasEscape());
- assertFalse(p.hasComment());
- assertFalse(p.hasNewlinesInValues());
+
org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_UNCOMPRESSED,
+ p.getFileCompressionType());
assertFalse(p.hasSchemaInferMaxRecords());
}
@Test
void fullyConfiguredRoundTripsThroughProto() throws
InvalidProtocolBufferException {
- CsvReadOptions opts =
- new CsvReadOptions()
- .hasHeader(false)
- .delimiter((byte) '|')
- .quote((byte) '\'')
- .terminator((byte) '\n')
- .escape((byte) '\\')
- .comment((byte) '#')
- .newlinesInValues(true)
- .schemaInferMaxRecords(10L)
- .fileExtension(".tsv")
- .fileCompressionType(CsvReadOptions.FileCompressionType.GZIP);
+ NdJsonReadOptions opts =
+ new NdJsonReadOptions()
+ .fileExtension(".ndjson")
+ .fileCompressionType(FileCompressionType.GZIP)
+ .schemaInferMaxRecords(50L);
- CsvReadOptionsProto p = CsvReadOptionsProto.parseFrom(opts.toBytes());
+ NdJsonReadOptionsProto p =
NdJsonReadOptionsProto.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());
+ assertEquals(".ndjson", p.getFileExtension());
+ assertEquals(
+
org.apache.datafusion.protobuf.FileCompressionType.FILE_COMPRESSION_TYPE_GZIP,
+ p.getFileCompressionType());
+ assertEquals(50L, p.getSchemaInferMaxRecords());
}
@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);
+ NdJsonReadOptions opts = new NdJsonReadOptions().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());
+ for (FileCompressionType t : FileCompressionType.values()) {
+ NdJsonReadOptionsProto p =
+ NdJsonReadOptionsProto.parseFrom(
+ new NdJsonReadOptions().fileCompressionType(t).toBytes());
assertEquals(
"FILE_COMPRESSION_TYPE_" + t.name(),
p.getFileCompressionType().name(),
"mismatch for " + t);
}
}
+
+ @Test
+ void schemaInferMaxRecordsRejectsNegative() {
+ // The proto wire field is uint64, so a negative long would be
reinterpreted
+ // as a huge unsigned value on the Rust side and silently expand schema
+ // inference across the full dataset. Reject at the Java setter instead.
+ NdJsonReadOptions opts = new NdJsonReadOptions();
+ assertThrows(IllegalArgumentException.class, () ->
opts.schemaInferMaxRecords(-1L));
+ assertThrows(IllegalArgumentException.class, () ->
opts.schemaInferMaxRecords(Long.MIN_VALUE));
+ }
}
diff --git
a/core/src/test/java/org/apache/datafusion/SessionContextJsonTest.java
b/core/src/test/java/org/apache/datafusion/SessionContextJsonTest.java
new file mode 100644
index 0000000..20c57b1
--- /dev/null
+++ b/core/src/test/java/org/apache/datafusion/SessionContextJsonTest.java
@@ -0,0 +1,148 @@
+/*
+ * 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.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.ipc.ArrowReader;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+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.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+class SessionContextJsonTest {
+
+ private static Path writeJson(Path dir, String name, String contents) throws
IOException {
+ Path file = dir.resolve(name);
+ Files.writeString(file, contents);
+ return file;
+ }
+
+ @Test
+ void registerJsonInfersSchemaAndCounts(@TempDir Path tempDir) throws
Exception {
+ Path file =
+ writeJson(
+ tempDir,
+ "people.json",
+ "{\"id\":1,\"name\":\"alice\"}\n"
+ + "{\"id\":2,\"name\":\"bob\"}\n"
+ + "{\"id\":3,\"name\":\"carol\"}\n");
+
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext()) {
+ ctx.registerJson("people", file.toAbsolutePath().toString());
+
+ try (DataFrame df = ctx.sql("SELECT COUNT(*) FROM people");
+ ArrowReader reader = df.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ BigIntVector count = (BigIntVector)
reader.getVectorSchemaRoot().getVector(0);
+ assertEquals(3L, count.get(0));
+ }
+
+ try (DataFrame df = ctx.sql("SELECT name FROM people WHERE id = 2");
+ ArrowReader reader = df.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ assertEquals(1, root.getRowCount());
+ VarCharVector names = (VarCharVector) root.getVector(0);
+ assertEquals("bob", new String(names.get(0)));
+ }
+ }
+ }
+
+ @Test
+ void readJsonWithExplicitSchema(@TempDir Path tempDir) throws Exception {
+ Path file =
+ writeJson(
+ tempDir, "headerless.json",
"{\"id\":10,\"name\":\"x\"}\n{\"id\":20,\"name\":\"y\"}\n");
+
+ Schema schema =
+ new Schema(
+ List.of(
+ new Field("id", FieldType.nullable(new ArrowType.Int(64,
true)), null),
+ new Field("name", FieldType.nullable(new ArrowType.Utf8()),
null)));
+
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext();
+ DataFrame df =
+ ctx.readJson(file.toAbsolutePath().toString(), new
NdJsonReadOptions().schema(schema));
+ ArrowReader reader = df.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ assertEquals(2, root.getRowCount());
+ assertEquals("id", root.getSchema().getFields().get(0).getName());
+ assertEquals("name", root.getSchema().getFields().get(1).getName());
+ }
+ }
+
+ @Test
+ void registerJsonWithCustomExtension(@TempDir Path tempDir) throws Exception
{
+ Path file = writeJson(tempDir, "data.ndjson",
"{\"x\":10,\"y\":20}\n{\"x\":30,\"y\":40}\n");
+
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext()) {
+ ctx.registerJson(
+ "t", file.toAbsolutePath().toString(), new
NdJsonReadOptions().fileExtension(".ndjson"));
+
+ try (DataFrame df = ctx.sql("SELECT SUM(x) + SUM(y) FROM t");
+ ArrowReader reader = df.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ BigIntVector v = (BigIntVector)
reader.getVectorSchemaRoot().getVector(0);
+ assertEquals(100L, v.get(0));
+ }
+ }
+ }
+
+ @Test
+ void registerJsonRejectsNullArguments() {
+ try (SessionContext ctx = new SessionContext()) {
+ NdJsonReadOptions opts = new NdJsonReadOptions();
+ assertThrows(IllegalArgumentException.class, () ->
ctx.registerJson(null, "/p"));
+ assertThrows(IllegalArgumentException.class, () -> ctx.registerJson("t",
null));
+ assertThrows(IllegalArgumentException.class, () ->
ctx.registerJson(null, "/p", opts));
+ assertThrows(IllegalArgumentException.class, () -> ctx.registerJson("t",
null, opts));
+ assertThrows(IllegalArgumentException.class, () -> ctx.registerJson("t",
"/p", null));
+ }
+ }
+
+ @Test
+ void readJsonRejectsNullArguments() {
+ try (SessionContext ctx = new SessionContext()) {
+ NdJsonReadOptions opts = new NdJsonReadOptions();
+ assertThrows(IllegalArgumentException.class, () -> ctx.readJson(null));
+ assertThrows(IllegalArgumentException.class, () -> ctx.readJson(null,
opts));
+ assertThrows(IllegalArgumentException.class, () -> ctx.readJson("/p",
null));
+ }
+ }
+}
diff --git a/native/build.rs b/native/build.rs
index 5a27cb0..7c164d5 100644
--- a/native/build.rs
+++ b/native/build.rs
@@ -18,7 +18,9 @@
fn main() {
const PROTOS: &[&str] = &[
"../proto/session_options.proto",
+ "../proto/file_compression_type.proto",
"../proto/csv_read_options.proto",
+ "../proto/json_read_options.proto",
"../proto/parquet_read_options.proto",
];
for p in PROTOS {
diff --git a/native/src/json.rs b/native/src/json.rs
new file mode 100644
index 0000000..43f0055
--- /dev/null
+++ b/native/src/json.rs
@@ -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.
+
+use
datafusion::datasource::file_format::file_compression_type::FileCompressionType;
+use datafusion::error::DataFusionError;
+use datafusion::prelude::JsonReadOptions;
+use datafusion::prelude::SessionContext;
+use jni::objects::{JByteArray, JClass, JString};
+use jni::sys::jlong;
+use jni::JNIEnv;
+use prost::Message;
+
+use crate::errors::{try_unwrap_or_throw, JniResult};
+use crate::proto_gen::{FileCompressionType as ProtoFileCompressionType,
NdJsonReadOptionsProto};
+use crate::runtime;
+use crate::schema::decode_optional_schema;
+
+fn with_json_options<R>(
+ env: &mut JNIEnv,
+ options_bytes: JByteArray,
+ schema_ipc_bytes: JByteArray,
+ f: impl FnOnce(JsonReadOptions) -> JniResult<R>,
+) -> JniResult<R> {
+ let bytes: Vec<u8> = env.convert_byte_array(&options_bytes)?;
+ let p = NdJsonReadOptionsProto::decode(bytes.as_slice())?;
+
+ let schema = decode_optional_schema(env, schema_ipc_bytes)?;
+
+ let compression = match p.file_compression_type() {
+ ProtoFileCompressionType::Unspecified => {
+ return Err("NdJsonReadOptionsProto.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 = JsonReadOptions::default()
+ .file_extension(&file_ext)
+ .file_compression_type(compression);
+
+ 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);
+ }
+
+ f(opts)
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_datafusion_SessionContext_registerJsonWithOptions<'local>(
+ mut env: JNIEnv<'local>,
+ _class: JClass<'local>,
+ handle: jlong,
+ name: JString<'local>,
+ path: JString<'local>,
+ options_bytes: JByteArray<'local>,
+ schema_ipc_bytes: JByteArray<'local>,
+) {
+ try_unwrap_or_throw(&mut env, (), |env| -> JniResult<()> {
+ 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 path: String = env.get_string(&path)?.into();
+ with_json_options(env, options_bytes, schema_ipc_bytes, |opts| {
+ runtime().block_on(async {
+ ctx.register_json(&name, &path, opts).await?;
+ Ok::<(), DataFusionError>(())
+ })?;
+ Ok(())
+ })
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn
Java_org_apache_datafusion_SessionContext_readJsonWithOptions<'local>(
+ mut env: JNIEnv<'local>,
+ _class: JClass<'local>,
+ handle: jlong,
+ path: JString<'local>,
+ options_bytes: JByteArray<'local>,
+ schema_ipc_bytes: JByteArray<'local>,
+) -> jlong {
+ try_unwrap_or_throw(&mut env, 0, |env| -> JniResult<jlong> {
+ if handle == 0 {
+ return Err("SessionContext handle is null".into());
+ }
+ let ctx = unsafe { &*(handle as *const SessionContext) };
+ let path: String = env.get_string(&path)?.into();
+ with_json_options(env, options_bytes, schema_ipc_bytes, |opts| {
+ let df = runtime().block_on(ctx.read_json(path, opts))?;
+ Ok(Box::into_raw(Box::new(df)) as jlong)
+ })
+ })
+}
diff --git a/native/src/lib.rs b/native/src/lib.rs
index 9041819..efd7801 100644
--- a/native/src/lib.rs
+++ b/native/src/lib.rs
@@ -17,6 +17,7 @@
mod csv;
mod errors;
+mod json;
mod proto;
mod schema;
diff --git a/proto/csv_read_options.proto b/proto/csv_read_options.proto
index 840867d..378c92a 100644
--- a/proto/csv_read_options.proto
+++ b/proto/csv_read_options.proto
@@ -19,6 +19,8 @@ syntax = "proto3";
package datafusion_java;
+import "file_compression_type.proto";
+
option java_package = "org.apache.datafusion.protobuf";
option java_multiple_files = true;
@@ -37,12 +39,3 @@ message CsvReadOptionsProto {
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/csv_read_options.proto b/proto/file_compression_type.proto
similarity index 65%
copy from proto/csv_read_options.proto
copy to proto/file_compression_type.proto
index 840867d..e02593f 100644
--- a/proto/csv_read_options.proto
+++ b/proto/file_compression_type.proto
@@ -22,22 +22,11 @@ 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;
-}
-
+// Compression of a file-format source. Shared by `csv_read_options.proto`
+// and `json_read_options.proto` (and any future format that exposes the
+// same set of compressions). Variants and tag numbers must match
+// DataFusion's `FileCompressionType` so the JNI dispatch can map values
+// directly without a translation table.
enum FileCompressionType {
FILE_COMPRESSION_TYPE_UNSPECIFIED = 0;
FILE_COMPRESSION_TYPE_UNCOMPRESSED = 1;
diff --git a/native/build.rs b/proto/json_read_options.proto
similarity index 52%
copy from native/build.rs
copy to proto/json_read_options.proto
index 5a27cb0..4a64f4d 100644
--- a/native/build.rs
+++ b/proto/json_read_options.proto
@@ -15,16 +15,23 @@
// specific language governing permissions and limitations
// under the License.
-fn main() {
- 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(PROTOS, &["../proto"]).expect("failed to
compile protos");
+syntax = "proto3";
+
+package datafusion_java;
+
+import "file_compression_type.proto";
+
+option java_package = "org.apache.datafusion.protobuf";
+option java_multiple_files = true;
+
+// Options used to read newline-delimited JSON (NDJSON) 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. The shared
+// `FileCompressionType` enum lives in its own proto file so both this and
+// csv_read_options.proto can reference it without one file depending on the
+// other.
+message NdJsonReadOptionsProto {
+ string file_extension = 1;
+ FileCompressionType file_compression_type = 2;
+ optional uint64 schema_infer_max_records = 3;
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
