Spark and Databricks

drex-spark calls Drex, Nace's decision models, from Spark. Its functions answer typed questions about each row and return calibrated probabilities as Spark values. See Questions for how to write them. There are three types:

• drex_choice picks one label from a set of options.
• drex_score rates the input on ordered levels.
• drex_noul gives the probability that a yes/no statement is true.

drex_decide asks several questions in one request and returns Drex's response as JSON, and drex_decide_image does the same with images. The functions run wherever PySpark 3.5 or later runs pandas UDFs: Databricks, EMR, Dataproc, Fabric or your own cluster. On a Databricks SQL warehouse, they are Unity Catalog SQL functions.

InstallCopy Anchor Link

pip install drex-spark

On Databricks, run %pip install drex-spark in a notebook, or add it to the cluster's or job's libraries as a PyPI package. It requires Python 3.10 and Spark 3.5 or later (Databricks Runtime 15 or later). Its source is nace-ai/drex-spark on GitHub, under the MIT license.

Set your API keyCopy Anchor Link

The functions read NACE_API_KEY, and NACE_BASE_URL if it's set, on the executors. The key is never a function argument, so it stays out of your code and query history. Set it from your platform's secret store:

• Databricks clusters: add spark.executorEnv.NACE_API_KEY {{secrets/<scope>/<key>}} to the cluster's Spark config. Serverless compute takes no Spark config, so its executors can't get the key: on serverless notebooks and jobs, call the Databricks SQL functions with spark.sql or %sql instead.
• Other Spark: pass --conf spark.executorEnv.NACE_API_KEY=... from your secret manager, or set it in the executors' environment. In local mode, the driver's environment is used.

Create a key as shown in Authentication. NACE_BASE_URL must use https://. http:// is accepted only for localhost, 127.0.0.1 and [::1].

UsageCopy Anchor Link

Ask one question per rowCopy Anchor Link

Each non-NULL row is one Drex request. A NULL input returns NULL and isn't sent or billed.

from drex_spark.functions import drex_choice, drex_noul

answers = tickets.select(
    "ticket_id",
    drex_choice(
        "body",
        "Which team should handle this support ticket?",
        {
            "billing": "Payments, invoices, refunds and charges",
            "technical": "Bugs, crashes and errors in the product",
            "sales": None,
        },
    ).alias("team"),
    drex_noul("body", "Is the customer asking for a refund?").alias("refund"),
)
answers.write.saveAsTable("ticket_answers")

Labels are a list (["billing", "technical"]), a list of {"label", "description"}, or an object of label to description, as a Python value or JSON text: 2 to 255 unique labels. A score's levels are a plain list in scale order, with no descriptions, so describe the scale in the instructions. Every function also takes model and concurrency keywords. The return type depends on the function:

probabilities always holds every option, sorted by label for a choice and by index for a score. Reading answers explains each field.

A STRING input is sent as text. A STRUCT, MAP or (on Spark 4) VARIANT is sent as structured state, so drex_noul(F.struct("subject", "body"), ...) asks about several columns at once.

The instructions, labels and questions are checked when you build the column, before any row is read. Spark SQL and Databricks SQL get them as columns, so they check each distinct value once per batch, and a value that fails a check fails the query.

Ask several questions at onceCopy Anchor Link

drex_decide sends every question in one request per row, so the row is sent and billed once instead of once per question:

from pyspark.sql import functions as F
from drex_spark.functions import drex_decide

questions = {
    "urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["can wait", "this week", "today"]},
    "refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"},
}
answers = tickets.select("ticket_id", drex_decide(F.to_json(F.struct("body")), questions).alias("response"))
answers.select(
    "ticket_id",
    F.get_json_object("response", "$.answers.urgency.score").cast("double").alias("urgency"),
    F.get_json_object("response", "$.answers.refund.noul").cast("double").alias("refund"),
)

Use thresholdsCopy Anchor Link

Every answer carries its probabilities, so you decide what is confident enough and leave the rest for a person or another model:

confident = answers.where(F.col("team.confidence") >= 0.8).select("ticket_id", "team.choice")

Choose the threshold from a labelled sample; Reading answers covers what each probability means.

Read imagesCopy Anchor Link

Drex Max reads images. Read them with Spark's binaryFile source and pass the content column, or an array of up to 8 of them. Prompts refer to them as "Image 1", "Image 2" and so on. PNG, JPEG and WebP are accepted, and Drex never fetches a URL.

from drex_spark.functions import drex_choice

scans = spark.read.format("binaryFile").load("/Volumes/main/docs/scans/*.jpg")
scans.select(
    "path",
    drex_choice(
        "content",
        "What kind of document is Image 1?",
        {"invoice": "A bill or invoice", "receipt": "A proof of payment", "other": None},
        model="drex-max",
    ).alias("kind"),
)

Images with a model that doesn't read them, the default included, fail the query. drex_decide_image defaults to drex-max. A row with more than 8 images, an image that isn't PNG, JPEG or WebP, or a request that would exceed the 4 MiB body limit returns NULL, so resize large images first. Drex Max's image limits list the rest.

Spark SQLCopy Anchor Link

register makes the five functions available to Spark SQL. Registered functions take every argument, and a NULL model means the default:

import drex_spark

drex_spark.register(spark)
spark.sql("""
    SELECT ticket_id, drex_choice(body, 'Which team should handle this ticket?', '["billing", "technical"]', NULL) AS team
    FROM tickets
""")

