feat: add experimental Jev SQL filtering (#9265)

* feat: add experimental Jev SQL filtering Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * feat: gate Jev filtering behind an opt-in Cargo feature Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * ci: verify Jev registration with default features Run the existing registry regression without the jev feature in both PR tests and merge-queue coverage, alongside the existing feature-enabled test runs. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * docs: clarify Jev concurrency scope and stabilization work Document the per-expression/batch concurrency bound and track process-wide limiting, rate-limit backoff, and request budgets and metrics as stabilization prerequisites. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * feat: enable renamed ai-functions feature by default Rename the Jev Cargo feature across the command, query, and function crates and enable it in their defaults. Update CI and documentation, retaining an isolated no-default-features registry check and the runtime API opt-in. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * ci: remove extra AI feature-off checks Use the regular AI-enabled unit and coverage runs for the default feature configuration. Keep feature-off validation available locally and update the usage guide to match. Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> * refactor: rename AI feature to ai_functions Signed-off-by: Lei, HUANG <ratuthomm@gmail.com> --------- Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
2026-10-16 17:02:33 +00:00 · 2026-09-22 03:03:23 +00:00
parent 16978cf6c2
commit b5199bc59a
14 changed files with 709 additions and 5 deletions
@@ -10,8 +10,9 @@ workspace = true
 [features]
 testing = []
 enterprise = []
-default = ["geo"]
+default = ["geo", "ai_functions"]
 geo = ["geohash", "h3o", "s2", "wkt", "geo-types", "dep:geo"]
+ai_functions = ["dep:futures", "dep:reqwest"]

 [dependencies]
 ahash.workspace = true
@@ -43,6 +44,7 @@ datafusion-pg-catalog.workspace = true
 datafusion-physical-expr.workspace = true
 datatypes.workspace = true
 derive_more.workspace = true
+futures = { workspace = true, optional = true }
 geo = { version = "0.29", optional = true }
 geo-types = { version = "0.7", optional = true }
 geohash = { version = "0.13", optional = true }
@@ -57,6 +59,7 @@ num = "0.4"
 num-traits = "0.2"
 paste.workspace = true
 regex.workspace = true
+reqwest = { workspace = true, optional = true }
 s2 = { version = "0.0.12", optional = true }
 serde.workspace = true
 serde_json.workspace = true
@@ -71,6 +74,7 @@ wkt = { version = "0.11", optional = true }

 [dev-dependencies]
 approx = "0.5"
+axum.workspace = true
 criterion.workspace = true
 futures.workspace = true
 pretty_assertions.workspace = true
@@ -80,3 +84,7 @@ tokio.workspace = true
 [[bench]]
 name = "uddsketch"
 harness = false
+
+[package.metadata.cargo-udeps.ignore]
+# Used only by the feature-gated Jev HTTP tests.
+development = ["axum"]
@@ -34,6 +34,8 @@ use crate::scalars::date::DateFunction;
 use crate::scalars::expression::ExpressionFunction;
 use crate::scalars::hll_count::HllCalcFunction;
 use crate::scalars::ip::IpFunctions;
+#[cfg(feature = "ai_functions")]
+use crate::scalars::jev::JevFunction;
 use crate::scalars::json::JsonFunction;
 use crate::scalars::matches::MatchesFunction;
 use crate::scalars::matches_term::MatchesTermFunction;
@@ -226,6 +228,8 @@ pub static FUNCTION_REGISTRY: LazyLock<Arc<FunctionRegistry>> = LazyLock::new(||
    // Full text search function
    MatchesFunction::register(&function_registry);
    MatchesTermFunction::register(&function_registry);
+    #[cfg(feature = "ai_functions")]
+    JevFunction::register(&function_registry);

    // System and administration functions
    SystemFunction::register(&function_registry);
@@ -379,6 +383,14 @@ mod tests {
        assert!(FUNCTION_REGISTRY.get_function("uddsketch_rank").is_some());
    }

+    #[test]
+    fn test_jev_registration_matches_feature() {
+        assert_eq!(
+            FUNCTION_REGISTRY.get_function("jev").is_some(),
+            cfg!(feature = "ai_functions")
+        );
+    }
+
    #[test]
    fn test_register_if_absent_registers_new_function() {
        let registry = FunctionRegistry::default();
@@ -17,6 +17,8 @@ pub(crate) mod date;
 pub mod expression;
 #[cfg(feature = "geo")]
 pub mod geo;
+#[cfg(feature = "ai_functions")]
+pub(crate) mod jev;
 pub mod json;
 pub mod matches;
 pub mod matches_term;
@@ -0,0 +1,209 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#[cfg(test)]
+mod tests;
+
+use std::fmt;
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::BooleanArray;
+use arrow::datatypes::DataType;
+use async_trait::async_trait;
+use datafusion_common::cast::{as_float64_array, as_string_view_array};
+use datafusion_common::utils::take_function_args;
+use datafusion_common::{Result, exec_datafusion_err, exec_err, not_impl_err};
+use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
+use futures::{StreamExt, TryStreamExt, stream};
+use reqwest::Client;
+use serde_json::{Value, json};
+
+use crate::function_factory::ScalarFunctionFactory;
+use crate::function_registry::FunctionRegistry;
+
+/// Experimental natural-language predicate: `jev(text, statement, threshold)`.
+/// Each non-null row is a Noul question; its probability is compared with `>=`.
+#[derive(PartialEq, Eq, Hash)]
+pub(crate) struct JevFunction {
+    signature: Signature,
+    enabled: bool,
+    api_key: Option<String>,
+    endpoint: String,
+    model: String,
+}
+
+impl JevFunction {
+    pub(crate) fn register(registry: &FunctionRegistry) {
+        registry.register(ScalarFunctionFactory {
+            name: "jev".to_string(),
+            factory: Arc::new(|_| AsyncScalarUDF::new(Arc::new(Self::default())).into_scalar_udf()),
+        });
+    }
+
+    async fn evaluate(&self, client: &Client, text: &str, statement: &str) -> Result<f64> {
+        let response: Value = client
+            .post(&self.endpoint)
+            .json(&json!({
+                "model": self.model,
+                "state": text,
+                "questions": {
+                    "matches": { "type": "noul", "instructions": statement }
+                }
+            }))
+            .send()
+            .await
+            .and_then(|response| response.error_for_status())
+            .map_err(|e| exec_datafusion_err!("jev request failed: {e}"))?
+            .json()
+            .await
+            .map_err(|e| exec_datafusion_err!("jev response is not valid JSON: {e}"))?;
+
+        let answer = &response["answers"]["matches"];
+        let probability = answer["noul"].as_f64().filter(|p| (0.0..=1.0).contains(p));
+        match (answer["type"].as_str(), probability) {
+            (Some("noul"), Some(probability)) => Ok(probability),
+            _ => exec_err!(
+                "jev response must contain answers.matches with type noul and a probability in [0, 1]"
+            ),
+        }
+    }
+
+    fn client(&self) -> Result<Client> {
+        if !self.enabled {
+            return exec_err!(
+                "jev is experimental; set GREPTIMEDB_EXPERIMENTAL_JEV=true to enable it"
+            );
+        }
+        let key = self
+            .api_key
+            .as_deref()
+            .filter(|key| !key.trim().is_empty())
+            .ok_or_else(|| {
+                exec_datafusion_err!("jev requires the JEV_API_KEY environment variable")
+            })?;
+        let mut authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {key}"))
+            .map_err(|_| exec_datafusion_err!("JEV_API_KEY is not a valid HTTP header value"))?;
+        authorization.set_sensitive(true);
+        let mut headers = reqwest::header::HeaderMap::new();
+        headers.insert(reqwest::header::AUTHORIZATION, authorization);
+        Client::builder()
+            .default_headers(headers)
+            .timeout(Duration::from_secs(30))
+            .build()
+            .map_err(|e| exec_datafusion_err!("failed to create jev HTTP client: {e}"))
+    }
+}
+
+impl Default for JevFunction {
+    fn default() -> Self {
+        Self {
+            // External model evaluations must not be constant-folded during planning.
+            signature: Signature::exact(
+                vec![DataType::Utf8View, DataType::Utf8View, DataType::Float64],
+                Volatility::Volatile,
+            ),
+            enabled: std::env::var("GREPTIMEDB_EXPERIMENTAL_JEV").as_deref() == Ok("true"),
+            api_key: std::env::var("JEV_API_KEY").ok(),
+            endpoint: std::env::var("JEV_ENDPOINT")
+                .unwrap_or_else(|_| "https://api.typesafe.ai/v1/systemone".to_string()),
+            model: std::env::var("JEV_MODEL").unwrap_or_else(|_| "jev-latest".to_string()),
+        }
+    }
+}
+
+impl fmt::Debug for JevFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // Credentials must never appear in plans or diagnostic output.
+        f.debug_struct("JevFunction")
+            .field("signature", &self.signature)
+            .field("enabled", &self.enabled)
+            .field("model", &self.model)
+            .finish_non_exhaustive()
+    }
+}
+
+impl ScalarUDFImpl for JevFunction {
+    fn name(&self) -> &str {
+        "jev"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Boolean)
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        not_impl_err!("jev can only be called from async contexts")
+    }
+}
+
+#[async_trait]
+impl AsyncScalarUDFImpl for JevFunction {
+    async fn invoke_async_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let arrays = args
+            .args
+            .into_iter()
+            .map(|arg| arg.into_array(args.number_rows))
+            .collect::<Result<Vec<_>>>()?;
+        let [texts, statements, thresholds] = take_function_args(self.name(), arrays)?;
+        let texts = as_string_view_array(&texts)?;
+        let statements = as_string_view_array(&statements)?;
+        let thresholds = as_float64_array(&thresholds)?;
+        let rows: Vec<_> = texts
+            .iter()
+            .zip(statements.iter())
+            .zip(thresholds.iter())
+            .map(|((text, statement), threshold)| Some((text?, statement?, threshold?)))
+            .collect();
+
+        // Validate the whole batch before making any billable requests.
+        for (_, _, threshold) in rows.iter().flatten() {
+            if !(0.0..=1.0).contains(threshold) {
+                return exec_err!("jev threshold must be finite and in [0, 1]");
+            }
+        }
+        if rows.iter().all(Option::is_none) {
+            return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null(
+                args.number_rows,
+            ))));
+        }
+
+        // Reuse connections within the batch; bound in-flight requests and preserve row order.
+        let client = self.client()?;
+        let requests: Vec<_> = rows
+            .into_iter()
+            .map(|row| {
+                let client = &client;
+                async move {
+                    match row {
+                        Some((text, statement, threshold)) => self
+                            .evaluate(client, text, statement)
+                            .await
+                            .map(|p| Some(p >= threshold)),
+                        None => Ok(None),
+                    }
+                }
+            })
+            .collect();
+        // This limit is per expression/batch invocation, not per query or process.
+        // Concurrent partitions and queries can each have their own in-flight requests.
+        let matches: Vec<Option<bool>> = stream::iter(requests).buffered(8).try_collect().await?;
+        Ok(ColumnarValue::Array(Arc::new(BooleanArray::from(matches))))
+    }
+}
@@ -0,0 +1,250 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use axum::http::{HeaderMap, StatusCode};
+use axum::routing::post;
+use axum::{Json, Router};
+use datafusion::assert_batches_eq;
+use datafusion::prelude::{SessionConfig, SessionContext};
+use tokio::net::TcpListener;
+use tokio::task::JoinHandle;
+
+use super::*;
+use crate::function::FunctionContext;
+use crate::function_registry::FUNCTION_REGISTRY;
+
+struct MockServer {
+    endpoint: String,
+    task: JoinHandle<()>,
+}
+
+impl MockServer {
+    async fn start(router: Router) -> Self {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let endpoint = format!("http://{}/v1/systemone", listener.local_addr().unwrap());
+        let task = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
+        Self { endpoint, task }
+    }
+
+    fn function(&self) -> JevFunction {
+        JevFunction {
+            enabled: true,
+            api_key: Some("test-key".to_string()),
+            endpoint: self.endpoint.clone(),
+            model: "test-model".to_string(),
+            ..Default::default()
+        }
+    }
+}
+
+impl Drop for MockServer {
+    fn drop(&mut self) {
+        self.task.abort();
+    }
+}
+
+fn context(function: JevFunction) -> SessionContext {
+    let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));
+    ctx.register_udf(AsyncScalarUDF::new(Arc::new(function)).into_scalar_udf());
+    ctx
+}
+
+#[tokio::test]
+async fn test_jev_sql_filters_using_noul_probability() {
+    let server = MockServer::start(Router::new().route(
+        "/v1/systemone",
+        post(
+            |headers: HeaderMap, Json(request): Json<Value>| async move {
+                assert_eq!(headers["authorization"], "Bearer test-key");
+                assert_eq!(request["model"], "test-model");
+                assert_eq!(request["questions"]["matches"]["type"], "noul");
+                // The SQL statement is passed through verbatim, not rewritten as a prompt.
+                assert_eq!(
+                    request["questions"]["matches"]["instructions"],
+                    "payment failed"
+                );
+                let probability = match request["state"].as_str().unwrap() {
+                    "failed" => {
+                        // Complete this row after later rows to exercise result alignment.
+                        tokio::time::sleep(Duration::from_millis(20)).await;
+                        0.95
+                    }
+                    "boundary" => 0.8,
+                    "recovered" => 0.1,
+                    unexpected => panic!("unexpected state: {unexpected}"),
+                };
+                Json(json!({"answers": {"matches": {"type": "noul", "noul": probability}}}))
+            },
+        ),
+    ))
+    .await;
+    let ctx = context(server.function());
+    let events = "(VALUES (1, 'failed'), (2, 'boundary'), (3, 'recovered'), (4, NULL)) AS events(id, message)";
+    let batches = ctx
+        .sql(&format!(
+            "SELECT id FROM {events} WHERE jev((message), 'payment failed', 0.8) ORDER BY id"
+        ))
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_batches_eq!(
+        ["+----+", "| id |", "+----+", "| 1  |", "| 2  |", "+----+"],
+        &batches
+    );
+
+    let batches = ctx
+        .sql(&format!(
+            "SELECT id, jev(message, 'payment failed', 0.8) AS matched FROM {events} ORDER BY id"
+        ))
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_batches_eq!(
+        [
+            "+----+---------+",
+            "| id | matched |",
+            "+----+---------+",
+            "| 1  | true    |",
+            "| 2  | true    |",
+            "| 3  | false   |",
+            "| 4  |         |",
+            "+----+---------+"
+        ],
+        &batches
+    );
+}
+
+#[tokio::test]
+async fn test_jev_nulls_and_invalid_thresholds_need_no_api_key() {
+    let ctx = context(JevFunction {
+        enabled: true,
+        api_key: None,
+        ..Default::default()
+    });
+    let batches = ctx
+        .sql("SELECT jev(NULL, 'condition', 0.8) AS a, jev('text', NULL, 0.8) AS b, jev('text', 'condition', NULL) AS c")
+        .await.unwrap().collect().await.unwrap();
+    assert_batches_eq!(
+        [
+            "+---+---+---+",
+            "| a | b | c |",
+            "+---+---+---+",
+            "|   |   |   |",
+            "+---+---+---+"
+        ],
+        &batches
+    );
+
+    for threshold in ["-0.1", "1.1", "'NaN'::DOUBLE", "'Infinity'::DOUBLE"] {
+        let error = ctx
+            .sql(&format!("SELECT jev('text', 'condition', {threshold})"))
+            .await
+            .unwrap()
+            .collect()
+            .await
+            .unwrap_err();
+        assert!(
+            error
+                .to_string()
+                .contains("jev threshold must be finite and in [0, 1]"),
+            "{error}"
+        );
+    }
+    let error = ctx
+        .sql("SELECT jev('text', 'condition', 0.8)")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap_err();
+    assert!(error.to_string().contains("JEV_API_KEY"), "{error}");
+
+    let ctx = context(JevFunction {
+        enabled: false,
+        api_key: None,
+        ..Default::default()
+    });
+    let error = ctx
+        .sql("SELECT jev('text', 'condition', 0.8)")
+        .await
+        .unwrap()
+        .collect()
+        .await
+        .unwrap_err();
+    assert!(
+        error
+            .to_string()
+            .contains("GREPTIMEDB_EXPERIMENTAL_JEV=true"),
+        "{error}"
+    );
+}
+
+#[tokio::test]
+async fn test_jev_bad_api_responses_fail_the_query() {
+    for (status, body, expected) in [
+        (StatusCode::TOO_MANY_REQUESTS, "{}", "jev request failed"),
+        (StatusCode::OK, "not json", "jev response is not valid JSON"),
+        (
+            StatusCode::OK,
+            r#"{"answers":{}}"#,
+            "jev response must contain",
+        ),
+        (
+            StatusCode::OK,
+            r#"{"answers":{"matches":{"type":"noul","noul":1.1}}}"#,
+            "jev response must contain",
+        ),
+        (
+            StatusCode::OK,
+            r#"{"answers":{"matches":{"type":"score","noul":0.9}}}"#,
+            "jev response must contain",
+        ),
+    ] {
+        let server = MockServer::start(
+            Router::new().route("/v1/systemone", post(move || async move { (status, body) })),
+        )
+        .await;
+        let ctx = context(server.function());
+        let error = ctx
+            .sql("SELECT jev('text', 'condition', 0.8)")
+            .await
+            .unwrap()
+            .collect()
+            .await
+            .unwrap_err();
+        assert!(error.to_string().contains(expected), "{error}");
+    }
+}
+
+#[tokio::test]
+#[ignore = "calls the real Jev API; requires JEV_API_KEY and GREPTIMEDB_EXPERIMENTAL_JEV=true"]
+async fn test_jev_live() {
+    let ctx = SessionContext::new();
+    let function = FUNCTION_REGISTRY.get_function("jev").unwrap();
+    ctx.register_udf(function.provide(FunctionContext::default()));
+    let batches = ctx.sql(
+        "SELECT id FROM (VALUES
+            (1, 'Payment failed permanently: all three retries exhausted; the payment is still unsuccessful.'),
+            (2, 'Payment succeeded on the second retry; the payment is complete.'),
+            (3, 'User logged in successfully.')
+         ) AS events(id, message)
+         WHERE jev(message, 'The event reports that a payment still failed after retries.', 0.8)
+         ORDER BY id",
+    ).await.unwrap().collect().await.unwrap();
+    assert_batches_eq!(["+----+", "| id |", "+----+", "| 1  |", "+----+"], &batches);
+}