Fabric Built-in Models with Spark SQL UDFs¶

Run batched Responses and Embeddings from Spark SQL on workers using Fabric's runtime-managed authentication. No Azure OpenAI resource, API key, or driver token broadcast is required. The compact Environment setup below uses openaivec==2.6.0 from PyPI.

Inference consumes Fabric capacity, and the Spark session also incurs compute usage. The input is six synthetic rows in two partitions; response batches and concurrency are deliberately small.

Prepare the Environment¶

1. In a dedicated Fabric Environment using Runtime 1.3, add openaivec==2.6.0 as the sole PyPI entry in External libraries, or import an external library YAML containing dependencies: [{pip: [openaivec==2.6.0]}]. Remove any older openaivec wheel from Custom libraries when switching to PyPI. Do not replace a shared Environment's library list.
2. Publish the Environment in Full mode, attach it to this notebook, and start a new session. The package declares its runtime dependencies, so individual package pins are not required. Do not install openaivec[spark]: Fabric supplies PySpark, SynapseML, and NotebookUtils.
3. Attach a Lakehouse as the default. The code cell writes a JSON report to Files/openaivec-example/spark-sql-udfs.json, replacing only that example report on reruns.

One library entry still installs the package's transitive dependencies. The 2.6.0rc1 candidate was validated with one custom wheel and zero external entries before release. Its 48 package source files matched on the driver and both workers; runtime dependency declarations are unchanged from 2.5.1. This is candidate evidence, not a separate run of the final PyPI package. Resolved versions can differ between publication dates.

The code uses the built-in gpt-5.1 and text-embedding-ada-002 models with API version 2025-04-01-preview. Check capacity, region, tenant permissions, and model availability first. Spark workers need the SynapseML get_openai_httpx_async_client() helper.

Fabric %pip installs on the driver and executors but is session-scoped and disabled in pipeline runs by default. !pip is driver-only; a published Environment is used here for reproducibility.

The candidate validation loaded OpenAI SDK 3.11.0 with aiohttp 3.14.3 on both the driver and workers. An earlier validation found older distribution metadata alongside Fabric overrides. This example's per-process versions records distribution metadata; loaded_versions records seven imported libraries, and aiohttp_socket_timeout_available checks the required transport capability. See the authentication guide for the candidate's live SQL and synthetic regression coverage and platform caveats.

Execute and Verify¶

setup_fabric(spark) must run before constructing UDFs. The UDFs capture model names and non-secret configuration; each partition creates its own async client inside its event loop and closes it on exit.

The following cell registers five functions with spark.udf.register and calls them through spark.sql: string and structured Responses, Embeddings, PreparedTask, and parsing. These session-local functions are also available to subsequent %%sql cells, not to the Lakehouse T-SQL analytics endpoint. Importing openaivec alone does not select authentication or register SQL functions.

The checks cover row-ID correspondence, partition-local duplicate results, response schemas, nonzero 1,536-dimensional embeddings, and a repeated SQL action. Arrow batches contain at most two rows to exercise reuse within a partition. The original Arrow batch setting is restored after execution.

In [ ]:

import importlib.metadata
import json
import os
import platform
from typing import Any

import aiohttp
import httpx
import notebookutils
import numpy as np
import openai
import pandas as pd
import pyarrow as pa
import pydantic
from pydantic import BaseModel, ConfigDict
from pyspark.sql import SparkSession

import openaivec
from openaivec.spark_ext import embeddings_udf, parse_udf, responses_udf, setup_fabric, task_udf


class EchoResult(BaseModel):
    model_config = ConfigDict(extra="forbid")
    original: str


def worker_versions(partition: int | str) -> dict[str, Any]:
    return {
        "partition": partition,
        "pid": os.getpid(),
        "python": platform.python_version(),
        "loaded_versions": {
            "aiohttp": aiohttp.__version__,
            "httpx": httpx.__version__,
            "numpy": np.__version__,
            "openai": openai.__version__,
            "pandas": pd.__version__,
            "pyarrow": pa.__version__,
            "pydantic": pydantic.__version__,
        },
        "aiohttp_socket_timeout_available": hasattr(aiohttp, "SocketTimeoutError"),
        "versions": {
            name: importlib.metadata.version(name)
            for name in (
                "openaivec",
                "openai",
                "httpx",
                "aiohttp",
                "numpy",
                "pandas",
                "pyarrow",
                "pydantic",
                "azure-identity",
                "ipywidgets",
                "tqdm",
                "typing-extensions",
            )
        },
    }


spark = SparkSession.getActiveSession()
if spark is None:
    raise RuntimeError("Run this example in a Fabric PySpark notebook.")
