Vortex Session (#5111)

# VortexSession

We have had a `VortexSession` object for a while, but never really made
use of it, in a large part because none of the Vortex components could
take a dependency without creating a cycle.

Since Vortex is incredibly extensible, the idea of the session is to
hold the plugin registries for arrays, layouts, expressions and so on.
While our existing solution does this, it created very fragile APIs
where it was possible to construct file readers/writers _without_ using
the main registry of plugins. Similarly, metrics registries were getting
lost and some metrics ended up in isolated `VortexMetrics::default()`
instances - never to see the light of day.

This PR introduces a top-level VortexSession create that essentially
acts as a type map. Components of Vortex can register session state, and
then read it back again later. This allows components to accept and hold
a VortexSession internally, while our root `vortex` crate performs a
default configuration.

# Runtime Handles

One of the most fragile bits of mis-configured API was the Vortex
runtime handle. This is our abstraction over runtimes (in practice, only
over Tokio or a custom smol-based CurrentThreadRuntime) and provides an
API for components to spawn CPU, I/O, blocking and other types of work.

This PR demonstrates the session API by holding a runtime handle and
plumbing this through APIs that formerly accepted a with_handle
argument.

# Language Bindings

This PR doesn't not completely migrate language bindings onto a
session-based API. Many bindings still use a global static session where
a session instance would be preferred.

---------

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

Cargo.lock

@@ -28,6 +28,7 @@ members = [
     "vortex-python",
     "vortex-scan",
     "vortex-scalar",
+    "vortex-session",
     "vortex-tui",
     "vortex-utils",
     "vortex-vector",
@@ -245,6 +246,7 @@ vortex-runend = { version = "0.1.0", path = "./encodings/runend", default-featur
 vortex-scalar = { version = "0.1.0", path = "./vortex-scalar", default-features = false }
 vortex-scan = { version = "0.1.0", path = "./vortex-scan", default-features = false }
 vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-features = false }
+vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false }
 vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false }
 vortex-tui = { version = "0.1.0", path = "./vortex-tui", default-features = false }
 vortex-utils = { version = "0.1.0", path = "./vortex-utils", default-features = false }

Cargo.toml

@@ -26,14 +26,14 @@ use tokio::fs::{OpenOptions, create_dir_all};
 use tracing::{Instrument, info, warn};
 use url::Url;
 use vortex::error::VortexExpect;
-use vortex::file::VortexWriteOptions;
+use vortex::file::WriteOptionsSessionExt;
 use vortex_datafusion::VortexFormat;
 
 use crate::conversions::parquet_to_vortex;
 #[cfg(feature = "lance")]
 use crate::utils;
 use crate::utils::file_utils::{idempotent, idempotent_async};
