Skip to content

Commit

Permalink
wire in and add a test
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Feb 2, 2024
1 parent ed3906c commit 200508b
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 1 deletion.
6 changes: 6 additions & 0 deletions datafusion/core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -521,6 +521,12 @@ pub mod functions {
pub use datafusion_functions::*;
}

/// re-export of [`datafusion_functions_array`] crate, if "array_expressions" feature is enabled
pub mod functions_array {
#[cfg(feature = "array_expressions")]
pub use datafusion_functions::*;
}

#[cfg(test)]
pub mod test;
pub mod test_util;
Expand Down
2 changes: 2 additions & 0 deletions datafusion/core/src/prelude.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,8 @@ pub use datafusion_expr::{
Expr,
};
pub use datafusion_functions::expr_fn::*;
#[cfg(feature = "array_expressions")]
pub use datafusion_functions_array::expr_fn::*;

pub use std::ops::Not;
pub use std::ops::{Add, Div, Mul, Neg, Rem, Sub};
Expand Down
30 changes: 29 additions & 1 deletion datafusion/core/tests/dataframe/dataframe_functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@ use arrow::{
array::{Int32Array, StringArray},
record_batch::RecordBatch,
};
use arrow_array::types::Int32Type;
use arrow_array::ListArray;
use arrow_schema::SchemaRef;
use std::sync::Arc;

Expand All @@ -40,6 +42,7 @@ fn test_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::Utf8, false),
Field::new("b", DataType::Int32, false),
Field::new("l", DataType::new_list(DataType::Int32, true), true),
]))
}

Expand All @@ -57,6 +60,12 @@ async fn create_test_table() -> Result<DataFrame> {
"123AbcDef",
])),
Arc::new(Int32Array::from(vec![1, 10, 10, 100])),
Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
Some(vec![Some(0), Some(1), Some(2)]),
None,
Some(vec![Some(3), None, Some(5)]),
Some(vec![Some(6), Some(7)]),
])),
],
)?;

Expand All @@ -67,7 +76,7 @@ async fn create_test_table() -> Result<DataFrame> {
ctx.table("test").await
}

/// Excutes an expression on the test dataframe as a select.
/// Executes an expression on the test dataframe as a select.
/// Compares formatted output of a record batch with an expected
/// vector of strings, using the assert_batch_eq! macro
macro_rules! assert_fn_batches {
Expand Down Expand Up @@ -841,3 +850,22 @@ async fn test_fn_decode() -> Result<()> {

Ok(())
}

#[tokio::test]
async fn test_fn_array_to_string() -> Result<()> {
let expr = array_to_string(col("l"), lit("***"));

let expected = [
"+-------------------------------------+",
"| array_to_string(test.l,Utf8(\"***\")) |",
"+-------------------------------------+",
"| 0***1***2 |",
"| |",
"| 3***5 |",
"| 6***7 |",
"+-------------------------------------+",
];
assert_fn_batches!(expr, expected);

Ok(())
}

0 comments on commit 200508b

Please sign in to comment.