DataFusion Session

Embed Algeon directly into DataFusion when the hosting Rust service owns request dispatch, authentication, catalog registration, and SessionContext lifecycles. Algeon adds native physical planning and GPU execution without altering host-level application concerns.

Dependencies

Use the release adapter and its compatible DataFusion dependency range:

[dependencies]
algeon-datafusion = { version = "0.1.1", features = ["cugraph", "cuvs", "cuml"] }
datafusion = "55.1"
futures = "0.3"

The dependency ranges permit Algeon 0.1.x and DataFusion 55.x updates. Build native libraries from the source tag matching the resolved Algeon version recorded in the application's Cargo.lock, such as v0.1.1 for 0.1.1. An Algeon version update requires native libraries from the corresponding release. Follow Building for an Application for the native build order and the application's build and runtime loader configuration. That guide pins 0.1.1. Update its source tag and adapter version together when using another Algeon release. The recipe requires the v0.1.1 tag and the crates.io packages.

The algeon-datafusion crate has no default features. Add cugraph for graph SQL, cuvs for vector SQL, cuml for machine-learning SQL, and nvml for optional device diagnostics. Omit all of them for relational cuDF execution only.

Build one backend

AlgeonGpuBackend is the process-scoped owner of CUDA devices, admission, per-device ledgers, and execution resources:

use datafusion::execution::SessionStateBuilder;
use datafusion::prelude::SessionContext;
use algeon_datafusion::backend::{
    GpuMemoryOwnership, AlgeonGpuBackend, AlgeonGpuDeviceConfig, AlgeonGpuDeviceProfile,
};

let device = AlgeonGpuDeviceConfig::new(AlgeonGpuDeviceProfile::new(0), GpuMemoryOwnership::WholeDeviceExclusive);
let backend = AlgeonGpuBackend::builder()
    .devices([device])
    .build()?;

let state = backend
    .install_on(SessionStateBuilder::new_with_default_features())?
    .build();

let ctx = SessionContext::new_with_state(state);

Every selected ordinal needs one profile and one matching ownership declaration. AlgeonGpuDeviceProfile::new(0) uses automatic capacity and one active attempt. Add explicit limits only when the target deployment requires custom bounds.

Only one live backend may exist in a process. Clone and share this backend across sessions and frontends; do not construct multiple backend instances within one process.

Add cuGraph, cuVS, and cuML SQL

install_on installs backend-bound native rules and coverage. It does not install cuGraph, cuVS, or cuML SQL. To enable the cugraph_* UDTFs on the session, chain try_with_cugraph_sql(config):

use algeon_datafusion::{
    cugraph_sql::CugraphSqlConfig,
    session::AlgeonSessionStateBuilderExt,
};

let state = backend
    .install_on(SessionStateBuilder::new_with_default_features())?
    .try_with_cugraph_sql(CugraphSqlConfig::default())?
    .build();

let ctx = SessionContext::new_with_state(state);

cuGraph UDTFs resolve the backend's admission and native execution configuration from session state when executed. To add the cuvs_* UDTFs, chain try_with_cuvs_sql() in the same manner. It takes no configuration. The UDTFs execute only when the crate was built with the cuvs feature. For cuml_* UDTFs, chain try_with_cuml_sql() as well. Execution requires the cuml feature. It installs no models root, so cuml_forest_predict is registered but stays unavailable with model_root_unconfigured. To enable it, chain try_with_cuml_sql_config instead with a CumlSqlConfig whose model_store: Option<CumlModelStoreConfig> names the root. try_with_cuml_sql() is equivalent to try_with_cuml_sql_config(CumlSqlConfig::default()):

use algeon_datafusion::cuml_sql::{CumlModelStoreConfig, CumlSqlConfig};

let state = backend
    .install_on(SessionStateBuilder::new_with_default_features())?
    .try_with_cuml_sql_config(CumlSqlConfig {
        model_store: Some(CumlModelStoreConfig::new("/srv/models")),
    })?
    .build();

