DF Test refactoring (#6223)

This PR attempts to make DF tests easier to read and write by moving
much of the boilerplate into `TestSessionContext` and using more SQL
(where its applicable).

---------

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

vortex-datafusion/src/lib.rs

@@ -11,6 +11,9 @@ use vortex::expr::stats::Precision;
 mod convert;
 mod persistent;
 
+#[cfg(test)]
+mod tests;
+
 pub use convert::exprs::ExpressionConvertor;
 pub use persistent::*;
 
@@ -46,3 +49,101 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod common_tests {
+    use std::sync::Arc;
+    use std::sync::LazyLock;
+
+    use datafusion::arrow::array::RecordBatch;
+    use datafusion::datasource::provider::DefaultTableFactory;
+    use datafusion::execution::SessionStateBuilder;
+    use datafusion::prelude::SessionContext;
+    use datafusion_catalog::TableProvider;
+    use datafusion_common::DFSchema;
+    use datafusion_common::GetExt;
+    use datafusion_expr::CreateExternalTable;
+    use object_store::ObjectStore;
+    use object_store::memory::InMemory;
+    use url::Url;
+    use vortex::VortexSessionDefault;
+    use vortex::array::ArrayRef;
+    use vortex::array::arrow::FromArrowArray;
+    use vortex::file::WriteOptionsSessionExt;
+    use vortex::io::ObjectStoreWriter;
+    use vortex::io::VortexWrite;
+    use vortex::session::VortexSession;
+
+    use crate::VortexFormatFactory;
+
+    static VX_SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::default);
+
+    pub struct TestSessionContext {
+        pub store: Arc<dyn ObjectStore>,
+        pub session: SessionContext,
+    }
+
+    impl Default for TestSessionContext {
+        fn default() -> Self {
+            let store = Arc::new(InMemory::new());
+            let factory = Arc::new(VortexFormatFactory::new());
+            let session_state_builder = SessionStateBuilder::new()
+                .with_default_features()
+                .with_table_factory(
+                    factory.get_ext().to_uppercase(),
+                    Arc::new(DefaultTableFactory::new()),
+                )
+                .with_file_formats(vec![factory])
+                .with_object_store(&Url::try_from("file://").unwrap(), store.clone());
+
+            let session: SessionContext =
+                SessionContext::new_with_state(session_state_builder.build()).enable_url_table();
+
+            Self { store, session }
+        }
+    }
+
+    impl TestSessionContext {
+        // Write arrow data into a vortex file.
+        pub async fn write_arrow_batch<P>(&self, path: P, batch: &RecordBatch) -> anyhow::Result<()>
+        where
+            P: Into<object_store::path::Path>,
+        {
+            let array = ArrayRef::from_arrow(batch, false)?;
+            let mut write = ObjectStoreWriter::new(self.store.clone(), &path.into()).await?;
+            VX_SESSION
+                .write_options()
+                .write(&mut write, array.to_array_stream())
+                .await?;
+            write.shutdown().await?;
+
+            Ok(())
+        }
+
+        /// Creates a ListingTable provider targeted at the provided path
+        pub async fn table_provider<S>(
+            &self,
+            name: &str,
+            location: impl Into<String>,
+            schema: S,
+        ) -> anyhow::Result<Arc<dyn TableProvider>>
+        where
+            DFSchema: TryFrom<S>,
+            anyhow::Error: From<<S as TryInto<DFSchema>>::Error>,
+        {
+            let factory = self.session.table_factory("VORTEX").unwrap();
+
+            let cmd = CreateExternalTable::builder(
+                name,
+                location.into(),
+                "vortex",
+                DFSchema::try_from(schema)?.into(),
+            )
+            .build();
+
+            let table = factory.create(&self.session.state(), &cmd).await?;
+
+            Ok(table)
+        }
+    }
+}

vortex-datafusion/src/persistent/format.rs

@@ -503,59 +503,45 @@ impl FileFormat for VortexFormat {
 
 #[cfg(test)]
 mod tests {
-    use datafusion::execution::SessionStateBuilder;
-    use datafusion::prelude::SessionContext;
-    use tempfile::TempDir;
 
     use super::*;
-    use crate::persistent::register_vortex_format_factory;
+    use crate::common_tests::TestSessionContext;
 
     #[tokio::test]
-    async fn create_table() {
-        let dir = TempDir::new().unwrap();
+    async fn create_table() -> anyhow::Result<()> {
+        let ctx = TestSessionContext::default();
 
-        let factory: VortexFormatFactory = VortexFormatFactory::new();
-        let mut session_state_builder = SessionStateBuilder::new().with_default_features();
-        register_vortex_format_factory(factory, &mut session_state_builder);
-        let session = SessionContext::new_with_state(session_state_builder.build());
-
-        let df = session
-            .sql(&format!(
+        ctx.session
+            .sql(
                 "CREATE EXTERNAL TABLE my_tbl \
                 (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
                 STORED AS vortex  \
-                LOCATION '{}'",
-                dir.path().to_str().unwrap()
-            ))
-            .await
-            .unwrap();
+                LOCATION 'table/'",
+            )
+            .await?;
 
-        assert_eq!(df.count().await.unwrap(), 0);
+        assert!(ctx.session.table_exist("my_tbl")?);
+
+        Ok(())
     }
 
     #[tokio::test]
-    async fn configure_format_source() {
-        let dir = TempDir::new().unwrap();
-
-        let factory = VortexFormatFactory::new();
-        let mut session_state_builder = SessionStateBuilder::new().with_default_features();
-        register_vortex_format_factory(factory, &mut session_state_builder);
-        let session = SessionContext::new_with_state(session_state_builder.build());
+    async fn configure_format_source() -> anyhow::Result<()> {
+        let ctx = TestSessionContext::default();
 
-        session
-            .sql(&format!(
+        ctx.session
+            .sql(
                 "CREATE EXTERNAL TABLE my_tbl \
                 (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
                 STORED AS vortex \
-                LOCATION '{}' \
+                LOCATION 'table/' \
                 OPTIONS( footer_initial_read_size_bytes '12345' );",
-                dir.path().to_str().unwrap()
-            ))
-            .await
-            .unwrap()
+            )
+            .await?
             .collect()
-            .await
-            .unwrap();
+            .await?;
+
+        Ok(())
     }
 
     #[test]