外部 AI 函数

通过将 TiDB Cloud Lake 与您自己的基础设施连接，构建强大的 AI/ML 能力。外部函数让您能够部署自定义模型、利用 GPU 加速，并与任何 ML 框架集成，同时确保数据安全。

关键能力

工作原理

1. Create AI Server：使用 Python 和 tidbcloudlake-udf 构建您的 AI/ML 服务器
2. Register Function：使用 CREATE FUNCTION 将您的服务器连接到 TiDB Cloud Lake
3. Use in SQL：直接在 SQL 查询中调用您的自定义 AI 函数

网络要求

外部函数使用 Arrow Flight over gRPC/HTTP2，而不是 REST。

1. 通过 HTTPS 暴露 UDF 服务器，并提供公共 TLS 证书以及 gRPC/HTTP2 支持。
2. 单击右上角的 Support，选择 Submit a request，并请求将服务器主机名添加到你的租户的 UDF server allowlist 中。
3. 如果你的防火墙受到严格限制，请提交支持请求以获取 TiDB Cloud Lake 出站地址，并在 TCP 443 上放行这些地址。

在主机被加入允许列表之前，CREATE FUNCTION 会因 Unallowed UDF server address 而失败；它还会验证远程 schema，因此该端点必须处于在线状态。

请将本地 UDFServer gRPC 监听器保持为私有；在 gRPC 负载均衡器处终止 TLS：

TiDB Cloud Lake → HTTPS/HTTP2 :443 → UDFServer gRPC :8815

示例：文本嵌入函数

# Simple embedding UDF server demo
from tidbcloudlake_udf import udf, UDFServer
from sentence_transformers import SentenceTransformer

# Load pre-trained model
model = SentenceTransformer('all-mpnet-base-v2')  # 768-dimensional vectors

@udf(
    input_types=["STRING"],
    result_type="ARRAY(FLOAT)",
)
def ai_embed_768(inputs: list[str], headers) -> list[list[float]]:
    """Generate 768-dimensional embeddings for input texts"""
    try:
        # Process inputs in a single batch
        embeddings = model.encode(inputs)
        # Convert to list format
        return [embedding.tolist() for embedding in embeddings]
    except Exception as e:
        print(f"Error generating embeddings: {e}")
        # Return empty lists in case of error
        return [[] for _ in inputs]

if __name__ == '__main__':
    print("Starting embedding UDF server on port 8815...")
    server = UDFServer("0.0.0.0:8815")
    server.add_function(ai_embed_768)
    server.serve()

-- Register the external function in TiDB Cloud Lake
CREATE OR REPLACE FUNCTION ai_embed_768 (STRING)
    RETURNS ARRAY(FLOAT)
    LANGUAGE PYTHON
    HANDLER = 'ai_embed_768'
    ADDRESS = 'https://your-ml-server.example.com';

-- Use the custom embedding in queries
SELECT
    id,
    title,
    cosine_distance(
        ai_embed_768(content),
        ai_embed_768('machine learning techniques')
    ) AS similarity
FROM articles
ORDER BY similarity ASC
LIMIT 5;