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

README.md

@@ -24,6 +24,61 @@ 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).
+
+`ai_generate` is a Spark column function that calls a Gemini model on every
+row, mirroring the semantics of BigQuery's
+[`AI.GENERATE`](https://cloud.google.com/bigquery/docs/reference/standard-sql/AI-functions#ai.generate).
+It returns a `STRUCT` with a `result`, the raw `full_response`, and a `status`
+that is empty on success.
+
+```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.concat(F.lit("Airport code for "), F.col("city"))),
+).select("city", "generated.result", "generated.status")
+
+result_df.show()
+# +-----+------+------+
+# | city|result|status|
+# +-----+------+------+
+# |Paris|   CDG|      |
+# |Tokyo|   HND|      |
+# +-----+------+------+
+```
+
+Pass `output_schema` to get typed columns back instead of a single string:
+
+```python
+df.withColumn(
+    "review",
+    ai_generate(
+        F.col("text"),
+        output_schema="sentiment STRING, score FLOAT64",
+    ),
+)
+```
+
+The same function can be registered for use from Spark SQL:
+
+```python
+from google.cloud.dataproc_ml.sql import ai_generate_udf
+
+spark.udf.register("ai_generate", ai_generate_udf())
+spark.sql("SELECT ai_generate('Airport code for Paris').result").show()
+```
+
 ### 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,44 @@
+# 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 mirror the BigQuery ``AI.*`` functions, so a query written for
+one engine reads the same on the other.
+
+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
+    ... )
+
+The same function can be registered for use from Spark SQL:
+
+    >>> from google.cloud.dataproc_ml.sql import ai_generate_udf
+    >>>
+    >>> spark.udf.register("ai_generate", ai_generate_udf())
+    >>> spark.sql("SELECT ai_generate(feedback).result FROM feedback_table")
+"""
+
+from ._ai_generate import ai_generate
+from ._ai_generate import ai_generate_udf
+
+__all__ = (
+    "ai_generate",
+    "ai_generate_udf",
+)