From e5fea8d7330a4df84ae2e67b82be0ef652f6e660 Mon Sep 17 00:00:00 2001 From: Sergey Zhukov Date: Mon, 24 Aug 2026 20:44:04 +0200 Subject: [PATCH 1/2] feat(dataframe): improve DataFrame::from_columns input types (#24630) --- datafusion/core/src/dataframe/mod.rs | 27 ++++++++------- datafusion/core/tests/dataframe/mod.rs | 47 ++++++++++++++------------ 2 files changed, 41 insertions(+), 33 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 1299e5fb7bd65..5b73404a29395 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -2614,7 +2614,7 @@ impl DataFrame { /// # async fn main() -> Result<()> { /// let id: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); /// let name: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar", "baz"])); - /// let df = DataFrame::from_columns(vec![("id", id), ("name", name)])?; + /// let df = DataFrame::from_columns([("id", id), ("name", name)])?; /// let expected = vec![ /// "+----+------+", /// "| id | name |", @@ -2628,17 +2628,20 @@ impl DataFrame { /// # Ok(()) /// # } /// ``` - pub fn from_columns(columns: Vec<(&str, ArrayRef)>) -> Result { - let fields = columns - .iter() - .map(|(name, array)| Field::new(*name, array.data_type().clone(), true)) - .collect::>(); - - let arrays = columns + pub fn from_columns(columns: I) -> Result + where + I: IntoIterator, + S: AsRef, + { + let (fields, arrays): (Vec<_>, Vec<_>) = columns .into_iter() - .map(|(_, array)| array) - .collect::>(); - + .map(|(name, array)| { + ( + Field::new(name.as_ref(), array.data_type().clone(), true), + array, + ) + }) + .unzip(); let schema = Arc::new(Schema::new(fields)); let batch = RecordBatch::try_new(schema, arrays)?; let ctx = SessionContext::new(); @@ -2695,7 +2698,7 @@ macro_rules! dataframe { use datafusion::prelude::DataFrame; use datafusion::common::test_util::IntoArrayRef; - let columns = vec![ + let columns = [ $( ($name, $data.into_array_ref()), )+ diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 44ab14e6cc137..30ea61f614cd4 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -6968,7 +6968,7 @@ async fn test_dataframe_from_columns() -> Result<()> { let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); - let df = DataFrame::from_columns(vec![ + let columns = [ ("bool", bools), ("i8", i8s), ("i16", i16s), @@ -6982,10 +6982,10 @@ async fn test_dataframe_from_columns() -> Result<()> { ("f32", f32s), ("f64", f64s), ("str", strings), - ])?; + ]; - assert_eq!(df.schema().fields().len(), 13); - assert_eq!(df.clone().count().await?, 3); + let df1 = DataFrame::from_columns(columns.clone())?; + let df2 = DataFrame::from_columns(columns.to_vec())?; let expected_types = [ ("bool", DataType::Boolean), @@ -7003,26 +7003,31 @@ async fn test_dataframe_from_columns() -> Result<()> { ("str", DataType::Utf8), ]; - let schema = df.schema(); + for df in [df1, df2] { + assert_eq!(df.schema().fields().len(), expected_types.len()); + assert_eq!(df.clone().count().await?, 3); - for (name, data_type) in expected_types { - assert_eq!(schema.field_with_name(None, name)?.data_type(), &data_type); - } + let schema = df.schema(); - let rows = df.sort(vec![col("i32").sort(true, true)])?; + for (name, data_type) in &expected_types { + assert_eq!(schema.field_with_name(None, *name)?.data_type(), data_type); + } - assert_batches_eq!( - &[ - "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", - "| bool | i8 | i16 | i32 | i64 | u8 | u16 | u32 | u64 | f16 | f32 | f64 | str |", - "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", - "| true | -1 | -1 | -1 | -1 | 0 | 0 | 0 | 0 | 1 | 1.0 | 1.0 | foo |", - "| false | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 | 2 | 2.0 | 2.0 | bar |", - "| true | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3.0 | 3.0 | |", - "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", - ], - &rows.collect().await? - ); + let rows = df.sort(vec![col("i32").sort(true, true)])?; + + assert_batches_eq!( + &[ + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| bool | i8 | i16 | i32 | i64 | u8 | u16 | u32 | u64 | f16 | f32 | f64 | str |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| true | -1 | -1 | -1 | -1 | 0 | 0 | 0 | 0 | 1 | 1.0 | 1.0 | foo |", + "| false | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 | 2 | 2.0 | 2.0 | bar |", + "| true | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3.0 | 3.0 | |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + ], + &rows.collect().await? + ); + } Ok(()) } From 41020b6f6c4d513030f06e8ae3c7f81d1ed232c8 Mon Sep 17 00:00:00 2001 From: Sergey Zhukov Date: Mon, 24 Aug 2026 21:16:47 +0200 Subject: [PATCH 2/2] fix cargo clippy in test_dataframe_from_columns --- datafusion/core/tests/dataframe/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 30ea61f614cd4..cf02786800f83 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -7010,7 +7010,7 @@ async fn test_dataframe_from_columns() -> Result<()> { let schema = df.schema(); for (name, data_type) in &expected_types { - assert_eq!(schema.field_with_name(None, *name)?.data_type(), data_type); + assert_eq!(schema.field_with_name(None, name)?.data_type(), data_type); } let rows = df.sort(vec![col("i32").sort(true, true)])?;