renato2099 commented on code in PR #19549:
URL: https://github.com/apache/datafusion/pull/19549#discussion_r2659188879
##########
datafusion/core/src/dataframe/mod.rs:
##########
@@ -2463,6 +2479,48 @@ impl DataFrame {
.collect()
}
+ /// Find qualified columns for this dataframe from names
+ ///
+ /// # Arguments
+ /// * `names` - Unqualified names to find.
+ ///
+ /// # Example
+ /// ```
+ /// # use datafusion::prelude::*;
+ /// # use datafusion::error::Result;
+ /// # use datafusion_common::ScalarValue;
+ /// # #[tokio::main]
+ /// # async fn main() -> Result<()> {
+ /// let ctx = SessionContext::new();
+ /// ctx.register_csv("first_table", "tests/data/example.csv",
CsvReadOptions::new())
+ /// .await?;
+ /// let df = ctx.table("first_table").await?;
+ /// ctx.register_csv("second_table", "tests/data/example.csv",
CsvReadOptions::new())
+ /// .await?;
+ /// let df2 = ctx.table("second_table").await?;
+ /// let join_expr = df.find_qualified_columns(&["a"])?.iter()
+ /// .zip(df2.find_qualified_columns(&["a"])?.iter())
+ /// .map(|(col1, col2)| col(*col1).eq(col(*col2)))
+ /// .collect::<Vec<Expr>>();
+ /// let df3 = df.join_on(df2, JoinType::Inner, join_expr)?;
+ /// # Ok(())
+ /// # }
+ /// ```
+ pub fn find_qualified_columns(
+ &self,
+ names: &[&str],
+ ) -> Result<Vec<(Option<&TableReference>, &FieldRef)>> {
+ let schema = self.logical_plan().schema();
+ names
+ .iter()
+ .map(|name| {
+ schema
+ .qualified_field_from_column(&Column::from_name(*name))
+ .map_err(|_| plan_datafusion_err!("Column '{}' not found",
name))
Review Comment:
much better than silently ignoring the error 👍
##########
datafusion/core/tests/dataframe/mod.rs:
##########
@@ -549,6 +550,107 @@ async fn drop_columns_with_empty_array() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn drop_columns_qualified() -> Result<()> {
+ // build plan using Table API
+ let mut t = test_table().await?;
+ t = t.select_columns(&["c1", "c2", "c11"])?;
+ let mut t2 = test_table_with_name("another_table").await?;
+ t2 = t2.select_columns(&["c1", "c2", "c11"])?;
+ let mut t3 = t.join_on(
+ t2,
+ JoinType::Inner,
+ [col("aggregate_test_100.c1").eq(col("another_table.c1"))],
+ )?;
+ t3 = t3.drop_columns(&["another_table.c2", "another_table.c11"])?;
Review Comment:
sorry I think I needed more time to look at the code, things look good to me
after checking `schema.qualified_field_from_column` and
`qualified_field_with_unqualified_name`, we will get an error from there if the
column is indeed ambiguous.
So the current implementation already supports the “error on ambiguity” mode
👍 (sorry for the noise)
##########
datafusion/core/tests/dataframe/mod.rs:
##########
@@ -549,6 +550,107 @@ async fn drop_columns_with_empty_array() -> Result<()> {
Ok(())
}
+#[tokio::test]
+async fn drop_columns_qualified() -> Result<()> {
+ // build plan using Table API
+ let mut t = test_table().await?;
+ t = t.select_columns(&["c1", "c2", "c11"])?;
+ let mut t2 = test_table_with_name("another_table").await?;
+ t2 = t2.select_columns(&["c1", "c2", "c11"])?;
+ let mut t3 = t.join_on(
+ t2,
+ JoinType::Inner,
+ [col("aggregate_test_100.c1").eq(col("another_table.c1"))],
+ )?;
+ t3 = t3.drop_columns(&["another_table.c2", "another_table.c11"])?;
Review Comment:
clarification question: what should happen if we do something like ?
```
t3 = t3.drop_columns(&["c1"])?;
```
i.e., attempting to drop an ambiguous column?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
