[sqllogictest] Define output types and check them in tests

datafusion/core/tests/sqllogictests/README.md

@@ -105,10 +105,13 @@ query <type_string> <sort_mode>
 
 - `test_name`: Uniquely identify the test name (arrow-datafusion only)
 - `type_string`: A short string that specifies the number of result columns and the expected datatype of each result column. There is one character in the <type_string> for each result column. The characters codes are:
-  - "T" for a text result,
-  - "I" for an integer result,
-  - "R" for a floating-point result,
-  - "?" for any other type.
+  - 'B' - **B**oolean,
+  - 'D' - **D**atetime,
+  - 'I' - **I**nteger,
+  - 'F' - **F**loating-point results,
+  - 'S' - **S**string,
+  - 'T' - **T**imestamp,
+  - "?" - any other types
 - `expected_result`: In the results section, some values are converted according to some rules:
   - floating point values are rounded to the scale of "12",
   - NULL values are rendered as `NULL`,

datafusion/core/tests/sqllogictests/src/engines/datafusion/create_table.rs

@@ -16,8 +16,9 @@
 // under the License.
 
 use super::error::Result;
+use crate::engines::datafusion::error::DFSqlLogicTestError;
 use crate::engines::datafusion::util::LogicTestContextProvider;
-use crate::{engines::datafusion::error::DFSqlLogicTestError, output::DFOutput};
+use crate::engines::output::DFOutput;
 use datafusion::datasource::MemTable;
 use datafusion::prelude::SessionContext;
 use datafusion_common::{DataFusionError, OwnedTableReference};

datafusion/core/tests/sqllogictests/src/engines/datafusion/insert.rs

@@ -16,7 +16,8 @@
 // under the License.
 
 use super::error::Result;
-use crate::{engines::datafusion::util::LogicTestContextProvider, output::DFOutput};
+use crate::engines::datafusion::util::LogicTestContextProvider;
+use crate::engines::output::DFOutput;
 use arrow::record_batch::RecordBatch;
 use datafusion::datasource::MemTable;
 use datafusion::prelude::SessionContext;

datafusion/core/tests/sqllogictests/src/engines/datafusion/mod.rs

@@ -18,7 +18,7 @@
 use std::path::PathBuf;
 use std::time::Duration;
 
-use crate::output::{DFColumnType, DFOutput};
+use crate::engines::output::{DFColumnType, DFOutput};
 
 use self::error::{DFSqlLogicTestError, Result};
 use async_trait::async_trait;

datafusion/core/tests/sqllogictests/src/engines/datafusion/normalize.rs

@@ -15,11 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::datatypes::SchemaRef;
 use arrow::{array, array::ArrayRef, datatypes::DataType, record_batch::RecordBatch};
 use datafusion::error::DataFusionError;
 use sqllogictest::DBOutput;
 
-use crate::output::{DFColumnType, DFOutput};
+use crate::engines::output::{DFColumnType, DFOutput};
 
 use super::super::conversion::*;
 use super::error::{DFSqlLogicTestError, Result};
@@ -35,10 +36,7 @@ pub fn convert_batches(batches: Vec<RecordBatch>) -> Result<DFOutput> {
     }
 
     let schema = batches[0].schema();
-
-    // TODO: report the the actual types of the result
-    // https://github.com/apache/arrow-datafusion/issues/4499
-    let types = vec![DFColumnType::Any; batches[0].num_columns()];
+    let types = convert_schema_to_types(&schema);
 
     let mut rows = vec![];
     for batch in batches {
@@ -125,3 +123,34 @@ pub fn cell_to_string(col: &ArrayRef, row: usize) -> Result<String> {
         .map_err(DFSqlLogicTestError::Arrow)
     }
 }
+
+fn convert_schema_to_types(schema: &SchemaRef) -> Vec<DFColumnType> {
+    schema
+        .fields()
+        .iter()
+        .map(|f| f.data_type())
+        .map(|data_type| match data_type {
+            DataType::Boolean => DFColumnType::Boolean,
+            DataType::Int8
+            | DataType::Int16
+            | DataType::Int32
+            | DataType::Int64
+            | DataType::UInt8
+            | DataType::UInt16
+            | DataType::UInt32
+            | DataType::UInt64 => DFColumnType::Integer,
+            DataType::Float16
+            | DataType::Float32
+            | DataType::Float64
+            | DataType::Decimal128(_, _)
+            | DataType::Decimal256(_, _) => DFColumnType::FloatingPoint,
+            DataType::Utf8 | DataType::LargeUtf8 => DFColumnType::String,
+            DataType::Date32
+            | DataType::Date64
+            | DataType::Time32(_)
+            | DataType::Time64(_) => DFColumnType::DateTime,
+            DataType::Timestamp(_, _) => DFColumnType::Timestamp,
+            _ => DFColumnType::Another,
+        })
+        .collect()
+}

