test: add sqllogictest runner

This change adds a sqllogictest runner + CLI which runs .slt files in the tests/sqllogictest directory. Right
now, the runner uses a hardcoded distributed config (ex. `with_network_shuffle_tasks(2)` etc.) but can be
extended in the future. Using sqllogictest will make it much easier to write tests (especially for `explain`
and `explain (analyze)`). `explain (analyze)` tests will be added in a future commit as it is still being
implemented.

Also, this change deletes `tests/distributed_aggregation.rs` and moves the test cases to `.slt` files.

Documentation:
- https://sqlite.org/sqllogictest/doc/trunk/about.wiki
- https://github.com/risinglightdb/sqllogictest-rs

Author: jayshrivastava

Cargo.lock

@@ -39,6 +39,10 @@ parquet = { version = "55.2.0", optional = true }
 arrow = { version = "55.2.0", optional = true }
 tokio-stream = { version = "0.1.17", optional = true }
 hyper-util = { version = "0.1.16", optional = true }
+sqllogictest = { version = "0.20", optional = true }
+regex = { version = "1.0", optional = true }
+clap = { version = "4.0", features = ["derive"], optional = true }
+env_logger = { version = "0.10", optional = true }
 pin-project = "1.1.10"
 
 [features]
@@ -50,6 +54,10 @@ integration = [
     "arrow",
     "tokio-stream",
     "hyper-util",
+    "sqllogictest",
+    "regex",
+    "clap",
+    "env_logger",
 ]
 
 tpch = ["integration"]

Cargo.toml