expected_version = importlib.metadata.version("openaivec")
setup_fabric(spark)
original_arrow_batch_size = spark.conf.get("spark.sql.execution.arrow.maxRecordsPerBatch")
inputs = [(0, "alpha"), (1, "beta"), (2, "alpha"), (3, "beta"), (4, "alpha"), (5, "beta")]
source = spark.createDataFrame(spark.sparkContext.parallelize(inputs, 2), "row_id long, text string")
source.createOrReplaceTempView("openaivec_example_inputs")
report: dict[str, Any] = {
    "probe": "spark-sql-udfs",
    "version": expected_version,
    "spark": spark.version,
    "requirements": importlib.metadata.requires("openaivec"),
    "driver": worker_versions("driver"),
    "workers": spark.sparkContext.parallelize([0, 1], 2).map(worker_versions).collect(),
    "completed": False,
    "operations": {},
}
assert all(worker["versions"]["openaivec"] == expected_version for worker in report["workers"])
assert all(worker["loaded_versions"] == report["driver"]["loaded_versions"] for worker in report["workers"])
assert report["driver"]["aiohttp_socket_timeout_available"]
assert all(worker["aiohttp_socket_timeout_available"] for worker in report["workers"])
options: dict[str, Any] = {
    "batch_size": 2,
    "max_concurrency": 1,
    "max_output_tokens": 512,
    "reasoning": {"effort": "none"},
}
echo_task = openaivec.PreparedTask(instructions="Copy each input exactly into original.", response_format=EchoResult)
operations = {
    "string": responses_udf("Return each input unchanged, with no explanation.", **options),
    "structured": responses_udf("Copy each input exactly into original.", response_format=EchoResult, **options),
    "embeddings": embeddings_udf(batch_size=2, max_concurrency=1),
    "task": task_udf(echo_task, **options),
    "parse": parse_udf("Copy each input exactly into original.", response_format=EchoResult, **options),
}
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "2")
try:
    for operation, udf in operations.items():
        report["active_operation"] = operation
        function_name = f"openaivec_example_{operation}"
        spark.udf.register(function_name, udf)
        query = (
            f"SELECT row_id, text, spark_partition_id() AS partition, {function_name}(text) AS prediction "
            "FROM openaivec_example_inputs"
        )
        result = spark.sql(query)
        rows = sorted(result.collect(), key=lambda row: row.row_id)
        assert [(row.row_id, row.text) for row in rows] == inputs
        assert len({row.partition for row in rows}) == 2
        predictions: dict[tuple[int, str], Any] = {}
        for row in rows:
            key = (row.partition, row.text)
            if key in predictions:
                assert predictions[key] == row.prediction
            predictions[key] = row.prediction
            if operation == "embeddings":
                assert len(row.prediction) == 1536
                assert any(value != 0 for value in row.prediction)
            elif operation == "string":
                assert row.text in row.prediction.lower()
            else:
                assert row.prediction.original == row.text
        report["operations"][operation] = {
            "sql": query,
            "row_count": len(rows),
            "partitions": sorted({row.partition for row in rows}),
            "schema": result.schema.jsonValue(),
            "duplicate_outputs_equal": True,
            "input_output_correspondence": True,
            "outputs": [
                {
                    "row_id": row.row_id,
                    "partition": row.partition,
                    "prediction": len(row.prediction)
                    if operation == "embeddings"
                    else row.prediction
                    if operation == "string"
                    else row.prediction.asDict(),
                }
                for row in rows
            ],
        }
    repeat_rows = spark.sql(
        "SELECT openaivec_example_embeddings(text) AS embedding FROM openaivec_example_inputs"
    ).collect()
    assert len(repeat_rows) == len(inputs)
    assert all(len(row.embedding) == 1536 for row in repeat_rows)
    report["repeated_action_success"] = True
    report["completed"] = True
finally:
    if original_arrow_batch_size is None:
        spark.conf.unset("spark.sql.execution.arrow.maxRecordsPerBatch")
    else:
        spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", original_arrow_batch_size)
    notebookutils.fs.mkdirs("Files/openaivec-example")
    notebookutils.fs.put("Files/openaivec-example/spark-sql-udfs.json", json.dumps(report), True)
    print(json.dumps(report, indent=2))

