feat(sql): ai_generate, a Spark column function for Gemini

README.md

@@ -24,6 +24,166 @@ pip install dataproc-ml
 
 Here are a couple of examples demonstrating how to use the handlers for distributed inference on a Spark DataFrame.
 
+### AI Functions: `ai_generate`
+
+> **Note:** `ai_generate` makes API calls to Vertex AI, which will incur costs.
+> Please review the [Vertex AI Generative
+> AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing).
+
+> **Note:** The AI functions work on Apache Spark 3.5 and later. The only thing
+> that needs Spark 4 is `parse_json` on `full_response`. Spark SQL settings
+> travel in a struct rather than as `name => value` named arguments, because
+> named arguments do not reach a Python UDF before Spark 4.0 — the query fails
+> at run time with `[INTERNAL_ERROR] Cannot evaluate expression`. Named
+> arguments do work on Spark 4 for anyone who prefers them.
+
+`ai_generate` is a Spark column function that calls a Gemini model on every
+row. It always returns the same struct, whatever the arguments:
+
+```
+STRUCT<result STRING, full_response STRING, status STRING>
+```
+
+`result` is the generated text, and `full_response` is the whole model response
+as JSON. `status` is empty when the row succeeded and otherwise describes the
+failure, led by its canonical error code, so `WHERE status <> ''` collects the
+rows that need attention. A row that fails reports it in `status` rather than
+failing the query; a mistake in the *call*, such as a misspelled model, fails
+the query immediately instead.
+
+```python
+from pyspark.sql import SparkSession, functions as F
+from google.cloud.dataproc_ml.sql import ai_generate
+
+spark = SparkSession.builder.getOrCreate()
+
+df = spark.createDataFrame([("Paris",), ("Tokyo",)], ["city"])
+
+result_df = df.withColumn(
+    "generated",
+    ai_generate([F.lit("Airport code for "), F.col("city")]),
+).select("city", "generated.result", "generated.status")
+
+result_df.show()
+# +-----+------+------+
+# | city|result|status|
+# +-----+------+------+
+# |Paris|   CDG|      |
+# |Tokyo|   HND|      |
+# +-----+------+------+
+```
+
+`full_response` is a JSON string, so it is readable on every supported
+version. On Spark 4 you can also query inside it without leaving SQL:
+
+```python
+# Spark 4 only: parse_json and variant_get were added in 4.0.
+df.withColumn("generated", ai_generate(F.col("city"))).selectExpr(
+    "variant_get(parse_json(generated.full_response),"
+    " '$.usage_metadata.total_token_count', 'int') AS tokens"
+)
+```
+
+Pass `output_schema` as a Spark DDL string to constrain the model to a shape.
+`result` then holds a JSON document matching it, which you can project with
+`from_json` when you want typed columns:
+
+```python
+schema = "sentiment STRING, score DOUBLE"
+
+df.withColumn(
+    "review",
+    ai_generate(F.col("text"), output_schema=schema),
+).select(F.from_json(F.col("review.result"), schema).alias("parsed"))
+```
+
+Model settings are passed as a JSON string matching the Vertex AI
+`generateContent` request body, so anything the API supports works without
+this library having to know about it:
+
+```python
+ai_generate(
+    F.col("text"),
+    model_params='{"generation_config": {"temperature": 0.0, "max_output_tokens": 512}}',
+)
+```
+
+`endpoint`, `model_params` and `output_schema` configure the call rather than
+the row, so they have to be constant for the query. `prompt` is the argument
+that varies.
+
+#### Prompts that include a file
+
+A prompt can mix text and files. Vertex AI reads the file itself, so nothing is
+downloaded into Spark, and the media type is inferred from the file extension
+unless you give it:
+
+```python
+from google.cloud.dataproc_ml.sql import ai_generate, file
+
+invoices.withColumn(
+    "extracted",
+    ai_generate(
+        [F.lit("Extract the invoice details."), file(F.col("uri"))],
+        output_schema="invoice_number STRING, total DOUBLE",
+    ),
+)
+```
+
+As with `prompt`, a plain string passed to `file` names a *column*; wrap a
+literal path in `F.lit`.
+
+#### Spark SQL
+
+The same function can be registered for use from Spark SQL. The call is
+`ai_generate(prompt [, endpoint] [, options])`; trailing arguments may be left
+out, and the `options` struct names each setting, so they can be written in any
+order and omitted individually:
+
+```python
+from google.cloud.dataproc_ml.sql import ai_generate_udf
+
+spark.udf.register("ai_generate", ai_generate_udf())
+
+spark.sql("""
+    SELECT ai_generate(
+        CONCAT('Airport code for ', city),
+        'gemini-3.6-flash',
+        named_struct('output_schema', 'code STRING')
+    ).result
+    FROM cities
+""").show()
+```
+
+`ai_generate(body)` is a complete call. Because `endpoint` is a string and
+`options` is a struct, the two are told apart by type, so there is never a
+`NULL` to pad with:
+
+```sql
+SELECT ai_generate(body, named_struct('output_schema', 'code STRING')).result
+FROM cities
+```
+
+A name that is not a setting is an error rather than being ignored, so a typo
+cannot quietly do nothing.
+
+In SQL a prompt is either a string or a struct whose fields are, in order, the
+parts of the prompt. A field holding `named_struct('uri', ...)` is a file, and
+a top-level `named_struct('uri', ...)` is itself a single file:
+
+```sql
+SELECT g.result
+FROM (
+    SELECT ai_generate(
+        struct('Extract the invoice details.',
+               named_struct('uri', uri)),
+        named_struct('output_schema', 'invoice_number STRING, total DOUBLE')
+    ) AS g
+    FROM invoices
+)
+WHERE g.status = ''
+```
+
 ### Generative AI (Gemini) Model Inference
 
 > **Note:** Using the `GenAiModelHandler` involves making API calls to 

