[Feature Request] Support structured outputs (JSON Schema) in OpenAIChatCompletion

Description

Add support for OpenAI's structured outputs (response_format: {"type": "json_schema", ...}) to OpenAIChatCompletion in both Python and Scala.

Currently, OpenAIChatCompletion returns stringified/unstructured output or standard JSON mode ({"type": "json_object"}). Introducing structured outputs ensures that LLM responses strictly conform to a provided JSON schema or a Spark StructType.

Motivation & Use Case

When performing distributed batch inference across large datasets in Spark:

LLM responses can deviate from expected formatting, breaking downstream transformations.
Users are forced to write custom PySpark UDFs using json.loads or Pydantic for schema validation, which adds serialization overhead and fails when the model outputs malformed JSON.
Native structured output support guarantees 100% adherence to schema, enabling zero-shot deterministic column extraction directly into the DataFrame.

Proposed Solution / API Design

Expose a setter method such as .setResponseFormat() or .setResponseSchema() that accepts a JSON schema string or a Spark StructType:

from synapse.ml.cognitive import OpenAIChatCompletion
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

schema = {
    "type": "object",
    "properties": {
        "sentiment": {"type": "string"},
        "score": {"type": "integer"}
    },
    "required": ["sentiment", "score"],
    "additionalProperties": False
}

chat = (OpenAIChatCompletion()
    .setSubscriptionKey("...")
    .setDeploymentName("gpt-4o")
    .setPromptCol("input_text")
    .setOutputCol("extracted_data")
    .setResponseSchema(schema))

Shresth Kumar Gupta (@Srazyy)

Hey Shresth Kumar Gupta (@Srazyy) 👋!
Thank you so much for reporting the issue/feature request 🚨.
Someone from SynapseML Team will be looking to triage this issue soon.
We appreciate your patience.

Shresth Kumar Gupta (@Srazyy)

Thank you Shresth Kumar Gupta (@Srazyy) for contributing to improving SynapseML and for explaining this use case!

The core structured-output support already exists in SynapseML 1.1.0 and later. It was added in #2428 and refined in #2436. OpenAIChatCompletion.setResponseFormat accepts a Python dictionary or Scala Map[String, Any], including type: "json_schema" and strict: true.

There are two separate steps:

Use setResponseFormat to request schema-constrained JSON from the model.
Use Spark's built-in from_json to turn the returned JSON text into a typed struct. This does not require a Python UDF or per-row json.loads.

Python

This example assumes an existing Spark session named spark and matching SynapseML Python/JVM packages. Configure AZURE_OPENAI_API_KEY, AZURE_OPENAI_SERVICE_NAME, and AZURE_OPENAI_DEPLOYMENT_NAME in your environment. The service name is the Azure OpenAI resource name, not its full URL. Choose a deployment whose model supports structured outputs.

import os

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, StringType, IntegerType
from synapse.ml.services.openai import OpenAIChatCompletion

schema = {
    "type": "object",
    "properties": {
        "sentiment": {"type": "string"},
        "score": {"type": "integer"},
    },
    "required": ["sentiment", "score"],
    "additionalProperties": False,
}

response_format = {
    "type": "json_schema",
    "json_schema": {
        "name": "sentiment_schema",
        "strict": True,
        "schema": schema,
    },
}

df = spark.createDataFrame(
    [("Extract sentiment and an integer score from 0 to 10: I love this product.",)],
    ["input_text"],
)
messages_df = df.withColumn(
    "messages",
    F.array(
        F.struct(
            F.lit("user").alias("role"),
            F.col("input_text").alias("content"),
        )
    ),
)

chat = (
    OpenAIChatCompletion()
    .setSubscriptionKey(os.environ["AZURE_OPENAI_API_KEY"])
    .setCustomServiceName(os.environ["AZURE_OPENAI_SERVICE_NAME"])
    .setDeploymentName(os.environ["AZURE_OPENAI_DEPLOYMENT_NAME"])
    .setMessagesCol("messages")
    .setResponseFormat(response_format)
    .setOutputCol("response")
    .setErrorCol("error")
)

