File-based table provider for Datafusion

Cargo.toml

@@ -86,7 +86,7 @@ mimalloc = "0.1.42"
 monoio = "0.2.3"
 num-traits = "0.2.18"
 num_enum = "0.7.2"
-object_store = "0.10.1"
+object_store = "0.10.2"
 parquet = "52.0.0"
 paste = "1.0.14"
 pin-project = "1.1.5"

bench-vortex/benches/datafusion_benchmark.rs

@@ -16,7 +16,7 @@ use lazy_static::lazy_static;
 use vortex::compress::CompressionStrategy;
 use vortex::encoding::EncodingRef;
 use vortex::{Array, Context, IntoArray, ToArrayData};
-use vortex_datafusion::{VortexMemTable, VortexMemTableOptions};
+use vortex_datafusion::memory::{VortexMemTable, VortexMemTableOptions};
 use vortex_dict::DictEncoding;
 use vortex_fastlanes::{BitPackedEncoding, DeltaEncoding, FoREncoding};
 use vortex_sampling_compressor::compressors::bitpacked::BitPackedCompressor;

bench-vortex/benches/random_access.rs

@@ -8,6 +8,7 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use mimalloc::MiMalloc;
 use object_store::aws::AmazonS3Builder;
 use object_store::local::LocalFileSystem;
+use object_store::ObjectStore;
 use tokio::runtime::Runtime;
 
 #[global_allocator]
@@ -31,7 +32,7 @@ fn random_access_vortex(c: &mut Criterion) {
             .iter(|| async { black_box(take_vortex_tokio(&taxi_vortex, &INDICES).await.unwrap()) })
     });
 
-    let local_fs = LocalFileSystem::new();
+    let local_fs = Arc::new(LocalFileSystem::new()) as Arc<dyn ObjectStore>;
     let local_fs_path = object_store::path::Path::from_filesystem_path(&taxi_vortex).unwrap();
     group.bench_function("localfs", |b| {
         b.to_async(Runtime::new().unwrap()).iter(|| async {
@@ -43,7 +44,7 @@ fn random_access_vortex(c: &mut Criterion) {
         })
     });
 
-    let r2_fs = AmazonS3Builder::from_env().build().unwrap();
+    let r2_fs = Arc::new(AmazonS3Builder::from_env().build().unwrap()) as Arc<dyn ObjectStore>;
     let r2_path =
         object_store::path::Path::from_url_path(taxi_vortex.file_name().unwrap().to_str().unwrap())
             .unwrap();

bench-vortex/src/reader.rs

@@ -159,8 +159,8 @@ pub async fn read_vortex_footer_format<R: VortexReadAt>(
     )
 }
 
