feat(write): add DataFrame.writeParquet with ParquetWriteOptions (#27)
diff --git a/native/src/lib.rs b/native/src/lib.rs
index fc3a983..4ea1a61 100644
--- a/native/src/lib.rs
+++ b/native/src/lib.rs
@@ -25,7 +25,9 @@
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;
+use datafusion::dataframe::DataFrameWriteOptions;
use datafusion::error::DataFusionError;
use datafusion::prelude::{ParquetReadOptions, SessionContext};
use jni::objects::{JByteArray, JClass, JObjectArray, JString};
@@ -196,6 +198,42 @@
}
#[no_mangle]
+pub extern "system" fn Java_org_apache_datafusion_DataFrame_writeParquetWithOptions<'local>(
+ mut env: JNIEnv<'local>,
+ _class: JClass<'local>,
+ handle: jlong,
+ path: JString<'local>,
+ compression: JString<'local>,
+ single_file_output_set: jboolean,
+ single_file_output_value: jboolean,
+) {
+ try_unwrap_or_throw(&mut env, (), |env| -> JniResult<()> {
+ if handle == 0 {
+ return Err("DataFrame handle is null".into());
+ }
+ let df = unsafe { &*(handle as *const DataFrame) }.clone();
+ let path: String = env.get_string(&path)?.into();
+
+ let mut write_opts = DataFrameWriteOptions::new();
+ if single_file_output_set != 0 {
+ write_opts = write_opts.with_single_file_output(single_file_output_value != 0);
+ }
+
+ let writer_opts: Option<TableParquetOptions> = if !compression.is_null() {
+ let c: String = env.get_string(&compression)?.into();
+ let mut tpo = TableParquetOptions::default();
+ tpo.global.compression = Some(c);
+ Some(tpo)
+ } else {
+ None
+ };
+
+ runtime().block_on(df.write_parquet(&path, write_opts, writer_opts))?;
+ Ok(())
+ })
+}
+
+#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_DataFrame_closeDataFrame<'local>(
mut env: JNIEnv<'local>,
_class: JClass<'local>,
diff --git a/src/main/java/org/apache/datafusion/DataFrame.java b/src/main/java/org/apache/datafusion/DataFrame.java
index 0bd77e7..244f6db 100644
--- a/src/main/java/org/apache/datafusion/DataFrame.java
+++ b/src/main/java/org/apache/datafusion/DataFrame.java
@@ -116,6 +116,36 @@
return new DataFrame(filterRows(nativeHandle, predicate));
}
+ /**
+ * Materialize this DataFrame as Parquet at {@code path}. The path is treated as a directory
+ * unless overridden via {@link ParquetWriteOptions#singleFileOutput(boolean)}. The receiver
+ * remains usable and must still be closed independently.
+ *
+ * @throws RuntimeException if the write fails.
+ */
+ public void writeParquet(String path) {
+ writeParquet(path, new ParquetWriteOptions());
+ }
+
+ /**
+ * Materialize this DataFrame as Parquet at {@code path} with the supplied {@link
+ * ParquetWriteOptions}. The receiver remains usable and must still be closed independently.
+ *
+ * @throws RuntimeException if the write fails (path inaccessible, invalid compression spec,
+ * etc.).
+ */
+ public void writeParquet(String path, ParquetWriteOptions options) {
+ if (nativeHandle == 0) {
+ throw new IllegalStateException("DataFrame is closed or already collected");
+ }
+ writeParquetWithOptions(
+ nativeHandle,
+ path,
+ options.compression(),
+ options.singleFileOutput() != null,
+ options.singleFileOutput() != null && options.singleFileOutput());
+ }
+
@Override
public void close() {
if (nativeHandle != 0) {
@@ -137,4 +167,11 @@
private static native long selectColumns(long handle, String[] columnNames);
private static native long filterRows(long handle, String predicate);
+
+ private static native void writeParquetWithOptions(
+ long handle,
+ String path,
+ String compression,
+ boolean singleFileOutputSet,
+ boolean singleFileOutputValue);
}
diff --git a/src/main/java/org/apache/datafusion/ParquetWriteOptions.java b/src/main/java/org/apache/datafusion/ParquetWriteOptions.java
new file mode 100644
index 0000000..1aa0d85
--- /dev/null
+++ b/src/main/java/org/apache/datafusion/ParquetWriteOptions.java
@@ -0,0 +1,65 @@
+/*
+ * 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;
+
+/**
+ * Configuration knobs for writing parquet, passed to {@link DataFrame#writeParquet(String,
+ * ParquetWriteOptions)}.
+ *
+ * <p>Mirrors a subset of DataFusion's {@code DataFrameWriteOptions} and {@code
+ * TableParquetOptions}. All setters return {@code this} for fluent chaining. Defaults: all fields
+ * {@code null} (meaning the DataFusion default is used).
+ *
+ * <p>Path semantics: when {@link #singleFileOutput(boolean)} is {@code true}, the path passed to
+ * {@code writeParquet} is the literal output filename. When left unset (the default), the path is a
+ * directory that DataFusion populates with one or more part-files.
+ */
+public final class ParquetWriteOptions {
+
+ private String compression;
+ private Boolean singleFileOutput;
+
+ /**
+ * Compression codec spec, passed verbatim to DataFusion. Examples: {@code "snappy"}, {@code
+ * "zstd(3)"}, {@code "gzip(6)"}, {@code "lz4"}, {@code "uncompressed"}. Invalid values surface as
+ * {@link RuntimeException} at write time.
+ */
+ public ParquetWriteOptions compression(String spec) {
+ this.compression = spec;
+ return this;
+ }
+
+ /**
+ * When {@code true}, write to a single file at the supplied path. When left unset (the default),
+ * the path is treated as a directory and DataFusion writes one or more part-files.
+ */
+ public ParquetWriteOptions singleFileOutput(boolean v) {
+ this.singleFileOutput = v;
+ return this;
+ }
+
+ String compression() {
+ return compression;
+ }
+
+ Boolean singleFileOutput() {
+ return singleFileOutput;
+ }
+}
diff --git a/src/test/java/org/apache/datafusion/DataFrameWriteParquetTest.java b/src/test/java/org/apache/datafusion/DataFrameWriteParquetTest.java
new file mode 100644
index 0000000..5be934e
--- /dev/null
+++ b/src/test/java/org/apache/datafusion/DataFrameWriteParquetTest.java
@@ -0,0 +1,112 @@
+/*
+ * 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.assertTrue;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.ipc.ArrowReader;
+import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+class DataFrameWriteParquetTest {
+
+ private static final Path LINEITEM = Path.of("tpch-data/sf1/lineitem.parquet");
+ private static final long LINEITEM_ROWS = 6_001_215L;
+
+ private static void assumeLineitem() {
+ Assumptions.assumeTrue(
+ Files.exists(LINEITEM), "TPC-H SF1 data not found; run `make tpch-data` first");
+ }
+
+ private static long countRowsAt(Path path) throws Exception {
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext()) {
+ ctx.registerParquet("t", path.toAbsolutePath().toString());
+ try (DataFrame df = ctx.sql("SELECT COUNT(*) FROM t");
+ ArrowReader reader = df.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ assertEquals(1, root.getRowCount());
+ return ((BigIntVector) root.getVector(0)).get(0);
+ }
+ }
+ }
+
+ @Test
+ void writeParquetRoundTripsRowCount(@TempDir Path tempDir) throws Exception {
+ assumeLineitem();
+ Path out = tempDir.resolve("out");
+
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.readParquet(LINEITEM.toAbsolutePath().toString())) {
+ df.writeParquet(out.toString());
+ }
+
+ assertEquals(LINEITEM_ROWS, countRowsAt(out));
+ }
+
+ @Test
+ void writeParquetSingleFileProducesOneFile(@TempDir Path tempDir) throws Exception {
+ assumeLineitem();
+ Path out = tempDir.resolve("out.parquet");
+
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.readParquet(LINEITEM.toAbsolutePath().toString())) {
+ df.writeParquet(out.toString(), new ParquetWriteOptions().singleFileOutput(true));
+ }
+
+ assertTrue(Files.isRegularFile(out), "expected single file at " + out);
+ assertEquals(LINEITEM_ROWS, countRowsAt(out));
+ }
+
+ @Test
+ void writeParquetWithCompressionRoundTrips(@TempDir Path tempDir) throws Exception {
+ assumeLineitem();
+ Path out = tempDir.resolve("zstd-out");
+
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.readParquet(LINEITEM.toAbsolutePath().toString())) {
+ df.writeParquet(out.toString(), new ParquetWriteOptions().compression("zstd(3)"));
+ }
+
+ assertEquals(LINEITEM_ROWS, countRowsAt(out));
+ }
+
+ @Test
+ void writeParquetRetainsDataFrame(@TempDir Path tempDir) throws Exception {
+ assumeLineitem();
+ Path out = tempDir.resolve("retained");
+
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.readParquet(LINEITEM.toAbsolutePath().toString())) {
+ df.writeParquet(out.toString());
+ assertEquals(LINEITEM_ROWS, df.count());
+ }
+ }
+}
diff --git a/src/test/java/org/apache/datafusion/ParquetWriteOptionsTest.java b/src/test/java/org/apache/datafusion/ParquetWriteOptionsTest.java
new file mode 100644
index 0000000..bf671ce
--- /dev/null
+++ b/src/test/java/org/apache/datafusion/ParquetWriteOptionsTest.java
@@ -0,0 +1,44 @@
+/*
+ * 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.assertNull;
+
+import org.junit.jupiter.api.Test;
+
+class ParquetWriteOptionsTest {
+
+ @Test
+ void defaultsAreAllNull() {
+ ParquetWriteOptions opts = new ParquetWriteOptions();
+ assertNull(opts.compression());
+ assertNull(opts.singleFileOutput());
+ }
+
+ @Test
+ void fluentSettersChainAndMutate() {
+ ParquetWriteOptions opts =
+ new ParquetWriteOptions().compression("zstd(3)").singleFileOutput(true);
+
+ assertEquals("zstd(3)", opts.compression());
+ assertEquals(Boolean.TRUE, opts.singleFileOutput());
+ }
+}