spark_schema = StructType([
    StructField("sentiment", StringType()),
    StructField("score", IntegerType()),
])

result = chat.transform(messages_df).withColumn(
    "extracted_data",
    F.from_json(
        F.col("response.choices")[0]["message"]["content"],
        spark_schema,
    ),
)

result.select(
    "extracted_data.*",
    "error",
    F.col("response.choices")[0]["finish_reason"].alias("finish_reason"),
).show(truncate=False)

Scala

The equivalent Scala example uses the same environment variables and an existing spark session:

import com.microsoft.azure.synapse.ml.services.openai.OpenAIChatCompletion
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import spark.implicits._

val schema: Map[String, Any] = Map(
  "type" -> "object",
  "properties" -> Map(
    "sentiment" -> Map("type" -> "string"),
    "score" -> Map("type" -> "integer")
  ),
  "required" -> Seq("sentiment", "score"),
  "additionalProperties" -> false
)

val responseFormat: Map[String, Any] = Map(
  "type" -> "json_schema",
  "json_schema" -> Map(
    "name" -> "sentiment_schema",
    "strict" -> true,
    "schema" -> schema
  )
)

val messagesDf = Seq(
  "Extract sentiment and an integer score from 0 to 10: I love this product."
).toDF("input_text").withColumn(
  "messages",
  array(struct(lit("user").as("role"), col("input_text").as("content")))
)

val chat = new OpenAIChatCompletion()
  .setSubscriptionKey(sys.env("AZURE_OPENAI_API_KEY"))
  .setCustomServiceName(sys.env("AZURE_OPENAI_SERVICE_NAME"))
  .setDeploymentName(sys.env("AZURE_OPENAI_DEPLOYMENT_NAME"))
  .setMessagesCol("messages")
  .setResponseFormat(responseFormat)
  .setOutputCol("response")
  .setErrorCol("error")

val sparkSchema = new StructType()
  .add("sentiment", StringType)
  .add("score", IntegerType)

val result = chat.transform(messagesDf).withColumn(
  "extracted_data",
  from_json(col("response.choices").getItem(0).getField("message").getField("content"), sparkSchema)
)

result.select(
  col("extracted_data.*"),
  col("error"),
  col("response.choices").getItem(0).getField("finish_reason").as("finish_reason")
).show(false)

What remains different from the proposed API

OpenAIChatCompletion takes setMessagesCol, not setPromptCol, and keeps the response envelope. The JSON text is in choices[0].message.content.
There is currently no setResponseSchema method or direct conversion from a Spark StructType to an OpenAI JSON Schema. The StructType above describes the Spark parsing step, not the model request.
Schema input to setResponseFormat must be a dictionary/Scala map, not a JSON-encoded string. If your schema starts as a JSON string, parse it once on the driver before configuring the transformer.
For prompt templates and built-in typed output, OpenAIPrompt also supports setResponseFormat, setPromptTemplate, setPostProcessing("json"), and setPostProcessingOptions({"jsonSchema": "sentiment STRING, score INT"}) in Python.

So the JSON Schema request support and UDF-free typed extraction are already available. Direct StructType conversion or a setResponseSchema convenience method would be separate enhancements.

I checked the released source and validated the published SynapseML 1.1.3 Python package and JVM jars on Spark 3.5.0. Thirteen focused tests passed with no skips, including four mock HTTP requests checking the exact strict-schema payload, typed extraction, OpenAIPrompt JSON post-processing, rejected input forms, and save/load. A separately compiled Scala check also passed. These were local mock-service checks, not live-model inference.

Strict schema support does not eliminate service errors, refusals, or truncated responses, and it does not guarantee that extracted values are factually correct. Check error, finish_reason, and the returned content before consuming a result.