-pub async fn take_vortex_object_store<O: ObjectStore>(
-    fs: &O,
+pub async fn take_vortex_object_store(
+    fs: &Arc<dyn ObjectStore>,
     path: &object_store::path::Path,
     indices: &[u64],
 ) -> VortexResult<Array> {

bench-vortex/src/tpch/mod.rs

@@ -10,7 +10,8 @@ use datafusion::prelude::{CsvReadOptions, ParquetReadOptions, SessionContext};
 use vortex::array::chunked::ChunkedArray;
 use vortex::arrow::FromArrowArray;
 use vortex::{Array, ArrayDType, ArrayData, IntoArray};
-use vortex_datafusion::{SessionContextExt, VortexMemTableOptions};
+use vortex_datafusion::memory::VortexMemTableOptions;
+use vortex_datafusion::SessionContextExt;
 
 use crate::idempotent_async;
 

vortex-array/src/array/struct_/mod.rs

@@ -1,6 +1,6 @@
 use serde::{Deserialize, Serialize};
 use vortex_dtype::{DType, FieldName, FieldNames, Nullability, StructDType};
-use vortex_error::{vortex_bail, VortexResult};
+use vortex_error::{vortex_bail, vortex_err, VortexResult};
 
 use crate::stats::{ArrayStatisticsCompute, StatsSet};
 use crate::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
@@ -97,12 +97,12 @@ impl StructArray {
         let mut children = Vec::with_capacity(projection.len());
         let mut names = Vec::with_capacity(projection.len());
 
-        for column_idx in projection {
+        for &column_idx in projection {
             children.push(
-                self.field(*column_idx)
-                    .expect("column must not exceed bounds"),
+                self.field(column_idx)
+                    .ok_or(vortex_err!(OutOfBounds: column_idx, 0, self.dtypes().len()))?,
             );
-            names.push(self.names()[*column_idx].clone());
+            names.push(self.names()[column_idx].clone());
         }
 
         StructArray::try_new(
@@ -124,7 +124,9 @@ impl ArrayVariants for StructArray {
 
 impl StructArrayTrait for StructArray {
     fn field(&self, idx: usize) -> Option<Array> {
-        self.array().child(idx, &self.dtypes()[idx], self.len())
+        self.dtypes()
+            .get(idx)
+            .and_then(|dtype| self.array().child(idx, dtype, self.len()))
     }
 }
 

vortex-datafusion/Cargo.toml

@@ -10,17 +10,23 @@ include.workspace = true
 edition.workspace = true
 rust-version.workspace = true
 
+[lib]
+name = "vortex_datafusion"
+path = "src/lib.rs"
+
 [dependencies]
 vortex-array = { path = "../vortex-array" }
 vortex-dtype = { path = "../vortex-dtype" }
 vortex-expr = { path = "../vortex-expr" }
-vortex-error = { path = "../vortex-error" }
+vortex-error = { path = "../vortex-error", features = ["datafusion"] }
+vortex-serde = { path = "../vortex-serde", features = ["object_store"] }
 vortex-scalar = { path = "../vortex-scalar", features = ["datafusion"] }
 
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
 async-trait = { workspace = true }
 
+chrono = { workspace = true }
 datafusion = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-expr = { workspace = true }
@@ -30,10 +36,16 @@ datafusion-physical-plan = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
 lazy_static = { workspace = true }
+log = { workspace = true }
+object_store = { workspace = true }
 pin-project = { workspace = true }
 
+
 [dev-dependencies]
-tokio = { workspace = true, features = ["test-util"] }
+anyhow = { workspace = true }
+url = "2"
+tempfile = "3"
+tokio = { workspace = true, features = ["test-util", "rt-multi-thread"] }
 
 [lints]
 workspace = true

vortex-datafusion/examples/table_provider.rs

@@ -0,0 +1,98 @@
+use std::sync::Arc;
+
+use arrow_schema::{DataType, Field, Schema};
+use datafusion::prelude::SessionContext;
+use datafusion_execution::object_store::ObjectStoreUrl;
+use object_store::local::LocalFileSystem;
+use object_store::path::Path;
+use object_store::ObjectStore;
+use tempfile::tempdir;
+use tokio::fs::OpenOptions;
+use url::Url;
+use vortex::array::chunked::ChunkedArray;
+use vortex::array::primitive::PrimitiveArray;
+use vortex::array::struct_::StructArray;
+use vortex::array::varbin::VarBinArray;
+use vortex::validity::Validity;
+use vortex::IntoArray;
+use vortex_datafusion::persistent::config::{VortexFile, VortexTableConfig};
+use vortex_datafusion::persistent::provider::VortexFileTableProvider;
+use vortex_serde::file::file_writer::FileWriter;
+
+#[tokio::main]
+async fn main() -> 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([
+        PrimitiveArray::from(vec![1u32, 2, 3, 4]).into_array(),
+        PrimitiveArray::from(vec![5u32, 6, 7, 8]).into_array(),
+    ])
+    .into_array();
+
+    let st = StructArray::try_new(
+        ["strings".into(), "numbers".into()].into(),
+        vec![strings, numbers],
+        8,
+        Validity::NonNullable,
+    )
+    .unwrap();
+
+    let filepath = temp_dir.path().join("a.vtx");
+
+    let f = OpenOptions::new()
+        .write(true)
+        .truncate(true)
+        .create(true)
+        .open(&filepath)
+        .await?;
+
+    let writer = FileWriter::new(f);
+    let writer = writer.write_array_columns(st.into_array()).await?;
+    writer.finalize().await?;
+
+    let f = tokio::fs::File::open(&filepath).await?;
+    let file_size = f.metadata().await?.len();
+
+    let object_store: Arc<dyn ObjectStore> = Arc::new(LocalFileSystem::new());
+    let url = ObjectStoreUrl::local_filesystem();
+
+    let p = Path::from_filesystem_path(filepath)?;
+
+    let config = VortexTableConfig::new(
+        Arc::new(Schema::new(vec![
+            Field::new("strings", DataType::Utf8, false),
+            Field::new("numbers", DataType::UInt32, false),
+        ])),
+        vec![VortexFile::new(p, file_size)],
+    );
+
+    let provider = Arc::new(VortexFileTableProvider::try_new(url, config)?);
+
+    let ctx = SessionContext::new();
+    ctx.register_table("vortex_tbl", Arc::clone(&provider) as _)?;
+
+    let url = Url::try_from("file://").unwrap();
+    ctx.register_object_store(&url, object_store);
+
+    run_query(&ctx, "SELECT * FROM vortex_tbl").await?;
+
+    Ok(())
+}
+
+async fn run_query(ctx: &SessionContext, query_string: impl AsRef<str>) -> anyhow::Result<()> {
+    let query_string = query_string.as_ref();
+
+    ctx.sql(&format!("EXPLAIN {query_string}"))
+        .await?
+        .show()
+        .await?;
+
+    ctx.sql(query_string).await?.show().await?;
+
+    Ok(())
+}