diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 55d3695a68d7a..466d60861eb12 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -22,55 +22,63 @@ 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::{ - 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::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, unqualified_field_not_found, + 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 @@ -305,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; @@ -2282,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 @@ -2745,6 +2779,398 @@ impl DataFrame { let df = ctx.read_batch(batch)?; Ok(df) } + + /// Append an Arrow array as a column to this [`DataFrame`]. + /// + /// 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. + /// + /// 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::*; + /// # 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(()) + /// # } + /// ``` + fn with_array_column(self, name: &str, array: ArrayRef) -> Result { + let (state, plan) = self.into_parts(); + 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, + name: String, + array: ArrayRef, + schema: SchemaRef, +} + +impl AddColumnProvider { + 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() + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + fields.push(Field::new(name, array.data_type().clone(), true)); + + Ok(Self { + input, + name: name.to_string(), + array, + 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.name.clone(), + Arc::clone(&self.array), + Arc::clone(&self.schema), + projection, + )?)) + } +} + +#[derive(Debug)] +struct AddColumnExec { + input: Arc, + name: String, + array: ArrayRef, + full_schema: SchemaRef, + projection: Option>, + cache: Arc, +} + +impl AddColumnExec { + fn new( + input: Arc, + name: String, + array: 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, + name, + array, + 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 { + write!(f, "AddColumnExec: {}", self.name) + } +} + +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, + 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.name.clone(), + Arc::clone(&self.array), + 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 array = Arc::clone(&self.array); + let full_schema = Arc::clone(&self.full_schema); + let projection = self.projection.clone(); + 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 array = Arc::clone(&array); + 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(); + 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/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 5db3d23ca1c8c..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, @@ -7363,6 +7363,133 @@ async fn test_dataframe_macro() -> Result<()> { Ok(()) } +#[tokio::test] +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])); + 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_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), + ("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_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_column("new_col", array_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_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_column("id", array_col(id)).unwrap_err(); + assert!(err.to_string().contains("already exists")); + Ok(()) +} + +#[tokio::test] +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_column("extra", array_col(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()?;