feat: harden inference path and add drop-in sql entry point

Cargo.toml

@@ -14,4 +14,4 @@ serde_json = "1"
 tokio = { version = "1", features = ["sync", "time"] }
 
 [dev-dependencies]
-tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1", features = ["macros", "rt-multi-thread", "test-util"] }

README.md

@@ -1,6 +1,11 @@
 # datafusion-jev
 
-Typed Jev decisions in DataFusion 55 SQL. A standalone extension crate; the host supplies the authenticated HTTP client.
+Ask questions about your text from SQL. `prompt_jev` sends each row to the Jev
+inference service and returns a typed answer: a probability, a choice from a list
+you define, or a score on a scale you define.
+
+Works with DataFusion 55. You bring the HTTP client and API key; the crate handles
+the SQL syntax, batching, retries, and result types.
 
 ```sql
 SELECT conversation_id,
@@ -14,38 +19,163 @@ SELECT conversation_id,
 FROM customer_conversations;
 ```
 
-## SQL interface
+## Setup
+
+Add the crate to your project, then follow three steps.
+
+**1. Provide the HTTP client.** Implement `JevClient` with whatever HTTP library
+you already use. POST the request body as JSON to
+`https://api.typesafe.ai/v1/systemone` with your API key in a header, and return
+the parsed JSON response. Keep the key on the server; it never appears in SQL.
+
+```rust
+use datafusion_jev::{JevClient, RequestError};
+
+#[derive(Debug)]              // must not print the API key
+struct MyClient { /* http client, api key */ }
+
+#[async_trait::async_trait]
+impl JevClient for MyClient {
+    async fn request(&self, body: serde_json::Value) -> Result<serde_json::Value, RequestError> {
+        // Send `body`. Map the outcome:
+        //   2xx           -> Ok(parsed JSON)
+        //   5xx / network -> Err(RequestError::Unavailable)   (retried)
+        //   4xx / other   -> Err(RequestError::Fatal(reason))  (fails the query)
+        todo!()
+    }
+}
+```
+
+**2. Register it once per session.**
+
+```rust
+let ctx = SessionContext::new();
+datafusion_jev::register(&ctx, Arc::new(MyClient::new()));
+```
+
+**3. Run queries through `datafusion_jev::sql`.**
+
+```rust
+let df = datafusion_jev::sql(&ctx, "SELECT prompt_jev(body, 'Is this urgent?') FROM tickets").await?;
+```
+
+Use this in place of `ctx.sql`. It accepts the same SQL and returns the same
+`DataFrame`, and `sql_with_options` takes DataFusion's `SQLOptions` just like
+`ctx.sql_with_options`. Plain `ctx.sql` does not understand the `prompt_jev`
+options syntax and will fail on it.
+
+If your application already parses SQL itself, call `datafusion_jev::rewrite` on
+the parsed statement before planning it instead.
+
+## Writing queries
+
+`prompt_jev(text, 'question', options...)`. The text can be any column or
+expression. The question and all options must be literals.
+
+### Yes/no probability (default)
+
+```sql
+SELECT prompt_jev(body, 'Is the customer asking for a refund?') AS p FROM tickets
+```
+
+Returns a `DOUBLE` from 0 to 1. Use it directly in filters:
+
+```sql
+WHERE prompt_jev(body, 'Is this urgent?') > 0.8
+```
+
+Optionally describe what `true` and `false` mean:
+
+```sql
+noul := [{label: 'true', description: 'Explicitly requests money back'},
+         {label: 'false', description: 'Anything else'}]
+```
+
+### Pick one option
+
+```sql
+SELECT prompt_jev(body, 'What is this about?', choice := ['billing', 'technical', 'sales']) AS c
+FROM tickets
+```
+
+Returns a struct. Read its fields with dot syntax:
+
+| Field | Type | Meaning |
+|---|---|---|
+| `c.choice` | `VARCHAR` | The selected label |
+| `c.confidence` | `DOUBLE` | How sure the model is, 0 to 1 |
+| `c.probabilities` | list of `{value, probability}` | One entry per option, in the order you listed them |
+
+Two to 255 options. Labels must be unique. Each option can be a plain string or
+`{label: '...', description: '...'}`.
+
+### Score on a scale
+
+```sql
+SELECT prompt_jev(body, 'How severe is the problem?', score := ['low', 'medium', 'high']) AS s
+FROM tickets
+```
+
+Returns a struct like `choice`, except the answer is a number:
+
+| Field | Type | Meaning |
+|---|---|---|
+| `s.score` | `DOUBLE` | Expected position on the scale. With three levels, 0 = low, 2 = high, 1.5 = between medium and high |
+| `s.confidence` | `DOUBLE` | How sure the model is, 0 to 1 |
+| `s.probabilities` | list of `{index, value, probability}` | One entry per level, in order |
+
+Two to 10 levels, listed from lowest to highest. `choice` and `score` cannot be
+combined in one call.
+
+### Batching
+
+```sql
+prompt_jev(body, 'question', batch_size := 8)
+```
 
