You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by xu...@apache.org on 2025/03/27 01:57:27 UTC
(datafusion) branch main updated: Migrate datasource tests to insta (#15258)
This is an automated email from the ASF dual-hosted git repository.
xudong963 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 3af31caa1f Migrate datasource tests to insta (#15258)
3af31caa1f is described below
commit 3af31caa1f294061316169f548f5d0d64a7f4294
Author: Shruti Sharma <98...@users.noreply.github.com>
AuthorDate: Thu Mar 27 07:27:23 2025 +0530
Migrate datasource tests to insta (#15258)
* insta migrate
* Update datafusion/core/src/datasource/physical_plan/parquet.rs
Co-authored-by: Dmitrii Blaginin <dm...@blaginin.me>
* More suggestions from code review
* Update snapshots
* fmt
---------
Co-authored-by: Andrew Lamb <an...@nerdnetworks.org>
Co-authored-by: Dmitrii Blaginin <dm...@blaginin.me>
---
datafusion/core/src/datasource/file_format/avro.rs | 31 +-
datafusion/core/src/datasource/file_format/csv.rs | 145 ++++---
datafusion/core/src/datasource/file_format/json.rs | 73 ++--
.../core/src/datasource/file_format/parquet.rs | 22 +-
datafusion/core/src/datasource/listing/table.rs | 92 ++--
datafusion/core/src/datasource/mod.rs | 18 +-
.../core/src/datasource/physical_plan/avro.rs | 93 ++--
.../core/src/datasource/physical_plan/csv.rs | 155 +++----
.../core/src/datasource/physical_plan/json.rs | 28 +-
.../core/src/datasource/physical_plan/parquet.rs | 483 ++++++++++-----------
datafusion/core/src/datasource/view_test.rs | 154 +++----
11 files changed, 629 insertions(+), 665 deletions(-)
diff --git a/datafusion/core/src/datasource/file_format/avro.rs b/datafusion/core/src/datasource/file_format/avro.rs
index 3723a728a0..a9516aad9e 100644
--- a/datafusion/core/src/datasource/file_format/avro.rs
+++ b/datafusion/core/src/datasource/file_format/avro.rs
@@ -28,8 +28,8 @@ mod tests {
};
use arrow::array::{as_string_array, Array};
use datafusion_catalog::Session;
+ use datafusion_common::test_util::batches_to_string;
use datafusion_common::{
- assert_batches_eq,
cast::{
as_binary_array, as_boolean_array, as_float32_array, as_float64_array,
as_int32_array, as_timestamp_microsecond_array,
@@ -41,6 +41,7 @@ mod tests {
use datafusion_execution::config::SessionConfig;
use datafusion_physical_plan::{collect, ExecutionPlan};
use futures::StreamExt;
+ use insta::assert_snapshot;
#[tokio::test]
async fn read_small_batches() -> Result<()> {
@@ -115,20 +116,20 @@ mod tests {
let batches = collect(exec, task_ctx).await?;
assert_eq!(batches.len(), 1);
- let expected = ["+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+",
- "| id | bool_col | tinyint_col | smallint_col | int_col | bigint_col | float_col | double_col | date_string_col | string_col | timestamp_col |",
- "+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+",
- "| 4 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30332f30312f3039 | 30 | 2009-03-01T00:00:00 |",
- "| 5 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30332f30312f3039 | 31 | 2009-03-01T00:01:00 |",
- "| 6 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30342f30312f3039 | 30 | 2009-04-01T00:00:00 |",
- "| 7 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30342f30312f3039 | 31 | 2009-04-01T00:01:00 |",
- "| 2 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30322f30312f3039 | 30 | 2009-02-01T00:00:00 |",
- "| 3 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30322f30312f3039 | 31 | 2009-02-01T00:01:00 |",
- "| 0 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30312f30312f3039 | 30 | 2009-01-01T00:00:00 |",
- "| 1 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30312f30312f3039 | 31 | 2009-01-01T00:01:00 |",
- "+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+"];
-
- assert_batches_eq!(expected, &batches);
+ assert_snapshot!(batches_to_string(&batches),@r###"
+ +----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+
+ | id | bool_col | tinyint_col | smallint_col | int_col | bigint_col | float_col | double_col | date_string_col | string_col | timestamp_col |
+ +----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+
+ | 4 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30332f30312f3039 | 30 | 2009-03-01T00:00:00 |
+ | 5 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30332f30312f3039 | 31 | 2009-03-01T00:01:00 |
+ | 6 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30342f30312f3039 | 30 | 2009-04-01T00:00:00 |
+ | 7 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30342f30312f3039 | 31 | 2009-04-01T00:01:00 |
+ | 2 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30322f30312f3039 | 30 | 2009-02-01T00:00:00 |
+ | 3 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30322f30312f3039 | 31 | 2009-02-01T00:01:00 |
+ | 0 | true | 0 | 0 | 0 | 0 | 0.0 | 0.0 | 30312f30312f3039 | 30 | 2009-01-01T00:00:00 |
+ | 1 | false | 1 | 1 | 1 | 10 | 1.1 | 10.1 | 30312f30312f3039 | 31 | 2009-01-01T00:01:00 |
+ +----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+
+ "###);
Ok(())
}
diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs
index ca22a85d35..309458975a 100644
--- a/datafusion/core/src/datasource/file_format/csv.rs
+++ b/datafusion/core/src/datasource/file_format/csv.rs
@@ -35,8 +35,8 @@ mod tests {
use datafusion_common::cast::as_string_array;
use datafusion_common::internal_err;
use datafusion_common::stats::Precision;
- use datafusion_common::test_util::arrow_test_data;
- use datafusion_common::{assert_batches_eq, Result};
+ use datafusion_common::test_util::{arrow_test_data, batches_to_string};
+ use datafusion_common::Result;
use datafusion_datasource::decoder::{
BatchDeserializer, DecoderDeserializer, DeserializerOutput,
};
@@ -57,6 +57,7 @@ mod tests {
use chrono::DateTime;
use futures::stream::BoxStream;
use futures::StreamExt;
+ use insta::assert_snapshot;
use object_store::local::LocalFileSystem;
use object_store::path::Path;
use object_store::{
@@ -557,15 +558,17 @@ mod tests {
.select_columns(&["c2", "c3"])?
.collect()
.await?;
- #[rustfmt::skip]
- let expected = ["+----+------+",
- "| c2 | c3 |",
- "+----+------+",
- "| 5 | 36 |",
- "| 5 | -31 |",
- "| 5 | -101 |",
- "+----+------+"];
- assert_batches_eq!(expected, &record_batch);
+
+ assert_snapshot!(batches_to_string(&record_batch), @r###"
+ +----+------+
+ | c2 | c3 |
+ +----+------+
+ | 5 | 36 |
+ | 5 | -31 |
+ | 5 | -101 |
+ +----+------+
+ "###);
+
Ok(())
}
@@ -671,13 +674,15 @@ mod tests {
let query_result = ctx.sql(query).await?.collect().await?;
let actual_partitions = count_query_csv_partitions(&ctx, query).await?;
- #[rustfmt::skip]
- let expected = ["+--------------+",
- "| sum(aggr.c2) |",
- "+--------------+",
- "| 285 |",
- "+--------------+"];
- assert_batches_eq!(expected, &query_result);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&query_result),@r###"
+ +--------------+
+ | sum(aggr.c2) |
+ +--------------+
+ | 285 |
+ +--------------+
+ "###);
+ }
+
assert_eq!(n_partitions, actual_partitions);
Ok(())
@@ -708,13 +713,15 @@ mod tests {
let query_result = ctx.sql(query).await?.collect().await?;
let actual_partitions = count_query_csv_partitions(&ctx, query).await?;
- #[rustfmt::skip]
- let expected = ["+--------------+",
- "| sum(aggr.c3) |",
- "+--------------+",
- "| 781 |",
- "+--------------+"];
- assert_batches_eq!(expected, &query_result);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&query_result),@r###"
+ +--------------+
+ | sum(aggr.c3) |
+ +--------------+
+ | 781 |
+ +--------------+
+ "###);
+ }
+
assert_eq!(1, actual_partitions); // Compressed csv won't be scanned in parallel
Ok(())
@@ -743,13 +750,15 @@ mod tests {
let query_result = ctx.sql(query).await?.collect().await?;
let actual_partitions = count_query_csv_partitions(&ctx, query).await?;
- #[rustfmt::skip]
- let expected = ["+--------------+",
- "| sum(aggr.c3) |",
- "+--------------+",
- "| 781 |",
- "+--------------+"];
- assert_batches_eq!(expected, &query_result);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&query_result),@r###"
+ +--------------+
+ | sum(aggr.c3) |
+ +--------------+
+ | 781 |
+ +--------------+
+ "###);
+ }
+
assert_eq!(1, actual_partitions); // csv won't be scanned in parallel when newlines_in_values is set
Ok(())
@@ -772,10 +781,10 @@ mod tests {
let query = "select * from empty where random() > 0.5;";
let query_result = ctx.sql(query).await?.collect().await?;
- #[rustfmt::skip]
- let expected = ["++",
- "++"];
- assert_batches_eq!(expected, &query_result);
+ assert_snapshot!(batches_to_string(&query_result),@r###"
+ ++
+ ++
+ "###);
Ok(())
}
@@ -797,10 +806,10 @@ mod tests {
let query = "select * from empty where random() > 0.5;";
let query_result = ctx.sql(query).await?.collect().await?;
- #[rustfmt::skip]
- let expected = ["++",
- "++"];
- assert_batches_eq!(expected, &query_result);
+ assert_snapshot!(batches_to_string(&query_result),@r###"
+ ++
+ ++
+ "###);
Ok(())
}
@@ -839,10 +848,10 @@ mod tests {
let query = "select * from empty where random() > 0.5;";
let query_result = ctx.sql(query).await?.collect().await?;
- #[rustfmt::skip]
- let expected = ["++",
- "++"];
- assert_batches_eq!(expected, &query_result);
+ assert_snapshot!(batches_to_string(&query_result),@r###"
+ ++
+ ++
+ "###);
Ok(())
}
@@ -891,13 +900,14 @@ mod tests {
let query_result = ctx.sql(query).await?.collect().await?;
let actual_partitions = count_query_csv_partitions(&ctx, query).await?;
- #[rustfmt::skip]
- let expected = ["+---------------------+",
- "| sum(empty.column_1) |",
- "+---------------------+",
- "| 10 |",
- "+---------------------+"];
- assert_batches_eq!(expected, &query_result);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&query_result),@r###"
+ +---------------------+
+ | sum(empty.column_1) |
+ +---------------------+
+ | 10 |
+ +---------------------+
+ "###);}
+
assert_eq!(n_partitions, actual_partitions); // Won't get partitioned if all files are empty
Ok(())
@@ -930,13 +940,6 @@ mod tests {
let query_result = ctx.sql(query).await?.collect().await?;
let actual_partitions = count_query_csv_partitions(&ctx, query).await?;
- #[rustfmt::skip]
- let expected = ["+-----------------------+",
- "| sum(one_col.column_1) |",
- "+-----------------------+",
- "| 50 |",
- "+-----------------------+"];
-
let file_size = std::fs::metadata("tests/data/one_col.csv")?.len() as usize;
// A 20-Byte file at most get partitioned into 20 chunks
let expected_partitions = if n_partitions <= file_size {
@@ -944,7 +947,16 @@ mod tests {
} else {
file_size
};
- assert_batches_eq!(expected, &query_result);
+
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&query_result),@r###"
+ +-----------------------+
+ | sum(one_col.column_1) |
+ +-----------------------+
+ | 50 |
+ +-----------------------+
+ "###);
+ }
+
assert_eq!(expected_partitions, actual_partitions);
Ok(())
@@ -975,13 +987,14 @@ mod tests {
let query_result = ctx.sql(query).await?.collect().await?;
let actual_partitions = count_query_csv_partitions(&ctx, query).await?;
- #[rustfmt::skip]
- let expected = ["+---------------+",
- "| sum_of_5_cols |",
- "+---------------+",
- "| 15 |",
- "+---------------+"];
- assert_batches_eq!(expected, &query_result);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&query_result),@r###"
+ +---------------+
+ | sum_of_5_cols |
+ +---------------+
+ | 15 |
+ +---------------+
+ "###);}
+
assert_eq!(n_partitions, actual_partitions);
Ok(())
diff --git a/datafusion/core/src/datasource/file_format/json.rs b/datafusion/core/src/datasource/file_format/json.rs
index ef9dd08c1c..d533dcf764 100644
--- a/datafusion/core/src/datasource/file_format/json.rs
+++ b/datafusion/core/src/datasource/file_format/json.rs
@@ -31,6 +31,7 @@ mod tests {
use arrow_schema::Schema;
use bytes::Bytes;
use datafusion_catalog::Session;
+ use datafusion_common::test_util::batches_to_string;
use datafusion_datasource::decoder::{
BatchDeserializer, DecoderDeserializer, DeserializerOutput,
};
@@ -42,11 +43,12 @@ mod tests {
use arrow::json::ReaderBuilder;
use arrow::util::pretty;
use datafusion_common::cast::as_int64_array;
+ use datafusion_common::internal_err;
use datafusion_common::stats::Precision;
- use datafusion_common::{assert_batches_eq, internal_err};
use datafusion_common::Result;
use futures::StreamExt;
+ use insta::assert_snapshot;
use object_store::local::LocalFileSystem;
use regex::Regex;
use rstest::rstest;
@@ -213,16 +215,14 @@ mod tests {
let result = ctx.sql(query).await?.collect().await?;
let actual_partitions = count_num_partitions(&ctx, query).await?;
- #[rustfmt::skip]
- let expected = [
- "+----------------------+",
- "| sum(json_parallel.a) |",
- "+----------------------+",
- "| -7 |",
- "+----------------------+"
- ];
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&result),@r###"
+ +----------------------+
+ | sum(json_parallel.a) |
+ +----------------------+
+ | -7 |
+ +----------------------+
+ "###);}
- assert_batches_eq!(expected, &result);
assert_eq!(n_partitions, actual_partitions);
Ok(())
@@ -246,13 +246,10 @@ mod tests {
let result = ctx.sql(query).await?.collect().await?;
- #[rustfmt::skip]
- let expected = [
- "++",
- "++",
- ];
-
- assert_batches_eq!(expected, &result);
+ assert_snapshot!(batches_to_string(&result),@r###"
+ ++
+ ++
+ "###);
Ok(())
}
@@ -284,17 +281,15 @@ mod tests {
}
assert_eq!(deserializer.next()?, DeserializerOutput::InputExhausted);
- let expected = [
- "+----+----+----+----+----+",
- "| c1 | c2 | c3 | c4 | c5 |",
- "+----+----+----+----+----+",
- "| 1 | 2 | 3 | 4 | 5 |",
- "| 6 | 7 | 8 | 9 | 10 |",
- "| 11 | 12 | 13 | 14 | 15 |",
- "+----+----+----+----+----+",
- ];
-
- assert_batches_eq!(expected, &[all_batches]);
+ assert_snapshot!(batches_to_string(&[all_batches]),@r###"
+ +----+----+----+----+----+
+ | c1 | c2 | c3 | c4 | c5 |
+ +----+----+----+----+----+
+ | 1 | 2 | 3 | 4 | 5 |
+ | 6 | 7 | 8 | 9 | 10 |
+ | 11 | 12 | 13 | 14 | 15 |
+ +----+----+----+----+----+
+ "###);
Ok(())
}
@@ -326,16 +321,14 @@ mod tests {
}
assert_eq!(deserializer.next()?, DeserializerOutput::RequiresMoreData);
- let expected = [
- "+----+----+----+----+----+",
- "| c1 | c2 | c3 | c4 | c5 |",
- "+----+----+----+----+----+",
- "| 1 | 2 | 3 | 4 | 5 |",
- "| 6 | 7 | 8 | 9 | 10 |",
- "+----+----+----+----+----+",
- ];
-
- assert_batches_eq!(expected, &[all_batches]);
+ insta::assert_snapshot!(fmt_batches(&[all_batches]),@r###"
+ +----+----+----+----+----+
+ | c1 | c2 | c3 | c4 | c5 |
+ +----+----+----+----+----+
+ | 1 | 2 | 3 | 4 | 5 |
+ | 6 | 7 | 8 | 9 | 10 |
+ +----+----+----+----+----+
+ "###);
Ok(())
}
@@ -349,4 +342,8 @@ mod tests {
.build_decoder()?;
Ok(DecoderDeserializer::new(JsonDecoder::new(decoder)))
}
+
+ fn fmt_batches(batches: &[RecordBatch]) -> String {
+ pretty::pretty_format_batches(batches).unwrap().to_string()
+ }
}
diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs
index f314ae54a1..27a7e7ae3c 100644
--- a/datafusion/core/src/datasource/file_format/parquet.rs
+++ b/datafusion/core/src/datasource/file_format/parquet.rs
@@ -125,8 +125,9 @@ mod tests {
};
use datafusion_common::config::{ParquetOptions, TableParquetOptions};
use datafusion_common::stats::Precision;
+ use datafusion_common::test_util::batches_to_string;
use datafusion_common::ScalarValue::Utf8;
- use datafusion_common::{assert_batches_eq, Result, ScalarValue};
+ use datafusion_common::{Result, ScalarValue};
use datafusion_datasource::file_format::FileFormat;
use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig};
use datafusion_datasource::{ListingTableUrl, PartitionedFile};
@@ -150,6 +151,7 @@ mod tests {
use datafusion_datasource::file_groups::FileGroup;
use futures::stream::BoxStream;
use futures::{Stream, StreamExt};
+ use insta::assert_snapshot;
use log::error;
use object_store::local::LocalFileSystem;
use object_store::ObjectMeta;
@@ -1091,10 +1093,11 @@ mod tests {
.expect("read_parquet should succeed");
let result = df.collect().await?;
- #[rustfmt::skip]
- let expected = ["++",
- "++"];
- assert_batches_eq!(expected, &result);
+
+ assert_snapshot!(batches_to_string(&result), @r###"
+ ++
+ ++
+ "###);
Ok(())
}
@@ -1121,10 +1124,11 @@ mod tests {
.expect("read_parquet should succeed");
let result = df.collect().await?;
- #[rustfmt::skip]
- let expected = ["++",
- "++"];
- assert_batches_eq!(expected, &result);
+
+ assert_snapshot!(batches_to_string(&result), @r###"
+ ++
+ ++
+ "###);
Ok(())
}
diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs
index 0d0dea910d..56d662571f 100644
--- a/datafusion/core/src/datasource/listing/table.rs
+++ b/datafusion/core/src/datasource/listing/table.rs
@@ -1191,14 +1191,12 @@ mod tests {
use crate::datasource::{provider_as_source, DefaultTableSource, MemTable};
use crate::execution::options::ArrowReadOptions;
use crate::prelude::*;
- use crate::{
- assert_batches_eq,
- test::{columns, object_store::register_test_store},
- };
+ use crate::test::{columns, object_store::register_test_store};
use arrow::compute::SortOptions;
use arrow::record_batch::RecordBatch;
use datafusion_common::stats::Precision;
+ use datafusion_common::test_util::batches_to_string;
use datafusion_common::{assert_contains, ScalarValue};
use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator};
use datafusion_physical_expr::PhysicalSortExpr;
@@ -2209,16 +2207,14 @@ mod tests {
// Execute the physical plan and collect the results
let res = collect(plan, session_ctx.task_ctx()).await?;
// Insert returns the number of rows written, in our case this would be 6.
- let expected = [
- "+-------+",
- "| count |",
- "+-------+",
- "| 20 |",
- "+-------+",
- ];
- // Assert that the batches read from the file match the expected result.
- assert_batches_eq!(expected, &res);
+ insta::allow_duplicates! {insta::assert_snapshot!(batches_to_string(&res),@r###"
+ +-------+
+ | count |
+ +-------+
+ | 20 |
+ +-------+
+ "###);}
// Read the records in the table
let batches = session_ctx
@@ -2226,16 +2222,14 @@ mod tests {
.await?
.collect()
.await?;
- let expected = [
- "+-------+",
- "| count |",
- "+-------+",
- "| 20 |",
- "+-------+",
- ];
- // Assert that the batches read from the file match the expected result.
- assert_batches_eq!(expected, &batches);
+ insta::allow_duplicates! {insta::assert_snapshot!(batches_to_string(&batches),@r###"
+ +-------+
+ | count |
+ +-------+
+ | 20 |
+ +-------+
+ "###);}
// Assert that `target_partition_number` many files were added to the table.
let num_files = tmp_dir.path().read_dir()?.count();
@@ -2249,17 +2243,14 @@ mod tests {
// Again, execute the physical plan and collect the results
let res = collect(plan, session_ctx.task_ctx()).await?;
- // Insert returns the number of rows written, in our case this would be 6.
- let expected = [
- "+-------+",
- "| count |",
- "+-------+",
- "| 20 |",
- "+-------+",
- ];
- // Assert that the batches read from the file match the expected result.
- assert_batches_eq!(expected, &res);
+ insta::allow_duplicates! {insta::assert_snapshot!(batches_to_string(&res),@r###"
+ +-------+
+ | count |
+ +-------+
+ | 20 |
+ +-------+
+ "###);}
// Read the contents of the table
let batches = session_ctx
@@ -2268,17 +2259,13 @@ mod tests {
.collect()
.await?;
- // Define the expected result after the second append.
- let expected = [
- "+-------+",
- "| count |",
- "+-------+",
- "| 40 |",
- "+-------+",
- ];
-
- // Assert that the batches read from the file after the second append match the expected result.
- assert_batches_eq!(expected, &batches);
+ insta::allow_duplicates! {insta::assert_snapshot!(batches_to_string(&batches),@r###"
+ +-------+
+ | count |
+ +-------+
+ | 40 |
+ +-------+
+ "###);}
// Assert that another `target_partition_number` many files were added to the table.
let num_files = tmp_dir.path().read_dir()?.count();
@@ -2334,16 +2321,15 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+-----+-----+---+",
- "| a | b | c |",
- "+-----+-----+---+",
- "| foo | bar | 1 |",
- "| foo | bar | 2 |",
- "| foo | bar | 3 |",
- "+-----+-----+---+",
- ];
- assert_batches_eq!(expected, &batches);
+ insta::allow_duplicates! {insta::assert_snapshot!(batches_to_string(&batches),@r###"
+ +-----+-----+---+
+ | a | b | c |
+ +-----+-----+---+
+ | foo | bar | 1 |
+ | foo | bar | 2 |
+ | foo | bar | 3 |
+ +-----+-----+---+
+ "###);}
Ok(())
}
diff --git a/datafusion/core/src/datasource/mod.rs b/datafusion/core/src/datasource/mod.rs
index b37198ff93..3d86e32816 100644
--- a/datafusion/core/src/datasource/mod.rs
+++ b/datafusion/core/src/datasource/mod.rs
@@ -60,7 +60,7 @@ mod tests {
use arrow::array::{Int32Array, StringArray};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
- use datafusion_common::assert_batches_sorted_eq;
+ use datafusion_common::test_util::batches_to_sort_string;
use datafusion_datasource::file_scan_config::FileScanConfig;
use datafusion_datasource::schema_adapter::{
DefaultSchemaAdapterFactory, SchemaAdapter, SchemaAdapterFactory, SchemaMapper,
@@ -138,15 +138,13 @@ mod tests {
let task_ctx = session_ctx.task_ctx();
let read = collect(parquet_exec, task_ctx).await.unwrap();
- let expected = [
- "+----+--------------+",
- "| id | extra_column |",
- "+----+--------------+",
- "| 1 | foo |",
- "+----+--------------+",
- ];
-
- assert_batches_sorted_eq!(expected, &read);
+ insta::assert_snapshot!(batches_to_sort_string(&read),@r###"
+ +----+--------------+
+ | id | extra_column |
+ +----+--------------+
+ | 1 | foo |
+ +----+--------------+
+ "###);
}
#[test]
diff --git a/datafusion/core/src/datasource/physical_plan/avro.rs b/datafusion/core/src/datasource/physical_plan/avro.rs
index 9fa2b3bc14..879fa45f4e 100644
--- a/datafusion/core/src/datasource/physical_plan/avro.rs
+++ b/datafusion/core/src/datasource/physical_plan/avro.rs
@@ -30,7 +30,8 @@ mod tests {
use crate::prelude::SessionContext;
use crate::test::object_store::local_unpartitioned_file;
use arrow::datatypes::{DataType, Field, SchemaBuilder};
- use datafusion_common::{assert_batches_eq, test_util, Result, ScalarValue};
+ use datafusion_common::test_util::batches_to_string;
+ use datafusion_common::{test_util, Result, ScalarValue};
use datafusion_datasource::file_format::FileFormat;
use datafusion_datasource::file_scan_config::FileScanConfig;
use datafusion_datasource::PartitionedFile;
@@ -40,6 +41,7 @@ mod tests {
use datafusion_physical_plan::ExecutionPlan;
use futures::StreamExt;
+ use insta::assert_snapshot;
use object_store::chunked::ChunkedStore;
use object_store::local::LocalFileSystem;
use object_store::ObjectStore;
@@ -102,22 +104,20 @@ mod tests {
.expect("plan iterator empty")
.expect("plan iterator returned an error");
- let expected = [
- "+----+----------+-------------+",
- "| id | bool_col | tinyint_col |",
- "+----+----------+-------------+",
- "| 4 | true | 0 |",
- "| 5 | false | 1 |",
- "| 6 | true | 0 |",
- "| 7 | false | 1 |",
- "| 2 | true | 0 |",
- "| 3 | false | 1 |",
- "| 0 | true | 0 |",
- "| 1 | false | 1 |",
- "+----+----------+-------------+",
- ];
-
- assert_batches_eq!(expected, &[batch]);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&[batch]), @r###"
+ +----+----------+-------------+
+ | id | bool_col | tinyint_col |
+ +----+----------+-------------+
+ | 4 | true | 0 |
+ | 5 | false | 1 |
+ | 6 | true | 0 |
+ | 7 | false | 1 |
+ | 2 | true | 0 |
+ | 3 | false | 1 |
+ | 0 | true | 0 |
+ | 1 | false | 1 |
+ +----+----------+-------------+
+ "###);}
let batch = results.next().await;
assert!(batch.is_none());
@@ -176,22 +176,20 @@ mod tests {
.expect("plan iterator empty")
.expect("plan iterator returned an error");
- let expected = [
- "+----+----------+-------------+-------------+",
- "| id | bool_col | tinyint_col | missing_col |",
- "+----+----------+-------------+-------------+",
- "| 4 | true | 0 | |",
- "| 5 | false | 1 | |",
- "| 6 | true | 0 | |",
- "| 7 | false | 1 | |",
- "| 2 | true | 0 | |",
- "| 3 | false | 1 | |",
- "| 0 | true | 0 | |",
- "| 1 | false | 1 | |",
- "+----+----------+-------------+-------------+",
- ];
-
- assert_batches_eq!(expected, &[batch]);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&[batch]), @r###"
+ +----+----------+-------------+-------------+
+ | id | bool_col | tinyint_col | missing_col |
+ +----+----------+-------------+-------------+
+ | 4 | true | 0 | |
+ | 5 | false | 1 | |
+ | 6 | true | 0 | |
+ | 7 | false | 1 | |
+ | 2 | true | 0 | |
+ | 3 | false | 1 | |
+ | 0 | true | 0 | |
+ | 1 | false | 1 | |
+ +----+----------+-------------+-------------+
+ "###);}
let batch = results.next().await;
assert!(batch.is_none());
@@ -251,21 +249,20 @@ mod tests {
.expect("plan iterator empty")
.expect("plan iterator returned an error");
- let expected = [
- "+----+----------+------------+-------------+",
- "| id | bool_col | date | tinyint_col |",
- "+----+----------+------------+-------------+",
- "| 4 | true | 2021-10-26 | 0 |",
- "| 5 | false | 2021-10-26 | 1 |",
- "| 6 | true | 2021-10-26 | 0 |",
- "| 7 | false | 2021-10-26 | 1 |",
- "| 2 | true | 2021-10-26 | 0 |",
- "| 3 | false | 2021-10-26 | 1 |",
- "| 0 | true | 2021-10-26 | 0 |",
- "| 1 | false | 2021-10-26 | 1 |",
- "+----+----------+------------+-------------+",
- ];
- assert_batches_eq!(expected, &[batch]);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&[batch]), @r###"
+ +----+----------+------------+-------------+
+ | id | bool_col | date | tinyint_col |
+ +----+----------+------------+-------------+
+ | 4 | true | 2021-10-26 | 0 |
+ | 5 | false | 2021-10-26 | 1 |
+ | 6 | true | 2021-10-26 | 0 |
+ | 7 | false | 2021-10-26 | 1 |
+ | 2 | true | 2021-10-26 | 0 |
+ | 3 | false | 2021-10-26 | 1 |
+ | 0 | true | 2021-10-26 | 0 |
+ | 1 | false | 2021-10-26 | 1 |
+ +----+----------+------------+-------------+
+ "###);}
let batch = results.next().await;
assert!(batch.is_none());
diff --git a/datafusion/core/src/datasource/physical_plan/csv.rs b/datafusion/core/src/datasource/physical_plan/csv.rs
index e80d04fe4b..1f5a395c05 100644
--- a/datafusion/core/src/datasource/physical_plan/csv.rs
+++ b/datafusion/core/src/datasource/physical_plan/csv.rs
@@ -36,6 +36,7 @@ mod tests {
use crate::prelude::SessionContext;
use crate::test::partitioned_file_groups;
use datafusion_common::test_util::arrow_test_data;
+ use datafusion_common::test_util::batches_to_string;
use datafusion_common::{assert_batches_eq, Result};
use datafusion_execution::config::SessionConfig;
use datafusion_physical_plan::metrics::MetricsSet;
@@ -49,6 +50,7 @@ mod tests {
use arrow::datatypes::*;
use bytes::Bytes;
+ use insta::assert_snapshot;
use object_store::chunked::ChunkedStore;
use object_store::local::LocalFileSystem;
use rstest::*;
@@ -122,20 +124,17 @@ mod tests {
assert_eq!(3, batch.num_columns());
assert_eq!(100, batch.num_rows());
- // slice of the first 5 lines
- let expected = [
- "+----+-----+------------+",
- "| c1 | c3 | c5 |",
- "+----+-----+------------+",
- "| c | 1 | 2033001162 |",
- "| d | -40 | 706441268 |",
- "| b | 29 | 994303988 |",
- "| a | -85 | 1171968280 |",
- "| b | -82 | 1824882165 |",
- "+----+-----+------------+",
- ];
-
- assert_batches_eq!(expected, &[batch.slice(0, 5)]);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&[batch.slice(0, 5)]), @r###"
+ +----+-----+------------+
+ | c1 | c3 | c5 |
+ +----+-----+------------+
+ | c | 1 | 2033001162 |
+ | d | -40 | 706441268 |
+ | b | 29 | 994303988 |
+ | a | -85 | 1171968280 |
+ | b | -82 | 1824882165 |
+ +----+-----+------------+
+ "###);}
Ok(())
}
@@ -183,20 +182,17 @@ mod tests {
assert_eq!(3, batch.num_columns());
assert_eq!(100, batch.num_rows());
- // slice of the first 5 lines
- let expected = [
- "+------------+----+-----+",
- "| c5 | c1 | c3 |",
- "+------------+----+-----+",
- "| 2033001162 | c | 1 |",
- "| 706441268 | d | -40 |",
- "| 994303988 | b | 29 |",
- "| 1171968280 | a | -85 |",
- "| 1824882165 | b | -82 |",
- "+------------+----+-----+",
- ];
-
- assert_batches_eq!(expected, &[batch.slice(0, 5)]);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&[batch.slice(0, 5)]), @r###"
+ +------------+----+-----+
+ | c5 | c1 | c3 |
+ +------------+----+-----+
+ | 2033001162 | c | 1 |
+ | 706441268 | d | -40 |
+ | 994303988 | b | 29 |
+ | 1171968280 | a | -85 |
+ | 1824882165 | b | -82 |
+ +------------+----+-----+
+ "###);}
Ok(())
}
@@ -246,17 +242,17 @@ mod tests {
assert_eq!(13, batch.num_columns());
assert_eq!(5, batch.num_rows());
- let expected = ["+----+----+-----+--------+------------+----------------------+-----+-------+------------+----------------------+-------------+---------------------+--------------------------------+",
- "| c1 | c2 | c3 | c4 | c5 | c6 | c7 | c8 | c9 | c10 | c11 | c12 | c13 |",
- "+----+----+-----+--------+------------+----------------------+-----+-------+------------+----------------------+-------------+---------------------+--------------------------------+",
- "| c | 2 | 1 | 18109 | 2033001162 | -6513304855495910254 | 25 | 43062 | 1491205016 | 5863949479783605708 | 0.110830784 | 0.9294097332465232 | 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW |",
- "| d | 5 | -40 | 22614 | 706441268 | -7542719935673075327 | 155 | 14337 | 3373581039 | 11720144131976083864 | 0.69632107 | 0.3114712539863804 | C2GT5KVyOPZpgKVl110TyZO0NcJ434 |",
- "| b | 1 | 29 | -18218 | 994303988 | 5983957848665088916 | 204 | 9489 | 3275293996 | 14857091259186476033 | 0.53840446 | 0.17909035118828576 | AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz |",
- "| a | 1 | -85 | -15154 | 1171968280 | 1919439543497968449 | 77 | 52286 | 774637006 | 12101411955859039553 | 0.12285209 | 0.6864391962767343 | 0keZ5G8BffGwgF2RwQD59TFzMStxCB |",
- "| b | 5 | -82 | 22080 | 1824882165 | 7373730676428214987 | 208 | 34331 | 3342719438 | 3330177516592499461 | 0.82634634 | 0.40975383525297016 | Ig1QcuKsjHXkproePdERo2w0mYzIqd |",
- "+----+----+-----+--------+------------+----------------------+-----+-------+------------+----------------------+-------------+---------------------+--------------------------------+"];
-
- assert_batches_eq!(expected, &[batch]);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&[batch]), @r###"
+ +----+----+-----+--------+------------+----------------------+-----+-------+------------+----------------------+-------------+---------------------+--------------------------------+
+ | c1 | c2 | c3 | c4 | c5 | c6 | c7 | c8 | c9 | c10 | c11 | c12 | c13 |
+ +----+----+-----+--------+------------+----------------------+-----+-------+------------+----------------------+-------------+---------------------+--------------------------------+
+ | c | 2 | 1 | 18109 | 2033001162 | -6513304855495910254 | 25 | 43062 | 1491205016 | 5863949479783605708 | 0.110830784 | 0.9294097332465232 | 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW |
+ | d | 5 | -40 | 22614 | 706441268 | -7542719935673075327 | 155 | 14337 | 3373581039 | 11720144131976083864 | 0.69632107 | 0.3114712539863804 | C2GT5KVyOPZpgKVl110TyZO0NcJ434 |
+ | b | 1 | 29 | -18218 | 994303988 | 5983957848665088916 | 204 | 9489 | 3275293996 | 14857091259186476033 | 0.53840446 | 0.17909035118828576 | AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz |
+ | a | 1 | -85 | -15154 | 1171968280 | 1919439543497968449 | 77 | 52286 | 774637006 | 12101411955859039553 | 0.12285209 | 0.6864391962767343 | 0keZ5G8BffGwgF2RwQD59TFzMStxCB |
+ | b | 5 | -82 | 22080 | 1824882165 | 7373730676428214987 | 208 | 34331 | 3342719438 | 3330177516592499461 | 0.82634634 | 0.40975383525297016 | Ig1QcuKsjHXkproePdERo2w0mYzIqd |
+ +----+----+-----+--------+------------+----------------------+-----+-------+------------+----------------------+-------------+---------------------+--------------------------------+
+ "###);}
Ok(())
}
@@ -365,19 +361,17 @@ mod tests {
assert_eq!(2, batch.num_columns());
assert_eq!(100, batch.num_rows());
- // slice of the first 5 lines
- let expected = [
- "+----+------------+",
- "| c1 | date |",
- "+----+------------+",
- "| c | 2021-10-26 |",
- "| d | 2021-10-26 |",
- "| b | 2021-10-26 |",
- "| a | 2021-10-26 |",
- "| b | 2021-10-26 |",
- "+----+------------+",
- ];
- assert_batches_eq!(expected, &[batch.slice(0, 5)]);
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&[batch.slice(0, 5)]), @r###"
+ +----+------------+
+ | c1 | date |
+ +----+------------+
+ | c | 2021-10-26 |
+ | d | 2021-10-26 |
+ | b | 2021-10-26 |
+ | a | 2021-10-26 |
+ | b | 2021-10-26 |
+ +----+------------+
+ "###);}
let metrics = csv.metrics().expect("doesn't found metrics");
let time_elapsed_processing = get_value(&metrics, "time_elapsed_processing");
@@ -501,16 +495,14 @@ mod tests {
let result = df.collect().await.unwrap();
- let expected = [
- "+---+---+",
- "| a | b |",
- "+---+---+",
- "| 1 | 2 |",
- "| 3 | 4 |",
- "+---+---+",
- ];
-
- assert_batches_eq!(expected, &result);
+ assert_snapshot!(batches_to_string(&result), @r###"
+ +---+---+
+ | a | b |
+ +---+---+
+ | 1 | 2 |
+ | 3 | 4 |
+ +---+---+
+ "###);
}
#[tokio::test]
@@ -532,16 +524,14 @@ mod tests {
let result = df.collect().await.unwrap();
- let expected = [
- "+---+---+",
- "| a | b |",
- "+---+---+",
- "| 1 | 2 |",
- "| 3 | 4 |",
- "+---+---+",
- ];
-
- assert_batches_eq!(expected, &result);
+ assert_snapshot!(batches_to_string(&result),@r###"
+ +---+---+
+ | a | b |
+ +---+---+
+ | 1 | 2 |
+ | 3 | 4 |
+ +---+---+
+ "###);
let e = session_ctx
.read_csv("memory:///", CsvReadOptions::new().terminator(Some(b'\n')))
@@ -571,17 +561,16 @@ mod tests {
.await?;
let df = ctx.sql(r#"select * from t1"#).await?.collect().await?;
- let expected = [
- "+------+--------+",
- "| col1 | col2 |",
- "+------+--------+",
- "| id0 | value0 |",
- "| id1 | value1 |",
- "| id2 | value2 |",
- "| id3 | value3 |",
- "+------+--------+",
- ];
- assert_batches_eq!(expected, &df);
+ assert_snapshot!(batches_to_string(&df),@r###"
+ +------+--------+
+ | col1 | col2 |
+ +------+--------+
+ | id0 | value0 |
+ | id1 | value1 |
+ | id2 | value2 |
+ | id3 | value3 |
+ +------+--------+
+ "###);
Ok(())
}
diff --git a/datafusion/core/src/datasource/physical_plan/json.rs b/datafusion/core/src/datasource/physical_plan/json.rs
index 041aacf687..398a842a4f 100644
--- a/datafusion/core/src/datasource/physical_plan/json.rs
+++ b/datafusion/core/src/datasource/physical_plan/json.rs
@@ -36,8 +36,8 @@ mod tests {
use crate::prelude::{CsvReadOptions, NdJsonReadOptions, SessionContext};
use crate::test::partitioned_file_groups;
use datafusion_common::cast::{as_int32_array, as_int64_array, as_string_array};
+ use datafusion_common::test_util::batches_to_string;
use datafusion_common::Result;
- use datafusion_common::{assert_batches_eq, assert_batches_sorted_eq};
use datafusion_datasource::file_compression_type::FileCompressionType;
use datafusion_datasource::file_format::FileFormat;
use datafusion_datasource_json::JsonFormat;
@@ -49,6 +49,7 @@ mod tests {
use arrow::datatypes::SchemaRef;
use arrow::datatypes::{Field, SchemaBuilder};
use datafusion_datasource::file_groups::FileGroup;
+ use insta::assert_snapshot;
use object_store::chunked::ChunkedStore;
use object_store::local::LocalFileSystem;
use object_store::ObjectStore;
@@ -136,19 +137,17 @@ mod tests {
let frame = ctx.read_json(path, read_options).await.unwrap();
let results = frame.collect().await.unwrap();
- assert_batches_eq!(
- &[
- "+-----+------------------+---------------+------+",
- "| a | b | c | d |",
- "+-----+------------------+---------------+------+",
- "| 1 | [2.0, 1.3, -6.1] | [false, true] | 4 |",
- "| -10 | [2.0, 1.3, -6.1] | [true, true] | 4 |",
- "| 2 | [2.0, , -6.1] | [false, ] | text |",
- "| | | | |",
- "+-----+------------------+---------------+------+",
- ],
- &results
- );
+ insta::allow_duplicates! {assert_snapshot!(batches_to_string(&results), @r###"
+ +-----+------------------+---------------+------+
+ | a | b | c | d |
+ +-----+------------------+---------------+------+
+ | 1 | [2.0, 1.3, -6.1] | [false, true] | 4 |
+ | -10 | [2.0, 1.3, -6.1] | [true, true] | 4 |
+ | 2 | [2.0, , -6.1] | [false, ] | text |
+ | | | | |
+ +-----+------------------+---------------+------+
+ "###);}
+
Ok(())
}
@@ -539,6 +538,7 @@ mod tests {
async fn test_json_with_repartitioning(
file_compression_type: FileCompressionType,
) -> Result<()> {
+ use datafusion_common::assert_batches_sorted_eq;
use datafusion_execution::config::SessionConfig;
let config = SessionConfig::new()
diff --git a/datafusion/core/src/datasource/physical_plan/parquet.rs b/datafusion/core/src/datasource/physical_plan/parquet.rs
index 9d21e56847..affb90b726 100644
--- a/datafusion/core/src/datasource/physical_plan/parquet.rs
+++ b/datafusion/core/src/datasource/physical_plan/parquet.rs
@@ -46,9 +46,8 @@ mod tests {
use arrow_schema::SchemaRef;
use bytes::{BufMut, BytesMut};
use datafusion_common::config::TableParquetOptions;
- use datafusion_common::{
- assert_batches_eq, assert_batches_sorted_eq, assert_contains, Result, ScalarValue,
- };
+ use datafusion_common::test_util::{batches_to_sort_string, batches_to_string};
+ use datafusion_common::{assert_contains, Result, ScalarValue};
use datafusion_datasource::file_format::FileFormat;
use datafusion_datasource::file_meta::FileMeta;
use datafusion_datasource::file_scan_config::FileScanConfig;
@@ -69,6 +68,8 @@ mod tests {
use chrono::{TimeZone, Utc};
use datafusion_datasource::file_groups::FileGroup;
use futures::StreamExt;
+ use insta;
+ use insta::assert_snapshot;
use object_store::local::LocalFileSystem;
use object_store::path::Path;
use object_store::{ObjectMeta, ObjectStore};
@@ -269,15 +270,15 @@ mod tests {
.round_trip(vec![batch.clone()])
.await;
let batches = rt.batches.unwrap();
- #[rustfmt::skip]
- let expected = [
- "+----+----+",
- "| c1 | c2 |",
- "+----+----+",
- "| 1 | |",
- "+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &batches);
+
+ insta::assert_snapshot!(batches_to_sort_string(&batches),@r###"
+ +----+----+
+ | c1 | c2 |
+ +----+----+
+ | 1 | |
+ +----+----+
+ "###);
+
let metrics = rt.parquet_exec.metrics().unwrap();
let metric = get_value(&metrics, "pushdown_rows_pruned");
assert_eq!(metric, 2, "Expected all rows to be pruned");
@@ -327,15 +328,15 @@ mod tests {
.round_trip(vec![batch.clone()])
.await;
let batches = rt.batches.unwrap();
- #[rustfmt::skip]
- let expected = [
- "+----+----+",
- "| c1 | c2 |",
- "+----+----+",
- "| 1 | |",
- "+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &batches);
+
+ insta::assert_snapshot!(batches_to_sort_string(&batches),@r###"
+ +----+----+
+ | c1 | c2 |
+ +----+----+
+ | 1 | |
+ +----+----+
+ "###);
+
let metrics = rt.parquet_exec.metrics().unwrap();
let metric = get_value(&metrics, "pushdown_rows_pruned");
assert_eq!(metric, 2, "Expected all rows to be pruned");
@@ -389,15 +390,15 @@ mod tests {
.round_trip(vec![batch.clone()])
.await;
let batches = rt.batches.unwrap();
- #[rustfmt::skip]
- let expected = [
- "+----+----+----+",
- "| c1 | c2 | c3 |",
- "+----+----+----+",
- "| 1 | | 7 |",
- "+----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &batches);
+
+ insta::assert_snapshot!(batches_to_sort_string(&batches),@r###"
+ +----+----+----+
+ | c1 | c2 | c3 |
+ +----+----+----+
+ | 1 | | 7 |
+ +----+----+----+
+ "###);
+
let metrics = rt.parquet_exec.metrics().unwrap();
let metric = get_value(&metrics, "pushdown_rows_pruned");
assert_eq!(metric, 2, "Expected all rows to be pruned");
@@ -451,15 +452,15 @@ mod tests {
.round_trip(vec![batch.clone()])
.await;
let batches = rt.batches.unwrap();
- #[rustfmt::skip]
- let expected = [
- "+----+----+----+",
- "| c1 | c2 | c3 |",
- "+----+----+----+",
- "| | | 7 |",
- "+----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &batches);
+
+ insta::assert_snapshot!(batches_to_sort_string(&batches),@r###"
+ +----+----+----+
+ | c1 | c2 | c3 |
+ +----+----+----+
+ | | | 7 |
+ +----+----+----+
+ "###);
+
let metrics = rt.parquet_exec.metrics().unwrap();
let metric = get_value(&metrics, "pushdown_rows_pruned");
assert_eq!(metric, 2, "Expected all rows to be pruned");
@@ -500,15 +501,15 @@ mod tests {
.await;
let batches = rt.batches.unwrap();
- #[rustfmt::skip]
- let expected = [
- "+----+----+----+",
- "| c1 | c2 | c3 |",
- "+----+----+----+",
- "| 1 | | 10 |",
- "+----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &batches);
+
+ insta::assert_snapshot!(batches_to_sort_string(&batches),@r###"
+ +----+----+----+
+ | c1 | c2 | c3 |
+ +----+----+----+
+ | 1 | | 10 |
+ +----+----+----+
+ "###);
+
let metrics = rt.parquet_exec.metrics().unwrap();
let metric = get_value(&metrics, "pushdown_rows_pruned");
assert_eq!(metric, 4, "Expected 4 rows to be pruned");
@@ -530,17 +531,17 @@ mod tests {
.await;
let batches = rt.batches.unwrap();
- #[rustfmt::skip]
- let expected = [
- "+----+----+----+",
- "| c1 | c2 | c3 |",
- "+----+----+----+",
- "| 3 | | 30 |",
- "| 4 | | 40 |",
- "| 5 | | 50 |",
- "+----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &batches);
+
+ insta::assert_snapshot!(batches_to_sort_string(&batches),@r###"
+ +----+----+----+
+ | c1 | c2 | c3 |
+ +----+----+----+
+ | 3 | | 30 |
+ | 4 | | 40 |
+ | 5 | | 50 |
+ +----+----+----+
+ "###);
+
let metrics = rt.parquet_exec.metrics().unwrap();
let metric = get_value(&metrics, "pushdown_rows_pruned");
assert_eq!(metric, 2, "Expected 2 rows to be pruned");
@@ -567,22 +568,22 @@ mod tests {
.round_trip_to_batches(vec![batch1, batch2, batch3])
.await
.unwrap();
- let expected = vec![
- "+-----+----+----+",
- "| c1 | c2 | c3 |",
- "+-----+----+----+",
- "| | | |",
- "| | | 20 |",
- "| | 2 | |",
- "| Foo | | |",
- "| Foo | | 10 |",
- "| Foo | 1 | |",
- "| bar | | |",
- "| bar | | |",
- "| bar | | |",
- "+-----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+
+ insta::assert_snapshot!(batches_to_sort_string(&read), @r###"
+ +-----+----+----+
+ | c1 | c2 | c3 |
+ +-----+----+----+
+ | | | |
+ | | | 20 |
+ | | 2 | |
+ | Foo | | |
+ | Foo | | 10 |
+ | Foo | 1 | |
+ | bar | | |
+ | bar | | |
+ | bar | | |
+ +-----+----+----+
+ "###);
}
#[tokio::test]
@@ -609,19 +610,19 @@ mod tests {
.round_trip_to_batches(vec![batch1, batch2])
.await
.unwrap();
- let expected = [
- "+-----+----+----+",
- "| c1 | c2 | c3 |",
- "+-----+----+----+",
- "| Foo | 1 | 10 |",
- "| | 2 | 20 |",
- "| bar | | |",
- "| Foo | 1 | 10 |",
- "| | 2 | 20 |",
- "| bar | | |",
- "+-----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+
+ insta::assert_snapshot!(batches_to_sort_string(&read),@r"
+ +-----+----+----+
+ | c1 | c2 | c3 |
+ +-----+----+----+
+ | | 2 | 20 |
+ | | 2 | 20 |
+ | Foo | 1 | 10 |
+ | Foo | 1 | 10 |
+ | bar | | |
+ | bar | | |
+ +-----+----+----+
+ ");
}
#[tokio::test]
@@ -644,19 +645,19 @@ mod tests {
.round_trip_to_batches(vec![batch1, batch2])
.await
.unwrap();
- let expected = [
- "+-----+----+----+",
- "| c1 | c3 | c2 |",
- "+-----+----+----+",
- "| Foo | 10 | |",
- "| | 20 | |",
- "| bar | | |",
- "| | 10 | 1 |",
- "| | 20 | 2 |",
- "| | | |",
- "+-----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+
+ insta::assert_snapshot!(batches_to_sort_string(&read),@r"
+ +-----+----+----+
+ | c1 | c3 | c2 |
+ +-----+----+----+
+ | | | |
+ | | 10 | 1 |
+ | | 20 | |
+ | | 20 | 2 |
+ | Foo | 10 | |
+ | bar | | |
+ +-----+----+----+
+ ");
}
#[tokio::test]
@@ -682,38 +683,32 @@ mod tests {
.round_trip_to_batches(vec![batch1, batch2])
.await
.unwrap();
- let expected = [
- "+-----+----+----+",
- "| c1 | c3 | c2 |",
- "+-----+----+----+",
- "| | | |",
- "| | 10 | 1 |",
- "| | 20 | |",
- "| | 20 | 2 |",
- "| Foo | 10 | |",
- "| bar | | |",
- "+-----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+
+ insta::assert_snapshot!(batches_to_sort_string(&read),@r###"
+ +-----+----+----+
+ | c1 | c3 | c2 |
+ +-----+----+----+
+ | | | |
+ | | 10 | 1 |
+ | | 20 | |
+ | | 20 | 2 |
+ | Foo | 10 | |
+ | bar | | |
+ +-----+----+----+
+ "###);
}
#[tokio::test]
async fn evolved_schema_intersection_filter_with_filter_pushdown() {
let c1: ArrayRef =
Arc::new(StringArray::from(vec![Some("Foo"), None, Some("bar")]));
-
let c2: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), Some(2), None]));
-
let c3: ArrayRef = Arc::new(Int8Array::from(vec![Some(10), Some(20), None]));
-
// batch1: c1(string), c3(int8)
let batch1 = create_batch(vec![("c1", c1), ("c3", c3.clone())]);
-
// batch2: c3(int8), c2(int64)
let batch2 = create_batch(vec![("c3", c3), ("c2", c2)]);
-
let filter = col("c2").eq(lit(2_i64)).or(col("c2").eq(lit(1_i64)));
-
// read/write them files:
let rt = RoundTrip::new()
.with_predicate(filter)
@@ -721,15 +716,14 @@ mod tests {
.round_trip(vec![batch1, batch2])
.await;
- let expected = [
- "+----+----+----+",
- "| c1 | c3 | c2 |",
- "+----+----+----+",
- "| | 10 | 1 |",
- "| | 20 | 2 |",
- "+----+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &rt.batches.unwrap());
+ insta::assert_snapshot!(batches_to_sort_string(&rt.batches.unwrap()), @r###"
+ +----+----+----+
+ | c1 | c3 | c2 |
+ +----+----+----+
+ | | 10 | 1 |
+ | | 20 | 2 |
+ +----+----+----+
+ "###);
let metrics = rt.parquet_exec.metrics().unwrap();
// Note there are were 6 rows in total (across three batches)
assert_eq!(get_value(&metrics, "pushdown_rows_pruned"), 4);
@@ -764,19 +758,19 @@ mod tests {
.round_trip_to_batches(vec![batch1, batch2])
.await
.unwrap();
- let expected = [
- "+-----+-----+",
- "| c1 | c4 |",
- "+-----+-----+",
- "| Foo | baz |",
- "| | boo |",
- "| bar | |",
- "| Foo | |",
- "| | |",
- "| bar | |",
- "+-----+-----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+
+ insta::assert_snapshot!(batches_to_sort_string(&read), @r###"
+ +-----+-----+
+ | c1 | c4 |
+ +-----+-----+
+ | | |
+ | | boo |
+ | Foo | |
+ | Foo | baz |
+ | bar | |
+ | bar | |
+ +-----+-----+
+ "###);
}
#[tokio::test]
@@ -838,19 +832,19 @@ mod tests {
// a null array, then the pruning predicate (currently) can not be applied.
// In a real query where this predicate was pushed down from a filter stage instead of created directly in the `DataSourceExec`,
// the filter stage would be preserved as a separate execution plan stage so the actual query results would be as expected.
- let expected = [
- "+-----+----+",
- "| c1 | c2 |",
- "+-----+----+",
- "| | |",
- "| | |",
- "| | 1 |",
- "| | 2 |",
- "| Foo | |",
- "| bar | |",
- "+-----+----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+
+ insta::assert_snapshot!(batches_to_sort_string(&read),@r###"
+ +-----+----+
+ | c1 | c2 |
+ +-----+----+
+ | | |
+ | | |
+ | | 1 |
+ | | 2 |
+ | Foo | |
+ | bar | |
+ +-----+----+
+ "###);
}
#[tokio::test]
@@ -875,14 +869,13 @@ mod tests {
.round_trip(vec![batch1, batch2])
.await;
- let expected = [
- "+----+----+",
- "| c1 | c2 |",
- "+----+----+",
- "| | 1 |",
- "+----+----+",
- ];
- assert_batches_sorted_eq!(expected, &rt.batches.unwrap());
+ insta::assert_snapshot!(batches_to_sort_string(&rt.batches.unwrap()), @r###"
+ +----+----+
+ | c1 | c2 |
+ +----+----+
+ | | 1 |
+ +----+----+
+ "###);
let metrics = rt.parquet_exec.metrics().unwrap();
// Note there are were 6 rows in total (across three batches)
assert_eq!(get_value(&metrics, "pushdown_rows_pruned"), 5);
@@ -936,25 +929,24 @@ mod tests {
.round_trip(vec![batch1, batch2, batch3, batch4])
.await;
- let expected = vec![
- "+------+----+",
- "| c1 | c2 |",
- "+------+----+",
- "| | 1 |",
- "| | 2 |",
- "| Bar | |",
- "| Bar | 2 |",
- "| Bar | 2 |",
- "| Bar2 | |",
- "| Bar3 | |",
- "| Foo | |",
- "| Foo | 1 |",
- "| Foo | 1 |",
- "| Foo2 | |",
- "| Foo3 | |",
- "+------+----+",
- ];
- assert_batches_sorted_eq!(expected, &rt.batches.unwrap());
+ insta::assert_snapshot!(batches_to_sort_string(&rt.batches.unwrap()), @r###"
+ +------+----+
+ | c1 | c2 |
+ +------+----+
+ | | 1 |
+ | | 2 |
+ | Bar | |
+ | Bar | 2 |
+ | Bar | 2 |
+ | Bar2 | |
+ | Bar3 | |
+ | Foo | |
+ | Foo | 1 |
+ | Foo | 1 |
+ | Foo2 | |
+ | Foo3 | |
+ +------+----+
+ "###);
let metrics = rt.parquet_exec.metrics().unwrap();
// There are 4 rows pruned in each of batch2, batch3, and
@@ -984,15 +976,14 @@ mod tests {
.await
.unwrap();
- let expected = [
- "+-----+----+",
- "| c1 | c2 |",
- "+-----+----+",
- "| Foo | 1 |",
- "| bar | |",
- "+-----+----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+ insta::assert_snapshot!(batches_to_sort_string(&read),@r###"
+ +-----+----+
+ | c1 | c2 |
+ +-----+----+
+ | Foo | 1 |
+ | bar | |
+ +-----+----+
+ "###);
}
#[tokio::test]
@@ -1015,16 +1006,15 @@ mod tests {
.await
.unwrap();
- let expected = [
- "+-----+----+",
- "| c1 | c2 |",
- "+-----+----+",
- "| | 2 |",
- "| Foo | 1 |",
- "| bar | |",
- "+-----+----+",
- ];
- assert_batches_sorted_eq!(expected, &read);
+ insta::assert_snapshot!(batches_to_sort_string(&read),@r###"
+ +-----+----+
+ | c1 | c2 |
+ +-----+----+
+ | | 2 |
+ | Foo | 1 |
+ | bar | |
+ +-----+----+
+ "###);
}
#[tokio::test]
@@ -1259,21 +1249,21 @@ mod tests {
let mut results = parquet_exec.execute(0, task_ctx)?;
let batch = results.next().await.unwrap()?;
assert_eq!(batch.schema().as_ref(), &expected_schema);
- let expected = [
- "+----+----------+-------------+-------+-----+",
- "| id | bool_col | tinyint_col | month | day |",
- "+----+----------+-------------+-------+-----+",
- "| 4 | true | 0 | 10 | 26 |",
- "| 5 | false | 1 | 10 | 26 |",
- "| 6 | true | 0 | 10 | 26 |",
- "| 7 | false | 1 | 10 | 26 |",
- "| 2 | true | 0 | 10 | 26 |",
- "| 3 | false | 1 | 10 | 26 |",
- "| 0 | true | 0 | 10 | 26 |",
- "| 1 | false | 1 | 10 | 26 |",
- "+----+----------+-------------+-------+-----+",
- ];
- assert_batches_eq!(expected, &[batch]);
+
+ assert_snapshot!(batches_to_string(&[batch]),@r###"
+ +----+----------+-------------+-------+-----+
+ | id | bool_col | tinyint_col | month | day |
+ +----+----------+-------------+-------+-----+
+ | 4 | true | 0 | 10 | 26 |
+ | 5 | false | 1 | 10 | 26 |
+ | 6 | true | 0 | 10 | 26 |
+ | 7 | false | 1 | 10 | 26 |
+ | 2 | true | 0 | 10 | 26 |
+ | 3 | false | 1 | 10 | 26 |
+ | 0 | true | 0 | 10 | 26 |
+ | 1 | false | 1 | 10 | 26 |
+ +----+----------+-------------+-------+-----+
+ "###);
let batch = results.next().await;
assert!(batch.is_none());
@@ -1344,17 +1334,14 @@ mod tests {
let metrics = rt.parquet_exec.metrics().unwrap();
- // assert the batches and some metrics
- #[rustfmt::skip]
- let expected = [
- "+-----+",
- "| int |",
- "+-----+",
- "| 4 |",
- "| 5 |",
- "+-----+"
- ];
- assert_batches_sorted_eq!(expected, &rt.batches.unwrap());
+ assert_snapshot!(batches_to_sort_string(&rt.batches.unwrap()),@r###"
+ +-----+
+ | int |
+ +-----+
+ | 4 |
+ | 5 |
+ +-----+
+ "###);
assert_eq!(get_value(&metrics, "page_index_rows_pruned"), 4);
assert_eq!(get_value(&metrics, "page_index_rows_matched"), 2);
assert!(
@@ -1418,10 +1405,14 @@ mod tests {
let metrics = rt.parquet_exec.metrics().unwrap();
// assert the batches and some metrics
- let expected = [
- "+-----+", "| c1 |", "+-----+", "| Foo |", "| zzz |", "+-----+",
- ];
- assert_batches_sorted_eq!(expected, &rt.batches.unwrap());
+ assert_snapshot!(batches_to_string(&rt.batches.unwrap()),@r###"
+ +-----+
+ | c1 |
+ +-----+
+ | Foo |
+ | zzz |
+ +-----+
+ "###);
// pushdown predicates have eliminated all 4 bar rows and the
// null row for 5 rows total
@@ -1770,14 +1761,13 @@ mod tests {
let sql = "select * from base_table where name='test02'";
let batch = ctx.sql(sql).await.unwrap().collect().await.unwrap();
assert_eq!(batch.len(), 1);
- let expected = [
- "+---------------------+----+--------+",
- "| struct | id | name |",
- "+---------------------+----+--------+",
- "| {id: 4, name: aaa2} | 2 | test02 |",
- "+---------------------+----+--------+",
- ];
- assert_batches_eq!(expected, &batch);
+ insta::assert_snapshot!(batches_to_string(&batch),@r###"
+ +---------------------+----+--------+
+ | struct | id | name |
+ +---------------------+----+--------+
+ | {id: 4, name: aaa2} | 2 | test02 |
+ +---------------------+----+--------+
+ "###);
Ok(())
}
@@ -1800,14 +1790,13 @@ mod tests {
let sql = "select * from base_table where name='test02'";
let batch = ctx.sql(sql).await.unwrap().collect().await.unwrap();
assert_eq!(batch.len(), 1);
- let expected = [
- "+---------------------+----+--------+",
- "| struct | id | name |",
- "+---------------------+----+--------+",
- "| {id: 4, name: aaa2} | 2 | test02 |",
- "+---------------------+----+--------+",
- ];
- assert_batches_eq!(expected, &batch);
+ insta::assert_snapshot!(batches_to_string(&batch),@r###"
+ +---------------------+----+--------+
+ | struct | id | name |
+ +---------------------+----+--------+
+ | {id: 4, name: aaa2} | 2 | test02 |
+ +---------------------+----+--------+
+ "###);
Ok(())
}
diff --git a/datafusion/core/src/datasource/view_test.rs b/datafusion/core/src/datasource/view_test.rs
index c3dd5a2dd9..85ad9ff664 100644
--- a/datafusion/core/src/datasource/view_test.rs
+++ b/datafusion/core/src/datasource/view_test.rs
@@ -20,10 +20,11 @@
#[cfg(test)]
mod tests {
use crate::error::Result;
+ use crate::execution::context::SessionConfig;
use crate::execution::options::ParquetReadOptions;
use crate::prelude::SessionContext;
use crate::test_util::parquet_test_data;
- use crate::{assert_batches_eq, execution::context::SessionConfig};
+ use datafusion_common::test_util::batches_to_string;
use datafusion_expr::{col, lit};
#[tokio::test]
@@ -45,9 +46,13 @@ mod tests {
.collect()
.await?;
- let expected = ["+---+", "| b |", "+---+", "| 2 |", "+---+"];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---+
+ | b |
+ +---+
+ | 2 |
+ +---+
+ "###);
Ok(())
}
@@ -91,16 +96,14 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------+---------+---------+",
- "| column1 | column2 | column3 |",
- "+---------+---------+---------+",
- "| 1 | 2 | 3 |",
- "| 4 | 5 | 6 |",
- "+---------+---------+---------+",
- ];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------+---------+---------+
+ | column1 | column2 | column3 |
+ +---------+---------+---------+
+ | 1 | 2 | 3 |
+ | 4 | 5 | 6 |
+ +---------+---------+---------+
+ "###);
let view_sql =
"CREATE VIEW replace_xyz AS SELECT * REPLACE (column1*2 as column1) FROM xyz";
@@ -112,16 +115,15 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------+---------+---------+",
- "| column1 | column2 | column3 |",
- "+---------+---------+---------+",
- "| 2 | 2 | 3 |",
- "| 8 | 5 | 6 |",
- "+---------+---------+---------+",
- ];
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------+---------+---------+
+ | column1 | column2 | column3 |
+ +---------+---------+---------+
+ | 2 | 2 | 3 |
+ | 8 | 5 | 6 |
+ +---------+---------+---------+
+ "###);
- assert_batches_eq!(expected, &results);
Ok(())
}
@@ -144,16 +146,14 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------------+",
- "| column1_alias |",
- "+---------------+",
- "| 1 |",
- "| 4 |",
- "+---------------+",
- ];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------------+
+ | column1_alias |
+ +---------------+
+ | 1 |
+ | 4 |
+ +---------------+
+ "###);
Ok(())
}
@@ -177,16 +177,14 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------------+---------------+",
- "| column2_alias | column1_alias |",
- "+---------------+---------------+",
- "| 2 | 1 |",
- "| 5 | 4 |",
- "+---------------+---------------+",
- ];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------------+---------------+
+ | column2_alias | column1_alias |
+ +---------------+---------------+
+ | 2 | 1 |
+ | 5 | 4 |
+ +---------------+---------------+
+ "###);
Ok(())
}
@@ -215,16 +213,14 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------+",
- "| column1 |",
- "+---------+",
- "| 1 |",
- "| 4 |",
- "+---------+",
- ];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------+
+ | column1 |
+ +---------+
+ | 1 |
+ | 4 |
+ +---------+
+ "###);
Ok(())
}
@@ -253,15 +249,13 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------+",
- "| column1 |",
- "+---------+",
- "| 4 |",
- "+---------+",
- ];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------+
+ | column1 |
+ +---------+
+ | 4 |
+ +---------+
+ "###);
Ok(())
}
@@ -293,16 +287,14 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------+---------+---------+",
- "| column2 | column1 | column3 |",
- "+---------+---------+---------+",
- "| 2 | 1 | 3 |",
- "| 5 | 4 | 6 |",
- "+---------+---------+---------+",
- ];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------+---------+---------+
+ | column2 | column1 | column3 |
+ +---------+---------+---------+
+ | 2 | 1 | 3 |
+ | 5 | 4 | 6 |
+ +---------+---------+---------+
+ "###);
Ok(())
}
@@ -450,16 +442,14 @@ mod tests {
.collect()
.await?;
- let expected = [
- "+---------+",
- "| column1 |",
- "+---------+",
- "| 1 |",
- "| 4 |",
- "+---------+",
- ];
-
- assert_batches_eq!(expected, &results);
+ insta::assert_snapshot!(batches_to_string(&results),@r###"
+ +---------+
+ | column1 |
+ +---------+
+ | 1 |
+ | 4 |
+ +---------+
+ "###);
Ok(())
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org