Parse Sqllogictest column types from physical schema (#11929)

* Parse Sqllogictest column types from physical schema

* Use execute_stream

Author: jonahgao

datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs

@@ -15,10 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::Arc;
 use std::{path::PathBuf, time::Duration};
 
 use arrow::record_batch::RecordBatch;
 use async_trait::async_trait;
+use datafusion::physical_plan::common::collect;
+use datafusion::physical_plan::execute_stream;
 use datafusion::prelude::SessionContext;
 use log::info;
 use sqllogictest::DBOutput;
@@ -69,9 +72,12 @@ impl sqllogictest::AsyncDB for DataFusion {
 
 async fn run_query(ctx: &SessionContext, sql: impl Into<String>) -> Result<DFOutput> {
     let df = ctx.sql(sql.into().as_str()).await?;
+    let task_ctx = Arc::new(df.task_ctx());
+    let plan = df.create_physical_plan().await?;
 
-    let types = normalize::convert_schema_to_types(df.schema().fields());
-    let results: Vec<RecordBatch> = df.collect().await?;
+    let stream = execute_stream(plan, task_ctx)?;
+    let types = normalize::convert_schema_to_types(stream.schema().fields());
+    let results: Vec<RecordBatch> = collect(stream).await?;
     let rows = normalize::convert_batches(results)?;
 
     if rows.is_empty() && types.is_empty() {

datafusion/sqllogictest/test_files/coalesce.slt

@@ -23,7 +23,7 @@ select coalesce(1, 2, 3);
 1
 
 # test with first null
-query ?T
+query IT
 select coalesce(null, 3, 2, 1), arrow_typeof(coalesce(null, 3, 2, 1));
 ----
 3 Int64
@@ -35,7 +35,7 @@ select coalesce(null, null);
 NULL
 
 # cast to float
-query IT
+query RT
 select
   coalesce(1, 2.0),
   arrow_typeof(coalesce(1, 2.0))
@@ -51,7 +51,7 @@ select
 ----
 2 Float64
 
-query IT
+query RT
 select
   coalesce(1, arrow_cast(2.0, 'Float32')),
   arrow_typeof(coalesce(1, arrow_cast(2.0, 'Float32')))
@@ -177,7 +177,7 @@ select
 2 Decimal256(22, 2)
 
 # coalesce string
-query T?
+query TT
 select
   coalesce('', 'test'),
   coalesce(null, 'test');
@@ -246,7 +246,7 @@ drop table test1
 statement ok
 create table t(c varchar) as values ('a'), (null);
 
-query TT
+query ?T
 select 
   coalesce(c, arrow_cast('b', 'Dictionary(Int32, Utf8)')),
   arrow_typeof(coalesce(c, arrow_cast('b', 'Dictionary(Int32, Utf8)')))
@@ -295,7 +295,7 @@ statement ok
 drop table t;
 
 # test dict(int32, int8)
-query I
+query ?
 select coalesce(34, arrow_cast(123, 'Dictionary(Int32, Int8)'));
 ----
 34

datafusion/sqllogictest/test_files/copy.slt

@@ -20,13 +20,13 @@ statement ok
 create table source_table(col1 integer, col2 varchar) as values (1, 'Foo'), (2, 'Bar');
 
 # Copy to directory as multiple files
-query IT
+query I
 COPY source_table TO 'test_files/scratch/copy/table/' STORED AS parquet OPTIONS ('format.compression' 'zstd(10)');
 ----
 2
 
 # Copy to directory as partitioned files
-query IT
+query I
 COPY source_table TO 'test_files/scratch/copy/partitioned_table1/' STORED AS parquet PARTITIONED BY (col2) OPTIONS ('format.compression' 'zstd(10)');
 ----
 2
@@ -53,7 +53,7 @@ select * from validate_partitioned_parquet_bar order by col1;
 2
 
 # Copy to directory as partitioned files
-query ITT
+query I
 COPY (values (1, 'a', 'x'), (2, 'b', 'y'), (3, 'c', 'z')) TO 'test_files/scratch/copy/partitioned_table2/' STORED AS parquet PARTITIONED BY (column2, column3)
 OPTIONS ('format.compression' 'zstd(10)');
 ----
@@ -81,7 +81,7 @@ select * from validate_partitioned_parquet_a_x order by column1;
 1
 
 # Copy to directory as partitioned files
-query TTT
+query I
 COPY (values ('1', 'a', 'x'), ('2', 'b', 'y'), ('3', 'c', 'z')) TO 'test_files/scratch/copy/partitioned_table3/' STORED AS parquet PARTITIONED BY (column1, column3)
 OPTIONS ('format.compression' 'zstd(10)');
 ----
@@ -167,7 +167,7 @@ physical_plan
 02)--MemoryExec: partitions=1, partition_sizes=[1]
 
 # Copy to directory as partitioned files with keep_partition_by_columns enabled
-query TT
+query I
 COPY (values ('1', 'a'), ('2', 'b'), ('3', 'c')) TO 'test_files/scratch/copy/partitioned_table4/' STORED AS parquet PARTITIONED BY (column1)
 OPTIONS (execution.keep_partition_by_columns true);
 ----
@@ -184,7 +184,7 @@ select column1, column2 from validate_partitioned_parquet4 order by column1,colu
 1 a
 
 # Copy more files to directory via query
-query IT
+query I
 COPY (select * from source_table UNION ALL select * from source_table) to 'test_files/scratch/copy/table/' STORED AS PARQUET;
 ----
 4
@@ -203,7 +203,7 @@ select * from validate_parquet;
 1 Foo
 2 Bar
 
-query ?
+query I
 copy (values (struct(timestamp '2021-01-01 01:00:01', 1)), (struct(timestamp '2022-01-01 01:00:01', 2)), 
 (struct(timestamp '2023-01-03 01:00:01', 3)), (struct(timestamp '2024-01-01 01:00:01', 4)))
 to 'test_files/scratch/copy/table_nested2/' STORED AS PARQUET;
@@ -221,7 +221,7 @@ select * from validate_parquet_nested2;
 {c0: 2023-01-03T01:00:01, c1: 3}
 {c0: 2024-01-01T01:00:01, c1: 4}
 
-query ??
+query I
 COPY 
 (values (struct ('foo', (struct ('foo', make_array(struct('a',1), struct('b',2))))), make_array(timestamp '2023-01-01 01:00:01',timestamp '2023-01-01 01:00:01')), 
 (struct('bar', (struct ('foo', make_array(struct('aa',10), struct('bb',20))))), make_array(timestamp '2024-01-01 01:00:01', timestamp '2024-01-01 01:00:01'))) 
@@ -239,7 +239,7 @@ select * from validate_parquet_nested;
 {c0: foo, c1: {c0: foo, c1: [{c0: a, c1: 1}, {c0: b, c1: 2}]}} [2023-01-01T01:00:01, 2023-01-01T01:00:01]
 {c0: bar, c1: {c0: foo, c1: [{c0: aa, c1: 10}, {c0: bb, c1: 20}]}} [2024-01-01T01:00:01, 2024-01-01T01:00:01]
 
-query ?
+query I
 copy (values ([struct('foo', 1), struct('bar', 2)])) 
 to 'test_files/scratch/copy/array_of_struct/'
 STORED AS PARQUET;
@@ -255,7 +255,7 @@ select * from validate_array_of_struct;
 ----
 [{c0: foo, c1: 1}, {c0: bar, c1: 2}]
 
-query ?
+query I
 copy (values (struct('foo', [1,2,3], struct('bar', [2,3,4])))) 
 to 'test_files/scratch/copy/struct_with_array/' STORED AS PARQUET;
 ----
@@ -272,7 +272,7 @@ select * from validate_struct_with_array;
 
 
 # Copy parquet with all supported statement overrides
-query IT
+query I
 COPY source_table
 TO 'test_files/scratch/copy/table_with_options/'
 STORED AS PARQUET
@@ -378,7 +378,7 @@ select * from validate_parquet_with_options;
 2 Bar
 
 # Copy from table to single file
-query IT
+query I
 COPY source_table to 'test_files/scratch/copy/table.parquet';
 ----
 2
@@ -394,7 +394,7 @@ select * from validate_parquet_single;
 2 Bar
 
 # copy from table to folder of compressed json files
-query IT
+query I
 COPY source_table  to 'test_files/scratch/copy/table_json_gz' STORED AS JSON OPTIONS ('format.compression' gzip);
 ----
 2
@@ -410,7 +410,7 @@ select * from validate_json_gz;
 2 Bar
 
 # copy from table to folder of compressed csv files
-query IT
+query I
 COPY source_table  to 'test_files/scratch/copy/table_csv' STORED AS CSV OPTIONS ('format.has_header' false, 'format.compression' gzip);
 ----
 2
@@ -426,7 +426,7 @@ select * from validate_csv;
 2 Bar
 
 # Copy from table to single csv
-query IT
+query I
 COPY source_table  to 'test_files/scratch/copy/table.csv';
 ----
 2
@@ -442,7 +442,7 @@ select * from validate_single_csv;
 2 Bar
 
 # Copy from table to folder of json
-query IT
+query I
 COPY source_table to 'test_files/scratch/copy/table_json' STORED AS JSON;
 ----
 2
@@ -458,7 +458,7 @@ select * from validate_json;
 2 Bar
 
 # Copy from table to single json file
-query IT
+query I
 COPY source_table  to 'test_files/scratch/copy/table.json' STORED AS JSON ;
 ----
 2
@@ -474,7 +474,7 @@ select * from validate_single_json;
 2 Bar
 
 # COPY csv files with all options set
-query IT
+query I
 COPY source_table
 to 'test_files/scratch/copy/table_csv_with_options'
 STORED AS CSV OPTIONS (
@@ -499,7 +499,7 @@ select * from validate_csv_with_options;
 2;Bar
 
 # Copy from table to single arrow file
-query IT
+query I
 COPY source_table to 'test_files/scratch/copy/table.arrow' STORED AS ARROW;
 ----
 2
@@ -517,7 +517,7 @@ select * from validate_arrow_file;
 2 Bar
 
 # Copy from dict encoded values to single arrow file
-query T?
+query I
 COPY (values 
 ('c', arrow_cast('foo', 'Dictionary(Int32, Utf8)')), ('d', arrow_cast('bar', 'Dictionary(Int32, Utf8)'))) 
 to 'test_files/scratch/copy/table_dict.arrow' STORED AS ARROW;
@@ -538,7 +538,7 @@ d bar
 
 
 # Copy from table to folder of json
-query IT
+query I
 COPY source_table to 'test_files/scratch/copy/table_arrow' STORED AS ARROW;
 ----
 2
@@ -556,7 +556,7 @@ select * from validate_arrow;
 # Format Options Support without the 'format.' prefix
 
 # Copy with format options for Parquet without the 'format.' prefix
-query IT
+query I
 COPY source_table TO 'test_files/scratch/copy/format_table.parquet'
 OPTIONS (
     compression snappy,
@@ -566,14 +566,14 @@ OPTIONS (
 2
 
 # Copy with format options for JSON without the 'format.' prefix
-query IT
+query I
 COPY source_table  to 'test_files/scratch/copy/format_table'
 STORED AS JSON OPTIONS (compression gzip);
 ----
 2
 
 # Copy with format options for CSV without the 'format.' prefix
-query IT
+query I
 COPY source_table to 'test_files/scratch/copy/format_table.csv'
 OPTIONS (
     has_header false,

datafusion/sqllogictest/test_files/csv_files.slt

@@ -115,14 +115,14 @@ CREATE TABLE src_table_2 (
 (7, 'ggg', 700, 2),
 (8, 'hhh', 800, 2);
 
-query ITII
+query I
 COPY  src_table_1 TO 'test_files/scratch/csv_files/csv_partitions/1.csv'
 STORED AS CSV;
 ----
 4
 
 
-query ITII
+query I
 COPY  src_table_2 TO 'test_files/scratch/csv_files/csv_partitions/2.csv'
 STORED AS CSV;
 ----
@@ -175,7 +175,7 @@ CREATE TABLE table_with_necessary_quoting (
 (4, 'h|h|h');
 
 # quote is required because `|` is delimiter and part of the data
-query IT
+query I
 COPY table_with_necessary_quoting TO 'test_files/scratch/csv_files/table_with_necessary_quoting.csv'
 STORED AS csv
 OPTIONS ('format.quote' '~',
@@ -247,7 +247,7 @@ id2 "value2"
 id3 "value3"
 
 # ensure that double quote option is used when writing to csv
-query TT
+query I
 COPY csv_with_double_quote TO 'test_files/scratch/csv_files/table_with_double_quotes.csv'
 STORED AS csv
 OPTIONS ('format.double_quote' 'true');
@@ -271,7 +271,7 @@ id2 "value2"
 id3 "value3"
 
 # ensure when double quote option is disabled that quotes are escaped instead
-query TT
+query I
 COPY csv_with_double_quote TO 'test_files/scratch/csv_files/table_with_escaped_quotes.csv'
 STORED AS csv
 OPTIONS ('format.double_quote' 'false', 'format.escape' '#');