This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 5498a02853 Add drop_columns to dataframe api (#11010)
5498a02853 is described below
commit 5498a02853021fc07a57743990d207e25ada27f4
Author: Bruce Ritchie <[email protected]>
AuthorDate: Fri Jun 21 12:14:57 2024 -0400
Add drop_columns to dataframe api (#11010)
* Add drop_columns to dataframe api #11007
* Prettier cleanup
* Added additional drop_columns tests and fixed issue with nonexistent
columns.
---
datafusion/core/src/dataframe/mod.rs | 169 +++++++++++++++++++++++++++++++++++
docs/source/user-guide/dataframe.md | 1 +
2 files changed, 170 insertions(+)
diff --git a/datafusion/core/src/dataframe/mod.rs
b/datafusion/core/src/dataframe/mod.rs
index b5c58eff57..43ac3992be 100644
--- a/datafusion/core/src/dataframe/mod.rs
+++ b/datafusion/core/src/dataframe/mod.rs
@@ -244,6 +244,42 @@ impl DataFrame {
})
}
+ /// Returns a new DataFrame containing all columns except the specified
columns.
+ ///
+ /// ```
+ /// # use datafusion::prelude::*;
+ /// # use datafusion::error::Result;
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<()> {
+ /// let ctx = SessionContext::new();
+ /// let df = ctx.read_csv("tests/data/example.csv",
CsvReadOptions::new()).await?;
+ /// let df = df.drop_columns(&["a"])?;
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub fn drop_columns(self, columns: &[&str]) -> Result<DataFrame> {
+ let fields_to_drop = columns
+ .iter()
+ .map(|name| {
+ self.plan
+ .schema()
+ .qualified_field_with_unqualified_name(name)
+ })
+ .filter(|r| r.is_ok())
+ .collect::<Result<Vec<_>>>()?;
+ let expr: Vec<Expr> = self
+ .plan
+ .schema()
+ .fields()
+ .into_iter()
+ .enumerate()
+ .map(|(idx, _)| self.plan.schema().qualified_field(idx))
+ .filter(|(qualifier, f)| !fields_to_drop.contains(&(*qualifier,
f)))
+ .map(|(qualifier, field)| Expr::Column(Column::from((qualifier,
field))))
+ .collect();
+ self.select(expr)
+ }
+
/// Expand each list element of a column to multiple rows.
#[deprecated(since = "37.0.0", note = "use unnest_columns instead")]
pub fn unnest_column(self, column: &str) -> Result<DataFrame> {
@@ -1799,6 +1835,139 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn drop_columns() -> Result<()> {
+ // build plan using Table API
+ let t = test_table().await?;
+ let t2 = t.drop_columns(&["c2", "c11"])?;
+ let plan = t2.plan.clone();
+
+ // build query using SQL
+ let sql_plan = create_plan(
+ "SELECT c1,c3,c4,c5,c6,c7,c8,c9,c10,c12,c13 FROM
aggregate_test_100",
+ )
+ .await?;
+
+ // the two plans should be identical
+ assert_same_plan(&plan, &sql_plan);
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn drop_columns_with_duplicates() -> Result<()> {
+ // build plan using Table API
+ let t = test_table().await?;
+ let t2 = t.drop_columns(&["c2", "c11", "c2", "c2"])?;
+ let plan = t2.plan.clone();
+
+ // build query using SQL
+ let sql_plan = create_plan(
+ "SELECT c1,c3,c4,c5,c6,c7,c8,c9,c10,c12,c13 FROM
aggregate_test_100",
+ )
+ .await?;
+
+ // the two plans should be identical
+ assert_same_plan(&plan, &sql_plan);
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn drop_columns_with_nonexistent_columns() -> Result<()> {
+ // build plan using Table API
+ let t = test_table().await?;
+ let t2 = t.drop_columns(&["canada", "c2", "rocks"])?;
+ let plan = t2.plan.clone();
+
+ // build query using SQL
+ let sql_plan = create_plan(
+ "SELECT c1,c3,c4,c5,c6,c7,c8,c9,c10,c11,c12,c13 FROM
aggregate_test_100",
+ )
+ .await?;
+
+ // the two plans should be identical
+ assert_same_plan(&plan, &sql_plan);
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn drop_columns_with_empty_array() -> Result<()> {
+ // build plan using Table API
+ let t = test_table().await?;
+ let t2 = t.drop_columns(&[])?;
+ let plan = t2.plan.clone();
+
+ // build query using SQL
+ let sql_plan = create_plan(
+ "SELECT c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,c11,c12,c13 FROM
aggregate_test_100",
+ )
+ .await?;
+
+ // the two plans should be identical
+ assert_same_plan(&plan, &sql_plan);
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn drop_with_quotes() -> Result<()> {
+ // define data with a column name that has a "." in it:
+ let array1: Int32Array = [1, 10].into_iter().collect();
+ let array2: Int32Array = [2, 11].into_iter().collect();
+ let batch = RecordBatch::try_from_iter(vec![
+ ("f\"c1", Arc::new(array1) as _),
+ ("f\"c2", Arc::new(array2) as _),
+ ])?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+
+ let df = ctx.table("t").await?.drop_columns(&["f\"c1"])?;
+
+ let df_results = df.collect().await?;
+
+ assert_batches_sorted_eq!(
+ [
+ "+------+",
+ "| f\"c2 |",
+ "+------+",
+ "| 2 |",
+ "| 11 |",
+ "+------+"
+ ],
+ &df_results
+ );
+
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn drop_with_periods() -> Result<()> {
+ // define data with a column name that has a "." in it:
+ let array1: Int32Array = [1, 10].into_iter().collect();
+ let array2: Int32Array = [2, 11].into_iter().collect();
+ let batch = RecordBatch::try_from_iter(vec![
+ ("f.c1", Arc::new(array1) as _),
+ ("f.c2", Arc::new(array2) as _),
+ ])?;
+
+ let ctx = SessionContext::new();
+ ctx.register_batch("t", batch)?;
+
+ let df = ctx.table("t").await?.drop_columns(&["f.c1"])?;
+
+ let df_results = df.collect().await?;
+
+ assert_batches_sorted_eq!(
+ ["+------+", "| f.c2 |", "+------+", "| 2 |", "| 11 |",
"+------+"],
+ &df_results
+ );
+
+ Ok(())
+ }
+
#[tokio::test]
async fn aggregate() -> Result<()> {
// build plan using DataFrame API
diff --git a/docs/source/user-guide/dataframe.md
b/docs/source/user-guide/dataframe.md
index c0210200a2..744a719e77 100644
--- a/docs/source/user-guide/dataframe.md
+++ b/docs/source/user-guide/dataframe.md
@@ -64,6 +64,7 @@ execution. The plan is evaluated (executed) when an action
method is invoked, su
| ------------------- |
------------------------------------------------------------------------------------------------------------------------------------------
|
| aggregate | Perform an aggregate query with optional grouping
expressions.
|
| distinct | Filter out duplicate rows.
|
+| drop_columns | Create a projection with all but the provided column
names.
|
| except | Calculate the exception of two DataFrames. The two
DataFrames must have exactly the same schema
|
| filter | Filter a DataFrame to only include rows that match the
specified filter expression.
|
| intersect | Calculate the intersection of two DataFrames. The two
DataFrames must have exactly the same schema
|
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