docs/api.rst

@@ -5,5 +5,9 @@ This page contains the auto-generated API reference for the public modules
 in the ``dataproc-ml`` library.
 
 .. automodule:: google.cloud.dataproc_ml.inference
+   :members:
+   :show-inheritance:
+
+.. automodule:: google.cloud.dataproc_ml.sql
    :members:
    :show-inheritance:

google/cloud/dataproc_ml/sql/__init__.py

@@ -0,0 +1,65 @@
+# Copyright 2026 Google LLC
+#
+# 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
+#
+#     https://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.
+
+"""Declarative AI functions for Spark DataFrames and Spark SQL.
+
+These functions apply Google's Gemini models to a DataFrame column without any
+of the usual plumbing: no client setup, no retry loops and no response
+parsing. They work on Apache Spark 3.5 and later.
+
+Example:
+    >>> from pyspark.sql.functions import col
+    >>> from google.cloud.dataproc_ml.sql import ai_generate
+    >>>
+    >>> summarized = df.withColumn(
+    ...     "summary", ai_generate(col("feedback")).result
+    ... )
+
+A prompt may include a file, which Vertex AI reads directly:
+
+    >>> from google.cloud.dataproc_ml.sql import ai_generate, file
+    >>>
+    >>> described = df.withColumn(
+    ...     "g", ai_generate([lit("Describe this image."), file(col("uri"))])
+    ... )
+
+The same function can be registered for use from Spark SQL, where the call is
+``ai_generate(prompt [, endpoint] [, options])`` and the settings travel in
+the ``options`` struct:
+
+    >>> from google.cloud.dataproc_ml.sql import ai_generate_udf
+    >>>
+    >>> spark.udf.register("ai_generate", ai_generate_udf())
+    >>> spark.sql('''
+    ...     SELECT ai_generate(
+    ...         concat('Summarize: ', feedback),
+    ...         named_struct('output_schema', 'summary STRING')
+    ...     ).result
+    ...     FROM feedback_table
+    ... ''')
+
+Every call returns ``STRUCT<result STRING, full_response STRING,
+status STRING>``. ``status`` is empty when the row succeeded, so the rows that
+need attention are collected with ``WHERE g.status <> ''``.
+"""
+
+from ._ai_generate import ai_generate
+from ._ai_generate import ai_generate_udf
+from ._ai_generate import file
+
+__all__ = (
+    "ai_generate",
+    "ai_generate_udf",
+    "file",
+)