CumlModelStoreConfig::new defaults max_model_bytes to CumlModelStoreConfig::DEFAULT_MAX_MODEL_BYTES (256 MiB). Set the field directly for a different cap. See cuVS execution and cuML execution.

Execute one attempt per query

For application-owned endpoints, create an attempt before planning, take its completion receiver before running the query, and read the report whether or not the query succeeded. The snippet assumes an orders table is already registered in ctx, for example with CREATE EXTERNAL TABLE ... STORED AS PARQUET or register_parquet:

use algeon_datafusion::attempt::{AttemptOptions, QueryAttemptReport};
use futures::TryStreamExt;
use std::sync::Arc;

let attempt = backend.begin_attempt(AttemptOptions {
    query_id: Some("request-42".to_owned()),
    ..Default::default()
})?;
let completion = attempt.completion();

let result = async {
    let stream = attempt.sql(&ctx, "SELECT SUM(amount) FROM orders").await?;
    stream.try_collect::<Vec<_>>().await
}
.await;

let report: Arc<QueryAttemptReport> = completion.await;
let batches = result?;

The attempt clones the session into invocation-local planning state. Its report records the final disposition, admission grant, timing, and terminal outcome. Completion seals on stream EOF, error, cancellation, or drop. Taking the receiver before planning and awaiting it after query completion guarantees the report remains accessible across both successful runs and error paths. Returning early with ? on the stream would drop the attempt and route the report to the abandoned-report queue.

Higher-level frameworks can manage this boundary. The Algeon Flight SQL server, for example, begins the attempt while issuing a statement ticket and carries it through DoGet. Call begin_attempt directly when an application directly coordinates query invocations.

Runtime requirements

The caller supplies the Tokio runtime that polls attempt.sql(...) and the result stream. The backend runs GPU work on threads it creates itself (Runtimes and threads).

• The runtime needs the Tokio time driver for admission waits and execution deadlines.
• algeon_server and the examples use the multi-thread flavor of #[tokio::main]; the crate's GPU tests use the current_thread flavor of #[tokio::test]. Use multi-thread in a serving process.
• An attempt's GPU work runs on the per-device algeon-gpu-{d} runtime and pins no caller worker. The attempt.sql(...) future and the stream need no spawn_blocking wrapper.
• Dropping a stream cancels its attempt without blocking the worker that drops it. An unread, undropped stream keeps its grant and device-resident result.
• close_and_wait_async waits on a blocking task. Only one backend can be live per process, and a replacement can be built after the close completes.

Admission configuration

Device capacity belongs to profiles; shared queue behavior belongs to AlgeonGpuQueuePolicy. This fragment replaces the backend construction shown above:

use algeon_datafusion::backend::{
    GpuDeviceMemoryLimit, GpuMemoryOwnership, AlgeonGpuBackend, AlgeonGpuDeviceConfig,
    AlgeonGpuDeviceProfile, AlgeonGpuQueuePolicy,
};

let device = AlgeonGpuDeviceConfig::new(
    AlgeonGpuDeviceProfile::new(0)
        .with_memory_limit(GpuDeviceMemoryLimit::Explicit {
            limit_bytes: 8 * 1024 * 1024 * 1024,
        })
        .with_backend_reserve_bytes(512 * 1024 * 1024)
        .with_default_query_limit_bytes(512 * 1024 * 1024)
        .with_max_active_attempts(10),
    GpuMemoryOwnership::WholeDeviceExclusive,
);
let backend = AlgeonGpuBackend::builder()
    .devices([device])
    .queue_policy(
        AlgeonGpuQueuePolicy::default()
            .with_max_queued_attempts(64)
            .with_max_queue_wait(std::time::Duration::from_secs(30))
    )
    .build()?;

Planning does not acquire a GPU grant. Admission begins when the first native stream is polled, and all native nodes in that attempt share the selected device and immutable grant. A full queue or expired wait returns structured service_overloaded backpressure. Eligible waiters are served in FIFO order. There is no overtaking policy.

Native execution policy is owned solely by the backend builder. Configure native execution with with_native_execution_config(...) on AlgeonGpuBackend::builder() before build().

