feat[ffi]: add a session used for caching (#3358)

For now this is a footer cache but we can cache more if we want too

Signed-off-by: Joe Isaacs <[email protected]>

---------

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

Cargo.lock

@@ -6,6 +6,7 @@
 
 #include "vortex.hpp"
 #include "vortex_error.hpp"
+#include "vortex_session.hpp"
 
 namespace vortex {
 
@@ -55,8 +56,8 @@ struct FileReader {
 		vx_file_reader_free(file);
 	}
 
-	static duckdb::unique_ptr<FileReader> Open(const vx_file_open_options *options) {
-		auto file = Try([&](auto err) { return vx_file_open_reader(options, err); });
+	static duckdb::unique_ptr<FileReader> Open(const vx_file_open_options *options, VortexSession &session) {
+		auto file = Try([&](auto err) { return vx_file_open_reader(options, session.session, err); });
 		return duckdb::make_uniq<FileReader>(file);
 	}
 

duckdb-vortex/src/include/vortex_common.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "vortex.hpp"
+#include "duckdb/storage/object_cache.hpp"
+
+namespace vortex {
+
+class VortexSession : public duckdb::ObjectCacheEntry {
+public:
+	VortexSession() : session(vx_session_create()) {
+	}
+
+	~VortexSession() override {
+		vx_session_free(session);
+	}
+
+	vx_session *session;
+
+	static std::string ObjectType() {
+		return "vortex_session_cache_metadata";
+	}
+
+	std::string GetObjectType() override {
+		return ObjectType();
+	}
+};
+
+} // namespace vortex

duckdb-vortex/src/include/vortex_session.hpp

@@ -13,6 +13,7 @@
 #include "duckdb/function/table_function.hpp"
 #include "duckdb/main/extension_util.hpp"
 #include "duckdb/common/file_system.hpp"
+#include "duckdb/storage/object_cache.hpp"
 
 #include "concurrentqueue.h"
 
@@ -22,6 +23,7 @@
 #include "vortex_scan.hpp"
 #include "vortex_common.hpp"
 #include "vortex_expr.hpp"
+#include "vortex_session.hpp"
 
 using namespace duckdb;
 
@@ -31,6 +33,9 @@ namespace vortex {
 /// file and its schema. This data is populated during the bind phase, which
 /// happens during the query planning phase.
 struct BindData : public TableFunctionData {
+	// Session used to caching
+	shared_ptr<VortexSession> session;
+
 	shared_ptr<MultiFileList> file_list;
 	vector<LogicalType> columns_types;
 	vector<string> column_names;
@@ -182,12 +187,12 @@ std::string EnsureFileProtocol(FileSystem &fs, const std::string &path) {
 	return prefix + absolute_path;
 }
 
-static unique_ptr<FileReader> OpenFile(const std::string &filename, vector<LogicalType> &column_types,
-                                       vector<string> &column_names) {
+static unique_ptr<FileReader> OpenFile(const std::string &filename, VortexSession &session,
+                                       vector<LogicalType> &column_types, vector<string> &column_names) {
 	vx_file_open_options options {
 	    .uri = filename.c_str(), .property_keys = nullptr, .property_vals = nullptr, .property_len = 0};
 
-	auto file = FileReader::Open(&options);
+	auto file = FileReader::Open(&options, session);
 	if (!file) {
 		throw IOException("Failed to open Vortex file: " + filename);
 	}
@@ -224,15 +229,15 @@ static void VerifyNewFile(const BindData &bind_data, vector<LogicalType> &column
 	}
 }
 
-static unique_ptr<FileReader> OpenFileAndVerify(FileSystem &fs, const std::string &filename,
+static unique_ptr<FileReader> OpenFileAndVerify(FileSystem &fs, VortexSession &session, const std::string &filename,
                                                 const BindData &bind_data) {
 	auto new_column_names = vector<string>();
 	new_column_names.reserve(bind_data.column_names.size());
 
 	auto new_column_types = vector<LogicalType>();
 	new_column_types.reserve(bind_data.columns_types.size());
 
-	auto file = OpenFile(EnsureFileProtocol(fs, filename), new_column_types, new_column_names);
+	auto file = OpenFile(EnsureFileProtocol(fs, filename), session, new_column_types, new_column_names);
 	VerifyNewFile(bind_data, new_column_types, new_column_names);
 	return file;
 }
@@ -383,7 +388,8 @@ static void VortexScanFunction(ClientContext &context, TableFunctionInput &data,
 			if (auto file_idx = global_state.next_file_idx.fetch_add(1);
 			    file_idx < global_state.expanded_files.size()) {
 				auto file_name = global_state.expanded_files[file_idx];
-				auto vortex_file = OpenFileAndVerify(FileSystem::GetFileSystem(context), file_name, bind_data);
+				auto vortex_file =
+				    OpenFileAndVerify(FileSystem::GetFileSystem(context), *bind_data.session, file_name, bind_data);
 				global_state.layout_readers[file_idx] = LayoutReader::CreateFromFile(vortex_file.get());
 				CreateScanPartitions(context, bind_data, global_state, local_state, file_idx, vortex_file);
 			}
@@ -412,12 +418,21 @@ static unique_ptr<FunctionData> VortexBind(ClientContext &context, TableFunction
 	auto result = make_uniq<BindData>();
 	result->arena = make_uniq<google::protobuf::Arena>();
 
+	const static string VortexExtensionKey = std::string("vortex_extension:vortex_session");
+	auto session = ObjectCache::GetObjectCache(context).Get<VortexSession>(VortexExtensionKey);
+	if (session == nullptr) {
+		ObjectCache::GetObjectCache(context).Put(VortexExtensionKey, make_shared_ptr<VortexSession>());
+		session = ObjectCache::GetObjectCache(context).Get<VortexSession>(VortexExtensionKey);
+	}
+
+	result->session = session;
+
 	auto file_glob = duckdb::vector<string> {input.inputs[0].GetValue<string>()};
 	result->file_list = make_shared_ptr<GlobMultiFileList>(context, file_glob, FileGlobOptions::DISALLOW_EMPTY);
 
 	// Open the first file to extract the schema.
 	auto filename = EnsureFileProtocol(FileSystem::GetFileSystem(context), result->file_list->GetFirstFile());
-	result->initial_file = OpenFile(filename, column_types, column_names);
+	result->initial_file = OpenFile(filename, *result->session, column_types, column_names);
 
 	result->column_names = column_names;
 	result->columns_types = column_types;

duckdb-vortex/src/vortex_scan.cpp

@@ -22,6 +22,7 @@ futures = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 mimalloc = { workspace = true, optional = true }
+moka = { workspace = true, features = ["sync"] }
 object_store = { workspace = true, features = ["aws", "azure", "gcp"] }
 paste = { workspace = true }
 prost = { workspace = true }

vortex-ffi/Cargo.toml

@@ -124,6 +124,12 @@ typedef struct vx_file_reader vx_file_reader;
  */
 typedef struct vx_layout_reader vx_layout_reader;
 
+/**
+ * An object that stores registries and caches.
+ * This should if possible be reused between queries in ann interactive session.
+ */
+typedef struct vx_session vx_session;
+
 /**
  * Options supplied for opening a file.
  */
@@ -381,6 +387,7 @@ void vx_error_free(struct vx_error *error);
  * Open a file at the given path on the file system.
  */
 struct vx_file_reader *vx_file_open_reader(const struct vx_file_open_options *options,
+                                           struct vx_session *session,
                                            struct vx_error **error);
 
 void vx_file_write_array(const char *path, struct vx_array *ffi_array, struct vx_error **error);
@@ -430,6 +437,16 @@ void vx_file_reader_free(struct vx_file_reader *file);
  */
 void vx_init_logging(enum vx_log_level level);
 
+/**
+ * Create a session to be used for the lifetime of an interactive session.
+ */
+struct vx_session *vx_session_create(void);
+
+/**
+ * Free a session
+ */
+void vx_session_free(struct vx_session *session);
+
 /**
  * Opens a writable array stream, where sink is used to push values into the stream.
  * To close the stream close the sink with `vx_array_sink_close`.

vortex-ffi/cinclude/vortex.h

@@ -1,7 +1,5 @@
 //! FFI interface for Vortex File I/O.
 
-#![allow(non_camel_case_types)]
-
 use std::ffi::{CStr, c_char, c_int, c_uint, c_ulong};
 use std::str::FromStr;
 use std::sync::Arc;
@@ -26,14 +24,17 @@ use vortex::proto::expr::Expr;
 
 use crate::array::{vx_array, vx_array_iterator};
 use crate::error::{try_or, vx_error};
+use crate::session::{FileKey, vx_session};
 use crate::{RUNTIME, to_string, to_string_vec};
 
 /// A file reader that can be used to read from a file.
+#[allow(non_camel_case_types)]
 pub struct vx_file_reader {
     pub inner: VortexFile,
 }
 
 /// A Vortex layout reader.
+#[allow(non_camel_case_types)]
 pub struct vx_layout_reader {
     pub inner: Arc<dyn LayoutReader>,
 }
@@ -121,6 +122,7 @@ impl vx_file_scan_options {
 #[unsafe(no_mangle)]
 pub unsafe extern "C-unwind" fn vx_file_open_reader(
     options: *const vx_file_open_options,
+    session: *mut vx_session,
     error: *mut *mut vx_error,
 ) -> *mut vx_file_reader {
     try_or(error, ptr::null_mut(), || {
@@ -133,19 +135,39 @@ pub unsafe extern "C-unwind" fn vx_file_open_reader(
         if options.uri.is_null() {
             vortex_bail!("null uri")
         }
-        let uri = unsafe { CStr::from_ptr(options.uri) }.to_string_lossy();
-        let uri: Url = uri.parse().vortex_expect("File_open: parse uri");
+        let uri_str = unsafe { CStr::from_ptr(options.uri) }.to_string_lossy();
+        let uri: Url = uri_str.parse().vortex_expect("File_open: parse uri");
 
         let prop_keys = unsafe { to_string_vec(options.property_keys, options.property_len) };
         let prop_vals = unsafe { to_string_vec(options.property_vals, options.property_len) };
 
         let object_store = make_object_store(&uri, &prop_keys, &prop_vals)?;
 
-        let inner = RUNTIME.block_on(async move {
-            VortexOpenOptions::file()
-                .open_object_store(&object_store, uri.path())
-                .await
-        })?;
+        let file = VortexOpenOptions::file();
+        let (file, cache_hit) = if let Some(footer) = unsafe { session.as_ref() }.and_then(|s| {
+            s.inner.get_footer(&FileKey {
+                location: uri_str.to_string(),
+            })
+        }) {
+            (file.with_footer(footer), true)
+        } else {
+            (file, false)
+        };
+
+        let inner = RUNTIME
+            .block_on(async move { file.open_object_store(&object_store, uri.path()).await })?;
+
+        if !cache_hit {
+            let _ = unsafe { session.as_ref() }.is_some_and(|s| {
+                s.inner.put_footer(
+                    FileKey {
+                        location: uri_str.to_string(),
+                    },
+                    inner.footer().clone(),
+                );
+                true
+            });
+        }
 
         Ok(Box::into_raw(Box::new(vx_file_reader { inner })))
     })

vortex-ffi/src/file.rs

@@ -9,6 +9,7 @@ mod duckdb;
 mod error;
 mod file;
 mod log;
+mod session;
 mod sink;
 
 use std::ffi::{CStr, c_char, c_int};

vortex-ffi/src/lib.rs

@@ -0,0 +1,85 @@
+use std::sync::Arc;
+
+use moka::sync::Cache;
+use vortex::aliases::DefaultHashBuilder;
+use vortex::dtype::DType;
+use vortex::file::{Footer, SegmentSpec};
+use vortex::layout::segments::SegmentId;
+use vortex::scalar::ScalarValue;
+use vortex::stats::{Precision, Stat};
+
+/// An object that stores registries and caches.
+/// This should if possible be reused between queries in ann interactive session.
+#[allow(non_camel_case_types)]
+pub struct vx_session {
+    pub inner: Arc<VortexSession>,
+}
+
+/// Create a session to be used for the lifetime of an interactive session.
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn vx_session_create() -> *mut vx_session {
+    Box::into_raw(Box::new(vx_session {
+        inner: Arc::new(VortexSession::new()),
+    }))
+}
+
+/// Free a session
+#[unsafe(no_mangle)]
+pub unsafe extern "C-unwind" fn vx_session_free(session: *mut vx_session) {
+    drop(unsafe { Box::from_raw(session) })
+}
+
+pub struct VortexSession {
+    file_cache: Cache<FileKey, Footer, DefaultHashBuilder>,
+}
+
+/// Cache key for a [`VortexFile`].
+#[derive(Hash, Eq, PartialEq, Debug, Clone)]
+pub struct FileKey {
+    // TODO: support last modified ts.
+    pub location: String,
+}
+
+impl VortexSession {
+    pub fn new() -> Self {
+        let file_cache = Cache::builder()
+            .max_capacity(64u64 * (1 << 20))
+            .eviction_listener(|k: Arc<FileKey>, _v: Footer, cause| {
+                log::trace!("Removed {:?} due to {:?}", k, cause);
+            })
+            .weigher(|_k, footer| u32::try_from(estimate_layout_size(footer)).unwrap_or(u32::MAX))
+            .build_with_hasher(DefaultHashBuilder::default());
+
+        Self { file_cache }
+    }
+
+    pub fn get_footer(&self, file_key: &FileKey) -> Option<Footer> {
+        self.file_cache.get(file_key)
+    }
+
+    pub fn put_footer(&self, file_key: FileKey, footer: Footer) {
+        self.file_cache.insert(file_key, footer)
+    }
+}
+
+// TODO(joe): unify with the df impl
+/// Approximate the in-memory size of a layout
+fn estimate_layout_size(footer: &Footer) -> usize {
+    let segments_size = footer.segment_map().len() * size_of::<SegmentSpec>();
+    let stats_size = footer
+        .statistics()
+        .iter()
+        .map(|v| {
+            v.iter()
+                .map(|_| size_of::<Stat>() + size_of::<Precision<ScalarValue>>())
+                .sum::<usize>()
+        })
+        .sum::<usize>();
+
+    let root_layout = footer.layout();
+    let layout_size = size_of::<DType>()
+        + root_layout.metadata().len()
+        + root_layout.segment_ids().len() * size_of::<SegmentId>();
+
+    segments_size + stats_size + layout_size
+}