skade-flight/src/engine.rs
//! The ONE long-lived DataFusion session over the warehouse.
//!
//! skade-sql built a fresh `SessionContext` per query and `collect()`ed the
//! result. Right for an embedded one-shot; wrong for a server that answers
//! Power BI refreshes all day. Here:
//!
//! - **One session**, built at open, held behind an `ArcSwap` so readers never
//! lock. The Iceberg providers snapshot the catalog when built.
//! - **Refresh, not rebuild-per-query**: `maybe_refresh` rebuilds the session
//! when it is older than `refresh_after` (default 5 s), so a new commit is
//! visible within that window and a burst of refreshes shares one rebuild.
//! The rebuild is a metadata read (the catalog's current pointers), not a
//! data read.
//! - **Plan once**: Flight SQL asks for the schema (`GetFlightInfo`) and then
//! the data (`DoGet`) as two RPCs. The statement is planned at the first —
//! logical AND physical, so the optimizer runs once — and cached (bounded,
//! keyed by the SQL text plus any bound parameters, cleared on refresh); the
//! second only executes the cached [`ExecutionPlan`]. Prepared statements
//! hit the same cache. [`Engine::plan_counts`] is the measurement: one
//! logical and one physical plan per distinct statement, not one per RPC.
//! - **Stream, never collect**: [`Engine::execute_stream`] hands the
//! `SendableRecordBatchStream` straight to the Flight encoder; batches leave
//! as scan partitions complete.
//! - **`target_partitions` = cores** so the Iceberg scans fan out (the
//! vendored provider gives DataFusion one partition per data file).
//!
//! Naming: the Iceberg catalog is registered as the DataFusion catalog
//! `skade`; each Iceberg namespace is a schema; skade's default namespace
//! (`main`) is the session default, so `SELECT * FROM t` resolves unqualified
//! and `skade.<ns>.<table>` always works.
//!
//! **Session views** (`CREATE VIEW` / `DROP VIEW`, TPC-H Q15's shape) are held
//! in a [`ViewStore`] outside the session and laid over the default schema by
//! an [`OverlayCatalog`], so they survive the rebuild and never reach the
//! catalog — see `views.rs` for the decision. DDL is classified BEFORE the plan
//! cache and applied only on the execution RPC; every other write is refused
//! with the rewrite in the message.
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use arc_swap::ArcSwap;
use arrow_schema::SchemaRef;
use datafusion::dataframe::DataFrame;
use datafusion::execution::SendableRecordBatchStream;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::scalar::ScalarValue;
use datafusion::prelude::{SessionConfig, SessionContext};
use iceberg::Catalog;
use skade::Warehouse;
use skade::connect::CatalogConnector;
use tokio::sync::Mutex;
use crate::error::{Error, Result};
use crate::views::{self, Classified, OverlayCatalog, ViewDdl, ViewStore};
/// The DataFusion catalog name the warehouse is mounted under.
pub const CATALOG: &str = "skade";
/// Upper bound on cached plans. Each is a `LogicalPlan` + schema; a BI
/// client's working set is tens of distinct statements, not thousands.
const PLAN_CACHE_CAP: usize = 256;
/// One table the session can see: `(schema, name, arrow schema)`.
#[derive(Debug, Clone)]
pub struct TableRef {
pub schema: String,
pub name: String,
pub arrow: SchemaRef,
}
/// Tunables. `Default` is what the binary and the tests use.
#[derive(Debug, Clone)]
pub struct EngineOptions {
/// Rebuild the session when older than this (a new commit becomes
/// visible within the window).
pub refresh_after: Duration,
/// DataFusion `target_partitions`.
pub target_partitions: usize,
/// DataFusion `batch_size`: rows per RecordBatch leaving the scan. Larger
/// batches mean fewer Flight frames and less per-batch overhead on the
/// wire; 8192 is DataFusion's default, 32768 measured better here.
pub batch_size: usize,
/// The Iceberg namespace unqualified names resolve in (the session's
/// default schema, where session views live). skade's own warehouses use
/// [`skade::DEFAULT_NAMESPACE`]; a foreign lake's namespace is its own.
pub default_namespace: String,
/// Refuse every write path — the update door and session-view DDL — with
/// [`Error::ReadOnly`]. For serving a lake another engine (Spark) writes.
pub read_only: bool,
}
impl Default for EngineOptions {
fn default() -> Self {
Self {
refresh_after: Duration::from_secs(5),
target_partitions: num_cpus(),
batch_size: 32_768,
default_namespace: skade::DEFAULT_NAMESPACE.to_string(),
read_only: false,
}
}
}
/// **A statement as the cache keys it**: the SQL text plus the parameter
/// values bound to its placeholders (`$1`…), in order. A prepared statement
/// executed with two different bindings is two plans — the optimizer folds a
/// bound literal, so the plans genuinely differ.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Statement {
pub sql: String,
pub params: Vec<ScalarValue>,
}
impl Statement {
pub fn sql(sql: impl Into<String>) -> Self {
Self { sql: sql.into(), params: Vec::new() }
}
}
/// A planned statement: its result schema, what executing it means, and —
/// once `GetFlightInfo` has asked — the physical plan `DoGet` executes.
pub struct Planned {
pub schema: SchemaRef,
/// The schema as the IPC message `FlightInfo.schema` carries — encoded
/// once per plan, cloned (refcounted `Bytes`) per GetFlightInfo.
pub schema_ipc: bytes::Bytes,
body: Body,
/// The optimized physical plan of a query, built once. `DataFrame::
/// execute_stream` would rebuild it (optimizer + physical planner) on
/// EVERY execute; the ticket path executes this instead.
physical: tokio::sync::OnceCell<Arc<dyn ExecutionPlan>>,
/// The placeholder types the planner inferred (`$1` → Int64 …), in
/// placeholder order. What a prepared statement advertises as its
/// `parameter_schema`.
pub parameters: Vec<(String, Option<arrow_schema::DataType>)>,
}
/// How many times the engine has planned. `logical` counts parse + resolve
/// (`create_logical_plan`); `physical` counts optimizer + physical planner
/// runs. The invariant the ticket path holds: one of each per distinct
/// [`Statement`], however many RPCs carry it.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct PlanCounts {
pub logical: u64,
pub physical: u64,
}
/// A query streams; a session-view DDL is applied (once, on execution) and
/// streams nothing.
enum Body {
Query(DataFrame),
View(ViewDdl),
}
pub struct Engine {
/// The catalog the session is built over — a skade warehouse's
/// `RedbCatalog`, or any `iceberg::Catalog` (an Iceberg REST catalog).
catalog: Arc<dyn Catalog>,
/// Present when opened over a skade [`Warehouse`] (held so its redb lock
/// lives as long as the engine); `None` for [`Engine::with_catalog`].
warehouse: Option<Arc<Warehouse>>,
ctx: ArcSwap<SessionContext>,
built_at: ArcSwap<Instant>,
opts: EngineOptions,
/// Serialises rebuilds so a burst of stale readers shares one.
rebuild: Mutex<()>,
/// Statement → plan, for the session in `ctx`. Cleared on every rebuild.
plans: StdMutex<HashMap<Statement, Arc<Planned>>>,
/// Session views, outside the session so they outlive its rebuilds.
views: Arc<ViewStore>,
logical_plans: AtomicU64,
physical_plans: AtomicU64,
}
impl Engine {
/// Build the session over an opened warehouse with default options.
pub async fn new(warehouse: Warehouse) -> Result<Self> {
Self::with_engine_options(warehouse, EngineOptions::default()).await
}
/// Back-compat shape: refresh window + partitions, default batch size.
pub async fn with_options(warehouse: Warehouse, refresh_after: Duration, target_partitions: usize) -> Result<Self> {
Self::with_engine_options(warehouse, EngineOptions { refresh_after, target_partitions, ..Default::default() }).await
}
pub async fn with_engine_options(warehouse: Warehouse, opts: EngineOptions) -> Result<Self> {
let warehouse = Arc::new(warehouse);
let catalog: Arc<dyn Catalog> = warehouse.catalog();
Self::build(catalog, Some(warehouse), opts).await
}
/// **Serve any Iceberg catalog** — an existing lake's catalog (opened by a
/// `skade::connect` connector, see [`with_connector`](Self::with_connector)),
/// or a catalog a library consumer already holds. No skade `Warehouse`
/// behind it: [`try_warehouse`](Self::try_warehouse) is `None`.
pub async fn with_catalog(catalog: Arc<dyn Catalog>, opts: EngineOptions) -> Result<Self> {
Self::build(catalog, None, opts).await
}
/// **Open through a connector** — `skade::connect::{SkadeLocal,
/// SkadeObjectStore, NessieV2, …}`: the one door every consumer shares. A
/// read-only connector forces `opts.read_only`; its kind and description
/// (never a secret) are logged.
pub async fn with_connector(connector: Arc<dyn CatalogConnector>, mut opts: EngineOptions) -> Result<Self> {
if connector.read_only() {
opts.read_only = true;
}
eprintln!(
"skade-flight: catalog {} — {}{}",
connector.kind(),
connector.describe(),
if opts.read_only { " (read-only)" } else { "" }
);
let catalog = connector.open().await?;
Self::with_catalog(catalog, opts).await
}
async fn build(catalog: Arc<dyn Catalog>, warehouse: Option<Arc<Warehouse>>, opts: EngineOptions) -> Result<Self> {
let views = Arc::new(ViewStore::default());
let ctx = build_session(&catalog, &opts, &views).await?;
Ok(Self {
catalog,
warehouse,
ctx: ArcSwap::from_pointee(ctx),
built_at: ArcSwap::from_pointee(Instant::now()),
opts,
rebuild: Mutex::new(()),
plans: StdMutex::new(HashMap::new()),
views,
logical_plans: AtomicU64::new(0),
physical_plans: AtomicU64::new(0),
})
}
/// The skade warehouse this engine was opened over.
///
/// # Panics
/// When the engine was built with [`Engine::with_catalog`] — a foreign
/// catalog has no skade `Warehouse`. Use [`try_warehouse`](Self::try_warehouse).
pub fn warehouse(&self) -> &Arc<Warehouse> {
self.warehouse.as_ref().expect(
"Engine::warehouse: this engine was built with Engine::with_catalog and has no skade Warehouse — use try_warehouse()",
)
}
/// The skade warehouse, if the engine was opened over one.
pub fn try_warehouse(&self) -> Option<&Arc<Warehouse>> {
self.warehouse.as_ref()
}
/// The Iceberg catalog the session is built over.
pub fn catalog(&self) -> &Arc<dyn Catalog> {
&self.catalog
}
/// Refuse a write path when the engine is read-only.
fn refuse_if_read_only(&self, what: &str) -> Result<()> {
if self.opts.read_only {
return Err(Error::ReadOnly(format!(
"{what} refused: this engine serves its catalog read-only (EngineOptions.read_only) — the lake's own writer owns every change"
)));
}
Ok(())
}
pub fn options(&self) -> &EngineOptions {
&self.opts
}
/// The current session (lock-free).
pub fn ctx(&self) -> Arc<SessionContext> {
self.ctx.load_full()
}
async fn rebuild_locked(&self) -> Result<()> {
let ctx = build_session(&self.catalog, &self.opts, &self.views).await?;
// Session views are re-planned against the NEW session so they read
// its snapshot; one that no longer plans is dropped, and said.
for line in views::replan(&self.views, &ctx, CATALOG, &self.opts.default_namespace).await {
eprintln!("skade-flight: {line}");
}
// Plans point at providers of the OLD session's snapshot: drop them
// before publishing the new session so no reader plans against one
// and executes against the other.
self.plans.lock().unwrap_or_else(|p| p.into_inner()).clear();
self.ctx.store(Arc::new(ctx));
self.built_at.store(Arc::new(Instant::now()));
Ok(())
}
/// Rebuild the session from the catalog's current state, unconditionally.
pub async fn refresh(&self) -> Result<()> {
let _g = self.rebuild.lock().await;
self.rebuild_locked().await
}
/// Rebuild only if the session is older than `refresh_after`. A second
/// caller arriving during a rebuild waits for it and then finds the
/// session fresh.
pub async fn maybe_refresh(&self) -> Result<()> {
if self.built_at.load().elapsed() < self.opts.refresh_after {
return Ok(());
}
let _g = self.rebuild.lock().await;
if self.built_at.load().elapsed() < self.opts.refresh_after {
return Ok(());
}
self.rebuild_locked().await
}
/// Plan a statement, through the cache. No execution — a session-view
/// DDL is only CLASSIFIED here (and never cached: each execution applies
/// it afresh, so `CREATE VIEW` twice is the error SQL says it is).
async fn planned(&self, stmt: &Statement) -> Result<Arc<Planned>> {
if let Some(p) = self.plans.lock().unwrap_or_else(|p| p.into_inner()).get(stmt) {
return Ok(p.clone());
}
let ctx = self.ctx();
let (plan, parameters) = self.logical(&ctx, &stmt.sql).await?;
let (schema, body) = match views::classify(plan, CATALOG, &self.opts.default_namespace)? {
Classified::Query(plan) => {
let plan = if !stmt.params.is_empty() {
plan.with_param_values(stmt.params.clone()).map_err(Error::from_df)?
} else if !parameters.is_empty() {
// Unbound placeholders plan (the schema is knowable) but
// cannot execute; DataFusion says so at execute time, which
// is late and reads as an internal error. Say it here, by name.
let names: Vec<&str> = parameters.iter().map(|(n, _)| n.as_str()).collect();
return Err(Error::Sql(format!(
"statement has unbound parameter(s) {}: bind them (DoPutPreparedStatementQuery) before executing",
names.join(", ")
)));
} else {
plan
};
let df = DataFrame::new(ctx.state(), plan);
// Decorated with the ODBC column metadata (columns.rs): this
// schema is what FlightInfo, the prepared dataset_schema and
// the stream's schema frame carry.
(crate::columns::decorate(df.schema().as_arrow(), None), Body::Query(df))
}
Classified::View(ddl) => (Arc::new(arrow_schema::Schema::empty()), Body::View(ddl)),
};
let opts = arrow_ipc::writer::IpcWriteOptions::default();
let arrow_flight::IpcMessage(schema_ipc) =
arrow_flight::SchemaAsIpc::new(schema.as_ref(), &opts).try_into().map_err(Error::Arrow)?;
let p = Arc::new(Planned { schema, schema_ipc, body, physical: tokio::sync::OnceCell::new(), parameters });
if matches!(p.body, Body::View(_)) {
return Ok(p);
}
let mut plans = self.plans.lock().unwrap_or_else(|p| p.into_inner());
if plans.len() >= PLAN_CACHE_CAP {
// Bounded, not LRU: a BI working set never reaches the cap in
// practice, and a full clear is cheaper than bookkeeping on every hit.
plans.clear();
}
plans.insert(stmt.clone(), p.clone());
Ok(p)
}
/// Parse + resolve, WITHOUT executing (`ctx.sql` would run a DDL right
/// here, on a GetFlightInfo that only asked what the statement looks
/// like), plus the placeholder types read off the plan before anything is
/// bound (bound, the placeholders are gone).
async fn logical(
&self,
ctx: &SessionContext,
sql: &str,
) -> Result<(datafusion::logical_expr::LogicalPlan, Vec<(String, Option<arrow_schema::DataType>)>)> {
// ODBC escapes and TOP n first (see odbc.rs); plain SQL passes byte-identical.
let rewritten = crate::odbc::rewrite(sql);
self.logical_plans.fetch_add(1, Ordering::Relaxed);
let plan = ctx.state().create_logical_plan(&rewritten).await.map_err(Error::from_df)?;
let mut parameters: Vec<(String, Option<arrow_schema::DataType>)> =
plan.get_parameter_types().map_err(Error::from_df)?.into_iter().collect();
parameters.sort_by_key(|(name, _)| placeholder_ordinal(name));
Ok((plan, parameters))
}
/// The physical plan of a planned query, built once per plan. A view DDL
/// has none (it is applied, not executed).
async fn physical(&self, p: &Planned) -> Result<Option<Arc<dyn ExecutionPlan>>> {
let Body::Query(df) = &p.body else { return Ok(None) };
p.physical
.get_or_try_init(|| async {
self.physical_plans.fetch_add(1, Ordering::Relaxed);
df.clone().create_physical_plan().await.map_err(Error::from_df)
})
.await
.map(|plan| Some(Arc::clone(plan)))
}
/// **Describe a statement for `GetFlightInfo`**: plan it fully (logical +
/// physical) so `DoGet` has nothing left to plan. Returns the cached plan.
pub async fn describe(&self, stmt: &Statement) -> Result<Arc<Planned>> {
let p = self.planned(stmt).await?;
self.physical(&p).await?;
Ok(p)
}
/// The placeholder types of a statement, in `$1, $2 …` order — what a
/// prepared statement's `parameter_schema` advertises — and its result
/// schema, what its `dataset_schema` advertises. Planned without binding
/// and without executing (a statement with unbound placeholders is
/// describable, not executable; a DDL is classified, not applied).
pub async fn prepared_shape(&self, sql: &str) -> Result<(SchemaRef, Vec<(String, Option<arrow_schema::DataType>)>)> {
let ctx = self.ctx();
let (plan, parameters) = self.logical(&ctx, sql).await?;
let schema = match views::classify(plan, CATALOG, &self.opts.default_namespace)? {
Classified::Query(plan) => crate::columns::decorate(plan.schema().as_arrow(), None),
Classified::View(_) => Arc::new(arrow_schema::Schema::empty()),
};
Ok((schema, parameters))
}
/// The plans made so far (see [`PlanCounts`]).
pub fn plan_counts(&self) -> PlanCounts {
PlanCounts {
logical: self.logical_plans.load(Ordering::Relaxed),
physical: self.physical_plans.load(Ordering::Relaxed),
}
}
/// Plan a statement (cached). What `GetFlightInfo` needs. A session-view
/// DDL has no DataFrame: it plans as an empty one.
pub async fn plan(&self, sql: &str) -> Result<DataFrame> {
match &self.planned(&Statement::sql(sql)).await?.body {
Body::Query(df) => Ok(df.clone()),
Body::View(_) => self.ctx().read_empty().map_err(Error::from_df),
}
}
/// Apply a session-view DDL, then forget every cached plan: a plan that
/// resolved the old view (or failed to resolve a name that now exists)
/// would otherwise be served from the cache.
async fn apply_view(&self, ddl: ViewDdl) -> Result<()> {
self.refuse_if_read_only("session-view DDL")?;
views::apply(&self.views, &self.ctx(), CATALOG, &self.opts.default_namespace, ddl).await?;
self.plans.lock().unwrap_or_else(|p| p.into_inner()).clear();
Ok(())
}
/// **The update door** (`CommandStatementUpdate`): a session-view DDL is
/// applied and reports 0 rows; a query is not an update; everything else
/// is the read-only refusal, with the rewrite in it.
pub async fn execute_update(&self, sql: &str) -> Result<i64> {
self.refuse_if_read_only("the update door (CommandStatementUpdate)")?;
let p = self.planned(&Statement::sql(sql)).await?;
match &p.body {
Body::View(ddl) => {
self.apply_view(ddl.clone()).await?;
Ok(0)
}
Body::Query(_) => Err(Error::Sql(
"that statement is a query, not an update — send it as a statement (GetFlightInfo + DoGet)".into(),
)),
}
}
/// The session views' names, sorted. Not what `GetTables` advertises.
pub fn view_names(&self) -> Vec<String> {
self.views.names()
}
/// The Arrow schema a statement returns, without running it.
pub async fn schema_of(&self, sql: &str) -> Result<SchemaRef> {
Ok(self.planned(&Statement::sql(sql)).await?.schema.clone())
}
/// The result schema as IPC bytes (`FlightInfo.schema`), plus the Arrow
/// schema. Both refcounted clones off the cached plan.
pub async fn schema_ipc_of(&self, sql: &str) -> Result<(SchemaRef, bytes::Bytes)> {
let p = self.planned(&Statement::sql(sql)).await?;
Ok((p.schema.clone(), p.schema_ipc.clone()))
}
/// Execute (planning through the cache), streaming. A session-view DDL
/// is applied here — this IS the execution — and streams nothing.
pub async fn execute_stream(&self, sql: &str) -> Result<SendableRecordBatchStream> {
let p = self.planned(&Statement::sql(sql)).await?;
match &p.body {
Body::Query(_) => {
let plan = fresh_instance(self.physical(&p).await?.expect("a query has a physical plan"))?;
datafusion::physical_plan::execute_stream(plan, self.ctx().task_ctx()).map_err(Error::from_df)
}
Body::View(ddl) => {
self.apply_view(ddl.clone()).await?;
self.ctx().read_empty().map_err(Error::from_df)?.execute_stream().await.map_err(Error::from_df)
}
}
}
/// Cached plan count (for tests and the bench).
pub fn cached_plans(&self) -> usize {
self.plans.lock().unwrap_or_else(|p| p.into_inner()).len()
}
/// Every table the session can see, with its Arrow schema. What
/// `GetTables` advertises — so a driver's `SELECT * FROM "t"` is exactly
/// what the planner resolves.
pub async fn tables(&self) -> Result<Vec<TableRef>> {
let ctx = self.ctx();
let mut out = Vec::new();
let Some(catalog) = ctx.catalog(CATALOG) else {
return Ok(out);
};
let mut schemas = catalog.schema_names();
schemas.sort();
for schema_name in schemas {
let Some(schema) = catalog.schema(&schema_name) else { continue };
let mut names = schema.table_names();
names.sort();
for name in names {
// iceberg-datafusion lists `t$snapshots` / `t$manifests` metadata
// tables beside every table. They stay queryable by name, but a BI
// user browsing the catalog wants the data tables, not two rows of
// plumbing per table.
if name.contains('$') {
continue;
}
// A session view is not a warehouse table (views.rs says why
// the navigator is not offered one).
if self.views.contains(&name) {
continue;
}
// A table listed a moment ago can vanish under a concurrent
// compaction swap; skip it rather than fail the listing.
let Ok(Some(provider)) = schema.table(&name).await else { continue };
let arrow = crate::columns::decorate(provider.schema().as_ref(), Some((CATALOG, &schema_name, &name)));
out.push(TableRef { schema: schema_name.clone(), name, arrow });
}
}
Ok(out)
}
/// The schema (namespace) names under the catalog.
pub fn schema_names(&self) -> Vec<String> {
let ctx = self.ctx();
let mut v = ctx.catalog(CATALOG).map(|c| c.schema_names()).unwrap_or_default();
v.sort();
v
}
/// The table's Iceberg identifier fields, in declaration order — its
/// logical primary key. Empty when the table declares none or does not
/// exist (a driver asking about an unknown table gets an empty set, not an
/// error; that is what `GetPrimaryKeys` semantics are).
pub async fn primary_keys(&self, schema: &str, table: &str) -> Result<Vec<crate::keys::PkColumn>> {
use iceberg::TableIdent;
// A flattened schema name (`a.b`) is the namespace's elements.
let ident = TableIdent::new(crate::namespaces::namespace_of(schema)?, table.to_string());
let t = match self.catalog.load_table(&ident).await {
Ok(t) => t,
Err(e) if e.kind() == iceberg::ErrorKind::TableNotFound => return Ok(Vec::new()),
Err(e) => return Err(e.into()),
};
let s = t.metadata().current_schema();
let mut out = Vec::new();
for (i, id) in s.identifier_field_ids().enumerate() {
if let Some(f) = s.field_by_id(id) {
out.push(crate::keys::PkColumn { column: f.name.clone(), sequence: i as i32 + 1 });
}
}
Ok(out)
}
}
async fn build_session(catalog: &Arc<dyn Catalog>, opts: &EngineOptions, views: &Arc<ViewStore>) -> Result<SessionContext> {
let config = SessionConfig::new()
.with_target_partitions(opts.target_partitions.max(1))
.with_batch_size(opts.batch_size.max(1024))
.with_default_catalog_and_schema(CATALOG, opts.default_namespace.as_str())
.with_information_schema(true);
// Through time_literal's constructor: its rule must run BEFORE TypeCoercion.
let ctx = crate::time_literal::session_context(config);
// Every namespace, at any depth, as one schema named `a.b` (namespaces.rs).
let (provider, skipped) = crate::namespaces::FlatNamespaceCatalog::try_new(Arc::clone(catalog)).await?;
for line in skipped {
eprintln!("skade-flight: {line}");
}
// The Iceberg catalog, with the session views laid over its default schema.
let overlay = OverlayCatalog::new(Arc::new(provider), &opts.default_namespace, Arc::clone(views));
ctx.register_catalog(CATALOG, Arc::new(overlay));
Ok(ctx)
}
fn num_cpus() -> usize {
std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1)
}
impl Engine {
/// Execute (planning through the cache) as one stream PER PARTITION, so
/// the caller can encode partitions on separate tasks. A sorted plan
/// yields one partition (order kept); an unsorted one yields as many as
/// the scan has data files.
pub async fn execute_stream_partitioned(&self, sql: &str) -> Result<Vec<SendableRecordBatchStream>> {
self.execute_statement_partitioned(&Statement::sql(sql)).await
}
/// [`Self::execute_stream_partitioned`] for a [`Statement`] (bound
/// parameters included). The plan cached by [`Self::describe`] is executed
/// as-is: nothing is parsed, optimized or planned here on a cache hit.
pub async fn execute_statement_partitioned(&self, stmt: &Statement) -> Result<Vec<SendableRecordBatchStream>> {
let p = self.planned(stmt).await?;
self.execute_planned(&p).await
}
/// Execute an already-described plan, one stream per partition. This is
/// the whole of `DoGet`: no parse, no optimizer, no physical planner. A
/// session-view DDL is applied here — this IS the execution — and streams
/// nothing (the encoder sends the empty schema frame and ends).
pub async fn execute_planned(&self, p: &Planned) -> Result<Vec<SendableRecordBatchStream>> {
match &p.body {
Body::Query(_) => {
let plan = fresh_instance(self.physical(p).await?.expect("a query has a physical plan"))?;
datafusion::physical_plan::execute_stream_partitioned(plan, self.ctx().task_ctx()).map_err(Error::from_df)
}
Body::View(ddl) => {
self.apply_view(ddl.clone()).await?;
Ok(Vec::new())
}
}
}
}
/// **A physical plan is re-executable only as a fresh instance.** Several
/// operators keep per-instance state across their partitions —
/// `RepartitionExec` hands out its channels once and panics ("partition not
/// used yet") on a second execute, `HashJoinExec` builds its hash table once.
/// Rebuilding the tree through `replace_children` re-instantiates every
/// operator (fresh state, same optimized shape, no planner run) and leaves the
/// cached instance untouched. The leaves (scans) carry no execute state and
/// are kept as they are.
fn fresh_instance(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> {
let children: Vec<Arc<dyn ExecutionPlan>> =
plan.children().into_iter().map(|c| fresh_instance(Arc::clone(c))).collect::<Result<_>>()?;
if children.is_empty() {
return Ok(plan);
}
// Same children shapes ⇒ the cached properties still hold; nothing is recomputed.
use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions};
plan.replace_children(children, ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep)).map_err(Error::from_df)
}
/// `$3` → 3, `$abc` / anything else → `u32::MAX` (sorted last, stably).
fn placeholder_ordinal(name: &str) -> u32 {
name.trim_start_matches('$').parse().unwrap_or(u32::MAX)
}