[sqllogictest] port tests in avro.rs to sqllogictest (#6362)

* feat: port tests in avro.rs to sqllogictest

* fix: add test setup for avro_query_multiple_files

* run cargo fmt

* Run hash collisions test with avro support too

* fix clippy

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: e1ijah1

.github/workflows/rust.yml

@@ -461,7 +461,7 @@ jobs:
       - name: Run tests
         run: |
           cd datafusion
-          cargo test --lib --tests --features=force_hash_collisions
+          cargo test --lib --tests --features=force_hash_collisions,avro
 
   cargo-toml-formatting-checks:
     name: check Cargo.toml formatting

datafusion/core/tests/sql/avro.rs

@@ -81,7 +81,6 @@ macro_rules! test_expression {
 pub mod aggregates;
 pub mod arrow_files;
 #[cfg(feature = "avro")]
-pub mod avro;
 pub mod create_drop;
 pub mod explain_analyze;
 pub mod expr;

datafusion/core/tests/sql/mod.rs

@@ -22,6 +22,7 @@ use std::thread;
 
 use log::info;
 use sqllogictest::strict_column_validator;
+use tempfile::TempDir;
 
 use datafusion::prelude::{SessionConfig, SessionContext};
 
@@ -83,7 +84,8 @@ async fn run_test_file(
     relative_path: PathBuf,
 ) -> Result<(), Box<dyn Error>> {
     info!("Running with DataFusion runner: {}", path.display());
-    let ctx = context_for_test_file(&relative_path).await;
+    let test_ctx = context_for_test_file(&relative_path).await;
+    let ctx = test_ctx.session_ctx().clone();
     let mut runner = sqllogictest::Runner::new(DataFusion::new(ctx, relative_path));
     runner.with_column_validator(strict_column_validator);
     runner.run_file_async(path).await?;
@@ -110,7 +112,8 @@ async fn run_complete_file(
 
     info!("Using complete mode to complete: {}", path.display());
 
-    let ctx = context_for_test_file(&relative_path).await;
+    let test_ctx = context_for_test_file(&relative_path).await;
+    let ctx = test_ctx.session_ctx().clone();
     let mut runner = sqllogictest::Runner::new(DataFusion::new(ctx, relative_path));
     let col_separator = " ";
     runner
@@ -160,27 +163,67 @@ fn read_dir_recursive<P: AsRef<Path>>(path: P) -> Box<dyn Iterator<Item = PathBu
 }
 
 /// Create a SessionContext, configured for the specific test
-async fn context_for_test_file(relative_path: &Path) -> SessionContext {
+async fn context_for_test_file(relative_path: &Path) -> TestContext {
     let config = SessionConfig::new()
         // hardcode target partitions so plans are deterministic
         .with_target_partitions(4);
 
-    let ctx = SessionContext::with_config(config);
+    let mut test_ctx = TestContext::new(SessionContext::with_config(config));
 
     match relative_path.file_name().unwrap().to_str().unwrap() {
         "aggregate.slt" => {
             info!("Registering aggregate tables");
-            setup::register_aggregate_tables(&ctx).await;
+            setup::register_aggregate_tables(test_ctx.session_ctx()).await;
         }
         "scalar.slt" => {
             info!("Registering scalar tables");
-            setup::register_scalar_tables(&ctx).await;
+            setup::register_scalar_tables(test_ctx.session_ctx()).await;
+        }
+        "avro.slt" => {
+            info!("Registering avro tables");
+            setup::register_avro_tables(&mut test_ctx).await;
         }
         _ => {
             info!("Using default SessionContext");
         }
     };
-    ctx
+    test_ctx
+}
+
+/// Context for running tests
+pub struct TestContext {
+    /// Context for running queries
+    ctx: SessionContext,
+    /// Temporary directory created and cleared at the end of the test
+    test_dir: Option<TempDir>,
+}
+
+impl TestContext {
+    fn new(ctx: SessionContext) -> Self {
+        Self {
+            ctx,
+            test_dir: None,
+        }
+    }
+
+    /// Enables the test directory feature. If not enabled,
+    /// calling `testdir_path` will result in a panic.
+    fn enable_testdir(&mut self) {
+        if self.test_dir.is_none() {
+            self.test_dir = Some(TempDir::new().expect("failed to create testdir"));
+        }
+    }
+
+    /// Returns the path to the test directory. Panics if the test
+    /// directory feature is not enabled via `enable_testdir`.
+    fn testdir_path(&self) -> &Path {
+        self.test_dir.as_ref().expect("testdir not enabled").path()
+    }
+
+    /// Returns a reference to the internal SessionContext
+    fn session_ctx(&self) -> &SessionContext {
+        &self.ctx
+    }
 }
 
 /// Parsed command line options

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

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use datafusion::prelude::AvroReadOptions;
 use datafusion::{
     arrow::{
         array::{
@@ -30,7 +31,40 @@ use datafusion::{
 };
 use std::sync::Arc;
 
-use crate::utils;
+use crate::{utils, TestContext};
+
+pub async fn register_avro_tables(ctx: &mut TestContext) {
+    register_avro_test_data(ctx).await;
+}
+
+async fn register_avro_test_data(ctx: &mut TestContext) {
+    ctx.enable_testdir();
+
+    let table_path = ctx.testdir_path().join("avro");
+    std::fs::create_dir(&table_path).expect("failed to create avro table path");
+
+    let testdata = datafusion::test_util::arrow_test_data();
+    let alltypes_plain_file = format!("{testdata}/avro/alltypes_plain.avro");
+    std::fs::copy(
+        &alltypes_plain_file,
+        format!("{}/alltypes_plain1.avro", table_path.display()),
+    )
+    .unwrap();
+    std::fs::copy(
+        &alltypes_plain_file,
+        format!("{}/alltypes_plain2.avro", table_path.display()),
+    )
+    .unwrap();
+
+    ctx.session_ctx()
+        .register_avro(
+            "alltypes_plain_multi_files",
+            table_path.display().to_string().as_str(),
+            AvroReadOptions::default(),
+        )
+        .await
+        .unwrap();
+}
 
 pub async fn register_aggregate_tables(ctx: &SessionContext) {
     register_aggregate_test_100(ctx).await;