Databricks SQLCopy Anchor Link

On a Databricks SQL warehouse, the functions are Unity Catalog SQL functions. Analysts call them from SQL with no cluster setup. Create them once, as the principal they should run as:

1. Store the key as a Unity Catalog secret, for example main.drex.nace_api_key (Catalog Explorer › Create › Secret). The functions run with their owner's permissions: the owner needs READ SECRET on the secret, and callers need only EXECUTE on the functions.

2. Allow network access. Python UDFs on a serverless SQL warehouse can't reach the internet until a workspace admin enables the Enable networking for isolated workloads in Serverless SQL Warehouses preview on the workspace's Previews page, then restart the warehouse: one that was running keeps refusing. If the workspace restricts serverless egress, allow console.nace.ai and pypi.org.

3. Create the functions. This prints the statements to run in the SQL editor:

   Copy Text

   pip install drex-spark
   python -m drex_spark.databricks --catalog main --schema drex --secret main.drex.nace_api_key > drex.sql

   Each function installs the same version of drex-spark from PyPI into its own environment. To install it from a volume instead, add --dependency /Volumes/<catalog>/<schema>/<volume>/<wheel>. The functions call https://console.nace.ai unless you pass --base-url.

SELECT
  ticket_id,
  main.drex.drex_choice(body, 'Which team should handle this ticket?', '["billing", "technical", "sales"]') AS team,
  main.drex.drex_noul(body, 'Is the customer asking for a refund?') AS refund
FROM tickets;

The labels and questions are JSON text, and model => 'drex-max' selects another model. Unity Catalog doesn't allow two functions with the same name, so input is text: ask about JSON state with drex_decide, and about images with drex_decide_image(state, images, questions), which takes an ARRAY<BINARY>. The functions run on pro and serverless SQL warehouses and on classic compute with Databricks Runtime 19 or later in standard access mode. They can't be called directly on dedicated access mode compute.

ModelsCopy Anchor Link

The default model is drex-v2.0, which reads text and JSON. Pass model="drex-max" for Drex Max, which also reads images. drex-v2.0 costs $0.03 and Drex Max $0.03 per million input tokens; see Pricing.

The model must be one that GET /v1/models lists for your key. An unknown model fails the query.

Concurrency and rate limitsCopy Anchor Link

spark.drex.concurrency is the job's budget of requests in flight, 8 by default. Each Python worker takes max(1, budget // parallelism) of it, with the parallelism read when you build the column or call register, so set the budget before. That fits the in-flight limit of a free account. It isn't a strict cap: with more workers than the budget, or on Spark Connect, where each worker gets the whole budget, more requests can be in flight, and Drex's 429s pace the rest:

On a paid account, raise it for the session or for one call:

spark.conf.set("spark.drex.concurrency", "16")
tickets.select(drex_noul("body", "Is this a complaint?", concurrency=16))

The Databricks SQL functions cap each batch at 2 requests in flight, since a warehouse runs many batches at once (about 20 on a Small serverless warehouse). A function can't see the warehouse's parallelism, so pass --concurrency when you create them to change the cap. Every API key on an account shares its limits (see Scope). Drex answers an account over its limit with a 429 and the wait: the request waits its turn, so a query over the limit runs at the account's pace instead of returning NULLs.

ErrorsCopy Anchor Link

• 401, 402, a 422 about the questions or model, an unknown model and a TLS failure fail the query. A 422 message lists each issue's path and message, and a 402 links to Billing.
• 429 and 529 mean Drex refused the request before doing any work, so they are retried after retry-after-ms for as long as Drex keeps answering the worker's other requests. A row returns NULL only if Drex answers none of them for 120 seconds.
• 408, 500, 502, 503, 504 and no response are retried up to 5 attempts within 120 seconds. A row whose retries run out returns NULL. Each wait is retry-after-ms (or retry-after), at most 30 seconds.
• A 422 about the row alone (too long, or issues only under state or images), a row over 4 MiB, and any other status, such as 403 or 404, return NULL, and the query continues.
• Each request times out after 60 seconds.

A NULL result is a row Drex didn't answer, not a negative answer. A drex_choice or drex_score row without an answer has every field NULL, so test team.choice IS NULL. Warnings go to the executor log. See Errors and retries.

What gets sent, and billedCopy Anchor Link

Every row Spark evaluates a Drex function on is a request. Every request Drex answers with a 200 is billed, whether or not the job keeps the row; a retry is billed twice only if a 200 was lost on its way back. Spark can evaluate a row more than once:

• A retried task, after a lost executor or a failed stage, sends its rows again.
• A sort over a Drex column, such as orderBy on an answer, evaluates every row twice, because Spark samples the input first.
• Every action on a DataFrame, such as a second count(), a self-join or a display(), runs it again.

Save the answers before you sort, filter or reuse them: write them to a table, or cache() the DataFrame and run count() once. Filter and limit the rows before calling Drex.

answers = tickets.limit(100).select("ticket_id", drex_noul("body", "Is this a complaint?").alias("complaint"))
answers.write.saveAsTable("ticket_complaints")
spark.table("ticket_complaints").orderBy("complaint")

CostCopy Anchor Link

Each row is billed for its input tokens at the model's price, and output tokens aren't billed. For 10,000,000 rows:

The same row counts a different number of tokens on each model, so read usage.input_tokens from drex_decide for a sample of rows on the model you plan to use before running a whole table.