Struct SqlEngine 

pub struct SqlEngine { /* private fields */ }

Implementations

impl SqlEngine

pub fn new() -> Self

Create a local SQL engine.

Window helper UDFs (tumble_start, tumble_end, hop_start, hop_end) are registered as part of construction. If registration fails the engine is still returned — non-window queries work — and a tracing::warn! is emitted. Use SqlEngine::try_new when callers need to surface the registration error.

DataFusion target_partitions defaults to the machine’s available CPU parallelism (matching a bare datafusion::SessionContext), overridable via KRISHIV_TARGET_PARALLELISM or SqlEngine::with_target_parallelism. Executor task engines deliberately scale this down to their per-slot share so concurrent tasks don’t oversubscribe the machine.

pub fn new_with_memory_limit(memory_limit_bytes: Option<usize>) -> Self

Create a local SQL engine whose DataFusion execution memory is capped at memory_limit_bytes.

When Some, the engine runs with a FairSpillPool of that size plus the default disk manager, so memory-intensive operators (sort, hash join, aggregation) spill to disk under pressure and queries that cannot spill fail with a resources-exhausted error instead of exhausting process memory. None keeps DataFusion’s default unbounded pool.

Shares SqlEngine::new’s fallback behavior for window helper UDF registration failures.

pub fn new_with_engine_memory(engine_memory: EngineMemory) -> Self

Create a local SQL engine over an explicit EngineMemory source.

Prefer this over new_with_memory_limit when several engines live in one process: passing each of them EngineMemory::Shared with the same pool bounds their combined execution memory, which per-engine limits cannot do.

Shares SqlEngine::new’s fallback behavior for window helper UDF registration failures.

pub fn try_new() -> SqlResult<Self>

Create a local SQL engine, propagating window helper registration errors.

Callers that need to abort startup when window functions cannot be registered should use this constructor.

pub fn with_in_memory_catalog( catalog: Arc<RwLock<InMemoryCatalog>>, ) -> SqlResult<Self>

Create an engine whose krishiv catalog resolves tables registered in InMemoryCatalog (P0-10).

pub fn with_target_parallelism(self, n: NonZeroUsize) -> Self

Set the DataFusion target_partitions parallelism level for this engine.

Higher values allow DataFusion to parallelise hash-join build, aggregation spilling, and parquet scans across more threads. Default: default_parallelism_from_env (available CPU parallelism, KRISHIV_TARGET_PARALLELISM override).

The new level applies to queries planned after this call; plans already produced (or cached logical plans re-planned physically) pick it up at physical planning time.

pub fn target_parallelism(&self) -> NonZeroUsize

Return the configured target_partitions parallelism level.

pub fn memory_limit_bytes(&self) -> Option<usize>

Return the DataFusion memory pool limit for this engine, if bounded.

pub fn session_context(&self) -> &SessionContext

Direct access to the underlying DataFusion session context.

Used by the distributed stage builder (ADR-0003) to create and round-trip physical plans; general query execution should go through SqlEngine::sql so engine-level rewrites and governance apply.

pub fn shuffle_partitions(&self) -> Option<u32>

Return the current shuffle.partitions override, if set via SET shuffle.partitions = N.

pub fn table_row_counts(&self) -> Arc<RwLock<HashMap<String, u64>>>

Return access to the table row-count registry.

Populated by register_parquet and register_record_batches with estimated row counts extracted from table-provider statistics. Used by SqlDataFrame::krishiv_logical_plan to annotate scan nodes.

pub fn registered_table_names(&self) -> Vec<String>

Return table/view names registered in the live DataFusion catalog.

Uses DataFusion’s catalog provider API directly instead of routing through SHOW TABLES, which requires optional information-schema support in some DataFusion configurations.

pub fn with_shuffle_partitions(self, n: Option<u32>) -> Self

Set an override for the shuffle partition count.

When n is Some, exchange and shuffle-write operations use n buckets instead of auto-sizing. Pass None to restore auto-sizing.

pub fn register_streaming_table( &self, name: &str, schema: SchemaRef, ) -> SqlResult<Arc<ContinuousTableInput>>

Register an unbounded continuous table, returning its typed input.

The returned input uses a bounded channel with capacity crate::streaming::CONTINUOUS_TABLE_CHANNEL_CAPACITY. When the consumer (the DataFusion query plan) is slower than the producer, ContinuousTableInput::send(...).await backpressures the producer, and ContinuousTableInput::try_send(...) returns a resource error rather than growing memory without limit. Use [register_streaming_table_with_capacity] for a non-default capacity.