-- `prompt_jev(text, 'question')`: a nullable `DOUBLE` between 0 and 1 (`noul`).
-- `choice := ['a', 'b']`: struct with `choice`, `probabilities`, `confidence`.
-- `score := ['low', 'medium', 'high']`: struct with weighted `score` (0–2 here), `probabilities`, `confidence`.
-- Criteria also accept `{label: '...', description: '...'}` literals. Descriptions may be NULL.
-- `noul := [...]` may attach descriptions to exactly `true` and `false`.
-- `batch_size := 1..64`, default 32. Only input varies per row; instructions and options currently require SQL literals, not computed constant expressions.
-- Choice requires 2–255 unique labels; score requires 2–10. These modes are mutually exclusive.
+Up to 64 rows are sent per request, default 32. Rows in the same request share
+context, which can occasionally influence answers. Use `batch_size := 1` when
+each row must be judged in complete isolation.
 
-Choice probabilities are `STRUCT(value VARCHAR, probability DOUBLE)[]`, in caller order. Score probabilities additionally have `index UINTEGER`. Confidence is the provider's confidence, not a locally inferred value.
+## What to expect
 
-This first version supports single-question calls. `questions := ...` and the JSON configuration escape hatch are not implemented. It follows MotherDuck's public interface for the supported subset; it is not MotherDuck code.
+- **NULL text** returns NULL without calling the service.
+- **Repeated text** is asked once and the answer is copied to every matching row.
+  Dedup works within each batch of up to 256 rows, so a constant like
+  `prompt_jev('hello', ...)` costs about one request per 256 rows scanned.
+- **Repeated calls** with identical arguments in one query are evaluated once.
+  Reading several fields from one result does not repeat the request.
+- **Limits.** Each row's text may be up to 64 KiB after JSON encoding. Longer text
+  fails the query. Requests are capped at 256 KiB.
+- **Outages.** Each request is tried three times with a 30-second timeout. If all
+  fail, the affected rows return NULL and the query completes. Bad credentials,
+  invalid options, or a malformed reply fail the query instead.
+- **Concurrency.** At most eight requests are in flight at a time per registered
+  client, across all queries.
+- **Cancelling** a query cancels its in-flight requests.
+- **No caching.** Every query calls the service again. Store results in a table
+  when you want to reuse them.
+- **Dialects.** The `:=` and `{label: ...}` syntax works with DataFusion's default
+  dialect and with DuckDB. If your session is set to PostgreSQL or MySQL,
+  `prompt_jev` queries fail with an error that says so.
 
-## Host integration
+Text columns may be `Utf8`, `LargeUtf8`, `Utf8View`, or dictionary-encoded.
 
-1. Implement `JevClient::request(Value)` using your existing HTTP transport. POST the JSON to `https://api.typesafe.ai/v1/systemone`, authenticate with a server-side API key, and return the parsed response. Apply your outbound policy; never put credentials in SQL. Distinguish retryable `Unavailable` from `Fatal` configuration/protocol errors.
-2. Call `datafusion_jev::register(&ctx, Arc::new(client))` once per client scope.
-3. Parse the user's SQL, call `datafusion_jev::rewrite(&mut datafusion_statement)` (which also handles EXPLAIN), and pass the rewritten AST to `SessionState::statement_to_plan`.
+## Compatibility
 
-The extension uses DataFusion's `AsyncScalarUDF`; it does not replace the host's query planner. Registration appends a physical optimizer rule that deduplicates repeated Jev expressions within an async execution node, preserving all output slots. SQL lowering preserves the input expression and encodes validated options as a constant. The private `__datafusion_jev` function is an implementation detail.
+The SQL interface follows MotherDuck's public `prompt_jev` for the supported
+subset. This is not MotherDuck code. Multi-question calls (`questions := ...`)
+and the JSON configuration form are not implemented yet.
 
-The host is responsible for credentials, authorization to use paid inference, endpoint configuration, HTTP status classification, usage accounting, and response-size limits. `Debug` implementations of clients must redact credentials. No live requests or credentials are required by this crate's tests.
+## How it works
 
