From 73c52c9a25284c41366cd4c97bb3aaf49a2d023b Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Sun, 21 Jun 2026 20:06:43 +0200 Subject: [PATCH] Slim data-lake: replace TableManager/FileManager/DataLake with native SessionContext Collapses beacon-data-lake's manager indirection onto DataFusion's native APIs, keeping only what DataFusion doesn't provide: persisting/restoring table definitions across restarts. - TableManager -> thin PersistentSchemaProvider that delegates the in-memory catalog to a native MemorySchemaProvider and only persists/removes tables:///table.json on register/deregister (deregister-then-register preserves overwrite-on-existing for MV refresh / Iceberg replace / alter). - Startup recovery -> free init_tables(ctx, schema) reusing loading/ordering. - FileManager -> free functions (create_listing_url, create_temp_output_file, list_datasets, list_dataset_schema); file_formats vec moves into Runtime. - DataLake struct -> register_object_stores(ctx, &ObjectStores) free function, fed by the runtime-owned ObjectStores (config-owned storage, not globals). - list_table_config reconstructs the TableDefinition from the live provider via definition_from_provider (no parallel definition registry). - Temp output files are created under the configured tmp dir; Output::parse takes the tmp_dir instead of a FileManager. Rebased onto current main (runtime-owned config / ObjectStores / per-format config). Workspace builds; beacon-data-lake (10), beacon-core (28) and beacon-api (9) tests pass. --- beacon-core/src/query/compiler.rs | 7 +- beacon-core/src/query/from.rs | 26 +- beacon-core/src/query/output.rs | 9 +- beacon-core/src/runtime.rs | 166 ++++++-- beacon-core/src/statement_plan/actions.rs | 2 +- beacon-data-lake/src/files/manager.rs | 147 -------- beacon-data-lake/src/files/mod.rs | 124 +++++- beacon-data-lake/src/lib.rs | 355 +++--------------- beacon-data-lake/src/table_runtime/mod.rs | 61 ++- .../src/table_runtime/ordering.rs | 39 ++ .../persistent_schema_provider.rs | 261 +++++++++++++ .../src/table_runtime/provider_factory.rs | 70 ---- .../src/table_runtime/schema_persistence.rs | 65 ++-- .../src/table_runtime/table_manager.rs | 268 ------------- 14 files changed, 704 insertions(+), 896 deletions(-) delete mode 100644 beacon-data-lake/src/files/manager.rs create mode 100644 beacon-data-lake/src/table_runtime/persistent_schema_provider.rs delete mode 100644 beacon-data-lake/src/table_runtime/provider_factory.rs delete mode 100644 beacon-data-lake/src/table_runtime/table_manager.rs diff --git a/beacon-core/src/query/compiler.rs b/beacon-core/src/query/compiler.rs index d824a9f1..3c9caf41 100644 --- a/beacon-core/src/query/compiler.rs +++ b/beacon-core/src/query/compiler.rs @@ -4,7 +4,6 @@ //! straight to DataFusion's SQL parser in `Runtime::plan_client_query`; only the //! JSON form is "compiled". -use beacon_data_lake::{FileManager, TableManager}; use datafusion::{logical_expr::LogicalPlan, prelude::SessionContext}; use crate::query::QueryBody; @@ -13,8 +12,6 @@ use crate::query::QueryBody; pub async fn compile_json_query( query_body: QueryBody, session: &SessionContext, - table_manager: &TableManager, - file_manager: &FileManager, ) -> anyhow::Result { // The runtime config is published as a SessionConfig extension; fall back to // defaults if absent (e.g. a bare session in a unit test). @@ -39,10 +36,10 @@ pub async fn compile_json_query( all_columns.extend(select_cols); } - from.init_builder(session, table_manager, file_manager, Some(&all_columns)) + from.init_builder(session, Some(&all_columns)) .await? } else { - from.init_builder(session, table_manager, file_manager, None) + from.init_builder(session, None) .await? }; diff --git a/beacon-core/src/query/from.rs b/beacon-core/src/query/from.rs index 476c7e25..6efaa2ff 100644 --- a/beacon-core/src/query/from.rs +++ b/beacon-core/src/query/from.rs @@ -4,7 +4,6 @@ use std::sync::Arc; -use beacon_data_lake::{FileManager, TableManager}; use beacon_datafusion_ext::file_collection::FileCollection; use beacon_arrow_odv::datafusion::OdvFormat; use beacon_arrow_csv::datafusion::CsvFormat; @@ -42,22 +41,18 @@ impl From { /// /// # Arguments /// * `session_context` - The DataFusion session context. - /// * `table_manager` - Table manager used to resolve table providers. - /// * `file_manager` - File manager used to resolve listing paths. /// /// # Returns /// * `LogicalPlanBuilder` for the specified source. pub async fn init_builder( &self, session_context: &SessionContext, - table_manager: &TableManager, - file_manager: &FileManager, projection: Option<&Vec>, ) -> datafusion::error::Result { match self { From::Table(name) => { - // Use a registered table. - if let Some(mut table) = table_manager.table_provider(name) { + // Use a registered table from the catalog. + if let Ok(mut table) = session_context.table_provider(name.as_str()).await { if let (Some(projection), Some(file_collection)) = (projection, table.as_any().downcast_ref::()) { @@ -75,9 +70,7 @@ impl From { } From::Format { format } => { // Use a file format as a table source. - let table_source = format - .as_table_source(session_context, file_manager) - .await?; + let table_source = format.as_table_source(session_context).await?; Ok(LogicalPlanBuilder::scan("tmp", table_source, None)?) } } @@ -121,17 +114,15 @@ impl FromFormat { /// /// # Arguments /// * `session_context` - The DataFusion session context. - /// * `file_manager` - File manager used for file path resolution. /// /// # Returns /// * `Arc` for the specified format. pub async fn as_table_source( &self, session_context: &SessionContext, - file_manager: &FileManager, ) -> datafusion::error::Result> { let file_format = self.file_format(session_context).await?; - let urls = self.listing_table_urls(file_manager)?; + let urls = self.listing_table_urls()?; // Create a FileCollection as the table provider. let table = @@ -167,11 +158,8 @@ impl FromFormat { } } - /// Resolves file paths to [`ListingTableUrl`]s using the data lake. - fn listing_table_urls( - &self, - file_manager: &FileManager, - ) -> datafusion::error::Result> { + /// Resolves file paths to [`ListingTableUrl`]s under the datasets store. + fn listing_table_urls(&self) -> datafusion::error::Result> { let paths = match self { FromFormat::Csv { paths, .. } | FromFormat::Parquet { paths } @@ -185,7 +173,7 @@ impl FromFormat { let mut listing_table_urls = Vec::with_capacity(paths.len()); for path in paths { - let url = file_manager.try_create_listing_url(path.to_string())?; + let url = beacon_data_lake::create_listing_url(path.to_string())?; listing_table_urls.push(url); } Ok(listing_table_urls) diff --git a/beacon-core/src/query/output.rs b/beacon-core/src/query/output.rs index 8d65c579..41593bc5 100644 --- a/beacon-core/src/query/output.rs +++ b/beacon-core/src/query/output.rs @@ -8,7 +8,6 @@ use std::sync::Arc; use beacon_arrow_netcdf::datafusion::{options::NetcdfOptions, NetCDFFormatFactory, NetcdfConfig}; use beacon_arrow_odv::datafusion::OdvFileFormatFactory; use beacon_arrow_odv::writer::OdvOptions; -use beacon_data_lake::FileManager; use beacon_arrow_csv::datafusion::CsvFormatFactory; use beacon_arrow_geoparquet::datafusion::{GeoParquetFormatFactory, GeoParquetOptions}; use beacon_arrow_ipc::datafusion::ArrowFormatFactory; @@ -33,8 +32,8 @@ impl Output { /// Parses the logical plan and prepares an output file in the specified format. /// /// # Arguments - /// * `_session_context` - DataFusion session context (unused). - /// * `file_manager` - FileManager instance for temporary file creation. + /// * `session_context` - DataFusion session context (provides the datasets store). + /// * `tmp_dir` - Directory the temporary output file is created in (the tmp store root). /// * `input_plan` - The logical plan to export. /// /// # Returns @@ -42,7 +41,7 @@ impl Output { pub async fn parse( &self, session_context: &SessionContext, - file_manager: &FileManager, + tmp_dir: &std::path::Path, input_plan: LogicalPlan, ) -> datafusion::error::Result<(LogicalPlan, QueryOutputFile)> { let datasets_store = session_context @@ -55,7 +54,7 @@ impl Output { ) })?; let file_type = self.format.file_type(datasets_store).await; - let temp_output = file_manager.try_create_temp_output_file(".tmp"); + let temp_output = beacon_data_lake::create_temp_output_file(tmp_dir, ".tmp"); let plan = LogicalPlanBuilder::copy_to( input_plan, temp_output.output_url(), diff --git a/beacon-core/src/runtime.rs b/beacon-core/src/runtime.rs index 247c3110..a8488bd0 100644 --- a/beacon-core/src/runtime.rs +++ b/beacon-core/src/runtime.rs @@ -6,9 +6,12 @@ use arrow::{ array::AsArray, datatypes::{SchemaRef, UInt64Type}, }; -use beacon_data_lake::{DataLake, FileManager, TableManager}; +use beacon_data_lake::{ + PersistentSchemaProvider, DATASETS_OBJECT_STORE_URL, TABLES_OBJECT_STORE_URL, +}; use beacon_datafusion_ext::{ - format_ext::DatasetMetadata, listing_table_factory_ext::ListingTableFactoryExt, + format_ext::{DatasetMetadata, FileFormatFactoryExt}, + listing_table_factory_ext::ListingTableFactoryExt, stats_cache::beacon_file_statistics_cache, }; use beacon_functions::function_doc::FunctionDoc; @@ -34,8 +37,7 @@ use crate::{ /// Beacon's single execution layer: startup, catalog access, queries, SQL, and files. pub struct Runtime { session_ctx: Arc, - table_manager: Arc, - file_manager: Arc, + file_formats: Vec>, listing_table_factory: Arc, query_metrics: Arc>>, /// The configuration this runtime was built with. Owned, not process-global. @@ -54,25 +56,25 @@ impl Runtime { config.runtime.vm_memory_size * 1024 * 1024, )); - let object_stores = - beacon_object_storage::ObjectStores::new(&config.storage).await?; + let object_stores = beacon_object_storage::ObjectStores::new(&config.storage).await?; let session_ctx = Self::init_ctx(memory_pool, object_stores.datasets.clone(), config.clone())?; - let data_lake = Arc::new( - DataLake::new(session_ctx.clone(), object_stores.clone(), config.clone()).await, - ); + beacon_data_lake::register_object_stores(&session_ctx, &object_stores)?; - let table_manager = data_lake.table_manager(); - let file_manager = data_lake.file_manager(); + let file_formats = beacon_data_lake::file_formats( + session_ctx.clone(), + object_stores.datasets.clone(), + &config, + )?; let mut table_functions = vec![]; table_functions.extend(beacon_functions::file_formats::register_table_functions( tokio::runtime::Handle::current(), session_ctx.clone(), - file_manager.data_object_store_url(), + DATASETS_OBJECT_STORE_URL.clone(), object_stores.datasets.clone(), - file_manager.file_formats().to_vec(), + file_formats.clone(), )); table_functions.extend(beacon_functions::metadata::register_metadata_functions( session_ctx.clone(), @@ -86,10 +88,15 @@ impl Runtime { ); } + let schema_provider = Arc::new(PersistentSchemaProvider::new( + tokio::runtime::Handle::current(), + session_ctx.clone(), + TABLES_OBJECT_STORE_URL.clone(), + )); session_ctx .catalog("beacon") .unwrap() - .register_schema("public", table_manager.clone())?; + .register_schema("public", schema_provider.clone())?; // Build the shared Iceberg catalog before discovering tables: startup // table discovery rebuilds Iceberg providers via the catalog, and the @@ -116,18 +123,17 @@ impl Runtime { session_ctx.register_udf(udf); } - table_manager.init_tables().await?; + beacon_data_lake::init_tables(&session_ctx, &schema_provider).await?; let listing_table_factory = Arc::new(ListingTableFactoryExt::new( - file_manager.data_object_store_url(), + DATASETS_OBJECT_STORE_URL.clone(), Arc::downgrade(&session_ctx), )); Ok(Self { session_ctx, - table_manager, + file_formats, listing_table_factory, - file_manager, query_metrics: Arc::new(Mutex::new(HashMap::new())), config, }) @@ -261,7 +267,7 @@ impl Runtime { // `Output::parse` wraps the (already validated) plan in a `COPY TO` the // temp file; this COPY is beacon-generated, so it is not re-validated. let (copy_plan, output_file) = output - .parse(self.session_ctx.as_ref(), self.file_manager.as_ref(), plan) + .parse(self.session_ctx.as_ref(), &self.config.storage.tmp_dir, plan) .await?; let output_file = QueryOutputFile::from(output_file); @@ -296,13 +302,7 @@ impl Runtime { match inner { crate::query::InnerQuery::Sql(sql) => self.lower_sql(&sql).await, crate::query::InnerQuery::Json(body) => { - crate::query::compile_json_query( - body, - self.session_ctx.as_ref(), - self.table_manager.as_ref(), - self.file_manager.as_ref(), - ) - .await + crate::query::compile_json_query(body, self.session_ctx.as_ref()).await } } } @@ -401,7 +401,11 @@ impl Runtime { } pub fn list_tables(&self) -> Vec { - self.table_manager.table_names() + self.session_ctx + .catalog("beacon") + .and_then(|catalog| catalog.schema("public")) + .map(|schema| schema.table_names()) + .unwrap_or_default() } /// Lists SQL catalogs visible to Flight SQL and other SQL-based clients. @@ -475,15 +479,15 @@ impl Runtime { } pub async fn list_table_config(&self, table_name: String) -> Option { - self.table_manager.list_table(&table_name).and_then( - |config| match TableConfigView::try_from(config) { - Ok(config) => Some(config), - Err(error) => { - tracing::error!(?error, "failed to map table config into API contract"); - None - } - }, - ) + let provider = self.session_ctx.table_provider(table_name.as_str()).await.ok()?; + let config = beacon_data_lake::definition_from_provider(&table_name, provider.as_ref()).ok()?; + match TableConfigView::try_from(config) { + Ok(config) => Some(config), + Err(error) => { + tracing::error!(?error, "failed to map table config into API contract"); + None + } + } } pub async fn list_table_schema(&self, table_name: String) -> Option { @@ -534,10 +538,10 @@ impl Runtime { offset: Option, limit: Option, ) -> anyhow::Result> { - Ok(self - .file_manager - .list_datasets(offset, limit, pattern) - .await?) + Ok( + beacon_data_lake::list_datasets(&self.session_ctx, &self.file_formats, offset, limit, pattern) + .await?, + ) } pub async fn total_datasets(&self) -> anyhow::Result { @@ -547,7 +551,7 @@ impl Runtime { } pub async fn list_dataset_schema(&self, file: String) -> anyhow::Result { - Ok(self.file_manager.list_dataset_schema(&file).await?) + Ok(beacon_data_lake::list_dataset_schema(&self.session_ctx, &file).await?) } pub async fn list_dataset_schema_view(&self, file: String) -> anyhow::Result { @@ -874,3 +878,83 @@ mod client_query_tests { ); } } + +#[cfg(test)] +mod restart_tests { + use super::Runtime; + use futures::TryStreamExt; + + async fn run_sql(runtime: &Runtime, sql: &str) { + runtime + .run_query(crate::query::Query::sql(sql.to_string()), true) + .await + .expect("sql should run") + .into_record_stream() + .expect("streamed result") + .try_collect::>() + .await + .expect("sql should drain"); + } + + async fn count_rows(runtime: &Runtime, sql: &str) -> usize { + runtime + .run_query(crate::query::Query::sql(sql.to_string()), true) + .await + .expect("sql should run") + .into_record_stream() + .expect("streamed result") + .try_collect::>() + .await + .expect("sql should drain") + .iter() + .map(|batch| batch.num_rows()) + .sum() + } + + /// Tables persisted by one runtime are rebuilt by a fresh runtime via + /// `init_tables`: the base table's data survives, and the dependent view is + /// rebuilt in dependency order so it resolves the base table registered ahead + /// of it. This exercises the persist-on-register + startup-recovery round trip + /// that replaces the old `TableManager` registry. + #[tokio::test(flavor = "multi_thread")] + async fn persisted_tables_survive_a_restart() { + let suffix = uuid::Uuid::new_v4().simple(); + let base = format!("restart_base_{suffix}"); + let view = format!("restart_view_{suffix}"); + + // Both runtimes share one config so they resolve the same on-disk tables + // store; the restart must rebuild from what the first runtime persisted. + let config = std::sync::Arc::new(beacon_config::Config::load().unwrap()); + + // First runtime: create a base table with data and a view over it. + let runtime = Runtime::new(config.clone()).await.expect("runtime should start"); + run_sql(&runtime, &format!("CREATE TABLE {base} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {base} VALUES (1), (2)")).await; + run_sql(&runtime, &format!("CREATE VIEW {view} AS SELECT a FROM {base}")).await; + drop(runtime); + + // A fresh runtime rebuilds the catalog purely from the persisted + // `tables:///table.json` definitions. + let restarted = Runtime::new(config).await.expect("runtime should restart"); + + assert_eq!( + count_rows(&restarted, &format!("SELECT * FROM {base}")).await, + 2, + "base table data should persist across a restart" + ); + assert_eq!( + count_rows(&restarted, &format!("SELECT * FROM {view}")).await, + 2, + "the view should be rebuilt and resolve its dependency after a restart" + ); + + // Cleanup so the shared on-disk tables store does not leak into other + // tests (best-effort; `DROP TABLE` deregisters either provider type). + let _ = restarted + .run_query(crate::query::Query::sql(format!("DROP TABLE {view}")), true) + .await; + let _ = restarted + .run_query(crate::query::Query::sql(format!("DROP TABLE {base}")), true) + .await; + } +} diff --git a/beacon-core/src/statement_plan/actions.rs b/beacon-core/src/statement_plan/actions.rs index 4f212a1c..42c6dd1f 100644 --- a/beacon-core/src/statement_plan/actions.rs +++ b/beacon-core/src/statement_plan/actions.rs @@ -204,7 +204,7 @@ pub(crate) async fn create_table( beacon_iceberg::create_iceberg_table(&catalog, &namespace, &table_name, &arrow_schema) .await?; - // Registration persists the table's `table.json` pointer via the TableManager. + // Registration persists the table's `table.json` pointer via the PersistentSchemaProvider. session.register_table(name.clone(), Arc::new(table))?; if is_ctas { diff --git a/beacon-data-lake/src/files/manager.rs b/beacon-data-lake/src/files/manager.rs deleted file mode 100644 index 3b9173db..00000000 --- a/beacon-data-lake/src/files/manager.rs +++ /dev/null @@ -1,147 +0,0 @@ -use std::{ - collections::HashMap, - path::{Path, PathBuf}, - sync::Arc, -}; - -use arrow::datatypes::SchemaRef; -use beacon_common::listing_url::parse_listing_table_url; -use beacon_datafusion_ext::file_collection::FileCollection; -use beacon_datafusion_ext::format_ext::{DatasetMetadata, FileFormatFactoryExt}; -use datafusion::{ - catalog::TableProvider, datasource::listing::ListingTableUrl, error::DataFusionError, - execution::object_store::ObjectStoreUrl, prelude::SessionContext, -}; -use futures::StreamExt; - -use crate::files::temp_output_file::TempOutputFile; - -pub struct FileManager { - session_context: Arc, - data_directory_store_url: ObjectStoreUrl, - file_formats: Vec>, - /// Directory temporary output files are created in. Must match the root of - /// the tmp object store (`tmp://`) so COPY-written bytes are visible when the - /// file is read back. See [`TempOutputFile::new`]. - tmp_dir: PathBuf, -} - -impl FileManager { - pub fn new( - session_context: Arc, - data_directory_store_url: ObjectStoreUrl, - file_formats: Vec>, - tmp_dir: PathBuf, - ) -> Self { - Self { - session_context, - data_directory_store_url, - file_formats, - tmp_dir, - } - } - - #[inline(always)] - pub fn try_create_listing_url( - &self, - path: String, - ) -> datafusion::error::Result { - parse_listing_table_url(&self.data_directory_store_url, &path) - } - - pub fn data_object_store_url(&self) -> ObjectStoreUrl { - self.data_directory_store_url.clone() - } - - pub fn file_formats(&self) -> &Vec> { - &self.file_formats - } - - pub fn try_create_temp_output_file(&self, extension: &str) -> TempOutputFile { - TempOutputFile::new(&self.tmp_dir, extension) - } - - pub async fn list_datasets( - &self, - offset: Option, - limit: Option, - pattern: Option, - ) -> datafusion::error::Result> { - let state = self.session_context.state(); - let object_store = self - .session_context - .runtime_env() - .object_store(self.data_directory_store_url.clone())?; - - let listing_url = - self.try_create_listing_url(pattern.unwrap_or_else(|| "*".to_string()))?; - - let mut objects = Vec::new(); - let mut entry_stream = listing_url - .list_all_files(&state, &object_store, "") - .await?; - - while let Some(entry) = entry_stream.next().await { - if let Ok(entry) = entry { - // Skip Beacon-internal storage (e.g. materialized view data) so it is - // not surfaced as a user dataset. - if entry.location.as_ref().starts_with("__beacon__/") { - continue; - } - objects.push(entry); - } - } - - let mut datasets = vec![]; - - for file_format in self.file_formats.iter() { - let format_datasets = file_format.discover_datasets(&objects)?; - datasets.extend(format_datasets); - } - - // Keep current pagination semantics to avoid behavior regressions. - let start = offset.unwrap_or(0); - let end = limit.map(|l| start + l).unwrap_or(datasets.len()); - let datasets = datasets.into_iter().skip(start).take(end - start).collect(); - - Ok(datasets) - } - - pub async fn list_dataset_schema( - &self, - file_pattern: &str, - ) -> datafusion::error::Result { - let session_state = self.session_context.state(); - let extension = if file_pattern.ends_with("zarr.json") { - "zarr.json".to_string() - } else if file_pattern.contains("/atlas.json") { - "atlas.json".to_string() - } else { - match Path::new(file_pattern).extension() { - Some(ext) => ext.to_string_lossy().to_string(), - None => { - return Err(DataFusionError::Plan(format!( - "No file extension found for {}. No file type information available.", - file_pattern - ))); - } - } - }; - - tracing::debug!("Interpreted file extension: {}", extension); - let listing_url = self.try_create_listing_url(file_pattern.to_string())?; - - let file_format_factory = session_state - .get_file_format_factory(&extension) - .ok_or_else(|| { - DataFusionError::Plan(format!("No file format reader found for {}", extension)) - })?; - let file_format = file_format_factory.create(&session_state, &HashMap::new())?; - tracing::debug!("Using file format: {:?}", file_format); - - let file_collection = - FileCollection::new(&session_state, file_format, vec![listing_url]).await?; - - Ok(file_collection.schema()) - } -} diff --git a/beacon-data-lake/src/files/mod.rs b/beacon-data-lake/src/files/mod.rs index 8bab53bd..237ac00b 100644 --- a/beacon-data-lake/src/files/mod.rs +++ b/beacon-data-lake/src/files/mod.rs @@ -1,2 +1,124 @@ -pub mod manager; +//! Dataset file helpers over the datasets object store. +//! +//! These are free functions over a [`SessionContext`]: resolving listing URLs, +//! creating temporary output files, and discovering datasets / inferring their +//! schemas. All paths are resolved relative to [`DATASETS_OBJECT_STORE_URL`]. + pub mod temp_output_file; + +use std::{collections::HashMap, path::Path, sync::Arc}; + +use arrow::datatypes::SchemaRef; +use beacon_common::listing_url::parse_listing_table_url; +use beacon_datafusion_ext::file_collection::FileCollection; +use beacon_datafusion_ext::format_ext::{DatasetMetadata, FileFormatFactoryExt}; +use datafusion::{ + catalog::TableProvider, datasource::listing::ListingTableUrl, error::DataFusionError, + prelude::SessionContext, +}; +use futures::StreamExt; + +use crate::DATASETS_OBJECT_STORE_URL; +use temp_output_file::TempOutputFile; + +/// Resolve a (possibly globbed) path to a [`ListingTableUrl`] under the datasets +/// object store. +#[inline] +pub fn create_listing_url(path: String) -> datafusion::error::Result { + parse_listing_table_url(&DATASETS_OBJECT_STORE_URL, &path) +} + +/// Create a temporary output file with the given extension, used to stage query +/// results before they are streamed back to the client. +/// +/// `tmp_dir` MUST be the directory the tmp object store +/// ([`crate::TMP_OBJECT_STORE_URL`]) is rooted at, so the COPY-written bytes are +/// visible when the file is read back. See [`TempOutputFile::new`]. +pub fn create_temp_output_file(tmp_dir: &Path, extension: &str) -> TempOutputFile { + TempOutputFile::new(tmp_dir, extension) +} + +/// Discover the datasets matching `pattern` (default `*`) under the datasets +/// object store, asking each registered file format which objects it owns. +pub async fn list_datasets( + session_ctx: &SessionContext, + file_formats: &[Arc], + offset: Option, + limit: Option, + pattern: Option, +) -> datafusion::error::Result> { + let state = session_ctx.state(); + let object_store = session_ctx + .runtime_env() + .object_store(&*DATASETS_OBJECT_STORE_URL)?; + + let listing_url = create_listing_url(pattern.unwrap_or_else(|| "*".to_string()))?; + + let mut objects = Vec::new(); + let mut entry_stream = listing_url.list_all_files(&state, &object_store, "").await?; + + while let Some(entry) = entry_stream.next().await { + if let Ok(entry) = entry { + // Skip Beacon-internal storage (e.g. materialized view data) so it is + // not surfaced as a user dataset. + if entry.location.as_ref().starts_with("__beacon__/") { + continue; + } + objects.push(entry); + } + } + + let mut datasets = vec![]; + + for file_format in file_formats.iter() { + let format_datasets = file_format.discover_datasets(&objects)?; + datasets.extend(format_datasets); + } + + // Keep current pagination semantics to avoid behavior regressions. + let start = offset.unwrap_or(0); + let end = limit.map(|l| start + l).unwrap_or(datasets.len()); + let datasets = datasets.into_iter().skip(start).take(end - start).collect(); + + Ok(datasets) +} + +/// Infer the Arrow schema of the dataset(s) matching `file_pattern` by resolving +/// the file format from the extension and reading the matching files. +pub async fn list_dataset_schema( + session_ctx: &SessionContext, + file_pattern: &str, +) -> datafusion::error::Result { + let session_state = session_ctx.state(); + let extension = if file_pattern.ends_with("zarr.json") { + "zarr.json".to_string() + } else if file_pattern.contains("/atlas.json") { + "atlas.json".to_string() + } else { + match Path::new(file_pattern).extension() { + Some(ext) => ext.to_string_lossy().to_string(), + None => { + return Err(DataFusionError::Plan(format!( + "No file extension found for {}. No file type information available.", + file_pattern + ))); + } + } + }; + + tracing::debug!("Interpreted file extension: {}", extension); + let listing_url = create_listing_url(file_pattern.to_string())?; + + let file_format_factory = session_state + .get_file_format_factory(&extension) + .ok_or_else(|| { + DataFusionError::Plan(format!("No file format reader found for {}", extension)) + })?; + let file_format = file_format_factory.create(&session_state, &HashMap::new())?; + tracing::debug!("Using file format: {:?}", file_format); + + let file_collection = + FileCollection::new(&session_state, file_format, vec![listing_url]).await?; + + Ok(file_collection.schema()) +} diff --git a/beacon-data-lake/src/lib.rs b/beacon-data-lake/src/lib.rs index 873b4a4d..5fa34285 100644 --- a/beacon-data-lake/src/lib.rs +++ b/beacon-data-lake/src/lib.rs @@ -1,62 +1,26 @@ -use std::{ - any::Any, - fmt::Debug, - sync::{Arc, LazyLock}, -}; +use std::sync::LazyLock; -use arrow::datatypes::SchemaRef; -use beacon_datafusion_ext::format_ext::DatasetMetadata; -use beacon_datafusion_ext::table_ext::TableDefinition; use beacon_object_storage::ObjectStores; -use datafusion::{ - catalog::{SchemaProvider, TableProvider}, - datasource::listing::ListingTableUrl, - error::DataFusionError, - execution::object_store::ObjectStoreUrl, - prelude::SessionContext, -}; +use datafusion::{execution::object_store::ObjectStoreUrl, prelude::SessionContext}; use url::Url; -use crate::files::temp_output_file::TempOutputFile; - -#[cfg(test)] -use object_store::path::PathPart; -#[cfg(test)] -use std::collections::HashMap; - pub mod file_formats; pub mod files; pub mod table; mod table_runtime; pub use file_formats::file_formats; - -pub use files::manager::FileManager; -pub use table_runtime::table_manager::TableManager; +pub use files::temp_output_file::TempOutputFile; +pub use files::{create_listing_url, create_temp_output_file, list_dataset_schema, list_datasets}; +pub use table_runtime::init_tables; +pub use table_runtime::persistent_schema_provider::PersistentSchemaProvider; +pub use table_runtime::schema_persistence::definition_from_provider; pub mod prelude { - pub use super::DataLake; - pub use super::FileManager; - pub use super::TableManager; pub use super::files::*; -} - -pub struct DataLake { - data_directory_store_url: ObjectStoreUrl, - table_directory_store_url: ObjectStoreUrl, - - table_manager: Arc, - file_manager: Arc, -} - -impl Debug for DataLake { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DataLake") - .field("data_directory_store_url", &self.data_directory_store_url) - .field("table_directory_store_url", &self.table_directory_store_url) - .field("table_count", &self.table_manager.table_names().len()) - .finish() - } + pub use super::{ + definition_from_provider, init_tables, register_object_stores, PersistentSchemaProvider, + }; } pub static DATASETS_OBJECT_STORE_URL: LazyLock = @@ -68,267 +32,40 @@ pub static TMP_OBJECT_STORE_URL: LazyLock = pub static INDEX_OBJECT_STORE_URL: LazyLock = LazyLock::new(|| ObjectStoreUrl::parse("index://").expect("Failed to parse index URL")); // ToDo: implement indexing on top of existing files utilizing the notified storage events. -impl DataLake { - #[cfg(test)] - fn table_directory_from_location( - location: &object_store::path::Path, - ) -> Option>> { - if location.filename() != Some("table.json") { - return None; - } - - let mut table_directory = location - .parts() - .map(|part| part.as_ref().to_string().into()) - .collect::>(); - table_directory.pop(); - - Some(table_directory) - } - - #[cfg(test)] - async fn order_tables( - tables: &HashMap>, - ) -> Vec> { - table_runtime::ordering::order_tables(tables).await - } - - #[inline(always)] - pub fn try_create_listing_url( - &self, - path: String, - ) -> datafusion::error::Result { - self.file_manager.try_create_listing_url(path) - } - - pub fn data_object_store_url(&self) -> ObjectStoreUrl { - self.file_manager.data_object_store_url() - } - - pub fn try_create_temp_output_file(&self, extension: &str) -> TempOutputFile { - self.file_manager.try_create_temp_output_file(extension) - } - - pub fn table_manager(&self) -> Arc { - self.table_manager.clone() - } - - pub fn file_manager(&self) -> Arc { - self.file_manager.clone() - } - - pub async fn new( - session_context: Arc, - object_stores: ObjectStores, - config: Arc, - ) -> Self { - // The runtime owns the object stores and passes them in; register each - // with the session context so DataFusion can resolve their URLs. - let datasets_object_store = object_stores.datasets.clone(); - let datasets_object_store_url = DATASETS_OBJECT_STORE_URL.clone(); - // Register the Beacon-internal store (rooted at the `__beacon__` prefix) - // used by materialized views to persist and read their data directly, - // bypassing the datasets store's user-facing hiding and metadata cache. - session_context.register_object_store( - &Url::parse(beacon_datafusion_ext::table_ext::INTERNAL_STORE_URL).unwrap(), - datasets_object_store.internal_store(), - ); - // Register datasets object store - session_context.register_object_store( - &Url::parse(datasets_object_store_url.as_str()).unwrap(), - datasets_object_store.clone(), - ); - // Register tables object store - let tables_object_store_url = TABLES_OBJECT_STORE_URL.clone(); - session_context.register_object_store( - &Url::parse(tables_object_store_url.as_str()).unwrap(), - object_stores.tables.clone(), - ); - // Register tmp object store - let tmp_object_store_url = TMP_OBJECT_STORE_URL.clone(); - session_context.register_object_store( - &Url::parse(tmp_object_store_url.as_str()).unwrap(), - object_stores.tmp.clone(), - ); - - let file_formats = - file_formats(session_context.clone(), datasets_object_store.clone(), &config).unwrap(); - let runtime_handle = tokio::runtime::Handle::current(); - - let table_manager = Arc::new(TableManager::new( - runtime_handle, - session_context.clone(), - datasets_object_store_url.clone(), - tables_object_store_url.clone(), - )); - let file_manager = Arc::new(FileManager::new( - session_context, - datasets_object_store_url.clone(), - file_formats, - config.storage.tmp_dir.clone(), - )); - - Self { - data_directory_store_url: datasets_object_store_url, - table_directory_store_url: tables_object_store_url, - table_manager, - file_manager, - } - } - - pub async fn init_tables(&self) -> anyhow::Result<()> { - self.table_manager.init_tables().await - } - - pub async fn list_datasets( - &self, - offset: Option, - limit: Option, - pattern: Option, - ) -> datafusion::error::Result> { - self.file_manager - .list_datasets(offset, limit, pattern) - .await - } - - pub async fn list_dataset_schema( - &self, - file_pattern: &str, - ) -> datafusion::error::Result { - self.file_manager.list_dataset_schema(file_pattern).await - } - - pub fn list_table_schema(&self, table_name: &str) -> Option { - self.table_manager.list_table_schema(table_name) - } - - pub fn list_table(&self, table_name: &str) -> Option> { - self.table_manager.list_table(table_name) - } -} - -#[async_trait::async_trait] -impl SchemaProvider for DataLake { - /// Returns true if table exist in the schema provider, false otherwise. - fn table_exist(&self, name: &str) -> bool { - self.table_manager.table_exist(name) - } - - /// Returns this `SchemaProvider` as [`Any`] so that it can be downcast to a - /// specific implementation. - fn as_any(&self) -> &dyn Any { - self - } - - /// Retrieves the list of available table names in this schema. - fn table_names(&self) -> Vec { - self.table_manager.table_names() - } - - /// Retrieves a specific table from the schema by name, if it exists, - /// otherwise returns `None`. - async fn table(&self, name: &str) -> Result>, DataFusionError> { - Ok(self.table_manager.table_provider(name)) - } - - fn register_table( - &self, - name: String, - table: Arc, - ) -> datafusion::error::Result>> { - self.table_manager.register_table(name, table) - } - - /// If supported by the implementation, removes the `name` table from this - /// schema and returns the previously registered [`TableProvider`], if any. - /// - /// If no `name` table exists, returns Ok(None). - #[allow(unused_variables)] - fn deregister_table( - &self, - name: &str, - ) -> datafusion::error::Result>> { - self.table_manager.deregister_table(name) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use beacon_datafusion_ext::table_ext::{ExternalTableDefinition, ViewTableDefinition}; - - #[test] - fn table_directory_from_location_extracts_parent_directory() { - let location = object_store::path::Path::from("folder/example/table.json"); - - let table_directory = DataLake::table_directory_from_location(&location) - .expect("table.json path should produce a directory"); - let parts = table_directory - .iter() - .map(|part| part.as_ref()) - .collect::>(); - - assert_eq!(parts, vec!["folder", "example"]); - } - - #[test] - fn table_directory_from_location_ignores_non_table_config_files() { - let location = object_store::path::Path::from("folder/example/not-a-table.json"); - - assert!(DataLake::table_directory_from_location(&location).is_none()); - } - - #[tokio::test] - async fn ordered_definition_views_follow_table_scan_dependencies() { - let mut tables: HashMap> = HashMap::new(); - - let base = ExternalTableDefinition { - name: "base_table".to_string(), - location: "dataset/base_table/*.parquet".to_string(), - file_type: "parquet".to_string(), - schema: Arc::new(datafusion::arrow::datatypes::Schema::empty()), - definition: None, - partition_cols: vec![], - options: HashMap::new(), - if_not_exists: false, - }; - - let view_a = ViewTableDefinition { - name: "view_a".to_string(), - definition: "SELECT * FROM base_table".to_string(), - dependencies: vec!["base_table".to_string()], - }; - - let view_b = ViewTableDefinition { - name: "view_b".to_string(), - definition: "SELECT * FROM view_a".to_string(), - dependencies: vec!["view_a".to_string()], - }; - - tables.insert(base.name.clone(), Arc::new(base)); - tables.insert(view_a.name.clone(), Arc::new(view_a)); - tables.insert(view_b.name.clone(), Arc::new(view_b)); - - let order = DataLake::order_tables(&tables).await; - let ordered_names = order - .iter() - .map(|table| table.table_name()) - .collect::>(); - - let base_pos = ordered_names - .iter() - .position(|name| *name == "base_table") - .expect("base table should be present"); - let view_a_pos = ordered_names - .iter() - .position(|name| *name == "view_a") - .expect("view_a should be present"); - let view_b_pos = ordered_names - .iter() - .position(|name| *name == "view_b") - .expect("view_b should be present"); - - assert!(base_pos < view_a_pos); - assert!(view_a_pos < view_b_pos); - } +/// Register beacon's custom object stores on the session context. +/// +/// The runtime owns the [`ObjectStores`] (built from the storage config, not a +/// process-global) and passes them in; this registers the `datasets://`, +/// `internal://` (materialized-view data), `tables://` (table definitions) and +/// `tmp://` (query output) URLs so DataFusion can resolve them. Must be called +/// before any table or dataset access. +pub fn register_object_stores( + session_context: &SessionContext, + object_stores: &ObjectStores, +) -> anyhow::Result<()> { + let datasets_object_store = object_stores.datasets.clone(); + // Register the Beacon-internal store (rooted at the `__beacon__` prefix) + // used by materialized views to persist and read their data directly, + // bypassing the datasets store's user-facing hiding and metadata cache. + session_context.register_object_store( + &Url::parse(beacon_datafusion_ext::table_ext::INTERNAL_STORE_URL).unwrap(), + datasets_object_store.internal_store(), + ); + // Register datasets object store + session_context.register_object_store( + &Url::parse(DATASETS_OBJECT_STORE_URL.as_str()).unwrap(), + datasets_object_store, + ); + // Register tables object store + session_context.register_object_store( + &Url::parse(TABLES_OBJECT_STORE_URL.as_str()).unwrap(), + object_stores.tables.clone(), + ); + // Register tmp object store + session_context.register_object_store( + &Url::parse(TMP_OBJECT_STORE_URL.as_str()).unwrap(), + object_stores.tmp.clone(), + ); + + Ok(()) } diff --git a/beacon-data-lake/src/table_runtime/mod.rs b/beacon-data-lake/src/table_runtime/mod.rs index 934c197d..47e64225 100644 --- a/beacon-data-lake/src/table_runtime/mod.rs +++ b/beacon-data-lake/src/table_runtime/mod.rs @@ -1,5 +1,62 @@ +use std::{collections::HashMap, sync::Arc}; + +use beacon_datafusion_ext::table_ext::TableDefinition; +use datafusion::prelude::SessionContext; + +use crate::{DATASETS_OBJECT_STORE_URL, TABLES_OBJECT_STORE_URL}; + pub mod loading; pub mod ordering; -pub mod provider_factory; +pub mod persistent_schema_provider; pub mod schema_persistence; -pub mod table_manager; + +pub use persistent_schema_provider::PersistentSchemaProvider; + +/// Rebuild the catalog from the persisted table definitions. +/// +/// Loads every `tables:///table.json`, orders the definitions so a table +/// is registered after the tables it depends on (base tables, then temporary +/// definitions, then views in dependency order), builds each provider against +/// the live session — so a view's defining query resolves the tables already +/// registered ahead of it — and inserts it into `schema` without re-persisting. +/// Finishes by ensuring the empty `default` table exists. +pub async fn init_tables( + session_ctx: &Arc, + schema: &PersistentSchemaProvider, +) -> anyhow::Result<()> { + tracing::info!("Initializing tables from object store"); + let tables_object_store = session_ctx + .runtime_env() + .object_store(&*TABLES_OBJECT_STORE_URL) + .map_err(|error| anyhow::anyhow!("Failed to get tables object store: {}", error))?; + + let discovered = loading::load_tables_from_object_store(tables_object_store).await; + let table_map = discovered + .into_iter() + .map(|table| (table.table_name().to_string(), table)) + .collect::>>(); + let ordered = ordering::order_tables(&table_map).await; + + for definition in ordered { + let table_name = definition.table_name().to_string(); + match definition + .build_provider(session_ctx.clone(), &DATASETS_OBJECT_STORE_URL) + .await + { + Ok(provider) => { + let _ = schema.insert_loaded(table_name.clone(), provider); + tracing::info!("Registered table '{}'", table_name); + } + Err(error) => { + tracing::error!( + "Failed to build provider for table '{}': {}. Skipping registration of this table.", + table_name, + error + ); + } + } + } + + schema.ensure_default_table(); + Ok(()) +} diff --git a/beacon-data-lake/src/table_runtime/ordering.rs b/beacon-data-lake/src/table_runtime/ordering.rs index 61ead83b..897f24f5 100644 --- a/beacon-data-lake/src/table_runtime/ordering.rs +++ b/beacon-data-lake/src/table_runtime/ordering.rs @@ -298,6 +298,45 @@ mod tests { assert!(position(&order, "view_parent") < position(&order, "view_child")); } + #[tokio::test] + async fn ordered_definition_views_follow_table_scan_dependencies() { + use beacon_datafusion_ext::table_ext::{ExternalTableDefinition, ViewTableDefinition}; + + let mut tables: HashMap> = HashMap::new(); + + let base = ExternalTableDefinition { + name: "base_table".to_string(), + location: "dataset/base_table/*.parquet".to_string(), + file_type: "parquet".to_string(), + schema: Arc::new(Schema::empty()), + definition: None, + partition_cols: vec![], + options: HashMap::new(), + if_not_exists: false, + }; + + let view_a = ViewTableDefinition { + name: "view_a".to_string(), + definition: "SELECT * FROM base_table".to_string(), + dependencies: vec!["base_table".to_string()], + }; + + let view_b = ViewTableDefinition { + name: "view_b".to_string(), + definition: "SELECT * FROM view_a".to_string(), + dependencies: vec!["view_a".to_string()], + }; + + tables.insert(base.name.clone(), Arc::new(base)); + tables.insert(view_a.name.clone(), Arc::new(view_a)); + tables.insert(view_b.name.clone(), Arc::new(view_b)); + + let order = order_tables(&tables).await; + + assert!(position(&order, "base_table") < position(&order, "view_a")); + assert!(position(&order, "view_a") < position(&order, "view_b")); + } + #[tokio::test] async fn order_tables_uses_table_type_for_view_classification() { let mut tables: HashMap> = HashMap::new(); diff --git a/beacon-data-lake/src/table_runtime/persistent_schema_provider.rs b/beacon-data-lake/src/table_runtime/persistent_schema_provider.rs new file mode 100644 index 00000000..a05d5830 --- /dev/null +++ b/beacon-data-lake/src/table_runtime/persistent_schema_provider.rs @@ -0,0 +1,261 @@ +//! The `beacon.public` schema provider. +//! +//! This is a thin wrapper around DataFusion's native [`MemorySchemaProvider`]: +//! the in-memory catalog (register/lookup/deregister) is delegated to it +//! verbatim. The wrapper exists only to add one side effect — persisting and +//! removing each table's `tables:///table.json` definition as it is +//! registered or deregistered — so that the catalog survives restarts (it is +//! rebuilt at startup by [`crate::init_tables`]). + +use std::{any::Any, sync::Arc}; + +use arrow::datatypes::Schema; +use datafusion::{ + catalog::{MemorySchemaProvider, SchemaProvider, TableProvider}, + datasource::empty::EmptyTable, + error::DataFusionError, + execution::object_store::ObjectStoreUrl, + prelude::SessionContext, +}; + +use super::schema_persistence::SchemaPersistenceService; + +/// Schema provider for `beacon.public` that persists table definitions on +/// registration and removes them on deregistration. +pub struct PersistentSchemaProvider { + inner: Arc, + runtime_handle: tokio::runtime::Handle, + session_context: Arc, + table_directory_store_url: ObjectStoreUrl, +} + +impl std::fmt::Debug for PersistentSchemaProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PersistentSchemaProvider") + .field("table_directory_store_url", &self.table_directory_store_url) + .field("table_count", &self.inner.table_names().len()) + .finish() + } +} + +impl PersistentSchemaProvider { + pub fn new( + runtime_handle: tokio::runtime::Handle, + session_context: Arc, + table_directory_store_url: ObjectStoreUrl, + ) -> Self { + Self { + inner: Arc::new(MemorySchemaProvider::new()), + runtime_handle, + session_context, + table_directory_store_url, + } + } + + fn schema_persistence_service(&self) -> SchemaPersistenceService { + SchemaPersistenceService::new( + self.session_context.clone(), + self.table_directory_store_url.clone(), + ) + } + + /// Register a provider that was loaded from a persisted definition, without + /// re-persisting it. Used by [`crate::init_tables`] during startup recovery. + pub fn insert_loaded( + &self, + name: String, + table: Arc, + ) -> datafusion::error::Result>> { + self.inner.register_table(name, table) + } + + /// Register the in-memory `default` table backed by an empty provider. + /// + /// The default table is not persisted; it is recreated on every startup so + /// queries against the configured default table always resolve. + pub fn ensure_default_table(&self) { + if self.inner.table_exist("default") { + return; + } + let provider: Arc = Arc::new(EmptyTable::new(Arc::new(Schema::empty()))); + let _ = self.inner.register_table("default".to_string(), provider); + } +} + +#[async_trait::async_trait] +impl SchemaProvider for PersistentSchemaProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn table_names(&self) -> Vec { + self.inner.table_names() + } + + fn table_exist(&self, name: &str) -> bool { + self.inner.table_exist(name) + } + + async fn table(&self, name: &str) -> Result>, DataFusionError> { + self.inner.table(name).await + } + + fn register_table( + &self, + name: String, + table: Arc, + ) -> datafusion::error::Result>> { + let handle = self.runtime_handle.clone(); + let persistence = self.schema_persistence_service(); + let persist_name = name.clone(); + let persist_table = table.clone(); + tokio::task::block_in_place(|| { + handle.block_on(async move { + persistence + .persist_provider_definition(&persist_name, persist_table.as_ref()) + .await + }) + })?; + + // DataFusion's `MemorySchemaProvider` refuses to overwrite an existing + // entry, but beacon registers a fresh provider over an existing name to + // swap it (materialized-view refresh, Iceberg replace/alter). Drop any + // prior entry first so registration overwrites, returning the old one. + let previous = self.inner.deregister_table(&name)?; + self.inner.register_table(name, table)?; + Ok(previous) + } + + fn deregister_table( + &self, + name: &str, + ) -> datafusion::error::Result>> { + let handle = self.runtime_handle.clone(); + let persistence = self.schema_persistence_service(); + let remove_name = name.to_string(); + tokio::task::block_in_place(|| { + handle.block_on(async move { persistence.remove_persisted_table(&remove_name).await }) + })?; + + self.inner.deregister_table(name) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::datasource::ViewTable; + use futures::StreamExt; + use object_store::{memory::InMemory, path::Path, ObjectStore, ObjectStoreExt}; + use url::Url; + + // The persistence side effect runs the async store I/O via `block_in_place`, + // which requires a multi-threaded runtime. + fn fixture() -> (PersistentSchemaProvider, Arc, Arc) { + let session_context = Arc::new(SessionContext::new()); + let tables_store = Arc::new(InMemory::new()); + let tables_url = ObjectStoreUrl::parse("tables://").expect("tables url should parse"); + session_context.register_object_store( + &Url::parse(tables_url.as_str()).expect("tables url should be valid"), + tables_store.clone(), + ); + let provider = PersistentSchemaProvider::new( + tokio::runtime::Handle::current(), + session_context.clone(), + tables_url, + ); + (provider, session_context, tables_store) + } + + async fn view(session_context: &SessionContext, sql: &str) -> Arc { + let plan = session_context + .state() + .create_logical_plan(sql) + .await + .expect("logical plan should be created"); + Arc::new(ViewTable::new(plan, Some(sql.to_string()))) + } + + #[tokio::test(flavor = "multi_thread")] + async fn register_persists_definition_and_registers_in_catalog() { + let (provider, ctx, store) = fixture(); + + let previous = provider + .register_table("v".to_string(), view(&ctx, "SELECT 1 AS x").await) + .expect("registration should succeed"); + + assert!(previous.is_none(), "first registration has no previous table"); + assert!(provider.table_exist("v")); + assert!( + store.get(&Path::from("v/table.json")).await.is_ok(), + "the definition should be persisted to the tables store" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn register_overwrites_existing_name() { + // DataFusion's `MemorySchemaProvider` errors when a name already exists; + // the wrapper must overwrite instead so materialized-view refresh and + // Iceberg replace/alter can swap a fresh provider under the same name. + let (provider, ctx, _store) = fixture(); + provider + .register_table("v".to_string(), view(&ctx, "SELECT 1 AS x").await) + .expect("first registration should succeed"); + + let previous = provider + .register_table("v".to_string(), view(&ctx, "SELECT 2 AS y").await) + .expect("re-registering an existing name should overwrite, not error"); + + assert!(previous.is_some(), "overwrite returns the replaced provider"); + let table = provider + .table("v") + .await + .expect("lookup should succeed") + .expect("table should be present"); + assert_eq!( + table.schema().field(0).name(), + "y", + "the catalog should resolve the newly registered provider" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn deregister_removes_from_catalog_and_store() { + let (provider, ctx, store) = fixture(); + provider + .register_table("v".to_string(), view(&ctx, "SELECT 1 AS x").await) + .expect("registration should succeed"); + + let removed = provider + .deregister_table("v") + .expect("deregistration should succeed"); + + assert!(removed.is_some(), "deregister returns the removed provider"); + assert!(!provider.table_exist("v")); + let mut listing = store.list(Some(&Path::from("v"))); + assert!( + listing.next().await.is_none(), + "the persisted definition should be removed" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn ensure_default_table_is_idempotent() { + let (provider, _ctx, _store) = fixture(); + assert!(!provider.table_exist("default")); + + provider.ensure_default_table(); + provider.ensure_default_table(); + + assert!(provider.table_exist("default")); + assert_eq!( + provider + .table_names() + .iter() + .filter(|name| name.as_str() == "default") + .count(), + 1, + "the default table should be registered exactly once" + ); + } +} diff --git a/beacon-data-lake/src/table_runtime/provider_factory.rs b/beacon-data-lake/src/table_runtime/provider_factory.rs deleted file mode 100644 index 7ba99274..00000000 --- a/beacon-data-lake/src/table_runtime/provider_factory.rs +++ /dev/null @@ -1,70 +0,0 @@ -use std::sync::Arc; - -use beacon_datafusion_ext::table_ext::TableDefinition; -use datafusion::{ - catalog::TableProvider, execution::object_store::ObjectStoreUrl, prelude::SessionContext, -}; - -#[derive(Clone)] -pub struct TableProviderFactory { - session_context: Arc, - data_directory_store_url: ObjectStoreUrl, -} - -impl TableProviderFactory { - pub fn new( - session_context: Arc, - data_directory_store_url: ObjectStoreUrl, - ) -> Self { - Self { - session_context, - data_directory_store_url, - } - } - - pub async fn build( - &self, - definition: &Arc, - ) -> anyhow::Result<(String, Arc)> { - let provider = definition - .build_provider(self.session_context.clone(), &self.data_directory_store_url) - .await?; - Ok((definition.table_name().to_string(), provider)) - } -} - -#[cfg(test)] -mod tests { - use super::TableProviderFactory; - use beacon_datafusion_ext::table_ext::{TableDefinition, ViewTableDefinition}; - use datafusion::{ - datasource::{TableType as ProviderTableType, ViewTable}, - execution::object_store::ObjectStoreUrl, - prelude::SessionContext, - }; - use std::sync::Arc; - - #[tokio::test] - async fn build_definition_table_returns_view_provider() { - let factory = TableProviderFactory::new( - Arc::new(SessionContext::new()), - ObjectStoreUrl::parse("datasets://").expect("datasets url should parse"), - ); - - let definition = ViewTableDefinition { - name: "view_definition".to_string(), - definition: "SELECT 1 AS col".to_string(), - dependencies: vec![], - }; - let table: Arc = Arc::new(definition.clone()); - - let (name, provider) = factory - .build(&table) - .await - .expect("definition provider should build"); - - assert_eq!(name, definition.table_name()); - assert_eq!(provider.table_type(), ProviderTableType::View); - assert!(provider.as_any().downcast_ref::().is_some()); - } -} diff --git a/beacon-data-lake/src/table_runtime/schema_persistence.rs b/beacon-data-lake/src/table_runtime/schema_persistence.rs index 683126b0..5c5e8aa9 100644 --- a/beacon-data-lake/src/table_runtime/schema_persistence.rs +++ b/beacon-data-lake/src/table_runtime/schema_persistence.rs @@ -80,34 +80,7 @@ impl SchemaPersistenceService { table_name: &str, table: &dyn TableProvider, ) -> datafusion::error::Result { - let definition: Arc = - if let Some(table) = table.as_any().downcast_ref::() { - Arc::new(table.definition().clone()) - } else if let Some(table) = table.as_any().downcast_ref::() { - let definition = table.definition(); - Arc::new(definition.clone()) - } else if let Some(table) = table.as_any().downcast_ref::() { - let definition = table.definition(); - Arc::new(definition.clone()) - } else if let Some(definition) = - beacon_datafusion_ext::remote::remote_table_definition(table) - { - Arc::new(definition) - } else if let Some(table) = table.as_any().downcast_ref::() { - let definition = - ViewTableDefinition::try_from_view(table_name, table).map_err(|error| { - DataFusionError::Plan(format!( - "Failed to create ViewTableDefinition for table {}: {}", - table_name, error - )) - })?; - Arc::new(definition) - } else { - return Err(DataFusionError::Plan(format!( - "Unsupported table provider type for table {}", - table_name - ))); - }; + let definition = definition_from_provider(table_name, table)?; let json = serde_json::to_string_pretty(&definition).map_err(|error| { DataFusionError::Plan(format!( @@ -150,6 +123,42 @@ impl SchemaPersistenceService { } } +/// Reconstruct a serializable [`TableDefinition`] from a live table provider by +/// downcasting it to one of beacon's managed provider types. +/// +/// This is the inverse of building a provider from a definition: it lets the +/// catalog recover a table's persisted spec (used both to persist the table and +/// to surface its configuration) without keeping a parallel registry of +/// definitions alongside the providers. +pub fn definition_from_provider( + table_name: &str, + table: &dyn TableProvider, +) -> datafusion::error::Result> { + if let Some(table) = table.as_any().downcast_ref::() { + Ok(Arc::new(table.definition().clone())) + } else if let Some(table) = table.as_any().downcast_ref::() { + Ok(Arc::new(table.definition().clone())) + } else if let Some(table) = table.as_any().downcast_ref::() { + Ok(Arc::new(table.definition().clone())) + } else if let Some(definition) = beacon_datafusion_ext::remote::remote_table_definition(table) { + Ok(Arc::new(definition)) + } else if let Some(table) = table.as_any().downcast_ref::() { + let definition = + ViewTableDefinition::try_from_view(table_name, table).map_err(|error| { + DataFusionError::Plan(format!( + "Failed to create ViewTableDefinition for table {}: {}", + table_name, error + )) + })?; + Ok(Arc::new(definition)) + } else { + Err(DataFusionError::Plan(format!( + "Unsupported table provider type for table {}", + table_name + ))) + } +} + #[cfg(test)] mod tests { use super::SchemaPersistenceService; diff --git a/beacon-data-lake/src/table_runtime/table_manager.rs b/beacon-data-lake/src/table_runtime/table_manager.rs deleted file mode 100644 index a369e10e..00000000 --- a/beacon-data-lake/src/table_runtime/table_manager.rs +++ /dev/null @@ -1,268 +0,0 @@ -use std::{any::Any, collections::HashMap, sync::Arc}; - -use arrow::datatypes::{Schema, SchemaRef}; -use beacon_datafusion_ext::table_ext::TableDefinition; -use datafusion::{ - catalog::{SchemaProvider, TableProvider}, - datasource::empty::EmptyTable, - error::DataFusionError, - execution::object_store::ObjectStoreUrl, - prelude::SessionContext, -}; -use object_store::ObjectStore; - -use super::{ - loading, ordering, provider_factory::TableProviderFactory, - schema_persistence::SchemaPersistenceService, -}; - -#[derive(Clone)] -struct TableRegistryEntry { - table: Option>, - provider: Arc, -} - -pub struct TableManager { - runtime_handle: tokio::runtime::Handle, - data_directory_store_url: ObjectStoreUrl, - table_directory_store_url: ObjectStoreUrl, - session_context: Arc, - registry: parking_lot::Mutex>, -} - -impl std::fmt::Debug for TableManager { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TableManager") - .field("data_directory_store_url", &self.data_directory_store_url) - .field("table_directory_store_url", &self.table_directory_store_url) - .field("table_count", &self.registry.lock().len()) - .finish() - } -} - -#[async_trait::async_trait] -impl SchemaProvider for TableManager { - fn table_exist(&self, name: &str) -> bool { - self.table_exist(name) - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn table_names(&self) -> Vec { - self.table_names() - } - - async fn table(&self, name: &str) -> Result>, DataFusionError> { - Ok(self.table_provider(name)) - } - - fn register_table( - &self, - name: String, - table: Arc, - ) -> datafusion::error::Result>> { - TableManager::register_table(self, name, table) - } - - fn deregister_table( - &self, - name: &str, - ) -> datafusion::error::Result>> { - TableManager::deregister_table(self, name) - } -} - -impl TableManager { - pub fn new( - runtime_handle: tokio::runtime::Handle, - session_context: Arc, - data_directory_store_url: ObjectStoreUrl, - table_directory_store_url: ObjectStoreUrl, - ) -> Self { - Self { - runtime_handle, - data_directory_store_url, - table_directory_store_url, - session_context, - registry: parking_lot::Mutex::new(HashMap::new()), - } - } - - fn schema_persistence_service(&self) -> SchemaPersistenceService { - SchemaPersistenceService::new( - self.session_context.clone(), - self.table_directory_store_url.clone(), - ) - } - - async fn register_tables(&self, tables_to_register: Vec>) { - let provider_factory = TableProviderFactory::new( - self.session_context.clone(), - self.data_directory_store_url.clone(), - ); - - for table in tables_to_register { - let table_name = table.table_name().to_string(); - - match provider_factory.build(&table).await { - Ok((name, provider)) => { - self.registry.lock().insert( - name.clone(), - TableRegistryEntry { - table: Some(table), - provider, - }, - ); - tracing::info!("Registered table '{}'", name); - } - Err(error) => { - tracing::error!( - "Failed to build provider for table '{}': {}. Skipping registration of this table.", - table_name, - error - ); - } - } - } - } - - async fn init_tables_impl( - tables_object_store_url: ObjectStoreUrl, - session_context: Arc, - ) -> anyhow::Result>> { - tracing::info!("Initializing tables from object store"); - let tables_object_store = session_context - .runtime_env() - .object_store(&tables_object_store_url) - .map_err(|error| anyhow::anyhow!("Failed to get tables object store: {}", error))?; - - let discovered_tables = Self::load_tables_from_object_store(tables_object_store).await; - Ok(discovered_tables) - } - - async fn load_tables_from_object_store( - tables_object_store: Arc, - ) -> Vec> { - loading::load_tables_from_object_store(tables_object_store).await - } - - async fn order_tables( - tables: &HashMap>, - ) -> Vec> { - ordering::order_tables(tables).await - } - - /// Registers the in-memory `default` table backed by an empty provider. - /// - /// The default table is not persisted; it is recreated on every startup so - /// queries against the configured default table always resolve. - fn ensure_default_table(&self) { - let mut registry = self.registry.lock(); - if registry.contains_key("default") { - return; - } - - let provider: Arc = - Arc::new(EmptyTable::new(Arc::new(Schema::empty()))); - registry.insert( - "default".to_string(), - TableRegistryEntry { - table: None, - provider, - }, - ); - } - - pub async fn init_tables(&self) -> anyhow::Result<()> { - let table_formats = Self::init_tables_impl( - self.table_directory_store_url.clone(), - self.session_context.clone(), - ) - .await?; - - let table_map = table_formats - .into_iter() - .map(|table| (table.table_name().to_string(), table)) - .collect::>(); - let ordered_table_formats = Self::order_tables(&table_map).await; - - self.registry.lock().clear(); - self.register_tables(ordered_table_formats).await; - - self.ensure_default_table(); - Ok(()) - } - - pub fn table_exist(&self, name: &str) -> bool { - self.registry.lock().contains_key(name) - } - - pub fn table_names(&self) -> Vec { - self.registry.lock().keys().cloned().collect() - } - - pub fn table_provider(&self, name: &str) -> Option> { - self.registry - .lock() - .get(name) - .map(|entry| entry.provider.clone()) - } - - pub fn list_table_schema(&self, table_name: &str) -> Option { - self.registry - .lock() - .get(table_name) - .map(|entry| entry.provider.schema()) - } - - pub fn list_table(&self, table_name: &str) -> Option> { - self.registry - .lock() - .get(table_name) - .and_then(|entry| entry.table.clone()) - } - - pub fn register_table( - &self, - name: String, - table: Arc, - ) -> datafusion::error::Result>> { - let handle = self.runtime_handle.clone(); - let persistence = self.schema_persistence_service(); - let persist_name = name.clone(); - let persist_table = table.clone(); - let cloned_table = table.clone(); - tokio::task::block_in_place(|| { - handle.block_on(async move { - persistence - .persist_provider_definition(&persist_name, persist_table.as_ref()) - .await - }) - })?; - - self.registry.lock().insert( - name, - TableRegistryEntry { - table: None, - provider: table, - }, - ); - - Ok(Some(cloned_table)) - } - - pub fn deregister_table( - &self, - name: &str, - ) -> datafusion::error::Result>> { - let handle = self.runtime_handle.clone(); - let persistence = self.schema_persistence_service(); - tokio::task::block_in_place(|| { - handle.block_on(async move { persistence.remove_persisted_table(name).await }) - })?; - - Ok(self.registry.lock().remove(name).map(|entry| entry.provider)) - } -}