datafusion/core/tests/sqllogictests/src/engines/mod.rs

@@ -17,4 +17,5 @@
 
 mod conversion;
 pub mod datafusion;
+mod output;
 pub mod postgres;

datafusion/core/tests/sqllogictests/src/engines/output.rs

@@ -0,0 +1,57 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use sqllogictest::{ColumnType, DBOutput};
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum DFColumnType {
+    Boolean,
+    DateTime,
+    Integer,
+    FloatingPoint,
+    String,
+    Timestamp,
+    Another,
+}
+
+impl ColumnType for DFColumnType {
+    fn from_char(value: char) -> Option<Self> {
+        match value {
+            'B' => Some(Self::Boolean),
+            'D' => Some(Self::DateTime),
+            'I' => Some(Self::Integer),
+            'F' => Some(Self::FloatingPoint),
+            'S' => Some(Self::String),
+            'T' => Some(Self::Timestamp),
+            _ => Some(Self::Another),
+        }
+    }
+
+    fn to_char(&self) -> char {
+        match self {
+            Self::Boolean => 'B',
+            Self::DateTime => 'D',
+            Self::Integer => 'I',
+            Self::FloatingPoint => 'F',
+            Self::String => 'S',
+            Self::Timestamp => 'T',
+            Self::Another => '?',
+        }
+    }
+}
+
+pub type DFOutput = DBOutput<DFColumnType>;

datafusion/core/tests/sqllogictests/src/engines/postgres/mod.rs

@@ -25,9 +25,8 @@ use log::debug;
 use sqllogictest::DBOutput;
 use tokio::task::JoinHandle;
 
-use crate::output::{DFColumnType, DFOutput};
-
 use super::conversion::*;
+use crate::engines::output::{DFColumnType, DFOutput};
 use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
 use postgres_types::Type;
 use rust_decimal::Decimal;
@@ -212,55 +211,15 @@ impl Drop for Postgres {
     }
 }
 