pub fn register_streaming_table_with_capacity( &self, name: &str, schema: SchemaRef, capacity: usize, ) -> SqlResult<Arc<ContinuousTableInput>>

Same as Self::register_streaming_table but with a caller-supplied channel capacity. Useful for tests that want to exercise the full/empty channel boundary without pushing CONTINUOUS_TABLE_CHANNEL_CAPACITY (64) batches.

pub fn register_kafka_source( &self, table_name: impl AsRef<str>, schema: SchemaRef, bootstrap_servers: impl Into<String>, topic: impl Into<String>, group_id: impl Into<String>, ) -> SqlResult<()>

Register a live Kafka/Redpanda topic as an unbounded streaming table.

This is the native Rust path — no Python bridge or external process required. Under the hood it creates an rdkafka consumer and wraps it in a DataFusion StreamingTable so normal SQL queries (SELECT, GROUP BY, windowed aggregations) work against the live topic.

Equivalent SQL DDL:

CREATE EXTERNAL TABLE <name> (<cols>) STORED AS KAFKA
  LOCATION '<topic>'
  OPTIONS ('bootstrap.servers' = '…', 'group.id' = '…');

pub async fn sql_to_kafka( &self, sql: impl AsRef<str>, bootstrap_servers: impl Into<String>, topic: impl Into<String>, ) -> SqlResult<u64>

Execute a SQL query and write every result row to a Kafka/Redpanda topic.

Each row is serialised as a JSON object using the same format as [KafkaSink]. The method blocks until the query stream ends and the producer queue is flushed, then returns the total number of rows written.

Note: If sql targets an unbounded streaming table (e.g. one registered via [register_kafka_source]) this call will never return. Use it with batch sources or add a LIMIT clause.

pub fn with_udf_limits(self, limits: ResourceLimits) -> Self

Configure this engine with explicit UDF resource limits (Track E). When set, calls to sql() and direct UDF syncs will use these budgets instead of unlimited defaults. Intended for job-specific engines.

pub fn is_streaming_source(&self, table_name: &str) -> bool

Returns true if table_name is registered as an unbounded streaming source.

pub fn register_streaming_source_name( &self, table_name: impl Into<String>, ) -> SqlResult<()>

Register a table name as a streaming source without creating a live connector.

This is the test-safe alternative to [register_kafka_source]: it marks table_name in the streaming_sources set so that is_streaming_query returns true for queries that reference it, without constructing any broker connection. Useful for unit tests where a live Kafka broker is not available and rdkafka’s log subsystem is not initialised. Returns SqlError::EmptyTableName if table_name is blank.

pub fn deregister_streaming_source(&self, name: &str) -> SqlResult<()>

Remove a streaming source registration.

Deregisters the table from DataFusion and removes it from the streaming- sources set. Invalidates the plan cache. Idempotent — deregistering a name that was never registered is not an error.

pub fn live_table_registry(&self) -> &Arc<LiveTableRegistry>

Shared live-table registry for CREATE LIVE TABLE DDL.

pub fn incremental_view_registry(&self) -> &Arc<IncrementalViewRegistry>

Shared incremental-view registry for CREATE INCREMENTAL VIEW DDL.

pub fn pipeline_registry(&self) -> &Arc<PipelineRegistry>

Shared pipeline registry for CREATE SOURCE / CREATE SINK DDL.

pub fn operation_registry(&self) -> &Arc<OperationRegistry>

Shared operation registry for cancellation and progress reporting.

pub fn deregister_table(&self, name: &str) -> SqlResult<()>

Drop a named table from the session context.

Idempotent — dropping a name that was never registered is not an error.

Also clears any streaming-source registration for name. This function used to drop the table from DataFusion and leave the name in streaming_sources forever, with has_streaming_sources latched true. Two consequences, and the second is the bad one:

• the set grew without bound for the life of the engine, and
• a name re-registered as an ordinary batch table stayed classified as streaming. DROP TABLE events; CREATE TABLE events AS SELECT … followed by a plain SELECT over events would be routed down the streaming path by [is_streaming_query], which decides execution mode purely from this set.

deregister_streaming_source always did this cleanup; the two paths simply disagreed, and the generic one is the one most callers reach for.