If your runtime includes an older SynapseML version, make sure both the Python package and the JVM jars are updated together.

References: original feature, follow-up, released implementation.

fix: preserve Spark port replay and exercise schema-only options

Shresth Kumar Gupta (@Srazyy)

Thank you again, Shresth Kumar Gupta (@Srazyy), for contributing to improving SynapseML.

The existing setResponseFormat API supports structured outputs, as described above. We can make the experience simpler by accepting the inner JSON Schema directly and adding the API-specific response-format envelope automatically. I've opened PR #2708 for this improvement: #2708.

The proposed setResponseSchema helper works with OpenAIPrompt, OpenAIChatCompletion, and OpenAIResponses, in both Python and Scala. Its defaults are name="response_schema" and strict=True.

These examples require a build containing PR #2708. The new setter is not in SynapseML 1.1.3. Existing released versions can continue using the full-envelope setResponseFormat examples in the earlier reply.

By "partial schema," we mean the inner JSON Schema without the outer response-format metadata. The helper preserves that schema; it does not infer required, add additionalProperties, convert a Spark StructType, or fill in missing schema definitions. For strict output, supply a schema supported by the selected model.

Python

Run in a configured Spark session. Set AZURE_OPENAI_ENDPOINT to the resource root, such as https://<resource>.openai.azure.com, and supply AZURE_OPENAI_API_KEY and AZURE_OPENAI_DEPLOYMENT through your environment or secret-management system. The deployment must support structured outputs and the selected API.

Shared setup:

import os
from pyspark.sql import SparkSession, functions as F
from synapse.ml.services.openai import (
    OpenAIChatCompletion,
    OpenAIPrompt,
    OpenAIResponses,
)

spark = SparkSession.builder.getOrCreate()
endpoint = os.environ["AZURE_OPENAI_ENDPOINT"].rstrip("/") + "/openai/v1"
key = os.environ["AZURE_OPENAI_API_KEY"]
deployment = os.environ["AZURE_OPENAI_DEPLOYMENT"]

schema = {
    "type": "object",
    "properties": {
        "sentiment": {"type": "string", "enum": ["positive", "negative", "neutral"]},
        "score": {"type": "integer"},
    },
    "required": ["sentiment", "score"],
    "additionalProperties": False,
}
result_schema = "sentiment STRING, score INT"
texts = spark.createDataFrame([("I love this product!",)], ["text"])
messages = texts.withColumn(
    "messages",
    F.array(F.struct(F.lit("user").alias("role"), F.col("text").alias("content"))),
)

OpenAIPrompt

OpenAIPrompt can parse the JSON into a Spark struct through its existing native JSON post-processing:

prompt = (
    OpenAIPrompt()
    .setUrl(endpoint)
    .setSubscriptionKey(key)
    .setDeploymentName(deployment)
    .setApiType("chat_completions")
    .setPromptTemplate(
        "Classify the sentiment as positive, negative or neutral "
        "and provide an integer score. Text: {text}"
    )
    .setResponseSchema(schema)
    .setPostProcessing("json")
    .setPostProcessingOptions({"jsonSchema": result_schema})
    .setOutputCol("result")
    .setErrorCol("error")
)
prompt.transform(texts).select("text", "result", "error").show(truncate=False)

Use .setApiType("responses") to send the same prompt through the Responses API instead. The schema setter and post-processing stay the same.

Chat Completions

Chat Completions uses setMessagesCol. Its output remains the response envelope; use Spark's from_json to parse the assistant's JSON text:

chat = (
    OpenAIChatCompletion()
    .setUrl(endpoint)
    .setSubscriptionKey(key)
    .setDeploymentName(deployment)
    .setMessagesCol("messages")
    .setResponseSchema(schema)
    .setOutputCol("result")
    .setErrorCol("error")
)
chat.transform(messages).select(
    "text",
    F.from_json(
        F.expr("try_element_at(result.choices, 1).message.content"), result_schema
    ).alias("parsed"),
    "error",
).show(truncate=False)

