From 06e62309848c6b19f813408296f3fb42cb2dd79c Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Sat, 15 Aug 2026 23:16:43 +0200 Subject: [PATCH] Change a setting with SET, not only with a restart Beacon read its configuration from the environment only. An operator had to redeploy the server to change one value. A client could not read a value back. The engine settings and the format settings now use a `beacon.*` namespace. The query engine owns that namespace, so a `SET` applies to the next query: SET beacon.netcdf.use_rust_reader = true; RESET beacon.netcdf.use_rust_reader; SHOW SETTINGS; `ALTER SYSTEM SET` also writes the value into `beacon.db`. The server applies that value again at the next start. The order of precedence is: persisted value, then environment variable, then default. A `SET` applies to the whole server. Only the super-user can change a setting. This is the scope that `SET datafusion.*` already had. Beacon runs one query engine for every connection. `SHOW SETTINGS` is readable for every authenticated user. It lists each setting, its value, its start value and its description. A setting that the server reads once at start stays an environment variable. `SET beacon.port` fails and names `BEACON_PORT`. `beacon.` is also a complete alias for the engine's `datafusion.` namespace. DataFusion 53 gives no alias mechanism, so Beacon rewrites the statement before the engine plans it. Fixes #359 --- CHANGELOG.md | 25 + Cargo.lock | 1 + beacon-db/beacon-core/src/lib.rs | 1 + .../beacon-core/src/parser/beacon_parser.rs | 74 +- beacon-db/beacon-core/src/parser/statement.rs | 37 + beacon-db/beacon-core/src/runtime.rs | 4 + beacon-db/beacon-core/src/runtime_builder.rs | 159 ++++- beacon-db/beacon-core/src/settings.rs | 96 ++- .../beacon-core/src/settings_persistence.rs | 244 +++++++ .../beacon-core/src/statement_plan/actions.rs | 114 +++- .../beacon-core/src/statement_plan/logical.rs | 100 +++ .../beacon-core/src/statement_plan/lower.rs | 5 + .../beacon-core/src/statement_plan/mod.rs | 62 +- .../src/statement_plan/physical.rs | 110 ++- .../src/statement_plan/query_planner.rs | 12 + .../src/statement_plan/settings.rs | 449 ++++++++++++ .../src/statement_plan/stream_coalescer.rs | 34 +- .../beacon-core/src/system_schema/mod.rs | 8 +- .../beacon-core/src/system_schema/settings.rs | 77 +++ beacon-db/beacon-core/tests/redb_tables.rs | 10 +- .../beacon-core/tests/runtime_settings.rs | 361 ++++++++++ beacon-db/beacon-core/tests/system_schema.rs | 1 + beacon-db/beacon-datafusion-ext/src/lib.rs | 1 + beacon-db/beacon-datafusion-ext/src/nd/mod.rs | 71 +- .../beacon-datafusion-ext/src/nd/optimizer.rs | 25 +- .../beacon-datafusion-ext/src/settings.rs | 640 ++++++++++++++++++ .../beacon-arrow-atlas/src/datafusion/mod.rs | 34 +- .../beacon-arrow-bbf/src/datafusion/mod.rs | 14 +- .../beacon-arrow-hdf5/src/format.rs | 41 +- .../beacon-arrow-netcdf/src/datafusion/mod.rs | 42 +- .../beacon-arrow-zarr/src/datafusion/mod.rs | 34 +- .../beacon-lance/src/config.rs | 191 ++++++ .../beacon-lance/src/io.rs | 125 ++-- .../beacon-lance/src/lib.rs | 59 +- .../beacon-lance/src/provider.rs | 27 +- .../beacon-lance/src/sink.rs | 9 +- beacon-server/beacon-server-config/Cargo.toml | 1 + beacon-server/beacon-server-config/src/lib.rs | 22 + beacon-server/beacon-server/src/server/mod.rs | 3 + docs/docs/2.0.0-rc2/server/configuration.md | 150 +++- integration-tests/test_settings.py | 147 ++++ 41 files changed, 3391 insertions(+), 229 deletions(-) create mode 100644 beacon-db/beacon-core/src/settings_persistence.rs create mode 100644 beacon-db/beacon-core/src/statement_plan/settings.rs create mode 100644 beacon-db/beacon-core/src/system_schema/settings.rs create mode 100644 beacon-db/beacon-core/tests/runtime_settings.rs create mode 100644 beacon-db/beacon-datafusion-ext/src/settings.rs create mode 100644 beacon-db/beacon-file-formats/beacon-lance/src/config.rs create mode 100644 integration-tests/test_settings.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 47a6b3b3..5db358bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,25 @@ tag. Releases before 2.0.0 are recorded in the ### Added +- **A setting changes with `SET`, not only with a restart.** Beacon's configuration was + environment-only, so tuning one knob meant a redeploy, and no client could read a value back. + The engine and format settings now live in a `beacon.*` namespace the query engine owns: + `SET beacon.netcdf.use_rust_reader = true` applies to the next query, `RESET` puts it back to + what the server started with, and `SHOW SETTINGS` lists every setting with its value, its + startup value and what it does — readable by any authenticated user, which is the piece that + made the settings discoverable at all. `ALTER SYSTEM SET` also writes the value into + `beacon.db`, so it survives a restart; precedence is persisted > environment > default. A `SET` + applies to the whole server and is super-user-only, which is the scope `SET datafusion.*` + already had — Beacon runs one engine for every connection, and this does not invent a + per-client session. A setting that is read once at startup (the port, the data directory, the + credentials, each reader-cache *size*) stays an environment variable and says so by name: + ``SET beacon.port`` answers with `BEACON_PORT` and "restart", rather than appearing to work. + `beacon.` is also a complete alias for the engine's own `datafusion.` namespace, so + `SET beacon.execution.batch_size` and `SET datafusion.execution.batch_size` are one option + under two spellings. Format settings keep their three layers, narrowest last: the environment + default, the `SET`, then a per-table `OPTIONS (...)`. A `SET` reaches a `read_netcdf(...)` + immediately; a registered external table builds its reader once, so `REFRESH` rebuilds it with + the current settings. See [Configuration](docs/docs/2.0.0-rc2/server/configuration.md). - **Zarr stores supply column ranges for file pruning.** A Zarr store recorded nothing in `beacon.system.file_stats`, so every query opened every store. It now reports a range per coordinate: an array of rank 0 or rank 1 is read and measured, and an array of rank 2 or higher — @@ -79,6 +98,12 @@ tag. Releases before 2.0.0 are recorded in the ### Fixed +- **`BEACON_ATLAS_*` and `BEACON_ENABLE_BBF_SPLIT_STREAMS_SLICE` now reach the reader.** Both were + parsed, validated and documented, and then dropped: the runtime built the Atlas and BBF formats + from `Default::default()` instead of the configuration, so setting either to a non-default value + did nothing. The `BEACON_LANCE_*` family had the opposite problem — it bypassed the configuration + entirely and read the process environment at each write, so it was absent from the reference. + Both are wired through the runtime now, and all of them are documented and settable at runtime. - **File statistics pruned no netCDF or HDF5 file.** The ranges were recorded and then never used. Pruning rewrites the file list of a built scan, and it looked for that list on the plan's root node. A netCDF or HDF5 scan is not that node: its arrays reach the plan encoded, so the format diff --git a/Cargo.lock b/Cargo.lock index e5acb4a3..35916df7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2156,6 +2156,7 @@ dependencies = [ "beacon-arrow-netcdf", "beacon-arrow-zarr", "beacon-common", + "beacon-lance", "envconfig", "object_store 0.13.2", "thiserror 2.0.20", diff --git a/beacon-db/beacon-core/src/lib.rs b/beacon-db/beacon-core/src/lib.rs index 2b19c3af..465e711f 100644 --- a/beacon-db/beacon-core/src/lib.rs +++ b/beacon-db/beacon-core/src/lib.rs @@ -13,6 +13,7 @@ pub mod runtime_builder; pub mod schema_persistence; pub(crate) mod secret_persistence; pub mod settings; +pub(crate) mod settings_persistence; mod statement_plan; pub(crate) mod system_schema; diff --git a/beacon-db/beacon-core/src/parser/beacon_parser.rs b/beacon-db/beacon-core/src/parser/beacon_parser.rs index 4e6aeae8..6b9f2eae 100644 --- a/beacon-db/beacon-core/src/parser/beacon_parser.rs +++ b/beacon-db/beacon-core/src/parser/beacon_parser.rs @@ -10,7 +10,7 @@ use datafusion::sql::{ use beacon_auth::{Privilege, PrivilegeTarget}; use super::statement::{ - AttachStatement, AuthStatement, BeaconStatement, CreateCrawlerStatement, CreateIndexStatement, + AlterSystemStatement, AttachStatement, AuthStatement, BeaconStatement, CreateCrawlerStatement, CreateIndexStatement, CreateMaterializedViewStatement, CreateSecretStatement, DetachStatement, DropCrawlerStatement, DropExtensionStatement, DropIndexStatement, DropSecretStatement, RefreshStatement, AnalyzeFilesStatement, RunCrawlerStatement, SetExtensionStatement, ShowExtensionsStatement, ShowIndexesStatement, @@ -111,6 +111,14 @@ impl<'a> BeaconParser<'a> { return self.parse_analyze_files(); } + if self.is_alter_system() { + return self.parse_alter_system(); + } + + if self.is_show_settings() { + return self.parse_show_settings(); + } + let df_statement = Box::new(self.df_parser.parse_statement()?); Ok(BeaconStatement::DFStatement(df_statement)) @@ -313,6 +321,70 @@ impl<'a> BeaconParser<'a> { })) } + /// Whether the next two tokens are `ALTER SYSTEM`. + /// + /// Both words are required, so `ALTER TABLE` still reaches DataFusion. + fn is_alter_system(&self) -> bool { + let t1 = &self.df_parser.parser.peek_nth_token(0).token; + let t2 = &self.df_parser.parser.peek_nth_token(1).token; + matches!(t1, Token::Word(w) if w.keyword == Keyword::ALTER) + && matches!(t2, Token::Word(w) if w.value.to_uppercase() == "SYSTEM") + } + + /// Parse: ALTER SYSTEM SET = | ALTER SYSTEM RESET + fn parse_alter_system(&mut self) -> Result { + self.df_parser.parser.next_token(); // ALTER + self.df_parser.parser.next_token(); // SYSTEM + + let is_reset = matches!( + &self.df_parser.parser.peek_nth_token(0).token, + Token::Word(w) if w.value.to_uppercase() == "RESET" + ); + if is_reset { + self.df_parser.parser.next_token(); // RESET + let key = self.parse_object_name()?; + return Ok(BeaconStatement::AlterSystem(AlterSystemStatement { + key, + value: None, + })); + } + + self.expect_keyword(Keyword::SET)?; + let key = self.parse_object_name()?; + self.expect_token(&Token::Eq)?; + // Every value is carried as a string; the config field parses it into its + // own type, exactly as it does for a plain `SET`. + let value = self.parse_string_value()?; + Ok(BeaconStatement::AlterSystem(AlterSystemStatement { + key, + value: Some(value), + })) + } + + /// Whether the next two tokens are `SHOW SETTINGS`. + fn is_show_settings(&self) -> bool { + let t1 = &self.df_parser.parser.peek_nth_token(0).token; + let t2 = &self.df_parser.parser.peek_nth_token(1).token; + matches!(t1, Token::Word(w) if w.value.to_uppercase() == "SHOW") + && matches!(t2, Token::Word(w) if w.value.to_uppercase() == "SETTINGS") + } + + /// Parse: SHOW SETTINGS + fn parse_show_settings(&mut self) -> Result { + self.df_parser.parser.next_token(); // SHOW + self.df_parser.parser.next_token(); // SETTINGS + Ok(BeaconStatement::ShowSettings) + } + + /// Consume the expected token or error. + fn expect_token(&mut self, token: &Token) -> Result<()> { + self.df_parser + .parser + .expect_token(token) + .map(|_| ()) + .map_err(|e| DataFusionError::External(Box::new(e))) + } + /// Consume the expected keyword or error. fn expect_keyword(&mut self, keyword: Keyword) -> Result<()> { self.df_parser diff --git a/beacon-db/beacon-core/src/parser/statement.rs b/beacon-db/beacon-core/src/parser/statement.rs index 99f4504b..5fc63aaa 100644 --- a/beacon-db/beacon-core/src/parser/statement.rs +++ b/beacon-db/beacon-core/src/parser/statement.rs @@ -27,6 +27,41 @@ pub enum BeaconStatement { DropSecret(DropSecretStatement), ShowSecrets, Summarize(SummarizeStatement), + AlterSystem(AlterSystemStatement), + ShowSettings, +} + +/// `ALTER SYSTEM SET = ` / `ALTER SYSTEM RESET ` +/// +/// The persistent half of `SET`: it applies to the running server *and* is +/// written into the database file, so a restart replays it. A plain `SET` is +/// live-only. +/// +/// Parsed by beacon rather than reaching DataFusion, because sqlparser 0.61 has +/// no `ALTER SYSTEM` at all. +#[derive(Debug, Clone)] +pub struct AlterSystemStatement { + /// The setting name as typed. It is resolved (aliases, startup-only + /// rejection) when the plan is built, so the parser stays free of any + /// knowledge of the settings themselves. + pub key: ObjectName, + /// The new value, or `None` for `RESET` — which drops the persisted value and + /// restores the one the runtime booted with. + pub value: Option, +} + +impl Display for AlterSystemStatement { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.value { + Some(value) => write!( + f, + "ALTER SYSTEM SET {} = '{}'", + self.key, + escape_sql_literal(value) + ), + None => write!(f, "ALTER SYSTEM RESET {}", self.key), + } + } } /// SUMMARIZE | SUMMARIZE @@ -404,6 +439,8 @@ impl Display for BeaconStatement { } Self::ShowSecrets => write!(f, "SHOW SECRETS"), Self::Summarize(s) => write!(f, "SUMMARIZE {}", s.source), + Self::AlterSystem(s) => write!(f, "{s}"), + Self::ShowSettings => write!(f, "SHOW SETTINGS"), } } } diff --git a/beacon-db/beacon-core/src/runtime.rs b/beacon-db/beacon-core/src/runtime.rs index af68282f..e01f65b9 100644 --- a/beacon-db/beacon-core/src/runtime.rs +++ b/beacon-db/beacon-core/src/runtime.rs @@ -544,6 +544,10 @@ impl Runtime { BeaconStatement::AnalyzeFiles(statement) => { Ok(crate::statement_plan::analyze_files_plan(statement)) } + BeaconStatement::AlterSystem(statement) => { + crate::statement_plan::alter_system_plan(statement) + } + BeaconStatement::ShowSettings => Ok(crate::statement_plan::show_settings_plan()), BeaconStatement::SetExtension(statement) => { Ok(crate::statement_plan::set_extension_plan(statement)) } diff --git a/beacon-db/beacon-core/src/runtime_builder.rs b/beacon-db/beacon-core/src/runtime_builder.rs index f9f3cfb8..dfcdc964 100644 --- a/beacon-db/beacon-core/src/runtime_builder.rs +++ b/beacon-db/beacon-core/src/runtime_builder.rs @@ -6,8 +6,8 @@ use std::{ use crate::crawler::{new_crawler_manager_handle, CrawlerConfig, CrawlerManager}; use crate::schema_persistence::{init_tables, PersistentSchemaProvider}; -use beacon_arrow_atlas::datafusion::AtlasFormatFactory; -use beacon_arrow_bbf::datafusion::BBFFormatFactory; +use beacon_arrow_atlas::datafusion::{AtlasConfig, AtlasFormatFactory}; +use beacon_arrow_bbf::datafusion::{BBFFormatFactory, BbfConfig}; use beacon_arrow_csv::datafusion::CsvFormatFactory; use beacon_arrow_geoparquet::datafusion::GeoParquetFormatFactory; use beacon_arrow_hdf5::Hdf5Config; @@ -20,8 +20,11 @@ use beacon_arrow_zarr::ZarrConfig; use beacon_auth::{ AuthContext, BasicAuthProvider, InMemoryUserStore, RoleProvider, RoleStore, UserDirectory, }; +use beacon_lance::LanceConfig; + use beacon_datafusion_ext::{ consts::{DEFAULT_DB_STORE_URL_OBJECT_URL, TMP_STORE_URL_OBJECT_URL}, + settings::{BeaconOptions, BootSettings}, format_ext::{new_file_format_registry_handle, FileFormatFactoryExt, FileFormatRegistry}, listing_factory::{DefaultStore, ListingFactory, RootStore}, listing_table_factory_ext::ListingTableFactoryExt, @@ -97,6 +100,9 @@ pub struct RuntimeBuilder { pub netcdf: NetcdfConfig, pub hdf5: Hdf5Config, pub zarr: ZarrConfig, + pub atlas: AtlasConfig, + pub bbf: BbfConfig, + pub lance: LanceConfig, pub auth_provider: Option>, pub secrets_encryption_key: Option<[u8; 32]>, @@ -241,6 +247,24 @@ impl RuntimeBuilder { self } + /// Replaces the whole Atlas reader configuration. + pub fn with_atlas_config(mut self, atlas: AtlasConfig) -> Self { + self.atlas = atlas; + self + } + + /// Replaces the whole Beacon Binary Format configuration. + pub fn with_bbf_config(mut self, bbf: BbfConfig) -> Self { + self.bbf = bbf; + self + } + + /// Replaces the whole managed-Lance configuration. + pub fn with_lance_config(mut self, lance: LanceConfig) -> Self { + self.lance = lance; + self + } + pub fn with_auth_provider(mut self, provider: Arc) -> Self { self.auth_provider = Some(provider); self @@ -683,10 +707,64 @@ async fn register_schema_provider( .await?; load_persisted_secrets_into_store(session_ctx).await; + apply_persisted_settings(session_ctx).await; Ok(()) } +/// Replay the settings an `ALTER SYSTEM SET` persisted into the database file. +/// +/// Applied *after* the environment built the session config, which is what makes +/// the precedence **persisted > environment > default**: an operator's runtime +/// change outlives a restart rather than silently reverting to the deployment's +/// variables. +/// +/// A setting that no longer applies (a key a later version removed or renamed) is +/// logged and skipped. The alternative — failing to open the database — would +/// leave an operator with a server that cannot start and cannot be fixed through +/// SQL, since the bad value is the reason SQL is unreachable. +async fn apply_persisted_settings(session_ctx: &Arc) { + let persistence = + crate::settings_persistence::SettingsPersistence::from_session(session_ctx); + let Some(store) = persistence.store().cloned() else { + return; + }; + + let settings = match crate::settings_persistence::load_persisted_settings(&store).await { + Ok(settings) => settings, + Err(error) => { + tracing::error!("failed to load persisted settings: {error:#}"); + return; + } + }; + + if settings.is_empty() { + return; + } + + let state_ref = session_ctx.state_ref(); + let mut state = state_ref.write(); + let config = state.config_mut(); + + for (key, value) in settings { + match config.options_mut().set(&key, &value) { + Ok(()) => tracing::info!("applied persisted setting {key} = {value}"), + Err(error) => { + tracing::error!("skipping persisted setting {key} = {value}: {error}") + } + } + } + + // Re-take the snapshot, so `RESET ` restores what this server actually + // started with. The first snapshot (in `build_session_config`) records the + // environment only; a persisted value is applied over it and is therefore + // *also* part of the startup state. Without this, `RESET` on a persisted key + // would silently drop to the environment value until the next restart put the + // persisted one back. + let boot = BootSettings::from_config(config).with_startup(config.options()); + config.set_extension(Arc::new(boot)); +} + /// Load any secrets persisted (encrypted) in the database file into the in-memory secret store. /// /// Requires both a master key and a persistence store (a file-backed database); with either @@ -778,9 +856,9 @@ fn register_file_formats( Arc::new(ZarrFormatFactory::new(builder.zarr.clone())), Arc::new(AtlasFormatFactory::new( Default::default(), - Default::default(), + builder.atlas.clone(), )), - Arc::new(BBFFormatFactory::new(Default::default())), + Arc::new(BBFFormatFactory::new(builder.bbf.clone())), Arc::new(GeoParquetFormatFactory::default()), Arc::new(NetCDFFormatFactory::new( listing_factory.clone(), @@ -861,11 +939,14 @@ fn build_session_state( // `EnforceDistribution` — without which a Final aggregate never merges its // partitions and `count(*)` returns one row per file group. Filter runs // before projection so the selection is established first. - if builder.nd_pipeline { - state_builder = state_builder - .with_physical_optimizer_rule(Arc::new(NdFilterPushdown::new())) - .with_physical_optimizer_rule(Arc::new(NdProjectionPushdown::new())); - } + // + // Always installed, and each rule reads `beacon.enable_nd_pipeline` on every + // plan: the optimizer chain is fixed once the session state is built, so a + // rule left out here could never be switched on by a later `SET`. A disabled + // rule returns the plan untouched. + state_builder = state_builder + .with_physical_optimizer_rule(Arc::new(NdFilterPushdown::new())) + .with_physical_optimizer_rule(Arc::new(NdProjectionPushdown::new())); // Make every partition merge order-preserving so query results are // reproducible run to run. Appended last so it sees the CoalescePartitionsExec @@ -936,13 +1017,19 @@ fn build_session_config( // registry erases the `Ext` type — which the external-table builder needs // to hand a natively-read format its root store. .with_extension(new_file_format_registry_handle()) - // Recovered by the JSON query compiler (default table, projection pushdown). - .with_extension(Arc::new(builder.sql.clone())) - // Recovered when a statement's result stream is built, to merge the small - // batches a plan emits into client-sized ones. - .with_extension(Arc::new(CoalesceSqlStream::new( - builder.sql.stream_coalesce, + // Where `ALTER SYSTEM SET` writes. Empty for an in-memory database, which + // has nowhere durable to keep a setting — as with persistent secrets. + .with_extension(Arc::new(crate::settings_persistence::SettingsPersistence::new( + builder.db_path.is_some().then(|| db_store.clone()), ))); + // The `beacon.*` config namespace: the query compiler, the stream coalescer + // and every format reader recover their settings from here. Unlike the typed + // extensions above it is a `ConfigExtension`, which is what lets + // `SET beacon.x = y` rewrite it on the live session. + config + .options_mut() + .extensions + .insert(build_beacon_options(builder)); config .options_mut() .execution @@ -960,9 +1047,51 @@ fn build_session_config( config.options_mut().optimizer.expand_views_at_output = true; config.options_mut().sql_parser.map_string_types_to_utf8view = false; + // Last, so it records the finished startup state of *both* namespaces. This is + // what `RESET ` restores to: the value the operator's environment + // supplied, rather than DataFusion's compiled default. + let boot = BootSettings::capture(config.options()); + config = config.with_extension(Arc::new(boot)); + Ok(config) } +/// The startup state of the `beacon.*` namespace: every builder-supplied value, +/// which on the server is every `BEACON_*` environment variable. +/// +/// One flat namespace rather than a per-subsystem extension each, because +/// DataFusion keys a `ConfigExtension` by its prefix and `beacon` can only be +/// claimed once. The per-crate config structs stay the crates' own shape; this +/// copies into and out of them. +fn build_beacon_options(builder: &RuntimeBuilder) -> BeaconOptions { + let mut options = BeaconOptions::default(); + builder.sql.apply_to(&mut options); + options.enable_nd_pipeline = builder.nd_pipeline; + + options.netcdf.use_reader_cache = builder.netcdf.use_reader_cache; + options.netcdf.enable_statistics = builder.netcdf.enable_statistics; + options.netcdf.use_rust_reader = builder.netcdf.use_rust_reader; + + options.hdf5.use_reader_cache = builder.hdf5.use_reader_cache; + options.hdf5.enable_statistics = builder.hdf5.enable_statistics; + options.hdf5.use_rust_reader = builder.hdf5.use_rust_reader; + + options.zarr.enable_statistics = builder.zarr.enable_statistics; + + options.atlas.use_reader_cache = builder.atlas.use_reader_cache; + options.atlas.use_pruning = builder.atlas.use_pruning; + + options.bbf.split_streams_slice = builder.bbf.split_streams_slice; + + options.lance.compression = builder.lance.compression.clone(); + options.lance.numeric_compression = builder.lance.numeric_compression.clone(); + options.lance.version = builder.lance.version.clone(); + options.lance.minichunk = builder.lance.minichunk.clone(); + options.lance.materialization = builder.lance.materialization.clone(); + + options +} + /// Build the [`ListingFactory`] that resolves dataset paths for this runtime, /// directly from the builder's [`RuntimeBuilder::default_store`]: /// diff --git a/beacon-db/beacon-core/src/settings.rs b/beacon-db/beacon-core/src/settings.rs index 7ae13742..68077d0f 100644 --- a/beacon-db/beacon-core/src/settings.rs +++ b/beacon-db/beacon-core/src/settings.rs @@ -2,11 +2,20 @@ //! //! Settings are supplied to [`RuntimeBuilder`](crate::runtime_builder::RuntimeBuilder) //! by the embedder — never read from a process-global — and are published on the -//! session config as extensions, so plan- and execution-time code (where no -//! `Runtime` handle exists) can recover them from the `SessionContext`. +//! session config, so plan- and execution-time code (where no `Runtime` handle +//! exists) can recover them from the `SessionContext`. +//! +//! The types here are the embedder-facing shape of the builder's settings. What +//! is published on the session is +//! [`BeaconOptions`](beacon_datafusion_ext::settings::BeaconOptions), the +//! `beacon.*` config namespace — so the same values an embedder supplies are also +//! what `SET beacon.default_table = 'obs'` writes, and a reader here sees the +//! change without a restart. use std::time::Duration; +use beacon_datafusion_ext::settings::BeaconOptions; +use datafusion::execution::context::SessionConfig; use datafusion::prelude::SessionContext; /// SQL-facing settings: how client queries are compiled and how their results are @@ -36,14 +45,43 @@ impl Default for SqlSettings { impl SqlSettings { /// Recovers the settings published on `session_ctx`, falling back to the - /// defaults if the extension is absent (e.g. a bare session in a unit test). + /// defaults if the namespace is absent (e.g. a bare session in a unit test). + /// + /// Read per statement, never cached: a `SET` rewrites the namespace on the + /// shared session, and the next query has to see it. pub fn from_session(session_ctx: &SessionContext) -> Self { - session_ctx - .state() - .config() - .get_extension::() - .map(|settings| (*settings).clone()) - .unwrap_or_default() + Self::from_config(session_ctx.state().config()) + } + + /// [`Self::from_session`] against a `SessionConfig` directly, for callers + /// holding a `&dyn Session` rather than a `SessionContext`. + pub fn from_config(config: &SessionConfig) -> Self { + Self::from_options(&BeaconOptions::from_config(config)) + } + + /// The subset of the `beacon.*` namespace these settings mirror. + pub fn from_options(options: &BeaconOptions) -> Self { + Self { + default_table: options.default_table.clone(), + enable_pushdown_projection: options.enable_pushdown_projection, + stream_coalesce: SqlStreamCoalesceSettings { + enabled: options.sql.stream_coalesce.enabled, + target_rows: options.sql.stream_coalesce.target_rows, + flush_timeout_ms: options.sql.stream_coalesce.flush_timeout_ms, + max_rows: options.sql.stream_coalesce.max_rows, + }, + } + } + + /// Writes these settings into `options`, which is how the builder's values + /// become the namespace's startup state. + pub fn apply_to(&self, options: &mut BeaconOptions) { + options.default_table = self.default_table.clone(); + options.enable_pushdown_projection = self.enable_pushdown_projection; + options.sql.stream_coalesce.enabled = self.stream_coalesce.enabled; + options.sql.stream_coalesce.target_rows = self.stream_coalesce.target_rows; + options.sql.stream_coalesce.flush_timeout_ms = self.stream_coalesce.flush_timeout_ms; + options.sql.stream_coalesce.max_rows = self.stream_coalesce.max_rows; } } @@ -93,8 +131,6 @@ impl SqlStreamCoalesceSettings { mod tests { use super::*; - use std::sync::Arc; - use datafusion::execution::context::SessionConfig; /// `0` is the documented "no timed flushing" sentinel, not a zero-length @@ -116,8 +152,8 @@ mod tests { } /// The settings the runtime publishes on the session config must be the ones - /// plan-/execution-time code recovers — this is the only channel between the - /// two, and the extension is keyed by type, so a mismatch is silent. + /// plan-/execution-time code recovers — the `beacon.*` namespace is the only + /// channel between the two. #[test] fn from_session_recovers_the_published_settings() { let published = SqlSettings { @@ -130,18 +166,46 @@ mod tests { max_rows: 9, }, }; - let config = SessionConfig::new().with_extension(Arc::new(published.clone())); + let mut options = BeaconOptions::default(); + published.apply_to(&mut options); + let mut config = SessionConfig::new(); + config.options_mut().extensions.insert(options); let session_ctx = SessionContext::new_with_config(config); assert_eq!(SqlSettings::from_session(&session_ctx), published); } /// A bare session (a unit test, or any context built without the runtime) - /// carries no extension; recovery must fall back to the defaults rather than + /// carries no namespace; recovery must fall back to the defaults rather than /// fail, since every caller treats the settings as always available. #[test] - fn from_session_falls_back_to_defaults_without_the_extension() { + fn from_session_falls_back_to_defaults_without_the_namespace() { let session_ctx = SessionContext::new(); assert_eq!(SqlSettings::from_session(&session_ctx), SqlSettings::default()); } + + /// A `SET` writes the namespace on the shared session, so a reader that ran + /// once must not have cached the old value. This pins the whole point of + /// moving off the type-keyed extension map. + #[test] + fn a_later_set_is_visible_to_the_next_read() { + let mut config = SessionConfig::new(); + config.options_mut().extensions.insert(BeaconOptions::default()); + let session_ctx = SessionContext::new_with_config(config); + + assert_eq!(SqlSettings::from_session(&session_ctx).default_table, "default"); + + session_ctx + .state_ref() + .write() + .config_mut() + .options_mut() + .set("beacon.default_table", "observations") + .unwrap(); + + assert_eq!( + SqlSettings::from_session(&session_ctx).default_table, + "observations" + ); + } } diff --git a/beacon-db/beacon-core/src/settings_persistence.rs b/beacon-db/beacon-core/src/settings_persistence.rs new file mode 100644 index 00000000..9c5d1d42 --- /dev/null +++ b/beacon-db/beacon-core/src/settings_persistence.rs @@ -0,0 +1,244 @@ +//! Persisting `ALTER SYSTEM SET` settings in the database file. +//! +//! A plain `SET` changes the live session and is lost on restart. `ALTER SYSTEM +//! SET` writes the value here as well, into the database's own store (redb for a +//! `beacon.db` file), and the runtime replays it at startup — so an operator can +//! turn a knob on a running server without a redeploy *and* without the change +//! quietly reverting at the next restart. +//! +//! One object per key, mirroring [`secret_persistence`](crate::secret_persistence): +//! a setting is small and rewritten often, which is the case redb reclaims pages +//! for, and a per-key object means two concurrent `ALTER SYSTEM SET`s on +//! different keys cannot clobber each other. +//! +//! Precedence at startup is **persisted > environment > default**: the +//! environment builds the session config, and these are applied over it. +//! +//! An in-memory runtime (no `db_path`) has nowhere to write, and persisting is +//! skipped — the same rule persisted secrets follow. + +use std::sync::Arc; + +use anyhow::Context as _; +use datafusion::prelude::SessionContext; +use futures::StreamExt as _; +use object_store::{ObjectStore, ObjectStoreExt as _, path::Path}; + +/// The store persisted settings are written to, published as a session +/// extension so `ALTER SYSTEM` can reach it with only a `SessionContext` in hand. +/// +/// `None` for an in-memory database, where persistence has nowhere durable to go +/// — the same rule persisted secrets follow. +#[derive(Debug, Clone, Default)] +pub struct SettingsPersistence(Option>); + +impl SettingsPersistence { + pub(crate) fn new(store: Option>) -> Self { + Self(store) + } + + /// The store, or `None` when this runtime persists nothing. + pub(crate) fn store(&self) -> Option<&Arc> { + self.0.as_ref() + } + + /// The persistence published on `session_ctx`, or an unavailable one for a + /// session beacon did not build. + pub(crate) fn from_session(session_ctx: &SessionContext) -> Self { + session_ctx + .state() + .config() + .get_extension::() + .map(|persistence| (*persistence).clone()) + .unwrap_or_default() + } +} + +/// The store prefix persisted settings live under. Distinct from the table +/// (`.../table.json`) and secret layouts, so the three never collide. +const SETTINGS_PREFIX: &str = "__beacon_settings__"; + +/// The on-disk form of one persisted setting. +/// +/// The key is stored alongside the value rather than only in the object name, so +/// a listing needs no unescaping and a future rename can migrate the file names +/// without losing what each object means. +#[derive(Debug, serde::Serialize, serde::Deserialize)] +struct PersistedSetting { + key: String, + value: String, +} + +/// The object a setting is stored at. +/// +/// Keys are `[a-z0-9_.]` by construction — they are resolved against +/// [`BeaconOptions`](beacon_datafusion_ext::settings::BeaconOptions) or +/// DataFusion's own option table before they reach here — so the key is a safe +/// path segment as it stands. +fn setting_path(key: &str) -> Path { + Path::from(format!("{SETTINGS_PREFIX}/{key}.json")) +} + +/// Write `key` = `value` into `store`, replacing any earlier value. +pub(crate) async fn persist_setting( + store: &Arc, + key: &str, + value: &str, +) -> anyhow::Result<()> { + let bytes = serde_json::to_vec(&PersistedSetting { + key: key.to_string(), + value: value.to_string(), + })?; + store + .put(&setting_path(key), bytes.into()) + .await + .with_context(|| format!("failed to persist setting '{key}'"))?; + Ok(()) +} + +/// Remove a persisted setting. A missing object is not an error: `ALTER SYSTEM +/// RESET` on a key that was never persisted is a no-op, not a failure. +pub(crate) async fn remove_persisted_setting( + store: &Arc, + key: &str, +) -> anyhow::Result<()> { + match store.delete(&setting_path(key)).await { + Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()), + Err(error) => Err(anyhow::anyhow!( + "failed to remove persisted setting '{key}': {error}" + )), + } +} + +/// Every persisted setting in `store`, as `(key, value)`. +/// +/// An entry that fails to parse is logged and skipped: a settings file that a +/// future version wrote differently must not stop the database from opening, +/// since the alternative is a server that cannot start and cannot be fixed +/// through SQL. +pub(crate) async fn load_persisted_settings( + store: &Arc, +) -> anyhow::Result> { + let mut settings = Vec::new(); + let mut listing = store.list(Some(&Path::from(SETTINGS_PREFIX))); + while let Some(entry) = listing.next().await { + let location = match entry { + Ok(meta) => meta.location, + Err(error) => { + tracing::error!("failed to list persisted settings: {error}"); + continue; + } + }; + match load_one(store, &location).await { + Ok(setting) => settings.push((setting.key, setting.value)), + Err(error) => { + tracing::error!("skipping persisted setting at {location}: {error:#}"); + } + } + } + // Deterministic order, so replaying them logs the same way every boot. + settings.sort(); + Ok(settings) +} + +async fn load_one( + store: &Arc, + location: &Path, +) -> anyhow::Result { + let bytes = store.get(location).await?.bytes().await?; + serde_json::from_slice(&bytes).context("parsing persisted setting") +} + +#[cfg(test)] +mod tests { + use super::*; + + fn store() -> Arc { + Arc::new(object_store::memory::InMemory::new()) + } + + #[tokio::test] + async fn a_persisted_setting_round_trips() { + let store = store(); + persist_setting(&store, "beacon.default_table", "observations") + .await + .unwrap(); + persist_setting(&store, "datafusion.execution.batch_size", "8192") + .await + .unwrap(); + + assert_eq!( + load_persisted_settings(&store).await.unwrap(), + vec![ + ( + "beacon.default_table".to_string(), + "observations".to_string() + ), + ( + "datafusion.execution.batch_size".to_string(), + "8192".to_string() + ), + ] + ); + } + + /// Setting the same key twice leaves one value, not two — the second + /// `ALTER SYSTEM SET` has to win rather than accumulate. + #[tokio::test] + async fn re_persisting_a_key_replaces_its_value() { + let store = store(); + persist_setting(&store, "beacon.default_table", "first") + .await + .unwrap(); + persist_setting(&store, "beacon.default_table", "second") + .await + .unwrap(); + + assert_eq!( + load_persisted_settings(&store).await.unwrap(), + vec![("beacon.default_table".to_string(), "second".to_string())] + ); + } + + /// `ALTER SYSTEM RESET` on a key that was never persisted is a no-op. + #[tokio::test] + async fn removing_an_absent_setting_succeeds() { + let store = store(); + remove_persisted_setting(&store, "beacon.default_table") + .await + .unwrap(); + + persist_setting(&store, "beacon.default_table", "observations") + .await + .unwrap(); + remove_persisted_setting(&store, "beacon.default_table") + .await + .unwrap(); + assert!(load_persisted_settings(&store).await.unwrap().is_empty()); + } + + /// One unreadable object must not hide the rest: a server whose settings file + /// a future version wrote differently still has to start. + #[tokio::test] + async fn a_corrupt_entry_is_skipped() { + let store = store(); + persist_setting(&store, "beacon.default_table", "observations") + .await + .unwrap(); + store + .put( + &Path::from(format!("{SETTINGS_PREFIX}/broken.json")), + bytes::Bytes::from_static(b"not json").into(), + ) + .await + .unwrap(); + + assert_eq!( + load_persisted_settings(&store).await.unwrap(), + vec![( + "beacon.default_table".to_string(), + "observations".to_string() + )] + ); + } +} diff --git a/beacon-db/beacon-core/src/statement_plan/actions.rs b/beacon-db/beacon-core/src/statement_plan/actions.rs index 057fc719..851f8596 100644 --- a/beacon-db/beacon-core/src/statement_plan/actions.rs +++ b/beacon-db/beacon-core/src/statement_plan/actions.rs @@ -180,6 +180,101 @@ pub(crate) async fn drop_secret( } } +/// `ALTER SYSTEM SET = ` / `ALTER SYSTEM RESET `. +/// +/// Two effects: the live session config changes, and the database file records +/// (or forgets) the value so a restart replays it. +/// +/// A `RESET` restores the value the runtime had before any persisted override — +/// its environment variable, or the compiled default — and drops that override. +/// +/// The order is: refuse what cannot be done, apply, then write. A runtime with no +/// durable store is rejected *before* the session changes, so a failed statement +/// leaves nothing behind; applying before the write then means an invalid value +/// is rejected before anything reaches disk. +pub(crate) async fn alter_system( + session: &Arc, + key: &str, + value: Option<&str>, +) -> anyhow::Result<()> { + let persistence = crate::settings_persistence::SettingsPersistence::from_session(session); + let Some(store) = persistence.store().cloned() else { + anyhow::bail!( + "cannot persist `{key}`: ALTER SYSTEM needs a file-backed database, not an \ + in-memory one — use `SET {key}` for a live-only change" + ); + }; + + let effective = match value { + Some(value) => value.to_string(), + None => pre_persisted_value(session, key)?, + }; + + session + .state_ref() + .write() + .config_mut() + .options_mut() + .set(key, &effective) + .map_err(|error| anyhow::anyhow!("cannot set `{key}`: {error}"))?; + + match value { + Some(value) => crate::settings_persistence::persist_setting(&store, key, value).await, + None => crate::settings_persistence::remove_persisted_setting(&store, key).await, + } +} + +/// The value `key` held *before* any persisted override — the environment's, or +/// the compiled default. +/// +/// Not the startup value a plain `RESET` uses: `ALTER SYSTEM RESET` is deleting +/// the persisted value, so restoring the startup state would put back the very +/// value it just removed. +fn pre_persisted_value(session: &Arc, key: &str) -> anyhow::Result { + beacon_datafusion_ext::settings::BootSettings::from_config(session.state().config()) + .environment(key) + .map(str::to_string) + .ok_or_else(|| { + anyhow::anyhow!("cannot reset `{key}`: it had no value when the server started") + }) +} + +/// `SHOW SETTINGS`: one row per runtime-settable setting, with the value it holds +/// now and the one the runtime booted with (what a `RESET` restores). +/// +/// Only the `beacon.*` namespace. The `datafusion.*` half is engine internals and +/// stays in `information_schema.df_settings`, which is where a DataFusion user +/// looks for it. +pub(crate) fn show_settings( + session: &Arc, +) -> anyhow::Result { + use arrow::array::{ArrayRef, StringArray}; + + let config = session.state(); + let config = config.config(); + let boot = beacon_datafusion_ext::settings::BootSettings::from_config(config); + + use datafusion::common::config::ExtensionOptions as _; + let mut entries = beacon_datafusion_ext::settings::BeaconOptions::from_config(config).entries(); + entries.sort_by(|a, b| a.key.cmp(&b.key)); + + let names: Vec<&str> = entries.iter().map(|entry| entry.key.as_str()).collect(); + let values: Vec> = entries.iter().map(|entry| entry.value.as_deref()).collect(); + let defaults: Vec> = entries.iter().map(|entry| boot.get(&entry.key)).collect(); + let descriptions: Vec<&str> = entries.iter().map(|entry| entry.description).collect(); + + let columns: Vec = vec![ + Arc::new(StringArray::from(names)), + Arc::new(StringArray::from(values)), + Arc::new(StringArray::from(defaults)), + Arc::new(StringArray::from(descriptions)), + ]; + Ok(arrow::array::RecordBatch::try_new( + super::logical::show_settings_arrow_schema(), + columns, + )?) +} + /// `SHOW SECRETS`: one row per secret — name, type, scope, and option *keys* (never values). pub(crate) async fn show_secrets( session: &Arc, @@ -529,9 +624,14 @@ pub(crate) async fn create_table( // backend (S3 only ever applies to the datasets store). let warehouse = lance_warehouse(session)?; let namespace = beacon_lance::beacon_namespace(); - let table = - beacon_lance::create_lance_table(warehouse.clone(), &namespace, &table_name, &arrow_schema) - .await?; + let table = beacon_lance::create_lance_table( + warehouse.clone(), + &namespace, + &table_name, + &arrow_schema, + &beacon_lance::LanceConfig::from_config(session.state().config()), + ) + .await?; let location = table.definition().location.clone(); let provider: Arc = Arc::new(table); @@ -635,7 +735,13 @@ pub(crate) async fn replace_table_contents( } None => { let stream = execute_stream(child, task_ctx.clone())?; - beacon_lance::replace_table_contents(&warehouse, &location, stream).await?; + beacon_lance::replace_table_contents( + &warehouse, + &location, + stream, + &beacon_lance::LanceConfig::from_config(task_ctx.session_config()), + ) + .await?; } } return Ok(()); diff --git a/beacon-db/beacon-core/src/statement_plan/logical.rs b/beacon-db/beacon-core/src/statement_plan/logical.rs index 759793f0..8273c838 100644 --- a/beacon-db/beacon-core/src/statement_plan/logical.rs +++ b/beacon-db/beacon-core/src/statement_plan/logical.rs @@ -945,6 +945,106 @@ fn show_extensions_df_schema() -> &'static DFSchemaRef { }) } +/// Logical node for `ALTER SYSTEM SET = ` and +/// `ALTER SYSTEM RESET `. +/// +/// `value: None` is the `RESET` form: drop the persisted value and restore what +/// the runtime booted with. +#[derive(Debug, PartialEq, Eq, PartialOrd, Hash)] +pub(crate) struct AlterSystemNode { + /// The setting name, already resolved (aliases applied, startup-only keys + /// rejected) — see `statement_plan::settings`. + pub(crate) key: String, + pub(crate) value: Option, +} + +impl AlterSystemNode { + pub(crate) fn new(key: String, value: Option) -> Self { + Self { key, value } + } +} + +impl UserDefinedLogicalNodeCore for AlterSystemNode { + fn name(&self) -> &str { + "AlterSystem" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + fn schema(&self) -> &DFSchemaRef { + empty_schema() + } + fn expressions(&self) -> Vec { + vec![] + } + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + match &self.value { + Some(value) => write!(f, "AlterSystem: SET {} = {value}", self.key), + None => write!(f, "AlterSystem: RESET {}", self.key), + } + } + fn with_exprs_and_inputs(&self, _exprs: Vec, _inputs: Vec) -> Result { + Ok(Self { + key: self.key.clone(), + value: self.value.clone(), + }) + } +} + +/// Logical node for `SHOW SETTINGS`. +#[derive(Debug, PartialEq, Eq, PartialOrd, Hash)] +pub(crate) struct ShowSettingsNode; + +impl UserDefinedLogicalNodeCore for ShowSettingsNode { + fn name(&self) -> &str { + "ShowSettings" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + fn schema(&self) -> &DFSchemaRef { + show_settings_df_schema() + } + fn expressions(&self) -> Vec { + vec![] + } + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "ShowSettings") + } + fn with_exprs_and_inputs(&self, _exprs: Vec, _inputs: Vec) -> Result { + Ok(Self) + } +} + +/// Arrow schema produced by `SHOW SETTINGS`. +/// +/// `default` is what the runtime booted with (the environment variable), which is +/// what a `RESET` restores — so an operator can see both what a setting is now +/// and what it would go back to. +pub(crate) fn show_settings_arrow_schema() -> Arc { + static SCHEMA: OnceLock> = OnceLock::new(); + SCHEMA + .get_or_init(|| { + Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, false), + Field::new("value", DataType::Utf8, true), + Field::new("default", DataType::Utf8, true), + Field::new("description", DataType::Utf8, false), + ])) + }) + .clone() +} + +fn show_settings_df_schema() -> &'static DFSchemaRef { + static SCHEMA: OnceLock = OnceLock::new(); + SCHEMA.get_or_init(|| { + Arc::new( + DFSchema::try_from(show_settings_arrow_schema().as_ref().clone()) + .expect("SHOW SETTINGS schema is valid"), + ) + }) +} + /// Logical node for `SET EXTENSION '' FOR
TO ''`. #[derive(Debug, PartialEq, Eq, PartialOrd, Hash)] pub(crate) struct SetExtensionNode { diff --git a/beacon-db/beacon-core/src/statement_plan/lower.rs b/beacon-db/beacon-core/src/statement_plan/lower.rs index 680931e2..5e6d2c6b 100644 --- a/beacon-db/beacon-core/src/statement_plan/lower.rs +++ b/beacon-db/beacon-core/src/statement_plan/lower.rs @@ -68,6 +68,11 @@ pub(crate) async fn lower_df_statement( session_ctx: &Arc, statement: datafusion::sql::parser::Statement, ) -> anyhow::Result { + // `SET`/`RESET`/`SHOW` over the `beacon.*` namespace, resolved here because + // `SHOW` validates its name during planning and `RESET` would otherwise + // restore DataFusion's compiled default over the operator's environment. + let statement = super::settings::rewrite_settings_statement(session_ctx, statement)?; + // DataFusion has no `ALTER TABLE` planning, so build the node from the AST. if let datafusion::sql::parser::Statement::Statement(sql_stmt) = &statement { if let SqlAstStatement::AlterTable(alter) = sql_stmt.as_ref() { diff --git a/beacon-db/beacon-core/src/statement_plan/mod.rs b/beacon-db/beacon-core/src/statement_plan/mod.rs index bdd08c8b..31b69598 100644 --- a/beacon-db/beacon-core/src/statement_plan/mod.rs +++ b/beacon-db/beacon-core/src/statement_plan/mod.rs @@ -20,6 +20,7 @@ mod lower; pub(crate) mod materialized_view; mod physical; mod query_planner; +mod settings; mod stream_coalescer; use std::collections::HashMap; @@ -33,7 +34,7 @@ use datafusion::{ }; use crate::parser::statement::{ - AttachStatement, AuthStatement, CreateCrawlerStatement, CreateIndexStatement, + AlterSystemStatement, AttachStatement, AuthStatement, CreateCrawlerStatement, CreateIndexStatement, CreateMaterializedViewStatement, CreateSecretStatement, DetachStatement, DropCrawlerStatement, DropExtensionStatement, DropIndexStatement, DropSecretStatement, RefreshStatement, AnalyzeFilesStatement, RunCrawlerStatement, SetExtensionStatement, ShowExtensionsStatement, ShowIndexesStatement, @@ -104,21 +105,46 @@ pub(crate) fn plan_produces_result_set(plan: &LogicalPlan) -> bool { } } -/// Whether `plan` contains any [`LogicalPlan::Extension`] node (all of beacon's -/// extension nodes are super-user-only operations). +/// Whether `plan` contains any privileged [`LogicalPlan::Extension`] node. +/// +/// Beacon's extension nodes are super-user-only as a class — they create tables, +/// write secrets, or list what other users did. [`is_public_node`] carves out the +/// few that document the engine rather than the instance. fn plan_contains_extension(plan: &LogicalPlan) -> anyhow::Result { let mut found = false; plan.apply(|node| { - if matches!(node, LogicalPlan::Extension(_)) { - found = true; - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) + if let LogicalPlan::Extension(extension) = node { + if !is_public_node(extension) { + found = true; + return Ok(TreeNodeRecursion::Stop); + } } + Ok(TreeNodeRecursion::Continue) })?; Ok(found) } +/// Whether an extension node is readable by any authenticated caller. +/// +/// The exemption is deliberately one node wide. `SHOW SETTINGS` lists the engine +/// knobs and their values — the same class of information as `SHOW FUNCTIONS`, +/// which [`Runtime::show_functions`](crate::runtime::Runtime::show_functions) +/// already runs as the engine rather than as the caller, so a user can discover +/// what the engine supports. It exposes no instance state: the node emits only +/// the `beacon.*` namespace, which holds no credential and no user data, and +/// every startup-only key (`BEACON_ADMIN_*`, `BEACON_OIDC_*`, `BEACON_S3_*`, +/// `BEACON_SECRETS_KEY`) is outside that namespace by construction. +/// +/// *Changing* a setting stays super-user-only: `AlterSystemNode` is not listed +/// here, and a plain `SET` is a `LogicalPlan::Statement`, which +/// `SQLOptions::with_allow_statements` already gates above. +fn is_public_node(extension: &Extension) -> bool { + extension + .node + .as_any() + .is::() +} + /// Late-initialized, weak handle to the [`SessionContext`] shared with the /// custom planner. /// @@ -481,6 +507,26 @@ pub(crate) fn show_crawlers_plan() -> LogicalPlan { }) } +/// Build the logical plan for `ALTER SYSTEM SET = ` / +/// `ALTER SYSTEM RESET `. +/// +/// The key is resolved here, so an alias (`beacon.batch_size`) and a startup-only +/// key (`beacon.port`) behave exactly as they do for a plain `SET` — the +/// difference between the two statements is persistence, not what a key means. +pub(crate) fn alter_system_plan(statement: AlterSystemStatement) -> anyhow::Result { + let key = settings::resolve_statement_key(&statement.key)?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(logical::AlterSystemNode::new(key, statement.value)), + })) +} + +/// Build the logical plan for `SHOW SETTINGS`. +pub(crate) fn show_settings_plan() -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(logical::ShowSettingsNode), + }) +} + /// Build the logical plan for `SET EXTENSION '' FOR
TO ''`. pub(crate) fn set_extension_plan(statement: SetExtensionStatement) -> LogicalPlan { LogicalPlan::Extension(Extension { diff --git a/beacon-db/beacon-core/src/statement_plan/physical.rs b/beacon-db/beacon-core/src/statement_plan/physical.rs index 650d986d..f680ba9a 100644 --- a/beacon-db/beacon-core/src/statement_plan/physical.rs +++ b/beacon-db/beacon-core/src/statement_plan/physical.rs @@ -32,8 +32,8 @@ use super::{ logical::{ analyze_files_arrow_schema, count_arrow_schema, run_crawler_arrow_schema, show_crawlers_arrow_schema, - show_indexes_arrow_schema, show_secrets_arrow_schema, AlterTableSpec, - Mutation, + show_indexes_arrow_schema, show_secrets_arrow_schema, show_settings_arrow_schema, + AlterTableSpec, Mutation, }, materialized_view, SessionCell, }; @@ -1556,3 +1556,109 @@ impl ExecutionPlan for ShowExtensionsExec { Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } } + +/// Physical node for `ALTER SYSTEM SET = ` / +/// `ALTER SYSTEM RESET `. +/// +/// Applies the value to the live session *and* writes it into the database file, +/// so the change takes effect now and survives a restart. +#[derive(Debug)] +pub(crate) struct AlterSystemExec { + key: String, + value: Option, + session: SessionCell, + cache: Arc, +} + +impl AlterSystemExec { + pub(crate) fn new(key: String, value: Option, session: SessionCell) -> Self { + Self { + key, + value, + session, + cache: Arc::new(side_effect_properties()), + } + } + fn fmt_label(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.value { + Some(value) => write!(f, "AlterSystemExec: SET {} = {value}", self.key), + None => write!(f, "AlterSystemExec: RESET {}", self.key), + } + } +} + +side_effect_exec!( + AlterSystemExec, + "AlterSystemExec", + |exec: &AlterSystemExec| { + let session = upgrade_session(&exec.session)?; + let key = exec.key.clone(); + let value = exec.value.clone(); + Ok(side_effect_stream(async move { + actions::alter_system(&session, &key, value.as_deref()) + .await + .map_err(to_df_err) + })) + } +); + +/// Physical node for `SHOW SETTINGS` — one row per runtime-settable setting. +#[derive(Debug)] +pub(crate) struct ShowSettingsExec { + session: SessionCell, + cache: Arc, +} + +impl ShowSettingsExec { + pub(crate) fn new(session: SessionCell) -> Self { + Self { + session, + cache: Arc::new(plan_properties(show_settings_arrow_schema())), + } + } +} + +impl DisplayAs for ShowSettingsExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ShowSettingsExec") + } + DisplayFormatType::TreeRender => write!(f, "ShowSettingsExec"), + } + } +} + +impl ExecutionPlan for ShowSettingsExec { + fn name(&self) -> &str { + "ShowSettingsExec" + } + fn as_any(&self) -> &dyn Any { + self + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let session = upgrade_session(&self.session)?; + let schema = show_settings_arrow_schema(); + let stream = + futures::stream::once( + async move { actions::show_settings(&session).map_err(to_df_err) }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} diff --git a/beacon-db/beacon-core/src/statement_plan/query_planner.rs b/beacon-db/beacon-core/src/statement_plan/query_planner.rs index 5785f29b..a6d541ab 100644 --- a/beacon-db/beacon-core/src/statement_plan/query_planner.rs +++ b/beacon-db/beacon-core/src/statement_plan/query_planner.rs @@ -309,6 +309,18 @@ impl ExtensionPlanner for BeaconExtensionPlanner { )))); } + if let Some(alter) = any.downcast_ref::() { + return Ok(Some(Arc::new(physical::AlterSystemExec::new( + alter.key.clone(), + alter.value.clone(), + session, + )))); + } + + if any.downcast_ref::().is_some() { + return Ok(Some(Arc::new(physical::ShowSettingsExec::new(session)))); + } + // Unrecognized node: let the default planner handle it. Ok(None) } diff --git a/beacon-db/beacon-core/src/statement_plan/settings.rs b/beacon-db/beacon-core/src/statement_plan/settings.rs new file mode 100644 index 00000000..195b2474 --- /dev/null +++ b/beacon-db/beacon-core/src/statement_plan/settings.rs @@ -0,0 +1,449 @@ +//! `SET` / `RESET` / `SHOW` for the `beacon.*` namespace. +//! +//! Everything here is an **AST rewrite**, applied in +//! [`lower_df_statement`](super::lower_df_statement) before DataFusion plans the +//! statement. The AST is the only seam that works for all three: +//! `show_variable_to_plan` validates a name at *plan* time, so a later rewrite +//! would come after the error, and DataFusion's `ConfigOptions::reset` refuses +//! any prefix that is not `datafusion`. +//! +//! Three things happen to a key: +//! +//! 1. A **startup-only** key is rejected, naming the environment variable to edit. +//! 2. A `beacon..*` key is rewritten to `datafusion.<...>`. +//! This is the prefix alias: an operator can spell every engine option in one +//! namespace, and `datafusion.*` keeps working unchanged. +//! 3. Everything else is left alone. A `beacon.*` key that names a real setting +//! routes into [`BeaconOptions`] on its own, because that type is a registered +//! `ConfigExtension`. +//! +//! `RESET` becomes a `SET` to the value the runtime booted with. DataFusion's own +//! `RESET` would restore *its* compiled default, discarding whatever the +//! operator's environment supplied — a `RESET beacon.netcdf.use_rust_reader` on a +//! server started with `BEACON_NETCDF_USE_RUST_READER=true` would silently turn +//! the Rust reader off. +//! +//! # Scope and privilege +//! +//! Unchanged from what `SET datafusion.*` already did: a `SET` applies to the one +//! shared session, so it takes effect for every later query and every user, and +//! `validate_query_plan` admits `LogicalPlan::Statement` only for a super-user. + +use beacon_datafusion_ext::settings::{BeaconOptions, BootSettings, startup_only_env_var}; +use datafusion::prelude::SessionContext; +use datafusion::sql::parser::{ResetStatement, Statement as DFStatement}; +use datafusion::sql::sqlparser::ast::{ + Expr as SqlExpr, Ident, ObjectName, Set, Statement as SqlAstStatement, Value, +}; + +/// The `beacon.` namespace, with its separator. +const BEACON_PREFIX: &str = "beacon."; + +/// The top-level sections of DataFusion's own `ConfigOptions`. +/// +/// `beacon.
.` is rewritten onto `datafusion.
.`, which +/// is what makes `beacon.` a complete alias rather than a second, partial +/// namespace. Kept as a list because DataFusion hardcodes these names in +/// `ConfigOptions::visit` and exposes no way to enumerate them. +const DATAFUSION_SECTIONS: &[&str] = &[ + "catalog", + "execution", + "optimizer", + "explain", + "sql_parser", + "format", + "runtime", +]; + +/// `beacon.*` names that stand for a DataFusion option beacon already sets from +/// an environment variable, so the SQL spelling matches the documented `BEACON_*` +/// one rather than exposing where the value happens to live. +const ALIASES: &[(&str, &str)] = &[("beacon.batch_size", "datafusion.execution.batch_size")]; + +/// Rewrite the settings statements in `statement`; everything else passes +/// through untouched. +pub(crate) fn rewrite_settings_statement( + session_ctx: &SessionContext, + statement: DFStatement, +) -> anyhow::Result { + match statement { + DFStatement::Reset(ResetStatement::Variable(name)) => { + reset_to_boot_value(session_ctx, name) + } + DFStatement::Statement(statement) => match *statement { + SqlAstStatement::Set(Set::SingleAssignment { + scope, + hivevar, + variable, + values, + }) => { + let variable = resolve_object_name(&variable)?; + Ok(single_assignment(scope, hivevar, variable, values)) + } + SqlAstStatement::ShowVariable { variable } => { + Ok(show_variable(resolve_show_variable(variable)?)) + } + other => Ok(DFStatement::Statement(Box::new(other))), + }, + other => Ok(other), + } +} + +/// `RESET ` as a `SET = `. +/// +/// Falls back to DataFusion's own `RESET` when nothing was recorded for the key — +/// an option whose boot value was unset has no string that would restore it, and +/// `ConfigOptions::reset` handles those correctly for the `datafusion.*` half. +fn reset_to_boot_value( + session_ctx: &SessionContext, + name: ObjectName, +) -> anyhow::Result { + let key = resolve_object_name(&name)?; + let boot = BootSettings::from_config(session_ctx.state().config()); + + match boot.get(&key) { + Some(value) => Ok(single_assignment( + None, + false, + key, + vec![SqlExpr::Value( + Value::SingleQuotedString(value.to_string()).into(), + )], + )), + None => Ok(DFStatement::Reset(ResetStatement::Variable(object_name( + &key, + )))), + } +} + +/// The key a `beacon.*` name resolves to, or an error explaining why it cannot +/// be set. A name outside the namespace is returned unchanged. +fn resolve_key(key: &str) -> anyhow::Result { + let key = key.to_ascii_lowercase(); + + let Some(rest) = key.strip_prefix(BEACON_PREFIX) else { + // `datafusion.*`, `timezone`, and anything else DataFusion owns. + return Ok(key); + }; + + // A real beacon setting: `ConfigOptions` routes it to the extension itself. + if BeaconOptions::has_key(&key) { + return Ok(key); + } + + if let Some((_, target)) = ALIASES.iter().find(|(alias, _)| *alias == key) { + return Ok(target.to_string()); + } + + // The prefix alias: `beacon.execution.batch_size` is + // `datafusion.execution.batch_size`. + let section = rest.split('.').next().unwrap_or_default(); + if DATAFUSION_SECTIONS.contains(§ion) { + return Ok(format!("datafusion.{rest}")); + } + + if let Some(env_var) = startup_only_env_var(&key) { + anyhow::bail!( + "`{key}` can only be set when the server starts: set the `{env_var}` \ + environment variable and restart" + ); + } + + anyhow::bail!( + "unknown setting `{key}`: `SHOW SETTINGS` lists every setting that can be changed \ + at runtime" + ) +} + +/// [`resolve_key`] over a dotted `ObjectName`, as `SET` and `RESET` carry it. +fn resolve_object_name(name: &ObjectName) -> anyhow::Result { + resolve_key(&join_parts(name)) +} + +/// The resolved key of an `ALTER SYSTEM` statement. +/// +/// The same resolution the live `SET` gets, so the two statements agree on what +/// a name means and differ only in whether the value is written to disk. +pub(crate) fn resolve_statement_key(name: &ObjectName) -> anyhow::Result { + resolve_object_name(name) +} + +/// [`resolve_key`] over the identifier list `SHOW` carries, preserving a trailing +/// `VERBOSE` (which DataFusion strips to widen the output, and which is not part +/// of the key). +fn resolve_show_variable(variable: Vec) -> anyhow::Result> { + let verbose = variable + .last() + .is_some_and(|ident| ident.value.eq_ignore_ascii_case("verbose")); + let (key_parts, suffix) = match verbose { + true => variable.split_at(variable.len() - 1), + false => (variable.as_slice(), &[][..]), + }; + + let key: String = key_parts + .iter() + .map(|ident| ident.value.as_str()) + .collect::>() + .join("."); + + // `SHOW ALL`, `SHOW TIMEZONE` and friends are DataFusion's, and a bare `SHOW` + // has nothing to resolve. + if key.is_empty() || !key.to_ascii_lowercase().starts_with(BEACON_PREFIX) { + return Ok(variable); + } + + let mut resolved = idents(&resolve_key(&key)?); + resolved.extend_from_slice(suffix); + Ok(resolved) +} + +/// The dotted string form of an object name, ignoring quoting. +fn join_parts(name: &ObjectName) -> String { + name.0 + .iter() + .map(|part| { + part.as_ident() + .map(|ident| ident.value.clone()) + .unwrap_or_default() + }) + .collect::>() + .join(".") +} + +fn idents(key: &str) -> Vec { + key.split('.').map(Ident::new).collect() +} + +fn object_name(key: &str) -> ObjectName { + ObjectName::from(idents(key)) +} + +fn show_variable(variable: Vec) -> DFStatement { + DFStatement::Statement(Box::new(SqlAstStatement::ShowVariable { variable })) +} + +fn single_assignment( + scope: Option, + hivevar: bool, + variable: String, + values: Vec, +) -> DFStatement { + DFStatement::Statement(Box::new(SqlAstStatement::Set(Set::SingleAssignment { + scope, + hivevar, + variable: object_name(&variable), + values, + }))) +} + +#[cfg(test)] +mod tests { + use super::*; + + use datafusion::execution::context::SessionConfig; + use datafusion::sql::parser::DFParser; + + fn parse(sql: &str) -> DFStatement { + DFParser::parse_sql(sql).unwrap().pop_front().unwrap() + } + + /// A session carrying the namespace and a boot snapshot, as the runtime + /// builds one. + fn session(default_table: &str) -> SessionContext { + let mut config = SessionConfig::new(); + config.options_mut().extensions.insert(BeaconOptions { + default_table: default_table.to_string(), + ..Default::default() + }); + config.options_mut().execution.batch_size = 4096; + let boot = BootSettings::capture(config.options()); + SessionContext::new_with_config(config.with_extension(std::sync::Arc::new(boot))) + } + + fn rewrite(sql: &str) -> anyhow::Result { + let ctx = session("observations"); + rewrite_settings_statement(&ctx, parse(sql)).map(|statement| statement.to_string()) + } + + /// The dotted key a rewritten `SHOW` carries. + /// + /// Asserted on the identifiers rather than the rendered statement: + /// `ShowVariable`'s `Display` joins its parts with a *space*, while + /// DataFusion's planner joins the same parts with a dot to look the key up. + /// The rendering would therefore make a correct rewrite look wrong. + fn show_key(sql: &str) -> anyhow::Result { + let ctx = session("observations"); + let statement = rewrite_settings_statement(&ctx, parse(sql))?; + let DFStatement::Statement(statement) = statement else { + panic!("`{sql}` did not stay a SHOW"); + }; + let SqlAstStatement::ShowVariable { variable } = *statement else { + panic!("`{sql}` did not stay a SHOW"); + }; + Ok(variable + .iter() + .map(|ident| ident.value.as_str()) + .collect::>() + .join(".")) + } + + #[test] + fn a_beacon_setting_passes_through_unchanged() { + assert_eq!( + rewrite("SET beacon.netcdf.use_rust_reader = true").unwrap(), + "SET beacon.netcdf.use_rust_reader = true" + ); + } + + /// The prefix alias: every DataFusion section is reachable under `beacon.`, + /// which is the whole point of the rewrite. + #[test] + fn a_datafusion_section_is_rewritten_onto_the_datafusion_prefix() { + for (input, expected) in [ + ( + "SET beacon.execution.batch_size = 8192", + "SET datafusion.execution.batch_size = 8192", + ), + ( + "SET beacon.optimizer.max_passes = 1", + "SET datafusion.optimizer.max_passes = 1", + ), + ( + "SET beacon.sql_parser.dialect = 'postgres'", + "SET datafusion.sql_parser.dialect = 'postgres'", + ), + ] { + assert_eq!(rewrite(input).unwrap(), expected, "for `{input}`"); + } + } + + /// `beacon.batch_size` is the documented `BEACON_BATCH_SIZE`, which the + /// runtime funnels into DataFusion's batch size. + #[test] + fn a_documented_alias_reaches_its_datafusion_option() { + assert_eq!( + rewrite("SET beacon.batch_size = 64000").unwrap(), + "SET datafusion.execution.batch_size = 64000" + ); + } + + /// `datafusion.*` keeps working exactly as before, and so does every + /// statement that is not a setting. + #[test] + fn statements_outside_the_namespace_are_untouched() { + for sql in [ + "SET datafusion.execution.batch_size = 8192", + "SET timezone = 'UTC'", + "SELECT 1", + ] { + assert_eq!(rewrite(sql).unwrap(), sql, "for `{sql}`"); + } + assert_eq!(show_key("SHOW ALL").unwrap(), "ALL"); + assert_eq!( + show_key("SHOW datafusion.execution.batch_size").unwrap(), + "datafusion.execution.batch_size" + ); + } + + /// A startup-only key would appear to work and change nothing, so it is + /// rejected — and the error has to name the variable to edit instead. + #[test] + fn a_startup_only_key_names_its_environment_variable() { + let err = rewrite("SET beacon.port = 1234").unwrap_err().to_string(); + assert!(err.contains("BEACON_PORT"), "unhelpful error: {err}"); + assert!(err.contains("restart"), "unhelpful error: {err}"); + + // Including the ones behind a family prefix. + let err = rewrite("SET beacon.s3.bucket = 'x'") + .unwrap_err() + .to_string(); + assert!(err.contains("BEACON_S3_*"), "unhelpful error: {err}"); + + // And the cache capacities, which look settable next to their siblings. + let err = rewrite("SET beacon.netcdf.reader_cache_size = 8") + .unwrap_err() + .to_string(); + assert!( + err.contains("BEACON_NETCDF_READER_CACHE_SIZE"), + "unhelpful error: {err}" + ); + } + + #[test] + fn an_unknown_key_points_at_show_settings() { + let err = rewrite("SET beacon.nope = 1").unwrap_err().to_string(); + assert!(err.contains("beacon.nope"), "unhelpful error: {err}"); + assert!(err.contains("SHOW SETTINGS"), "unhelpful error: {err}"); + } + + /// `SHOW` is validated at plan time against the key `entries()` advertises, + /// so an aliased name has to be rewritten before planning, not after. + #[test] + fn show_resolves_the_same_names_as_set() { + assert_eq!( + show_key("SHOW beacon.execution.batch_size").unwrap(), + "datafusion.execution.batch_size" + ); + assert_eq!( + show_key("SHOW beacon.default_table").unwrap(), + "beacon.default_table" + ); + // A trailing VERBOSE widens the output; it is not part of the key. + assert_eq!( + show_key("SHOW beacon.batch_size VERBOSE").unwrap(), + "datafusion.execution.batch_size.VERBOSE" + ); + } + + /// The case beacon has to own: DataFusion's `RESET` restores *its* compiled + /// default, which would discard the operator's environment value. Restoring + /// the recorded boot value is the whole reason `RESET` is intercepted. + #[test] + fn reset_restores_the_value_the_runtime_booted_with() { + assert_eq!( + rewrite("RESET beacon.default_table").unwrap(), + "SET beacon.default_table = 'observations'" + ); + // …including through the alias, where DataFusion's compiled default + // (8192) differs from what this runtime started with. + assert_eq!( + rewrite("RESET beacon.batch_size").unwrap(), + "SET datafusion.execution.batch_size = '4096'" + ); + } + + /// A key with no recorded boot value (an option that started unset) falls + /// back to DataFusion's own `RESET`, which handles those. + #[test] + fn reset_falls_back_when_nothing_was_recorded() { + let ctx = SessionContext::new(); + let rewritten = + rewrite_settings_statement(&ctx, parse("RESET datafusion.execution.time_zone")) + .unwrap(); + assert_eq!( + rewritten.to_string(), + "RESET datafusion.execution.time_zone" + ); + } + + /// A startup-only key must be refused on `RESET` too, not only on `SET`. + #[test] + fn reset_rejects_a_startup_only_key() { + let err = rewrite("RESET beacon.port").unwrap_err().to_string(); + assert!(err.contains("BEACON_PORT"), "unhelpful error: {err}"); + } + + /// Beacon turns off DataFusion's identifier normalization, so a key typed in + /// upper case would otherwise miss both the settings table and the alias list. + #[test] + fn keys_are_matched_without_regard_to_case() { + assert_eq!( + rewrite("SET BEACON.EXECUTION.BATCH_SIZE = 8192").unwrap(), + "SET datafusion.execution.batch_size = 8192" + ); + assert_eq!( + show_key("SHOW Beacon.Default_Table").unwrap(), + "beacon.default_table" + ); + } +} diff --git a/beacon-db/beacon-core/src/statement_plan/stream_coalescer.rs b/beacon-db/beacon-core/src/statement_plan/stream_coalescer.rs index ab763e67..2409be16 100644 --- a/beacon-db/beacon-core/src/statement_plan/stream_coalescer.rs +++ b/beacon-db/beacon-core/src/statement_plan/stream_coalescer.rs @@ -11,9 +11,10 @@ use crate::settings::SqlStreamCoalesceSettings; /// least [`SqlStreamCoalesceSettings::target_rows`] rows before they reach a /// client. /// -/// Published on the session config as an extension by the runtime builder, so +/// Configured through the `beacon.*` namespace on the session config, so /// execution-time code recovers it with [`CoalesceSqlStream::from_session`] -/// rather than threading a `Runtime` handle through. +/// rather than threading a `Runtime` handle through — and picks up a +/// `SET beacon.sql.stream_coalesce.*` on the very next statement. #[derive(Debug, Clone, Copy, Default)] pub(crate) struct CoalesceSqlStream { settings: SqlStreamCoalesceSettings, @@ -24,15 +25,10 @@ impl CoalesceSqlStream { Self { settings } } - /// Recovers the coalescer published on `session_ctx`, falling back to the - /// defaults if the extension is absent (a session beacon did not build). + /// Recovers the coalescer configured on `session_ctx`, falling back to the + /// defaults if the namespace is absent (a session beacon did not build). pub(crate) fn from_session(session_ctx: &SessionContext) -> Self { - session_ctx - .state() - .config() - .get_extension::() - .map(|coalescer| *coalescer) - .unwrap_or_default() + Self::new(crate::settings::SqlSettings::from_session(session_ctx).stream_coalesce) } /// Wraps `stream` in the coalescing adapter. Returns `stream` unchanged when @@ -251,18 +247,24 @@ mod tests { assert_eq!(output_batches[1].num_rows(), 2_000); } - /// The coalescer is recovered from the session extension the runtime builder - /// publishes, and falls back to the defaults on a session without one. + /// The coalescer is recovered from the `beacon.*` namespace the runtime + /// builder publishes, and falls back to the defaults on a session without one. #[tokio::test] - async fn reads_settings_from_the_session_extension() { + async fn reads_settings_from_the_session_namespace() { let settings = SqlStreamCoalesceSettings { enabled: true, target_rows: 30_000, flush_timeout_ms: 10_000, max_rows: 200_000, }; - let config = - SessionConfig::new().with_extension(Arc::new(CoalesceSqlStream::new(settings))); + let mut options = beacon_datafusion_ext::settings::BeaconOptions::default(); + crate::settings::SqlSettings { + stream_coalesce: settings, + ..Default::default() + } + .apply_to(&mut options); + let mut config = SessionConfig::new(); + config.options_mut().extensions.insert(options); let session_ctx = SessionContext::new_with_config(config); let output_batches = CoalesceSqlStream::from_session(&session_ctx) @@ -280,7 +282,7 @@ mod tests { assert_eq!( CoalesceSqlStream::from_session(&SessionContext::new()).settings, SqlStreamCoalesceSettings::default(), - "a session without the extension should fall back to the defaults" + "a session without the namespace should fall back to the defaults" ); } } diff --git a/beacon-db/beacon-core/src/system_schema/mod.rs b/beacon-db/beacon-core/src/system_schema/mod.rs index 7ebaa672..9cb55686 100644 --- a/beacon-db/beacon-core/src/system_schema/mod.rs +++ b/beacon-db/beacon-core/src/system_schema/mod.rs @@ -1,7 +1,8 @@ //! `beacon.system`: runtime introspection exposed as ordinary SQL tables. //! //! What the runtime knows about itself that is not user data lives here — the -//! auth directory and the recorded query metrics — so it is reachable through +//! auth directory, the recorded query metrics, the runtime settings — so it is +//! reachable through //! the one query endpoint rather than through a typed method on //! [`Runtime`](crate::runtime::Runtime) and a bespoke HTTP route per item. //! @@ -28,6 +29,7 @@ mod auth; mod file_stats; +mod settings; mod table; pub(crate) use file_stats::{FileStatisticsFunc, FileStatisticsTable}; @@ -106,6 +108,10 @@ impl SystemSchemaProvider { "file_stats_segments".to_string(), Arc::new(file_stats::segments_table(file_stats)), ); + tables.insert( + "settings".to_string(), + Arc::new(settings::settings_table(session.clone())), + ); Self { tables, session } } diff --git a/beacon-db/beacon-core/src/system_schema/settings.rs b/beacon-db/beacon-core/src/system_schema/settings.rs new file mode 100644 index 00000000..f06ea8a3 --- /dev/null +++ b/beacon-db/beacon-core/src/system_schema/settings.rs @@ -0,0 +1,77 @@ +//! `beacon.system.settings` — every runtime-settable setting, as SQL. +//! +//! The table form of `SHOW SETTINGS`, for a client that would rather filter and +//! join than parse a `SHOW`. Same rows, same columns. +//! +//! Unlike the statement, this table is super-user-only, because everything in +//! `beacon.system` is (see the module docs and +//! [`authorize_logical_plan`](crate::statement_plan::authorize_logical_plan)). +//! A regular user reads the same values through `SHOW SETTINGS`. + +use std::sync::Arc; + +use arrow::{ + array::{ArrayRef, StringArray}, + datatypes::{DataType, Field, Schema, SchemaRef}, + record_batch::RecordBatch, +}; +use beacon_datafusion_ext::settings::{BeaconOptions, BootSettings}; +use datafusion::common::Result as DFResult; +use datafusion::common::config::ExtensionOptions as _; + +use super::table::{Snapshot, SystemTable}; +use crate::statement_plan::{SessionCell, upgrade_session}; + +fn settings_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, false), + Field::new("value", DataType::Utf8, true), + // What the runtime booted with — the `BEACON_*` variable, or the compiled + // default. This is what `RESET ` restores. + Field::new("default", DataType::Utf8, true), + Field::new("description", DataType::Utf8, false), + ])) +} + +/// `beacon.system.settings` — one row per `beacon.*` setting. +/// +/// Snapshotted per scan through the session, so a `SET` in one statement is +/// visible to a `SELECT` in the next. +pub(super) fn settings_table(session: SessionCell) -> SystemTable { + let snapshot: Snapshot = Arc::new(move || { + let session = session.clone(); + Box::pin(async move { + let Ok(session) = upgrade_session(&session, "beacon.system.settings") else { + return empty_batch(); + }; + let state = session.state(); + let config = state.config(); + let boot = BootSettings::from_config(config); + + let mut entries = BeaconOptions::from_config(config).entries(); + entries.sort_by(|left, right| left.key.cmp(&right.key)); + + let names: Vec<&str> = entries.iter().map(|entry| entry.key.as_str()).collect(); + let values: Vec> = + entries.iter().map(|entry| entry.value.as_deref()).collect(); + let defaults: Vec> = + entries.iter().map(|entry| boot.get(&entry.key)).collect(); + let descriptions: Vec<&str> = entries.iter().map(|entry| entry.description).collect(); + + let columns: Vec = vec![ + Arc::new(StringArray::from(names)), + Arc::new(StringArray::from(values)), + Arc::new(StringArray::from(defaults)), + Arc::new(StringArray::from(descriptions)), + ]; + Ok(RecordBatch::try_new(settings_schema(), columns)?) + }) + }); + SystemTable::new(settings_schema(), snapshot) +} + +/// No rows, for a torn-down runtime — the same answer the other system tables +/// give when their source is unavailable. +fn empty_batch() -> DFResult { + Ok(RecordBatch::new_empty(settings_schema())) +} diff --git a/beacon-db/beacon-core/tests/redb_tables.rs b/beacon-db/beacon-core/tests/redb_tables.rs index 9e423067..8975934a 100644 --- a/beacon-db/beacon-core/tests/redb_tables.rs +++ b/beacon-db/beacon-core/tests/redb_tables.rs @@ -59,7 +59,13 @@ async fn lance_table_lifecycle_over_redb() { let warehouse = Arc::new(LanceWarehouse::new(tables.clone())); let namespace = beacon_namespace(); - let table = create_lance_table(warehouse.clone(), &namespace, "orders", &sample_schema()) + let table = create_lance_table( + warehouse.clone(), + &namespace, + "orders", + &sample_schema(), + &Default::default(), + ) .await .expect("create table on redb"); let location = table.definition().location.clone(); @@ -87,7 +93,7 @@ async fn lance_table_lifecycle_over_redb() { .execute_stream() .await .unwrap(); - replace_table_contents(&warehouse, &location, keep) + replace_table_contents(&warehouse, &location, keep, &Default::default()) .await .expect("replace (delete) on redb"); assert_eq!( diff --git a/beacon-db/beacon-core/tests/runtime_settings.rs b/beacon-db/beacon-core/tests/runtime_settings.rs new file mode 100644 index 00000000..18b8acd8 --- /dev/null +++ b/beacon-db/beacon-core/tests/runtime_settings.rs @@ -0,0 +1,361 @@ +//! `SET` / `RESET` / `ALTER SYSTEM` / `SHOW SETTINGS` end to end, through a real +//! runtime. +//! +//! The unit tests cover the pieces (the namespace round-trips, the AST rewrite +//! resolves a key). These prove the pieces are joined: a `SET` reaches the shared +//! session, a later statement sees it, an `ALTER SYSTEM SET` survives a restart, +//! and the privilege boundary holds. +//! +//! Every assertion reads a value back through SQL rather than through the +//! `Runtime`, which exposes no config getter on purpose — and reading a setting +//! from a SQL client is the feature under test. + +mod common; + +use arrow::array::{Array as _, StringArray}; +use arrow::record_batch::RecordBatch; +use beacon_core::AuthIdentity; +use common::TestRuntime; + +/// The single string in a one-row, one-column result. +fn scalar_str(batches: &[RecordBatch]) -> String { + let column = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("a string column"); + assert_eq!(column.len(), 1, "expected exactly one row"); + column.value(0).to_string() +} + +/// Every value in the first column, across batches. +fn column_strings(batches: &[RecordBatch], index: usize) -> Vec { + batches + .iter() + .flat_map(|batch| { + let column = batch + .column(index) + .as_any() + .downcast_ref::() + .expect("a string column"); + (0..column.len()) + .map(|row| column.value(row).to_string()) + .collect::>() + }) + .collect() +} + +/// The value a setting currently holds, read the way a client would. +async fn setting(rt: &TestRuntime, name: &str) -> String { + let rows = rt + .sql(&format!( + "SELECT value FROM information_schema.df_settings WHERE name = '{name}'" + )) + .await; + scalar_str(&rows) +} + +/// A non-super-user identity. Beacon's super-user is a single configured +/// credential, so any other authenticated principal is non-super by construction. +fn regular_user() -> AuthIdentity { + AuthIdentity { + username: "reader".to_string(), + roles: vec![], + is_super_user: false, + } +} + +/// The whole point of moving the settings onto a `ConfigExtension`: a `SET` has to +/// change what the *next* statement sees, on the one shared session. +#[tokio::test(flavor = "multi_thread")] +async fn set_changes_what_a_later_statement_reads() { + let rt = common::runtime("set-live").await; + + assert_eq!(setting(&rt, "beacon.default_table").await, "default"); + + rt.sql("SET beacon.default_table = 'observations'").await; + assert_eq!(setting(&rt, "beacon.default_table").await, "observations"); + + rt.sql("SET beacon.sql.stream_coalesce.target_rows = 1024") + .await; + assert_eq!( + setting(&rt, "beacon.sql.stream_coalesce.target_rows").await, + "1024" + ); +} + +/// The `beacon.` prefix reaches DataFusion's own options — the alias this rewrite +/// exists for. `SHOW` has to agree with `SET`, since both resolve a name through +/// the same path. +#[tokio::test(flavor = "multi_thread")] +async fn the_beacon_prefix_reaches_datafusion_options() { + let rt = common::runtime("prefix-alias").await; + + rt.sql("SET beacon.execution.batch_size = 8192").await; + assert_eq!( + setting(&rt, "datafusion.execution.batch_size").await, + "8192" + ); + + // Both spellings of `SHOW` name the same option. + let via_beacon = rt.sql("SHOW beacon.execution.batch_size").await; + let via_datafusion = rt.sql("SHOW datafusion.execution.batch_size").await; + assert_eq!(scalar_str(&via_beacon), "datafusion.execution.batch_size"); + assert_eq!( + scalar_str(&via_datafusion), + "datafusion.execution.batch_size" + ); + + // The documented `BEACON_BATCH_SIZE` spelling lands in the same option. + rt.sql("SET beacon.batch_size = 4096").await; + assert_eq!( + setting(&rt, "datafusion.execution.batch_size").await, + "4096" + ); + + // …and `datafusion.*` keeps working unchanged. + rt.sql("SET datafusion.execution.batch_size = 2048").await; + assert_eq!( + setting(&rt, "datafusion.execution.batch_size").await, + "2048" + ); +} + +/// A beacon setting has to appear in `information_schema.df_settings` under its +/// fully qualified name — which is also what makes `SHOW ` resolve at all. +#[tokio::test(flavor = "multi_thread")] +async fn beacon_settings_are_visible_to_show() { + let rt = common::runtime("show-key").await; + + rt.sql("SET beacon.netcdf.use_rust_reader = true").await; + assert_eq!(setting(&rt, "beacon.netcdf.use_rust_reader").await, "true"); + + let rows = rt.sql("SHOW beacon.netcdf.use_rust_reader").await; + assert_eq!(scalar_str(&rows), "beacon.netcdf.use_rust_reader"); + + let count = rt + .sql("SELECT count(*) FROM information_schema.df_settings WHERE name LIKE 'beacon.%'") + .await; + assert!( + common::scalar_i64(&count) > 10, + "the whole namespace should be listed" + ); +} + +/// `RESET` restores the value the runtime *booted* with, not DataFusion's compiled +/// default — the reason beacon intercepts `RESET` rather than delegating it. +#[tokio::test(flavor = "multi_thread")] +async fn reset_restores_the_runtimes_own_default() { + // A runtime whose batch size differs from DataFusion's compiled default. + let rt = common::runtime_with("reset-boot", |builder| builder.with_batch_size(12_345)).await; + + rt.sql("SET beacon.batch_size = 999").await; + assert_eq!(setting(&rt, "datafusion.execution.batch_size").await, "999"); + + rt.sql("RESET beacon.batch_size").await; + assert_eq!( + setting(&rt, "datafusion.execution.batch_size").await, + "12345", + "RESET must restore the runtime's configured value, not DataFusion's default" + ); + + // The same holds for a beacon-native setting. + rt.sql("SET beacon.netcdf.use_rust_reader = true").await; + rt.sql("RESET beacon.netcdf.use_rust_reader").await; + assert_eq!(setting(&rt, "beacon.netcdf.use_rust_reader").await, "false"); +} + +/// A startup-only key would look like it worked and change nothing, so it is +/// refused — with the variable to edit. +#[tokio::test(flavor = "multi_thread")] +async fn a_startup_only_setting_is_refused() { + let rt = common::runtime("startup-only").await; + + let error = rt + .try_sql("SET beacon.port = 1234") + .await + .expect_err("a startup-only key must be refused") + .to_string(); + assert!(error.contains("BEACON_PORT"), "unhelpful error: {error}"); + + let error = rt + .try_sql("SET beacon.nonsense = 1") + .await + .expect_err("an unknown key must be refused") + .to_string(); + assert!(error.contains("SHOW SETTINGS"), "unhelpful error: {error}"); +} + +/// `ALTER SYSTEM SET` is the persistent half: it applies now *and* replays at the +/// next boot, which is what an operator on Docker or Kubernetes needs. +#[tokio::test(flavor = "multi_thread")] +async fn alter_system_survives_a_restart() { + let rt = common::restartable_runtime("alter-system", |b| b).await; + + rt.sql("ALTER SYSTEM SET beacon.default_table = 'observations'") + .await; + rt.sql("ALTER SYSTEM SET beacon.netcdf.use_rust_reader = 'true'") + .await; + + // Applied to the live session straight away. + assert_eq!(setting(&rt, "beacon.default_table").await, "observations"); + + let rt = rt.restart().await; + + assert_eq!( + setting(&rt, "beacon.default_table").await, + "observations", + "an ALTER SYSTEM value must outlive a restart" + ); + assert_eq!(setting(&rt, "beacon.netcdf.use_rust_reader").await, "true"); +} + +/// After a restart, a persisted value *is* what the server started with, so a +/// plain `RESET` has to return to it — not skip past it to the environment's +/// value, which the next restart would immediately override again. +#[tokio::test(flavor = "multi_thread")] +async fn reset_returns_to_a_persisted_value_after_a_restart() { + let rt = common::restartable_runtime("reset-vs-persisted", |b| b).await; + + rt.sql("ALTER SYSTEM SET beacon.default_table = 'observations'") + .await; + let rt = rt.restart().await; + + rt.sql("SET beacon.default_table = 'scratch'").await; + assert_eq!(setting(&rt, "beacon.default_table").await, "scratch"); + + rt.sql("RESET beacon.default_table").await; + assert_eq!( + setting(&rt, "beacon.default_table").await, + "observations", + "RESET must return to the startup state, which includes the persisted value" + ); +} + +/// `ALTER SYSTEM RESET` drops the persisted value, so the next boot goes back to +/// what the environment supplied. +#[tokio::test(flavor = "multi_thread")] +async fn alter_system_reset_forgets_the_persisted_value() { + let rt = common::restartable_runtime("alter-system-reset", |b| b).await; + + rt.sql("ALTER SYSTEM SET beacon.default_table = 'observations'") + .await; + let rt = rt.restart().await; + assert_eq!(setting(&rt, "beacon.default_table").await, "observations"); + + rt.sql("ALTER SYSTEM RESET beacon.default_table").await; + // Restored live… + assert_eq!(setting(&rt, "beacon.default_table").await, "default"); + // …and no longer replayed at boot. + let rt = rt.restart().await; + assert_eq!(setting(&rt, "beacon.default_table").await, "default"); +} + +/// A plain `SET` is live-only. Without this the two statements would be the same +/// thing and the split would be pointless. +#[tokio::test(flavor = "multi_thread")] +async fn a_plain_set_does_not_survive_a_restart() { + let rt = common::restartable_runtime("set-not-persisted", |b| b).await; + + rt.sql("SET beacon.default_table = 'observations'").await; + assert_eq!(setting(&rt, "beacon.default_table").await, "observations"); + + let rt = rt.restart().await; + assert_eq!(setting(&rt, "beacon.default_table").await, "default"); +} + +/// An in-memory runtime has nowhere to persist to, and has to say so rather than +/// accept a value it would silently lose. +#[tokio::test(flavor = "multi_thread")] +async fn alter_system_refuses_an_in_memory_runtime() { + let rt = common::runtime("alter-system-in-memory").await; + + let error = rt + .try_sql("ALTER SYSTEM SET beacon.default_table = 'observations'") + .await + .expect_err("an in-memory runtime cannot persist") + .to_string(); + assert!(error.contains("in-memory"), "unhelpful error: {error}"); + + // …and the refusal leaves nothing behind. A statement that reported failure + // must not have changed the session on its way to the error. + assert_eq!(setting(&rt, "beacon.default_table").await, "default"); +} + +/// `SHOW SETTINGS` documents the engine, so any authenticated caller can read it — +/// the issue's "a user cannot discover which settings exist". Changing one stays +/// super-user-only. +#[tokio::test(flavor = "multi_thread")] +async fn show_settings_is_readable_but_set_is_not() { + let rt = common::runtime("settings-privileges").await; + + let rows = rt.sql_as("SHOW SETTINGS", regular_user()).await; + let names = column_strings(&rows, 0); + assert!(names.iter().any(|name| name == "beacon.default_table")); + assert!( + names.iter().all(|name| name.starts_with("beacon.")), + "SHOW SETTINGS must expose only the beacon namespace" + ); + + for sql in [ + "SET beacon.default_table = 'observations'", + "ALTER SYSTEM SET beacon.default_table = 'observations'", + // The table form lives in `beacon.system`, which is super-user-only. + "SELECT * FROM beacon.system.settings", + ] { + let error = rt + .try_sql_as(sql, regular_user()) + .await + .err() + .unwrap_or_else(|| panic!("`{sql}` must be refused for a regular user")) + .to_string(); + assert!( + error.contains("permitted") || error.contains("permission"), + "`{sql}` should fail as a privilege error, got: {error}" + ); + } + + // None of the refused statements changed anything. + assert_eq!(setting(&rt, "beacon.default_table").await, "default"); +} + +/// `SHOW SETTINGS` reports the live value and the one a `RESET` would restore, so +/// an operator can see both without running the reset. +#[tokio::test(flavor = "multi_thread")] +async fn show_settings_reports_the_value_and_the_boot_default() { + let rt = common::runtime("settings-columns").await; + + rt.sql("SET beacon.default_table = 'observations'").await; + + let rows = rt + .sql( + "SELECT value, \"default\" FROM beacon.system.settings \ + WHERE name = 'beacon.default_table'", + ) + .await; + assert_eq!(column_strings(&rows, 0), vec!["observations".to_string()]); + assert_eq!(column_strings(&rows, 1), vec!["default".to_string()]); +} + +/// Beacon's new statements must not shadow the SQL they resemble. `ALTER TABLE` +/// and `SET ` predate this feature and have to keep working. +#[tokio::test(flavor = "multi_thread")] +async fn the_new_statements_do_not_shadow_existing_sql() { + let rt = common::runtime("no-shadowing").await; + + rt.sql("SET datafusion.execution.batch_size = 8192").await; + rt.sql("SET timezone = 'UTC'").await; + rt.sql("SHOW TABLES").await; + + // `ALTER TABLE` on a missing table must fail as a *table* error, proving the + // `ALTER SYSTEM` peek did not swallow it. + let error = rt + .try_sql("ALTER TABLE nope ADD COLUMN x INT") + .await + .expect_err("the table does not exist") + .to_string(); + assert!( + !error.contains("SHOW SETTINGS") && !error.contains("startup"), + "ALTER TABLE was mistaken for ALTER SYSTEM: {error}" + ); +} diff --git a/beacon-db/beacon-core/tests/system_schema.rs b/beacon-db/beacon-core/tests/system_schema.rs index eb8e3ddf..d090e500 100644 --- a/beacon-db/beacon-core/tests/system_schema.rs +++ b/beacon-db/beacon-core/tests/system_schema.rs @@ -47,6 +47,7 @@ async fn system_tables_are_listed_in_information_schema() { "file_stats_segments", "query_metrics", "roles", + "settings", "users" ], "the system schema should expose exactly these tables" diff --git a/beacon-db/beacon-datafusion-ext/src/lib.rs b/beacon-db/beacon-datafusion-ext/src/lib.rs index 2ba75c44..0510023a 100644 --- a/beacon-db/beacon-datafusion-ext/src/lib.rs +++ b/beacon-db/beacon-datafusion-ext/src/lib.rs @@ -10,6 +10,7 @@ pub mod nd; pub mod object_store_registry; pub mod remote; pub mod secrets; +pub mod settings; pub mod stats_cache; pub mod table_ext; pub mod type_widening; diff --git a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs index 7cee86e4..c1f9c110 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/mod.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/mod.rs @@ -53,6 +53,19 @@ mod tests { use super::exec::{NdBroadcastExec, NdSourceExec}; use super::*; + /// Config with the nd rewrite switched on. Both rules read + /// `beacon.enable_nd_pipeline` on every plan (they are always installed, so a + /// later `SET` can turn them on), and it is off by default — so a bare + /// `ConfigOptions` would make every rule here a no-op. + fn nd_pipeline_on() -> datafusion::common::config::ConfigOptions { + let mut options = datafusion::common::config::ConfigOptions::default(); + options.extensions.insert(crate::settings::BeaconOptions { + enable_nd_pipeline: true, + ..Default::default() + }); + options + } + fn dims(spec: &[(&str, usize)]) -> Dimensions { Dimensions::try_new( spec.iter() @@ -421,7 +434,7 @@ mod tests { let expected = run(original.clone()).await.unwrap(); let optimized = NdProjectionPushdown::new() - .optimize(original, &ConfigOptions::default()) + .optimize(original, &nd_pipeline_on()) .unwrap(); // Schema is preserved (the rule reports schema_check = true). @@ -442,6 +455,56 @@ mod tests { assert_eq!(actual, expected); } + /// Both rules are installed unconditionally so a later + /// `SET beacon.enable_nd_pipeline = true` can switch them on — which only + /// works if a rule left off returns the plan exactly as it found it. + #[tokio::test] + async fn rules_are_inert_when_the_nd_pipeline_is_off() { + use datafusion::common::config::ConfigOptions; + use datafusion::physical_optimizer::PhysicalOptimizerRule; + use datafusion::physical_plan::displayable; + + let schema = test_source().schema(); + let exprs = projection_exprs(&schema); + let original: Arc = Arc::new( + ProjectionExec::try_new( + exprs + .iter() + .cloned() + .map(|(expr, alias)| ProjectionExpr { expr, alias }), + Arc::new(NdBroadcastExec::try_new(test_source()).unwrap()), + ) + .unwrap(), + ); + let before = displayable(original.as_ref()).indent(true).to_string(); + + // A default `ConfigOptions` carries no beacon namespace at all, which is + // the other way the flag reads as off. + for config in [ConfigOptions::default(), { + let mut off = ConfigOptions::default(); + off.extensions.insert(crate::settings::BeaconOptions::default()); + off + }] { + let untouched = NdProjectionPushdown::new() + .optimize(original.clone(), &config) + .unwrap(); + assert_eq!( + displayable(untouched.as_ref()).indent(true).to_string(), + before, + "the projection rule rewrote a plan with the nd pipeline off" + ); + + let untouched = NdFilterPushdown::new() + .optimize(original.clone(), &config) + .unwrap(); + assert_eq!( + displayable(untouched.as_ref()).indent(true).to_string(), + before, + "the filter rule rewrote a plan with the nd pipeline off" + ); + } + } + /// The rule leaves a projection in place when any expression is not /// element-wise (here a volatile scalar function): no `NdProjectionExec`. #[tokio::test] @@ -500,7 +563,7 @@ mod tests { ); let optimized = NdProjectionPushdown::new() - .optimize(original, &ConfigOptions::default()) + .optimize(original, &nd_pipeline_on()) .unwrap(); let rendered = displayable(optimized.as_ref()).indent(true).to_string(); assert!( @@ -747,7 +810,7 @@ mod tests { let expected = run(original.clone()).await.unwrap(); let optimized = NdFilterPushdown::new() - .optimize(original, &ConfigOptions::default()) + .optimize(original, &nd_pipeline_on()) .unwrap(); assert_eq!(optimized.schema(), original_schema); @@ -837,7 +900,7 @@ mod tests { ); let optimized = NdFilterPushdown::new() - .optimize(original, &ConfigOptions::default()) + .optimize(original, &nd_pipeline_on()) .unwrap(); let rendered = displayable(optimized.as_ref()).indent(true).to_string(); diff --git a/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs b/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs index 2b16818e..ffe3d508 100644 --- a/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs +++ b/beacon-db/beacon-datafusion-ext/src/nd/optimizer.rs @@ -35,6 +35,21 @@ use datafusion::physical_plan::projection::ProjectionExec; use datafusion::logical_expr::Volatility; use super::exec::{NdBroadcastExec, NdFilterExec, NdProjectionExec}; +use crate::settings::BeaconOptions; + +/// Whether the node-rewriting nd optimization is on for this plan. +/// +/// Both rules are always installed and check here, rather than being installed +/// conditionally at startup: the optimizer chain is fixed when the session state +/// is built, so a rule that is absent can never be switched on again, and +/// `SET beacon.enable_nd_pipeline = true` would need a restart to take effect. +/// The base nd pipeline runs either way; this only gates the rewrite. +fn nd_pipeline_enabled(config: &ConfigOptions) -> bool { + config + .extensions + .get::() + .is_some_and(|options| options.enable_nd_pipeline) +} /// Sinks element-wise `ProjectionExec`s below an [`NdBroadcastExec`] into an /// [`NdProjectionExec`], so they evaluate before broadcasting. @@ -51,8 +66,11 @@ impl PhysicalOptimizerRule for NdProjectionPushdown { fn optimize( &self, plan: Arc, - _config: &ConfigOptions, + config: &ConfigOptions, ) -> Result> { + if !nd_pipeline_enabled(config) { + return Ok(plan); + } plan.transform_down(|node| { let Some(projection) = node.as_any().downcast_ref::() else { return Ok(Transformed::no(node)); @@ -138,8 +156,11 @@ impl PhysicalOptimizerRule for NdFilterPushdown { fn optimize( &self, plan: Arc, - _config: &ConfigOptions, + config: &ConfigOptions, ) -> Result> { + if !nd_pipeline_enabled(config) { + return Ok(plan); + } plan.transform_down(|node| { let Some(filter) = node.as_any().downcast_ref::() else { return Ok(Transformed::no(node)); diff --git a/beacon-db/beacon-datafusion-ext/src/settings.rs b/beacon-db/beacon-datafusion-ext/src/settings.rs new file mode 100644 index 00000000..ae19639e --- /dev/null +++ b/beacon-db/beacon-datafusion-ext/src/settings.rs @@ -0,0 +1,640 @@ +//! The `beacon.*` configuration namespace. +//! +//! [`BeaconOptions`] is a DataFusion [`ConfigExtension`], which is what makes +//! `SET beacon.netcdf.use_rust_reader = true` reach beacon at all: DataFusion's +//! `SET` writes `ConfigOptions`, and `ConfigOptions` routes a namespaced key to +//! the extension registered under its prefix. A value set this way is visible to +//! every later query, because plan- and execution-time code reads the options off +//! the session it is handed rather than off a snapshot taken at startup. +//! +//! This is deliberately *not* `SessionConfig::with_extension`, the `TypeId`-keyed +//! map beacon used before. That map is invisible to `SET`, so a setting published +//! there could only ever be read, never changed. +//! +//! # Scope +//! +//! A `SET` applies to the whole server, not to one client: beacon runs one shared +//! `SessionContext` for every transport and every user. That is why `SET` is +//! super-user-only (`validate_query_plan`), exactly as `SET datafusion.*` already +//! was. +//! +//! # Layers +//! +//! A format setting has three layers, narrowest last: +//! +//! 1. the runtime default, from the `BEACON_*` environment variable, +//! 2. this namespace, changed with `SET`, +//! 3. the per-table `CREATE EXTERNAL TABLE ... OPTIONS (...)` override. + +use std::any::Any; +use std::collections::HashMap; + +use datafusion::catalog::Session; +use datafusion::common::config::{ + ConfigEntry, ConfigExtension, ConfigField, ConfigOptions, ExtensionOptions, Visit, +}; +use datafusion::common::config_namespace; +use datafusion::error::Result as DFResult; +use datafusion::execution::context::SessionConfig; + +/// The namespace every setting in this module is addressed under. +pub const BEACON_PREFIX: &str = "beacon"; + +config_namespace! { + /// Result-stream coalescing: the small record batches a plan emits are merged + /// into client-sized ones before they leave the server. + pub struct StreamCoalesceOptions { + /// Whether to coalesce at all. Disabled passes batches through untouched. + pub enabled: bool, default = true + + /// Buffer batches until at least this many rows have accumulated. + pub target_rows: usize, default = 64 * 1024 + + /// Flush a non-empty buffer after this long, even below `target_rows`, so a + /// slow-producing plan stays responsive. `0` disables the timeout. + pub flush_timeout_ms: u64, default = 25 + + /// Hard upper bound on a buffered batch, so one oversized input batch cannot + /// grow the buffer without limit. + pub max_rows: usize, default = 256 * 1024 + } +} + +config_namespace! { + /// How a client query is compiled and how its result is streamed back. + pub struct SqlOptions { + /// Result-stream coalescing. + pub stream_coalesce: StreamCoalesceOptions, default = Default::default() + } +} + +config_namespace! { + /// NetCDF reader settings. Each is also a per-table `OPTIONS (...)` key. + pub struct NetcdfOptions { + /// Whether reads consult the shared reader cache. The cache's *capacity* is + /// fixed when the runtime starts (`BEACON_NETCDF_READER_CACHE_SIZE`). + pub use_reader_cache: bool, default = true + + /// Whether to compute per-file statistics during planning, used to prune the + /// files a query cannot match. Needs the pure-Rust reader. + pub enable_statistics: bool, default = true + + /// Whether reads go through the pure-Rust reader instead of netcdf-c. + pub use_rust_reader: bool, default = false + } +} + +config_namespace! { + /// HDF5 reader settings. Each is also a per-table `OPTIONS (...)` key. + pub struct Hdf5Options { + /// Whether reads consult the shared reader cache. Only the pure-Rust reader + /// has a cache of its own; under netcdf-c the netCDF cache applies instead. + pub use_reader_cache: bool, default = true + + /// Whether to compute per-file statistics during planning. Needs the + /// pure-Rust reader. + pub enable_statistics: bool, default = true + + /// Whether reads go through the pure-Rust reader instead of netcdf-c. + pub use_rust_reader: bool, default = false + } +} + +config_namespace! { + /// Zarr reader settings. + pub struct ZarrOptions { + /// Whether to compute per-store statistics during planning. A store answers + /// from its metadata where it can, and otherwise reads only its rank-0 and + /// rank-1 arrays. + pub enable_statistics: bool, default = true + } +} + +config_namespace! { + /// Atlas reader settings. Each is also a per-table `OPTIONS (...)` key. + pub struct AtlasOptions { + /// Whether reads consult the shared reader cache. The cache's *capacity* is + /// fixed when the runtime starts (`BEACON_ATLAS_READER_CACHE_SIZE`). + pub use_reader_cache: bool, default = true + + /// Whether a predicate scan drops the datasets that cannot match before + /// reading them. A pure optimization. + pub use_pruning: bool, default = true + } +} + +config_namespace! { + /// Beacon Binary Format settings. + pub struct BbfOptions { + /// Whether to split each record batch into `batch_size`-row slices, which + /// bounds peak memory on a wide table. + pub split_streams_slice: bool, default = false + } +} + +config_namespace! { + /// Managed-Lance settings. An empty value means "leave it to Lance", which is + /// what an unset `BEACON_LANCE_*` variable meant. + /// + /// The first four apply when beacon *writes* a Lance table, so they change the + /// files a later `CREATE TABLE`/`INSERT` produces, never the ones already on + /// disk. `materialization` is a read setting and applies to the next scan. + pub struct LanceOptions { + /// Block compression for string columns: `fsst`, `zstd`, `lz4`, or `none`. + pub compression: String, default = String::new() + + /// Block compression for numeric columns: `zstd`, `lz4`, or `none`. `none` + /// also disables bitpacking and RLE, which usually measures *larger*. + pub numeric_compression: String, default = String::new() + + /// Lance file format version: `2.0`, `2.1`, or `2.2`. + pub version: String, default = String::new() + + /// Minichunk size in bytes. Needs `version` = `2.2` to take effect. + pub minichunk: String, default = String::new() + + /// Column materialization on a scan: `late` or `early`. + pub materialization: String, default = String::new() + } +} + +config_namespace! { + /// Every runtime-settable beacon setting. + /// + /// Registered on the session as a [`ConfigExtension`], so `SET beacon.x = y`, + /// `RESET beacon.x`, `SHOW beacon.x` and `information_schema.df_settings` all + /// work against it with no further wiring. + pub struct BeaconOptions { + /// The table a JSON query without a `from` resolves against. SQL always + /// names its own source, so this only affects the JSON query API. + pub default_table: String, default = "default".to_string() + + /// Whether the JSON query compiler pushes the selected columns into the scan. + pub enable_pushdown_projection: bool, default = true + + /// Whether the N-dimensional pipeline optimizer sinks element-wise + /// projections and filters below the grid broadcast. The base nd pipeline + /// always runs; this only enables the node-rewriting optimization. + pub enable_nd_pipeline: bool, default = false + + /// How client queries are compiled and streamed. + pub sql: SqlOptions, default = Default::default() + + /// NetCDF reader settings. + pub netcdf: NetcdfOptions, default = Default::default() + + /// HDF5 reader settings. + pub hdf5: Hdf5Options, default = Default::default() + + /// Zarr reader settings. + pub zarr: ZarrOptions, default = Default::default() + + /// Atlas reader settings. + pub atlas: AtlasOptions, default = Default::default() + + /// Beacon Binary Format settings. + pub bbf: BbfOptions, default = Default::default() + + /// Managed-Lance settings. + pub lance: LanceOptions, default = Default::default() + } +} + +impl ConfigExtension for BeaconOptions { + const PREFIX: &'static str = BEACON_PREFIX; +} + +impl ExtensionOptions for BeaconOptions { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn cloned(&self) -> Box { + Box::new(self.clone()) + } + + /// The key arrives without the namespace: DataFusion splits `beacon.netcdf.x` + /// once and hands the extension `netcdf.x`. + fn set(&mut self, key: &str, value: &str) -> DFResult<()> { + ConfigField::set(self, key, value) + } + + /// Emits **fully qualified** keys (`beacon.netcdf.use_rust_reader`). + /// + /// This is why the trait is written out rather than generated by + /// `extensions_options!`, whose `entries()` emits the bare field name. The + /// qualified form is what `information_schema.df_settings` displays, and what + /// `SHOW ` validates a name against — a bare `use_rust_reader` would make + /// `SHOW beacon.netcdf.use_rust_reader` fail as an unknown variable. + fn entries(&self) -> Vec { + struct Collector(Vec); + + impl Visit for Collector { + fn some( + &mut self, + key: &str, + value: V, + description: &'static str, + ) { + self.0.push(ConfigEntry { + key: key.to_string(), + value: Some(value.to_string()), + description, + }); + } + + fn none(&mut self, key: &str, description: &'static str) { + self.0.push(ConfigEntry { + key: key.to_string(), + value: None, + description, + }); + } + } + + let mut collector = Collector(Vec::new()); + self.visit(&mut collector, BEACON_PREFIX, ""); + collector.0 + } +} + +impl BeaconOptions { + /// The options published on `config`, or `None` when the namespace is absent. + /// + /// Absent means the session was not built by beacon's runtime — a bare + /// `SessionContext` in a unit test, or an embedder wiring a format factory up + /// by hand. A caller that holds its own configuration should fall back to it + /// rather than to the compiled defaults, which is why this is separate from + /// [`Self::from_config`]. + pub fn try_from_config(config: &SessionConfig) -> Option { + config.options().extensions.get::().cloned() + } + + /// [`Self::try_from_config`] for the `&dyn Session` a `TableProvider` or + /// `FileFormatFactory` is handed, where no `SessionContext` is in reach. + pub fn try_from_session(session: &dyn Session) -> Option { + Self::try_from_config(session.config()) + } + + /// The options published on `config`, or the compiled defaults when the + /// namespace is absent. + pub fn from_config(config: &SessionConfig) -> Self { + Self::try_from_config(config).unwrap_or_default() + } + + /// [`Self::from_config`] for the `&dyn Session` a `TableProvider` or + /// `FileFormatFactory` is handed, where no `SessionContext` is in reach. + pub fn from_session(session: &dyn Session) -> Self { + Self::from_config(session.config()) + } + + /// Whether `key` names a setting in this namespace, with or without the + /// `beacon.` prefix. + pub fn has_key(key: &str) -> bool { + let qualified = match key.strip_prefix("beacon.") { + Some(_) => key.to_string(), + None => format!("{BEACON_PREFIX}.{key}"), + }; + Self::default() + .entries() + .iter() + .any(|entry| entry.key == qualified) + } + + /// Every key in this namespace, fully qualified and sorted. + pub fn keys() -> Vec { + let mut keys: Vec = Self::default() + .entries() + .into_iter() + .map(|entry| entry.key) + .collect(); + keys.sort(); + keys + } +} + +/// The value every setting held when the runtime started, before any `SET`. +/// +/// `RESET beacon.x` restores from here rather than from DataFusion's +/// `ConfigOptions::reset`, which would reinstate DataFusion's *compiled* default +/// and silently discard the value the operator's environment supplied. Published +/// as a plain typed session extension, since it never changes after startup. +/// +/// Two layers, because the two `RESET` statements return to different places: +/// +/// * `startup` — what the server actually came up with: the environment, then any +/// `ALTER SYSTEM SET` value replayed over it. Plain `RESET` restores this. +/// * `environment` — the same snapshot *before* the replay. `ALTER SYSTEM RESET` +/// restores this, since it is deleting the persisted value and must not put it +/// straight back. +/// +/// On a runtime with nothing persisted the two are identical. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct BootSettings { + startup: HashMap, + environment: HashMap, +} + +impl BootSettings { + /// Snapshots every set option, beacon's namespace and DataFusion's alike, + /// as both layers — the state before any persisted value is replayed. + /// + /// An option with no value (an unset `Option`) is skipped: there is no + /// string that would restore it, so those fall back to DataFusion's own + /// `RESET`. + pub fn capture(options: &ConfigOptions) -> Self { + let values = Self::values(options); + Self { + startup: values.clone(), + environment: values, + } + } + + /// This snapshot with `startup` re-taken from `options`, keeping the original + /// `environment` layer. Called once, after the persisted settings are + /// replayed, so plain `RESET` returns to the state the server came up with. + pub fn with_startup(&self, options: &ConfigOptions) -> Self { + Self { + startup: Self::values(options), + environment: self.environment.clone(), + } + } + + fn values(options: &ConfigOptions) -> HashMap { + options + .entries() + .into_iter() + .filter_map(|entry| entry.value.map(|value| (entry.key, value))) + .collect() + } + + /// The value of `key` before any persisted override — what an + /// `ALTER SYSTEM RESET` returns to. + pub fn environment(&self, key: &str) -> Option<&str> { + self.environment.get(key).map(String::as_str) + } + + /// The value of `key` the server actually started with — what a plain `RESET` + /// returns to. + pub fn get(&self, key: &str) -> Option<&str> { + self.startup.get(key).map(String::as_str) + } + + /// The options published on `config`, or an empty snapshot for a session + /// beacon did not build. + pub fn from_config(config: &SessionConfig) -> Self { + config + .get_extension::() + .map(|boot| (*boot).clone()) + .unwrap_or_default() + } +} + +/// The `BEACON_*` variable behind a `beacon.*` key that can only be set at +/// startup, or `None` when the key is not one of those. +/// +/// Every one of these decides something built once — a socket, a directory, a +/// credential, a thread pool, a cache's capacity — so a `SET` would appear to +/// work and change nothing. Rejecting it by name, with the variable to edit, is +/// the useful answer. +pub fn startup_only_env_var(key: &str) -> Option<&'static str> { + let key = key.strip_prefix("beacon.").unwrap_or(key); + let var = match key { + "port" => "BEACON_PORT", + "host" => "BEACON_HOST", + "worker_threads" => "BEACON_WORKER_THREADS", + "log_level" => "BEACON_LOG_LEVEL", + "base_path" => "BEACON_BASE_PATH", + "web_ui_dir" => "BEACON_WEB_UI_DIR", + "max_upload_bytes" => "BEACON_MAX_UPLOAD_BYTES", + "vm_memory_size" => "BEACON_VM_MEMORY_SIZE", + "enable_sql" => "BEACON_ENABLE_SQL", + "enable_sys_info" => "BEACON_ENABLE_SYS_INFO", + "data_dir" => "BEACON_DATA_DIR", + "secrets_key" => "BEACON_SECRETS_KEY", + "stats_cache_capacity" => "BEACON_STATS_CACHE_CAPACITY", + "netcdf.reader_cache_size" => "BEACON_NETCDF_READER_CACHE_SIZE", + "hdf5.reader_cache_size" => "BEACON_HDF5_READER_CACHE_SIZE", + "atlas.reader_cache_size" => "BEACON_ATLAS_READER_CACHE_SIZE", + "admin.username" => "BEACON_ADMIN_USERNAME", + "admin.password" => "BEACON_ADMIN_PASSWORD", + "auth.enforce" => "BEACON_AUTH_ENFORCE", + "auth.anonymous_enabled" => "BEACON_AUTH_ANONYMOUS_ENABLED", + "crawler.enable" => "BEACON_CRAWLER_ENABLE", + "crawler.default_interval_secs" => "BEACON_CRAWLER_DEFAULT_INTERVAL_SECS", + "file_stats.enable" => "BEACON_FILE_STATS_ENABLE", + "file_stats.interval_secs" => "BEACON_FILE_STATS_INTERVAL_SECS", + "file_stats.on_startup" => "BEACON_FILE_STATS_ON_STARTUP", + "file_stats.concurrency" => "BEACON_FILE_STATS_CONCURRENCY", + "file_stats.batch_files" => "BEACON_FILE_STATS_BATCH_FILES", + other => return startup_only_family(other), + }; + Some(var) +} + +/// The startup-only keys that share a prefix, matched as a family so a whole +/// group answers with one entry rather than twenty. +fn startup_only_family(key: &str) -> Option<&'static str> { + let (prefix, _) = key.split_once('.')?; + let var = match prefix { + "s3" => "BEACON_S3_*", + "oidc" => "BEACON_OIDC_*", + "flight_sql" => "BEACON_FLIGHT_SQL_*", + "cors" => "BEACON_CORS_*", + "api" => "BEACON_API_*", + "file_stats" => "BEACON_FILE_STATS_*", + "crawler" => "BEACON_CRAWLER_*", + _ => return None, + }; + Some(var) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Every key must be addressable by the exact name `entries()` advertises. + /// `SHOW ` validates against that list and `df_settings` displays it, so + /// a key that reports one name and accepts another is broken in both. + #[test] + fn every_advertised_key_round_trips() { + let mut options = BeaconOptions::default(); + for key in BeaconOptions::keys() { + let unqualified = key + .strip_prefix("beacon.") + .unwrap_or_else(|| panic!("`{key}` is not in the beacon namespace")); + // A value every field type accepts: `bool` takes "true", the numeric + // fields would not, so probe with the type's own current rendering. + let current = options + .entries() + .into_iter() + .find(|entry| entry.key == key) + .and_then(|entry| entry.value) + .unwrap_or_default(); + ExtensionOptions::set(&mut options, unqualified, ¤t) + .unwrap_or_else(|e| panic!("`{key}` is advertised but not settable: {e}")); + } + } + + #[test] + fn keys_are_fully_qualified_and_cover_every_namespace() { + let keys = BeaconOptions::keys(); + assert!(keys.iter().all(|key| key.starts_with("beacon."))); + for expected in [ + "beacon.default_table", + "beacon.enable_pushdown_projection", + "beacon.enable_nd_pipeline", + "beacon.sql.stream_coalesce.target_rows", + "beacon.netcdf.use_rust_reader", + "beacon.hdf5.use_rust_reader", + "beacon.zarr.enable_statistics", + "beacon.atlas.use_pruning", + "beacon.bbf.split_streams_slice", + "beacon.lance.materialization", + ] { + assert!(keys.iter().any(|key| key == expected), "missing {expected}"); + } + } + + #[test] + fn set_changes_the_value_entries_reports() { + let mut options = BeaconOptions::default(); + assert!(!options.netcdf.use_rust_reader); + + ExtensionOptions::set(&mut options, "netcdf.use_rust_reader", "true").unwrap(); + assert!(options.netcdf.use_rust_reader); + + let entry = options + .entries() + .into_iter() + .find(|entry| entry.key == "beacon.netcdf.use_rust_reader") + .expect("key is advertised"); + assert_eq!(entry.value.as_deref(), Some("true")); + } + + #[test] + fn unknown_key_is_rejected() { + let mut options = BeaconOptions::default(); + let err = ExtensionOptions::set(&mut options, "netcdf.nope", "true").unwrap_err(); + assert!(err.to_string().contains("nope"), "unhelpful error: {err}"); + } + + /// A `SET` writes through `ConfigOptions`, so the extension has to be reachable + /// by its prefix — the whole point of registering it as a `ConfigExtension`. + #[test] + fn config_options_routes_the_beacon_prefix() { + let mut options = ConfigOptions::default(); + options.extensions.insert(BeaconOptions::default()); + + options + .set("beacon.sql.stream_coalesce.target_rows", "1024") + .unwrap(); + + let beacon = options.extensions.get::().unwrap(); + assert_eq!(beacon.sql.stream_coalesce.target_rows, 1024); + } + + /// `SHOW ALL` and `information_schema.df_settings` both read `entries()`, so the + /// beacon keys have to show up there alongside DataFusion's own. + #[test] + fn beacon_keys_appear_in_config_options_entries() { + let mut options = ConfigOptions::default(); + options.extensions.insert(BeaconOptions::default()); + + let keys: Vec = options.entries().into_iter().map(|e| e.key).collect(); + assert!(keys.iter().any(|k| k == "beacon.default_table")); + assert!(keys.iter().any(|k| k == "datafusion.execution.batch_size")); + } + + #[test] + fn has_key_accepts_qualified_and_bare_names() { + assert!(BeaconOptions::has_key("beacon.default_table")); + assert!(BeaconOptions::has_key("default_table")); + assert!(!BeaconOptions::has_key("beacon.port")); + assert!(!BeaconOptions::has_key("datafusion.execution.batch_size")); + } + + #[test] + fn boot_settings_record_the_values_the_runtime_started_with() { + let mut options = ConfigOptions::default(); + options.extensions.insert(BeaconOptions::default()); + options.execution.batch_size = 4096; + options.set("beacon.default_table", "observations").unwrap(); + + let boot = BootSettings::capture(&options); + + // Both namespaces are captured, so `RESET` restores an operator's + // environment value rather than DataFusion's compiled default. + assert_eq!(boot.get("beacon.default_table"), Some("observations")); + assert_eq!(boot.get("datafusion.execution.batch_size"), Some("4096")); + assert_eq!(boot.get("beacon.nope"), None); + } + + /// The two `RESET` statements return to different places, so the snapshot + /// keeps both layers. Getting this wrong is silent: `ALTER SYSTEM RESET` would + /// restore the very value it just deleted. + #[test] + fn the_snapshot_separates_the_environment_from_a_persisted_value() { + let mut options = ConfigOptions::default(); + options.extensions.insert(BeaconOptions::default()); + options.set("beacon.default_table", "from_env").unwrap(); + + // At build time both layers are the environment's. + let boot = BootSettings::capture(&options); + assert_eq!(boot.get("beacon.default_table"), Some("from_env")); + assert_eq!(boot.environment("beacon.default_table"), Some("from_env")); + + // A persisted value is replayed over it, and only `startup` moves. + options.set("beacon.default_table", "persisted").unwrap(); + let boot = boot.with_startup(&options); + assert_eq!( + boot.get("beacon.default_table"), + Some("persisted"), + "a plain RESET returns to what the server came up with" + ); + assert_eq!( + boot.environment("beacon.default_table"), + Some("from_env"), + "an ALTER SYSTEM RESET returns to the environment, not the value it is deleting" + ); + } + + #[test] + fn startup_only_keys_name_their_variable() { + assert_eq!(startup_only_env_var("beacon.port"), Some("BEACON_PORT")); + assert_eq!(startup_only_env_var("port"), Some("BEACON_PORT")); + assert_eq!( + startup_only_env_var("beacon.netcdf.reader_cache_size"), + Some("BEACON_NETCDF_READER_CACHE_SIZE") + ); + // Families answer for every member. + assert_eq!( + startup_only_env_var("beacon.s3.bucket"), + Some("BEACON_S3_*") + ); + assert_eq!( + startup_only_env_var("beacon.flight_sql.port"), + Some("BEACON_FLIGHT_SQL_*") + ); + // A settable key must not be claimed as startup-only. + assert_eq!(startup_only_env_var("beacon.netcdf.use_rust_reader"), None); + assert_eq!(startup_only_env_var("beacon.default_table"), None); + } + + /// The two tables must not overlap: a key claimed as startup-only would be + /// rejected even though the extension can set it. + #[test] + fn no_settable_key_is_also_startup_only() { + for key in BeaconOptions::keys() { + assert_eq!( + startup_only_env_var(&key), + None, + "`{key}` is settable but also listed as startup-only" + ); + } + } +} diff --git a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs index 5ee50e30..146d1482 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-atlas/src/datafusion/mod.rs @@ -12,6 +12,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use beacon_common::super_typing::super_type_schema; use beacon_datafusion_ext::format_ext::{DatasetMetadata, FileFormatFactoryExt}; +use beacon_datafusion_ext::settings::BeaconOptions; use datafusion::{ catalog::{Session, memory::DataSourceExec}, common::{GetExt, Statistics, exec_datafusion_err}, @@ -121,6 +122,27 @@ impl AtlasFormatFactory { } } + /// The settings this format starts from, for a format built on `state`. + /// + /// The runtime's `BEACON_ATLAS_*` values seed `beacon.atlas.*` at startup, so + /// reading the namespace gives the runtime default until an operator runs a + /// `SET`, and the new value afterwards. A session beacon did not build carries + /// no namespace and falls back to this factory's own config. + fn session_config(&self, state: &dyn Session) -> AtlasConfig { + let Some(atlas) = BeaconOptions::try_from_session(state).map(|beacon| beacon.atlas) else { + // No namespace: this session was not built by beacon's runtime, so + // this factory's own config is the only configuration there is. + return self.config.clone(); + }; + AtlasConfig { + use_reader_cache: atlas.use_reader_cache, + use_pruning: atlas.use_pruning, + // Not settable: the cache is built once, at the capacity the runtime + // started with. + reader_cache_size: self.config.reader_cache_size, + } + } + /// Build an [`AtlasFormat`] with the given per-table effective settings, /// wiring in the shared reader cache when caching is enabled. fn build_format( @@ -139,14 +161,16 @@ impl AtlasFormatFactory { impl FileFormatFactory for AtlasFormatFactory { fn create( &self, - _state: &dyn Session, + state: &dyn Session, format_options: &std::collections::HashMap, ) -> datafusion::error::Result> { - // Per-table overrides from `CREATE EXTERNAL TABLE ... OPTIONS (...)`, - // defaulting to the runtime config. + // Three layers, narrowest last: the runtime default, the session's + // `beacon.atlas.*` namespace (what `SET` writes), then the per-table + // `CREATE EXTERNAL TABLE ... OPTIONS (...)`. + let session_config = self.session_config(state); let mut options = self.options.clone(); - let mut use_reader_cache = self.config.use_reader_cache; - let mut use_pruning = self.config.use_pruning; + let mut use_reader_cache = session_config.use_reader_cache; + let mut use_pruning = session_config.use_pruning; if let Some(value) = format_options.get("read_dimensions") { options.read_dimensions = Some( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs index 52c6d972..5a2fb214 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-bbf/src/datafusion/mod.rs @@ -7,6 +7,7 @@ use beacon_binary_format::{ use beacon_common::file_descriptors::file_open_parallelism; use beacon_common::super_typing::super_type_schema; use beacon_datafusion_ext::format_ext::{DatasetMetadata, FileFormatFactoryExt}; +use beacon_datafusion_ext::settings::BeaconOptions; use datafusion::{ catalog::{Session, memory::DataSourceExec}, common::{GetExt, Statistics}, @@ -69,12 +70,17 @@ impl GetExt for BBFFormatFactory { impl FileFormatFactory for BBFFormatFactory { fn create( &self, - _state: &dyn Session, + state: &dyn Session, format_options: &HashMap, ) -> datafusion::error::Result> { - // Per-table override from `CREATE EXTERNAL TABLE ... OPTIONS (...)`, - // defaulting to the runtime config. - let mut split_streams_slice = self.config.split_streams_slice; + // Three layers, narrowest last: the runtime default, the session's + // `beacon.bbf.*` namespace (what `SET` writes), then the per-table + // `CREATE EXTERNAL TABLE ... OPTIONS (...)`. + // No namespace means this session was not built by beacon's runtime, so + // this factory's own config is the only configuration there is. + let mut split_streams_slice = BeaconOptions::try_from_session(state) + .map(|beacon| beacon.bbf.split_streams_slice) + .unwrap_or(self.config.split_streams_slice); if let Some(value) = format_options.get("split_streams_slice") { split_streams_slice = parse_bool_option("split_streams_slice", value)?; } diff --git a/beacon-db/beacon-file-formats/beacon-arrow-hdf5/src/format.rs b/beacon-db/beacon-file-formats/beacon-arrow-hdf5/src/format.rs index 6b02fcd6..c43b0708 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-hdf5/src/format.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-hdf5/src/format.rs @@ -15,6 +15,7 @@ use beacon_arrow_netcdf::datafusion::{statistics, NetCDFFormatFactory, NetcdfFor use beacon_common::super_typing::super_type_schema; use beacon_datafusion_ext::format_ext::{DatasetMetadata, FileFormatFactoryExt}; use beacon_datafusion_ext::listing_factory::ListingFactory; +use beacon_datafusion_ext::settings::BeaconOptions; use datafusion::{ catalog::{memory::DataSourceExec, Session}, common::{exec_datafusion_err, GetExt, Statistics}, @@ -108,23 +109,49 @@ impl Hdf5FormatFactory { /// default. pub fn uses_rust_reader( &self, + state: &dyn Session, format_options: &HashMap, ) -> datafusion::error::Result { match format_options.get("use_rust_reader") { Some(value) => parse_bool_option("use_rust_reader", value), - None => Ok(self.config.use_rust_reader), + None => Ok(self.session_config(state).use_rust_reader), } } - /// Resolve the per-table options against the runtime config. + /// The settings this format starts from, for a format built on `state`. + /// + /// The runtime's `BEACON_HDF5_*` values seed `beacon.hdf5.*` at startup, so + /// reading the namespace gives the runtime default until an operator runs a + /// `SET`, and the new value afterwards. A session beacon did not build carries + /// no namespace and falls back to this factory's own config. + fn session_config(&self, state: &dyn Session) -> Hdf5Config { + let Some(hdf5) = BeaconOptions::try_from_session(state).map(|beacon| beacon.hdf5) else { + // No namespace: this session was not built by beacon's runtime, so + // this factory's own config is the only configuration there is. + return self.config.clone(); + }; + Hdf5Config { + use_rust_reader: hdf5.use_rust_reader, + use_reader_cache: hdf5.use_reader_cache, + enable_statistics: hdf5.enable_statistics, + // Not settable: the cache is built once, at the capacity the runtime + // started with. + reader_cache_size: self.config.reader_cache_size, + } + } + + /// Resolve the per-table options against the session's settings, which the + /// runtime config seeded. fn effective_options( &self, + state: &dyn Session, format_options: &HashMap, ) -> datafusion::error::Result { + let config = self.session_config(state); let mut options = EffectiveOptions { - use_rust_reader: self.config.use_rust_reader, - use_reader_cache: self.config.use_reader_cache, - enable_statistics: self.config.enable_statistics, + use_rust_reader: config.use_rust_reader, + use_reader_cache: config.use_reader_cache, + enable_statistics: config.enable_statistics, read_dimensions: None, }; @@ -174,7 +201,7 @@ impl FileFormatFactory for Hdf5FormatFactory { state: &dyn Session, format_options: &HashMap, ) -> datafusion::error::Result> { - let options = self.effective_options(format_options)?; + let options = self.effective_options(state, format_options)?; // netcdf-c: hand the whole call to the netCDF factory, exactly as this // crate did before a second reader existed. if !options.use_rust_reader { @@ -214,7 +241,7 @@ impl FileFormatFactoryExt for Hdf5FormatFactory { url: &ListingTableUrl, listing: &ListingFactory, ) -> datafusion::error::Result> { - if self.uses_rust_reader(format_options)? { + if self.uses_rust_reader(state, format_options)? { return self.create(state, format_options); } self.inner diff --git a/beacon-db/beacon-file-formats/beacon-arrow-netcdf/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-netcdf/src/datafusion/mod.rs index 50e84730..74fcfa3f 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-netcdf/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-netcdf/src/datafusion/mod.rs @@ -6,6 +6,7 @@ use arrow::datatypes::SchemaRef; use beacon_common::super_typing::super_type_schema; use beacon_datafusion_ext::format_ext::{DatasetMetadata, FileFormatFactoryExt}; use beacon_datafusion_ext::listing_factory::ListingFactory; +use beacon_datafusion_ext::settings::BeaconOptions; use beacon_datafusion_ext::unique_values::UniqueValuesExec; use datafusion::{ catalog::{memory::DataSourceExec, Session}, @@ -116,6 +117,33 @@ impl NetCDFFormatFactory { } } + /// The settings this format starts from, for a format built on `session`. + /// + /// The runtime's `BEACON_NETCDF_*` values seed `beacon.netcdf.*` at startup, + /// so reading the namespace here yields the runtime default until an operator + /// runs `SET beacon.netcdf.use_rust_reader = true`, and the new value from + /// then on. A session beacon did not build carries no namespace, and falls + /// back to the config this factory was constructed with. + /// + /// A per-table `OPTIONS (...)` value still wins over both — see + /// [`FileFormatFactory::create`]. + fn session_config(&self, session: &dyn Session) -> NetcdfConfig { + let Some(netcdf) = BeaconOptions::try_from_session(session).map(|beacon| beacon.netcdf) + else { + // No namespace: this session was not built by beacon's runtime, so + // this factory's own config is the only configuration there is. + return self.config.clone(); + }; + NetcdfConfig { + use_reader_cache: netcdf.use_reader_cache, + enable_statistics: netcdf.enable_statistics, + use_rust_reader: netcdf.use_rust_reader, + // Not settable: the cache is built once, at the capacity the runtime + // started with. + reader_cache_size: self.config.reader_cache_size, + } + } + /// Build a [`NetcdfFormat`] with the given per-table effective settings, /// wiring in the shared reader cache when caching is enabled. fn build_format( @@ -151,15 +179,17 @@ fn access_for(use_rust_reader: bool) -> FileAccess { impl FileFormatFactory for NetCDFFormatFactory { fn create( &self, - _state: &dyn Session, + state: &dyn Session, format_options: &std::collections::HashMap, ) -> datafusion::error::Result> { - // Per-table overrides from `CREATE EXTERNAL TABLE ... OPTIONS (...)`, - // defaulting to the runtime config. + // Three layers, narrowest last: the runtime default, the session's + // `beacon.netcdf.*` namespace (what `SET` writes), then the per-table + // `CREATE EXTERNAL TABLE ... OPTIONS (...)`. + let session_config = self.session_config(state); let mut options = self.options.clone(); - let mut use_reader_cache = self.config.use_reader_cache; - let mut enable_statistics = self.config.enable_statistics; - let mut use_rust_reader = self.config.use_rust_reader; + let mut use_reader_cache = session_config.use_reader_cache; + let mut enable_statistics = session_config.enable_statistics; + let mut use_rust_reader = session_config.use_rust_reader; if let Some(value) = format_options.get("read_dimensions") { options.read_dimensions = Some( diff --git a/beacon-db/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs b/beacon-db/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs index 2b6cc38c..71e509a9 100644 --- a/beacon-db/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs +++ b/beacon-db/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs @@ -10,6 +10,7 @@ use std::{any::Any, sync::Arc}; use arrow::datatypes::SchemaRef; use beacon_common::super_typing::super_type_schema; use beacon_datafusion_ext::format_ext::{DatasetMetadata, FileFormatFactoryExt}; +use beacon_datafusion_ext::settings::BeaconOptions; use datafusion::{ catalog::{Session, memory::DataSourceExec}, common::{GetExt, Statistics}, @@ -75,11 +76,12 @@ impl GetExt for ZarrFormatFactory { impl FileFormatFactory for ZarrFormatFactory { fn create( &self, - _state: &dyn Session, + state: &dyn Session, format_options: &std::collections::HashMap, ) -> datafusion::error::Result> { - // Per-table overrides from `CREATE EXTERNAL TABLE ... OPTIONS (...)`, - // defaulting to the runtime config. + // Three layers, narrowest last: the runtime default, the session's + // `beacon.zarr.*` namespace (what `SET` writes), then the per-table + // `CREATE EXTERNAL TABLE ... OPTIONS (...)`. let read_dimensions = format_options.get("read_dimensions").map(|value| { value .split(',') @@ -87,7 +89,11 @@ impl FileFormatFactory for ZarrFormatFactory { .filter(|s| !s.is_empty()) .collect() }); - let mut enable_statistics = self.config.enable_statistics; + // No namespace means this session was not built by beacon's runtime, so + // this factory's own config is the only configuration there is. + let mut enable_statistics = BeaconOptions::try_from_session(state) + .map(|beacon| beacon.zarr.enable_statistics) + .unwrap_or(self.config.enable_statistics); if let Some(value) = format_options.get("enable_statistics") { enable_statistics = parse_bool_option("enable_statistics", value)?; } @@ -463,17 +469,29 @@ mod tests { ctx.register_table("gridded", Arc::new(table)).unwrap(); } - /// A session with the nd projection-pushdown rule registered — the same - /// wiring beacon-core installs, so a `SELECT`-with-computed-column plan gets - /// the projection sunk below the broadcast. + /// A session with the nd projection-pushdown rule registered *and* enabled — + /// the same wiring beacon-core installs, so a `SELECT`-with-computed-column + /// plan gets the projection sunk below the broadcast. + /// + /// The rule is always installed and reads `beacon.enable_nd_pipeline` on each + /// plan (so a `SET` can switch it on without a restart), and that flag is off + /// by default — hence the namespace here, not just the rule. fn ctx_with_pushdown() -> SessionContext { use datafusion::execution::session_state::SessionStateBuilder; use datafusion::prelude::SessionConfig; // Single partition so row order is deterministic (the differential tests // compare results positionally). + let mut config = SessionConfig::new().with_target_partitions(1); + config + .options_mut() + .extensions + .insert(beacon_datafusion_ext::settings::BeaconOptions { + enable_nd_pipeline: true, + ..Default::default() + }); let state = SessionStateBuilder::new() - .with_config(SessionConfig::new().with_target_partitions(1)) + .with_config(config) .with_default_features() .with_physical_optimizer_rule(Arc::new( beacon_datafusion_ext::nd::NdProjectionPushdown::new(), diff --git a/beacon-db/beacon-file-formats/beacon-lance/src/config.rs b/beacon-db/beacon-file-formats/beacon-lance/src/config.rs new file mode 100644 index 00000000..8979c705 --- /dev/null +++ b/beacon-db/beacon-file-formats/beacon-lance/src/config.rs @@ -0,0 +1,191 @@ +//! [`LanceConfig`]: the runtime settings of the managed Lance table engine. + +use beacon_datafusion_ext::settings::BeaconOptions; +use datafusion::catalog::Session; +use datafusion::execution::context::SessionConfig; +use lance::dataset::scanner::MaterializationStyle; +use lance_encoding::version::LanceFileVersion; + +/// Runtime configuration for managed Lance tables. +/// +/// Plain data with sensible defaults; the caller populates it (there is no +/// environment parsing here, so the crate stays reusable and the host decides +/// where the values come from). An empty string means "leave it to Lance", which +/// is what an unset `BEACON_LANCE_*` variable used to mean. +/// +/// The first four are *write* settings: they shape the files a `CREATE TABLE` or +/// `INSERT` produces, and never the files already on disk. `materialization` is a +/// read setting and applies to the next scan. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct LanceConfig { + /// Block compression for string and binary columns: `fsst`, `zstd`, `lz4`, or + /// `none`. Empty leaves them uncompressed. + pub compression: String, + /// Block compression for numeric columns: `zstd`, `lz4`, or `none`. Empty + /// keeps Lance's default. + pub numeric_compression: String, + /// File format version for new data: `2.0`, `2.1` or `2.2`. Empty means the + /// beacon default, `2.2`. + pub version: String, + /// Minichunk size in bytes for fixed-width columns. Empty leaves it to Lance. + pub minichunk: String, + /// Column materialization on a filtered scan: `late` or `early`. Empty keeps + /// beacon's projection-width rule. + pub materialization: String, +} + +impl LanceConfig { + /// The `beacon.lance.*` settings on `config`, or the defaults when the + /// namespace is absent (a session beacon did not build). + /// + /// Read at each write and each scan, never cached, so + /// `SET beacon.lance.version = '2.1'` applies to the next statement. + pub fn from_config(config: &SessionConfig) -> Self { + let lance = BeaconOptions::from_config(config).lance; + Self { + compression: lance.compression, + numeric_compression: lance.numeric_compression, + version: lance.version, + minichunk: lance.minichunk, + materialization: lance.materialization, + } + } + + /// [`Self::from_config`] for the `&dyn Session` a `TableProvider` is handed. + pub fn from_session(session: &dyn Session) -> Self { + Self::from_config(session.config()) + } + + /// Compression scheme for string columns, or `None` when unset or `none`. + /// + /// Compression is applied only to string/binary columns, and only when asked + /// for. Measured on a 20M-row ClickBench subset: + /// * no compression : 5.15GB, string scan 126ms, int scan 12.5ms + /// * zstd (all cols): 4.13GB, string scan 1334ms, int scan 42.1ms + /// Block-compressing numerics is a bad trade (3x slower scans for little + /// size), so numeric columns are never compressed here regardless. + pub(crate) fn string_compression(&self) -> Option<&str> { + non_empty(&self.compression).filter(|v| !v.eq_ignore_ascii_case("none")) + } + + /// Compression scheme for numeric columns, or `None` for Lance's default. + /// + /// From file version 2.2 Lance block-compresses any buffer over 32KB by + /// default, which shrinks the dataset but slows numeric scans. Setting `none` + /// opts numeric columns back out while leaving strings compressed. + pub(crate) fn numeric_compression(&self) -> Option<&str> { + non_empty(&self.numeric_compression) + } + + /// File format version for new data. + /// + /// Beacon writes 2.2, not Lance's 2.1 default. 2.2 is a stable version (Lance + /// only treats `>= Next` as unstable) and adds RLE for whole blocks plus + /// automatic block compression for buffers over 32KB. Measured on a 100M-row + /// ClickBench table: 27GB -> 21.9GB (-19%) with query time unchanged + /// (96.99s vs 95.01s over the 43-query suite, within run-to-run noise). + /// + /// Appends to an existing table keep that table's own version, so this only + /// affects newly created tables. An unrecognized value falls back to 2.2 + /// rather than failing a write. + pub(crate) fn storage_version(&self) -> Option { + match non_empty(&self.version) { + Some("2.0") => Some(LanceFileVersion::V2_0), + Some("2.1") => Some(LanceFileVersion::V2_1), + _ => Some(LanceFileVersion::V2_2), + } + } + + /// Minichunk size in bytes, as the string Lance's field metadata expects, or + /// `None` when unset or not a number. + pub(crate) fn minichunk_size(&self) -> Option { + non_empty(&self.minichunk)? + .parse::() + .ok() + .map(|n| n.to_string()) + } + + /// Materialization override, or `None` to keep beacon's projection-width rule. + pub(crate) fn materialization_style(&self) -> Option { + match non_empty(&self.materialization) { + Some("late") => Some(MaterializationStyle::AllLate), + Some("early") => Some(MaterializationStyle::AllEarly), + _ => None, + } + } +} + +/// The trimmed value, or `None` when it is empty — the "unset" spelling in this +/// config, since a config namespace has no null. +fn non_empty(value: &str) -> Option<&str> { + let trimmed = value.trim(); + (!trimmed.is_empty()).then_some(trimmed) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn empty_means_unset() { + let config = LanceConfig::default(); + assert_eq!(config.string_compression(), None); + assert_eq!(config.numeric_compression(), None); + assert_eq!(config.minichunk_size(), None); + assert!(config.materialization_style().is_none()); + // The version is the one setting with a beacon default rather than a + // Lance one, so "unset" still resolves to 2.2. + assert_eq!(config.storage_version(), Some(LanceFileVersion::V2_2)); + } + + /// `none` is how an operator turns string compression off, and it must not + /// reach Lance as a scheme name. + #[test] + fn none_disables_string_compression() { + let config = LanceConfig { + compression: "none".to_string(), + ..Default::default() + }; + assert_eq!(config.string_compression(), None); + + let config = LanceConfig { + compression: "zstd".to_string(), + ..Default::default() + }; + assert_eq!(config.string_compression(), Some("zstd")); + } + + /// `none` on numerics is meaningful — it opts them out of Lance's automatic + /// block compression — so unlike strings it is passed through. + #[test] + fn none_is_passed_through_for_numeric_compression() { + let config = LanceConfig { + numeric_compression: "none".to_string(), + ..Default::default() + }; + assert_eq!(config.numeric_compression(), Some("none")); + } + + #[test] + fn version_and_minichunk_parse() { + let config = LanceConfig { + version: "2.0".to_string(), + minichunk: " 65536 ".to_string(), + materialization: "late".to_string(), + ..Default::default() + }; + assert_eq!(config.storage_version(), Some(LanceFileVersion::V2_0)); + assert_eq!(config.minichunk_size().as_deref(), Some("65536")); + assert!(matches!( + config.materialization_style(), + Some(MaterializationStyle::AllLate) + )); + + // A non-numeric minichunk is ignored rather than failing a write. + let config = LanceConfig { + minichunk: "big".to_string(), + ..Default::default() + }; + assert_eq!(config.minichunk_size(), None); + } +} diff --git a/beacon-db/beacon-file-formats/beacon-lance/src/io.rs b/beacon-db/beacon-file-formats/beacon-lance/src/io.rs index c8c43f22..79c525e1 100644 --- a/beacon-db/beacon-file-formats/beacon-lance/src/io.rs +++ b/beacon-db/beacon-file-formats/beacon-lance/src/io.rs @@ -14,7 +14,8 @@ use futures::StreamExt; use lance::dataset::write::InsertBuilder; use lance::dataset::{WriteMode, WriteParams}; use lance::session::Session; -use lance_encoding::version::LanceFileVersion; + +use crate::config::LanceConfig; /// Map an Arrow data type to one Lance can store. Lance 7.x does not support the /// Arrow "view" types (`Utf8View`/`BinaryView`) that DataFusion 53 produces for @@ -27,8 +28,9 @@ pub(crate) fn lance_compatible_type(data_type: &DataType) -> DataType { } } -/// A Lance-writable version of `schema` (view types widened to non-view). -pub(crate) fn lance_compatible_schema(schema: &Schema) -> SchemaRef { +/// A Lance-writable version of `schema` (view types widened to non-view), with +/// `config`'s encoding hints attached as Arrow field metadata. +pub(crate) fn lance_compatible_schema(schema: &Schema, config: &LanceConfig) -> SchemaRef { // Lance takes per-field encoding hints from Arrow field metadata. // // Lance used to store this data about 2x larger than parquet (27GB vs 14GB on @@ -48,22 +50,22 @@ pub(crate) fn lance_compatible_schema(schema: &Schema) -> SchemaRef { // // Note this also preserves any metadata already on the field: rebuilding with // `Field::new` alone silently dropped it. - let compression = lance_string_compression(); - let numeric_compression = lance_numeric_compression(); - let minichunk = lance_minichunk_size(); + let compression = config.string_compression(); + let numeric_compression = config.numeric_compression(); + let minichunk = config.minichunk_size(); let fields = schema .fields() .iter() .map(|f| { let mut metadata = f.metadata().clone(); - if let Some(scheme) = compression.as_deref() { + if let Some(scheme) = compression { if is_string_like(f.data_type()) { metadata .entry(COMPRESSION_META_KEY.to_string()) .or_insert_with(|| scheme.to_string()); } } - if let Some(scheme) = numeric_compression.as_deref() { + if let Some(scheme) = numeric_compression { if !is_string_like(f.data_type()) { metadata .entry(COMPRESSION_META_KEY.to_string()) @@ -93,75 +95,11 @@ pub(crate) fn lance_compatible_schema(schema: &Schema) -> SchemaRef { /// Arrow field-metadata key Lance reads for a per-column compression scheme. const COMPRESSION_META_KEY: &str = "lance-encoding:compression"; -/// Optional compression scheme for managed Lance *string* columns. -/// -/// Compression is applied only to string/binary columns, and only when asked -/// for. Measured on a 20M-row ClickBench subset: -/// * no compression : 5.15GB, string scan 126ms, int scan 12.5ms -/// * zstd (all cols): 4.13GB, string scan 1334ms, int scan 42.1ms -/// Block-compressing numerics is a bad trade (3x slower scans for little size), -/// so numeric columns are never compressed here regardless of the setting. -/// -/// `BEACON_LANCE_COMPRESSION=fsst|zstd|lz4` opts in; unset or `none` disables. -fn lance_string_compression() -> Option { - match std::env::var("BEACON_LANCE_COMPRESSION") { - Ok(v) if v.trim().is_empty() || v.eq_ignore_ascii_case("none") => None, - Ok(v) => Some(v), - Err(_) => None, - } -} - -/// Optional compression scheme for managed Lance *numeric* columns. -/// -/// From file version 2.2 Lance block-compresses any buffer over 32KB by default, -/// which shrinks the dataset but slows numeric scans. Setting `none` here opts -/// numeric columns back out while leaving strings compressed. -/// -/// `BEACON_LANCE_NUMERIC_COMPRESSION=none|zstd|lz4`; unset means Lance's default. -fn lance_numeric_compression() -> Option { - let v = std::env::var("BEACON_LANCE_NUMERIC_COMPRESSION").ok()?; - let v = v.trim(); - (!v.is_empty()).then(|| v.to_string()) -} - /// Arrow field-metadata key for Lance's minichunk size (its decode unit for /// fixed-width columns). Sizes of 32KB and up need file version 2.2 or later; /// on 2.1 Lance logs a warning and ignores them. const MINICHUNK_SIZE_META_KEY: &str = "lance-encoding:minichunk-size"; -/// Lance file format version for new data. -/// -/// Beacon writes 2.2, not Lance's 2.1 default. 2.2 is a stable version (Lance -/// only treats `>= Next` as unstable) and adds RLE for whole blocks plus -/// automatic block compression for buffers over 32KB. Measured on a 100M-row -/// ClickBench table: 27GB -> 21.9GB (-19%) with query time unchanged -/// (96.99s vs 95.01s over the 43-query suite, within run-to-run noise). -/// -/// Appends to an existing table keep that table's own version, so this only -/// affects newly created tables. -/// -/// `BEACON_LANCE_VERSION=2.0|2.1|2.2` overrides. -fn lance_storage_version() -> Option { - match std::env::var("BEACON_LANCE_VERSION") - .ok() - .as_deref() - .map(str::trim) - { - Some("2.0") => Some(LanceFileVersion::V2_0), - Some("2.1") => Some(LanceFileVersion::V2_1), - _ => Some(LanceFileVersion::V2_2), - } -} - -/// Optional minichunk size, in bytes, for fixed-width columns. -/// -/// `BEACON_LANCE_MINICHUNK=`; needs `BEACON_LANCE_VERSION=2.2` to take -/// effect above 32KB. -fn lance_minichunk_size() -> Option { - let v = std::env::var("BEACON_LANCE_MINICHUNK").ok()?; - v.trim().parse::().ok().map(|n| n.to_string()) -} - /// True for the types where block compression can pay for itself. fn is_string_like(dt: &arrow::datatypes::DataType) -> bool { use arrow::datatypes::DataType::*; @@ -231,8 +169,9 @@ pub async fn write_stream( session: Arc, rows: SendableRecordBatchStream, kind: WriteKind, + config: &LanceConfig, ) -> anyhow::Result { - let target = lance_compatible_schema(&rows.schema()); + let target = lance_compatible_schema(&rows.schema(), config); // Count rows + coerce view types as batches stream past, without collecting. let written = Arc::new(AtomicU64::new(0)); @@ -249,7 +188,7 @@ pub async fn write_stream( let params = WriteParams { mode: kind.into(), session: Some(session), - data_storage_version: lance_storage_version(), + data_storage_version: config.storage_version(), ..Default::default() }; InsertBuilder::new(uri) @@ -298,7 +237,7 @@ mod tests { Field::new("name", DataType::Utf8View, true), Field::new("blob", DataType::BinaryView, true), ]); - let widened = lance_compatible_schema(&schema); + let widened = lance_compatible_schema(&schema, &LanceConfig::default()); assert_eq!(widened.field(0).data_type(), &DataType::Int64); assert_eq!(widened.field(1).data_type(), &DataType::Utf8); assert!(widened.field(1).is_nullable()); @@ -307,6 +246,40 @@ mod tests { assert!(!widened.field(0).is_nullable()); } + /// The encoding hints reach Lance as field metadata, and only on the column + /// kinds each applies to — string compression on strings, numeric compression + /// and minichunking on everything else. + #[test] + fn config_encoding_hints_land_on_the_right_columns() { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + ]); + let config = LanceConfig { + compression: "zstd".to_string(), + numeric_compression: "lz4".to_string(), + minichunk: "65536".to_string(), + ..Default::default() + }; + + let hinted = lance_compatible_schema(&schema, &config); + let numeric = hinted.field(0).metadata().clone(); + let string = hinted.field(1).metadata().clone(); + + assert_eq!(numeric.get(COMPRESSION_META_KEY).map(String::as_str), Some("lz4")); + assert_eq!( + numeric.get(MINICHUNK_SIZE_META_KEY).map(String::as_str), + Some("65536") + ); + assert_eq!(string.get(COMPRESSION_META_KEY).map(String::as_str), Some("zstd")); + assert_eq!(string.get(MINICHUNK_SIZE_META_KEY), None); + + // The default config attaches nothing at all. + let plain = lance_compatible_schema(&schema, &LanceConfig::default()); + assert!(plain.field(0).metadata().is_empty()); + assert!(plain.field(1).metadata().is_empty()); + } + #[test] fn write_kind_maps_to_lance_write_mode() { assert!(matches!(WriteMode::from(WriteKind::Create), WriteMode::Create)); @@ -339,7 +312,7 @@ mod tests { ) .unwrap(); - let target = lance_compatible_schema(&source_schema); + let target = lance_compatible_schema(&source_schema, &LanceConfig::default()); let coerced = coerce_batch(&batch, &target).unwrap(); assert_eq!(coerced.schema().field(1).data_type(), &DataType::Utf8); assert_eq!(coerced.num_rows(), 2); diff --git a/beacon-db/beacon-file-formats/beacon-lance/src/lib.rs b/beacon-db/beacon-file-formats/beacon-lance/src/lib.rs index dac90b68..eb3868d0 100644 --- a/beacon-db/beacon-file-formats/beacon-lance/src/lib.rs +++ b/beacon-db/beacon-file-formats/beacon-lance/src/lib.rs @@ -15,6 +15,7 @@ //! torn reads; a per-dataset write lock serializes writers. pub mod alter; +pub mod config; pub mod definition; pub mod index; pub mod io; @@ -31,6 +32,7 @@ use futures::StreamExt; use object_store::{ObjectStore, ObjectStoreExt}; pub use alter::{SchemaChange, alter_table}; +pub use config::LanceConfig; pub use definition::LanceTableDefinition; pub use index::{ IndexInfo, ScalarIndexKind, create_default_indexes, create_index, drop_index, list_indices, @@ -48,13 +50,14 @@ pub async fn create_lance_table( namespace: &[String], name: &str, arrow_schema: &ArrowSchema, + config: &LanceConfig, ) -> anyhow::Result { let uri = warehouse.table_uri(namespace, name); tracing::info!(namespace = ?namespace, table = name, uri = %uri, "creating Lance table"); // Store the Lance-writable schema (view types widened) so the provider's // schema matches the written dataset. - let schema = io::lance_compatible_schema(arrow_schema); + let schema = io::lance_compatible_schema(arrow_schema, config); { let lock = warehouse.lock(&uri); let _guard = lock.lock().await; @@ -65,6 +68,7 @@ pub async fn create_lance_table( warehouse.session(), io::empty_stream(schema.clone()), WriteKind::Create, + config, ) .await?; } @@ -108,12 +112,20 @@ pub async fn replace_table_contents( warehouse: &LanceWarehouse, uri: &str, new_rows: SendableRecordBatchStream, + config: &LanceConfig, ) -> anyhow::Result<()> { tracing::info!(uri = %uri, "replacing Lance table contents"); let lock = warehouse.lock(uri); let _guard = lock.lock().await; - io::write_stream(uri, warehouse.session(), new_rows, WriteKind::Overwrite).await?; + io::write_stream( + uri, + warehouse.session(), + new_rows, + WriteKind::Overwrite, + config, + ) + .await?; Ok(()) } @@ -163,7 +175,13 @@ mod tests { let warehouse = test_warehouse(&dir); let namespace = beacon_namespace(); - let table = create_lance_table(warehouse.clone(), &namespace, "orders", &sample_schema()) + let table = create_lance_table( + warehouse.clone(), + &namespace, + "orders", + &sample_schema(), + &LanceConfig::default(), + ) .await .expect("table should be created"); let location = table.definition().location.clone(); @@ -191,7 +209,7 @@ mod tests { .execute_stream() .await .unwrap(); - replace_table_contents(&warehouse, &location, keep) + replace_table_contents(&warehouse, &location, keep, &LanceConfig::default()) .await .expect("replace should succeed"); assert_eq!( @@ -215,6 +233,7 @@ mod tests { &namespace, "discovered", &sample_schema(), + &LanceConfig::default(), ) .await .unwrap(); @@ -265,7 +284,13 @@ mod tests { let warehouse = test_warehouse(&dir); let namespace = beacon_namespace(); - let table = create_lance_table(warehouse.clone(), &namespace, "orders", &sample_schema()) + let table = create_lance_table( + warehouse.clone(), + &namespace, + "orders", + &sample_schema(), + &LanceConfig::default(), + ) .await .unwrap(); let location = table.definition().location.clone(); @@ -349,7 +374,13 @@ mod tests { let warehouse = test_warehouse(&dir); let namespace = beacon_namespace(); - let table = create_lance_table(warehouse.clone(), &namespace, "orders", &sample_schema()) + let table = create_lance_table( + warehouse.clone(), + &namespace, + "orders", + &sample_schema(), + &LanceConfig::default(), + ) .await .unwrap(); let location = table.definition().location.clone(); @@ -397,12 +428,12 @@ mod tests { let warehouse = test_warehouse(&dir); let namespace = beacon_namespace(); - create_lance_table(warehouse.clone(), &namespace, "dupe", &sample_schema()) + create_lance_table(warehouse.clone(), &namespace, "dupe", &sample_schema(), &LanceConfig::default()) .await .expect("first create succeeds"); // A second CREATE at the same location must not silently clobber the // dataset (WriteKind::Create errors when a dataset is already present). - let err = create_lance_table(warehouse.clone(), &namespace, "dupe", &sample_schema()) + let err = create_lance_table(warehouse.clone(), &namespace, "dupe", &sample_schema(), &LanceConfig::default()) .await .expect_err("second create must fail"); assert!( @@ -424,7 +455,7 @@ mod tests { Field::new("blob", DataType::BinaryView, true), ]); - let table = create_lance_table(warehouse.clone(), &beacon_namespace(), "views", &schema) + let table = create_lance_table(warehouse.clone(), &beacon_namespace(), "views", &schema, &LanceConfig::default()) .await .unwrap(); use datafusion::catalog::TableProvider; @@ -439,7 +470,7 @@ mod tests { let warehouse = test_warehouse(&dir); let namespace = beacon_namespace(); - let table = create_lance_table(warehouse.clone(), &namespace, "gone", &sample_schema()) + let table = create_lance_table(warehouse.clone(), &namespace, "gone", &sample_schema(), &LanceConfig::default()) .await .unwrap(); let location = table.definition().location.clone(); @@ -471,7 +502,13 @@ mod tests { let warehouse = test_warehouse(&dir); let namespace = beacon_namespace(); - let table = create_lance_table(warehouse.clone(), &namespace, "orders", &sample_schema()) + let table = create_lance_table( + warehouse.clone(), + &namespace, + "orders", + &sample_schema(), + &LanceConfig::default(), + ) .await .unwrap(); let location = table.definition().location.clone(); diff --git a/beacon-db/beacon-file-formats/beacon-lance/src/provider.rs b/beacon-db/beacon-file-formats/beacon-lance/src/provider.rs index 34fd1f06..5ea557ae 100644 --- a/beacon-db/beacon-file-formats/beacon-lance/src/provider.rs +++ b/beacon-db/beacon-file-formats/beacon-lance/src/provider.rs @@ -29,6 +29,7 @@ use lance::session::Session as LanceSession; use crate::definition::LanceTableDefinition; use crate::io::WriteKind; use crate::sink::LanceDataSink; +use crate::config::LanceConfig; use crate::warehouse::LanceWarehouse; /// A beacon-managed Lance table provider. @@ -103,6 +104,7 @@ async fn scan_fragment_group( range: std::ops::Range, projection: Option<&Vec>, filters: &[Expr], + config: &LanceConfig, ) -> DataFusionResult> { let projected_columns = projection.map_or_else(|| schema.fields().len(), |p| p.len()); let mut scan = dataset.scan(); @@ -153,7 +155,7 @@ async fn scan_fragment_group( // time): above `LATE_MATERIALIZATION_MIN_COLUMNS` the downside is a bounded // ~10-25% on unselective filters and the upside is several fold, while // narrow projections keep Lance's heuristic, where early costs little. - match lance_materialization_style() { + match config.materialization_style() { Some(style) => scan.materialization_style(style), None if projected_columns > LATE_MATERIALIZATION_MIN_COLUMNS => { scan.materialization_style(MaterializationStyle::AllLate) @@ -173,22 +175,6 @@ async fn scan_fragment_group( /// See the crossover measurements in `scan_fragment_group`. const LATE_MATERIALIZATION_MIN_COLUMNS: usize = 16; -/// Override for Lance's column materialization heuristic. -/// -/// `BEACON_LANCE_MATERIALIZATION=late|early` forces all columns one way; unset -/// keeps Lance's per-column heuristic. -fn lance_materialization_style() -> Option { - match std::env::var("BEACON_LANCE_MATERIALIZATION") - .ok() - .as_deref() - .map(str::trim) - { - Some("late") => Some(MaterializationStyle::AllLate), - Some("early") => Some(MaterializationStyle::AllEarly), - _ => None, - } -} - #[async_trait] impl TableProvider for LanceTable { fn as_any(&self) -> &dyn Any { @@ -223,6 +209,10 @@ impl TableProvider for LanceTable { return provider.scan(state, projection, filters, limit).await; } + // Read per scan, so `SET beacon.lance.materialization` applies to the next + // query rather than to the next restart. + let config = LanceConfig::from_session(state); + // Spread fragments over at most `target` groups. let groups = target.min(n_frags); let per = n_frags.div_ceil(groups); @@ -231,7 +221,8 @@ impl TableProvider for LanceTable { while start < n_frags { let end = (start + per).min(n_frags); plans.push( - scan_fragment_group(&dataset, &self.schema, start..end, projection, filters).await?, + scan_fragment_group(&dataset, &self.schema, start..end, projection, filters, &config) + .await?, ); start = end; } diff --git a/beacon-db/beacon-file-formats/beacon-lance/src/sink.rs b/beacon-db/beacon-file-formats/beacon-lance/src/sink.rs index a578a339..3d46eb47 100644 --- a/beacon-db/beacon-file-formats/beacon-lance/src/sink.rs +++ b/beacon-db/beacon-file-formats/beacon-lance/src/sink.rs @@ -17,6 +17,7 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_plan::metrics::MetricsSet; use datafusion::physical_plan::{DisplayAs, DisplayFormatType}; +use crate::config::LanceConfig; use crate::io::{write_stream, WriteKind}; use crate::warehouse::LanceWarehouse; @@ -68,13 +69,17 @@ impl DataSink for LanceDataSink { async fn write_all( &self, data: SendableRecordBatchStream, - _context: &Arc, + context: &Arc, ) -> DataFusionResult { + // The encoding settings come off the task's session config, so a + // `SET beacon.lance.*` shapes the files this write produces. + let config = LanceConfig::from_config(context.session_config()); + // Serialize writers to this dataset across the (async) write, then stream // the input directly into Lance — no full-table buffering. let lock = self.warehouse.lock(&self.uri); let _guard = lock.lock().await; - write_stream(&self.uri, self.warehouse.session(), data, self.kind) + write_stream(&self.uri, self.warehouse.session(), data, self.kind, &config) .await .map_err(|e| DataFusionError::External(e.into())) } diff --git a/beacon-server/beacon-server-config/Cargo.toml b/beacon-server/beacon-server-config/Cargo.toml index 4ea2bcc3..bfeeb906 100644 --- a/beacon-server/beacon-server-config/Cargo.toml +++ b/beacon-server/beacon-server-config/Cargo.toml @@ -21,3 +21,4 @@ beacon-arrow-hdf5 = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-h beacon-arrow-zarr = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-zarr" } beacon-arrow-atlas = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-atlas" } beacon-arrow-bbf = { path = "../../beacon-db/beacon-file-formats/beacon-arrow-bbf" } +beacon-lance = { path = "../../beacon-db/beacon-file-formats/beacon-lance" } diff --git a/beacon-server/beacon-server-config/src/lib.rs b/beacon-server/beacon-server-config/src/lib.rs index c430cb16..98aecdce 100644 --- a/beacon-server/beacon-server-config/src/lib.rs +++ b/beacon-server/beacon-server-config/src/lib.rs @@ -15,6 +15,7 @@ pub use beacon_arrow_bbf::datafusion::BbfConfig; pub use beacon_arrow_hdf5::Hdf5Config; pub use beacon_arrow_netcdf::datafusion::NetcdfConfig; pub use beacon_arrow_zarr::ZarrConfig; +pub use beacon_lance::LanceConfig; pub use beacon_common::FileStatsConfig; pub use beacon_common::CrawlerConfig; @@ -33,6 +34,7 @@ pub struct Config { pub zarr: ZarrConfig, pub atlas: AtlasConfig, pub bbf: BbfConfig, + pub lance: LanceConfig, pub crawler: CrawlerConfig, pub file_stats: FileStatsConfig, pub api_docs: ApiDocsConfig, @@ -472,6 +474,19 @@ struct RawConfig { #[envconfig(from = "BEACON_ENABLE_BBF_SPLIT_STREAMS_SLICE", default = "false")] bbf_split_streams_slice: bool, + // Managed Lance tables. Empty means "leave it to Lance"; each is also + // settable at runtime as `beacon.lance.*`. + #[envconfig(from = "BEACON_LANCE_COMPRESSION", default = "")] + lance_compression: String, + #[envconfig(from = "BEACON_LANCE_NUMERIC_COMPRESSION", default = "")] + lance_numeric_compression: String, + #[envconfig(from = "BEACON_LANCE_VERSION", default = "")] + lance_version: String, + #[envconfig(from = "BEACON_LANCE_MINICHUNK", default = "")] + lance_minichunk: String, + #[envconfig(from = "BEACON_LANCE_MATERIALIZATION", default = "")] + lance_materialization: String, + // Base64-encoded 32-byte master key for encrypting persisted secrets // (external-database credentials). Optional; validated in `Config::load`. #[envconfig(from = "BEACON_SECRETS_KEY")] @@ -628,6 +643,13 @@ impl From for Config { bbf: BbfConfig { split_streams_slice: raw.bbf_split_streams_slice, }, + lance: LanceConfig { + compression: raw.lance_compression, + numeric_compression: raw.lance_numeric_compression, + version: raw.lance_version, + minichunk: raw.lance_minichunk, + materialization: raw.lance_materialization, + }, file_stats: FileStatsConfig { enable: raw.file_stats_enable, interval_secs: raw.file_stats_interval_secs, diff --git a/beacon-server/beacon-server/src/server/mod.rs b/beacon-server/beacon-server/src/server/mod.rs index b3eeaac7..c6809f9b 100644 --- a/beacon-server/beacon-server/src/server/mod.rs +++ b/beacon-server/beacon-server/src/server/mod.rs @@ -264,6 +264,9 @@ async fn build_runtime( .with_netcdf_config(config.netcdf.clone()) .with_hdf5_config(config.hdf5.clone()) .with_zarr_config(config.zarr.clone()) + .with_atlas_config(config.atlas.clone()) + .with_bbf_config(config.bbf.clone()) + .with_lance_config(config.lance.clone()) .with_sql_settings(SqlSettings { default_table: config.sql.default_table.clone(), enable_pushdown_projection: config.sql.enable_pushdown_projection, diff --git a/docs/docs/2.0.0-rc2/server/configuration.md b/docs/docs/2.0.0-rc2/server/configuration.md index e07433e8..6901cfe8 100644 --- a/docs/docs/2.0.0-rc2/server/configuration.md +++ b/docs/docs/2.0.0-rc2/server/configuration.md @@ -1,12 +1,16 @@ --- -description: Full reference of the BEACON_* environment variables. It covers the server, engine, storage, S3, Flight SQL, crawlers and formats, with their defaults. +description: Full reference of the BEACON_* environment variables. It covers the server, engine, storage, S3, Flight SQL, crawlers and formats, with their defaults, and the SQL statements that change a setting without a restart. --- # Configuration -You configure Beacon with **environment variables** only. There is no -configuration file. Beacon reads every option below from the environment at -startup. An unset variable takes the default from this page. +You configure Beacon with **environment variables**. There is no configuration +file. Beacon reads every option below from the environment at startup. An unset +variable takes the default from this page. + +Many engine and format settings also change **at runtime**, through SQL, with no +restart. See [Change a setting at runtime](#change-a-setting-at-runtime). The +table for each section marks which ones. ::: info Every setting uses a `BEACON_*` name. The S3 credential variables are the @@ -14,6 +18,91 @@ exception. They use the standard `AWS_*` names, so they work with your AWS tools See [S3 object storage](#s3-object-storage). ::: +## Change a setting at runtime + +`SET` changes a setting on the running server: + +```sql +SET beacon.netcdf.use_rust_reader = true; +SET beacon.sql.stream_coalesce.target_rows = 131072; +SET beacon.default_table = 'observations'; +``` + +A setting name is its `BEACON_*` variable, lowercased, under the `beacon.` +namespace: `BEACON_NETCDF_USE_RUST_READER` becomes `beacon.netcdf.use_rust_reader`. + +`SHOW SETTINGS` lists every setting you can change, with its value, the value the +server started with, and what it does. Any authenticated user can read it. + +```sql +SHOW SETTINGS; +SHOW beacon.default_table; +SELECT * FROM beacon.system.settings WHERE name LIKE 'beacon.netcdf.%'; +``` + +### Scope + +A `SET` applies to the **whole server**, not to one client. Beacon runs one query +engine for every connection, so the next query from any user sees the new value. +Only the super-user may change a setting. This is the scope `SET` has always had +for the engine options below. + +### Persistence + +`SET` lasts until the server stops. `ALTER SYSTEM SET` also writes the value into +the database file, and the server applies it again at the next start: + +```sql +ALTER SYSTEM SET beacon.netcdf.use_rust_reader = 'true'; +``` + +`RESET` puts a setting back to the value the server started with. `ALTER SYSTEM +RESET` also removes the stored value: + +```sql +RESET beacon.netcdf.use_rust_reader; +ALTER SYSTEM RESET beacon.netcdf.use_rust_reader; +``` + +Order of precedence, strongest first: `ALTER SYSTEM SET` > environment variable > +default. `ALTER SYSTEM` needs a file-backed database. An in-memory database +refuses it. + +### Engine options + +The `beacon.` namespace also covers the query engine's own options. Both spellings +work: + +```sql +SET beacon.execution.batch_size = 8192; +SET datafusion.execution.batch_size = 8192; -- the same option +``` + +`SELECT * FROM information_schema.df_settings` lists both namespaces. + +### Settings you cannot change at runtime + +A setting that Beacon uses once, at startup, stays an environment variable. This +covers the port, the host, the worker threads, the data directory, the S3 store, +the credentials, the Flight SQL server, CORS, the crawler, the file statistics, +and each reader cache **size**. A `SET` on one of these fails and names the +variable to edit: + +``` +`beacon.port` can only be set when the server starts: set the `BEACON_PORT` +environment variable and restart +``` + +### What a change reaches + +A `SET` on a format setting applies to the next `read_netcdf(...)`, +`read_zarr(...)` or similar call. A **registered** external table keeps the +settings it was created with, because Beacon builds its reader once. Run +`REFRESH
` to rebuild it with the current settings. + +A per-table `CREATE EXTERNAL TABLE ... OPTIONS (...)` value always wins over a +`SET`, which in turn wins over the environment variable. + ## Server | Variable | Default | Description | @@ -74,10 +163,10 @@ rejects that `CREATE`. Beacon never writes plaintext. | --- | --- | --- | | `BEACON_ENABLE_SQL` | `true` | Enable the raw SQL query interface. Set to `false` to disable it (the JSON query API stays available). | | `BEACON_VM_MEMORY_SIZE` | `8192` | Working memory (MB) available to the query engine. More is better for larger datasets and memory-heavy operations such as spatial joins and `GROUP BY`. | -| `BEACON_DEFAULT_TABLE` | `default` | Table queried when a request omits the source. Only applies to the JSON query API, SQL queries must always specify a source. | -| `BEACON_ENABLE_PUSHDOWN_PROJECTION` | `true` | Push column projection down into file readers so only requested columns are decoded. | -| `BEACON_ENABLE_ND_PIPELINE` | `false` | Enable the N-dimensional pipeline optimizer for zarr/netcdf reads: sink element-wise projections below the grid broadcast so `lat * 2` and similar run on the coordinate axis instead of the full cross-product. The base nd pipeline always runs; this only enables the node-rewriting optimization. | -| `BEACON_BATCH_SIZE` | `64000` | Batch size, in rows, for NetCDF reads (local and MPIO). | +| `BEACON_DEFAULT_TABLE` | `default` | Table queried when a request omits the source. Only applies to the JSON query API, SQL queries must always specify a source. Settable at runtime as `beacon.default_table`. | +| `BEACON_ENABLE_PUSHDOWN_PROJECTION` | `true` | Push column projection down into file readers so only requested columns are decoded. Settable at runtime as `beacon.enable_pushdown_projection`. | +| `BEACON_ENABLE_ND_PIPELINE` | `false` | Enable the N-dimensional pipeline optimizer for zarr/netcdf reads: sink element-wise projections below the grid broadcast so `lat * 2` and similar run on the coordinate axis instead of the full cross-product. The base nd pipeline always runs; this only enables the node-rewriting optimization. Settable at runtime as `beacon.enable_nd_pipeline`. | +| `BEACON_BATCH_SIZE` | `64000` | Batch size, in rows, for NetCDF reads (local and MPIO). Settable at runtime as `beacon.batch_size`. | | `BEACON_STATS_CACHE_CAPACITY` | `10000` | Maximum number of per-file statistics entries cached for query pruning. Read once at startup. | ### SQL result-stream coalescing @@ -88,10 +177,10 @@ many small batches. | Variable | Default | Description | | --- | --- | --- | -| `BEACON_SQL_STREAM_COALESCE_ENABLED` | `true` | Enable coalescing of the SQL result stream. | -| `BEACON_SQL_STREAM_COALESCE_TARGET_ROWS` | `65536` | Target rows per coalesced batch. | -| `BEACON_SQL_STREAM_COALESCE_FLUSH_TIMEOUT_MS` | `25` | Max time (ms) to wait while accumulating rows before flushing a partial batch. | -| `BEACON_SQL_STREAM_COALESCE_MAX_ROWS` | `262144` | Hard upper bound on rows per coalesced batch. | +| `BEACON_SQL_STREAM_COALESCE_ENABLED` | `true` | Enable coalescing of the SQL result stream. Settable at runtime as `beacon.sql.stream_coalesce.enabled`. | +| `BEACON_SQL_STREAM_COALESCE_TARGET_ROWS` | `65536` | Target rows per coalesced batch. Settable at runtime as `beacon.sql.stream_coalesce.target_rows`. | +| `BEACON_SQL_STREAM_COALESCE_FLUSH_TIMEOUT_MS` | `25` | Max time (ms) to wait while accumulating rows before flushing a partial batch. Settable at runtime as `beacon.sql.stream_coalesce.flush_timeout_ms`. | +| `BEACON_SQL_STREAM_COALESCE_MAX_ROWS` | `262144` | Hard upper bound on rows per coalesced batch. Settable at runtime as `beacon.sql.stream_coalesce.max_rows`. | ## Arrow Flight SQL @@ -216,10 +305,10 @@ change them. | Variable | Default | Description | | --- | --- | --- | -| `BEACON_NETCDF_ENABLE_STATISTICS` | `true` | Compute and cache per-file statistics used for query pruning. | -| `BEACON_NETCDF_USE_READER_CACHE` | `true` | Cache opened NetCDF readers in memory. | +| `BEACON_NETCDF_ENABLE_STATISTICS` | `true` | Compute and cache per-file statistics used for query pruning. Settable at runtime as `beacon.netcdf.enable_statistics`. | +| `BEACON_NETCDF_USE_READER_CACHE` | `true` | Cache opened NetCDF readers in memory. Settable at runtime as `beacon.netcdf.use_reader_cache`. | | `BEACON_NETCDF_READER_CACHE_SIZE` | `128` | Max NetCDF reader entries to keep cached. | -| `BEACON_NETCDF_USE_RUST_READER` | `false` | Read NetCDF with the pure-Rust reader instead of the netCDF-C library. | +| `BEACON_NETCDF_USE_RUST_READER` | `false` | Read NetCDF with the pure-Rust reader instead of the netCDF-C library. Settable at runtime as `beacon.netcdf.use_rust_reader`. | ### HDF5 @@ -229,9 +318,9 @@ reader flag, so you can move one format at a time. | Variable | Default | Description | | --- | --- | --- | -| `BEACON_HDF5_USE_RUST_READER` | `false` | Read HDF5 with the pure-Rust reader instead of the netCDF-C library. | -| `BEACON_HDF5_ENABLE_STATISTICS` | `true` | Compute per-file statistics used for query pruning. Needs the pure-Rust reader. | -| `BEACON_HDF5_USE_READER_CACHE` | `true` | Cache opened HDF5 readers in memory. | +| `BEACON_HDF5_USE_RUST_READER` | `false` | Read HDF5 with the pure-Rust reader instead of the netCDF-C library. Settable at runtime as `beacon.hdf5.use_rust_reader`. | +| `BEACON_HDF5_ENABLE_STATISTICS` | `true` | Compute per-file statistics used for query pruning. Needs the pure-Rust reader. Settable at runtime as `beacon.hdf5.enable_statistics`. | +| `BEACON_HDF5_USE_READER_CACHE` | `true` | Cache opened HDF5 readers in memory. Settable at runtime as `beacon.hdf5.use_reader_cache`. | | `BEACON_HDF5_READER_CACHE_SIZE` | `128` | Max HDF5 reader entries to keep cached. | The pure-Rust reader also reads two layouts the netCDF data model cannot express: a nested group, @@ -242,7 +331,7 @@ and a compound dataset. See | Variable | Default | Description | | --- | --- | --- | -| `BEACON_ZARR_ENABLE_STATISTICS` | `true` | Compute per-file statistics used for query pruning. | +| `BEACON_ZARR_ENABLE_STATISTICS` | `true` | Compute per-file statistics used for query pruning. Settable at runtime as `beacon.zarr.enable_statistics`. | A store answers from its `actual_range` metadata where it can. Where it cannot, it reads only its rank-0 and rank-1 arrays — the coordinates a `WHERE` clause names. A data grid of rank 2 or higher @@ -256,14 +345,33 @@ which values a store holds, so a store may hold values outside them. | Variable | Default | Description | | --- | --- | --- | -| `BEACON_ATLAS_USE_READER_CACHE` | `true` | Cache opened Atlas store readers in memory, avoiding re-opening the same `atlas.json` across queries. | +| `BEACON_ATLAS_USE_READER_CACHE` | `true` | Cache opened Atlas store readers in memory, avoiding re-opening the same `atlas.json` across queries. Settable at runtime as `beacon.atlas.use_reader_cache`. | | `BEACON_ATLAS_READER_CACHE_SIZE` | `32` | Max Atlas reader entries to keep cached. | +| `BEACON_ATLAS_USE_PRUNING` | `true` | Drop the datasets a predicate cannot match, from the collection's statistics, before reading them. A pure optimization: off trades throughput for skipping the pruning-index build. Settable at runtime as `beacon.atlas.use_pruning`. | ### Beacon Binary Format (BBF) | Variable | Default | Description | | --- | --- | --- | -| `BEACON_ENABLE_BBF_SPLIT_STREAMS_SLICE` | `false` | Split large batches into smaller slices for better memory use and parallelism on BBF queries. | +| `BEACON_ENABLE_BBF_SPLIT_STREAMS_SLICE` | `false` | Split large batches into smaller slices for better memory use and parallelism on BBF queries. Settable at runtime as `beacon.bbf.split_streams_slice`. | + +### Managed Lance tables + +These tune the [managed table](/docs/2.0.0-rc2/server/index) engine. An empty +value leaves the choice to Lance. + +The first four apply when Beacon **writes** a table. They shape the files a +`CREATE TABLE` or `INSERT` produces, and never the files already on disk. Change +one and rewrite the table to apply it. `BEACON_LANCE_MATERIALIZATION` is a read +setting and applies to the next query. + +| Variable | Default | Description | +| --- | --- | --- | +| `BEACON_LANCE_COMPRESSION` | _(none)_ | Block compression for string and binary columns: `fsst`, `zstd`, `lz4`, or `none`. Numeric columns are never compressed by this setting. Settable at runtime as `beacon.lance.compression`. | +| `BEACON_LANCE_NUMERIC_COMPRESSION` | _(Lance default)_ | Block compression for numeric columns: `zstd`, `lz4`, or `none`. `none` also turns off bitpacking and RLE, which usually measures **larger**. Settable at runtime as `beacon.lance.numeric_compression`. | +| `BEACON_LANCE_VERSION` | `2.2` | Lance file format version for new data: `2.0`, `2.1` or `2.2`. An append keeps the version of the table it appends to. Settable at runtime as `beacon.lance.version`. | +| `BEACON_LANCE_MINICHUNK` | _(Lance default)_ | Minichunk size, in bytes, for fixed-width columns. A size of 32KB and up needs version `2.2`. Settable at runtime as `beacon.lance.minichunk`. | +| `BEACON_LANCE_MATERIALIZATION` | _(width rule)_ | Column materialization on a filtered scan: `late` or `early`. Unset keeps Beacon's rule, which reads late above 16 projected columns. Settable at runtime as `beacon.lance.materialization`. | ## API documentation metadata diff --git a/integration-tests/test_settings.py b/integration-tests/test_settings.py new file mode 100644 index 00000000..8f1f60da --- /dev/null +++ b/integration-tests/test_settings.py @@ -0,0 +1,147 @@ +"""Runtime settings: `SET`, `RESET`, `ALTER SYSTEM` and `SHOW SETTINGS`. + +These guard the SQL surface of issue #359 against a live server — that a setting +changes without a restart, that the `beacon.` prefix reaches DataFusion's own +options, and that a setting nobody can change at runtime says so. + +`SET` is server-global and super-user-only, so the writes here go through the +admin credential and every test puts what it touched back. +""" + +from __future__ import annotations + +import pytest + + +def setting(client, name: str) -> str: + """The value a setting currently holds, read the way a client would. + + `information_schema` is super-user-only, hence `admin=True`. + """ + return client.scalar( + f"SELECT value FROM information_schema.df_settings WHERE name = '{name}'", + admin=True, + ) + + +@pytest.fixture +def restore_settings(client): + """Puts the settings this module touches back the way it found them. + + A `SET` applies to the whole server, so a test that left one changed would + leak into every later test in the session. + """ + keys = [ + "beacon.default_table", + "beacon.sql.stream_coalesce.target_rows", + "beacon.netcdf.use_rust_reader", + "beacon.batch_size", + ] + yield + for key in keys: + client.execute(f"RESET {key}") + + +def test_set_changes_a_setting_without_a_restart(client, restore_settings): + assert setting(client, "beacon.default_table") == "default" + + client.execute("SET beacon.default_table = 'observations'") + assert setting(client, "beacon.default_table") == "observations" + + +def test_set_and_show_agree(client, restore_settings): + client.execute("SET beacon.sql.stream_coalesce.target_rows = 1024") + + rows = client.sql_rows("SHOW beacon.sql.stream_coalesce.target_rows", admin=True) + # header + one row: the name and its value. + assert len(rows) == 2 + assert rows[1][0] == "beacon.sql.stream_coalesce.target_rows" + assert rows[1][1] == "1024" + + +def test_the_beacon_prefix_reaches_datafusion_options(client, restore_settings): + """`beacon.` is a complete alias for `datafusion.`, and both keep working.""" + client.execute("SET beacon.execution.batch_size = 8192") + assert setting(client, "datafusion.execution.batch_size") == "8192" + + # The documented BEACON_BATCH_SIZE spelling lands in the same option. + client.execute("SET beacon.batch_size = 4096") + assert setting(client, "datafusion.execution.batch_size") == "4096" + + client.execute("SET datafusion.execution.batch_size = 2048") + assert setting(client, "datafusion.execution.batch_size") == "2048" + + +def test_reset_restores_the_startup_value(client, restore_settings): + client.execute("SET beacon.netcdf.use_rust_reader = true") + assert setting(client, "beacon.netcdf.use_rust_reader") == "true" + + client.execute("RESET beacon.netcdf.use_rust_reader") + assert setting(client, "beacon.netcdf.use_rust_reader") == "false" + + +def test_a_query_still_runs_with_a_changed_setting(client, sample_data, restore_settings): + """Changing a knob must not change an answer.""" + obs = "read_parquet(['obs/*.parquet'])" + before = client.count(f"SELECT * FROM {obs}") + + client.execute("SET beacon.sql.stream_coalesce.target_rows = 1024") + client.execute("SET beacon.execution.batch_size = 512") + + assert client.count(f"SELECT * FROM {obs}") == before + + +def test_a_startup_only_setting_is_refused(client): + with pytest.raises(Exception) as excinfo: + client.execute("SET beacon.port = 1234") + assert "BEACON_PORT" in str(excinfo.value) + + +def test_an_unknown_setting_points_at_show_settings(client): + with pytest.raises(Exception) as excinfo: + client.execute("SET beacon.nonsense = 1") + assert "SHOW SETTINGS" in str(excinfo.value) + + +def test_show_settings_is_readable_without_admin(client): + """The issue's "a user cannot discover which settings exist": `SHOW SETTINGS` + is the one settings surface a non-super-user can read.""" + rows = client.sql_rows("SHOW SETTINGS", admin=False) + names = [row[0] for row in rows[1:]] + + assert "beacon.default_table" in names + assert "beacon.netcdf.use_rust_reader" in names + # Only the beacon namespace: the DataFusion half stays in df_settings. + assert all(name.startswith("beacon.") for name in names) + # Every row carries a description, so the listing documents itself. + assert all(row[3] for row in rows[1:]) + + +def test_changing_a_setting_needs_admin(client): + with pytest.raises(Exception): + client.execute("SET beacon.default_table = 'observations'", admin=False) + assert setting(client, "beacon.default_table") == "default" + + +def test_alter_system_applies_immediately(client): + """The persistent form still takes effect now. The restart half is covered by + the Rust suite, which can restart a runtime in-process.""" + try: + client.execute("ALTER SYSTEM SET beacon.default_table = 'observations'") + assert setting(client, "beacon.default_table") == "observations" + finally: + client.execute("ALTER SYSTEM RESET beacon.default_table") + assert setting(client, "beacon.default_table") == "default" + + +def test_beacon_system_settings_table(client, restore_settings): + """The table form of `SHOW SETTINGS`, for a client that would rather filter.""" + client.execute("SET beacon.default_table = 'observations'") + + rows = client.sql_rows( + 'SELECT value, "default" FROM beacon.system.settings ' + "WHERE name = 'beacon.default_table'", + admin=True, + ) + assert rows[1][0] == "observations" + assert rows[1][1] == "default"