-## Execution and failures
+`prompt_jev` is rewritten at parse time into a private async scalar function
+whose options are a validated constant. Inference runs through DataFusion's
+`AsyncScalarUDF`, so no blocking threads are involved. A physical optimizer rule
+collapses repeated identical calls in one projection into a single evaluation.
 
-- Up to eight requests execute concurrently per registration, shared across partitions and queries. Futures are awaited inline and are dropped on query cancellation.
-- Input is evaluated in bounded batches of at most 256 rows. NULL input skips inference and returns NULL.
-- Batches share a state object and use a separately scoped question for each row. This can affect predictions: choose `batch_size := 1` for isolated state per row.
-- Input is capped at 64 KiB per row; encoded requests at 256 KiB. Requests use `jev-latest`.
-- Transient failures get three attempts with a 30-second timeout per attempt and bounded backoff. Exhaustion returns NULL for affected rows. Invalid configuration, authentication, or malformed answers fail the query.
-- No persistent result cache: materialize results using the host's result-to-table facility when you want to reuse them without inference.
+Your application stays in charge of credentials, who may run paid inference,
+endpoint configuration, and usage accounting.
 
 ## Development
 
@@ -55,4 +185,5 @@ cargo clippy --all-targets -- -D warnings
 cargo test
 ```
 
-Tests execute real DataFusion plans against a mock provider, covering the public syntax, result fields, batching, NULLs, filtering, and error behavior.
+Tests run real DataFusion plans against a mock service. No network access or
+credentials are needed.

src/lib.rs

@@ -1,30 +1,44 @@
-//! Jev decisions for DataFusion 55. Parse with `rewrite_statement`, then plan
-//! using a context configured by `register`. Network IO uses DataFusion's native
-//! asynchronous UDF execution; no blocking threads or detached inference tasks.
+//! `prompt_jev` for DataFusion 55: ask questions about text from SQL and get
+//! typed answers.
+//!
+//! 1. Implement [`JevClient`] with your HTTP client and API key.
+//! 2. Call [`register`] once on your `SessionContext`.
+//! 3. Run queries with [`sql()`] instead of `SessionContext::sql`.
+//!
+//! See the README for the SQL syntax and result types.
 mod optimizer;
 pub mod sql;
 mod udf;
 use async_trait::async_trait;
-use datafusion::prelude::SessionContext;
+use datafusion::{
+    common::plan_datafusion_err,
+    execution::context::SQLOptions,
+    prelude::{DataFrame, SessionContext},
+};
 use serde_json::Value;
 pub use sql::rewrite_statement;
 use std::sync::Arc;
 
-/// Transport errors never carry input text, credentials, or response bodies.
+/// How a request failed. Never include input text, credentials, or response
+/// bodies in these errors; they can surface in query error messages.
 #[derive(Debug)]
 pub enum RequestError {
-    /// Retried; exhausted requests produce NULLs for their input rows.
+    /// Network error, timeout, or 5xx. Retried; if retries run out the
+    /// affected rows become NULL and the query still completes.
     Unavailable,
-    /// Authentication, invalid configuration, or malformed provider responses.
+    /// Authentication, invalid request, or malformed response. Fails the query.
     Fatal(String),
 }
-/// Supplied by the host so credentials and outbound policy stay host-owned.
+/// Your HTTP client for the Jev service. POST `body` as JSON to the endpoint
+/// with your API key and return the parsed JSON response. The `Debug`
+/// implementation must not print credentials.
 #[async_trait]
 pub trait JevClient: Send + Sync + std::fmt::Debug {
     async fn request(&self, body: Value) -> Result<Value, RequestError>;
 }
-/// Register the async implementation. Install once per host/client scope.
-/// Public SQL must first pass through `rewrite_statement`.
+/// Make `prompt_jev` available in this session. Call once per `SessionContext`.
+/// Then run queries with [`sql()`], or call [`rewrite`] before planning if you
+/// parse SQL yourself.
 pub fn register(ctx: &SessionContext, client: Arc<dyn JevClient>) {
     ctx.register_udf(udf::function(client));
     let state_ref = ctx.state_ref();
@@ -34,7 +48,9 @@ pub fn register(ctx: &SessionContext, client: Arc<dyn JevClient>) {
         .build();
 }
 
-/// Rewrite DataFusion statements, including its EXPLAIN wrapper.
+/// For applications that parse SQL themselves: rewrite `prompt_jev` calls in a
+/// parsed statement (including inside EXPLAIN) so DataFusion can plan it.
+/// Safe to call on statements without `prompt_jev`, and safe to call twice.
 pub fn rewrite(
     statement: &mut datafusion::sql::parser::Statement,
 ) -> datafusion::common::Result<()> {
@@ -44,3 +60,40 @@ pub fn rewrite(
         _ => Ok(()),
     }
 }
+
+/// Run a SQL query that may use `prompt_jev`. Use in place of
+/// `SessionContext::sql`; same input, same `DataFrame` result.
+pub async fn sql(ctx: &SessionContext, text: &str) -> datafusion::common::Result<DataFrame> {
+    sql_with_options(ctx, text, SQLOptions::new()).await
+}
+
+/// [`sql()`] with DataFusion's `SQLOptions`, for example to forbid DDL.
+/// Use in place of `SessionContext::sql_with_options`.
+pub async fn sql_with_options(
+    ctx: &SessionContext,
+    text: &str,
+    options: SQLOptions,
+) -> datafusion::common::Result<DataFrame> {
+    let state = ctx.state();
+    let dialect = state.config().options().sql_parser.dialect;
+    // Strict dialects either reject `:=` outright or parse it as an ordinary
+    // expression, so the failure can surface from parsing or from the rewrite.
+    let hint = |e: datafusion::common::DataFusionError| {
+        let lenient = matches!(
+            dialect,
+            datafusion::config::Dialect::Generic | datafusion::config::Dialect::DuckDB
+        );
+        if lenient || !text.to_lowercase().contains("prompt_jev") {
+            return e;
+        }
+        plan_datafusion_err!(
+            "{e}. prompt_jev syntax (':=' options and '{{label: ...}}' literals) requires \
+             the generic (default) or duckdb SQL dialect; the session uses {dialect}"
+        )
+    };
+    let mut statement = state.sql_to_statement(text, &dialect).map_err(hint)?;
+    rewrite(&mut statement).map_err(hint)?;
+    let plan = state.statement_to_plan(statement).await?;
+    options.verify_plan(&plan)?;
+    ctx.execute_logical_plan(plan).await
+}

src/optimizer.rs

@@ -44,7 +44,7 @@ impl PhysicalOptimizerRule for DeduplicateJev {
                     .downcast_ref::<ScalarFunctionExpr>()
                     .is_some_and(|f| f.fun().name() == "__datafusion_jev");
                 let existing = if is_jev {
-                    unique.iter().position(|x| x.func == expr.func.clone())
+                    unique.iter().position(|x| x.func.eq(&expr.func))
                 } else {
                     None
                 };

src/sql.rs

@@ -78,21 +78,43 @@ fn criteria(e: &Expr) -> Result<Vec<Criterion>> {
             "prompt_jev criteria must be a constant array"
         ));
     };
-    a.elem.iter().map(|e| {
-        if let Expr::Dictionary(fields) = e {
-            let mut label = None; let mut description = None; let mut seen = HashSet::new();
+    a.elem
+        .iter()
+        .map(|e| {
+            let Expr::Dictionary(fields) = e else {
+                return Ok(Criterion {
+                    label: string(e)?,
+                    description: None,
+                });
+            };
+            let mut label = None;
+            let mut description = None;
+            let mut seen = HashSet::new();
             for field in fields {
                 let key = field.key.value.as_str();
-                if !seen.insert(key) { return Err(plan_datafusion_err!("duplicate criterion field")); }
+                if !seen.insert(key) {
+                    return Err(plan_datafusion_err!("duplicate criterion field"));
+                }
                 match key {
                     "label" => label = Some(string(&field.value)?),
-                    "description" => { if !matches!(field.value.as_ref(), Expr::Value(v) if v.value == Value::Null) { description=Some(string(&field.value)?); } },
+                    "description" => {
+                        let null = matches!(
+                            field.value.as_ref(),
+                            Expr::Value(v) if v.value == Value::Null
+                        );
+                        if !null {
+                            description = Some(string(&field.value)?);
+                        }
+                    }
                     _ => return Err(plan_datafusion_err!("unknown criterion field: {key}")),
                 }
             }
-            Ok(Criterion { label: label.ok_or_else(|| plan_datafusion_err!("criterion requires label"))?, description })
-        } else { Ok(Criterion { label: string(e)?, description: None }) }
-    }).collect()
+            Ok(Criterion {
+                label: label.ok_or_else(|| plan_datafusion_err!("criterion requires label"))?,
+                description,
+            })
+        })
+        .collect()
 }
 /// Call after parsing and before DataFusion plans the statement. Idempotent.
 pub fn rewrite_statement(statement: &mut ast::Statement) -> Result<()> {