Responses

OpenAIResponses also uses setMessagesCol. The JSON text is in an output_text content item. This expression finds that item without assuming that the first output entry contains text:

responses = (
    OpenAIResponses()
    .setUrl(endpoint)
    .setSubscriptionKey(key)
    .setDeploymentName(deployment)
    .setMessagesCol("messages")
    .setResponseSchema(schema)
    .setOutputCol("result")
    .setErrorCol("error")
)
response_text = F.expr(
    "try_element_at(filter(flatten(transform("
    "filter(result.output, x -> x.content IS NOT NULL), x -> x.content)), "
    "x -> x.type = 'output_text'), 1).text"
)
responses.transform(messages).select(
    "text",
    F.from_json(response_text, result_schema).alias("parsed"),
    "error",
).show(truncate=False)

To customize the metadata on any of these stages, use .setResponseSchema(schema, name="review_schema", strict=False). Omitting those arguments requests the default name and strict output.

Scala equivalents

Scala

The same three APIs accept a Map[String, Any]. Use ListMap where JSON property order matters.

Shared setup, using the same environment variables:

import com.microsoft.azure.synapse.ml.services.openai.{
  OpenAIChatCompletion, OpenAIPrompt, OpenAIResponses
}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.{array, col, expr, from_json, lit, struct}
import org.apache.spark.sql.types.StructType
import scala.collection.immutable.ListMap

val spark = SparkSession.builder().getOrCreate()
import spark.implicits._

val endpoint = sys.env("AZURE_OPENAI_ENDPOINT").stripSuffix("/") + "/openai/v1"
val key = sys.env("AZURE_OPENAI_API_KEY")
val deployment = sys.env("AZURE_OPENAI_DEPLOYMENT")
val schema: Map[String, Any] = ListMap(
  "type" -> "object",
  "properties" -> ListMap(
    "sentiment" -> Map("type" -> "string", "enum" -> Seq("positive", "negative", "neutral")),
    "score" -> Map("type" -> "integer")
  ),
  "required" -> Seq("sentiment", "score"),
  "additionalProperties" -> false
)
val resultSchema = "sentiment STRING, score INT"
val resultType = StructType.fromDDL(resultSchema)
val texts = Seq("I love this product!").toDF("text")
val messages = texts.withColumn(
  "messages", array(struct(lit("user").as("role"), col("text").as("content")))
)

OpenAIPrompt

val prompt = new OpenAIPrompt()
  .setUrl(endpoint)
  .setSubscriptionKey(key)
  .setDeploymentName(deployment)
  .setApiType("chat_completions")
  .setPromptTemplate(
    "Classify the sentiment as positive, negative or neutral " +
      "and provide an integer score. Text: {text}"
  )
  .setResponseSchema(schema)
  .setPostProcessing("json")
  .setPostProcessingOptions(Map("jsonSchema" -> resultSchema))
  .setOutputCol("result")
  .setErrorCol("error")

prompt.transform(texts).select("text", "result", "error").show(truncate = false)

As in Python, .setApiType("responses") selects the Responses API.

Chat Completions

val chat = new OpenAIChatCompletion()
  .setUrl(endpoint)
  .setSubscriptionKey(key)
  .setDeploymentName(deployment)
  .setMessagesCol("messages")
  .setResponseSchema(schema)
  .setOutputCol("result")
  .setErrorCol("error")

chat.transform(messages).select(
  col("text"),
  from_json(expr("try_element_at(result.choices, 1).message.content"), resultType).as("parsed"),
  col("error")
).show(truncate = false)

Responses

val responses = new OpenAIResponses()
  .setUrl(endpoint)
  .setSubscriptionKey(key)
  .setDeploymentName(deployment)
  .setMessagesCol("messages")
  .setResponseSchema(schema)
  .setOutputCol("result")
  .setErrorCol("error")