DataFusion sessions carry standard DataFusion configuration only. Native execution policy and graph input bounds cannot be altered via SQL SET or session extensions.

Choose an execution mode

Set AlgeonGpuExecutionMode on the shared backend before installing sessions:

Rust mode	Behavior	Use it when
NativePreferred (default)	Select supported relational SQL for native GPU execution. Retain DataFusion for unsupported or unselected candidates. Explicit GPU UDTFs still run on GPU.	Accelerate relational SQL together with graph, vector, and machine learning functions on the GPU.
FunctionsOnly	Keep ordinary relational SQL on DataFusion CPU. Run explicit cugraph_*, cuvs_*, and cuml_* UDTFs on GPU.	Execute scans, joins, filters, and aggregations on DataFusion CPU while delegating GPU UDTFs to hardware.
NativeRequired	Attempt native relational selection, then reject any completed plan containing DataFusion CPU execution.	Enforce plans containing zero DataFusion CPU execution operators.

This replaces the backend construction above, using the same device:

use algeon_datafusion::planner::AlgeonGpuExecutionMode;

let backend = AlgeonGpuBackend::builder()
    .with_execution_mode(AlgeonGpuExecutionMode::FunctionsOnly)
    .devices([device])
    .build()?;

Install the required UDTF families with the features and session methods shown above. FunctionsOnly changes relational planning. It does not register UDTFs or remove their GPU runtime and backend requirements. The Flight SQL equivalents are documented in Configuration.

Compose CPU and GPU work

In FunctionsOnly, a CPU relation can feed a GPU UDTF whose Arrow results return to CPU SQL. A GPU result can also pass through CPU SQL into another GPU UDTF. cuGraph inputs remain registered table or view names; cuVS and cuML inputs remain parenthesized subqueries. Supported compositions share the same backend admission and immutable query grant.

CPU-to-GPU input boundaries stage Arrow data in host memory before upload. The current staging path does not spill: cuGraph host input and output staging each default to a 1 GiB ceiling. cuVS host input staging has a cumulative 1 GiB ceiling across one operation's inputs. cuML uses the same staging boundary. DataFusion's memory pool can impose a lower limit. These are host staging limits, separate from the GPU query cap. Keep intermediate inputs bounded. See the runnable PageRank example for PageRank with a CPU join and ordering.

DataFusion fallback and GPU-only plans

An explicit GPU UDTF has no CPU algorithm fallback. Unsupported UDTF inputs, memory exhaustion, cancellation, and runtime failures remain errors. After GPU execution starts, Algeon never replays the query on CPU. See Handling Errors for the structured error contract. NativeRequired checks the completed plan without replanning under a different policy.

Use GPU Coverage Validation to inspect the final planning disposition before execution. validate_query(&ctx, sql) is the Rust entry point; algeon_explain_coverage(...) exposes the same planning evidence through SQL.

Observe and shut down

Per-attempt evidence comes from QueryAttemptReport. The backend's diagnostics() reads the admission snapshot and the device ledger, cache, memory-failure, and execution-pool reports in one call. backend_metrics_snapshot() and device_telemetry_snapshots() stay separate. At shutdown, stop accepting work, drain streams and sessions, then call close_and_wait on the shared backend.

See the runnable Parquet query example for the complete per-query path.

Direct optimizer install

The raw session extension is useful for planning and coverage tests:

use datafusion::execution::SessionStateBuilder;
use algeon_datafusion::{
    planner::AlgeonNativeOptimizerConfig,
    session::AlgeonSessionStateBuilderExt,
};

let state = SessionStateBuilder::new_with_default_features()
    .with_algeon_native(AlgeonNativeOptimizerConfig::default())
    .build();

This path has no admission owner. Because native execution fails locally with gpu_backend_required, production GPU execution must always install an AlgeonGpuBackend.

The public optimizer rule name is also its deduplication key. Use ALGEON_NATIVE_PHYSICAL_OPTIMIZER_RULE_NAME when registering or deduplicating the rule.