feat(dataframe): expose sort and repartition (#66)
diff --git a/core/src/main/java/org/apache/datafusion/DataFrame.java b/core/src/main/java/org/apache/datafusion/DataFrame.java
index 3669b55..38ec2f5 100644
--- a/core/src/main/java/org/apache/datafusion/DataFrame.java
+++ b/core/src/main/java/org/apache/datafusion/DataFrame.java
@@ -335,6 +335,87 @@
}
/**
+ * Order the rows by the supplied sort keys. Each {@link SortExpr} names a column and a direction
+ * ({@link SortExpr#asc(String)} / {@link SortExpr#desc(String)}); call {@link
+ * SortExpr#nullsFirst(boolean)} to override null placement.
+ *
+ * <p>An empty {@code exprs} array is a no-op (matches DataFusion's {@code sort(vec![])}). The
+ * receiver remains usable and must still be closed independently.
+ *
+ * @throws IllegalArgumentException if {@code exprs} or any element is {@code null}.
+ * @throws RuntimeException if a sort column does not exist in this DataFrame's schema.
+ */
+ public DataFrame sort(SortExpr... exprs) {
+ if (nativeHandle == 0) {
+ throw new IllegalStateException("DataFrame is closed or already collected");
+ }
+ if (exprs == null) {
+ throw new IllegalArgumentException("sort exprs must be non-null");
+ }
+ String[] columns = new String[exprs.length];
+ boolean[] ascending = new boolean[exprs.length];
+ boolean[] nullsFirst = new boolean[exprs.length];
+ for (int i = 0; i < exprs.length; i++) {
+ SortExpr e = exprs[i];
+ if (e == null) {
+ throw new IllegalArgumentException("sort exprs[" + i + "] must be non-null");
+ }
+ columns[i] = e.column();
+ ascending[i] = e.ascending();
+ nullsFirst[i] = e.nullsFirst();
+ }
+ return new DataFrame(sortRows(nativeHandle, columns, ascending, nullsFirst));
+ }
+
+ /**
+ * Repartition this DataFrame using a round-robin scheme across {@code numPartitions} output
+ * partitions. The receiver remains usable and must still be closed independently.
+ *
+ * @throws IllegalArgumentException if {@code numPartitions <= 0}.
+ * @throws RuntimeException if the underlying repartition plan rejects the request.
+ */
+ public DataFrame repartitionRoundRobin(int numPartitions) {
+ if (nativeHandle == 0) {
+ throw new IllegalStateException("DataFrame is closed or already collected");
+ }
+ if (numPartitions <= 0) {
+ throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions);
+ }
+ return new DataFrame(repartitionRoundRobinRows(nativeHandle, numPartitions));
+ }
+
+ /**
+ * Repartition this DataFrame by hashing the named columns into {@code numPartitions} output
+ * partitions. v1 supports column-name keys only; expression keys are deferred until the Java
+ * binding gains an {@code Expr} builder. The receiver remains usable and must still be closed
+ * independently.
+ *
+ * @throws IllegalArgumentException if {@code numPartitions <= 0}, {@code columns} is {@code null}
+ * or empty, or any element of {@code columns} is {@code null}.
+ * @throws RuntimeException if a partition column does not exist in this DataFrame's schema.
+ */
+ public DataFrame repartitionHash(int numPartitions, String... columns) {
+ if (nativeHandle == 0) {
+ throw new IllegalStateException("DataFrame is closed or already collected");
+ }
+ if (numPartitions <= 0) {
+ throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions);
+ }
+ if (columns == null) {
+ throw new IllegalArgumentException("repartitionHash columns must be non-null");
+ }
+ if (columns.length == 0) {
+ throw new IllegalArgumentException("repartitionHash requires at least one column");
+ }
+ for (int i = 0; i < columns.length; i++) {
+ if (columns[i] == null) {
+ throw new IllegalArgumentException("repartitionHash columns[" + i + "] must be non-null");
+ }
+ }
+ return new DataFrame(repartitionHashRows(nativeHandle, numPartitions, columns));
+ }
+
+ /**
* Equi-join this DataFrame with {@code right} on the named columns, using the given {@link
* JoinType}. The receiver and {@code right} both remain usable and must still be closed
* independently.
@@ -587,6 +668,13 @@
private static native long unnestColumns(long handle, String[] columns, boolean preserveNulls);
+ private static native long sortRows(
+ long handle, String[] columns, boolean[] ascending, boolean[] nullsFirst);
+
+ private static native long repartitionRoundRobinRows(long handle, int numPartitions);
+
+ private static native long repartitionHashRows(long handle, int numPartitions, String[] columns);
+
private static native long joinDataFrame(
long leftHandle,
long rightHandle,
diff --git a/core/src/main/java/org/apache/datafusion/SortExpr.java b/core/src/main/java/org/apache/datafusion/SortExpr.java
new file mode 100644
index 0000000..e94997d
--- /dev/null
+++ b/core/src/main/java/org/apache/datafusion/SortExpr.java
@@ -0,0 +1,82 @@
+/*
+ * 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;
+
+/**
+ * A single sort key passed to {@link DataFrame#sort(SortExpr...)}, mirroring DataFusion's {@code
+ * expr::Sort{ expr, asc, nulls_first }}. Build via {@link #asc(String)} or {@link #desc(String)};
+ * tweak null placement via {@link #nullsFirst(boolean)}.
+ *
+ * <p>v1 accepts column names only -- the {@code column} string is interpreted as a column
+ * reference, not a SQL expression. Complex sort keys (e.g. {@code "a + b"}) are intentionally
+ * deferred until the Java binding gains an {@code Expr} builder.
+ *
+ * <p>Defaults match DataFusion: {@link #asc(String)} sorts ascending with NULLs last; {@link
+ * #desc(String)} sorts descending with NULLs first.
+ */
+public final class SortExpr {
+
+ private final String column;
+ private final boolean ascending;
+ private boolean nullsFirst;
+
+ private SortExpr(String column, boolean ascending, boolean nullsFirst) {
+ this.column = column;
+ this.ascending = ascending;
+ this.nullsFirst = nullsFirst;
+ }
+
+ /** Sort the given column ascending, with NULLs placed last. */
+ public static SortExpr asc(String column) {
+ if (column == null) {
+ throw new IllegalArgumentException("SortExpr column must be non-null");
+ }
+ return new SortExpr(column, true, false);
+ }
+
+ /** Sort the given column descending, with NULLs placed first. */
+ public static SortExpr desc(String column) {
+ if (column == null) {
+ throw new IllegalArgumentException("SortExpr column must be non-null");
+ }
+ return new SortExpr(column, false, true);
+ }
+
+ /** Override the default NULL placement. {@code true} = NULLs first, {@code false} = last. */
+ public SortExpr nullsFirst(boolean v) {
+ this.nullsFirst = v;
+ return this;
+ }
+
+ /** The column name this sort key references. */
+ public String column() {
+ return column;
+ }
+
+ /** {@code true} if ascending, {@code false} if descending. */
+ public boolean ascending() {
+ return ascending;
+ }
+
+ /** {@code true} if NULLs are placed first, {@code false} if last. */
+ public boolean nullsFirst() {
+ return nullsFirst;
+ }
+}
diff --git a/core/src/test/java/org/apache/datafusion/DataFrameTransformationsTest.java b/core/src/test/java/org/apache/datafusion/DataFrameTransformationsTest.java
index 495f4ce..ab977a9 100644
--- a/core/src/test/java/org/apache/datafusion/DataFrameTransformationsTest.java
+++ b/core/src/test/java/org/apache/datafusion/DataFrameTransformationsTest.java
@@ -135,6 +135,9 @@
assertThrows(IllegalStateException.class, () -> df.withColumnRenamed("x", "y"));
assertThrows(IllegalStateException.class, () -> df.withColumn("y", "x + 1"));
assertThrows(IllegalStateException.class, () -> df.unnestColumns("x"));
+ assertThrows(IllegalStateException.class, () -> df.sort(SortExpr.asc("x")));
+ assertThrows(IllegalStateException.class, () -> df.repartitionRoundRobin(1));
+ assertThrows(IllegalStateException.class, () -> df.repartitionHash(1, "x"));
assertThrows(IllegalStateException.class, df::count);
assertThrows(IllegalStateException.class, df::show);
assertThrows(IllegalStateException.class, () -> df.show(5));
@@ -158,6 +161,9 @@
assertThrows(IllegalStateException.class, () -> df.withColumnRenamed("x", "y"));
assertThrows(IllegalStateException.class, () -> df.withColumn("y", "x + 1"));
assertThrows(IllegalStateException.class, () -> df.unnestColumns("x"));
+ assertThrows(IllegalStateException.class, () -> df.sort(SortExpr.asc("x")));
+ assertThrows(IllegalStateException.class, () -> df.repartitionRoundRobin(1));
+ assertThrows(IllegalStateException.class, () -> df.repartitionHash(1, "x"));
assertThrows(IllegalStateException.class, df::count);
assertThrows(IllegalStateException.class, df::show);
assertThrows(IllegalStateException.class, () -> df.show(5));
@@ -504,4 +510,219 @@
() -> df.unnestColumns(new UnnestOptions(), (String[]) null));
}
}
+
+ // -- sort -----------------------------------------------------------------
+
+ @Test
+ void sortAscOrdersAscending() throws Exception {
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT * FROM (VALUES (3), (1), (2)) AS t(x)");
+ DataFrame sorted = df.sort(SortExpr.asc("x"));
+ ArrowReader reader = sorted.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ BigIntVector x = (BigIntVector) root.getVector("x");
+ assertEquals(3, x.getValueCount());
+ assertEquals(1L, x.get(0));
+ assertEquals(2L, x.get(1));
+ assertEquals(3L, x.get(2));
+ }
+ }
+
+ @Test
+ void sortDescOrdersDescending() throws Exception {
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT * FROM (VALUES (3), (1), (2)) AS t(x)");
+ DataFrame sorted = df.sort(SortExpr.desc("x"));
+ ArrowReader reader = sorted.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ BigIntVector x = (BigIntVector) root.getVector("x");
+ assertEquals(3, x.getValueCount());
+ assertEquals(3L, x.get(0));
+ assertEquals(2L, x.get(1));
+ assertEquals(1L, x.get(2));
+ }
+ }
+
+ @Test
+ void sortMultipleKeysMixedDirection() throws Exception {
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT * FROM (VALUES (1, 1), (1, 2), (2, 1)) AS t(x, y)");
+ DataFrame sorted = df.sort(SortExpr.asc("x"), SortExpr.desc("y"));
+ ArrowReader reader = sorted.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ BigIntVector x = (BigIntVector) root.getVector("x");
+ BigIntVector y = (BigIntVector) root.getVector("y");
+ // (1,2), (1,1), (2,1) -- x asc, y desc within each x group.
+ assertEquals(1L, x.get(0));
+ assertEquals(2L, y.get(0));
+ assertEquals(1L, x.get(1));
+ assertEquals(1L, y.get(1));
+ assertEquals(2L, x.get(2));
+ assertEquals(1L, y.get(2));
+ }
+ }
+
+ @Test
+ void sortNullsFirstPlacesNullsFirst() throws Exception {
+ String sql =
+ "SELECT * FROM (VALUES (CAST(1 AS BIGINT)), (CAST(NULL AS BIGINT)),"
+ + " (CAST(2 AS BIGINT))) AS t(x)";
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql(sql);
+ DataFrame sorted = df.sort(SortExpr.asc("x").nullsFirst(true));
+ ArrowReader reader = sorted.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ BigIntVector x = (BigIntVector) root.getVector("x");
+ assertEquals(3, x.getValueCount());
+ assertTrue(x.isNull(0), "expected NULL first");
+ assertEquals(1L, x.get(1));
+ assertEquals(2L, x.get(2));
+ }
+ }
+
+ @Test
+ void sortNullsLastPlacesNullsLast() throws Exception {
+ String sql =
+ "SELECT * FROM (VALUES (CAST(1 AS BIGINT)), (CAST(NULL AS BIGINT)),"
+ + " (CAST(2 AS BIGINT))) AS t(x)";
+ try (BufferAllocator allocator = new RootAllocator();
+ SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql(sql);
+ DataFrame sorted = df.sort(SortExpr.asc("x").nullsFirst(false));
+ ArrowReader reader = sorted.collect(allocator)) {
+ assertTrue(reader.loadNextBatch());
+ VectorSchemaRoot root = reader.getVectorSchemaRoot();
+ BigIntVector x = (BigIntVector) root.getVector("x");
+ assertEquals(3, x.getValueCount());
+ assertEquals(1L, x.get(0));
+ assertEquals(2L, x.get(1));
+ assertTrue(x.isNull(2), "expected NULL last");
+ }
+ }
+
+ @Test
+ void sortIsNonDestructive() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame source = ctx.sql("SELECT * FROM (VALUES (3), (1), (2)) AS t(x)")) {
+ try (DataFrame sorted = source.sort(SortExpr.asc("x"))) {
+ assertEquals(3L, sorted.count());
+ }
+ assertEquals(3L, source.count());
+ }
+ }
+
+ @Test
+ void sortRejectsNullArgs() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT 1 AS x")) {
+ assertThrows(IllegalArgumentException.class, () -> df.sort((SortExpr[]) null));
+ assertThrows(IllegalArgumentException.class, () -> df.sort(new SortExpr[] {null}));
+ }
+ }
+
+ @Test
+ void sortUnknownColumnThrows() {
+ // Sort plan-builds lazily; DataFusion may not validate column names until the plan is
+ // actually executed. Drive a count() to force planning + optimization.
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT 1 AS x")) {
+ assertThrows(
+ RuntimeException.class,
+ () -> {
+ try (DataFrame sorted = df.sort(SortExpr.asc("not_a_column"))) {
+ sorted.count();
+ }
+ });
+ }
+ }
+
+ // -- repartitionRoundRobin -------------------------------------------------
+
+ @Test
+ void repartitionRoundRobinPreservesRows() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4), (5)) AS t(x)");
+ DataFrame repartitioned = source.repartitionRoundRobin(3)) {
+ assertEquals(5L, repartitioned.count());
+ }
+ }
+
+ @Test
+ void repartitionRoundRobinIsNonDestructive() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")) {
+ try (DataFrame rp = source.repartitionRoundRobin(2)) {
+ assertEquals(3L, rp.count());
+ }
+ assertEquals(3L, source.count());
+ }
+ }
+
+ @Test
+ void repartitionRoundRobinRejectsNonPositive() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT 1 AS x")) {
+ assertThrows(IllegalArgumentException.class, () -> df.repartitionRoundRobin(0));
+ assertThrows(IllegalArgumentException.class, () -> df.repartitionRoundRobin(-1));
+ }
+ }
+
+ // -- repartitionHash -------------------------------------------------------
+
+ @Test
+ void repartitionHashPreservesRows() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3), (4), (5)) AS t(x)");
+ DataFrame repartitioned = source.repartitionHash(3, "x")) {
+ assertEquals(5L, repartitioned.count());
+ }
+ }
+
+ @Test
+ void repartitionHashIsNonDestructive() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame source = ctx.sql("SELECT * FROM (VALUES (1), (2), (3)) AS t(x)")) {
+ try (DataFrame rp = source.repartitionHash(2, "x")) {
+ assertEquals(3L, rp.count());
+ }
+ assertEquals(3L, source.count());
+ }
+ }
+
+ @Test
+ void repartitionHashRejectsBadArgs() {
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT 1 AS x")) {
+ assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(0, "x"));
+ assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(-1, "x"));
+ assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(2, (String[]) null));
+ assertThrows(IllegalArgumentException.class, () -> df.repartitionHash(2));
+ assertThrows(
+ IllegalArgumentException.class, () -> df.repartitionHash(2, new String[] {null}));
+ }
+ }
+
+ @Test
+ void repartitionHashUnknownColumnThrows() {
+ // Repartition plan-builds lazily; DataFusion may not validate column names until the plan
+ // is actually executed. Drive a count() to force planning + optimization.
+ try (SessionContext ctx = new SessionContext();
+ DataFrame df = ctx.sql("SELECT 1 AS x")) {
+ assertThrows(
+ RuntimeException.class,
+ () -> {
+ try (DataFrame rp = df.repartitionHash(2, "not_a_column")) {
+ rp.count();
+ }
+ });
+ }
+ }
}
diff --git a/core/src/test/java/org/apache/datafusion/SortExprTest.java b/core/src/test/java/org/apache/datafusion/SortExprTest.java
new file mode 100644
index 0000000..7287205
--- /dev/null
+++ b/core/src/test/java/org/apache/datafusion/SortExprTest.java
@@ -0,0 +1,64 @@
+/*
+ * 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 org.junit.jupiter.api.Test;
+
+class SortExprTest {
+
+ @Test
+ void ascDefaultsToNullsLast() {
+ SortExpr s = SortExpr.asc("a");
+ assertTrue(s.ascending());
+ assertFalse(s.nullsFirst());
+ }
+
+ @Test
+ void descDefaultsToNullsFirst() {
+ SortExpr s = SortExpr.desc("a");
+ assertFalse(s.ascending());
+ assertTrue(s.nullsFirst());
+ }
+
+ @Test
+ void nullsFirstSetterRoundTrips() {
+ SortExpr s = SortExpr.asc("a").nullsFirst(true);
+ assertTrue(s.nullsFirst());
+ s.nullsFirst(false);
+ assertFalse(s.nullsFirst());
+ }
+
+ @Test
+ void factoryRejectsNullColumn() {
+ assertThrows(IllegalArgumentException.class, () -> SortExpr.asc(null));
+ assertThrows(IllegalArgumentException.class, () -> SortExpr.desc(null));
+ }
+
+ @Test
+ void accessorReturnsColumnName() {
+ assertEquals("foo", SortExpr.asc("foo").column());
+ assertEquals("bar", SortExpr.desc("bar").column());
+ }
+}
diff --git a/native/src/lib.rs b/native/src/lib.rs
index ccfb38e..b59d4fd 100644
--- a/native/src/lib.rs
+++ b/native/src/lib.rs
@@ -49,10 +49,10 @@
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::execution::SendableRecordBatchStream;
use datafusion::logical_expr::Expr;
-use datafusion::logical_expr::{ScalarUDF, Signature};
+use datafusion::logical_expr::{col, Partitioning, ScalarUDF, Signature, SortExpr};
use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
use futures::StreamExt;
-use jni::objects::{JByteArray, JClass, JObject, JObjectArray, JString};
+use jni::objects::{JBooleanArray, JByteArray, JClass, JObject, JObjectArray, JString};
use jni::sys::{jboolean, jbyte, jbyteArray, jint, jlong};
use jni::JNIEnv;
use jni::JavaVM;
@@ -727,6 +727,91 @@
}
#[no_mangle]
+pub extern "system" fn Java_org_apache_datafusion_DataFrame_sortRows<'local>(
+ mut env: JNIEnv<'local>,
+ _class: JClass<'local>,
+ handle: jlong,
+ columns: JObjectArray<'local>,
+ ascending: JBooleanArray<'local>,
+ nulls_first: JBooleanArray<'local>,
+) -> jlong {
+ try_unwrap_or_throw(&mut env, 0, |env| -> JniResult<jlong> {
+ if handle == 0 {
+ return Err("DataFrame handle is null".into());
+ }
+ let df = unsafe { &*(handle as *const DataFrame) }.clone();
+
+ let len = env.get_array_length(&columns)? as usize;
+ let mut names: Vec<String> = Vec::with_capacity(len);
+ for i in 0..len {
+ let elem = env.get_object_array_element(&columns, i as i32)?;
+ let jstr: JString = elem.into();
+ names.push(env.get_string(&jstr)?.into());
+ }
+
+ // Decode the two parallel boolean arrays via primitive-array region copies.
+ let mut asc_buf = vec![0u8; len];
+ env.get_boolean_array_region(&ascending, 0, &mut asc_buf)?;
+ let mut nf_buf = vec![0u8; len];
+ env.get_boolean_array_region(&nulls_first, 0, &mut nf_buf)?;
+
+ let sort_exprs: Vec<SortExpr> = names
+ .into_iter()
+ .zip(asc_buf.into_iter().zip(nf_buf))
+ .map(|(name, (asc, nf))| SortExpr::new(col(&name), asc != 0, nf != 0))
+ .collect();
+
+ let new_df = df.sort(sort_exprs)?;
+ Ok(Box::into_raw(Box::new(new_df)) as jlong)
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn Java_org_apache_datafusion_DataFrame_repartitionRoundRobinRows<'local>(
+ mut env: JNIEnv<'local>,
+ _class: JClass<'local>,
+ handle: jlong,
+ num_partitions: jint,
+) -> jlong {
+ try_unwrap_or_throw(&mut env, 0, |_env| -> JniResult<jlong> {
+ if handle == 0 {
+ return Err("DataFrame handle is null".into());
+ }
+ let df = unsafe { &*(handle as *const DataFrame) }.clone();
+ let new_df = df.repartition(Partitioning::RoundRobinBatch(num_partitions as usize))?;
+ Ok(Box::into_raw(Box::new(new_df)) as jlong)
+ })
+}
+
+#[no_mangle]
+pub extern "system" fn Java_org_apache_datafusion_DataFrame_repartitionHashRows<'local>(
+ mut env: JNIEnv<'local>,
+ _class: JClass<'local>,
+ handle: jlong,
+ num_partitions: jint,
+ columns: JObjectArray<'local>,
+) -> jlong {
+ try_unwrap_or_throw(&mut env, 0, |env| -> JniResult<jlong> {
+ if handle == 0 {
+ return Err("DataFrame handle is null".into());
+ }
+ let df = unsafe { &*(handle as *const DataFrame) }.clone();
+
+ let len = env.get_array_length(&columns)?;
+ let mut owned: Vec<String> = Vec::with_capacity(len as usize);
+ for i in 0..len {
+ let elem = env.get_object_array_element(&columns, i)?;
+ let jstr: JString = elem.into();
+ owned.push(env.get_string(&jstr)?.into());
+ }
+ let exprs = owned.iter().map(|s| col(s.as_str())).collect();
+
+ let new_df = df.repartition(Partitioning::Hash(exprs, num_partitions as usize))?;
+ Ok(Box::into_raw(Box::new(new_df)) as jlong)
+ })
+}
+
+#[no_mangle]
pub extern "system" fn Java_org_apache_datafusion_DataFrame_writeParquetWithOptions<'local>(
mut env: JNIEnv<'local>,
_class: JClass<'local>,
