feat(sql): add ai_search, the Spark counterpart of BigQuery AI.SEARCH

README.md

@@ -24,6 +24,194 @@ 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()
+```
+
+### AI Functions: `ai_embed`
+
+> **Note:** `ai_embed` 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_embed` turns a text column into dense embedding vectors, mirroring the
+semantics of BigQuery's
+[`AI.EMBED`](https://cloud.google.com/bigquery/docs/reference/standard-sql/AI-functions#ai.embed).
+It returns a `STRUCT` with a `result`, an `ARRAY<DOUBLE>` holding the
+embedding, and a `status` that is empty on success. The rows of a batch are
+embedded together, in as few requests as the endpoint allows.
+
+```python
+from pyspark.sql import SparkSession, functions as F
+from google.cloud.dataproc_ml.sql import ai_embed
+
+spark = SparkSession.builder.getOrCreate()
+
+df = spark.createDataFrame(
+    [("The cat slept on the windowsill.",), ("Rates rose again.",)],
+    ["text"],
+)
+
+result_df = df.withColumn("embedded", ai_embed(F.col("text"))).select(
+    F.size("embedded.result").alias("dimensions"), "embedded.status"
+)
+
+result_df.show()
+# +----------+------+
+# |dimensions|status|
+# +----------+------+
+# |       768|      |
+# |       768|      |
+# +----------+------+
+```
+
+Pass `task_type` to tailor the embeddings to how they will be used, and
+`model_params` for anything else the endpoint accepts:
+
+```python
+df.withColumn(
+    "embedded",
+    ai_embed(
+        F.col("text"),
+        task_type="RETRIEVAL_DOCUMENT",
+        title="Product reviews",
+        model_params={"outputDimensionality": 256},
+    ),
+)
+```
+
+The same function can be registered for use from Spark SQL:
+
+```python
+from google.cloud.dataproc_ml.sql import ai_embed_udf
+
+spark.udf.register("ai_embed", ai_embed_udf())
+spark.sql("SELECT ai_embed(text).result AS vector FROM documents").show()
+```
+
+### AI Functions: `ai_search`
+
+> **Note:** `ai_search` 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_search` finds the rows closest in meaning to a search term, mirroring the
+semantics of BigQuery's
+[`AI.SEARCH`](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-search).
+It takes the whole DataFrame rather than a column and returns the same two
+columns BigQuery does: `base`, a struct of every column of the table searched,
+and `distance`. A smaller distance always means a closer match.
+
+```python
+from pyspark.sql import SparkSession
+from google.cloud.dataproc_ml.sql import ai_search
+
+spark = SparkSession.builder.getOrCreate()
+
+products = spark.createDataFrame(
+    [
+        ("Rocket kite", "A huge stunt kite that loops in a light breeze."),
+        ("Desk lamp", "An adjustable lamp with a warm dimmable bulb."),
+        ("Bath ducks", "A pack of six squeaky ducks for the tub."),
+    ],
+    ["name", "description"],
+)
+
+ai_search(products, "description", "A really fun toy", top_k=2).select(
+    "base.name", "distance"
+).show()
+# +-----------+------------------+
+# |       name|          distance|
+# +-----------+------------------+
+# | Bath ducks|0.7721...         |
+# |Rocket kite|0.8034...         |
+# +-----------+------------------+
+```
+
+BigQuery searches a table whose embeddings are maintained for it in a
+generated column. Spark has no such column, so `ai_search` looks for the
+embeddings in three places, in order:
+
+1. the column named by `embedding_column`;
+2. a column named `<column_to_search>_embedding`, which is the name
+   BigQuery's own examples give the generated column;
+3. failing both, it generates them by calling `ai_embed` on
+   `column_to_search`.
+
+An existing embedding column may hold either a plain `ARRAY<DOUBLE>` or the
+struct that `ai_embed` returns. Prefer to embed the table once and keep the
+result, as in the third case every search re-embeds every row:
+
+```python
+from google.cloud.dataproc_ml.sql import ai_embed
+
+embedded = products.withColumn(
+    "description_embedding", ai_embed("description")
+).cache()
+
+ai_search(embedded, "description", "Something to read by", top_k=1)
+```
+
+Pass `distance_type` to choose how distance is measured, one of `EUCLIDEAN`
+(the default, as in BigQuery), `COSINE` or `DOT_PRODUCT`. Because a smaller
+distance always means a closer match, `DOT_PRODUCT`, which grows with
+similarity, is returned negated.
+
+Rows whose own embedding is missing are skipped. A search term that cannot be
+embedded raises instead, since there is no result to return without it. This
+also matches BigQuery.
+
 ### 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,64 @@
+# 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")
+
+Embeddings follow the same shape:
+
+    >>> from google.cloud.dataproc_ml.sql import ai_embed
+    >>>
+    >>> embedded = df.withColumn("vector", ai_embed(col("body")).result)
+
+Semantic search takes the whole DataFrame rather than a column, and returns
+the matching rows:
+
+    >>> from google.cloud.dataproc_ml.sql import ai_search
+    >>>
+    >>> hits = ai_search(df, "body", "I cannot sign in", top_k=3)
+    >>> hits.select("base.id", "base.body", "distance")
+"""
+
+from ._ai_embed import ai_embed
+from ._ai_embed import ai_embed_udf
+from ._ai_generate import ai_generate
+from ._ai_generate import ai_generate_udf
+from ._ai_search import ai_search
+
+__all__ = (
+    "ai_embed",
+    "ai_embed_udf",
+    "ai_generate",
+    "ai_generate_udf",
+    "ai_search",
+)