import importlib.metadata import json import os import platform from typing import Any import aiohttp import httpx import notebookutils import numpy as np import openai import pandas as pd import pyarrow as pa import pydantic from pydantic import BaseModel, ConfigDict from pyspark.sql import SparkSession import openaivec from openaivec.spark_ext import embeddings_udf, parse_udf, responses_udf, setup_fabric, task_udf class EchoResult(BaseModel): model_config = ConfigDict(extra="forbid") original: str def worker_versions(partition: int | str) -> dict[str, Any]: return { "partition": partition, "pid": os.getpid(), "python": platform.python_version(), "loaded_versions": { "aiohttp": aiohttp.__version__, "httpx": httpx.__version__, "numpy": np.__version__, "openai": openai.__version__, "pandas": pd.__version__, "pyarrow": pa.__version__, "pydantic": pydantic.__version__, }, "aiohttp_socket_timeout_available": hasattr(aiohttp, "SocketTimeoutError"), "versions": { name: importlib.metadata.version(name) for name in ( "openaivec", "openai", "httpx", "aiohttp", "numpy", "pandas", "pyarrow", "pydantic", "azure-identity", "ipywidgets", "tqdm", "typing-extensions", ) }, } spark = SparkSession.getActiveSession() if spark is None: raise RuntimeError("Run this example in a Fabric PySpark notebook.") expected_version = importlib.metadata.version("openaivec") setup_fabric(spark) original_arrow_batch_size = spark.conf.get("spark.sql.execution.arrow.maxRecordsPerBatch") inputs = [(0, "alpha"), (1, "beta"), (2, "alpha"), (3, "beta"), (4, "alpha"), (5, "beta")] source = spark.createDataFrame(spark.sparkContext.parallelize(inputs, 2), "row_id long, text string") source.createOrReplaceTempView("openaivec_example_inputs") report: dict[str, Any] = { "probe": "spark-sql-udfs", "version": expected_version, "spark": spark.version, "requirements": importlib.metadata.requires("openaivec"), "driver": worker_versions("driver"), "workers": spark.sparkContext.parallelize([0, 1], 2).map(worker_versions).collect(), "completed": False, "operations": {}, } assert all(worker["versions"]["openaivec"] == expected_version for worker in report["workers"]) assert all(worker["loaded_versions"] == report["driver"]["loaded_versions"] for worker in report["workers"]) assert report["driver"]["aiohttp_socket_timeout_available"] assert all(worker["aiohttp_socket_timeout_available"] for worker in report["workers"]) options: dict[str, Any] = { "batch_size": 2, "max_concurrency": 1, "max_output_tokens": 512, "reasoning": {"effort": "none"}, } echo_task = openaivec.PreparedTask(instructions="Copy each input exactly into original.", response_format=EchoResult) operations = { "string": responses_udf("Return each input unchanged, with no explanation.", **options), "structured": responses_udf("Copy each input exactly into original.", response_format=EchoResult, **options), "embeddings": embeddings_udf(batch_size=2, max_concurrency=1), "task": task_udf(echo_task, **options), "parse": parse_udf("Copy each input exactly into original.", response_format=EchoResult, **options), } spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "2") try: for operation, udf in operations.items(): report["active_operation"] = operation function_name = f"openaivec_example_{operation}" spark.udf.register(function_name, udf) query = ( f"SELECT row_id, text, spark_partition_id() AS partition, {function_name}(text) AS prediction " "FROM openaivec_example_inputs" ) result = spark.sql(query) rows = sorted(result.collect(), key=lambda row: row.row_id) assert [(row.row_id, row.text) for row in rows] == inputs assert len({row.partition for row in rows}) == 2 predictions: dict[tuple[int, str], Any] = {} for row in rows: key = (row.partition, row.text) if key in predictions: assert predictions[key] == row.prediction predictions[key] = row.prediction if operation == "embeddings": assert len(row.prediction) == 1536 assert any(value != 0 for value in row.prediction) elif operation == "string": assert row.text in row.prediction.lower() else: assert row.prediction.original == row.text report["operations"][operation] = { "sql": query, "row_count": len(rows), "partitions": sorted({row.partition for row in rows}), "schema": result.schema.jsonValue(), "duplicate_outputs_equal": True, "input_output_correspondence": True, "outputs": [ { "row_id": row.row_id, "partition": row.partition, "prediction": len(row.prediction) if operation == "embeddings" else row.prediction if operation == "string" else row.prediction.asDict(), } for row in rows ], } repeat_rows = spark.sql( "SELECT openaivec_example_embeddings(text) AS embedding FROM openaivec_example_inputs" ).collect() assert len(repeat_rows) == len(inputs) assert all(len(row.embedding) == 1536 for row in repeat_rows) report["repeated_action_success"] = True report["completed"] = True finally: if original_arrow_batch_size is None: spark.conf.unset("spark.sql.execution.arrow.maxRecordsPerBatch") else: spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", original_arrow_batch_size) notebookutils.fs.mkdirs("Files/openaivec-example") notebookutils.fs.put("Files/openaivec-example/spark-sql-udfs.json", json.dumps(report), True) print(json.dumps(report, indent=2))

Interpret the Report¶

A successful run reports probe: "spark-sql-udfs", completed: true, five entries in operations, and repeated_action_success: true. Each operation records its SQL query and six rows across partitions 0 and 1. Driver and worker versions and the package's dependency declarations are recorded without credentials. Distribution metadata can differ from loaded modules in Fabric; compare loaded_versions and the actual UDF checks.

Spark DataFrames have no implicit global row order; this example compares results using row_id. Deduplication is local to each partition, not global across the Spark job. Each new action may incur fresh inference calls. The repeated action does not prove long-running token refresh or identical worker-process reuse. Scheduled execution identities and whole-platform dependency consistency require separate validation.

Recreate UDFs after changing model or authentication settings. Responses are stateless (store=False); stored-response IDs are not supported by this route.