-use crate::{CompactionStrategy, Format};
+use crate::{CompactionStrategy, Format, SESSION};
 
 pub static HITS_SCHEMA: LazyLock<Schema> = LazyLock::new(|| {
     use DataType::*;
@@ -204,7 +204,7 @@ pub async fn convert_parquet_to_vortex(
                             .open(&vtx_file)
                             .await?;
 
-                        let write_options = compaction.apply_options(VortexWriteOptions::default());
+                        let write_options = compaction.apply_options(SESSION.write_options());
 
                         write_options.write(&mut f, array_stream).await?;
 
@@ -250,7 +250,7 @@ pub async fn register_vortex_files(
     glob_pattern: Option<Pattern>,
 ) -> anyhow::Result<()> {
     let vortex_path = input_path.join(&format!("{}/", Format::OnDiskVortex.name()))?;
-    let format = Arc::new(VortexFormat::default());
+    let format = Arc::new(VortexFormat::new(SESSION.clone()));
 
     info!(
         "Registering table from {vortex_path} with glob {:?}",
@@ -283,7 +283,7 @@ pub async fn register_vortex_compact_files(
     glob_pattern: Option<Pattern>,
 ) -> anyhow::Result<()> {
     let vortex_compact_path = input_path.join(&format!("{}/", Format::VortexCompact.name()))?;
-    let format = Arc::new(VortexFormat::default());
+    let format = Arc::new(VortexFormat::new(SESSION.clone()));
 
     info!(
         "Registering vortex-compact table from {vortex_compact_path} with glob {:?}",

bench-vortex/src/clickbench/clickbench_data.rs

@@ -7,20 +7,23 @@ use std::sync::Arc;
 use bytes::Bytes;
 use futures::{StreamExt, pin_mut};
 use vortex::Array;
-use vortex::file::{VortexOpenOptions, VortexWriteOptions};
+use vortex::file::{OpenOptionsSessionExt, WriteOptionsSessionExt};
+
+use crate::SESSION;
 
 #[inline(never)]
 pub async fn vortex_compress_write(array: &dyn Array, buf: &mut Vec<u8>) -> anyhow::Result<u64> {
     let mut cursor = Cursor::new(buf);
-    VortexWriteOptions::default()
+    SESSION
+        .write_options()
         .write(&mut cursor, array.to_array_stream())
         .await?;
     Ok(cursor.position())
 }
 
 #[inline(never)]
 pub async fn vortex_decompress_read(buf: Bytes) -> anyhow::Result<usize> {
-    let scan = VortexOpenOptions::new().open_buffer(buf)?.scan()?;
+    let scan = SESSION.open_options().open_buffer(buf)?.scan()?;
     let schema = Arc::new(scan.dtype()?.to_arrow_schema()?);
 
     let stream = scan.into_record_batch_stream(schema)?;

bench-vortex/src/compress/vortex.rs

@@ -17,6 +17,7 @@ use vortex_datafusion::VortexFormat;
 #[cfg(feature = "lance")]
 use {crate::Format, lance::datafusion::LanceTableProvider, lance::dataset::Dataset};
 
+use crate::SESSION;
 use crate::datasets::BenchmarkDataset;
 
 pub async fn register_parquet_files(
@@ -85,7 +86,7 @@ pub async fn register_vortex_files(
                 &file_url,
                 glob.as_ref().map(|g| g.as_str()).unwrap_or("")
             );
-            let format = Arc::new(VortexFormat::default());
+            let format = Arc::new(VortexFormat::new(SESSION.clone()));
             let table_url = ListingTableUrl::try_new(file_url.clone(), glob)?;
             let config = ListingTableConfig::new(table_url).with_listing_options(
                 ListingOptions::new(format).with_session_config_options(session.state().config()),
@@ -133,7 +134,7 @@ pub async fn register_vortex_compact_files(
                 &file_url,
                 glob.as_ref().map(|g| g.as_str()).unwrap_or("")
             );
-            let format = Arc::new(VortexFormat::default());
+            let format = Arc::new(VortexFormat::new(SESSION.clone()));
             let table_url = ListingTableUrl::try_new(file_url.clone(), glob)?;
             let config = ListingTableConfig::new(table_url).with_listing_options(
                 ListingOptions::new(format).with_session_config_options(session.state().config()),

bench-vortex/src/datasets/file.rs

@@ -8,7 +8,7 @@ use async_trait::async_trait;
 use tokio::fs::File as TokioFile;
 use tokio::io::AsyncWriteExt;
 use vortex::ArrayRef;
-use vortex::file::{VortexOpenOptions, VortexWriteOptions};
+use vortex::file::{OpenOptionsSessionExt, WriteOptionsSessionExt};
 use vortex::stream::ArrayStreamExt;
 #[cfg(feature = "lance")]
 use {
@@ -21,7 +21,7 @@ use {
 use crate::conversions::parquet_to_vortex;
 use crate::datasets::Dataset;
 use crate::datasets::data_downloads::download_data;
-use crate::{CompactionStrategy, IdempotentPath, idempotent_async};
+use crate::{CompactionStrategy, IdempotentPath, SESSION, idempotent_async};
 
 pub struct TaxiData;
 
@@ -45,7 +45,8 @@ pub async fn taxi_data_parquet() -> Result<PathBuf> {
 
 pub async fn fetch_taxi_data() -> Result<ArrayRef> {
     let vortex_data = taxi_data_vortex().await?;
-    Ok(VortexOpenOptions::new()
+    Ok(SESSION
+        .open_options()
         .open(vortex_data)
         .await?
         .scan()?
@@ -58,7 +59,8 @@ pub async fn taxi_data_vortex() -> Result<PathBuf> {
     idempotent_async("taxi/taxi.vortex", |output_fname| async move {
         let buf = output_fname.to_path_buf();
         let mut output_file = TokioFile::create(output_fname).await?;
-        VortexWriteOptions::default()
+        SESSION
+            .write_options()
             .write(
                 &mut output_file,
                 parquet_to_vortex(taxi_data_parquet().await?)?,
@@ -76,8 +78,7 @@ pub async fn taxi_data_vortex_compact() -> Result<PathBuf> {
         let mut output_file = TokioFile::create(output_fname).await?;
 
         // This is the only difference to `taxi_data_vortex`.
-        let write_options =
-            CompactionStrategy::Compact.apply_options(VortexWriteOptions::default());
+        let write_options = CompactionStrategy::Compact.apply_options(SESSION.write_options());
 
         write_options
             .write(

bench-vortex/src/datasets/taxi_data.rs

@@ -8,12 +8,12 @@ use glob::glob;
 use vortex::arrays::ChunkedArray;
 use vortex::dtype::Nullability::NonNullable;
 use vortex::expr::{col, pack};
-use vortex::file::VortexOpenOptions;
+use vortex::file::OpenOptionsSessionExt;
 use vortex::{Array, ArrayRef, IntoArray, ToCanonical};
 
 use crate::datasets::Dataset;
 use crate::tpch::tpchgen::{TpchGenOptions, generate_tpch_tables};
-use crate::{Format, IdempotentPath};
+use crate::{Format, IdempotentPath, SESSION};
 
 pub struct TPCHLCommentChunked;
 
@@ -44,7 +44,7 @@ impl Dataset for TPCHLCommentChunked {
                 .to_string_lossy()
                 .as_ref(),
         )? {
-            let file = VortexOpenOptions::new().open(path?).await?;
+            let file = SESSION.open_options().open(path?).await?;
             let file_chunks: Vec<_> = file
                 .scan()?
                 .with_projection(pack(vec![("l_comment", col("l_comment"))], NonNullable))

bench-vortex/src/datasets/tpch_l_comment.rs

@@ -4,13 +4,13 @@
 use async_trait::async_trait;
 use tokio::fs::File;
 use vortex::ArrayRef;
-use vortex::file::{VortexOpenOptions, VortexWriteOptions};
+use vortex::file::{OpenOptionsSessionExt, WriteOptionsSessionExt};
 use vortex::stream::ArrayStreamExt;
 
 use crate::conversions::parquet_to_vortex;
 use crate::datasets::Dataset;
 use crate::datasets::data_downloads::download_data;
-use crate::{IdempotentPath, idempotent_async};
+use crate::{IdempotentPath, SESSION, idempotent_async};
 
 /// Datasets which can be downloaded over HTTP in Parquet format.
 ///
@@ -57,7 +57,8 @@ impl Dataset for DownloadableDataset {
         let vortex = dir.join(format!("{}.vortex", self.name()));
         download_data(parquet.clone(), self.parquet_url()).await?;
         idempotent_async(&vortex, async |path| -> anyhow::Result<()> {
-            VortexWriteOptions::default()
+            SESSION
+                .write_options()
                 .write(
                     &mut File::create(path)
                         .await
@@ -70,7 +71,8 @@ impl Dataset for DownloadableDataset {
         })
         .await?;
 
-        Ok(VortexOpenOptions::new()
+        Ok(SESSION
+            .open_options()
             .open(vortex.as_path())
             .await?
             .scan()?

bench-vortex/src/downloadable_dataset.rs

@@ -14,13 +14,13 @@ use log::info;
 use parquet::arrow::async_writer::AsyncFileWriter;
 use url::Url;
 use vortex::compressor::CompactCompressor;
-use vortex::file::{VortexWriteOptions, WriteStrategyBuilder};
+use vortex::file::{WriteOptionsSessionExt, WriteStrategyBuilder};
 use vortex_datafusion::VortexFormat;
 
 use crate::benchmark_trait::Benchmark;
 use crate::conversions::parquet_to_vortex;
 use crate::engines::EngineCtx;
-use crate::{BenchmarkDataset, Format, Target, idempotent_async};
+use crate::{BenchmarkDataset, Format, SESSION, Target, idempotent_async};
 
 /// URL to the sample file
 const SAMPLE_URL: &str = "https://huggingface.co/datasets/HuggingFaceFW/fineweb/resolve/v1.4.0/sample/10BT/001_00000.parquet";
@@ -146,7 +146,8 @@ impl Benchmark for Fineweb {
                     info!("Converting FineWeb to Vortex with default compressor");
                     let array_stream = parquet_to_vortex(parquet)?;
                     let w = tokio::fs::File::create(vortex_path).await?;
-                    VortexWriteOptions::default()
+                    SESSION
+                        .write_options()
                         .write(w, array_stream)
                         .await
                         .map_err(|e| anyhow::anyhow!("Failed to write to VortexWriter: {e}"))
@@ -158,7 +159,8 @@ impl Benchmark for Fineweb {
                     info!("Converting FineWeb to Vortex with Compact compressor");
                     let array_stream = parquet_to_vortex(parquet)?;
                     let w = tokio::fs::File::create(vortex_path).await?;
-                    VortexWriteOptions::default()
+                    SESSION
+                        .write_options()
                         .with_strategy(
                             WriteStrategyBuilder::new()
                                 .with_compressor(CompactCompressor::default())
@@ -221,7 +223,9 @@ pub async fn register_table(
         .with_listing_options(
             ListingOptions::new(match format {
                 Format::Parquet => Arc::from(ParquetFormat::new()),
-                Format::OnDiskVortex | Format::VortexCompact => Arc::from(VortexFormat::default()),
+                Format::OnDiskVortex | Format::VortexCompact => {
+                    Arc::from(VortexFormat::new(SESSION.clone()))
+                }
                 _ => anyhow::bail!("unsupported format for `fineweb` bench: {}", format),
             })
             .with_session_config_options(session.state().config()),