oxidelake_runtime/
session.rs

1//! The user-facing session: one query API over embedded and cluster execution.
2
3use std::sync::Arc;
4
5use ballista::prelude::SessionContextExt;
6use ballista_core::extension::SessionConfigExt;
7use datafusion::dataframe::DataFrame;
8use datafusion::execution::SessionStateBuilder;
9use datafusion::physical_plan::displayable;
10use datafusion::prelude::{SessionConfig, SessionContext};
11use oxidelake_compute::{local_backend, oxide_udfs};
12use oxidelake_core::telemetry::TelemetryHub;
13use oxidelake_core::{BackendKind, EngineError};
14use oxidelake_planner::{HardwarePlacementRule, physical_optimizer_rules};
15use oxidelake_storage::{
16    default_object_store, register_local_store, register_parquet_table, with_gpu_batch_size,
17    with_pruning,
18};
19
20use crate::cluster;
21
22/// Where a session executes.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub enum SessionMode {
25    /// In-process DataFusion with the placement rule targeting the local backend.
26    Embedded {
27        /// The backend detected on this machine.
28        target: BackendKind,
29    },
30    /// A Ballista cluster reached through its scheduler URL (`df://host:port`).
31    Cluster {
32        /// The scheduler URL.
33        scheduler_url: String,
34    },
35}
36
37/// An OxideLake session.
38pub struct OxideSession {
39    ctx: SessionContext,
40    mode: SessionMode,
41    telemetry: Arc<TelemetryHub>,
42}
43
44impl OxideSession {
45    /// Creates an embedded session: Parquet pruning on, the local object store
46    /// registered, the SQL UDFs, and the placement rule targeting the detected
47    /// backend.
48    pub fn local() -> Result<Self, EngineError> {
49        Self::local_with_target(local_backend()?.kind())
50    }
51
52    /// An embedded session whose placement rule targets `target` instead of
53    /// the detected backend. Placement is a planning decision: operators still
54    /// select the real local backend at `execute()` time and fall back to the
55    /// CPU reference, so planning for an absent GPU is safe — it is exactly
56    /// what cluster executors do with the scheduler's plans.
57    pub fn local_with_target(target: BackendKind) -> Result<Self, EngineError> {
58        let telemetry = TelemetryHub::new();
59        let rule = HardwarePlacementRule::new(target).with_telemetry(Arc::clone(&telemetry));
60        let mut config = with_pruning(SessionConfig::new());
61        if target.is_gpu() {
62            config = with_gpu_batch_size(config);
63        }
64        let state = SessionStateBuilder::new()
65            .with_default_features()
66            .with_config(config)
67            .with_physical_optimizer_rules(physical_optimizer_rules(rule))
68            .build();
69        let ctx = SessionContext::new_with_state(state);
70        register_local_store(&ctx, default_object_store());
71        for udf in oxide_udfs() {
72            ctx.register_udf(udf.as_ref().clone());
73        }
74        Ok(Self {
75            ctx,
76            mode: SessionMode::Embedded { target },
77            telemetry,
78        })
79    }
80
81    /// Connects to a Ballista scheduler (`df://host:port`). The session carries
82    /// OxideLake's plan codec so `Gpu*Exec` nodes survive the trip to executors,
83    /// and the SQL UDFs so queries plan client-side; placement itself happens
84    /// on the scheduler.
85    pub async fn connect(scheduler_url: &str) -> Result<Self, EngineError> {
86        let config = with_pruning(SessionConfig::new_with_ballista())
87            .with_ballista_physical_extension_codec(cluster::oxide_codec());
88        let state = SessionStateBuilder::new()
89            .with_default_features()
90            .with_config(config)
91            .build();
92        let ctx = SessionContext::remote_with_state(scheduler_url, state).await?;
93        for udf in oxide_udfs() {
94            ctx.register_udf(udf.as_ref().clone());
95        }
96        Ok(Self {
97            ctx,
98            mode: SessionMode::Cluster {
99                scheduler_url: scheduler_url.to_owned(),
100            },
101            telemetry: TelemetryHub::new(),
102        })
103    }
104
105    /// The execution mode.
106    pub fn mode(&self) -> &SessionMode {
107        &self.mode
108    }
109
110    /// The underlying DataFusion context.
111    pub fn ctx(&self) -> &SessionContext {
112        &self.ctx
113    }
114
115    /// The telemetry hub for this session (local process only).
116    pub fn telemetry(&self) -> &Arc<TelemetryHub> {
117        &self.telemetry
118    }
119
120    /// Plans a SQL statement into a lazily executed [`DataFrame`].
121    pub async fn sql(&self, query: &str) -> Result<DataFrame, EngineError> {
122        Ok(self.ctx.sql(query).await?)
123    }
124
125    /// Registers a Parquet file or directory as `name`.
126    pub async fn register_parquet(&self, name: &str, path: &str) -> Result<(), EngineError> {
127        register_parquet_table(&self.ctx, name, path).await
128    }
129
130    /// The indented physical plan for `query`, with placement tags in embedded
131    /// mode. In cluster mode this is the client-side plan (`DistributedQueryExec`);
132    /// the scheduler's plan is what carries the tags there.
133    pub async fn explain(&self, query: &str) -> Result<String, EngineError> {
134        let plan = self.ctx.sql(query).await?.create_physical_plan().await?;
135        Ok(displayable(plan.as_ref()).indent(true).to_string())
136    }
137}
138
139impl std::fmt::Debug for OxideSession {
140    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
141        f.debug_struct("OxideSession")
142            .field("mode", &self.mode)
143            .field("session_id", &self.ctx.session_id())
144            .finish_non_exhaustive()
145    }
146}
147
148#[cfg(test)]
149#[allow(clippy::unwrap_used, clippy::expect_used)]
150mod tests {
151    use super::*;
152
153    #[tokio::test]
154    async fn embedded_session_runs_sql_and_reports_mode() {
155        let session = OxideSession::local().unwrap();
156        assert!(matches!(session.mode(), SessionMode::Embedded { .. }));
157        let batches = session
158            .sql("SELECT 1 + 1 AS two")
159            .await
160            .unwrap()
161            .collect()
162            .await
163            .unwrap();
164        assert_eq!(batches.len(), 1);
165        assert_eq!(batches[0].num_rows(), 1);
166        let text = session.explain("SELECT 1 + 1 AS two").await.unwrap();
167        assert!(text.contains("ProjectionExec"), "{text}");
168    }
169}
oxidelake_runtime/session.rs

oxidelake_runtime/
session.rs
