This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 5224194944 Migrate dataframe tests to `insta` (#15262)
5224194944 is described below
commit 5224194944960c539f637991b22912bb8f61698e
Author: jsai28 <[email protected]>
AuthorDate: Mon Mar 17 12:44:03 2025 -0600
Migrate dataframe tests to `insta` (#15262)
* convert to assert_snapshot
* edit comments
* formatting
* formatting
* ran cargo fmt --all
* add get_batches_with_limit function
---
.../core/tests/dataframe/dataframe_functions.rs | 1320 +++++++++++---------
datafusion/core/tests/dataframe/describe.rs | 99 +-
2 files changed, 765 insertions(+), 654 deletions(-)
diff --git a/datafusion/core/tests/dataframe/dataframe_functions.rs
b/datafusion/core/tests/dataframe/dataframe_functions.rs
index fec3ab786f..c763d4c8de 100644
--- a/datafusion/core/tests/dataframe/dataframe_functions.rs
+++ b/datafusion/core/tests/dataframe/dataframe_functions.rs
@@ -28,13 +28,13 @@ use std::sync::Arc;
use datafusion::error::Result;
use datafusion::prelude::*;
-
-use datafusion::assert_batches_eq;
+use datafusion_common::test_util::batches_to_string;
use datafusion_common::{DFSchema, ScalarValue};
use datafusion_expr::expr::Alias;
use datafusion_expr::{table_scan, ExprSchemable, LogicalPlanBuilder};
use datafusion_functions_aggregate::expr_fn::{approx_median,
approx_percentile_cont};
use datafusion_functions_nested::map::map;
+use insta::assert_snapshot;
fn test_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
@@ -75,34 +75,32 @@ async fn create_test_table() -> Result<DataFrame> {
}
/// 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 {
- ($EXPR:expr, $EXPECTED: expr) => {
- assert_fn_batches!($EXPR, $EXPECTED, 10)
- };
- ($EXPR:expr, $EXPECTED: expr, $LIMIT: expr) => {
- let df = create_test_table().await?;
- let df = df.select(vec![$EXPR])?.limit(0, Some($LIMIT))?;
- let batches = df.collect().await?;
-
- assert_batches_eq!($EXPECTED, &batches);
- };
+async fn get_batches(expr: Expr) -> Result<Vec<RecordBatch>> {
+ get_batches_with_limit(expr, 10).await
+}
+
+async fn get_batches_with_limit(expr: Expr, limit: usize) ->
Result<Vec<RecordBatch>> {
+ let df = create_test_table().await?;
+ let df = df.select(vec![expr])?.limit(0, Some(limit))?;
+ df.collect().await
}
#[tokio::test]
async fn test_fn_ascii() -> Result<()> {
let expr = ascii(col("a"));
- let expected = [
- "+---------------+",
- "| ascii(test.a) |",
- "+---------------+",
- "| 97 |",
- "+---------------+",
- ];
-
- assert_fn_batches!(expr, expected, 1);
+ let batches = get_batches_with_limit(expr, 1).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +---------------+
+ | ascii(test.a) |
+ +---------------+
+ | 97 |
+ +---------------+
+ "
+ );
Ok(())
}
@@ -111,17 +109,21 @@ async fn test_fn_ascii() -> Result<()> {
async fn test_fn_bit_length() -> Result<()> {
let expr = bit_length(col("a"));
- let expected = [
- "+--------------------+",
- "| bit_length(test.a) |",
- "+--------------------+",
- "| 48 |",
- "| 48 |",
- "| 48 |",
- "| 72 |",
- "+--------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +--------------------+
+ | bit_length(test.a) |
+ +--------------------+
+ | 48 |
+ | 48 |
+ | 48 |
+ | 72 |
+ +--------------------+
+ "
+ );
Ok(())
}
@@ -130,15 +132,17 @@ async fn test_fn_bit_length() -> Result<()> {
async fn test_fn_btrim() -> Result<()> {
let expr = btrim(vec![lit(" a b c ")]);
- let expected = [
- "+-----------------------------------------+",
- "| btrim(Utf8(\" a b c \")) |",
- "+-----------------------------------------+",
- "| a b c |",
- "+-----------------------------------------+",
- ];
+ let batches = get_batches_with_limit(expr, 1).await?;
- assert_fn_batches!(expr, expected, 1);
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-----------------------------------------+
+ | btrim(Utf8(" a b c ")) |
+ +-----------------------------------------+
+ | a b c |
+ +-----------------------------------------+
+ "#);
Ok(())
}
@@ -147,18 +151,21 @@ async fn test_fn_btrim() -> Result<()> {
async fn test_fn_btrim_with_chars() -> Result<()> {
let expr = btrim(vec![col("a"), lit("ab")]);
- let expected = [
- "+--------------------------+",
- "| btrim(test.a,Utf8(\"ab\")) |",
- "+--------------------------+",
- "| cDEF |",
- "| c123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+--------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +--------------------------+
+ | btrim(test.a,Utf8("ab")) |
+ +--------------------------+
+ | cDEF |
+ | c123 |
+ | CBAdef |
+ | 123AbcDef |
+ +--------------------------+
+ "#
+ );
Ok(())
}
@@ -167,18 +174,20 @@ async fn test_fn_btrim_with_chars() -> Result<()> {
async fn test_fn_nullif() -> Result<()> {
let expr = nullif(col("a"), lit("abcDEF"));
- let expected = [
- "+-------------------------------+",
- "| nullif(test.a,Utf8(\"abcDEF\")) |",
- "+-------------------------------+",
- "| |",
- "| abc123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+-------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-------------------------------+
+ | nullif(test.a,Utf8("abcDEF")) |
+ +-------------------------------+
+ | |
+ | abc123 |
+ | CBAdef |
+ | 123AbcDef |
+ +-------------------------------+
+ "#);
Ok(())
}
@@ -187,18 +196,20 @@ async fn test_fn_nullif() -> Result<()> {
async fn test_fn_arrow_cast() -> Result<()> {
let expr = arrow_typeof(arrow_cast(col("b"), lit("Float64")));
- let expected = [
- "+--------------------------------------------------+",
- "| arrow_typeof(arrow_cast(test.b,Utf8(\"Float64\"))) |",
- "+--------------------------------------------------+",
- "| Float64 |",
- "| Float64 |",
- "| Float64 |",
- "| Float64 |",
- "+--------------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +--------------------------------------------------+
+ | arrow_typeof(arrow_cast(test.b,Utf8("Float64"))) |
+ +--------------------------------------------------+
+ | Float64 |
+ | Float64 |
+ | Float64 |
+ | Float64 |
+ +--------------------------------------------------+
+ "#);
Ok(())
}
@@ -215,21 +226,24 @@ async fn test_nvl() -> Result<()> {
)
.alias("nvl_expr");
- let expected = [
- "+-------------+",
- "| nvl_expr |",
- "+-------------+",
- "| TURNED_NULL |",
- "| abc123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+-------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-------------+
+ | nvl_expr |
+ +-------------+
+ | TURNED_NULL |
+ | abc123 |
+ | CBAdef |
+ | 123AbcDef |
+ +-------------+
+ ");
Ok(())
}
+
#[tokio::test]
async fn test_nvl2() -> Result<()> {
let lit_null = lit(ScalarValue::Utf8(None));
@@ -243,18 +257,20 @@ async fn test_nvl2() -> Result<()> {
)
.alias("nvl2_expr");
- let expected = [
- "+-------------+",
- "| nvl2_expr |",
- "+-------------+",
- "| TURNED_NULL |",
- "| NON_NULL |",
- "| NON_NULL |",
- "| NON_NULL |",
- "+-------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-------------+
+ | nvl2_expr |
+ +-------------+
+ | TURNED_NULL |
+ | NON_NULL |
+ | NON_NULL |
+ | NON_NULL |
+ +-------------+
+ ");
Ok(())
}
@@ -262,18 +278,20 @@ async fn test_nvl2() -> Result<()> {
async fn test_fn_arrow_typeof() -> Result<()> {
let expr = arrow_typeof(col("l"));
- let expected = [
-
"+------------------------------------------------------------------------------------------------------------------+",
- "| arrow_typeof(test.l)
|",
-
"+------------------------------------------------------------------------------------------------------------------+",
- "| List(Field { name: \"item\", data_type: Int32, nullable: true,
dict_id: 0, dict_is_ordered: false, metadata: {} }) |",
- "| List(Field { name: \"item\", data_type: Int32, nullable: true,
dict_id: 0, dict_is_ordered: false, metadata: {} }) |",
- "| List(Field { name: \"item\", data_type: Int32, nullable: true,
dict_id: 0, dict_is_ordered: false, metadata: {} }) |",
- "| List(Field { name: \"item\", data_type: Int32, nullable: true,
dict_id: 0, dict_is_ordered: false, metadata: {} }) |",
-
"+------------------------------------------------------------------------------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+
+------------------------------------------------------------------------------------------------------------------+
+ | arrow_typeof(test.l)
|
+
+------------------------------------------------------------------------------------------------------------------+
+ | List(Field { name: "item", data_type: Int32, nullable: true, dict_id: 0,
dict_is_ordered: false, metadata: {} }) |
+ | List(Field { name: "item", data_type: Int32, nullable: true, dict_id: 0,
dict_is_ordered: false, metadata: {} }) |
+ | List(Field { name: "item", data_type: Int32, nullable: true, dict_id: 0,
dict_is_ordered: false, metadata: {} }) |
+ | List(Field { name: "item", data_type: Int32, nullable: true, dict_id: 0,
dict_is_ordered: false, metadata: {} }) |
+
+------------------------------------------------------------------------------------------------------------------+
+ "#);
Ok(())
}
@@ -282,18 +300,20 @@ async fn test_fn_arrow_typeof() -> Result<()> {
async fn test_fn_struct() -> Result<()> {
let expr = r#struct(vec![col("a"), col("b")]);
- let expected = [
- "+--------------------------+",
- "| struct(test.a,test.b) |",
- "+--------------------------+",
- "| {c0: abcDEF, c1: 1} |",
- "| {c0: abc123, c1: 10} |",
- "| {c0: CBAdef, c1: 10} |",
- "| {c0: 123AbcDef, c1: 100} |",
- "+--------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +--------------------------+
+ | struct(test.a,test.b) |
+ +--------------------------+
+ | {c0: abcDEF, c1: 1} |
+ | {c0: abc123, c1: 10} |
+ | {c0: CBAdef, c1: 10} |
+ | {c0: 123AbcDef, c1: 100} |
+ +--------------------------+
+ ");
Ok(())
}
@@ -302,18 +322,20 @@ async fn test_fn_struct() -> Result<()> {
async fn test_fn_named_struct() -> Result<()> {
let expr = named_struct(vec![lit("column_a"), col("a"), lit("column_b"),
col("b")]);
- let expected = [
- "+---------------------------------------------------------------+",
- "| named_struct(Utf8(\"column_a\"),test.a,Utf8(\"column_b\"),test.b)
|",
- "+---------------------------------------------------------------+",
- "| {column_a: abcDEF, column_b: 1} |",
- "| {column_a: abc123, column_b: 10} |",
- "| {column_a: CBAdef, column_b: 10} |",
- "| {column_a: 123AbcDef, column_b: 100} |",
- "+---------------------------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +---------------------------------------------------------------+
+ | named_struct(Utf8("column_a"),test.a,Utf8("column_b"),test.b) |
+ +---------------------------------------------------------------+
+ | {column_a: abcDEF, column_b: 1} |
+ | {column_a: abc123, column_b: 10} |
+ | {column_a: CBAdef, column_b: 10} |
+ | {column_a: 123AbcDef, column_b: 100} |
+ +---------------------------------------------------------------+
+ "#);
Ok(())
}
@@ -322,18 +344,20 @@ async fn test_fn_named_struct() -> Result<()> {
async fn test_fn_coalesce() -> Result<()> {
let expr = coalesce(vec![lit(ScalarValue::Utf8(None)), lit("ab")]);
- let expected = [
- "+---------------------------------+",
- "| coalesce(Utf8(NULL),Utf8(\"ab\")) |",
- "+---------------------------------+",
- "| ab |",
- "| ab |",
- "| ab |",
- "| ab |",
- "+---------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +---------------------------------+
+ | coalesce(Utf8(NULL),Utf8("ab")) |
+ +---------------------------------+
+ | ab |
+ | ab |
+ | ab |
+ | ab |
+ +---------------------------------+
+ "#);
Ok(())
}
@@ -342,18 +366,18 @@ async fn test_fn_coalesce() -> Result<()> {
async fn test_fn_approx_median() -> Result<()> {
let expr = approx_median(col("b"));
- let expected = [
- "+-----------------------+",
- "| approx_median(test.b) |",
- "+-----------------------+",
- "| 10 |",
- "+-----------------------+",
- ];
-
let df = create_test_table().await?;
let batches = df.aggregate(vec![], vec![expr]).unwrap().collect().await?;
- assert_batches_eq!(expected, &batches);
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-----------------------+
+ | approx_median(test.b) |
+ +-----------------------+
+ | 10 |
+ +-----------------------+
+ ");
Ok(())
}
@@ -362,18 +386,18 @@ async fn test_fn_approx_median() -> Result<()> {
async fn test_fn_approx_percentile_cont() -> Result<()> {
let expr = approx_percentile_cont(col("b"), lit(0.5), None);
- let expected = [
- "+---------------------------------------------+",
- "| approx_percentile_cont(test.b,Float64(0.5)) |",
- "+---------------------------------------------+",
- "| 10 |",
- "+---------------------------------------------+",
- ];
-
let df = create_test_table().await?;
let batches = df.aggregate(vec![], vec![expr]).unwrap().collect().await?;
- assert_batches_eq!(expected, &batches);
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +---------------------------------------------+
+ | approx_percentile_cont(test.b,Float64(0.5)) |
+ +---------------------------------------------+
+ | 10 |
+ +---------------------------------------------+
+ ");
// the arg2 parameter is a complex expr, but it can be evaluated to the
literal value
let alias_expr = Expr::Alias(Alias::new(
@@ -383,31 +407,34 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
));
let expr = approx_percentile_cont(col("b"), alias_expr, None);
let df = create_test_table().await?;
- let expected = [
- "+--------------------------------------+",
- "| approx_percentile_cont(test.b,arg_2) |",
- "+--------------------------------------+",
- "| 10 |",
- "+--------------------------------------+",
- ];
let batches = df.aggregate(vec![], vec![expr]).unwrap().collect().await?;
- assert_batches_eq!(expected, &batches);
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +--------------------------------------+
+ | approx_percentile_cont(test.b,arg_2) |
+ +--------------------------------------+
+ | 10 |
+ +--------------------------------------+
+ "
+ );
// with number of centroids set
let expr = approx_percentile_cont(col("b"), lit(0.5), Some(lit(2)));
- let expected = [
- "+------------------------------------------------------+",
- "| approx_percentile_cont(test.b,Float64(0.5),Int32(2)) |",
- "+------------------------------------------------------+",
- "| 30 |",
- "+------------------------------------------------------+",
- ];
let df = create_test_table().await?;
let batches = df.aggregate(vec![], vec![expr]).unwrap().collect().await?;
- assert_batches_eq!(expected, &batches);
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +------------------------------------------------------+
+ | approx_percentile_cont(test.b,Float64(0.5),Int32(2)) |
+ +------------------------------------------------------+
+ | 30 |
+ +------------------------------------------------------+
+ ");
Ok(())
}
@@ -417,18 +444,20 @@ async fn test_fn_approx_percentile_cont() -> Result<()> {
async fn test_fn_character_length() -> Result<()> {
let expr = character_length(col("a"));
- let expected = [
- "+--------------------------+",
- "| character_length(test.a) |",
- "+--------------------------+",
- "| 6 |",
- "| 6 |",
- "| 6 |",
- "| 9 |",
- "+--------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +--------------------------+
+ | character_length(test.a) |
+ +--------------------------+
+ | 6 |
+ | 6 |
+ | 6 |
+ | 9 |
+ +--------------------------+
+ ");
Ok(())
}
@@ -437,15 +466,17 @@ async fn test_fn_character_length() -> Result<()> {
async fn test_fn_chr() -> Result<()> {
let expr = chr(lit(128175));
- let expected = [
- "+--------------------+",
- "| chr(Int32(128175)) |",
- "+--------------------+",
- "| 💯 |",
- "+--------------------+",
- ];
+ let batches = get_batches_with_limit(expr, 1).await?;
- assert_fn_batches!(expr, expected, 1);
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +--------------------+
+ | chr(Int32(128175)) |
+ +--------------------+
+ | 💯 |
+ +--------------------+
+ ");
Ok(())
}
@@ -454,18 +485,20 @@ async fn test_fn_chr() -> Result<()> {
async fn test_fn_initcap() -> Result<()> {
let expr = initcap(col("a"));
- let expected = [
- "+-----------------+",
- "| initcap(test.a) |",
- "+-----------------+",
- "| Abcdef |",
- "| Abc123 |",
- "| Cbadef |",
- "| 123abcdef |",
- "+-----------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-----------------+
+ | initcap(test.a) |
+ +-----------------+
+ | Abcdef |
+ | Abc123 |
+ | Cbadef |
+ | 123abcdef |
+ +-----------------+
+ ");
Ok(())
}
@@ -475,18 +508,20 @@ async fn test_fn_initcap() -> Result<()> {
async fn test_fn_left() -> Result<()> {
let expr = left(col("a"), lit(3));
- let expected = [
- "+-----------------------+",
- "| left(test.a,Int32(3)) |",
- "+-----------------------+",
- "| abc |",
- "| abc |",
- "| CBA |",
- "| 123 |",
- "+-----------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-----------------------+
+ | left(test.a,Int32(3)) |
+ +-----------------------+
+ | abc |
+ | abc |
+ | CBA |
+ | 123 |
+ +-----------------------+
+ ");
Ok(())
}
@@ -495,18 +530,20 @@ async fn test_fn_left() -> Result<()> {
async fn test_fn_lower() -> Result<()> {
let expr = lower(col("a"));
- let expected = [
- "+---------------+",
- "| lower(test.a) |",
- "+---------------+",
- "| abcdef |",
- "| abc123 |",
- "| cbadef |",
- "| 123abcdef |",
- "+---------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +---------------+
+ | lower(test.a) |
+ +---------------+
+ | abcdef |
+ | abc123 |
+ | cbadef |
+ | 123abcdef |
+ +---------------+
+ ");
Ok(())
}
@@ -516,18 +553,20 @@ async fn test_fn_lower() -> Result<()> {
async fn test_fn_lpad() -> Result<()> {
let expr = lpad(vec![col("a"), lit(10)]);
- let expected = [
- "+------------------------+",
- "| lpad(test.a,Int32(10)) |",
- "+------------------------+",
- "| abcDEF |",
- "| abc123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +------------------------+
+ | lpad(test.a,Int32(10)) |
+ +------------------------+
+ | abcDEF |
+ | abc123 |
+ | CBAdef |
+ | 123AbcDef |
+ +------------------------+
+ ");
Ok(())
}
@@ -537,18 +576,20 @@ async fn test_fn_lpad() -> Result<()> {
async fn test_fn_lpad_with_string() -> Result<()> {
let expr = lpad(vec![col("a"), lit(10), lit("*")]);
- let expected = [
- "+----------------------------------+",
- "| lpad(test.a,Int32(10),Utf8(\"*\")) |",
- "+----------------------------------+",
- "| ****abcDEF |",
- "| ****abc123 |",
- "| ****CBAdef |",
- "| *123AbcDef |",
- "+----------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +----------------------------------+
+ | lpad(test.a,Int32(10),Utf8("*")) |
+ +----------------------------------+
+ | ****abcDEF |
+ | ****abc123 |
+ | ****CBAdef |
+ | *123AbcDef |
+ +----------------------------------+
+ "#);
Ok(())
}
@@ -557,15 +598,17 @@ async fn test_fn_lpad_with_string() -> Result<()> {
async fn test_fn_ltrim() -> Result<()> {
let expr = ltrim(vec![lit(" a b c ")]);
- let expected = [
- "+-----------------------------------------+",
- "| ltrim(Utf8(\" a b c \")) |",
- "+-----------------------------------------+",
- "| a b c |",
- "+-----------------------------------------+",
- ];
+ let batches = get_batches_with_limit(expr, 1).await?;
- assert_fn_batches!(expr, expected, 1);
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-----------------------------------------+
+ | ltrim(Utf8(" a b c ")) |
+ +-----------------------------------------+
+ | a b c |
+ +-----------------------------------------+
+ "#);
Ok(())
}
@@ -574,18 +617,20 @@ async fn test_fn_ltrim() -> Result<()> {
async fn test_fn_ltrim_with_columns() -> Result<()> {
let expr = ltrim(vec![col("a")]);
- let expected = [
- "+---------------+",
- "| ltrim(test.a) |",
- "+---------------+",
- "| abcDEF |",
- "| abc123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+---------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +---------------+
+ | ltrim(test.a) |
+ +---------------+
+ | abcDEF |
+ | abc123 |
+ | CBAdef |
+ | 123AbcDef |
+ +---------------+
+ ");
Ok(())
}
@@ -595,18 +640,20 @@ async fn test_fn_ltrim_with_columns() -> Result<()> {
async fn test_fn_md5() -> Result<()> {
let expr = md5(col("a"));
- let expected = [
- "+----------------------------------+",
- "| md5(test.a) |",
- "+----------------------------------+",
- "| ea2de8bd80f3a1f52c754214fc9b0ed1 |",
- "| e99a18c428cb38d5f260853678922e03 |",
- "| 11ed4a6e9985df40913eead67f022e27 |",
- "| 8f5e60e523c9253e623ae38bb58c399a |",
- "+----------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +----------------------------------+
+ | md5(test.a) |
+ +----------------------------------+
+ | ea2de8bd80f3a1f52c754214fc9b0ed1 |
+ | e99a18c428cb38d5f260853678922e03 |
+ | 11ed4a6e9985df40913eead67f022e27 |
+ | 8f5e60e523c9253e623ae38bb58c399a |
+ +----------------------------------+
+ ");
Ok(())
}
@@ -616,33 +663,37 @@ async fn test_fn_md5() -> Result<()> {
async fn test_fn_regexp_like() -> Result<()> {
let expr = regexp_like(col("a"), lit("[a-z]"), None);
- let expected = [
- "+-----------------------------------+",
- "| regexp_like(test.a,Utf8(\"[a-z]\")) |",
- "+-----------------------------------+",
- "| true |",
- "| true |",
- "| true |",
- "| true |",
- "+-----------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-----------------------------------+
+ | regexp_like(test.a,Utf8("[a-z]")) |
+ +-----------------------------------+
+ | true |
+ | true |
+ | true |
+ | true |
+ +-----------------------------------+
+ "#);
let expr = regexp_like(col("a"), lit("abc"), Some(lit("i")));
- let expected = [
- "+-------------------------------------------+",
- "| regexp_like(test.a,Utf8(\"abc\"),Utf8(\"i\")) |",
- "+-------------------------------------------+",
- "| true |",
- "| true |",
- "| false |",
- "| true |",
- "+-------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-------------------------------------------+
+ | regexp_like(test.a,Utf8("abc"),Utf8("i")) |
+ +-------------------------------------------+
+ | true |
+ | true |
+ | false |
+ | true |
+ +-------------------------------------------+
+ "#);
Ok(())
}
@@ -652,33 +703,37 @@ async fn test_fn_regexp_like() -> Result<()> {
async fn test_fn_regexp_match() -> Result<()> {
let expr = regexp_match(col("a"), lit("[a-z]"), None);
- let expected = [
- "+------------------------------------+",
- "| regexp_match(test.a,Utf8(\"[a-z]\")) |",
- "+------------------------------------+",
- "| [a] |",
- "| [a] |",
- "| [d] |",
- "| [b] |",
- "+------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +------------------------------------+
+ | regexp_match(test.a,Utf8("[a-z]")) |
+ +------------------------------------+
+ | [a] |
+ | [a] |
+ | [d] |
+ | [b] |
+ +------------------------------------+
+ "#);
let expr = regexp_match(col("a"), lit("[A-Z]"), Some(lit("i")));
- let expected = [
- "+----------------------------------------------+",
- "| regexp_match(test.a,Utf8(\"[A-Z]\"),Utf8(\"i\")) |",
- "+----------------------------------------------+",
- "| [a] |",
- "| [a] |",
- "| [C] |",
- "| [A] |",
- "+----------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +----------------------------------------------+
+ | regexp_match(test.a,Utf8("[A-Z]"),Utf8("i")) |
+ +----------------------------------------------+
+ | [a] |
+ | [a] |
+ | [C] |
+ | [A] |
+ +----------------------------------------------+
+ "#);
Ok(())
}
@@ -688,33 +743,37 @@ async fn test_fn_regexp_match() -> Result<()> {
async fn test_fn_regexp_replace() -> Result<()> {
let expr = regexp_replace(col("a"), lit("[a-z]"), lit("x"),
Some(lit("g")));
- let expected = [
- "+----------------------------------------------------------+",
- "| regexp_replace(test.a,Utf8(\"[a-z]\"),Utf8(\"x\"),Utf8(\"g\")) |",
- "+----------------------------------------------------------+",
- "| xxxDEF |",
- "| xxx123 |",
- "| CBAxxx |",
- "| 123AxxDxx |",
- "+----------------------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +----------------------------------------------------------+
+ | regexp_replace(test.a,Utf8("[a-z]"),Utf8("x"),Utf8("g")) |
+ +----------------------------------------------------------+
+ | xxxDEF |
+ | xxx123 |
+ | CBAxxx |
+ | 123AxxDxx |
+ +----------------------------------------------------------+
+ "#);
let expr = regexp_replace(col("a"), lit("[a-z]"), lit("x"), None);
- let expected = [
- "+------------------------------------------------+",
- "| regexp_replace(test.a,Utf8(\"[a-z]\"),Utf8(\"x\")) |",
- "+------------------------------------------------+",
- "| xbcDEF |",
- "| xbc123 |",
- "| CBAxef |",
- "| 123AxcDef |",
- "+------------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +------------------------------------------------+
+ | regexp_replace(test.a,Utf8("[a-z]"),Utf8("x")) |
+ +------------------------------------------------+
+ | xbcDEF |
+ | xbc123 |
+ | CBAxef |
+ | 123AxcDef |
+ +------------------------------------------------+
+ "#);
Ok(())
}
@@ -723,18 +782,20 @@ async fn test_fn_regexp_replace() -> Result<()> {
async fn test_fn_replace() -> Result<()> {
let expr = replace(col("a"), lit("abc"), lit("x"));
- let expected = [
- "+---------------------------------------+",
- "| replace(test.a,Utf8(\"abc\"),Utf8(\"x\")) |",
- "+---------------------------------------+",
- "| xDEF |",
- "| x123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+---------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +---------------------------------------+
+ | replace(test.a,Utf8("abc"),Utf8("x")) |
+ +---------------------------------------+
+ | xDEF |
+ | x123 |
+ | CBAdef |
+ | 123AbcDef |
+ +---------------------------------------+
+ "#);
Ok(())
}
@@ -743,18 +804,20 @@ async fn test_fn_replace() -> Result<()> {
async fn test_fn_repeat() -> Result<()> {
let expr = repeat(col("a"), lit(2));
- let expected = [
- "+-------------------------+",
- "| repeat(test.a,Int32(2)) |",
- "+-------------------------+",
- "| abcDEFabcDEF |",
- "| abc123abc123 |",
- "| CBAdefCBAdef |",
- "| 123AbcDef123AbcDef |",
- "+-------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-------------------------+
+ | repeat(test.a,Int32(2)) |
+ +-------------------------+
+ | abcDEFabcDEF |
+ | abc123abc123 |
+ | CBAdefCBAdef |
+ | 123AbcDef123AbcDef |
+ +-------------------------+
+ ");
Ok(())
}
@@ -764,18 +827,20 @@ async fn test_fn_repeat() -> Result<()> {
async fn test_fn_reverse() -> Result<()> {
let expr = reverse(col("a"));
- let expected = [
- "+-----------------+",
- "| reverse(test.a) |",
- "+-----------------+",
- "| FEDcba |",
- "| 321cba |",
- "| fedABC |",
- "| feDcbA321 |",
- "+-----------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-----------------+
+ | reverse(test.a) |
+ +-----------------+
+ | FEDcba |
+ | 321cba |
+ | fedABC |
+ | feDcbA321 |
+ +-----------------+
+ ");
Ok(())
}
@@ -785,18 +850,20 @@ async fn test_fn_reverse() -> Result<()> {
async fn test_fn_right() -> Result<()> {
let expr = right(col("a"), lit(3));
- let expected = [
- "+------------------------+",
- "| right(test.a,Int32(3)) |",
- "+------------------------+",
- "| DEF |",
- "| 123 |",
- "| def |",
- "| Def |",
- "+------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +------------------------+
+ | right(test.a,Int32(3)) |
+ +------------------------+
+ | DEF |
+ | 123 |
+ | def |
+ | Def |
+ +------------------------+
+ ");
Ok(())
}
@@ -806,18 +873,20 @@ async fn test_fn_right() -> Result<()> {
async fn test_fn_rpad() -> Result<()> {
let expr = rpad(vec![col("a"), lit(11)]);
- let expected = [
- "+------------------------+",
- "| rpad(test.a,Int32(11)) |",
- "+------------------------+",
- "| abcDEF |",
- "| abc123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +------------------------+
+ | rpad(test.a,Int32(11)) |
+ +------------------------+
+ | abcDEF |
+ | abc123 |
+ | CBAdef |
+ | 123AbcDef |
+ +------------------------+
+ ");
Ok(())
}
@@ -827,18 +896,20 @@ async fn test_fn_rpad() -> Result<()> {
async fn test_fn_rpad_with_characters() -> Result<()> {
let expr = rpad(vec![col("a"), lit(11), lit("x")]);
- let expected = [
- "+----------------------------------+",
- "| rpad(test.a,Int32(11),Utf8(\"x\")) |",
- "+----------------------------------+",
- "| abcDEFxxxxx |",
- "| abc123xxxxx |",
- "| CBAdefxxxxx |",
- "| 123AbcDefxx |",
- "+----------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +----------------------------------+
+ | rpad(test.a,Int32(11),Utf8("x")) |
+ +----------------------------------+
+ | abcDEFxxxxx |
+ | abc123xxxxx |
+ | CBAdefxxxxx |
+ | 123AbcDefxx |
+ +----------------------------------+
+ "#);
Ok(())
}
@@ -848,18 +919,20 @@ async fn test_fn_rpad_with_characters() -> Result<()> {
async fn test_fn_sha224() -> Result<()> {
let expr = sha224(col("a"));
- let expected = [
- "+----------------------------------------------------------+",
- "| sha224(test.a) |",
- "+----------------------------------------------------------+",
- "| 8b9ef961d2b19cfe7ee2a8452e3adeea98c7b22954b4073976bf80ee |",
- "| 5c69bb695cc29b93d655e1a4bb5656cda624080d686f74477ea09349 |",
- "| b3b3783b7470594e7ddb845eca0aec5270746dd6d0bc309bb948ceab |",
- "| fc8a30d59386d78053328440c6670c3b583404a905cbe9bbd491a517 |",
- "+----------------------------------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +----------------------------------------------------------+
+ | sha224(test.a) |
+ +----------------------------------------------------------+
+ | 8b9ef961d2b19cfe7ee2a8452e3adeea98c7b22954b4073976bf80ee |
+ | 5c69bb695cc29b93d655e1a4bb5656cda624080d686f74477ea09349 |
+ | b3b3783b7470594e7ddb845eca0aec5270746dd6d0bc309bb948ceab |
+ | fc8a30d59386d78053328440c6670c3b583404a905cbe9bbd491a517 |
+ +----------------------------------------------------------+
+ ");
Ok(())
}
@@ -868,17 +941,20 @@ async fn test_fn_sha224() -> Result<()> {
async fn test_fn_split_part() -> Result<()> {
let expr = split_part(col("a"), lit("b"), lit(1));
- let expected = [
- "+---------------------------------------+",
- "| split_part(test.a,Utf8(\"b\"),Int32(1)) |",
- "+---------------------------------------+",
- "| a |",
- "| a |",
- "| CBAdef |",
- "| 123A |",
- "+---------------------------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +---------------------------------------+
+ | split_part(test.a,Utf8("b"),Int32(1)) |
+ +---------------------------------------+
+ | a |
+ | a |
+ | CBAdef |
+ | 123A |
+ +---------------------------------------+
+ "#);
Ok(())
}
@@ -887,18 +963,20 @@ async fn test_fn_split_part() -> Result<()> {
async fn test_fn_starts_with() -> Result<()> {
let expr = starts_with(col("a"), lit("abc"));
- let expected = [
- "+---------------------------------+",
- "| starts_with(test.a,Utf8(\"abc\")) |",
- "+---------------------------------+",
- "| true |",
- "| true |",
- "| false |",
- "| false |",
- "+---------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +---------------------------------+
+ | starts_with(test.a,Utf8("abc")) |
+ +---------------------------------+
+ | true |
+ | true |
+ | false |
+ | false |
+ +---------------------------------+
+ "#);
Ok(())
}
@@ -907,18 +985,20 @@ async fn test_fn_starts_with() -> Result<()> {
async fn test_fn_ends_with() -> Result<()> {
let expr = ends_with(col("a"), lit("DEF"));
- let expected = [
- "+-------------------------------+",
- "| ends_with(test.a,Utf8(\"DEF\")) |",
- "+-------------------------------+",
- "| true |",
- "| false |",
- "| false |",
- "| false |",
- "+-------------------------------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-------------------------------+
+ | ends_with(test.a,Utf8("DEF")) |
+ +-------------------------------+
+ | true |
+ | false |
+ | false |
+ | false |
+ +-------------------------------+
+ "#);
Ok(())
}
@@ -928,17 +1008,20 @@ async fn test_fn_ends_with() -> Result<()> {
async fn test_fn_strpos() -> Result<()> {
let expr = strpos(col("a"), lit("f"));
- let expected = [
- "+--------------------------+",
- "| strpos(test.a,Utf8(\"f\")) |",
- "+--------------------------+",
- "| 0 |",
- "| 0 |",
- "| 6 |",
- "| 9 |",
- "+--------------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +--------------------------+
+ | strpos(test.a,Utf8("f")) |
+ +--------------------------+
+ | 0 |
+ | 0 |
+ | 6 |
+ | 9 |
+ +--------------------------+
+ "#);
Ok(())
}
@@ -948,17 +1031,20 @@ async fn test_fn_strpos() -> Result<()> {
async fn test_fn_substr() -> Result<()> {
let expr = substr(col("a"), lit(2));
- let expected = [
- "+-------------------------+",
- "| substr(test.a,Int32(2)) |",
- "+-------------------------+",
- "| bcDEF |",
- "| bc123 |",
- "| BAdef |",
- "| 23AbcDef |",
- "+-------------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +-------------------------+
+ | substr(test.a,Int32(2)) |
+ +-------------------------+
+ | bcDEF |
+ | bc123 |
+ | BAdef |
+ | 23AbcDef |
+ +-------------------------+
+ ");
Ok(())
}
@@ -966,18 +1052,20 @@ async fn test_fn_substr() -> Result<()> {
#[tokio::test]
async fn test_cast() -> Result<()> {
let expr = cast(col("b"), DataType::Float64);
- let expected = [
- "+--------+",
- "| test.b |",
- "+--------+",
- "| 1.0 |",
- "| 10.0 |",
- "| 10.0 |",
- "| 100.0 |",
- "+--------+",
- ];
-
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +--------+
+ | test.b |
+ +--------+
+ | 1.0 |
+ | 10.0 |
+ | 10.0 |
+ | 100.0 |
+ +--------+
+ ");
Ok(())
}
@@ -986,17 +1074,20 @@ async fn test_cast() -> Result<()> {
async fn test_fn_to_hex() -> Result<()> {
let expr = to_hex(col("b"));
- let expected = [
- "+----------------+",
- "| to_hex(test.b) |",
- "+----------------+",
- "| 1 |",
- "| a |",
- "| a |",
- "| 64 |",
- "+----------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +----------------+
+ | to_hex(test.b) |
+ +----------------+
+ | 1 |
+ | a |
+ | a |
+ | 64 |
+ +----------------+
+ ");
Ok(())
}
@@ -1006,17 +1097,20 @@ async fn test_fn_to_hex() -> Result<()> {
async fn test_fn_translate() -> Result<()> {
let expr = translate(col("a"), lit("bc"), lit("xx"));
- let expected = [
- "+-----------------------------------------+",
- "| translate(test.a,Utf8(\"bc\"),Utf8(\"xx\")) |",
- "+-----------------------------------------+",
- "| axxDEF |",
- "| axx123 |",
- "| CBAdef |",
- "| 123AxxDef |",
- "+-----------------------------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-----------------------------------------+
+ | translate(test.a,Utf8("bc"),Utf8("xx")) |
+ +-----------------------------------------+
+ | axxDEF |
+ | axx123 |
+ | CBAdef |
+ | 123AxxDef |
+ +-----------------------------------------+
+ "#);
Ok(())
}
@@ -1025,17 +1119,20 @@ async fn test_fn_translate() -> Result<()> {
async fn test_fn_upper() -> Result<()> {
let expr = upper(col("a"));
- let expected = [
- "+---------------+",
- "| upper(test.a) |",
- "+---------------+",
- "| ABCDEF |",
- "| ABC123 |",
- "| CBADEF |",
- "| 123ABCDEF |",
- "+---------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r"
+ +---------------+
+ | upper(test.a) |
+ +---------------+
+ | ABCDEF |
+ | ABC123 |
+ | CBADEF |
+ | 123ABCDEF |
+ +---------------+
+ ");
Ok(())
}
@@ -1044,17 +1141,20 @@ async fn test_fn_upper() -> Result<()> {
async fn test_fn_encode() -> Result<()> {
let expr = encode(col("a"), lit("hex"));
- let expected = [
- "+----------------------------+",
- "| encode(test.a,Utf8(\"hex\")) |",
- "+----------------------------+",
- "| 616263444546 |",
- "| 616263313233 |",
- "| 434241646566 |",
- "| 313233416263446566 |",
- "+----------------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +----------------------------+
+ | encode(test.a,Utf8("hex")) |
+ +----------------------------+
+ | 616263444546 |
+ | 616263313233 |
+ | 434241646566 |
+ | 313233416263446566 |
+ +----------------------------+
+ "#);
Ok(())
}
@@ -1070,17 +1170,20 @@ async fn test_fn_decode() -> Result<()> {
// so it looks like nothing is done
.cast_to(&DataType::Utf8, &df_schema)?;
- let expected = [
- "+------------------------------------------------+",
- "| decode(encode(test.a,Utf8(\"hex\")),Utf8(\"hex\")) |",
- "+------------------------------------------------+",
- "| abcDEF |",
- "| abc123 |",
- "| CBAdef |",
- "| 123AbcDef |",
- "+------------------------------------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +------------------------------------------------+
+ | decode(encode(test.a,Utf8("hex")),Utf8("hex")) |
+ +------------------------------------------------+
+ | abcDEF |
+ | abc123 |
+ | CBAdef |
+ | 123AbcDef |
+ +------------------------------------------------+
+ "#);
Ok(())
}
@@ -1089,17 +1192,20 @@ async fn test_fn_decode() -> Result<()> {
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);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+ +-------------------------------------+
+ | array_to_string(test.l,Utf8("***")) |
+ +-------------------------------------+
+ | 0***1***2 |
+ | |
+ | 3***5 |
+ | 6***7 |
+ +-------------------------------------+
+ "#);
Ok(())
}
@@ -1110,17 +1216,20 @@ async fn test_fn_map() -> Result<()> {
vec![lit("a"), lit("b"), lit("c")],
vec![lit(1), lit(2), lit(3)],
);
- let expected = [
-
"+---------------------------------------------------------------------------------------+",
- "|
map(make_array(Utf8(\"a\"),Utf8(\"b\"),Utf8(\"c\")),make_array(Int32(1),Int32(2),Int32(3)))
|",
-
"+---------------------------------------------------------------------------------------+",
- "| {a: 1, b: 2, c: 3}
|",
- "| {a: 1, b: 2, c: 3}
|",
- "| {a: 1, b: 2, c: 3}
|",
- "| {a: 1, b: 2, c: 3}
|",
-
"+---------------------------------------------------------------------------------------+",
- ];
- assert_fn_batches!(expr, expected);
+ let batches = get_batches(expr).await?;
+
+ assert_snapshot!(
+ batches_to_string(&batches),
+ @r#"
+
+---------------------------------------------------------------------------------------+
+ |
map(make_array(Utf8("a"),Utf8("b"),Utf8("c")),make_array(Int32(1),Int32(2),Int32(3)))
|
+
+---------------------------------------------------------------------------------------+
+ | {a: 1, b: 2, c: 3}
|
+ | {a: 1, b: 2, c: 3}
|
+ | {a: 1, b: 2, c: 3}
|
+ | {a: 1, b: 2, c: 3}
|
+
+---------------------------------------------------------------------------------------+
+ "#);
Ok(())
}
@@ -1145,13 +1254,14 @@ async fn test_count_wildcard() -> Result<()> {
.build()
.unwrap();
- let expected = "Sort: count(*) ASC NULLS LAST [count(*):Int64]\
- \n Projection: count(*) [count(*):Int64]\
- \n Aggregate: groupBy=[[test.b]], aggr=[[count(Int64(1)) AS count(*)]]
[b:UInt32, count(*):Int64]\
- \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]";
-
let formatted_plan = plan.display_indent_schema().to_string();
- assert_eq!(formatted_plan, expected);
+ assert_snapshot!(formatted_plan,
+ @r"
+ Sort: count(*) ASC NULLS LAST [count(*):Int64]
+ Projection: count(*) [count(*):Int64]
+ Aggregate: groupBy=[[test.b]], aggr=[[count(Int64(1)) AS count(*)]]
[b:UInt32, count(*):Int64]
+ TableScan: test [a:UInt32, b:UInt32, c:UInt32]
+ ");
Ok(())
}
diff --git a/datafusion/core/tests/dataframe/describe.rs
b/datafusion/core/tests/dataframe/describe.rs
index 9321481efb..9bd69dfa72 100644
--- a/datafusion/core/tests/dataframe/describe.rs
+++ b/datafusion/core/tests/dataframe/describe.rs
@@ -15,11 +15,10 @@
// specific language governing permissions and limitations
// under the License.
-use datafusion::{
- assert_batches_eq,
- prelude::{ParquetReadOptions, SessionContext},
-};
+use datafusion::prelude::{ParquetReadOptions, SessionContext};
+use datafusion_common::test_util::batches_to_string;
use datafusion_common::{test_util::parquet_test_data, Result};
+use insta::assert_snapshot;
#[tokio::test]
async fn describe() -> Result<()> {
@@ -33,21 +32,21 @@ async fn describe() -> Result<()> {
.collect()
.await?;
- #[rustfmt::skip]
- let expected = [
-
"+------------+-------------------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-----------------+------------+-------------------------+--------------------+-------------------+",
- "| describe | id | bool_col | tinyint_col |
smallint_col | int_col | bigint_col | float_col
| double_col | date_string_col | string_col | timestamp_col
| year | month |",
-
"+------------+-------------------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-----------------+------------+-------------------------+--------------------+-------------------+",
- "| count | 7300.0 | 7300 | 7300.0 |
7300.0 | 7300.0 | 7300.0 | 7300.0
| 7300.0 | 7300 | 7300 | 7300
| 7300.0 | 7300.0 |",
- "| null_count | 0.0 | 0 | 0.0 |
0.0 | 0.0 | 0.0 | 0.0
| 0.0 | 0 | 0 | 0
| 0.0 | 0.0 |",
- "| mean | 3649.5 | null | 4.5 |
4.5 | 4.5 | 45.0 |
4.949999964237213 | 45.45 | null | null | null
| 2009.5 | 6.526027397260274 |",
- "| std | 2107.472815166704 | null | 2.8724780750809518 |
2.8724780750809518 | 2.8724780750809518 | 28.724780750809533 |
3.1597258182544645 | 29.012028558317645 | null | null | null
| 0.5000342500942125 | 3.44808750051728 |",
- "| min | 0.0 | null | 0.0 |
0.0 | 0.0 | 0.0 | 0.0
| 0.0 | 01/01/09 | 0 | 2008-12-31T23:00:00
| 2009.0 | 1.0 |",
- "| max | 7299.0 | null | 9.0 |
9.0 | 9.0 | 90.0 |
9.899999618530273 | 90.89999999999999 | 12/31/10 | 9 |
2010-12-31T04:09:13.860 | 2010.0 | 12.0 |",
- "| median | 3649.0 | null | 4.0 |
4.0 | 4.0 | 45.0 |
4.949999809265137 | 45.45 | null | null | null
| 2009.0 | 7.0 |",
-
"+------------+-------------------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-----------------+------------+-------------------------+--------------------+-------------------+",
- ];
- assert_batches_eq!(expected, &describe_record_batch);
+ assert_snapshot!(
+ batches_to_string(&describe_record_batch),
+ @r"
+
+------------+-------------------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-----------------+------------+-------------------------+--------------------+-------------------+
+ | describe | id | bool_col | tinyint_col |
smallint_col | int_col | bigint_col | float_col
| double_col | date_string_col | string_col | timestamp_col
| year | month |
+
+------------+-------------------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-----------------+------------+-------------------------+--------------------+-------------------+
+ | count | 7300.0 | 7300 | 7300.0 | 7300.0
| 7300.0 | 7300.0 | 7300.0 |
7300.0 | 7300 | 7300 | 7300 |
7300.0 | 7300.0 |
+ | null_count | 0.0 | 0 | 0.0 | 0.0
| 0.0 | 0.0 | 0.0 | 0.0
| 0 | 0 | 0 | 0.0
| 0.0 |
+ | mean | 3649.5 | null | 4.5 | 4.5
| 4.5 | 45.0 | 4.949999964237213 |
45.45 | null | null | null |
2009.5 | 6.526027397260274 |
+ | std | 2107.472815166704 | null | 2.8724780750809518 |
2.8724780750809518 | 2.8724780750809518 | 28.724780750809533 |
3.1597258182544645 | 29.012028558317645 | null | null | null
| 0.5000342500942125 | 3.44808750051728 |
+ | min | 0.0 | null | 0.0 | 0.0
| 0.0 | 0.0 | 0.0 | 0.0
| 01/01/09 | 0 | 2008-12-31T23:00:00 |
2009.0 | 1.0 |
+ | max | 7299.0 | null | 9.0 | 9.0
| 9.0 | 90.0 | 9.899999618530273 |
90.89999999999999 | 12/31/10 | 9 | 2010-12-31T04:09:13.860 |
2010.0 | 12.0 |
+ | median | 3649.0 | null | 4.0 | 4.0
| 4.0 | 45.0 | 4.949999809265137 |
45.45 | null | null | null |
2009.0 | 7.0 |
+
+------------+-------------------+----------+--------------------+--------------------+--------------------+--------------------+--------------------+--------------------+-----------------+------------+-------------------------+--------------------+-------------------+
+ ");
Ok(())
}
@@ -63,21 +62,22 @@ async fn describe_boolean_binary() -> Result<()> {
.await?
.collect()
.await?;
- #[rustfmt::skip]
- let expected = [
- "+------------+------+------+",
- "| describe | a | b |",
- "+------------+------+------+",
- "| count | 1 | 1 |",
- "| null_count | 0 | 0 |",
- "| mean | null | null |",
- "| std | null | null |",
- "| min | a | null |",
- "| max | a | null |",
- "| median | null | null |",
- "+------------+------+------+"
- ];
- assert_batches_eq!(expected, &result);
+
+ assert_snapshot!(
+ batches_to_string(&result),
+ @r"
+ +------------+------+------+
+ | describe | a | b |
+ +------------+------+------+
+ | count | 1 | 1 |
+ | null_count | 0 | 0 |
+ | mean | null | null |
+ | std | null | null |
+ | min | a | null |
+ | max | a | null |
+ | median | null | null |
+ +------------+------+------+
+ ");
Ok(())
}
@@ -93,21 +93,22 @@ async fn describe_null() -> Result<()> {
.await?
.collect()
.await?;
- #[rustfmt::skip]
- let expected = [
- "+------------+------+------+",
- "| describe | a | b |",
- "+------------+------+------+",
- "| count | 1 | 0 |",
- "| null_count | 0 | 1 |",
- "| mean | null | null |",
- "| std | null | null |",
- "| min | a | null |",
- "| max | a | null |",
- "| median | null | null |",
- "+------------+------+------+"
- ];
- assert_batches_eq!(expected, &result);
+
+ assert_snapshot!(
+ batches_to_string(&result),
+ @r"
+ +------------+------+------+
+ | describe | a | b |
+ +------------+------+------+
+ | count | 1 | 0 |
+ | null_count | 0 | 1 |
+ | mean | null | null |
+ | std | null | null |
+ | min | a | null |
+ | max | a | null |
+ | median | null | null |
+ +------------+------+------+
+ ");
Ok(())
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