@@ -0,0 +1,98 @@
+use clap::Parser;
+use datafusion_distributed::test_utils::sqllogictest::DatafusionDistributedDB;
+use sqllogictest::Runner;
+use std::path::PathBuf;
+
+#[derive(Parser)]
+#[command(name = "logictest")]
+#[command(
+    about = "A SQLLogicTest runner for DataFusion Distributed. Docs: https://sqlite.org/sqllogictest/doc/trunk/about.wiki"
+)]
+struct Args {
+    /// Files or directories to run
+    #[arg(required = true)]
+    files: Vec<PathBuf>,
+
+    /// Update test files with actual output rather than verifying the existing output
+    #[arg(long = "override")]
+    override_mode: bool,
+
+    /// Number of workers
+    #[arg(long, default_value = "3")]
+    num_workers: usize,
+}
+
+async fn run<D, M>(paths: Vec<PathBuf>, runner: &mut Runner<D, M>, override_mode: bool)
+where
+    D: sqllogictest::AsyncDB,
+    M: sqllogictest::MakeConnection<Conn = D>,
+{
+    let mut queue = paths;
+    let mut idx = 0;
+    while idx < queue.len() {
+        let file_path = &queue[idx];
+        idx += 1;
+        if !file_path.is_file() {
+            queue.extend(
+                expand_directory(file_path).await.unwrap_or_else(|_| {
+                    panic!("Failed to expand directory: {}", file_path.display())
+                }),
+            );
+            continue;
+        }
+        let file_path_str = file_path.to_str().expect("Invalid file path");
+
+        let result = match override_mode {
+            true => {
+                runner
+                    .update_test_file(
+                        file_path_str,
+                        " ",
+                        sqllogictest::default_validator,
+                        sqllogictest::default_column_validator,
+                    )
+                    .await
+            }
+
+            false => runner.run_file_async(file_path).await.map_err(|e| e.into()),
+        };
+        match result {
+            Ok(_) => println!("🟢 Success: {}", file_path.display()),
+            Err(e) => eprintln!("🔴 Failure: {}:\n{e}", file_path.display()),
+        }
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args = Args::parse();
+
+    let mut runner =
+        Runner::new(
+            move || async move { Ok(DatafusionDistributedDB::new(args.num_workers).await) },
+        );
+
+    run(args.files, &mut runner, args.override_mode).await;
+
+    Ok(())
+}
+
+async fn expand_directory(dir_path: &PathBuf) -> Result<Vec<PathBuf>, Box<dyn std::error::Error>> {
+    let mut entries: Vec<_> = std::fs::read_dir(dir_path)?
+        .filter_map(|entry| entry.ok())
+        .filter(|entry| {
+            entry
+                .path()
+                .extension()
+                .and_then(|ext| ext.to_str())
+                .map(|ext| ext == "slt")
+                .unwrap_or(false)
+        })
+        .map(|entry| entry.path())
+        .collect();
+
+    // Sort entries for consistent order
+    entries.sort();
+
+    Ok(entries)
+}

src/bin/logictest.rs

@@ -5,4 +5,5 @@ pub mod metrics;
 pub mod mock_exec;
 pub mod parquet;
 pub mod session_context;
+pub mod sqllogictest;
 pub mod tpch;

src/test_utils/mod.rs

@@ -0,0 +1,145 @@
+use crate::DefaultSessionBuilder;
+use crate::DistributedPhysicalOptimizerRule;
+use crate::test_utils::localhost::start_localhost_context;
+use crate::test_utils::parquet::register_parquet_tables;
+use async_trait::async_trait;
+use datafusion::arrow::array::RecordBatch;
+use datafusion::arrow::array::{ArrayRef, StringArray};
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::arrow::util::display::array_value_to_string;
+use datafusion::common::runtime::JoinSet;
+use datafusion::error::DataFusionError;
+use datafusion::execution::context::SessionContext;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::displayable;
+use datafusion::physical_plan::execution_plan::collect;
+use sqllogictest::{AsyncDB, DBOutput, DefaultColumnType};
+use std::sync::Arc;
+
+pub struct DatafusionDistributedDB {
+    ctx: SessionContext,
+    _guard: JoinSet<()>,
+}
+
+impl DatafusionDistributedDB {
+    pub async fn new(num_nodes: usize) -> Self {
+        let (ctx, _guard) = start_localhost_context(num_nodes, DefaultSessionBuilder).await;
+        register_parquet_tables(&ctx).await.unwrap();
+        Self { ctx, _guard }
+    }
+
+    pub fn optimize_distributed(
+        plan: Arc<dyn ExecutionPlan>,
+    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
+        DistributedPhysicalOptimizerRule::default()
+            .with_network_shuffle_tasks(2)
+            .with_network_coalesce_tasks(2)
+            .optimize(plan, &Default::default())
+    }
+
+    fn convert_batches_to_output(
+        &self,
+        batches: Vec<RecordBatch>,
+    ) -> Result<DBOutput<DefaultColumnType>, datafusion::error::DataFusionError> {
+        if batches.is_empty() {
+            return Ok(DBOutput::Rows {
+                types: vec![],
+                rows: vec![],
+            });
+        }
+
+        let num_columns = batches[0].num_columns();
+        let column_types = vec![DefaultColumnType::Text; num_columns]; // Everything as text
+
+        let mut rows = Vec::new();
+        for batch in batches {
+            for row_idx in 0..batch.num_rows() {
+                let mut row = Vec::new();
+                for col_idx in 0..batch.num_columns() {
+                    let column = batch.column(col_idx);
+                    let value = array_value_to_string(column, row_idx)
+                        .map_err(|e| DataFusionError::ArrowError(Box::new(e), None))?;
+                    row.push(value);
+                }
+                rows.push(row);
+            }
+        }
+
+        Ok(DBOutput::Rows {
+            types: column_types,
+            rows,
+        })
+    }
+
+    async fn handle_explain_analyze(
+        &mut self,
+        _sql: &str,
+    ) -> Result<DBOutput<DefaultColumnType>, datafusion::error::DataFusionError> {
+        unimplemented!();
+    }
+
+    async fn handle_explain(
+        &mut self,
+        sql: &str,
+    ) -> Result<DBOutput<DefaultColumnType>, datafusion::error::DataFusionError> {
+        let query = sql.trim_start_matches("EXPLAIN").trim();
+        let df = self.ctx.sql(query).await?;
+        let physical_plan = df.create_physical_plan().await?;
+
+        let physical_distributed = Self::optimize_distributed(physical_plan)?;
+
+        let physical_distributed_str = displayable(physical_distributed.as_ref())
+            .indent(true)
+            .to_string();
+
+        let lines: Vec<String> = physical_distributed_str
+            .lines()
+            .map(|s| s.to_string())
+            .collect();
+        let schema = Arc::new(Schema::new(vec![Field::new("plan", DataType::Utf8, false)]));
+        let batch =
+            RecordBatch::try_new(schema, vec![Arc::new(StringArray::from(lines)) as ArrayRef])?;
+
+        self.convert_batches_to_output(vec![batch])
+    }
+}
+
+#[async_trait]
+impl AsyncDB for DatafusionDistributedDB {
+    type Error = datafusion::error::DataFusionError;
+    type ColumnType = DefaultColumnType;
+
+    async fn run(&mut self, sql: &str) -> Result<DBOutput<Self::ColumnType>, Self::Error> {
+        let sql = sql.trim();
+
+        // Ignore DDL/DML
+        if sql.to_uppercase().starts_with("CREATE")
+            || sql.to_uppercase().starts_with("INSERT")
+            || sql.to_uppercase().starts_with("DROP")
+        {
+            return Ok(DBOutput::StatementComplete(0));
+        }
+
+        if sql.to_uppercase().starts_with("EXPLAIN ANALYZE") {
+            return self.handle_explain_analyze(sql).await;
+        }
+
+        if sql.to_uppercase().starts_with("EXPLAIN") {
+            return self.handle_explain(sql).await;
+        }
+
+        // Default: Execute SELECT statement
+        let df = self.ctx.sql(sql).await?;
+        let task_ctx = Arc::new(df.task_ctx());
+        let plan = df.create_physical_plan().await?;
+        let distributed_plan = Self::optimize_distributed(plan)?;
+        let batches = collect(distributed_plan, task_ctx).await?;
+
+        self.convert_batches_to_output(batches)
+    }
+
+    fn engine_name(&self) -> &str {
+        "datafusion-distributed"
+    }
+}