-macro_rules! make_string {
-    ($row:ident, $idx:ident, $t:ty) => {{
-        let value: Option<$t> = $row.get($idx);
-        match value {
-            Some(value) => value.to_string(),
-            None => NULL_STR.to_string(),
-        }
-    }};
-    ($row:ident, $idx:ident, $t:ty, $convert:ident) => {{
-        let value: Option<$t> = $row.get($idx);
-        match value {
-            Some(value) => $convert(value).to_string(),
-            None => NULL_STR.to_string(),
-        }
-    }};
-}
-
-fn cell_to_string(row: &Row, column: &Column, idx: usize) -> String {
-    match column.type_().clone() {
-        Type::CHAR => make_string!(row, idx, i8),
-        Type::INT2 => make_string!(row, idx, i16),
-        Type::INT4 => make_string!(row, idx, i32),
-        Type::INT8 => make_string!(row, idx, i64),
-        Type::NUMERIC => make_string!(row, idx, Decimal, decimal_to_str),
-        Type::DATE => make_string!(row, idx, NaiveDate),
-        Type::TIME => make_string!(row, idx, NaiveTime),
-        Type::TIMESTAMP => {
-            let value: Option<NaiveDateTime> = row.get(idx);
-            value
-                .map(|d| format!("{d:?}"))
-                .unwrap_or_else(|| "NULL".to_string())
-        }
-        Type::BOOL => make_string!(row, idx, bool, bool_to_str),
-        Type::BPCHAR | Type::VARCHAR | Type::TEXT => {
-            make_string!(row, idx, &str, varchar_to_str)
-        }
-        Type::FLOAT4 => make_string!(row, idx, f32, f32_to_str),
-        Type::FLOAT8 => make_string!(row, idx, f64, f64_to_str),
-        Type::REGTYPE => make_string!(row, idx, PgRegtype),
-        _ => unimplemented!("Unsupported type: {}", column.type_().name()),
-    }
-}
-
 #[async_trait]
 impl sqllogictest::AsyncDB for Postgres {
     type Error = Error;
     type ColumnType = DFColumnType;
 
-    async fn run(&mut self, sql: &str) -> Result<DFOutput, Self::Error> {
+    async fn run(
+        &mut self,
+        sql: &str,
+    ) -> Result<DBOutput<Self::ColumnType>, Self::Error> {
         println!(
             "[{}] Running query: \"{}\"",
             self.relative_path.display(),
@@ -290,27 +249,20 @@ impl sqllogictest::AsyncDB for Postgres {
             return Ok(DBOutput::StatementComplete(0));
         }
         let rows = self.client.query(sql, &[]).await?;
-        let output = rows
-            .iter()
-            .map(|row| {
-                row.columns()
-                    .iter()
-                    .enumerate()
-                    .map(|(idx, column)| cell_to_string(row, column, idx))
-                    .collect::<Vec<String>>()
-            })
-            .collect::<Vec<_>>();
 
-        if output.is_empty() {
-            let stmt = self.client.prepare(sql).await?;
-            Ok(DBOutput::Rows {
-                types: vec![DFColumnType::Any; stmt.columns().len()],
-                rows: vec![],
-            })
+        if rows.is_empty() {
+            return Ok(DBOutput::StatementComplete(0));
         } else {
+            let types = convert_types(
+                rows[0]
+                    .columns()
+                    .iter()
+                    .map(|c| c.type_().clone())
+                    .collect(),
+            );
             Ok(DBOutput::Rows {
-                types: vec![DFColumnType::Any; output[0].len()],
-                rows: output,
+                types,
+                rows: convert_rows(rows),
             })
         }
     }
@@ -319,3 +271,73 @@ impl sqllogictest::AsyncDB for Postgres {
         "postgres"
     }
 }
+
+fn convert_rows(rows: Vec<Row>) -> Vec<Vec<String>> {
+    rows.iter()
+        .map(|row| {
+            row.columns()
+                .iter()
+                .enumerate()
+                .map(|(idx, column)| cell_to_string(row, column, idx))
+                .collect::<Vec<String>>()
+        })
+        .collect::<Vec<_>>()
+}
+
+macro_rules! make_string {
+    ($row:ident, $idx:ident, $t:ty) => {{
+        let value: Option<$t> = $row.get($idx);
+        match value {
+            Some(value) => value.to_string(),
+            None => NULL_STR.to_string(),
+        }
+    }};
+    ($row:ident, $idx:ident, $t:ty, $convert:ident) => {{
+        let value: Option<$t> = $row.get($idx);
+        match value {
+            Some(value) => $convert(value).to_string(),
+            None => NULL_STR.to_string(),
+        }
+    }};
+}
+
+fn cell_to_string(row: &Row, column: &Column, idx: usize) -> String {
+    match column.type_().clone() {
+        Type::CHAR => make_string!(row, idx, i8),
+        Type::INT2 => make_string!(row, idx, i16),
+        Type::INT4 => make_string!(row, idx, i32),
+        Type::INT8 => make_string!(row, idx, i64),
+        Type::NUMERIC => make_string!(row, idx, Decimal, decimal_to_str),
+        Type::DATE => make_string!(row, idx, NaiveDate),
+        Type::TIME => make_string!(row, idx, NaiveTime),
+        Type::TIMESTAMP => {
+            let value: Option<NaiveDateTime> = row.get(idx);
+            value
+                .map(|d| format!("{d:?}"))
+                .unwrap_or_else(|| "NULL".to_string())
+        }
+        Type::BOOL => make_string!(row, idx, bool, bool_to_str),
+        Type::BPCHAR | Type::VARCHAR | Type::TEXT => {
+            make_string!(row, idx, &str, varchar_to_str)
+        }
+        Type::FLOAT4 => make_string!(row, idx, f32, f32_to_str),
+        Type::FLOAT8 => make_string!(row, idx, f64, f64_to_str),
+        Type::REGTYPE => make_string!(row, idx, PgRegtype),
+        _ => unimplemented!("Unsupported type: {}", column.type_().name()),
+    }
+}
+
+fn convert_types(types: Vec<Type>) -> Vec<DFColumnType> {
+    types
+        .into_iter()
+        .map(|t| match t {
+            Type::BOOL => DFColumnType::Boolean,
+            Type::INT2 | Type::INT4 | Type::INT8 => DFColumnType::Integer,
+            Type::BPCHAR | Type::VARCHAR | Type::TEXT => DFColumnType::String,
+            Type::FLOAT4 | Type::FLOAT8 | Type::NUMERIC => DFColumnType::FloatingPoint,
+            Type::DATE | Type::TIME => DFColumnType::DateTime,
+            Type::TIMESTAMP => DFColumnType::Timestamp,
+            _ => DFColumnType::Another,
+        })
+        .collect()
+}

datafusion/core/tests/sqllogictests/src/main.rs

@@ -19,14 +19,14 @@ use std::error::Error;
 use std::path::{Path, PathBuf};
 
 use log::info;
+use sqllogictest::strict_column_validator;
 
 use datafusion::prelude::{SessionConfig, SessionContext};
 
 use crate::engines::datafusion::DataFusion;
 use crate::engines::postgres::Postgres;
 
 mod engines;
-mod output;
 mod setup;
 mod utils;
 
@@ -68,6 +68,7 @@ async fn run_test_file(
     info!("Running with DataFusion runner: {}", path.display());
     let ctx = context_for_test_file(&relative_path).await;
     let mut runner = sqllogictest::Runner::new(DataFusion::new(ctx, relative_path));
+    runner.with_column_validator(strict_column_validator);
     runner.run_file_async(path).await?;
     Ok(())
 }