Add substrait roundtrip option in sqllogictests

Author: gabotechs

Cargo.lock

@@ -45,6 +45,7 @@ chrono = { workspace = true, optional = true }
 clap = { version = "4.5.39", features = ["derive", "env"] }
 datafusion = { workspace = true, default-features = true, features = ["avro"] }
 datafusion-spark = { workspace = true, default-features = true }
+datafusion-substrait = { path = "../substrait" }
 futures = { workspace = true }
 half = { workspace = true, default-features = true }
 indicatif = "0.17"

datafusion/sqllogictest/Cargo.toml

@@ -21,7 +21,8 @@ use datafusion::common::utils::get_available_parallelism;
 use datafusion::common::{exec_err, DataFusionError, Result};
 use datafusion_sqllogictest::{
     df_value_validator, read_dir_recursive, setup_scratch_dir, should_skip_file,
-    should_skip_record, value_normalizer, DataFusion, Filter, TestContext,
+    should_skip_record, value_normalizer, DataFusion, DataFusionSubstraitRoundTrip,
+    Filter, TestContext,
 };
 use futures::stream::StreamExt;
 use indicatif::{
@@ -102,6 +103,11 @@ async fn run_tests() -> Result<()> {
         // to stdout and return OK so they can continue listing other tests.
         return Ok(());
     }
+    if options.substrait_round_trip && (options.postgres_runner || options.complete) {
+        let msg = "--substrait-round-trip option is not supported with --postgres-runner or --complete";
+        return Err(DataFusionError::External(msg.into()));
+    }
+
     options.warn_on_ignored();
 
     #[cfg(feature = "postgres")]
@@ -138,8 +144,22 @@ async fn run_tests() -> Result<()> {
             let filters = options.filters.clone();
 
             SpawnedTask::spawn(async move {
-                match (options.postgres_runner, options.complete) {
-                    (false, false) => {
+                match (
+                    options.postgres_runner,
+                    options.complete,
+                    options.substrait_round_trip,
+                ) {
+                    (_, _, true) => {
+                        run_test_file_substrait_round_trip(
+                            test_file,
+                            validator,
+                            m_clone,
+                            m_style_clone,
+                            filters.as_ref(),
+                        )
+                        .await?
+                    }
+                    (false, false, _) => {
                         run_test_file(
                             test_file,
                             validator,
@@ -149,11 +169,11 @@ async fn run_tests() -> Result<()> {
                         )
                         .await?
                     }
-                    (false, true) => {
+                    (false, true, _) => {
                         run_complete_file(test_file, validator, m_clone, m_style_clone)
                             .await?
                     }
-                    (true, false) => {
+                    (true, false, _) => {
                         run_test_file_with_postgres(
                             test_file,
                             validator,
@@ -163,7 +183,7 @@ async fn run_tests() -> Result<()> {
                         )
                         .await?
                     }
-                    (true, true) => {
+                    (true, true, _) => {
                         run_complete_file_with_postgres(
                             test_file,
                             validator,
@@ -210,6 +230,45 @@ async fn run_tests() -> Result<()> {
     }
 }
 
+async fn run_test_file_substrait_round_trip(
+    test_file: TestFile,
+    validator: Validator,
+    mp: MultiProgress,
+    mp_style: ProgressStyle,
+    filters: &[Filter],
+) -> Result<()> {
+    let TestFile {
+        path,
+        relative_path,
+    } = test_file;
+    let Some(test_ctx) = TestContext::try_new_for_test_file(&relative_path).await else {
+        info!("Skipping: {}", path.display());
+        return Ok(());
+    };
+    setup_scratch_dir(&relative_path)?;
+
+    let count: u64 = get_record_count(&path, "DatafusionSubstraitRoundTrip".to_string());
+    let pb = mp.add(ProgressBar::new(count));
+
+    pb.set_style(mp_style);
+    pb.set_message(format!("{:?}", &relative_path));
+
+    let mut runner = sqllogictest::Runner::new(|| async {
+        Ok(DataFusionSubstraitRoundTrip::new(
+            test_ctx.session_ctx().clone(),
+            relative_path.clone(),
+            pb.clone(),
+        ))
+    });
+    runner.add_label("DatafusionSubstraitRoundTrip");
+    runner.with_column_validator(strict_column_validator);
+    runner.with_normalizer(value_normalizer);
+    runner.with_validator(validator);
+    let res = run_file_in_runner(path, runner, filters).await;
+    pb.finish_and_clear();
+    res
+}
+
 async fn run_test_file(
     test_file: TestFile,
     validator: Validator,
@@ -578,6 +637,12 @@ struct Options {
     )]
     postgres_runner: bool,
 
+    #[clap(
+        long,
+        help = "Before executing each query, convert its logical plan to Substrait and from Substrait back to its logical plan"
+    )]
+    substrait_round_trip: bool,
+
     #[clap(long, env = "INCLUDE_SQLITE", help = "Include sqlite files")]
     include_sqlite: bool,
 

datafusion/sqllogictest/bin/sqllogictests.rs

@@ -0,0 +1,3 @@
+mod runner;
+
+pub use runner::*;

datafusion/sqllogictest/src/engines/datafusion_substrait_roundtrip_engine/mod.rs

@@ -0,0 +1,154 @@
+// 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 std::sync::Arc;
+use std::{path::PathBuf, time::Duration};
+
+use crate::engines::datafusion_engine::Result;
+use crate::engines::output::{DFColumnType, DFOutput};
+use crate::{convert_batches, convert_schema_to_types, DFSqlLogicTestError};
+use arrow::record_batch::RecordBatch;
+use async_trait::async_trait;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::common::collect;
+use datafusion::physical_plan::execute_stream;
+use datafusion::prelude::SessionContext;
+use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+use datafusion_substrait::logical_plan::producer::to_substrait_plan;
+use indicatif::ProgressBar;
+use log::Level::{Debug, Info};
+use log::{debug, log_enabled, warn};
+use sqllogictest::DBOutput;
+use tokio::time::Instant;
+
+pub struct DataFusionSubstraitRoundTrip {
+    ctx: SessionContext,
+    relative_path: PathBuf,
+    pb: ProgressBar,
+}
+
+impl DataFusionSubstraitRoundTrip {
+    pub fn new(ctx: SessionContext, relative_path: PathBuf, pb: ProgressBar) -> Self {
+        Self {
+            ctx,
+            relative_path,
+            pb,
+        }
+    }
+
+    fn update_slow_count(&self) {
+        let msg = self.pb.message();
+        let split: Vec<&str> = msg.split(" ").collect();
+        let mut current_count = 0;
+
+        if split.len() > 2 {
+            // third match will be current slow count
+            current_count = split[2].parse::<i32>().unwrap();
+        }
+
+        current_count += 1;
+
+        self.pb
+            .set_message(format!("{} - {} took > 500 ms", split[0], current_count));
+    }
+}
+
+#[async_trait]
+impl sqllogictest::AsyncDB for DataFusionSubstraitRoundTrip {
+    type Error = DFSqlLogicTestError;
+    type ColumnType = DFColumnType;
+
+    async fn run(&mut self, sql: &str) -> Result<DFOutput> {
+        if log_enabled!(Debug) {
+            debug!(
+                "[{}] Running query: \"{}\"",
+                self.relative_path.display(),
+                sql
+            );
+        }
+
+        let start = Instant::now();
+        let result = run_query_substrait_round_trip(&self.ctx, sql).await;
+        let duration = start.elapsed();
+
+        if duration.gt(&Duration::from_millis(500)) {
+            self.update_slow_count();
+        }
+
+        self.pb.inc(1);
+
+        if log_enabled!(Info) && duration.gt(&Duration::from_secs(2)) {
+            warn!(
+                "[{}] Running query took more than 2 sec ({duration:?}): \"{sql}\"",
+                self.relative_path.display()
+            );
+        }
+
+        result
+    }
+
+    /// Engine name of current database.
+    fn engine_name(&self) -> &str {
+        "DataFusionSubstraitRoundTrip"
+    }
+
+    /// [`DataFusion`] calls this function to perform sleep.
+    ///
+    /// The default implementation is `std::thread::sleep`, which is universal to any async runtime
+    /// but would block the current thread. If you are running in tokio runtime, you should override
+    /// this by `tokio::time::sleep`.
+    async fn sleep(dur: Duration) {
+        tokio::time::sleep(dur).await;
+    }
+
+    async fn shutdown(&mut self) {}
+}
+
+async fn run_query_substrait_round_trip(
+    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 state = ctx.state();
+    let round_tripped_plan = match df.logical_plan() {
+        // Substrait does not handle these plans
+        LogicalPlan::Ddl(_)
+        | LogicalPlan::Explain(_)
+        | LogicalPlan::Dml(_)
+        | LogicalPlan::Copy(_)
+        | LogicalPlan::Statement(_) => df.logical_plan().clone(),
+        // For any other plan, convert to Substrait
+        logical_plan => {
+            let plan = to_substrait_plan(logical_plan, &state)?;
+            from_substrait_plan(&state, &plan).await?
+        }
+    };
+
+    let physical_plan = state.create_physical_plan(&round_tripped_plan).await?;
+    let stream = execute_stream(physical_plan, task_ctx)?;
+    let types = convert_schema_to_types(stream.schema().fields());
+    let results: Vec<RecordBatch> = collect(stream).await?;
+    let rows = convert_batches(results, false)?;
+
+    if rows.is_empty() && types.is_empty() {
+        Ok(DBOutput::StatementComplete(0))
+    } else {
+        Ok(DBOutput::Rows { types, rows })
+    }
+}

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

@@ -18,12 +18,14 @@
 /// Implementation of sqllogictest for datafusion.
 mod conversion;
 mod datafusion_engine;
+mod datafusion_substrait_roundtrip_engine;
 mod output;
 
 pub use datafusion_engine::convert_batches;
 pub use datafusion_engine::convert_schema_to_types;
 pub use datafusion_engine::DFSqlLogicTestError;
 pub use datafusion_engine::DataFusion;
+pub use datafusion_substrait_roundtrip_engine::DataFusionSubstraitRoundTrip;
 pub use output::DFColumnType;
 pub use output::DFOutput;
 

datafusion/sqllogictest/src/engines/mod.rs

@@ -34,6 +34,7 @@ pub use engines::DFColumnType;
 pub use engines::DFOutput;
 pub use engines::DFSqlLogicTestError;
 pub use engines::DataFusion;
+pub use engines::DataFusionSubstraitRoundTrip;
 
 #[cfg(feature = "postgres")]
 pub use engines::Postgres;