[arrow-datafusion] branch master updated: Fix `DataFrame::with_column` to handle creating column names with a period (#3700)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 152896103 Fix `DataFrame::with_column` to handle creating column names with a period (#3700)
152896103 is described below

commit 15289610318e4acad48e40f5adabe0c5a9e8f9b9
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Fri Oct 7 15:51:24 2022 -0400

    Fix `DataFrame::with_column` to handle creating column names with a period (#3700)
    
    * Fix `DataFrame::with_column` to handle creating column names with a period
    
    * Remove old comments
---
 datafusion/core/src/dataframe.rs | 45 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/datafusion/core/src/dataframe.rs b/datafusion/core/src/dataframe.rs
index 0642e556b..06768b563 100644
--- a/datafusion/core/src/dataframe.rs
+++ b/datafusion/core/src/dataframe.rs
@@ -38,6 +38,7 @@ use crate::physical_plan::{execute_stream, execute_stream_partitioned, Execution
 use crate::prelude::SessionContext;
 use crate::scalar::ScalarValue;
 use async_trait::async_trait;
+use datafusion_common::Column;
 use parking_lot::RwLock;
 use parquet::file::properties::WriterProperties;
 use std::any::Any;
@@ -675,7 +676,10 @@ impl DataFrame {
                     col_exists = true;
                     new_column.clone()
                 } else {
-                    col(f.name())
+                    Expr::Column(Column {
+                        relation: None,
+                        name: f.name().into(),
+                    })
                 }
             })
             .collect();
@@ -831,6 +835,7 @@ mod tests {
     use crate::physical_plan::ColumnarValue;
     use crate::test_util;
     use crate::{assert_batches_sorted_eq, execution::context::SessionContext};
+    use arrow::array::Int32Array;
     use arrow::datatypes::DataType;
     use datafusion_expr::{
         avg, cast, count, count_distinct, create_udf, lit, max, min, sum,
@@ -1391,9 +1396,9 @@ mod tests {
         ctx.register_batch("test", data)?;
 
         let sql = r#"
-        SELECT 
+        SELECT
             COUNT(1)
-        FROM 
+        FROM
             test
         GROUP BY
             column_1"#;
@@ -1404,6 +1409,40 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn with_column_name() -> Result<()> {
+        // define data with a column name that has a "." in it:
+        let array: Int32Array = [1, 10].into_iter().collect();
+        let batch =
+            RecordBatch::try_from_iter(vec![("f.c1", Arc::new(array) as _)]).unwrap();
+
+        let ctx = SessionContext::new();
+        ctx.register_batch("t", batch).unwrap();
+
+        let df = ctx
+            .table("t")
+            .unwrap()
+            // try and create a column with a '.' in it
+            .with_column("f.c2", lit("hello"))
+            .unwrap();
+
+        let df_results = df.collect().await.unwrap();
+
+        assert_batches_sorted_eq!(
+            vec![
+                "+------+-------+",
+                "| f.c1 | f.c2  |",
+                "+------+-------+",
+                "| 1    | hello |",
+                "| 10   | hello |",
+                "+------+-------+",
+            ],
+            &df_results
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn cache_test() -> Result<()> {
         let df = test_table()