Fix datafusion example and also make it a test (#3134)

Honestly not sure if datafusion's listing logic changed in regareds to
file extensions or if I just made some last-second change and broke a
working example before committing it.

Author: AdamGS

vortex-datafusion/examples/vortex_table.rs

@@ -4,11 +4,8 @@ use datafusion::datasource::listing::{
     ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
 };
 use datafusion::prelude::SessionContext;
-use object_store::ObjectStore;
-use object_store::local::LocalFileSystem;
 use tempfile::tempdir;
 use tokio::fs::OpenOptions;
-use url::Url;
 use vortex_array::arrays::{ChunkedArray, StructArray, VarBinArray};
 use vortex_array::stream::ArrayStreamArrayExt;
 use vortex_array::validity::Validity;
@@ -40,7 +37,7 @@ async fn main() -> anyhow::Result<()> {
         Validity::NonNullable,
     )?;
 
-    let filepath = temp_dir.path().join("a.vtx");
+    let filepath = temp_dir.path().join("a.vortex");
 
     let f = OpenOptions::new()
         .write(true)
@@ -54,19 +51,14 @@ async fn main() -> anyhow::Result<()> {
         .await?;
 
     let ctx = SessionContext::new();
-
-    let object_store: Arc<dyn ObjectStore> = Arc::new(LocalFileSystem::new());
-    let url: Url = Url::try_from("file://")?;
-    ctx.register_object_store(&url, object_store);
-
     let format = Arc::new(VortexFormat::default());
     let table_url = ListingTableUrl::parse(
         filepath
             .to_str()
             .ok_or_else(|| vortex_err!("Path is not valid UTF-8"))?,
     )?;
     let config = ListingTableConfig::new(table_url)
-        .with_listing_options(ListingOptions::new(format as _))
+        .with_listing_options(ListingOptions::new(format))
         .infer_schema(&ctx.state())
         .await?;
 

vortex-datafusion/src/persistent/mod.rs

@@ -26,3 +26,83 @@ fn register_vortex_format_factory(
         file_formats.push(std::sync::Arc::new(factory));
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use datafusion::datasource::listing::{
+        ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
+    };
+    use datafusion::prelude::SessionContext;
+    use tempfile::tempdir;
+    use tokio::fs::OpenOptions;
+    use vortex_array::arrays::{ChunkedArray, StructArray, VarBinArray};
+    use vortex_array::stream::ArrayStreamArrayExt;
+    use vortex_array::validity::Validity;
+    use vortex_array::{Array, IntoArray};
+    use vortex_buffer::buffer;
+    use vortex_error::vortex_err;
+    use vortex_file::VortexWriteOptions;
+
+    use crate::persistent::VortexFormat;
+
+    #[tokio::test]
+    async fn query_file() -> anyhow::Result<()> {
+        let temp_dir = tempdir()?;
+        let strings = ChunkedArray::from_iter([
+            VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
+            VarBinArray::from(vec!["ab", "foo", "bar", "baz"]).into_array(),
+        ])
+        .into_array();
+
+        let numbers = ChunkedArray::from_iter([
+            buffer![1u32, 2, 3, 4].into_array(),
+            buffer![5u32, 6, 7, 8].into_array(),
+        ])
+        .into_array();
+
+        let st = StructArray::try_new(
+            ["strings".into(), "numbers".into()].into(),
+            vec![strings, numbers],
+            8,
+            Validity::NonNullable,
+        )?;
+
+        let filepath = temp_dir.path().join("data.vortex");
+
+        let f = OpenOptions::new()
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(&filepath)
+            .await?;
+
+        VortexWriteOptions::default()
+            .write(f, st.to_array_stream())
+            .await?;
+
+        let ctx = SessionContext::default();
+        let format = Arc::new(VortexFormat::default());
+        let table_url = ListingTableUrl::parse(
+            temp_dir
+                .path()
+                .to_str()
+                .ok_or_else(|| vortex_err!("Path is not valid UTF-8"))?,
+        )?;
+        assert!(table_url.is_collection());
+
+        let config = ListingTableConfig::new(table_url)
+            .with_listing_options(ListingOptions::new(format))
+            .infer_schema(&ctx.state())
+            .await?;
+
+        let listing_table = Arc::new(ListingTable::try_new(config)?);
+
+        ctx.register_table("vortex_tbl", listing_table as _)?;
+        let row_count = ctx.table("vortex_tbl").await?.count().await?;
+        assert_eq!(row_count, 8);
+
+        Ok(())
+    }
+}