pub fn register_table_udf_fn( &self, name: impl Into<String>, schema: Schema, f: impl Fn(&[ScalarValue]) -> Result<RecordBatch, UdfError> + Send + Sync + 'static, ) -> SqlResult<()>

Register a table UDF backed by a Rust closure.

The closure receives literal arguments passed by the SQL caller as ScalarValue values and returns an Arrow RecordBatch. Non-literal arguments are rejected because they cannot be evaluated safely at the synchronous DataFusion table-function boundary. schema describes the output columns.

Example

engine.register_table_udf_fn(
    "generate_ints",
    Schema::new(vec![Field::new("n", DataType::Int64, false)]),
    |args| {
        let count = match args.first() {
            Some(ScalarValue::Int64(n)) => *n,
            _ => 10,
        };
        let arr = Int64Array::from((0..count).collect::<Vec<_>>());
        Ok(RecordBatch::try_from_iter([("n", Arc::new(arr) as _)])?)
    },
)?;

pub fn is_streaming_query(&self, sql: &str) -> SqlResult<bool>

Returns true if any table referenced in sql is a registered streaming source.

pub fn krishiv_catalog(&self) -> Option<&Arc<RwLock<InMemoryCatalog>>>

Shared Krishiv catalog backing this engine, if configured.

pub async fn register_iceberg_rest_catalog_from_env( &self, ) -> Result<bool, String>

Register a platform Iceberg REST catalog from the KRISHIV_ICEBERG_REST_* environment (URI / WAREHOUSE / TOKEN / NAME) when KRISHIV_ICEBERG_REST_URI is set, so governed catalog.namespace.table references resolve on this engine. Returns whether a catalog was registered.

Unlike the Flight host’s registration (which only fires in InProcess mode), this is callable from the per-task engines the executor builds for coordinated batch SQL — closing the coordinator-mode catalog gap so SELECT … FROM krishiv.<ns>.<table> works on the split (coordinator + platformd) topology. s3:// warehouses are read through KrishivStorage. No-op returning Ok(false) when built without the rest-catalog feature.

pub fn with_udf_registry(self, registry: Arc<RwLock<UdfRegistry>>) -> Self

Share a session UDF registry so scalar UDFs are visible in SQL.

pub fn clear_plan_cache(&self)

Expose cache invalidation for tests and external callers that register tables through a different path.

pub async fn register_python_udfs_from_sql( &self, sql: &str, ) -> SqlResult<String>

Register all scalar UDFs from the attached registry with DataFusion. Uses unlimited defaults (backward compat). Extract Python UDF/UDAF directive comments from sql, register each on this engine, and return sql with the directives removed. Handles both /* krishiv-register-python-udf:name:in,…:out:pickle_b64 */ (scalar, via Self::register_python_udf) and /* krishiv-register-python-udaf:… */ (aggregate, via Self::register_python_udaf). This is how a Python UDF shipped inside a fragment’s SQL reaches the executor’s fresh per-task engine. No-op for SQL that carries none.

pub async fn register_python_udf( &self, name: &str, pickle: &[u8], input_types: &[String], output_type: &str, ) -> SqlResult<()>

Register a distributed Python UDF from cloudpickled bytes. The callable runs in a shared python3 worker subprocess (crate::python_udf) — so it works on the Rust executors with no embedded interpreter. input_types / output_type are Arrow type names. Requires a UDF registry on this engine (executors’ task engines set one).

pub async fn register_python_udaf( &self, name: &str, pickle: &[u8], input_types: &[String], output_type: &str, ) -> SqlResult<()>

Register a distributed Python aggregate UDF (GROUPED_AGG) from cloudpickled bytes. The callable buffers a group’s rows and runs once at finalize inside the shared python3 worker, so it works on the Rust executors and merges correctly across partitions/executors (two-phase aggregation). input_types/output_type are Arrow type names. Requires a UDF registry on this engine (executors’ task engines set one).

pub async fn sync_scalar_udfs(&self) -> SqlResult<()>

pub async fn sync_scalar_udfs_with_limits( &self, limits: ResourceLimits, ) -> SqlResult<()>

Register scalar UDFs with explicit ResourceLimits for sandbox enforcement. Callers that have a job context (scheduler, runner, api session for a job) should use this and pass limits derived from the JobSpec (memory + time cap). This is the concrete Track E seam from job limits to UDF execution.

pub async fn sync_scalar_udfs_with_limits_for_profile( &self, limits: ResourceLimits, profile: DurabilityProfile, ) -> SqlResult<()>

Register scalar UDFs using a caller-resolved durability profile.

pub async fn sync_scalar_udfs_with_limits_for_policy( &self, limits: ResourceLimits, policy: NativeScalarUdfPolicy, ) -> SqlResult<()>

Register scalar UDFs using a caller-snapshotted policy decision.

pub async fn sync_aggregate_udfs(&self) -> SqlResult<()>

Register aggregate UDFs from the attached registry (P1-21).

pub async fn sync_table_udfs(&self) -> SqlResult<()>

Register table UDFs from the attached registry (P1-21).

pub async fn sync_all_udfs(&self) -> SqlResult<()>

Sync all UDF categories, respecting any limits configured on this engine (Track E).

pub async fn register_parquet( &self, table_name: impl AsRef<str>, path: impl AsRef<Path>, ) -> SqlResult<()>

Register a local Parquet path as a table.

pub async fn register_parquet_with_primary_key<S: AsRef<str>>( &self, table_name: impl AsRef<str>, path: impl AsRef<Path>, primary_key: &[S], ) -> SqlResult<()>

As Self::register_parquet, declaring a primary key for the table.

Parquet carries no key, so this declaration is the only way the optimizer can learn that one column determines the others — which is what unlocks functional-dependency group-by pruning and the late-materialisation join-back (crate::late_materialize, measured at 14.8x on TPC-H q10 at SF100).

Informational and unverified, exactly as Spark/Databricks RELY: nothing reads the data to check it, and a key that is not unique or not non-null produces wrong answers, not slow ones.

This exists so the embedded engine can be told the same thing the distributed one is told through BatchSqlTable::primary_key. A declaration that reached only one of the two would mean the two topologies silently plan different queries — the hardest kind of difference to notice, because both still return correct results.

pub async fn read_parquet( &self, path: impl AsRef<Path>, ) -> SqlResult<SqlDataFrame>

Create a DataFrame by reading a local Parquet path directly.

pub async fn register_record_batches( &self, table_name: impl AsRef<str>, batches: Vec<RecordBatch>, ) -> SqlResult<()>

Register an in-memory table from Arrow record batches.

The schema is inferred from the first batch. An empty batches slice registers a table with no rows using the provided schema if the batches are non-empty, or is a no-op if empty.

pub async fn read_parquet_with_options( &self, path: impl AsRef<Path>, opts: &ParquetReaderOptions, ) -> SqlResult<SqlDataFrame>

Create a DataFrame by reading a local Parquet path with typed options.

pub async fn read_csv(&self, path: impl AsRef<Path>) -> SqlResult<SqlDataFrame>

Create a DataFrame by reading a local CSV path directly.

pub async fn read_csv_with_options( &self, path: impl AsRef<Path>, opts: &CsvReaderOptions, ) -> SqlResult<SqlDataFrame>

Create a DataFrame by reading a local CSV path with typed options.

pub async fn read_json(&self, path: impl AsRef<Path>) -> SqlResult<SqlDataFrame>

Create a DataFrame by reading a local JSON/NDJSON path directly.

pub async fn read_delta( &self, path: impl AsRef<str>, version: Option<i64>, ) -> SqlResult<SqlDataFrame>

Read a local Delta table directory into a DataFrame.

pub async fn read_hudi( &self, path: impl AsRef<str>, query_type: HudiQueryType, begin_instant: Option<&str>, ) -> SqlResult<SqlDataFrame>

Read a Hudi table directory.

pub fn sql<'a>( &'a self, query: impl AsRef<str> + Send + 'a, ) -> Pin<Box<dyn Future<Output = SqlResult<SqlDataFrame>> + Send + 'a>>

Plan a SQL query with DataFusion.

Returns a boxed future ON PURPOSE (do not revert to async fn): this body is one of the largest async state machines in the workspace, and as an opaque impl Future its Send proof leaked into every CONSUMER crate — rustc’s old trait solver re-derived it per call site, which drove krishiv-executor’s front-end to 20+ hour compiles (95% of CPU in ObligationForest::process_obligations, bisected 2026-07-19). Boxing at the definition pays that proof once, here.

pub async fn execute_with_timeout( &self, query: impl AsRef<str> + Send, timeout_ms: u64, ) -> SqlResult<SqlDataFrame>

Execute a SQL query with a timeout.

Returns SqlError::Timeout if timeout_ms elapses before the query produces a result. The underlying DataFusion task is abandoned (not cancelled at the engine level) when the timeout fires; its resources are released when the spawned task eventually completes.

pub async fn execute_with_operation_id( &self, operation_id: u64, query: impl AsRef<str> + Send, cancelled_ids: &OperationRegistry, ) -> SqlResult<TaggedQueryResult>

Execute a SQL query tagged with a caller-supplied operation ID.

The operation ID is recorded in the returned TaggedQueryResult and can be used to correlate logs, metrics, and cancellation requests. If cancelled_ids contains operation_id before execution begins the function returns SqlError::OperationCancelled immediately.

Trait Implementations

impl Clone for SqlEngine

fn clone(&self) -> SqlEngine

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for SqlEngine

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for SqlEngine

fn default() -> Self

Returns the “default value” for a type.