val responseText = expr(
  "try_element_at(filter(flatten(transform(" +
    "filter(result.output, x -> x.content IS NOT NULL), x -> x.content)), " +
    "x -> x.type = 'output_text'), 1).text"
)
responses.transform(messages).select(
  col("text"), from_json(responseText, resultType).as("parsed"), col("error")
).show(truncate = false)

Scala also supports .setResponseSchema(schema, strict = false) and .setResponseSchema(schema, name = "review_schema", strict = false).

These examples use Spark-native parsing rather than Python UDFs. Applications should inspect the error column and handle null parsed results, refusals, and incomplete responses; strict output is not a guarantee that every request produces a usable result.

The exact Python and Scala examples above passed local Spark execution against the PR's generated wrappers and JVM classes, with loopback HTTP checks for request formats and typed results. The updated notebook adds eight live-service cases covering defaults and custom non-strict options across the four routes, including both OpenAIPrompt modes.

Live validation passed on the final head: the updated OpenAI.ipynb passed Databricks CI on commit 10bead3b18, executing all eight new-option cases. This is actual Azure inference, not just configuration or loopback validation. Final CI checks are green, with all 65 jobs passing. This does not claim Fabric runtime validation; the fork PR's pipeline excludes Fabric E2E.

Static-schema caching and format compatibility

PR #2708 now also caches unchanged static response-format JSON per stage/executor instance. It preserves column-driven formats, row-dependent Responses verbosity, and existing request-helper overrides. Mutable Scala sequences remain uncached to avoid stale requests.

"text" and "json_object" still use setResponseFormat and remain plain selectors. setResponseSchema wraps an inner schema; existing named/full formats keep using setResponseFormat with their name and strictness preserved. Neither path infers missing constraints for raw schemas.

On the current head, 74 Scala tests, 20 Python tests with 116 subtests, and all eight exact notebook cases passed locally against loopback HTTP. A request-encoding benchmark measured 36-79% fewer allocations across 1-, 10-, and 100-property schemas. This is not an LLM/network latency claim, and the schema still accompanies every request.

The review reports now follow master's layout under reviews/pr-2708/. This review-only update is commit f10ee0aa328d7c5aeab300e449b0c80010450391; every non-review file is unchanged from the tested cache implementation.

Full current-head CI passed: Azure build 236151515 completed all 65 jobs successfully. The affected Scala suites passed 74/74 tests with no skips, and OpenAI.ipynb passed all eight live inference cases on Databricks, covering Chat Completions, Responses, and both Prompt API modes with default and custom non-strict settings. Current-head automated review is clear; only required human approval remains.

The helper remains unreleased while the PR is open; the existing setResponseFormat examples above remain available in released builds. This validation does not claim a Fabric AI Functions runtime test.

style: format schema-only notebook assertions for CI

feat: add schema-only OpenAI structured output configuration

fix: preserve Spark port replay and exercise schema-only options

style: format schema-only notebook assertions for CI

perf: cache static OpenAI response-format JSON

completed PR

Navigation Menu

Repository navigation

[Feature Request] Support structured outputs (JSON Schema) in OpenAIChatCompletion #2700

Description

Motivation & Use Case

Proposed Solution / API Design

Python

Scala

What remains different from the proposed API

Python

OpenAIPrompt

Chat Completions

Responses

Scala

OpenAIPrompt

Chat Completions

Responses

Static-schema caching and format compatibility

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Feature Request] Support structured outputs (JSON Schema) in OpenAIChatCompletion #2700

Description

Description

Motivation & Use Case

Proposed Solution / API Design

Activity

github-actions commented on Sep 6, 2026

ranadeepsingh commented on Sep 15, 2026

Python

Scala

What remains different from the proposed API

ranadeepsingh commented on Sep 15, 2026

Python

OpenAIPrompt

Chat Completions

Responses

Scala

OpenAIPrompt

Chat Completions

Responses

Static-schema caching and format compatibility

ranadeepsingh commented on Sep 17, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions