From 6d3f3791d523b8ff52160343278b64d7a2c89afc Mon Sep 17 00:00:00 2001 From: Sergey Zhukov Date: Mon, 14 Sep 2026 20:07:18 +0200 Subject: [PATCH 1/2] feat(dataframe): add column to a dataframe (#25298) --- datafusion/core/src/dataframe/mod.rs | 317 ++++++++++++++++++++++++- datafusion/core/tests/dataframe/mod.rs | 128 ++++++++++ 2 files changed, 442 insertions(+), 3 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 55d3695a68d7a..643088d26d684 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -36,13 +36,18 @@ use crate::logical_expr::{ Expr, JoinType, LogicalPlan, LogicalPlanBuilder, LogicalPlanBuilderOptions, Partitioning, TableType, col, ident, }; +use crate::physical_expr::EquivalenceProperties; use crate::physical_plan::{ - ExecutionPlan, SendableRecordBatchStream, collect, collect_partitioned, - execute_stream, execute_stream_partitioned, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, Partitioning as PhysicalPartitioning, PhysicalExpr, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, + coalesce_partitions::CoalescePartitionsExec, collect, collect_partitioned, + execute_stream, execute_stream_partitioned, stream::RecordBatchStreamAdapter, }; use crate::prelude::SessionContext; use std::borrow::Cow; use std::collections::{HashMap, HashSet}; +use std::fmt::{self, Formatter}; use std::sync::Arc; use arrow::array::{Array, ArrayRef, Int64Array, StringArray}; @@ -54,7 +59,8 @@ use datafusion_common::config::{CsvOptions, JsonOptions}; use datafusion_common::{ Column, DFSchema, DataFusionError, ParamValues, ScalarValue, SchemaError, TableReference, UnnestOptions, exec_err, internal_datafusion_err, not_impl_err, - plan_datafusion_err, plan_err, unqualified_field_not_found, + plan_datafusion_err, plan_err, project_schema, tree_node::TreeNodeRecursion, + unqualified_field_not_found, }; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{ @@ -70,6 +76,7 @@ use datafusion_functions_aggregate::expr_fn::{ use async_trait::async_trait; use datafusion_catalog::Session; use datafusion_expr::extension_types::DFArrayFormatterFactory; +use futures::StreamExt; use futures::future::BoxFuture; /// Contains options that control how data is @@ -2745,6 +2752,310 @@ impl DataFrame { let df = ctx.read_batch(batch)?; Ok(df) } + + /// Append named Arrow arrays as columns to this [`DataFrame`]. + /// + /// This does not execute the current plan. The arrays are attached when the + /// returned DataFrame is collected. Each array must have the same length as + /// the number of rows this DataFrame produces. + /// + /// See [`Self::with_column`] to add a column from an [`Expr`]. + /// + /// # Example + /// + /// ``` + /// use std::sync::Arc; + /// use arrow::array::{ArrayRef, Int32Array, StringArray}; + /// use datafusion::prelude::DataFrame; + /// # use datafusion::error::Result; + /// # use datafusion_common::assert_batches_sorted_eq; + /// # #[tokio::main] + /// # 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([("id", id), ("name", name)])?; + /// + /// let extra: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); + /// let df = df.with_array_columns([("extra", extra)])?; + /// + /// let expected = vec![ + /// "+----+------+-------+", + /// "| id | name | extra |", + /// "+----+------+-------+", + /// "| 1 | foo | 10 |", + /// "| 2 | bar | 20 |", + /// "| 3 | baz | 30 |", + /// "+----+------+-------+", + /// ]; + /// # assert_batches_sorted_eq!(expected, &df.collect().await?); + /// # Ok(()) + /// # } + /// ``` + pub fn with_array_columns<'a, I>(self, columns: I) -> Result + where + I: IntoIterator, + { + let added: Vec<(String, ArrayRef)> = columns + .into_iter() + .map(|(name, array)| (name.to_string(), array)) + .collect(); + + if added.is_empty() { + return Ok(self); + } + + if let Some((_, first)) = added.first() { + let expected = first.len(); + for (name, array) in &added { + if array.len() != expected { + return plan_err!( + "Column '{name}' has length {}, expected {expected}", + array.len() + ); + } + } + } + + let (state, plan) = self.into_parts(); + let provider = Arc::new(AddColumnProvider::try_new(plan, added)?); + let plan = + LogicalPlanBuilder::scan(UNNAMED_TABLE, provider_as_source(provider), None)? + .build()?; + + Ok(DataFrame::new(state, plan)) + } +} + +#[derive(Debug)] +struct AddColumnProvider { + input: LogicalPlan, + added: Vec<(String, ArrayRef)>, + schema: SchemaRef, +} + +impl AddColumnProvider { + fn try_new(input: LogicalPlan, added: Vec<(String, ArrayRef)>) -> Result { + let mut fields: Vec = input + .schema() + .as_arrow() + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + + for (name, array) in &added { + if input.schema().has_column_with_unqualified_name(name) { + return plan_err!("Column '{name}' already exists"); + } + fields.push(Field::new(name, array.data_type().clone(), true)); + } + + Ok(Self { + input, + added, + schema: Arc::new(Schema::new(fields)), + }) + } +} + +#[async_trait] +impl TableProvider for AddColumnProvider { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&[usize]>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + let input = state.create_physical_plan(&self.input).await?; + Ok(Arc::new(AddColumnExec::new( + input, + self.added.clone(), + Arc::clone(&self.schema), + projection, + )?)) + } +} + +#[derive(Debug)] +struct AddColumnExec { + input: Arc, + added: Vec<(String, ArrayRef)>, + full_schema: SchemaRef, + projection: Option>, + cache: Arc, +} + +impl AddColumnExec { + fn new( + input: Arc, + added: Vec<(String, ArrayRef)>, + full_schema: SchemaRef, + projection: Option<&[usize]>, + ) -> Result { + let projected_schema = project_schema(&full_schema, projection)?; + Ok(Self { + cache: Arc::new(Self::compute_properties(projected_schema, input.as_ref())), + input, + added, + full_schema, + projection: projection.map(|p| p.to_vec()), + }) + } + + fn compute_properties( + schema: SchemaRef, + input: &dyn ExecutionPlan, + ) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + PhysicalPartitioning::UnknownPartitioning(1), + input.pipeline_behavior(), + input.boundedness(), + ) + } +} + +impl DisplayAs for AddColumnExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + let names: Vec<&str> = self.added.iter().map(|(n, _)| n.as_str()).collect(); + write!(f, "AddColumnExec: {}", names.join(", ")) + } +} + +impl ExecutionPlan for AddColumnExec { + fn name(&self) -> &'static str { + "AddColumnExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + if children.len() != 1 { + return plan_err!("AddColumnExec expects exactly one child"); + } + + let input = children.swap_remove(0); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input, + added: self.added.clone(), + full_schema: Arc::clone(&self.full_schema), + projection: self.projection.clone(), + cache: Arc::clone(&self.cache), + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(Self::new( + input, + self.added.clone(), + Arc::clone(&self.full_schema), + self.projection.as_deref(), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + if partition != 0 { + return exec_err!("AddColumnExec only produces partition 0"); + } + + // A single ArrayRef is one global vector, so the input must be one stream. + let input = if self.input.output_partitioning().partition_count() == 1 { + Arc::clone(&self.input) + } else { + Arc::new(CoalescePartitionsExec::new(Arc::clone(&self.input))) + }; + + let stream = input.execute(0, context)?; + let added = self.added.clone(); + let full_schema = Arc::clone(&self.full_schema); + let projection = self.projection.clone(); + let expected_rows = added.first().map(|(_, a)| a.len()).unwrap_or(0); + let out_schema = self.schema(); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + out_schema, + futures::stream::try_unfold((stream, 0usize), move |(mut stream, offset)| { + let added = added.clone(); + let full_schema = Arc::clone(&full_schema); + let projection = projection.clone(); + + async move { + match stream.next().await { + Some(Ok(batch)) => { + let n = batch.num_rows(); + if offset + n > expected_rows { + return exec_err!( + "Added column has {expected_rows} rows, \ + dataframe produced at least {}", + offset + n + ); + } + + let mut columns = batch.columns().to_vec(); + for (_, array) in &added { + columns.push(array.slice(offset, n)); + } + let full = RecordBatch::try_new(full_schema, columns)?; + let batch = match &projection { + Some(indices) => full.project(indices)?, + None => full, + }; + Ok(Some((batch, (stream, offset + n)))) + } + Some(Err(e)) => Err(e), + None if offset != expected_rows => exec_err!( + "Added column has {expected_rows} rows, \ + dataframe produced {offset}" + ), + None => Ok(None), + } + } + }), + ))) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } /// Create a DataFrame from column names and values. diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 5db3d23ca1c8c..74a897249ef2e 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -7363,6 +7363,134 @@ async fn test_dataframe_macro() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_dataframe_with_array_columns() -> Result<()> { + let df = dataframe!("id" => [1_i32, 2, 3])?; + + let bools: ArrayRef = Arc::new(BooleanArray::from(vec![true, false, true])); + let i8s: ArrayRef = Arc::new(Int8Array::from(vec![-1, 0, 1])); + let i16s: ArrayRef = Arc::new(Int16Array::from(vec![-1, 0, 1])); + let i32s: ArrayRef = Arc::new(Int32Array::from(vec![-1, 0, 1])); + let i64s: ArrayRef = Arc::new(Int64Array::from(vec![-1, 0, 1])); + let u8s: ArrayRef = Arc::new(UInt8Array::from(vec![0, 1, 2])); + let u16s: ArrayRef = Arc::new(UInt16Array::from(vec![0, 1, 2])); + let u32s: ArrayRef = Arc::new(UInt32Array::from(vec![0, 1, 2])); + let u64s: ArrayRef = Arc::new(UInt64Array::from(vec![0, 1, 2])); + let f16s: ArrayRef = Arc::new(Float16Array::from(vec![ + half::f16::from_f64(1.0), + half::f16::from_f64(2.0), + half::f16::from_f64(3.0), + ])); + let f32s: ArrayRef = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0])); + let f64s: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); + let strings: ArrayRef = + Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); + + let df = df.with_array_columns([ + ("bool", bools), + ("i8", i8s), + ("i16", i16s), + ("i32", i32s), + ("i64", i64s), + ("u8", u8s), + ("u16", u16s), + ("u32", u32s), + ("u64", u64s), + ("f16", f16s), + ("f32", f32s), + ("f64", f64s), + ("str", strings), + ])?; + + let expected_types = [ + ("id", DataType::Int32), + ("bool", DataType::Boolean), + ("i8", DataType::Int8), + ("i16", DataType::Int16), + ("i32", DataType::Int32), + ("i64", DataType::Int64), + ("u8", DataType::UInt8), + ("u16", DataType::UInt16), + ("u32", DataType::UInt32), + ("u64", DataType::UInt64), + ("f16", DataType::Float16), + ("f32", DataType::Float32), + ("f64", DataType::Float64), + ("str", DataType::Utf8), + ]; + + assert_eq!(df.schema().fields().len(), expected_types.len()); + assert_eq!(df.clone().count().await?, 3); + + let schema = df.schema(); + for (name, data_type) in &expected_types { + assert_eq!(schema.field_with_name(None, name)?.data_type(), data_type); + } + + let rows = df.sort(vec![col("id").sort(true, true)])?; + assert_batches_eq!( + &[ + "+----+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| id | bool | i8 | i16 | i32 | i64 | u8 | u16 | u32 | u64 | f16 | f32 | f64 | str |", + "+----+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| 1 | true | -1 | -1 | -1 | -1 | 0 | 0 | 0 | 0 | 1 | 1.0 | 1.0 | foo |", + "| 2 | false | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 | 2 | 2.0 | 2.0 | bar |", + "| 3 | true | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3.0 | 3.0 | |", + "+----+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + ], + &rows.collect().await? + ); + + Ok(()) +} + +#[tokio::test] +async fn test_dataframe_with_array_columns_then_filter() -> Result<()> { + let df = dataframe!("id" => [1, 2, 3], "data" => [42, 43, 44])?; + let new_col: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar", "baz"])); + + let df = df + .with_array_columns([("new_col", new_col)])? + .filter(col("id").gt(lit(1)))?; + + assert_eq!( + df.schema().field_with_name(None, "new_col")?.data_type(), + &DataType::Utf8 + ); + assert_eq!(df.clone().count().await?, 2); + assert_batches_eq!( + &[ + "+----+------+---------+", + "| id | data | new_col |", + "+----+------+---------+", + "| 2 | 43 | bar |", + "| 3 | 44 | baz |", + "+----+------+---------+", + ], + &df.collect().await? + ); + Ok(()) +} + +#[test] +fn test_dataframe_with_array_columns_duplicate_name() -> Result<()> { + let df = dataframe!("id" => [1, 2, 3])?; + let id: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); + let err = df.with_array_columns([("id", id)]).unwrap_err(); + assert!(err.to_string().contains("already exists")); + Ok(()) +} + +#[tokio::test] +async fn test_dataframe_with_array_columns_length_mismatch() -> Result<()> { + let df = dataframe!("id" => [1, 2, 3])?; + let too_short: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let df = df.with_array_columns([("extra", too_short)])?; + let err = df.collect().await.unwrap_err(); + assert!(err.to_string().contains("rows")); + Ok(()) +} + #[tokio::test] async fn test_copy_schema() -> Result<()> { let tmp_dir = TempDir::new()?; From ef4c56079b80ede5aeee2d2916dc2eab6f2ce9a7 Mon Sep 17 00:00:00 2001 From: Sergey Zhukov Date: Wed, 16 Sep 2026 19:33:42 +0200 Subject: [PATCH 2/2] reuse with_column method --- datafusion/core/src/dataframe/mod.rs | 277 +++++++++++++++++-------- datafusion/core/src/prelude.rs | 2 +- datafusion/core/tests/dataframe/mod.rs | 45 ++-- 3 files changed, 219 insertions(+), 105 deletions(-) diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 643088d26d684..466d60861eb12 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -22,19 +22,19 @@ mod parquet; use crate::arrow::record_batch::RecordBatch; use crate::arrow::util::pretty; -use crate::datasource::file_format::csv::CsvFormatFactory; -use crate::datasource::file_format::format_as_file_type; -use crate::datasource::file_format::json::JsonFormatFactory; use crate::datasource::{ - DefaultTableSource, MemTable, TableProvider, provider_as_source, + DefaultTableSource, MemTable, TableProvider, + file_format::{csv::CsvFormatFactory, format_as_file_type, json::JsonFormatFactory}, + provider_as_source, }; use crate::error::Result; -use crate::execution::FunctionRegistry; -use crate::execution::context::{SessionState, TaskContext}; -use crate::logical_expr::utils::find_window_exprs; +use crate::execution::{ + FunctionRegistry, + context::{SessionState, TaskContext}, +}; use crate::logical_expr::{ Expr, JoinType, LogicalPlan, LogicalPlanBuilder, LogicalPlanBuilderOptions, - Partitioning, TableType, col, ident, + Partitioning, TableType, col, ident, utils::find_window_exprs, }; use crate::physical_expr::EquivalenceProperties; use crate::physical_plan::{ @@ -45,39 +45,40 @@ use crate::physical_plan::{ execute_stream, execute_stream_partitioned, stream::RecordBatchStreamAdapter, }; use crate::prelude::SessionContext; -use std::borrow::Cow; -use std::collections::{HashMap, HashSet}; -use std::fmt::{self, Formatter}; -use std::sync::Arc; - use arrow::array::{Array, ArrayRef, Int64Array, StringArray}; use arrow::compute::{cast, concat}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::util::display::{ArrayFormatter, FormatOptions}; use arrow_schema::FieldRef; -use datafusion_common::config::{CsvOptions, JsonOptions}; +use datafusion_catalog::Session; use datafusion_common::{ Column, DFSchema, DataFusionError, ParamValues, ScalarValue, SchemaError, - TableReference, UnnestOptions, exec_err, internal_datafusion_err, not_impl_err, - plan_datafusion_err, plan_err, project_schema, tree_node::TreeNodeRecursion, + TableReference, UnnestOptions, + config::{CsvOptions, JsonOptions}, + exec_err, internal_datafusion_err, not_impl_err, plan_datafusion_err, plan_err, + project_schema, + tree_node::TreeNodeRecursion, unqualified_field_not_found, }; -use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{ - ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, case, - dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION, + ColumnarValue, ExplainOption, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, + Signature, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, Volatility, case, + dml::InsertOp, extension_types::DFArrayFormatterFactory, is_null, lit, + select_expr::SelectExpr, utils::COUNT_STAR_EXPANSION, }; -use datafusion_functions::core::coalesce; -use datafusion_functions::math::nanvl; +use datafusion_functions::{core::coalesce, math::nanvl}; use datafusion_functions_aggregate::expr_fn::{ avg, count, max, median, min, stddev, sum, }; use async_trait::async_trait; -use datafusion_catalog::Session; -use datafusion_expr::extension_types::DFArrayFormatterFactory; use futures::StreamExt; use futures::future::BoxFuture; +use std::borrow::Cow; +use std::collections::{HashMap, HashSet}; +use std::fmt::{self, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; /// Contains options that control how data is /// written out from a DataFrame @@ -312,6 +313,7 @@ impl DataFrame { /// Filter the DataFrame by column. Returns a new DataFrame only containing the /// specified columns. /// + /// # Example /// ``` /// # use datafusion::prelude::*; /// # use datafusion::error::Result; @@ -2289,21 +2291,46 @@ impl DataFrame { /// Add or replace a column in the DataFrame. /// + /// The column can be created from a DataFusion expression or from a + /// pre-materialized in-memory Arrow array. + /// /// # Example /// ``` + /// # use std::sync::Arc; + /// # use arrow::array::{ArrayRef, Int32Array}; /// # use datafusion::prelude::*; /// # use datafusion::error::Result; + /// # use datafusion_common::assert_batches_sorted_eq; /// # #[tokio::main] /// # async fn main() -> Result<()> { /// let ctx = SessionContext::new(); /// let df = ctx /// .read_csv("tests/data/example.csv", CsvReadOptions::new()) /// .await?; + /// + /// // Add a column computed from existing DataFrame columns. /// let df = df.with_column("ab_sum", col("a") + col("b"))?; + /// + /// // Add a column from a pre-materialized Arrow array. + /// let values: ArrayRef = Arc::new(Int32Array::from(vec![42])); + /// let df = df.with_column("extra", array_col(values))?; + /// + /// let expected = vec![ + /// "+---+---+---+--------+-------+", + /// "| a | b | c | ab_sum | extra |", + /// "+---+---+---+--------+-------+", + /// "| 1 | 2 | 3 | 3 | 42 |", + /// "+---+---+---+--------+-------+", + /// ]; + /// # assert_batches_sorted_eq!(expected, &df.collect().await?); /// # Ok(()) /// # } /// ``` pub fn with_column(self, name: &str, expr: Expr) -> Result { + if let Some(array) = take_array_col(&expr) { + return self.with_array_column(name, array); + } + let window_func_exprs = find_window_exprs([&expr]); let original_names: HashSet = self @@ -2753,20 +2780,20 @@ impl DataFrame { Ok(df) } - /// Append named Arrow arrays as columns to this [`DataFrame`]. + /// Append an Arrow array as a column to this [`DataFrame`]. /// - /// This does not execute the current plan. The arrays are attached when the - /// returned DataFrame is collected. Each array must have the same length as + /// This does not execute the current plan. The array is attached when the + /// returned DataFrame is collected. The array must have the same length as /// the number of rows this DataFrame produces. /// - /// See [`Self::with_column`] to add a column from an [`Expr`]. + /// Called from [`Self::with_column`] when the expression is [`array_col`]. /// /// # Example /// /// ``` /// use std::sync::Arc; /// use arrow::array::{ArrayRef, Int32Array, StringArray}; - /// use datafusion::prelude::DataFrame; + /// use datafusion::prelude::*; /// # use datafusion::error::Result; /// # use datafusion_common::assert_batches_sorted_eq; /// # #[tokio::main] @@ -2776,7 +2803,7 @@ impl DataFrame { /// let df = DataFrame::from_columns([("id", id), ("name", name)])?; /// /// let extra: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); - /// let df = df.with_array_columns([("extra", extra)])?; + /// let df = df.with_column("extra", array_col(extra))?; /// /// let expected = vec![ /// "+----+------+-------+", @@ -2791,50 +2818,141 @@ impl DataFrame { /// # Ok(()) /// # } /// ``` - pub fn with_array_columns<'a, I>(self, columns: I) -> Result - where - I: IntoIterator, - { - let added: Vec<(String, ArrayRef)> = columns - .into_iter() - .map(|(name, array)| (name.to_string(), array)) - .collect(); - - if added.is_empty() { - return Ok(self); - } - - if let Some((_, first)) = added.first() { - let expected = first.len(); - for (name, array) in &added { - if array.len() != expected { - return plan_err!( - "Column '{name}' has length {}, expected {expected}", - array.len() - ); - } - } - } - + fn with_array_column(self, name: &str, array: ArrayRef) -> Result { let (state, plan) = self.into_parts(); - let provider = Arc::new(AddColumnProvider::try_new(plan, added)?); + let provider = Arc::new(AddColumnProvider::try_new(plan, name, array)?); let plan = LogicalPlanBuilder::scan(UNNAMED_TABLE, provider_as_source(provider), None)? .build()?; - Ok(DataFrame::new(state, plan)) } } +/// Create an [`Expr`] that attaches an Arrow [`ArrayRef`] as a DataFrame column. +/// +/// Pass the result to [`DataFrame::with_column`]. Unlike [`lit`] or nested +/// constructors such as make_array, values are zipped onto existing rows +/// positionally. +/// +/// # Example +/// +/// ``` +/// use std::sync::Arc; +/// use arrow::array::{ArrayRef, Int32Array, StringArray}; +/// use datafusion::prelude::*; +/// # use datafusion::error::Result; +/// # use datafusion_common::assert_batches_sorted_eq; +/// # #[tokio::main] +/// # 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([("id", id), ("name", name)])?; +/// +/// let extra: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); +/// let df = df.with_column("extra", array_col(extra))?; +/// +/// let expected = vec![ +/// "+----+------+-------+", +/// "| id | name | extra |", +/// "+----+------+-------+", +/// "| 1 | foo | 10 |", +/// "| 2 | bar | 20 |", +/// "| 3 | baz | 30 |", +/// "+----+------+-------+", +/// ]; +/// # assert_batches_sorted_eq!(expected, &df.collect().await?); +/// # Ok(()) +/// # } +/// ``` +pub fn array_col(array: ArrayRef) -> Expr { + ScalarUDF::new_from_impl(ArrayCol::new(array)).call(vec![]) +} + +/// Marker UDF used by [`array_col`]. [`DataFrame::with_column`] intercepts this +/// expression and zips the array onto existing rows. Evaluating it as a normal +/// scalar function is not supported. +#[derive(Debug)] +struct ArrayCol { + array: ArrayRef, + signature: Signature, +} + +impl ArrayCol { + fn new(array: ArrayRef) -> Self { + Self { + array, + signature: Signature::nullary(Volatility::Volatile), + } + } +} + +impl PartialEq for ArrayCol { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.array, &other.array) + } +} + +impl Eq for ArrayCol {} + +impl Hash for ArrayCol { + fn hash(&self, state: &mut H) { + self.array.data_type().hash(state); + self.array.len().hash(state); + Arc::as_ptr(&self.array).cast::<()>().hash(state); + } +} + +impl ScalarUDFImpl for ArrayCol { + fn name(&self) -> &str { + "array_col" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(self.array.data_type().clone()) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + exec_err!("array_col() can only be used with DataFrame::with_column") + } + + fn should_evaluate_const(&self, _args: &[&ScalarValue]) -> bool { + false + } +} + +/// Extract the array from an [`array_col`] expression, unwrapping aliases. +fn take_array_col(mut expr: &Expr) -> Option { + while let Expr::Alias(alias) = expr { + expr = alias.expr.as_ref(); + } + match expr { + Expr::ScalarFunction(func) if func.args.is_empty() => func + .func + .inner() + .downcast_ref::() + .map(|array_col| Arc::clone(&array_col.array)), + _ => None, + } +} + #[derive(Debug)] struct AddColumnProvider { input: LogicalPlan, - added: Vec<(String, ArrayRef)>, + name: String, + array: ArrayRef, schema: SchemaRef, } impl AddColumnProvider { - fn try_new(input: LogicalPlan, added: Vec<(String, ArrayRef)>) -> Result { + fn try_new(input: LogicalPlan, name: &str, array: ArrayRef) -> Result { + if input.schema().has_column_with_unqualified_name(name) { + return plan_err!("Column '{name}' already exists"); + } + let mut fields: Vec = input .schema() .as_arrow() @@ -2842,17 +2960,12 @@ impl AddColumnProvider { .iter() .map(|f| f.as_ref().clone()) .collect(); - - for (name, array) in &added { - if input.schema().has_column_with_unqualified_name(name) { - return plan_err!("Column '{name}' already exists"); - } - fields.push(Field::new(name, array.data_type().clone(), true)); - } + fields.push(Field::new(name, array.data_type().clone(), true)); Ok(Self { input, - added, + name: name.to_string(), + array, schema: Arc::new(Schema::new(fields)), }) } @@ -2878,7 +2991,8 @@ impl TableProvider for AddColumnProvider { let input = state.create_physical_plan(&self.input).await?; Ok(Arc::new(AddColumnExec::new( input, - self.added.clone(), + self.name.clone(), + Arc::clone(&self.array), Arc::clone(&self.schema), projection, )?)) @@ -2888,7 +3002,8 @@ impl TableProvider for AddColumnProvider { #[derive(Debug)] struct AddColumnExec { input: Arc, - added: Vec<(String, ArrayRef)>, + name: String, + array: ArrayRef, full_schema: SchemaRef, projection: Option>, cache: Arc, @@ -2897,7 +3012,8 @@ struct AddColumnExec { impl AddColumnExec { fn new( input: Arc, - added: Vec<(String, ArrayRef)>, + name: String, + array: ArrayRef, full_schema: SchemaRef, projection: Option<&[usize]>, ) -> Result { @@ -2905,7 +3021,8 @@ impl AddColumnExec { Ok(Self { cache: Arc::new(Self::compute_properties(projected_schema, input.as_ref())), input, - added, + name, + array, full_schema, projection: projection.map(|p| p.to_vec()), }) @@ -2926,8 +3043,7 @@ impl AddColumnExec { impl DisplayAs for AddColumnExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { - let names: Vec<&str> = self.added.iter().map(|(n, _)| n.as_str()).collect(); - write!(f, "AddColumnExec: {}", names.join(", ")) + write!(f, "AddColumnExec: {}", self.name) } } @@ -2961,14 +3077,16 @@ impl ExecutionPlan for AddColumnExec { match options.children_properties { ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { input, - added: self.added.clone(), + name: self.name.clone(), + array: Arc::clone(&self.array), full_schema: Arc::clone(&self.full_schema), projection: self.projection.clone(), cache: Arc::clone(&self.cache), })), ChildrenPropertiesMode::Recompute => Ok(Arc::new(Self::new( input, - self.added.clone(), + self.name.clone(), + Arc::clone(&self.array), Arc::clone(&self.full_schema), self.projection.as_deref(), )?)), @@ -3002,16 +3120,16 @@ impl ExecutionPlan for AddColumnExec { }; let stream = input.execute(0, context)?; - let added = self.added.clone(); + let array = Arc::clone(&self.array); let full_schema = Arc::clone(&self.full_schema); let projection = self.projection.clone(); - let expected_rows = added.first().map(|(_, a)| a.len()).unwrap_or(0); + let expected_rows = array.len(); let out_schema = self.schema(); Ok(Box::pin(RecordBatchStreamAdapter::new( out_schema, futures::stream::try_unfold((stream, 0usize), move |(mut stream, offset)| { - let added = added.clone(); + let array = Arc::clone(&array); let full_schema = Arc::clone(&full_schema); let projection = projection.clone(); @@ -3026,11 +3144,8 @@ impl ExecutionPlan for AddColumnExec { offset + n ); } - let mut columns = batch.columns().to_vec(); - for (_, array) in &added { - columns.push(array.slice(offset, n)); - } + columns.push(array.slice(offset, n)); let full = RecordBatch::try_new(full_schema, columns)?; let batch = match &projection { Some(indices) => full.project(indices)?, diff --git a/datafusion/core/src/prelude.rs b/datafusion/core/src/prelude.rs index 31d9d7eb471f0..05643717bfda8 100644 --- a/datafusion/core/src/prelude.rs +++ b/datafusion/core/src/prelude.rs @@ -26,7 +26,7 @@ //! ``` pub use crate::dataframe; -pub use crate::dataframe::DataFrame; +pub use crate::dataframe::{DataFrame, array_col}; pub use crate::execution::context::{SQLOptions, SessionConfig, SessionContext}; pub use crate::execution::options::{ AvroReadOptions, CsvReadOptions, JsonReadOptions, ParquetReadOptions, diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 74a897249ef2e..2666ac773b65b 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -59,7 +59,7 @@ use datafusion::error::Result; use datafusion::execution::context::SessionContext; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::logical_expr::{ColumnarValue, Volatility}; -use datafusion::prelude::{CsvReadOptions, JoinType, ParquetReadOptions}; +use datafusion::prelude::{CsvReadOptions, JoinType, ParquetReadOptions, array_col}; use datafusion::test_util::{ parquet_test_data, populate_csv_partitions, register_aggregate_csv, test_table, test_table_with_cache_factory, test_table_with_name, @@ -7364,7 +7364,7 @@ async fn test_dataframe_macro() -> Result<()> { } #[tokio::test] -async fn test_dataframe_with_array_columns() -> Result<()> { +async fn test_dataframe_with_column_array_col() -> Result<()> { let df = dataframe!("id" => [1_i32, 2, 3])?; let bools: ArrayRef = Arc::new(BooleanArray::from(vec![true, false, true])); @@ -7386,21 +7386,20 @@ async fn test_dataframe_with_array_columns() -> Result<()> { let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); - let df = df.with_array_columns([ - ("bool", bools), - ("i8", i8s), - ("i16", i16s), - ("i32", i32s), - ("i64", i64s), - ("u8", u8s), - ("u16", u16s), - ("u32", u32s), - ("u64", u64s), - ("f16", f16s), - ("f32", f32s), - ("f64", f64s), - ("str", strings), - ])?; + let df = df + .with_column("bool", array_col(bools))? + .with_column("i8", array_col(i8s))? + .with_column("i16", array_col(i16s))? + .with_column("i32", array_col(i32s))? + .with_column("i64", array_col(i64s))? + .with_column("u8", array_col(u8s))? + .with_column("u16", array_col(u16s))? + .with_column("u32", array_col(u32s))? + .with_column("u64", array_col(u64s))? + .with_column("f16", array_col(f16s))? + .with_column("f32", array_col(f32s))? + .with_column("f64", array_col(f64s))? + .with_column("str", array_col(strings))?; let expected_types = [ ("id", DataType::Int32), @@ -7445,12 +7444,12 @@ async fn test_dataframe_with_array_columns() -> Result<()> { } #[tokio::test] -async fn test_dataframe_with_array_columns_then_filter() -> Result<()> { +async fn test_dataframe_with_column_array_col_then_filter() -> Result<()> { let df = dataframe!("id" => [1, 2, 3], "data" => [42, 43, 44])?; let new_col: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar", "baz"])); let df = df - .with_array_columns([("new_col", new_col)])? + .with_column("new_col", array_col(new_col))? .filter(col("id").gt(lit(1)))?; assert_eq!( @@ -7473,19 +7472,19 @@ async fn test_dataframe_with_array_columns_then_filter() -> Result<()> { } #[test] -fn test_dataframe_with_array_columns_duplicate_name() -> Result<()> { +fn test_dataframe_with_column_array_col_duplicate_name() -> Result<()> { let df = dataframe!("id" => [1, 2, 3])?; let id: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); - let err = df.with_array_columns([("id", id)]).unwrap_err(); + let err = df.with_column("id", array_col(id)).unwrap_err(); assert!(err.to_string().contains("already exists")); Ok(()) } #[tokio::test] -async fn test_dataframe_with_array_columns_length_mismatch() -> Result<()> { +async fn test_dataframe_with_column_array_col_length_mismatch() -> Result<()> { let df = dataframe!("id" => [1, 2, 3])?; let too_short: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); - let df = df.with_array_columns([("extra", too_short)])?; + let df = df.with_column("extra", array_col(too_short))?; let err = df.collect().await.unwrap_err(); assert!(err.to_string().contains("rows")); Ok(())