feat(dataframe): add dropColumns method

Author: andygrove

native/src/lib.rs

@@ -281,6 +281,33 @@ pub extern "system" fn Java_org_apache_datafusion_DataFrame_distinctRows<'local>
     })
 }
 
+#[no_mangle]
+pub extern "system" fn Java_org_apache_datafusion_DataFrame_dropColumns<'local>(
+    mut env: JNIEnv<'local>,
+    _class: JClass<'local>,
+    handle: jlong,
+    column_names: 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(&column_names)?;
+        let mut owned: Vec<String> = Vec::with_capacity(len as usize);
+        for i in 0..len {
+            let elem = env.get_object_array_element(&column_names, i)?;
+            let jstr: JString = elem.into();
+            owned.push(env.get_string(&jstr)?.into());
+        }
+        let refs: Vec<&str> = owned.iter().map(String::as_str).collect();
+
+        let new_df = df.drop_columns(&refs)?;
+        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

@@ -152,6 +152,17 @@ public DataFrame distinct() {
     return new DataFrame(distinctRows(nativeHandle));
   }
 
+  /**
+   * Drop the named columns. The inverse of {@link #select(String...)}. The receiver remains usable
+   * and must still be closed independently.
+   */
+  public DataFrame dropColumns(String... columnNames) {
+    if (nativeHandle == 0) {
+      throw new IllegalStateException("DataFrame is closed or already collected");
+    }
+    return new DataFrame(dropColumns(nativeHandle, columnNames));
+  }
+
   /**
    * Materialize this DataFrame as Parquet at {@code path}. The path is treated as a directory
    * unless overridden via {@link ParquetWriteOptions#singleFileOutput(boolean)}. The receiver
@@ -208,6 +219,8 @@ public void close() {
 
   private static native long distinctRows(long handle);
 
+  private static native long dropColumns(long handle, String[] columnNames);
+
   private static native void writeParquetWithOptions(
       long handle,
       String path,

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

@@ -253,4 +253,45 @@ void distinctIsNonDestructive() {
       assertEquals(3L, source.count());
     }
   }
+
+  @Test
+  void dropColumnsRemovesNamedColumns() throws Exception {
+    try (BufferAllocator allocator = new RootAllocator();
+        SessionContext ctx = new SessionContext();
+        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b, 3 AS c");
+        DataFrame dropped = source.dropColumns("b");
+        ArrowReader reader = dropped.collect(allocator)) {
+      assertTrue(reader.loadNextBatch());
+      VectorSchemaRoot root = reader.getVectorSchemaRoot();
+      assertArrayEquals(
+          new String[] {"a", "c"},
+          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
+    }
+  }
+
+  @Test
+  void dropColumnsIsNonDestructive() {
+    try (SessionContext ctx = new SessionContext();
+        DataFrame source = ctx.sql("SELECT 1 AS a, 2 AS b")) {
+      try (DataFrame dropped = source.dropColumns("a")) {
+        assertEquals(1L, dropped.count());
+      }
+      assertEquals(1L, source.count());
+    }
+  }
+
+  @Test
+  void dropColumnsSilentlyIgnoresUnknownNames() throws Exception {
+    try (BufferAllocator allocator = new RootAllocator();
+        SessionContext ctx = new SessionContext();
+        DataFrame df = ctx.sql("SELECT 1 AS x");
+        DataFrame dropped = df.dropColumns("not_a_column");
+        ArrowReader reader = dropped.collect(allocator)) {
+      assertTrue(reader.loadNextBatch());
+      VectorSchemaRoot root = reader.getVectorSchemaRoot();
+      assertArrayEquals(
+          new String[] {"x"},
+          root.getSchema().getFields().stream().map(f -> f.getName()).toArray(String[]::new));
+    }
+  }
 }