experiment: eagerly decode dict on read

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

Cargo.lock

@@ -44,6 +44,7 @@ vortex-utils = { workspace = true, features = ["dashmap"] }
 anyhow = { workspace = true }
 datafusion = { workspace = true }
 insta = { workspace = true }
+mimalloc = "0.1"
 rstest = { workspace = true }
 tempfile = { workspace = true }
 tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "fs"] }

vortex-datafusion/Cargo.toml

@@ -0,0 +1,77 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use datafusion::datasource::provider::DefaultTableFactory;
+use datafusion::execution::SessionStateBuilder;
+use datafusion::prelude::SessionContext;
+use datafusion_common::GetExt;
+use datafusion_execution::config::SessionConfig;
+use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+use futures::StreamExt;
+use mimalloc::MiMalloc;
+use vortex_datafusion::VortexFormatFactory;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+fn get_session_context() -> SessionContext {
+    let rt_builder = RuntimeEnvBuilder::new();
+
+    let rt = rt_builder
+        .build_arc()
+        .expect("could not build runtime environment");
+
+    let factory = VortexFormatFactory::new();
+
+    let mut session_state_builder = SessionStateBuilder::new()
+        .with_config(SessionConfig::default())
+        .with_runtime_env(rt)
+        .with_default_features();
+
+    if let Some(table_factories) = session_state_builder.table_factories() {
+        table_factories.insert(
+            GetExt::get_ext(&factory).to_uppercase(), // Has to be uppercase
+            Arc::new(DefaultTableFactory::new()),
+        );
+    }
+
+    if let Some(file_formats) = session_state_builder.file_formats() {
+        file_formats.push(Arc::new(factory));
+    }
+
+    SessionContext::new_with_state(session_state_builder.build())
+}
+
+#[tokio::main]
+pub async fn main() -> anyhow::Result<()> {
+    let ctx = get_session_context();
+
+    ctx.sql(
+        r#"
+    CREATE EXTERNAL TABLE hits
+    STORED AS VORTEX
+    LOCATION '/Volumes/Code/vortex/bench-vortex/data/clickbench_partitioned/vortex-file-compressed/'
+    "#,
+    )
+    .await?;
+
+    let start = std::time::Instant::now();
+
+    let mut stream = ctx
+        .sql("select * from hits")
+        .await?
+        .execute_stream()
+        .await?;
+
+    while let Some(batch) = stream.next().await.transpose()? {
+        // Discard the batches
+        drop(batch);
+    }
+
+    let elapsed = start.elapsed();
+    println!("scanned 14GB of clickbench data in {elapsed:?}");
+
+    Ok(())
+}

vortex-datafusion/examples/read_clickbench.rs

@@ -8,7 +8,7 @@ use std::sync::{Arc, OnceLock};
 use futures::future::BoxFuture;
 use futures::{FutureExt, TryFutureExt, try_join};
 use vortex_array::compute::{MinMaxResult, min_max, take};
-use vortex_array::{ArrayRef, MaskFuture};
+use vortex_array::{Array, ArrayRef, IntoArray, MaskFuture};
 use vortex_dict::DictArray;
 use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexError, VortexExpect, VortexResult};
@@ -76,6 +76,7 @@ impl DictReader {
                     )
                     .vortex_expect("must construct dict values array evaluation")
                     .map_err(Arc::new)
+                    .map_ok(|arr| arr.to_canonical().into_array())
                     .boxed()
                     .shared()
             })