feat(dataframe): add withColumnRenamed method

Author: andygrove

native/src/lib.rs

@@ -308,6 +308,26 @@ pub extern "system" fn Java_org_apache_datafusion_DataFrame_dropColumns<'local>(
     })
 }
 
+#[no_mangle]
+pub extern "system" fn Java_org_apache_datafusion_DataFrame_renameColumn<'local>(
+    mut env: JNIEnv<'local>,
+    _class: JClass<'local>,
+    handle: jlong,
+    old_name: JString<'local>,
+    new_name: JString<'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 old: String = env.get_string(&old_name)?.into();
+        let new: String = env.get_string(&new_name)?.into();
+        let new_df = df.with_column_renamed(&old, &new)?;
+        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>,

src/main/java/org/apache/datafusion/DataFrame.java

@@ -163,6 +163,14 @@ public DataFrame dropColumns(String... columnNames) {
     return new DataFrame(dropColumns(nativeHandle, columnNames));
   }
 
+  /** Rename a column. The receiver remains usable and must still be closed independently. */
+  public DataFrame withColumnRenamed(String oldName, String newName) {
+    if (nativeHandle == 0) {
+      throw new IllegalStateException("DataFrame is closed or already collected");
+    }
+    return new DataFrame(renameColumn(nativeHandle, oldName, newName));
+  }
+
   /**
    * Materialize this DataFrame as Parquet at {@code path}. The path is treated as a directory
    * unless overridden via {@link ParquetWriteOptions#singleFileOutput(boolean)}. The receiver
@@ -221,6 +229,8 @@ public void close() {
 
   private static native long dropColumns(long handle, String[] columnNames);
 
+  private static native long renameColumn(long handle, String oldName, String newName);
+
   private static native void writeParquetWithOptions(
       long handle,
       String path,

src/test/java/org/apache/datafusion/DataFrameTransformationsTest.java

@@ -294,4 +294,53 @@ void dropColumnsSilentlyIgnoresUnknownNames() throws Exception {
           root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
     }
   }
+
+  @Test
+  void withColumnRenamedChangesColumnName() throws Exception {
+    try (BufferAllocator allocator = new RootAllocator();
+        SessionContext ctx = new SessionContext();
+        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b");
+        DataFrame renamed = source.withColumnRenamed("a", "alpha");
+        ArrowReader reader = renamed.collect(allocator)) {
+      assertTrue(reader.loadNextBatch());
+      VectorSchemaRoot root = reader.getVectorSchemaRoot();
+      assertArrayEquals(
+          new String[] {"alpha", "b"},
+          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
+    }
+  }
+
+  @Test
+  void withColumnRenamedIsNonDestructive() throws Exception {
+    try (BufferAllocator allocator = new RootAllocator();
+        SessionContext ctx = new SessionContext();
+        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b")) {
+      try (DataFrame renamed = source.withColumnRenamed("a", "alpha")) {
+        assertEquals(1L, renamed.count());
+      }
+      try (DataFrame again = source.select("a");
+          ArrowReader reader = again.collect(allocator)) {
+        assertTrue(reader.loadNextBatch());
+        VectorSchemaRoot root = reader.getVectorSchemaRoot();
+        assertArrayEquals(
+            new String[] {"a"},
+            root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
+      }
+    }
+  }
+
+  @Test
+  void withColumnRenamedUnknownColumnIsNoOp() throws Exception {
+    try (BufferAllocator allocator = new RootAllocator();
+        SessionContext ctx = new SessionContext();
+        DataFrame df = ctx.sql("SELECT 1 AS x");
+        DataFrame renamed = df.withColumnRenamed("not_a_column", "y");
+        ArrowReader reader = renamed.collect(allocator)) {
+      assertTrue(reader.loadNextBatch());
+      VectorSchemaRoot root = reader.getVectorSchemaRoot();
+      assertArrayEquals(
+          new String[] {"x"},
+          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
+    }
+  }
 }