diff --git a/README.md b/README.md index 2a5cab8..afb8d10 100644 --- a/README.md +++ b/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` 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 `_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` 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 diff --git a/docs/api.rst b/docs/api.rst index 2155aa0..bdfbd43 100644 --- a/docs/api.rst +++ b/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: \ No newline at end of file diff --git a/google/cloud/dataproc_ml/sql/__init__.py b/google/cloud/dataproc_ml/sql/__init__.py new file mode 100644 index 0000000..59d4412 --- /dev/null +++ b/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", +) diff --git a/google/cloud/dataproc_ml/sql/_ai_embed.py b/google/cloud/dataproc_ml/sql/_ai_embed.py new file mode 100644 index 0000000..095fd57 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_embed.py @@ -0,0 +1,463 @@ +# 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. + +"""``AI.EMBED``: dense text embeddings over Spark columns.""" + +import asyncio +import logging +import math +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple + +import pandas as pd +from pyspark.sql import Column +from pyspark.sql.pandas.functions import pandas_udf +from pyspark.sql.types import ( + ArrayType, + DoubleType, + StringType, + StructField, + StructType, +) + +from google.cloud.dataproc_ml.sql import _client +from google.cloud.dataproc_ml.sql import _columns +from google.cloud.dataproc_ml.sql import _model_params +from google.cloud.dataproc_ml.sql import _retry + +logger = logging.getLogger(__name__) + +#: The model used when no endpoint is given. +DEFAULT_ENDPOINT = "text-embedding-005" + +# Fixed request settings, matching ai_generate. BigQuery does not expose these +# either; they are constants so that the function behaves sensibly on a large +# cluster without asking the caller to tune anything. +_MAX_CONCURRENT_REQUESTS = 5 +_MAX_ATTEMPTS = 5 +_BASE_RETRY_DELAY_SECONDS = 1.0 +_MAX_RETRY_DELAY_SECONDS = 60.0 + +# The endpoint rejects a request carrying more instances than this with +# "250 instance(s) is allowed per prediction". A Spark Arrow batch holds +# 10,000 rows by default, so splitting is required rather than optional. +_MAX_INSTANCES_PER_REQUEST = 250 + +#: Field names of :data:`RETURN_TYPE`. Not private to the module because +#: ``ai_search`` reads the embeddings back out of a column of this shape. +RESULT_FIELD = "result" +STATUS_FIELD = "status" + +# The endpoint rejects a request whose only instance is the empty string, and, +# worse, silently drops an empty instance from a request that has other +# instances too, returning fewer embeddings than were asked for. Empty values +# are therefore answered here and never sent. The wording follows the shape +# used by _retry for a request error. +EMPTY_CONTENT_STATUS = ( + "INVALID_ARGUMENT: The content to embed is empty. The model requires at " + "least one character." +) + +ContentArg = _columns.ColumnArg + +#: The struct returned for every row, mirroring BigQuery's ``AI.EMBED``. Unlike +#: ``AI.GENERATE`` there is no ``full_response`` field: an embedding response +#: carries nothing beyond the vector that a caller could use. +RETURN_TYPE = StructType( + [ + StructField(RESULT_FIELD, ArrayType(DoubleType()), nullable=True), + StructField(STATUS_FIELD, StringType(), nullable=True), + ] +) + + +def ai_embed( + content: ContentArg, + *, + endpoint: str = DEFAULT_ENDPOINT, + task_type: Optional[str] = None, + title: Optional[str] = None, + model_params: Optional[Dict[str, Any]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _embed_content=None, +) -> Column: + """Generates a dense embedding vector for every row. + + This is the Spark counterpart of BigQuery's ``AI.EMBED`` function and + returns the same struct, so queries can be moved between the two engines + with little change. It takes the same arguments, except that the Vertex AI + call is authenticated with Application Default Credentials rather than a + BigQuery connection, so ``project`` and ``location`` take the place of + ``connection_id``. + + The content follows the usual PySpark convention: a string names a column, + and literal text must be wrapped in :func:`pyspark.sql.functions.lit`. A + sequence of parts is concatenated in order. + + Unlike :func:`~google.cloud.dataproc_ml.sql.ai_generate`, which calls the + model once per row, the rows of an Arrow batch are embedded together in as + few requests as the endpoint allows. Requests are retried with exponential + backoff and are issued a few at a time per worker. + + .. note:: + Using this function incurs Vertex AI charges. See the + `Vertex AI Generative AI pricing + `_ page. + + Args: + content: The column to embed, the name of that column, or a sequence + of parts to concatenate. + endpoint: The embedding model to call. Defaults to + ``text-embedding-005``. + task_type: The downstream use of the embeddings, which lets the model + tailor them, for example ``'RETRIEVAL_DOCUMENT'`` or + ``'SEMANTIC_SIMILARITY'``. Embeddings meant to be compared with + each other must be generated with the same value. + title: The title of the document being embedded, which improves + quality. Only valid together with + ``task_type='RETRIEVAL_DOCUMENT'``. + model_params: Additional request settings, following the Vertex AI + ``embedContent`` request body without its ``contents`` field. The + useful one is ``{"outputDimensionality": 256}``, which shortens + the returned vector. + project: The Google Cloud project billed for the calls. Inferred from + the environment when not set. + location: The Vertex AI location. Defaults to ``global``. + + Returns: + A column holding a struct with the fields ``result``, an + ``ARRAY`` holding the embedding, and ``status``, which is + empty when the row succeeded and holds the error message otherwise. + + Example: + >>> df.withColumn("e", ai_embed(sf.col("body"))).select( + ... "e.result", "e.status" + ... ) + >>> df.withColumn( + ... "e", + ... ai_embed( + ... "body", + ... task_type="RETRIEVAL_DOCUMENT", + ... model_params={"outputDimensionality": 256}, + ... ), + ... ) + """ + embed = ai_embed_udf( + endpoint=endpoint, + task_type=task_type, + title=title, + model_params=model_params, + project=project, + location=location, + _embed_content=_embed_content, + ) + return embed(_columns.to_string_column(content, argument_name="content")) + + +def ai_embed_udf( + *, + endpoint: str = DEFAULT_ENDPOINT, + task_type: Optional[str] = None, + title: Optional[str] = None, + model_params: Optional[Dict[str, Any]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _embed_content=None, +): + """Builds a reusable ``AI.EMBED`` user defined function. + + Use this when the same configuration is applied in several places, or to + make the function available to Spark SQL: + + >>> spark.udf.register("ai_embed", ai_embed_udf()) + >>> spark.sql("SELECT ai_embed(body).result FROM docs") + + Args: + endpoint: See :func:`ai_embed`. + task_type: See :func:`ai_embed`. + title: See :func:`ai_embed`. + model_params: See :func:`ai_embed`. + project: See :func:`ai_embed`. + location: See :func:`ai_embed`. + _embed_content: Internal. Replaces the call to Vertex AI with an + awaitable of the same shape, so that tests can run the whole Spark + path without credentials or network access. + + Returns: + A pandas user defined function taking a single content column. + + Raises: + ValueError: If any argument is invalid. All validation happens here, on + the driver, so that mistakes surface before a stage is scheduled. + """ + if not isinstance(endpoint, str) or not endpoint.strip(): + raise ValueError("endpoint must be a non-empty model name.") + + flat_params = _model_params.normalize_embed_params( + model_params, task_type, title + ) + + def _predict( + batch_iterator: Iterator[pd.Series], + ) -> Iterator[pd.DataFrame]: + """Runs on the executor, once per task.""" + # Resolved once per task and shared with every other task running in + # this worker process. + embed_content = _embed_content or _client.get_embed_content( + project, location + ) + config = _model_params.build_embed_config(flat_params) + + for batch in batch_iterator: + rows = _client.run_coroutine( + _embed_batch( + embed_content=embed_content, + contents=batch.tolist(), + endpoint=endpoint, + config=config, + ) + ) + yield pd.DataFrame( + rows, + columns=[field.name for field in RETURN_TYPE.fields], + index=batch.index, + ) + + return pandas_udf(_predict, returnType=RETURN_TYPE) + + +def embed_one( + *, + embed_content, + content: str, + endpoint: str, + config, +) -> List[float]: + """Embeds a single value on the driver, raising if it cannot be embedded. + + The column functions report a failed row in its ``status`` field so that + one bad row never fails a stage. A single value embedded on the driver has + no such row to report against, and a caller that asked for one embedding + and got none cannot carry on, so this raises instead. + + Retries and error wording are shared with the column function. + + Args: + embed_content: The awaitable Vertex AI call. + content: The text to embed. + endpoint: The embedding model to call. + config: The embedding config. + + Returns: + The embedding vector. + + Raises: + ValueError: If the content is empty or the model reported an error. + """ + rows = _client.run_coroutine( + _embed_batch( + embed_content=embed_content, + contents=[content], + endpoint=endpoint, + config=config, + ) + ) + row = rows[0] + if row[STATUS_FIELD]: + raise ValueError(f"Could not embed {content!r}: {row[STATUS_FIELD]}") + vector = row[RESULT_FIELD] + if not vector: + raise ValueError( + f"The embedding model {endpoint} returned no vector for " + f"{content!r}." + ) + return vector + + +async def _embed_batch( + *, + embed_content, + contents: List[Any], + endpoint: str, + config, +) -> List[Dict[str, Any]]: + """Embeds every row of an Arrow batch, in as few requests as allowed.""" + rows: List[Optional[Dict[str, Any]]] = [None] * len(contents) + sendable: List[Tuple[int, str]] = [] + + for index, content in enumerate(contents): + if _is_missing(content): + # A null produces a null embedding without calling the model. + rows[index] = _row(result=None, status="") + continue + text = _to_text(content) + if not text: + rows[index] = _row(result=None, status=EMPTY_CONTENT_STATUS) + continue + sendable.append((index, text)) + + semaphore = asyncio.Semaphore(_MAX_CONCURRENT_REQUESTS) + + async def _one(chunk: List[Tuple[int, str]]) -> None: + async with semaphore: + await _embed_chunk( + embed_content=embed_content, + chunk=chunk, + endpoint=endpoint, + config=config, + rows=rows, + ) + + await asyncio.gather(*[_one(chunk) for chunk in _chunks(sendable)]) + + # Every row is filled in by one of the branches above. Dropping or + # reordering here would shift every later row onto the wrong embedding, so + # the invariant is checked rather than assumed. + missing = [index for index, row in enumerate(rows) if row is None] + if missing: + raise RuntimeError( + f"No embedding was produced for rows {missing}. Refusing to " + "return a batch whose rows may be misaligned." + ) + return rows + + +def _chunks( + sendable: Sequence[Tuple[int, str]], +) -> List[List[Tuple[int, str]]]: + """Splits the rows to send into requests the endpoint will accept.""" + return [ + list(sendable[start : start + _MAX_INSTANCES_PER_REQUEST]) + for start in range(0, len(sendable), _MAX_INSTANCES_PER_REQUEST) + ] + + +async def _embed_chunk( + *, + embed_content, + chunk: List[Tuple[int, str]], + endpoint: str, + config, + rows: List[Optional[Dict[str, Any]]], +) -> None: + """Embeds one request worth of rows, retrying transient failures. + + A request either succeeds for all of its rows or fails for all of them, so + a terminal failure is recorded against every row it carried. That keeps the + per-row reporting of BigQuery's ``AI.EMBED`` while still sending the rows + together. + """ + texts = [text for _, text in chunk] + + for attempt in range(_MAX_ATTEMPTS): + try: + response = await embed_content( + model=endpoint, contents=texts, config=config + ) + except Exception as e: # pylint: disable=broad-except + # Anything that is not a backend failure, such as a bug in this + # library, is left to fail the task instead of being hidden in the + # status field. + if not _retry.is_reportable(e): + raise + is_last_attempt = attempt == _MAX_ATTEMPTS - 1 + if _retry.is_retryable(e) and not is_last_attempt: + delay = _retry.full_jitter_delay( + attempt, + _BASE_RETRY_DELAY_SECONDS, + _MAX_RETRY_DELAY_SECONDS, + ) + logger.warning( + "Retrying request to %s in %.2fs after error: %s", + endpoint, + delay, + e, + ) + await asyncio.sleep(delay) + continue + status = _retry.to_status_message(e) + for index, _ in chunk: + rows[index] = _row(result=None, status=status) + return + + _record(response, chunk, rows) + return + + # Unreachable: the loop either records a result or retries. + raise AssertionError("The retry loop ended without producing a result.") + + +def _record( + response, + chunk: List[Tuple[int, str]], + rows: List[Optional[Dict[str, Any]]], +) -> None: + """Writes the embeddings of one response back to their rows. + + Raises: + RuntimeError: If the response does not carry exactly one embedding per + row that was sent. + """ + embeddings = getattr(response, "embeddings", None) or [] + if len(embeddings) != len(chunk): + # The rows are matched to their embeddings by position, so a response + # of a different length would silently attach the wrong vector to + # every row after the first discrepancy. Corrupt embeddings are far + # worse than a failed query, so this is raised rather than reported in + # status. It is not expected to happen: the one known cause, an empty + # instance being dropped, is handled before the request is sent. + raise RuntimeError( + f"The embedding model returned {len(embeddings)} embeddings for " + f"{len(chunk)} inputs. Refusing to guess which embedding belongs " + "to which row." + ) + + for (index, _), embedding in zip(chunk, embeddings): + values = getattr(embedding, "values", None) + rows[index] = _row( + result=None if values is None else [float(v) for v in values], + status="", + ) + + +def _row(*, result: Optional[List[float]], status: str) -> Dict[str, Any]: + """Builds one output row.""" + return {RESULT_FIELD: result, STATUS_FIELD: status} + + +def _is_missing(content: Any) -> bool: + """Reports whether a value is null. + + Spark nulls arrive as ``None``, as NaN in a numeric batch, or as + ``pandas.NA`` from a nullable column. Every other value, including a + number or a timestamp, is real content that the caller renders as text. + """ + if content is None or content is pd.NA: + return True + return isinstance(content, float) and math.isnan(content) + + +def _to_text(content: Any) -> str: + """Renders content that Spark has not already cast to a string. + + A whole number is rendered without a fraction. Arrow hands a nullable + integer column to pandas as floats, so a BIGINT 2 arrives here as 2.0 as + soon as any row in the batch is null. Rendering that as "2.0" would let + one row's content depend on whether an unrelated row happened to be null. + """ + if isinstance(content, str): + return content + if isinstance(content, float) and float(content).is_integer(): + return str(int(content)) + return str(content) diff --git a/google/cloud/dataproc_ml/sql/_ai_generate.py b/google/cloud/dataproc_ml/sql/_ai_generate.py new file mode 100644 index 0000000..ef80a50 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -0,0 +1,551 @@ +# 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. + +"""``AI.GENERATE``: generative text and structured output over Spark columns.""" + +import asyncio +import json +import logging +import math +from typing import Any, Dict, Iterator, List, Optional, Union + +import pandas as pd +from pyspark.sql import Column +from pyspark.sql.pandas.functions import pandas_udf +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + StringType, + StructField, + StructType, +) + +from google.cloud.dataproc_ml.sql import _client +from google.cloud.dataproc_ml.sql import _columns +from google.cloud.dataproc_ml.sql import _model_params +from google.cloud.dataproc_ml.sql import _output_schema as _output_schema_util +from google.cloud.dataproc_ml.sql import _retry + +try: + # Spark 4 models semi-structured data as VARIANT, which is the closest + # equivalent of the JSON type that BigQuery's AI functions return. + from pyspark.sql.types import VariantType, VariantVal # pylint: disable=ungrouped-imports + + _VARIANT_SUPPORTED = True +except ImportError: # pragma: no cover - depends on the Spark version + _VARIANT_SUPPORTED = False + + +logger = logging.getLogger(__name__) + +#: The model used when no endpoint is given. +DEFAULT_ENDPOINT = "gemini-3.6-flash" + +# Fixed request settings. BigQuery does not expose these either; they are +# constants here so that the function behaves sensibly on a large cluster +# without asking the caller to tune anything. +_MAX_CONCURRENT_REQUESTS = 5 +_MAX_ATTEMPTS = 5 +_BASE_RETRY_DELAY_SECONDS = 1.0 +_MAX_RETRY_DELAY_SECONDS = 60.0 + +_RESULT_FIELD = "result" +_FULL_RESPONSE_FIELD = "full_response" +_STATUS_FIELD = "status" + +# Spellings accepted when a BOOL field comes back as text rather than as a +# JSON boolean. Anything else is treated as an answer that does not fit the +# schema and becomes null. +_TRUE_WORDS = frozenset({"true", "t", "yes", "y", "1"}) +_FALSE_WORDS = frozenset({"false", "f", "no", "n", "0"}) + +PromptArg = _columns.ColumnArg + + +def ai_generate( + prompt: PromptArg, + *, + endpoint: str = DEFAULT_ENDPOINT, + model_params: Optional[Dict[str, Any]] = None, + output_schema: Optional[Union[str, StructType]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _generate_content=None, +) -> Column: + """Generates text for every row using a Gemini model. + + This is the Spark counterpart of BigQuery's ``AI.GENERATE`` function and + returns the same struct, so queries can be moved between the two engines + with little change. It takes the same arguments, except that the Vertex AI + call is authenticated with Application Default Credentials rather than a + BigQuery connection, so ``project`` and ``location`` take the place of + ``connection_id``. + + The prompt follows the usual PySpark convention: a string names a column, + and literal text must be wrapped in :func:`pyspark.sql.functions.lit`. A + sequence of parts is concatenated in order, like BigQuery's struct prompt. + + Requests are retried with exponential backoff and are issued a few at a + time per worker. Like BigQuery, this function does not expose those + settings. + + .. note:: + Using this function incurs Vertex AI charges. See the + `Vertex AI Generative AI pricing + `_ page. + + Args: + prompt: The prompt column, the name of a prompt column, or a sequence + of parts to concatenate. + endpoint: The Gemini model to call. Defaults to ``gemini-3.6-flash``. + model_params: Additional request settings, following the Vertex AI + ``generateContent`` request body without its ``contents`` field, + for example ``{"generation_config": {"temperature": 0.2}}``. + output_schema: An optional description of the fields the model should + return, for example ``"sentiment STRING, score INT64"``. When it + is set, the ``result`` field is replaced by one field per declared + field, sorted by name as BigQuery does. + project: The Google Cloud project billed for the calls. Inferred from + the environment when not set. + location: The Vertex AI location. Defaults to ``global``, which is + where the current Gemini models are served. Set it only when a + particular region is required, and check that the endpoint is + available there. + + Returns: + A column holding a struct. Without ``output_schema`` the fields are + ``result``, ``full_response`` and ``status``; with it, the declared + fields replace ``result``. ``full_response`` holds the complete model + response, typed as VARIANT on Spark 4 and as a JSON string on older + versions. ``status`` is empty when the row succeeded and holds the + error message otherwise. + + Example: + >>> df.withColumn("g", ai_generate(sf.col("review"))).select( + ... "g.result", "g.status" + ... ) + >>> df.withColumn( + ... "g", + ... ai_generate( + ... [sf.lit("Summarize in 10 words: "), sf.col("body")], + ... output_schema="summary STRING, sentiment STRING", + ... ), + ... ) + """ + generate = ai_generate_udf( + endpoint=endpoint, + model_params=model_params, + output_schema=output_schema, + project=project, + location=location, + _generate_content=_generate_content, + ) + return generate(_to_prompt_column(prompt)) + + +def ai_generate_udf( + *, + endpoint: str = DEFAULT_ENDPOINT, + model_params: Optional[Dict[str, Any]] = None, + output_schema: Optional[Union[str, StructType]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _generate_content=None, +): + """Builds a reusable ``AI.GENERATE`` user defined function. + + Use this when the same configuration is applied in several places, or to + make the function available to Spark SQL: + + >>> spark.udf.register("ai_generate", ai_generate_udf()) + >>> spark.sql("SELECT ai_generate(review).result FROM reviews") + + Args: + endpoint: See :func:`ai_generate`. + model_params: See :func:`ai_generate`. + output_schema: See :func:`ai_generate`. + project: See :func:`ai_generate`. + location: See :func:`ai_generate`. + _generate_content: Internal. Replaces the call to Vertex AI with an + awaitable of the same shape, so that tests can run the whole Spark + path without credentials or network access. + + Returns: + A pandas user defined function taking a single prompt column. + + Raises: + ValueError: If any argument is invalid. All validation happens here, on + the driver, so that mistakes surface before a stage is scheduled. + """ + if not isinstance(endpoint, str) or not endpoint.strip(): + raise ValueError("endpoint must be a non-empty model name.") + + parsed_schema = ( + _output_schema_util.parse_output_schema(output_schema) + if output_schema is not None + else None + ) + response_schema = ( + _output_schema_util.to_response_schema(parsed_schema) + if parsed_schema is not None + else None + ) + flat_params = _model_params.normalize_model_params( + model_params, response_schema + ) + return_type = _build_return_type(parsed_schema) + + def _predict( + batch_iterator: Iterator[pd.Series], + ) -> Iterator[pd.DataFrame]: + """Runs on the executor, once per task.""" + # Resolved once per task and shared with every other task running in + # this worker process. + generate_content = _generate_content or _client.get_generate_content( + project, location + ) + config = _model_params.build_generation_config(flat_params) + + for batch in batch_iterator: + rows = _client.run_coroutine( + _generate_batch( + generate_content=generate_content, + prompts=batch.tolist(), + endpoint=endpoint, + config=config, + schema=parsed_schema, + ) + ) + yield pd.DataFrame( + rows, + columns=[field.name for field in return_type.fields], + index=batch.index, + ) + + return pandas_udf(_predict, returnType=return_type) + + +def _to_prompt_column(prompt: PromptArg) -> Column: + """Builds the single string column that holds each row's prompt. + + Every path casts to STRING. A prompt column of another type would + otherwise reach the worker as a non-string value and be treated as a null + prompt, turning a whole column into nulls without any error. + """ + return _columns.to_string_column(prompt, argument_name="prompt") + + +def _build_return_type(parsed_schema: Optional[StructType]) -> StructType: + """Builds the struct returned by the function.""" + if parsed_schema is None: + fields = [StructField(_RESULT_FIELD, StringType())] + else: + fields = list(parsed_schema.fields) + for reserved in (_FULL_RESPONSE_FIELD, _STATUS_FIELD): + if any(field.name == reserved for field in fields): + raise ValueError( + f"output_schema must not declare a '{reserved}' field; it " + "is always added to the result." + ) + # Every field is nullable: a row that fails reports the failure in status + # and leaves the generated fields empty. + fields = [ + StructField(field.name, field.dataType, nullable=True) + for field in fields + ] + fields.append(StructField(_FULL_RESPONSE_FIELD, _full_response_data_type())) + fields.append(StructField(_STATUS_FIELD, StringType())) + return StructType(fields) + + +def _full_response_data_type() -> DataType: + """Returns the type of the ``full_response`` field. + + BigQuery returns the raw model response as JSON. Spark 4 has an equivalent + semi-structured type, VARIANT, which allows ``variant_get`` and + ``to_json`` to be used on the column. On older Spark versions, which have + no such type, the response is returned as a JSON string instead. + """ + return VariantType() if _VARIANT_SUPPORTED else StringType() + + +async def _generate_batch( + *, + generate_content, + prompts: List[Any], + endpoint: str, + config, + schema: Optional[StructType], +) -> List[Dict[str, Any]]: + """Generates a response for every prompt in an Arrow batch.""" + semaphore = asyncio.Semaphore(_MAX_CONCURRENT_REQUESTS) + + async def _one(prompt: Any) -> Dict[str, Any]: + async with semaphore: + return await _generate_one( + generate_content=generate_content, + prompt=prompt, + endpoint=endpoint, + config=config, + schema=schema, + ) + + return await asyncio.gather(*[_one(prompt) for prompt in prompts]) + + +async def _generate_one( + *, + generate_content, + prompt: Any, + endpoint: str, + config, + schema: Optional[StructType], +) -> Dict[str, Any]: + """Generates the response for one row, retrying transient failures.""" + if _is_missing(prompt): + # A null prompt produces a null result without calling the model. + return _row(schema, result=None, full_response=None, status="") + + # ``ai_generate`` casts its prompt column to STRING on the driver, but the + # function returned by ``ai_generate_udf`` can also be registered and + # called straight from Spark SQL, where no such cast is inserted. Render + # the value here so that a non-string column, say an id, still produces a + # prompt rather than being dropped. + text = _to_text(prompt) + + for attempt in range(_MAX_ATTEMPTS): + try: + response = await generate_content( + model=endpoint, contents=text, config=config + ) + return _to_row(response, schema) + except Exception as e: # pylint: disable=broad-except + # Anything that is not a backend failure, such as a bug in this + # library, is left to fail the task instead of being hidden in the + # status field. + if not _retry.is_reportable(e): + raise + is_last_attempt = attempt == _MAX_ATTEMPTS - 1 + if _retry.is_retryable(e) and not is_last_attempt: + delay = _retry.full_jitter_delay( + attempt, + _BASE_RETRY_DELAY_SECONDS, + _MAX_RETRY_DELAY_SECONDS, + ) + logger.warning( + "Retrying request to %s in %.2fs after error: %s", + endpoint, + delay, + e, + ) + await asyncio.sleep(delay) + continue + return _row( + schema, + result=None, + full_response=None, + status=_retry.to_status_message(e), + ) + + # Unreachable: the loop either returns a row or retries. + raise AssertionError("The retry loop ended without producing a row.") + + +def _is_missing(prompt: Any) -> bool: + """Reports whether a prompt is null. + + Spark nulls arrive as ``None``, as NaN in a numeric batch, or as + ``pandas.NA`` from a nullable column. Every other value, including a + number or a timestamp, is a real prompt that the caller renders as text. + """ + if prompt is None or prompt is pd.NA: + return True + return isinstance(prompt, float) and math.isnan(prompt) + + +def _to_text(prompt: Any) -> str: + """Renders a prompt that Spark has not already cast to a string. + + A whole number is rendered without a fraction. Arrow hands a nullable + integer column to pandas as floats, so a BIGINT 2 arrives here as 2.0 as + soon as any row in the batch is null. Rendering that as "2.0" would let + one row's prompt depend on whether an unrelated row happened to be null. + """ + if isinstance(prompt, str): + return prompt + if isinstance(prompt, float) and float(prompt).is_integer(): + return str(int(prompt)) + return str(prompt) + + +def _to_row(response, schema: Optional[StructType]) -> Dict[str, Any]: + """Converts a model response into one output row.""" + full_response = _dump_response(response) + text = _extract_text(response) + + if schema is None: + return _row(schema, result=text, full_response=full_response, status="") + + if text is None: + # The call succeeded but produced no usable text, for example because + # the response was filtered. This is not an error. + return _row(schema, result=None, full_response=full_response, status="") + + try: + payload = json.loads(text) + except (TypeError, ValueError): + return _row( + schema, + result=None, + full_response=full_response, + status=_retry.PARSE_FAILURE_STATUS, + ) + + if not isinstance(payload, dict): + return _row( + schema, + result=None, + full_response=full_response, + status=_retry.PARSE_FAILURE_STATUS, + ) + + row = { + field.name: _coerce(payload.get(field.name), field.dataType) + for field in schema.fields + } + row[_FULL_RESPONSE_FIELD] = full_response + row[_STATUS_FIELD] = "" + return row + + +def _row( + schema: Optional[StructType], + *, + result: Optional[str], + full_response: Any, + status: str, +) -> Dict[str, Any]: + """Builds an output row, leaving the generated fields empty if needed.""" + if schema is None: + row: Dict[str, Any] = {_RESULT_FIELD: result} + else: + row = {field.name: None for field in schema.fields} + row[_FULL_RESPONSE_FIELD] = full_response + row[_STATUS_FIELD] = status + return row + + +def _extract_text(response) -> Optional[str]: + """Returns the generated text, following BigQuery's extraction rules. + + Only the first candidate is considered, reasoning parts are skipped, and + the first remaining text part wins. + """ + candidates = getattr(response, "candidates", None) or [] + if not candidates: + return None + content = getattr(candidates[0], "content", None) + for part in getattr(content, "parts", None) or []: + if getattr(part, "thought", None): + continue + text = getattr(part, "text", None) + if isinstance(text, str): + return text + return None + + +def _dump_response(response) -> Any: + """Serializes the whole response, as BigQuery's ``full_response`` does. + + Args: + response: The model response. + + Returns: + A Spark variant value, or a JSON string on Spark versions without the + variant type. + """ + try: + serialized = response.model_dump_json(exclude_none=True) + except Exception: # pylint: disable=broad-except + logger.debug("Could not serialize the model response.", exc_info=True) + serialized = "{}" + return _to_full_response(serialized) + + +def _to_full_response(serialized: str) -> Any: + """Converts serialized JSON into the value stored in ``full_response``.""" + if not _VARIANT_SUPPORTED: + return serialized + try: + return VariantVal.parseJson(serialized) + except Exception: # pylint: disable=broad-except + logger.debug("Could not encode the response as a variant.") + return VariantVal.parseJson("{}") + + +def _coerce(value: Any, data_type: DataType) -> Any: + """Converts a JSON value to the type declared in the output schema.""" + if value is None: + return None + try: + if isinstance(data_type, StringType): + return value if isinstance(value, str) else json.dumps(value) + if isinstance(data_type, BooleanType): + return _to_boolean(value) + if isinstance(data_type, (LongType, IntegerType)): + return int(value) + if isinstance(data_type, (DoubleType, FloatType)): + return float(value) + if isinstance(data_type, ArrayType): + if not isinstance(value, list): + return None + return [_coerce(item, data_type.elementType) for item in value] + if isinstance(data_type, StructType): + if not isinstance(value, dict): + return None + return { + field.name: _coerce(value.get(field.name), field.dataType) + for field in data_type.fields + } + except (TypeError, ValueError): + return None + return None + + +def _to_boolean(value: Any) -> Optional[bool]: + """Converts a JSON value to a boolean. + + ``bool()`` on its own is wrong here, because every non-empty string is + truthy: a model answering with the text "false" would be recorded as True. + A value that does not clearly denote a boolean becomes null, which is how + the rest of this module reports an answer that does not fit the requested + schema. + """ + if isinstance(value, bool): + return value + if isinstance(value, str): + text = value.strip().lower() + if text in _TRUE_WORDS: + return True + if text in _FALSE_WORDS: + return False + return None + if isinstance(value, (int, float)): + return bool(value) + return None diff --git a/google/cloud/dataproc_ml/sql/_ai_search.py b/google/cloud/dataproc_ml/sql/_ai_search.py new file mode 100644 index 0000000..c8b3bcd --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_search.py @@ -0,0 +1,428 @@ +# 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. + +"""``AI.SEARCH``: semantic search over a DataFrame.""" + +import logging +from typing import Any, Dict, List, Optional, Union + +from pyspark.sql import Column +from pyspark.sql import DataFrame +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf +from pyspark.sql.types import ArrayType +from pyspark.sql.types import DoubleType +from pyspark.sql.types import StructType + +from google.cloud.dataproc_ml.sql import _ai_embed +from google.cloud.dataproc_ml.sql import _client +from google.cloud.dataproc_ml.sql import _model_params + +logger = logging.getLogger(__name__) + +#: The number of results returned when ``top_k`` is not given, matching +#: BigQuery. +DEFAULT_TOP_K = 10 + +#: The distance used when ``distance_type`` is not given, matching BigQuery. +DEFAULT_DISTANCE_TYPE = "EUCLIDEAN" + +#: The distances BigQuery's ``AI.SEARCH`` accepts. +DISTANCE_TYPES = ("EUCLIDEAN", "COSINE", "DOT_PRODUCT") + +#: The columns of the returned DataFrame, matching BigQuery's ``AI.SEARCH``. +BASE_COLUMN = "base" +DISTANCE_COLUMN = "distance" + +#: BigQuery searches a table whose embeddings are maintained for it in a +#: generated column conventionally named after the column it embeds. The same +#: name is looked for here, so a DataFrame that carries its embeddings needs no +#: extra argument. +EMBEDDING_COLUMN_SUFFIX = "_embedding" + +# Holds each row's vector between the two projections below. Prefixed so that +# it cannot collide with a column of the table being searched. +_VECTOR_COLUMN = "__dataproc_ml_ai_search_vector" + +_VECTOR_TYPE = ArrayType(DoubleType()) + +BaseTable = Union[DataFrame, str] + + +def ai_search( + base_table: BaseTable, + column_to_search: str, + query: str, + *, + top_k: int = DEFAULT_TOP_K, + distance_type: str = DEFAULT_DISTANCE_TYPE, + embedding_column: Optional[str] = None, + endpoint: str = _ai_embed.DEFAULT_ENDPOINT, + task_type: Optional[str] = None, + model_params: Optional[Dict[str, Any]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _embed_content=None, +) -> DataFrame: + """Returns the rows closest in meaning to a search term. + + This is the Spark counterpart of BigQuery's ``AI.SEARCH`` function and + returns the same two columns, ``base`` and ``distance``, so a query can be + moved between the two engines with little change. + + BigQuery searches a table whose embeddings are maintained for it, in a + generated column. Spark has no such column, so the embeddings are found in + one of three ways, in order: + + 1. the column named by ``embedding_column``; + 2. a column named ``_embedding``, which is the name + BigQuery's own examples give the generated column; + 3. failing both, they are generated here by calling + :func:`~google.cloud.dataproc_ml.sql.ai_embed` on + ``column_to_search``. + + An existing embedding column may hold either a plain ``ARRAY`` 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. + + The search term itself is embedded once, on the driver. Unlike the column + functions, which report a failed row in its ``status`` field, a search term + that cannot be embedded raises: there is no result to return without it. + This also matches BigQuery. Rows whose own embedding is missing are skipped + rather than failing the search. + + .. note:: + Using this function incurs Vertex AI charges. See the + `Vertex AI Generative AI pricing + `_ page. + + Args: + base_table: The DataFrame to search, or the name of a table or view to + read it from. + column_to_search: The name of the column holding the text to search. + As in BigQuery this names the source column, not the embeddings. + query: The search term. + top_k: The number of rows to return, nearest first. Defaults to 10. + distance_type: How to measure the distance between two embeddings, one + of ``'EUCLIDEAN'``, ``'COSINE'`` or ``'DOT_PRODUCT'``. Defaults to + ``'EUCLIDEAN'``. + embedding_column: The name of a column holding the embeddings of + ``column_to_search``, when it is not named for the column it + embeds. + endpoint: The embedding model to call. Defaults to + ``text-embedding-005``. It must be the model the table's own + embeddings were generated with, since embeddings from two models + cannot be compared. + task_type: The ``task_type`` of the search term, and of the table's + embeddings when they are generated here. See + :func:`~google.cloud.dataproc_ml.sql.ai_embed`. + model_params: Additional request settings. See + :func:`~google.cloud.dataproc_ml.sql.ai_embed`. + project: The Google Cloud project billed for the calls. Inferred from + the environment when not set. + location: The Vertex AI location. Defaults to ``global``. + _embed_content: Internal. Replaces the call to Vertex AI with an + awaitable of the same shape, so that tests can run the whole Spark + path without credentials or network access. + + Returns: + A DataFrame of at most ``top_k`` rows ordered by ``distance``, holding + ``base``, a struct of every column of the table searched, and + ``distance``, the distance between the row and the search term. A + smaller distance always means a closer match, so ``DOT_PRODUCT`` + returns the negated inner product. + + Raises: + ValueError: If any argument is invalid, or if the search term cannot + be embedded. + + Example: + >>> hits = ai_search(products, "description", "A really fun toy") + >>> hits.select("base.name", "base.description", "distance").show() + + Searching a table that already carries its embeddings, which is both + cheaper and what BigQuery does: + + >>> embedded = products.withColumn( + ... "description_embedding", ai_embed("description") + ... ).cache() + >>> ai_search(embedded, "description", "A really fun toy", top_k=3) + """ + table = _resolve_base_table(base_table) + distance_type = _validate_distance_type(distance_type) + _validate_top_k(top_k) + _validate_query(query) + _validate_column_to_search(table, column_to_search) + + flat_params = _model_params.normalize_embed_params(model_params, task_type) + + query_vector = _ai_embed.embed_one( + embed_content=( + _embed_content or _client.get_embed_content(project, location) + ), + content=query, + endpoint=endpoint, + config=_model_params.build_embed_config(flat_params), + ) + + vector = _row_vector( + table, + column_to_search=column_to_search, + embedding_column=embedding_column, + endpoint=endpoint, + task_type=task_type, + model_params=model_params, + project=project, + location=location, + _embed_content=_embed_content, + ) + + base_columns = list(table.columns) + if _VECTOR_COLUMN in base_columns: + raise ValueError( + f"{_VECTOR_COLUMN} is reserved by ai_search and cannot be a " + "column of the table being searched." + ) + + # The vector is attached in its own projection and read back by name. When + # it is generated here that expression is a call to the model, and Spark + # will not fold a projection whose output is read more than once into the + # one that reads it, so the distance below costs one call per row rather + # than one per mention. + scored = table.withColumn(_VECTOR_COLUMN, vector).select( + sf.struct(*[sf.col(name) for name in base_columns]).alias(BASE_COLUMN), + _distance(sf.col(_VECTOR_COLUMN), query_vector, distance_type).alias( + DISTANCE_COLUMN + ), + ) + + # A row with no embedding, or one of a different width, scores null and is + # dropped. BigQuery likewise skips rows whose embedding is missing. + # + # Those rows are sorted to the end and dropped after the limit rather than + # before it, which returns the same rows: the limit keeps the nearest + # matches, and only reaches a null once it has run out of real ones. The + # order matters because a filter placed before the limit is pushed down + # past the projection above, which would evaluate the embeddings a second + # time for every row. Nothing is pushed past a limit. + return ( + scored.orderBy(sf.col(DISTANCE_COLUMN).asc_nulls_last()) + .limit(top_k) + .where(sf.col(DISTANCE_COLUMN).isNotNull()) + ) + + +def _resolve_base_table(base_table: BaseTable) -> DataFrame: + """Returns the DataFrame to search. + + Raises: + ValueError: If the argument is neither a DataFrame nor a table name, + or if no Spark session is available to resolve a name. + """ + if isinstance(base_table, DataFrame): + return base_table + if isinstance(base_table, str): + if not base_table.strip(): + raise ValueError("base_table must be a non-empty table name.") + spark = SparkSession.getActiveSession() + if spark is None: + raise ValueError( + f"Cannot read the table {base_table!r} because there is no " + "active Spark session. Pass a DataFrame instead." + ) + return spark.table(base_table) + raise ValueError( + "base_table must be a DataFrame or the name of a table, got " + f"{type(base_table).__name__}." + ) + + +def _validate_distance_type(distance_type: str) -> str: + """Returns the distance type, normalized to upper case. + + Raises: + ValueError: If the distance type is not one of :data:`DISTANCE_TYPES`. + """ + if ( + not isinstance(distance_type, str) + or distance_type.upper() not in DISTANCE_TYPES + ): + supported = ", ".join(DISTANCE_TYPES) + raise ValueError( + f"Unsupported distance_type: {distance_type!r}. Supported values " + f"are {supported}." + ) + return distance_type.upper() + + +def _validate_top_k(top_k: int) -> None: + """Checks ``top_k``. + + Raises: + ValueError: If it is not a positive whole number. + """ + # A bool is an int in Python, and top_k=True is a mistake, not a request + # for one row. + if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k < 1: + raise ValueError(f"top_k must be a positive integer, got {top_k!r}.") + + +def _validate_query(query: str) -> None: + """Checks the search term. + + Raises: + ValueError: If it is not a non-empty string. + """ + if not isinstance(query, str) or not query.strip(): + raise ValueError( + f"query must be a non-empty search term, got {query!r}." + ) + + +def _validate_column_to_search(table: DataFrame, name: str) -> None: + """Checks that the column to search exists. + + Raises: + ValueError: If the name is empty or names no column of the table. + """ + if not isinstance(name, str) or not name.strip(): + raise ValueError( + "column_to_search must be the name of a column, got " f"{name!r}." + ) + if name not in table.columns: + available = ", ".join(table.columns) + raise ValueError( + f"column_to_search {name!r} is not a column of the table being " + f"searched. Its columns are {available}." + ) + + +def _row_vector( + table: DataFrame, + *, + column_to_search: str, + embedding_column: Optional[str], + endpoint: str, + task_type: Optional[str], + model_params: Optional[Dict[str, Any]], + project: Optional[str], + location: Optional[str], + _embed_content, +) -> Column: + """Returns the column holding each row's embedding. + + Raises: + ValueError: If a named embedding column is missing or holds neither an + array nor the struct returned by ``ai_embed``. + """ + if embedding_column is not None: + if embedding_column not in table.columns: + available = ", ".join(table.columns) + raise ValueError( + f"embedding_column {embedding_column!r} is not a column of " + f"the table being searched. Its columns are {available}." + ) + return _existing_vector(table, embedding_column) + + conventional = f"{column_to_search}{EMBEDDING_COLUMN_SUFFIX}" + if conventional in table.columns: + logger.info( + "Searching the embeddings in %s. Pass embedding_column to use a " + "different column.", + conventional, + ) + return _existing_vector(table, conventional) + + logger.info( + "The table carries no embeddings of %s, so they will be generated " + "for every row of this search. Add a %s column to avoid paying for " + "them again on the next search.", + column_to_search, + conventional, + ) + return _ai_embed.ai_embed( + column_to_search, + endpoint=endpoint, + task_type=task_type, + model_params=model_params, + project=project, + location=location, + _embed_content=_embed_content, + ).getField(_ai_embed.RESULT_FIELD) + + +def _existing_vector(table: DataFrame, name: str) -> Column: + """Returns the embeddings held in an existing column. + + Raises: + ValueError: If the column holds neither an array of numbers nor the + struct returned by ``ai_embed``. + """ + data_type = table.schema[name].dataType + if isinstance(data_type, StructType): + if _ai_embed.RESULT_FIELD not in data_type.fieldNames(): + raise ValueError( + f"The embedding column {name} is a struct without a " + f"{_ai_embed.RESULT_FIELD} field, so it does not hold the " + "output of ai_embed." + ) + return sf.col(name).getField(_ai_embed.RESULT_FIELD).cast(_VECTOR_TYPE) + if isinstance(data_type, ArrayType): + return sf.col(name).cast(_VECTOR_TYPE) + raise ValueError( + f"The embedding column {name} is a {data_type.simpleString()}. It " + "must be an array of numbers, or the struct returned by ai_embed." + ) + + +def _distance( + vector: Column, query_vector: List[float], distance_type: str +) -> Column: + """Returns the distance between each row's embedding and the search term. + + Smaller always means nearer, so ``DOT_PRODUCT``, which grows with + similarity, is negated. That keeps one ordering for every distance and + lets ``top_k`` mean the same thing throughout. + + Rows whose embedding is null, or of a different width from the search + term, score null: ``zip_with`` pads the shorter array with nulls, which + carry through the sum. They are dropped rather than compared against a + vector they cannot meaningfully be compared with. + """ + query = sf.array(*[sf.lit(float(value)) for value in query_vector]) + + if distance_type == "EUCLIDEAN": + squares = sf.zip_with(vector, query, lambda x, y: (x - y) * (x - y)) + return sf.sqrt(_total(squares)) + + products = sf.zip_with(vector, query, lambda x, y: x * y) + if distance_type == "DOT_PRODUCT": + return -_total(products) + + # COSINE. The search term's norm is a constant, so only the row's own norm + # is computed in Spark. A vector of zeroes has no direction to compare, so + # it scores null and the row is dropped. The division is the tolerant one + # because Spark raises on a division by zero under ANSI mode, which is the + # default from Spark 4, and one such row would otherwise fail the search. + query_norm = sum(value * value for value in query_vector) ** 0.5 + row_norm = sf.sqrt(_total(sf.transform(vector, lambda x: x * x))) + similarity = sf.try_divide(_total(products), row_norm * sf.lit(query_norm)) + return sf.lit(1.0) - similarity + + +def _total(values: Column) -> Column: + """Sums an array of doubles, yielding null if any element is null.""" + return sf.aggregate( + values, sf.lit(0.0), lambda accumulator, value: accumulator + value + ) diff --git a/google/cloud/dataproc_ml/sql/_client.py b/google/cloud/dataproc_ml/sql/_client.py new file mode 100644 index 0000000..427d811 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_client.py @@ -0,0 +1,207 @@ +# 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. + +"""Process-level Vertex AI clients and event loop used by the AI functions. + +Spark imports a module only once per Python worker process, so the state held +here is shared by every task the worker runs. That is what lets one HTTPS +connection pool serve all the partitions a worker processes, instead of being +rebuilt for each one. +""" + +import asyncio +import logging +import os +import threading +from typing import Awaitable, Dict, Optional, Tuple, TypeVar + +logger = logging.getLogger(__name__) + +_T = TypeVar("_T") + +# The location used when neither the caller nor the environment supplies one. +# The global endpoint is used because the current Gemini models are served +# there and it spreads load across regions; a single region can be requested +# with the location argument when data residency requires it. +DEFAULT_LOCATION = "global" + +_LOOP_LOCK = threading.Lock() +_LOOP: Optional[asyncio.AbstractEventLoop] = None + +_CLIENT_LOCK = threading.Lock() +_CLIENTS: Dict[Tuple[Optional[str], Optional[str]], object] = {} + +# Counts how often a cached client was reused. Only read by tests, to confirm +# that clients really are shared across batches and tasks. +_client_cache_hits = 0 + + +def get_event_loop() -> asyncio.AbstractEventLoop: + """Returns the event loop shared by every task in this worker process. + + The asynchronous Vertex AI client owns an ``httpx.AsyncClient`` bound to + the loop that created it. Spark calls a pandas UDF once per Arrow batch, so + running :func:`asyncio.run` per batch would build and tear down a loop for + every batch and leave any cached client bound to a closed one. Instead a + single loop runs on a daemon thread for the life of the process. + + Returns: + The running event loop. + """ + global _LOOP # pylint: disable=global-statement + with _LOOP_LOCK: + if _LOOP is None or _LOOP.is_closed(): + _LOOP = asyncio.new_event_loop() + threading.Thread( + target=_LOOP.run_forever, + name="dataproc-ml-ai-functions", + daemon=True, + ).start() + return _LOOP + + +def run_coroutine(coro: Awaitable[_T]) -> _T: + """Runs a coroutine on the shared loop and waits for its result. + + Args: + coro: The coroutine to execute. + + Returns: + Whatever the coroutine returns. + """ + return asyncio.run_coroutine_threadsafe(coro, get_event_loop()).result() + + +def resolve_project_and_location( + project: Optional[str], location: Optional[str] +) -> Tuple[Optional[str], str]: + """Fills in the project from the environment and the location from the + default when they are unset. + + Workers on Dataproc, Dataproc Serverless and GKE resolve their identity + from Application Default Credentials, so the project normally does not + have to be passed explicitly. + + The location deliberately ignores ``GOOGLE_CLOUD_LOCATION`` and + ``GOOGLE_CLOUD_REGION``. A Dataproc cluster sets those to the region it + runs in, which is rarely a statement about where a model should be + called, and the current Gemini models are served only from the global + endpoint. Honoring them would route the default model to a region that + does not serve it, and because a failed row reports its error in + ``status`` rather than raising, that would surface as a column of NULLs. + A region can still be requested explicitly with the ``location`` + argument. + + Args: + project: An explicitly configured project, or None. + location: An explicitly configured location, or None. + + Returns: + A (project, location) tuple. The project may still be None when it + cannot be discovered, in which case the client reports the problem. + """ + project = project or os.environ.get("GOOGLE_CLOUD_PROJECT") + + if not project: + try: + # Imported lazily: only needed when the project is not configured. + import google.auth # pylint: disable=import-outside-toplevel + + _, project = google.auth.default() + except Exception as e: # pylint: disable=broad-except + logger.debug("Could not infer the project from ADC: %s", e) + + return project, location or DEFAULT_LOCATION + + +def get_client(project: Optional[str], location: Optional[str]): + """Returns a client for the given project and location, creating it once. + + The first call in a worker process builds the client; later calls in that + process, including calls from other Spark tasks, reuse it. + + Args: + project: The Google Cloud project, or None to infer it. + location: The Vertex AI location, or None to infer it. + + Returns: + A ``google.genai.Client`` configured for Vertex AI. + """ + global _client_cache_hits # pylint: disable=global-statement + + key = (project, location) + with _CLIENT_LOCK: + client = _CLIENTS.get(key) + if client is not None: + _client_cache_hits += 1 + return client + + resolved_project, resolved_location = resolve_project_and_location( + project, location + ) + logger.info( + "Creating a Vertex AI client for project=%s location=%s", + resolved_project, + resolved_location, + ) + + # Imported lazily so that importing this package does not pull in the + # Vertex AI dependency until an AI function is evaluated. + from google import genai # pylint: disable=import-outside-toplevel + + client = genai.Client( + vertexai=True, + project=resolved_project, + location=resolved_location, + ) + _CLIENTS[key] = client + return client + + +def get_generate_content(project: Optional[str], location: Optional[str]): + """Returns the asynchronous ``generate_content`` of the cached client. + + Args: + project: The Google Cloud project, or None to infer it. + location: The Vertex AI location, or None to infer it. + + Returns: + An awaitable callable accepting ``model``, ``contents`` and ``config``. + """ + return get_client(project, location).aio.models.generate_content + + +def get_embed_content(project: Optional[str], location: Optional[str]): + """Returns the asynchronous ``embed_content`` of the cached client. + + Shares the client, and therefore the connection pool, with + :func:`get_generate_content`, so a query calling both functions opens one + set of connections rather than two. + + Args: + project: The Google Cloud project, or None to infer it. + location: The Vertex AI location, or None to infer it. + + Returns: + An awaitable callable accepting ``model``, ``contents`` and ``config``. + """ + return get_client(project, location).aio.models.embed_content + + +def clear_cache() -> None: + """Drops all cached clients. Intended for tests.""" + global _client_cache_hits # pylint: disable=global-statement + with _CLIENT_LOCK: + _CLIENTS.clear() + _client_cache_hits = 0 diff --git a/google/cloud/dataproc_ml/sql/_columns.py b/google/cloud/dataproc_ml/sql/_columns.py new file mode 100644 index 0000000..bbd4b98 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_columns.py @@ -0,0 +1,70 @@ +# 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. + +"""Column handling shared by the AI functions. + +Every AI function takes its input the same way: a column, the name of a column, +or a sequence of parts to concatenate. They all need that input to arrive at the +worker as a string, so the conversion lives here rather than in each function. +""" + +from typing import Sequence, Union + +from pyspark.sql import Column +from pyspark.sql import functions as sf +from pyspark.sql.types import StringType + +#: A column, the name of a column, or a sequence of either. +ColumnArg = Union[Column, str, Sequence[Union[Column, str]]] + + +def to_string_column(value: ColumnArg, *, argument_name: str) -> Column: + """Builds the single string column holding each row's input. + + Every path casts to STRING. An input column of another type would otherwise + reach the worker as a non-string value and be treated as missing, turning a + whole column into nulls without any error. + + Args: + value: The column, the name of a column, or a sequence of parts. + argument_name: The name of the caller's argument, used in error + messages so that they name the argument the user actually passed. + + Returns: + A single column of type STRING. + + Raises: + ValueError: If the value is not a column, a column name, or a non-empty + sequence of those. + """ + if isinstance(value, Column): + return value.cast(StringType()) + if isinstance(value, str): + return sf.col(value).cast(StringType()) + if isinstance(value, Sequence): + parts = list(value) + if not parts: + raise ValueError(f"{argument_name} must contain at least one part.") + columns = [ + part if isinstance(part, Column) else sf.col(part) for part in parts + ] + if len(columns) == 1: + return columns[0].cast(StringType()) + # Concatenation matches BigQuery, where the parts of a struct argument + # are combined in the order they are given. + return sf.concat(*[column.cast(StringType()) for column in columns]) + raise ValueError( + f"{argument_name} must be a Column, a column name, or a sequence of " + f"those, got {type(value).__name__}." + ) diff --git a/google/cloud/dataproc_ml/sql/_model_params.py b/google/cloud/dataproc_ml/sql/_model_params.py new file mode 100644 index 0000000..cdcbec0 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_model_params.py @@ -0,0 +1,261 @@ +# 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. + +"""Translation of the ``model_params`` option into a generation config. + +``model_params`` follows the Vertex AI ``generateContent`` request body, minus +the ``contents`` field which is supplied by the prompt. That body nests the +generation settings under ``generation_config`` whereas the client library +takes them as a flat configuration object, so this module flattens the request +body and validates it. +""" + +import re +from typing import Any, Dict, Optional + +_CAMEL_BOUNDARY = re.compile(r"(? str: + """Converts a camelCase request body key to its snake_case equivalent.""" + return _CAMEL_BOUNDARY.sub("_", key).lower() + + +def normalize_model_params( + model_params: Optional[Dict[str, Any]], + response_schema: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """Flattens ``model_params`` into generation config keyword arguments. + + Both camelCase and snake_case keys are accepted, matching the behavior of + the JSON request body. + + Args: + model_params: The user supplied request body, or None. + response_schema: A response schema derived from ``output_schema``, or + None when the function returns free form text. + + Returns: + A dictionary of keyword arguments for the generation config. + + Raises: + ValueError: If the parameters contain ``contents``, declare a response + schema that conflicts with ``output_schema``, or do not describe a + valid request body. + """ + params = model_params or {} + if not isinstance(params, dict): + raise ValueError( + "model_params must be a dictionary matching the generateContent " + f"request body, got {type(params).__name__}." + ) + + flat: Dict[str, Any] = {} + for key, value in params.items(): + snake_key = _to_snake_case(str(key)) + if snake_key == "contents": + raise ValueError(_CONTENTS_ERROR) + if snake_key in _GENERATION_CONFIG_KEYS: + if not isinstance(value, dict): + raise ValueError( + "model_params.generation_config must be a dictionary, got " + f"{type(value).__name__}." + ) + for nested_key, nested_value in value.items(): + flat[_to_snake_case(str(nested_key))] = nested_value + else: + flat[snake_key] = value + + if response_schema is not None: + if flat.get("response_schema") or flat.get("response_json_schema"): + raise ValueError(_BOTH_SCHEMAS_ERROR) + flat["response_mime_type"] = "application/json" + flat["response_schema"] = response_schema + + # Building the config here surfaces invalid parameters on the driver, + # before any executor work is scheduled. + build_generation_config(flat) + return flat + + +def build_generation_config(flat_params: Dict[str, Any]): + """Builds the client library generation config. + + Args: + flat_params: Keyword arguments produced by + :func:`normalize_model_params`. + + Returns: + A ``google.genai.types.GenerateContentConfig``. + + Raises: + ValueError: If the parameters do not describe a valid request body. + """ + # Imported lazily to keep the Vertex AI dependency off the import path of + # the package itself. + from google.genai import types # pylint: disable=import-outside-toplevel + + try: + return types.GenerateContentConfig(**flat_params) + except Exception as e: + raise ValueError( + "The model_params is invalid and does not conform to the " + "generateContent request body. See " + "https://cloud.google.com/vertex-ai/docs/reference/rest/v1/" + f"projects.locations.publishers.models/generateContent. Cause: {e}" + ) from e + + +def normalize_embed_params( + model_params: Optional[Dict[str, Any]], + task_type: Optional[str] = None, + title: Optional[str] = None, +) -> Dict[str, Any]: + """Flattens the embedding ``model_params`` into config keyword arguments. + + BigQuery exposes ``task_type`` and ``title`` as their own arguments while + everything else, notably ``outputDimensionality``, goes through + ``model_params``. The client library takes all of them as one flat config, + so they are merged here. Both camelCase and snake_case keys are accepted. + + Args: + model_params: The user supplied parameters, or None. + task_type: The ``task_type`` argument, or None. + title: The ``title`` argument, or None. + + Returns: + A dictionary of keyword arguments for the embedding config. + + Raises: + ValueError: If the parameters contain ``contents``, set a value that is + also given as its own argument, use an unknown ``task_type``, set + ``title`` without ``task_type='RETRIEVAL_DOCUMENT'``, or otherwise + do not describe a valid request. + """ + params = model_params or {} + if not isinstance(params, dict): + raise ValueError( + "model_params must be a dictionary matching the embedContent " + f"request body, got {type(params).__name__}." + ) + + flat: Dict[str, Any] = {} + for key, value in params.items(): + snake_key = _to_snake_case(str(key)) + if snake_key == "contents": + raise ValueError(_EMBED_CONTENTS_ERROR) + flat[snake_key] = value + + for name, value in (("task_type", task_type), ("title", title)): + if value is None: + continue + if flat.get(name) is not None: + raise ValueError( + f"{name} can only be set either in the model_params or in the " + f"{name} argument, not both." + ) + flat[name] = value + + _validate_task_type(flat) + + # Building the config here surfaces invalid parameters on the driver, + # before any executor work is scheduled. + build_embed_config(flat) + return flat + + +def _validate_task_type(flat: Dict[str, Any]) -> None: + """Checks ``task_type`` and the ``title`` that depends on it.""" + task_type = flat.get("task_type") + if task_type is not None: + if ( + not isinstance(task_type, str) + or task_type.upper() not in TASK_TYPES + ): + supported = ", ".join(sorted(TASK_TYPES)) + raise ValueError( + f"Unsupported task_type: {task_type!r}. Supported values are " + f"{supported}." + ) + # Normalized so that a lowercase spelling, which reads naturally in + # SQL, reaches the API in the form it expects. + flat["task_type"] = task_type.upper() + + if flat.get("title") is not None and flat.get("task_type") != ( + "RETRIEVAL_DOCUMENT" + ): + raise ValueError( + "title can only be set when task_type is 'RETRIEVAL_DOCUMENT'." + ) + + +def build_embed_config(flat_params: Dict[str, Any]): + """Builds the client library embedding config. + + Args: + flat_params: Keyword arguments produced by + :func:`normalize_embed_params`. + + Returns: + A ``google.genai.types.EmbedContentConfig``. + + Raises: + ValueError: If the parameters do not describe a valid request. + """ + # Imported lazily to keep the Vertex AI dependency off the import path of + # the package itself. + from google.genai import types # pylint: disable=import-outside-toplevel + + try: + return types.EmbedContentConfig(**flat_params) + except Exception as e: + raise ValueError( + "The model_params is invalid and does not conform to the " + "embedContent request body. See " + "https://cloud.google.com/vertex-ai/docs/reference/rest/v1/" + f"projects.locations.publishers.models/embedContent. Cause: {e}" + ) from e diff --git a/google/cloud/dataproc_ml/sql/_output_schema.py b/google/cloud/dataproc_ml/sql/_output_schema.py new file mode 100644 index 0000000..d6c8f64 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_output_schema.py @@ -0,0 +1,241 @@ +# 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. + +"""Translation of the ``output_schema`` option into Spark and Gemini schemas. + +A single schema declaration drives two things: the Spark ``StructType`` that +the AI function returns, and the response schema sent to the model so that it +replies with conforming JSON. +""" + +import re +from typing import Any, Dict, Union + +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DataType, + DoubleType, + FloatType, + IntegerType, + LongType, + StringType, + StructField, + StructType, +) + +# BigQuery spells some scalar types differently from Spark. Both spellings are +# accepted so that a schema can be copied straight from a BigQuery query. +_TYPE_ALIASES = { + "INT64": "BIGINT", + "FLOAT64": "DOUBLE", + "BOOL": "BOOLEAN", +} + +# An alias is only rewritten where a type may appear, which is recognized by +# what surrounds it. Matching the bare name would also rewrite a *field* named +# after a type: "int64 INT64" would become "BIGINT BIGINT", quietly renaming +# the field to BIGINT instead of failing. A type is preceded by a space, a +# colon or "<", and is followed by the end of the field or by one of the +# modifiers that Spark allows after a type. +_ALIAS_PATTERNS = tuple( + ( + re.compile( + rf"(?<=[\s:<]){bigquery_name}" + r"(?=\s*(?:,|>|\)|$|NOT\s+NULL|COMMENT\b))", + re.IGNORECASE, + ), + spark_name, + ) + for bigquery_name, spark_name in _TYPE_ALIASES.items() +) + +# Matches the field options syntax that BigQuery supports but this library does +# not parse yet, e.g. ``OPTIONS(description = '...')``. +_OPTIONS_PATTERN = re.compile(r"\bOPTIONS\s*\(", re.IGNORECASE) + +# Mirrors BigQuery's limit on how deeply an output schema may nest. +MAX_NESTING_DEPTH = 5 + +_SCALAR_RESPONSE_TYPES = { + StringType: "STRING", + BooleanType: "BOOLEAN", + LongType: "INTEGER", + IntegerType: "INTEGER", + DoubleType: "NUMBER", + FloatType: "NUMBER", +} + +_SUPPORTED_TYPE_NAMES = "STRING, INT64, FLOAT64, BOOL, ARRAY and STRUCT" + + +def parse_output_schema( + output_schema: Union[str, StructType], +) -> StructType: + """Parses the ``output_schema`` option into a Spark ``StructType``. + + Args: + output_schema: Either a comma separated list of fields, such as + ``"sentiment STRING, score INT64"``, or an already built + ``StructType``. + + Returns: + The validated struct, with fields sorted by name. + + Raises: + ValueError: If the schema cannot be parsed, uses an unsupported type, + or nests too deeply. + """ + if isinstance(output_schema, StructType): + struct = output_schema + elif isinstance(output_schema, str): + struct = _parse_ddl(output_schema) + else: + raise ValueError( + "output_schema must be a string or a pyspark StructType, got " + f"{type(output_schema).__name__}." + ) + + if not struct.fields: + raise ValueError("output_schema must declare at least one field.") + + _validate(struct, depth=1) + return _sorted_struct(struct) + + +def _parse_ddl(output_schema: str) -> StructType: + """Parses the schema string, translating BigQuery type names first.""" + if _OPTIONS_PATTERN.search(output_schema): + raise ValueError( + "The OPTIONS clause in output_schema is not supported yet. " + "Remove it and describe the fields in the prompt instead." + ) + + normalized = output_schema + for pattern, spark_name in _ALIAS_PATTERNS: + normalized = pattern.sub(spark_name, normalized) + + try: + parsed = StructType.fromDDL(normalized) + except Exception as e: + raise ValueError( + "Failed to parse the option `output_schema`. It must be a valid " + "comma separated list of field names and types, for example " + f"'sentiment STRING, score INT64'. Got: {output_schema!r}." + ) from e + + if not isinstance(parsed, StructType): + raise ValueError( + "Failed to parse the option `output_schema` into a struct: " + f"{output_schema!r}." + ) + return parsed + + +def _validate(data_type: DataType, depth: int) -> None: + """Recursively checks that only supported types are used.""" + if depth > MAX_NESTING_DEPTH: + raise ValueError( + "Option `output_schema` has too many nested levels. Maximum " + f"allowed is {MAX_NESTING_DEPTH}." + ) + + if isinstance(data_type, StructType): + names = [field.name for field in data_type.fields] + duplicates = {name for name in names if names.count(name) > 1} + if duplicates: + names = ", ".join(sorted(duplicates)) + raise ValueError( + f"Option `output_schema` has duplicate field names: {names}." + ) + for field in data_type.fields: + _validate(field.dataType, depth + 1) + elif isinstance(data_type, ArrayType): + _validate(data_type.elementType, depth + 1) + elif type(data_type) not in _SCALAR_RESPONSE_TYPES: + raise ValueError( + "Unsupported field type in output_schema: " + f"{data_type.simpleString()}." + f" Only {_SUPPORTED_TYPE_NAMES} are supported." + ) + + +def _sorted_struct(struct: StructType) -> StructType: + """Returns the struct with its fields, and nested fields, sorted by name. + + BigQuery returns the fields of a custom output schema in alphabetical + order rather than declaration order, and this library matches it so that + the same query produces the same column layout on both engines. + """ + return StructType( + [ + StructField( + field.name, + _sorted_type(field.dataType), + nullable=field.nullable, + ) + for field in sorted(struct.fields, key=lambda f: f.name) + ] + ) + + +def _sorted_type(data_type: DataType) -> DataType: + """Applies :func:`_sorted_struct` to any nested structs.""" + if isinstance(data_type, StructType): + return _sorted_struct(data_type) + if isinstance(data_type, ArrayType): + return ArrayType( + _sorted_type(data_type.elementType), + containsNull=data_type.containsNull, + ) + return data_type + + +def to_response_schema(struct: StructType) -> Dict[str, Any]: + """Converts a Spark struct into a Gemini response schema. + + Args: + struct: The parsed output schema. + + Returns: + A schema dictionary suitable for the ``response_schema`` generation + setting. + """ + return _to_response_schema(struct) + + +def _to_response_schema(data_type: DataType) -> Dict[str, Any]: + """Recursively converts a Spark type into a Gemini response schema.""" + if isinstance(data_type, StructType): + required = [ + field.name for field in data_type.fields if not field.nullable + ] + schema: Dict[str, Any] = { + "type": "OBJECT", + "properties": { + field.name: _to_response_schema(field.dataType) + for field in data_type.fields + }, + } + if required: + schema["required"] = required + return schema + + if isinstance(data_type, ArrayType): + return { + "type": "ARRAY", + "items": _to_response_schema(data_type.elementType), + } + + return {"type": _SCALAR_RESPONSE_TYPES[type(data_type)]} diff --git a/google/cloud/dataproc_ml/sql/_retry.py b/google/cloud/dataproc_ml/sql/_retry.py new file mode 100644 index 0000000..5f3c24d --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_retry.py @@ -0,0 +1,187 @@ +# 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. + +"""Retry policy and error reporting for the AI functions. + +Transient backend failures are retried with full-jitter exponential backoff. +Errors that survive the retry budget are reported per row in the ``status`` +field of the result, using the same wording as the equivalent BigQuery AI +function, so that a single bad row never fails an entire Spark stage. +""" + +import logging +import random +from typing import Optional + +logger = logging.getLogger(__name__) + +# HTTP statuses that are worth retrying: rate limiting, transient backend +# failures and gateway timeouts. +RETRYABLE_HTTP_CODES = frozenset({408, 429, 500, 503, 504}) + +# Canonical error codes whose messages BigQuery passes through verbatim; these +# describe a problem with the request itself, so the detail is actionable. +_PASSTHROUGH_CODES = frozenset( + { + "FAILED_PRECONDITION", + "PERMISSION_DENIED", + "NOT_FOUND", + "INVALID_ARGUMENT", + "OUT_OF_RANGE", + } +) + +_QUOTA_MESSAGE_PREFIX = "Quota exceeded for metric " + +_HTTP_TO_CANONICAL_CODE = { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 408: "DEADLINE_EXCEEDED", + 409: "ALREADY_EXISTS", + 413: "OUT_OF_RANGE", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 500: "INTERNAL", + 501: "NOT_IMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", +} + +# Reported when a structured response cannot be parsed as JSON. +PARSE_FAILURE_STATUS = ( + "Failed to parse the response - please check the full_response field." +) + + +def _api_error_types(): + """Returns the google-genai API error types, imported lazily.""" + from google.genai import errors # pylint: disable=import-outside-toplevel + + return errors.APIError + + +def _transport_error_types(): + """Returns the transport error types treated as retryable.""" + import httpx # pylint: disable=import-outside-toplevel + + return (httpx.TransportError, TimeoutError) + + +def is_reportable(exception: BaseException) -> bool: + """Reports whether an exception should become a row-level ``status``. + + Only failures that come from the model backend are reported per row. + Anything else, such as a bug in this library, is left to propagate so that + it is not silently hidden in a column of an otherwise successful query. + + Args: + exception: The exception raised while calling the model. + + Returns: + True if the exception should be converted into a status message. + """ + return isinstance( + exception, (_api_error_types(),) + _transport_error_types() + ) + + +def is_retryable(exception: BaseException) -> bool: + """Reports whether a failed request is worth retrying. + + Args: + exception: The exception raised while calling the model. + + Returns: + True for rate limiting, transient backend errors and transport errors. + """ + if isinstance(exception, _transport_error_types()): + return True + if isinstance(exception, _api_error_types()): + return _http_code(exception) in RETRYABLE_HTTP_CODES + return False + + +def _http_code(exception: BaseException) -> Optional[int]: + """Extracts the HTTP status code from an API error, if it has one.""" + code = getattr(exception, "code", None) + return code if isinstance(code, int) else None + + +def _canonical_code(exception: BaseException) -> str: + """Determines the canonical error code, e.g. ``RESOURCE_EXHAUSTED``.""" + status = getattr(exception, "status", None) + if isinstance(status, str) and status: + return status + code = _http_code(exception) + if code is not None and code in _HTTP_TO_CANONICAL_CODE: + return _HTTP_TO_CANONICAL_CODE[code] + if isinstance(exception, _transport_error_types()): + return "UNAVAILABLE" + return "UNKNOWN" + + +def to_status_message(exception: BaseException) -> str: + """Converts a model error into the row's ``status`` value. + + The wording matches BigQuery's AI functions: request errors are surfaced + verbatim, quota messages are passed through, and everything else is + summarized as a retryable error. + + Args: + exception: The exception raised while calling the model. + + Returns: + A human readable status message. Never empty. + """ + code = _canonical_code(exception) + message = getattr(exception, "message", None) or str(exception) + + if code in _PASSTHROUGH_CODES: + return f"{code}: {message}" + if code == "RESOURCE_EXHAUSTED" and message.startswith( + _QUOTA_MESSAGE_PREFIX + ): + return message + return ( + f"A retryable error occurred: {code} error from remote " + "service/endpoint." + ) + + +def full_jitter_delay( + attempt: int, + base_delay_seconds: float, + max_delay_seconds: float, + rng: Optional[random.Random] = None, +) -> float: + """Computes a full-jitter exponential backoff delay. + + The delay is drawn uniformly from ``[0, min(cap, base * 2 ** attempt)]``. + Randomising the whole interval, rather than only a fraction of it, stops + the tasks of a large cluster from retrying in lockstep after a shared + throttling event. + + Args: + attempt: The zero-based index of the attempt that just failed. + base_delay_seconds: The delay unit that is doubled on each attempt. + max_delay_seconds: The upper bound on the exponential term. + rng: Random source, overridable for tests. + + Returns: + The number of seconds to sleep before the next attempt. + """ + ceiling = min(max_delay_seconds, base_delay_seconds * (2**attempt)) + return (rng or random).uniform(0.0, max(0.0, ceiling)) diff --git a/pyproject.toml b/pyproject.toml index 4666af4..9381ea2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ requires-python = ">=3.12" dependencies = [ "google-cloud-aiplatform>=1.121.0, <2.0.0", "google-cloud-storage>=2.19.0, <3.0.0", + "google-genai>=1.38.0, <3.0.0", "pandas>=2.2.2, <3.0.0", "pyarrow>=18.1.0, <20.0.0", "tenacity>=8.5.0, <9.0.0", diff --git a/tests/integration/sql/__init__.py b/tests/integration/sql/__init__.py new file mode 100644 index 0000000..75edc9b --- /dev/null +++ b/tests/integration/sql/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/integration/sql/test_ai_embed.py b/tests/integration/sql/test_ai_embed.py new file mode 100644 index 0000000..e32ee2e --- /dev/null +++ b/tests/integration/sql/test_ai_embed.py @@ -0,0 +1,201 @@ +# 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. + +"""Integration tests that call Vertex AI for real. + +They rely on Application Default Credentials and, unless GOOGLE_CLOUD_PROJECT +is set, on the project discovered from the environment. The location is not +taken from the environment: these tests exercise the default global endpoint. +Running them incurs Vertex AI charges. +""" + +import math +import unittest + +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf + +from google.cloud.dataproc_ml.sql import _ai_embed as ai_embed_module +from google.cloud.dataproc_ml.sql import ai_embed, ai_embed_udf + +# The number of values text-embedding-005 returns when nothing else is asked. +_DEFAULT_DIMENSIONS = 768 + + +def cosine_similarity(left, right): + """Returns how close two embeddings are, between -1 and 1.""" + dot = sum(a * b for a, b in zip(left, right)) + norm = math.sqrt(sum(a * a for a in left)) * math.sqrt( + sum(b * b for b in right) + ) + return dot / norm + + +class TestAiEmbed(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.spark = SparkSession.builder.getOrCreate() + cls.sentences = cls.spark.createDataFrame( + [ + ("cat", "The cat slept on the warm windowsill."), + ("kitten", "A kitten dozed in the afternoon sun."), + ("finance", "The central bank raised interest rates today."), + ("cooking", "Simmer the sauce for twenty minutes."), + ("travel", "The night train reaches Oslo before dawn."), + ("sports", "She scored twice in the second half."), + ("weather", "Heavy rain is expected across the region."), + # Several rows per batch, spread over more than one task. + ], + ["topic", "text"], + ).repartition(3) + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + def test_embeds_every_row(self): + rows = ( + self.sentences.withColumn("e", ai_embed(sf.col("text"))) + .select("topic", "e.result", "e.status") + .collect() + ) + + self.assertEqual(len(rows), 7) + for row in rows: + with self.subTest(topic=row["topic"]): + self.assertEqual(row["status"], "") + self.assertEqual(len(row["result"]), _DEFAULT_DIMENSIONS) + + def test_related_sentences_embed_closer_together(self): + # The real check that a row kept its own vector: if the rows were + # shuffled against their embeddings this would not hold. + vectors = { + row["topic"]: row["result"] + for row in self.sentences.withColumn("e", ai_embed(sf.col("text"))) + .select("topic", "e.result") + .collect() + } + + related = cosine_similarity(vectors["cat"], vectors["kitten"]) + unrelated = cosine_similarity(vectors["cat"], vectors["finance"]) + self.assertGreater(related, unrelated) + + def test_output_dimensionality_shortens_the_vector(self): + row = ( + self.sentences.limit(1) + .withColumn( + "e", + ai_embed( + sf.col("text"), model_params={"outputDimensionality": 256} + ), + ) + .select("e.result", "e.status") + .first() + ) + + self.assertEqual(row["status"], "") + self.assertEqual(len(row["result"]), 256) + + def test_task_type_and_title_are_accepted(self): + rows = ( + self.sentences.limit(2) + .withColumn( + "e", + ai_embed( + sf.col("text"), + task_type="RETRIEVAL_DOCUMENT", + title="Sentences", + ), + ) + .select("e.result", "e.status") + .collect() + ) + + for row in rows: + self.assertEqual(row["status"], "") + self.assertEqual(len(row["result"]), _DEFAULT_DIMENSIONS) + + def test_more_rows_than_one_request_allows(self): + # The endpoint accepts at most 250 instances per request while an + # Arrow batch holds far more, so this exercises the splitting against + # the real limit. + limit = ai_embed_module._MAX_INSTANCES_PER_REQUEST + many = self.spark.range(limit + 50).withColumn( + "text", sf.concat(sf.lit("Row number "), sf.col("id")) + ) + + rows = ( + many.withColumn("e", ai_embed(sf.col("text"))) + .select("id", "e.result", "e.status") + .collect() + ) + + self.assertEqual(len(rows), limit + 50) + for row in rows: + self.assertEqual(row["status"], "") + self.assertEqual(len(row["result"]), _DEFAULT_DIMENSIONS) + + def test_null_and_empty_rows_are_answered_without_the_model(self): + df = self.spark.createDataFrame( + [("Oslo",), (None,), ("",)], "text: string" + ) + rows = ( + df.withColumn("e", ai_embed(sf.col("text"))) + .select("text", "e.result", "e.status") + .collect() + ) + + by_text = {row["text"]: row for row in rows} + self.assertEqual(len(by_text["Oslo"]["result"]), _DEFAULT_DIMENSIONS) + self.assertIsNone(by_text[None]["result"]) + self.assertEqual(by_text[None]["status"], "") + self.assertIsNone(by_text[""]["result"]) + self.assertIn("INVALID_ARGUMENT", by_text[""]["status"]) + + def test_unknown_model_is_reported_per_row(self): + # A bad endpoint must not fail the query: every row reports the error + # in its status, exactly as BigQuery's AI.EMBED does. + rows = ( + self.sentences.limit(2) + .withColumn( + "e", + ai_embed( + sf.col("text"), endpoint="text-embedding-does-not-exist" + ), + ) + .select("e.result", "e.status") + .collect() + ) + + self.assertEqual(len(rows), 2) + for row in rows: + self.assertIsNone(row["result"]) + self.assertIn("NOT_FOUND", row["status"]) + + def test_registered_for_spark_sql(self): + self.spark.udf.register("ai_embed", ai_embed_udf()) + self.sentences.createOrReplaceTempView("sentences") + + rows = self.spark.sql( + "SELECT topic, ai_embed(text).result AS vector FROM sentences" + ).collect() + + self.assertEqual(len(rows), 7) + for row in rows: + self.assertEqual(len(row["vector"]), _DEFAULT_DIMENSIONS) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/sql/test_ai_generate.py b/tests/integration/sql/test_ai_generate.py new file mode 100644 index 0000000..d197c3b --- /dev/null +++ b/tests/integration/sql/test_ai_generate.py @@ -0,0 +1,196 @@ +# 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. + +"""Integration tests that call Vertex AI for real. + +They rely on Application Default Credentials and, unless GOOGLE_CLOUD_PROJECT +is set, on the project discovered from the environment. The location is not +taken from the environment: these tests exercise the default global endpoint. +Running them incurs Vertex AI charges. +""" + +import json +import unittest + +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf + +from google.cloud.dataproc_ml.sql import ai_generate, ai_generate_udf +from google.cloud.dataproc_ml.sql import _ai_generate as ai_generate_module + +_AIRPORT_PROMPT = ( + "What is the airport code of the largest airport in {city}? " + "Answer with the three letter code only." +) + + +class TestAiGenerate(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.spark = SparkSession.builder.getOrCreate() + cls.cities = cls.spark.createDataFrame( + [ + ("Bengaluru", "BLR"), + ("London", "LHR"), + ("San Francisco", "SFO"), + ("Paris", "CDG"), + ("Tokyo", "HND"), + ("Sydney", "SYD"), + ("New York", "JFK"), + ], + ["city", "expected"], + # Several rows per batch, spread over more than one task. + ).repartition(3) + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + def _prompt(self): + before, after = _AIRPORT_PROMPT.split("{city}", maxsplit=1) + return [sf.lit(before), sf.col("city"), sf.lit(after)] + + def test_generates_an_answer_for_every_row(self): + rows = ( + self.cities.withColumn("g", ai_generate(self._prompt())) + .select("city", "expected", "g.result", "g.status") + .collect() + ) + + self.assertEqual(len(rows), 7) + for row in rows: + with self.subTest(city=row["city"]): + self.assertEqual(row["status"], "") + self.assertIsNotNone(row["result"]) + self.assertIn(row["expected"], row["result"].upper()) + + def test_full_response_carries_usage_metadata(self): + result = self.cities.limit(1).withColumn( + "g", ai_generate(self._prompt()) + ) + payload_column = ( + "to_json(g.full_response)" + if ai_generate_module._VARIANT_SUPPORTED + else "g.full_response" + ) + payload = json.loads( + result.selectExpr(f"{payload_column} AS payload").first()["payload"] + ) + + self.assertIn("candidates", payload) + self.assertIn("usage_metadata", payload) + + def test_model_params_are_applied(self): + result = ( + self.cities.limit(2) + .withColumn( + "g", + ai_generate( + self._prompt(), + model_params={ + "generation_config": { + "temperature": 0.0, + "thinking_config": {"thinking_budget": 0}, + } + }, + ), + ) + .select("g.result", "g.status") + .collect() + ) + + for row in result: + self.assertEqual(row["status"], "") + self.assertIsNotNone(row["result"]) + + def test_structured_output(self): + reviews = self.spark.createDataFrame( + [ + ("The battery lasts two days and the screen is gorgeous.",), + ("It broke after a week and support never replied.",), + ], + ["review"], + ) + + result = reviews.withColumn( + "g", + ai_generate( + [ + sf.lit("Analyze this product review: "), + sf.col("review"), + ], + output_schema=( + "sentiment STRING, score INT64, topics ARRAY" + ), + model_params={"generation_config": {"temperature": 0.0}}, + ), + ) + + # BigQuery returns custom schema fields sorted by name. + self.assertEqual( + [field.name for field in result.schema["g"].dataType.fields], + ["score", "sentiment", "topics", "full_response", "status"], + ) + + rows = result.select( + "g.sentiment", "g.score", "g.topics", "g.status" + ).collect() + self.assertEqual(len(rows), 2) + for row in rows: + with self.subTest(sentiment=row["sentiment"]): + self.assertEqual(row["status"], "") + self.assertIsInstance(row["sentiment"], str) + self.assertIsInstance(row["score"], int) + self.assertTrue(row["topics"]) + + def test_unknown_model_is_reported_per_row(self): + # A bad endpoint must not fail the query: every row reports the error + # in its status, exactly as BigQuery's AI.GENERATE does. + rows = ( + self.cities.limit(2) + .withColumn( + "g", + ai_generate( + self._prompt(), + endpoint="gemini-does-not-exist", + ), + ) + .select("g.result", "g.status") + .collect() + ) + + self.assertEqual(len(rows), 2) + for row in rows: + self.assertIsNone(row["result"]) + self.assertIn("NOT_FOUND", row["status"]) + + def test_registered_for_spark_sql(self): + self.spark.udf.register("ai_generate", ai_generate_udf()) + self.cities.createOrReplaceTempView("cities") + + rows = self.spark.sql( + "SELECT city," + " ai_generate(CONCAT('Airport code of the largest airport in '," + " city, '? Three letters only.')).result AS code" + " FROM cities" + ).collect() + + self.assertEqual(len(rows), 7) + for row in rows: + self.assertIsNotNone(row["code"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/integration/sql/test_ai_search.py b/tests/integration/sql/test_ai_search.py new file mode 100644 index 0000000..5feba7c --- /dev/null +++ b/tests/integration/sql/test_ai_search.py @@ -0,0 +1,174 @@ +# 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. + +"""Integration tests that call Vertex AI for real. + +They rely on Application Default Credentials and, unless GOOGLE_CLOUD_PROJECT +is set, on the project discovered from the environment. The location is not +taken from the environment: these tests exercise the default global endpoint. +Running them incurs Vertex AI charges. +""" + +import unittest + +from pyspark.sql import SparkSession + +from google.cloud.dataproc_ml.sql import _ai_search as ai_search_module +from google.cloud.dataproc_ml.sql import ai_embed, ai_search + +# Products that a reader can rank by hand. Each pair of rows is close on one +# subject and far apart on every other, so a search for one subject has an +# unambiguous answer. +PRODUCTS = [ + (1, "Stunt kite", "A huge kite that loops and dives in a light breeze."), + (2, "Bath ducks", "Six squeaky rubber ducks for the tub."), + (3, "Desk lamp", "An adjustable lamp with a warm dimmable bulb."), + (4, "Floor lamp", "A tall standing light for a reading corner."), + (5, "Chef's knife", "A forged carbon steel blade for the kitchen."), + (6, "Cutting board", "An end grain maple block for chopping."), +] + +# The rows a reader would pick for each search, in no particular order. +EXPECTED = { + "a really fun toy": {1, 2}, + "something to read by in the evening": {3, 4}, + "preparing dinner": {5, 6}, +} + + +class TestAiSearch(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.spark = SparkSession.builder.getOrCreate() + cls.products = cls.spark.createDataFrame( + PRODUCTS, ["id", "name", "description"] + ).repartition(3) + # Embedded once and kept, both because it is what the function is + # meant to be used with and because it keeps these tests cheap. + cls.embedded = cls.products.withColumn( + "description_embedding", ai_embed("description") + ).cache() + cls.embedded.count() + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + def ids(self, hits): + """Returns the ids of the rows found, in the order they came back.""" + return [row["base"]["id"] for row in hits.collect()] + + def test_finds_the_rows_a_reader_would_pick(self): + for query, expected in EXPECTED.items(): + with self.subTest(query=query): + hits = ai_search(self.embedded, "description", query, top_k=2) + self.assertEqual(set(self.ids(hits)), expected) + + def test_every_distance_agrees_on_the_nearest_row(self): + query = "preparing dinner" + nearest = set() + for distance_type in ai_search_module.DISTANCE_TYPES: + with self.subTest(distance_type=distance_type): + hits = ai_search( + self.embedded, + "description", + query, + top_k=2, + distance_type=distance_type, + ) + found = set(self.ids(hits)) + self.assertEqual(found, EXPECTED[query]) + nearest.add(tuple(self.ids(hits))) + # The three distances rank the same two rows. They may disagree on + # the order within the pair, which is not something to assert on. + + def test_the_returned_columns_match_bigquery(self): + hits = ai_search(self.embedded, "description", "a lamp", top_k=1) + self.assertEqual(hits.columns, ["base", "distance"]) + row = hits.collect()[0] + self.assertEqual( + set(row["base"].asDict()), + {"id", "name", "description", "description_embedding"}, + ) + self.assertIsInstance(row["distance"], float) + + def test_top_k_bounds_the_result(self): + hits = ai_search(self.embedded, "description", "a lamp", top_k=3) + self.assertEqual(hits.count(), 3) + + def test_the_distance_never_decreases(self): + distances = [ + row["distance"] + for row in ai_search( + self.embedded, "description", "a lamp", top_k=6 + ).collect() + ] + self.assertEqual(distances, sorted(distances)) + + def test_embeds_the_table_when_it_carries_no_embeddings(self): + # The same answer, paid for with a call per row instead of one call. + hits = ai_search( + self.products, "description", "a really fun toy", top_k=2 + ) + self.assertEqual(set(self.ids(hits)), EXPECTED["a really fun toy"]) + + def test_a_named_embedding_column_is_used(self): + renamed = self.embedded.withColumnRenamed( + "description_embedding", "vectors" + ) + hits = ai_search( + renamed, + "description", + "preparing dinner", + top_k=2, + embedding_column="vectors", + ) + self.assertEqual(set(self.ids(hits)), EXPECTED["preparing dinner"]) + + def test_a_named_table_is_searched(self): + self.embedded.createOrReplaceTempView("products_for_search") + hits = ai_search( + "products_for_search", "description", "a lamp", top_k=2 + ) + self.assertEqual(set(self.ids(hits)), {3, 4}) + + def test_a_shorter_vector_still_searches(self): + # outputDimensionality has to match on both sides, so the table is + # re-embedded at the same width as the search term. + params = {"outputDimensionality": 256} + hits = ai_search( + self.products, + "description", + "preparing dinner", + top_k=2, + model_params=params, + ) + self.assertEqual(set(self.ids(hits)), EXPECTED["preparing dinner"]) + + def test_an_unknown_endpoint_fails_the_search(self): + # Unlike a row, a search term that cannot be embedded has nowhere to + # report a failure, so it raises. + with self.assertRaises(ValueError) as caught: + ai_search( + self.embedded, + "description", + "a lamp", + endpoint="text-embedding-does-not-exist", + ) + self.assertIn("a lamp", str(caught.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/__init__.py b/tests/unit/sql/__init__.py new file mode 100644 index 0000000..75edc9b --- /dev/null +++ b/tests/unit/sql/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/tests/unit/sql/test_ai_embed.py b/tests/unit/sql/test_ai_embed.py new file mode 100644 index 0000000..0820c0d --- /dev/null +++ b/tests/unit/sql/test_ai_embed.py @@ -0,0 +1,484 @@ +# 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. + +import asyncio +import unittest +from unittest import mock + +from google.genai import errors, types +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf +from pyspark.sql.types import ArrayType, DoubleType + +# pylint: disable=ungrouped-imports +from google.cloud.dataproc_ml.sql import _ai_embed as ai_embed_module +from google.cloud.dataproc_ml.sql import ai_embed, ai_embed_udf + +# pylint: enable=ungrouped-imports + + +def embedding_of(text): + """Builds a vector that can be traced back to the text it came from. + + The first value is the length of the text, so a test can check that a row + was matched with its own embedding rather than a neighbour's. + """ + return types.ContentEmbedding(values=[float(len(text)), 0.5]) + + +def embeds_each_content(recorded=None): + """Returns a fake ``embed_content`` that answers every instance sent. + + The returned coroutine function is defined here rather than at module + level so that it is serialized by value and can run inside a Spark worker. + + Args: + recorded: An optional list that each request's instances are appended + to. Left unset when the fake has to cross into a worker. + + Returns: + An awaitable callable with the same shape as the client method. + """ + + async def _embed_content(*, model, contents, config): + del model, config + if recorded is not None: + recorded.append(list(contents)) + return types.EmbedContentResponse( + embeddings=[embedding_of(text) for text in contents] + ) + + return _embed_content + + +def api_error(code, status, message): + """Builds an API error the way the client library reports one.""" + return errors.APIError( + code, {"error": {"code": code, "status": status, "message": message}} + ) + + +def replies_with(*outcomes): + """Returns a fake ``embed_content`` that replays the given outcomes. + + Args: + *outcomes: Responses or exceptions, consumed in order. The last one is + repeated once exhausted. A response of ``None`` stands for "embed + whatever was sent", which keeps a success after a retry readable. + + Returns: + An awaitable callable with the same shape as the client method. + """ + remaining = list(outcomes) + + async def _embed_content(*, model, contents, config): + del model, config + outcome = remaining.pop(0) if len(remaining) > 1 else remaining[0] + if isinstance(outcome, BaseException): + raise outcome + if outcome is None: + return types.EmbedContentResponse( + embeddings=[embedding_of(text) for text in contents] + ) + return outcome + + return _embed_content + + +class SparkTestCase(unittest.TestCase): + """Base class holding a Spark session shared by every test module.""" + + @classmethod + def setUpClass(cls): + cls.spark = ( + SparkSession.builder.master("local[2]") + .appName("dataproc-ml-sql-tests") + .config("spark.sql.execution.arrow.maxRecordsPerBatch", "2") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + + +class TestReturnType(SparkTestCase): + + def test_fields_match_bigquery(self): + # AI.EMBED returns only the vector and the status. There is no + # full_response field, unlike AI.GENERATE. + self.assertEqual( + [field.name for field in ai_embed_module.RETURN_TYPE.fields], + ["result", "status"], + ) + + def test_the_result_is_an_array_of_doubles(self): + # BigQuery returns ARRAY, which is a double in Spark. + self.assertEqual( + ai_embed_module.RETURN_TYPE["result"].dataType, + ArrayType(DoubleType()), + ) + + +class TestValidation(SparkTestCase): + """Configuration mistakes are reported on the driver, before any job.""" + + def test_empty_endpoint_is_rejected(self): + with self.assertRaises(ValueError): + ai_embed_udf(endpoint=" ") + + def test_contents_in_model_params_is_rejected(self): + with self.assertRaises(ValueError): + ai_embed_udf(model_params={"contents": ["hi"]}) + + def test_unknown_task_type_is_rejected(self): + with self.assertRaises(ValueError): + ai_embed_udf(task_type="RETRIEVAL_EVERYTHING") + + def test_title_needs_a_document_task_type(self): + with self.assertRaises(ValueError): + ai_embed_udf(title="A title") + with self.assertRaises(ValueError): + ai_embed_udf(task_type="SEMANTIC_SIMILARITY", title="A title") + + def test_a_task_type_given_twice_is_rejected(self): + with self.assertRaises(ValueError): + ai_embed_udf( + task_type="CLUSTERING", + model_params={"taskType": "CLUSTERING"}, + ) + + def test_a_valid_configuration_is_accepted(self): + ai_embed_udf( + task_type="retrieval_document", + title="A title", + model_params={"outputDimensionality": 256}, + ) + + +class TestChunking(SparkTestCase): + """Requests are split to stay within what the endpoint accepts.""" + + def test_no_request_exceeds_the_documented_limit(self): + # The endpoint rejects more than 250 instances, and an Arrow batch + # holds 10,000 rows by default, so this is a correctness requirement. + sendable = [(index, "text") for index in range(600)] + sizes = [len(chunk) for chunk in ai_embed_module._chunks(sendable)] + self.assertEqual(sizes, [250, 250, 100]) + + def test_everything_fits_in_one_request_when_small(self): + sendable = [(index, "text") for index in range(3)] + self.assertEqual(len(ai_embed_module._chunks(sendable)), 1) + + def test_nothing_to_send_makes_no_request(self): + self.assertEqual(ai_embed_module._chunks([]), []) + + +class TestEmbedBatch(SparkTestCase): + """The per batch behavior, exercised without Spark for precision.""" + + def _run( + self, + embed_content, + contents, + *, + max_attempts=5, + max_instances=250, + ): + # The retry and batching settings are module constants. They are + # patched here so that the tests neither sleep nor depend on the + # shipped values. + with mock.patch.multiple( + ai_embed_module, + _MAX_ATTEMPTS=max_attempts, + _BASE_RETRY_DELAY_SECONDS=0.0, + _MAX_RETRY_DELAY_SECONDS=0.0, + _MAX_INSTANCES_PER_REQUEST=max_instances, + ): + return asyncio.run( + ai_embed_module._embed_batch( + embed_content=embed_content, + contents=contents, + endpoint="text-embedding-005", + config=None, + ) + ) + + def test_every_row_keeps_its_own_embedding(self): + rows = self._run(embeds_each_content(), ["a", "bb", "ccc"]) + self.assertEqual([row["result"][0] for row in rows], [1.0, 2.0, 3.0]) + self.assertEqual([row["status"] for row in rows], ["", "", ""]) + + def test_null_rows_skip_the_model(self): + recorded = [] + rows = self._run( + embeds_each_content(recorded), [None, "bb", float("nan")] + ) + self.assertEqual(recorded, [["bb"]]) + self.assertIsNone(rows[0]["result"]) + self.assertEqual(rows[0]["status"], "") + self.assertEqual(rows[1]["result"][0], 2.0) + self.assertIsNone(rows[2]["result"]) + self.assertEqual(rows[2]["status"], "") + + def test_empty_content_is_answered_without_a_request(self): + # The endpoint silently drops an empty instance from a request that + # carries others, which would shift every later row onto the wrong + # embedding. Empty values are therefore never sent. + recorded = [] + rows = self._run(embeds_each_content(recorded), ["a", "", "ccc"]) + self.assertEqual(recorded, [["a", "ccc"]]) + self.assertEqual(rows[0]["result"][0], 1.0) + self.assertIsNone(rows[1]["result"]) + self.assertEqual( + rows[1]["status"], ai_embed_module.EMPTY_CONTENT_STATUS + ) + self.assertEqual(rows[2]["result"][0], 3.0) + + def test_whitespace_is_real_content(self): + # Only a truly empty instance is rejected by the endpoint, so a + # whitespace string is sent as it is rather than trimmed away. + recorded = [] + rows = self._run(embeds_each_content(recorded), [" "]) + self.assertEqual(recorded, [[" "]]) + self.assertEqual(rows[0]["status"], "") + + def test_a_batch_is_split_into_several_requests(self): + recorded = [] + contents = [str(index) for index in range(5)] + rows = self._run( + embeds_each_content(recorded), contents, max_instances=2 + ) + self.assertEqual(sorted(len(call) for call in recorded), [1, 2, 2]) + self.assertEqual(len(rows), 5) + self.assertTrue(all(row["status"] == "" for row in rows)) + + def test_a_non_string_value_is_rendered_as_text(self): + # Spark SQL can call the registered function on any column, and no + # cast is inserted there. A whole number carried as a float, which is + # how Arrow delivers a nullable integer column, keeps its integer + # spelling so that one row does not depend on another being null. + recorded = [] + self._run(embeds_each_content(recorded), [42, 2.0, 3.5, True]) + self.assertEqual(recorded, [["42", "2", "3.5", "True"]]) + + def test_transient_errors_are_retried_then_succeed(self): + rows = self._run( + replies_with( + api_error(429, "RESOURCE_EXHAUSTED", "slow down"), + api_error(503, "UNAVAILABLE", "backend down"), + None, + ), + ["a"], + ) + self.assertEqual(rows[0]["result"][0], 1.0) + self.assertEqual(rows[0]["status"], "") + + def test_exhausted_retries_report_the_error(self): + rows = self._run( + replies_with(api_error(429, "RESOURCE_EXHAUSTED", "slow down")), + ["a"], + max_attempts=3, + ) + self.assertIsNone(rows[0]["result"]) + self.assertEqual( + rows[0]["status"], + "A retryable error occurred: RESOURCE_EXHAUSTED error from remote " + "service/endpoint.", + ) + + def test_a_failed_request_is_reported_for_all_of_its_rows(self): + # A request either succeeds or fails as a whole, so the failure is + # attributed to every row it carried, and only to those. + async def _fail_the_second_chunk(*, model, contents, config): + del model, config + if "c" in contents: + raise api_error(404, "NOT_FOUND", "Publisher Model not found") + return types.EmbedContentResponse( + embeddings=[embedding_of(text) for text in contents] + ) + + rows = self._run( + _fail_the_second_chunk, ["a", "b", "c", "d"], max_instances=2 + ) + self.assertEqual([row["status"] for row in rows[:2]], ["", ""]) + for row in rows[2:]: + self.assertIsNone(row["result"]) + self.assertEqual( + row["status"], "NOT_FOUND: Publisher Model not found" + ) + + def test_request_errors_are_not_retried(self): + attempts = [] + + async def _fail(*, model, contents, config): + del model, contents, config + attempts.append(1) + raise api_error(404, "NOT_FOUND", "Publisher Model was not found") + + self._run(_fail, ["a"]) + self.assertEqual(len(attempts), 1) + + def test_library_errors_are_not_hidden_in_status(self): + async def _bug(*, model, contents, config): + del model, contents, config + raise KeyError("a bug in the library") + + with self.assertRaises(KeyError): + self._run(_bug, ["a"]) + + def test_a_short_response_fails_rather_than_misaligning_rows(self): + # Rows are matched to embeddings by position, so a response of a + # different length would attach the wrong vector to a row. That is + # worse than a failed query. + async def _drop_one(*, model, contents, config): + del model, config + return types.EmbedContentResponse( + embeddings=[embedding_of(text) for text in contents[:-1]] + ) + + with self.assertRaises(RuntimeError): + self._run(_drop_one, ["a", "b"]) + + +class TestSparkExecution(SparkTestCase): + """The function as users see it, running through real Spark tasks.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.df = cls.spark.createDataFrame( + [("Bengaluru",), ("London",), ("Tokyo",), ("Paris",), ("Oslo",)], + ["city"], + ).repartition(2) + + def test_schema_and_dot_access(self): + result = self.df.withColumn( + "e", ai_embed(sf.col("city"), _embed_content=embeds_each_content()) + ) + self.assertEqual( + [field.name for field in result.schema["e"].dataType.fields], + ["result", "status"], + ) + + rows = result.select("city", "e.result", "e.status").collect() + self.assertEqual(len(rows), 5) + for row in rows: + self.assertEqual(row["result"][0], float(len(row["city"]))) + self.assertEqual(row["status"], "") + + def test_a_string_names_a_column(self): + rows = ( + self.df.withColumn( + "e", ai_embed("city", _embed_content=embeds_each_content()) + ) + .select("city", "e.result") + .collect() + ) + for row in rows: + self.assertEqual(row["result"][0], float(len(row["city"]))) + + def test_content_parts_are_concatenated(self): + rows = ( + self.df.withColumn( + "e", + ai_embed( + [sf.lit("City: "), sf.col("city")], + _embed_content=embeds_each_content(), + ), + ) + .select("city", "e.result") + .collect() + ) + for row in rows: + self.assertEqual( + row["result"][0], float(len("City: ") + len(row["city"])) + ) + + def test_null_and_empty_rows_produce_null_results(self): + df = self.spark.createDataFrame( + [("Oslo",), (None,), ("",)], "city: string" + ) + rows = ( + df.withColumn( + "e", + ai_embed(sf.col("city"), _embed_content=embeds_each_content()), + ) + .select("city", "e.result", "e.status") + .collect() + ) + by_city = {row["city"]: row for row in rows} + self.assertEqual(by_city["Oslo"]["result"][0], 4.0) + self.assertIsNone(by_city[None]["result"]) + self.assertEqual(by_city[None]["status"], "") + self.assertIsNone(by_city[""]["result"]) + self.assertEqual( + by_city[""]["status"], ai_embed_module.EMPTY_CONTENT_STATUS + ) + + def test_registered_for_spark_sql(self): + self.spark.udf.register( + "ai_embed", ai_embed_udf(_embed_content=embeds_each_content()) + ) + self.df.createOrReplaceTempView("cities") + + rows = self.spark.sql( + "SELECT city, ai_embed(city).result AS vector," + " ai_embed(city).status AS status FROM cities" + ).collect() + + self.assertEqual(len(rows), 5) + for row in rows: + self.assertEqual(row["vector"][0], float(len(row["city"]))) + self.assertEqual(row["status"], "") + + def test_a_non_string_column_works_in_spark_sql(self): + # Spark SQL applies the registered function to the column as declared, + # with no cast inserted, so the worker has to render the value. + self.spark.udf.register( + "ai_embed", ai_embed_udf(_embed_content=embeds_each_content()) + ) + numbers = self.spark.createDataFrame( + [(1,), (22,), (None,)], "n: bigint" + ) + numbers.createOrReplaceTempView("numbers") + + rows = self.spark.sql( + "SELECT n, ai_embed(n).result AS vector FROM numbers" + ).collect() + + vectors = {row["n"]: row["vector"] for row in rows} + self.assertEqual(vectors[1][0], 1.0) + self.assertEqual(vectors[22][0], 2.0) + self.assertIsNone(vectors[None]) + + def test_failed_rows_do_not_fail_the_query(self): + result = self.df.withColumn( + "e", + ai_embed( + sf.col("city"), + _embed_content=replies_with( + api_error(404, "NOT_FOUND", "Publisher Model not found") + ), + ), + ) + + rows = result.select("e.result", "e.status").collect() + self.assertEqual(len(rows), 5) + for row in rows: + self.assertIsNone(row["result"]) + self.assertEqual( + row["status"], "NOT_FOUND: Publisher Model not found" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_ai_generate.py b/tests/unit/sql/test_ai_generate.py new file mode 100644 index 0000000..52ad61c --- /dev/null +++ b/tests/unit/sql/test_ai_generate.py @@ -0,0 +1,566 @@ +# 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. + +import asyncio +import json +import unittest +from unittest import mock + +from google.genai import errors, types +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf +from pyspark.sql.types import ( + BooleanType, + LongType, + StringType, + StructField, + StructType, +) + +# pylint: disable=ungrouped-imports +from google.cloud.dataproc_ml.sql import _ai_generate as ai_generate_module +from google.cloud.dataproc_ml.sql import ai_generate, ai_generate_udf + +# pylint: enable=ungrouped-imports + + +def response_with_text(text, thought=None): + """Builds a model response containing a single answer part.""" + parts = [] + if thought is not None: + parts.append(types.Part(text=thought, thought=True)) + parts.append(types.Part(text=text)) + return types.GenerateContentResponse( + candidates=[ + types.Candidate(content=types.Content(role="model", parts=parts)) + ] + ) + + +def empty_response(): + """Builds a response with no usable text, as a filtered answer looks.""" + return types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(role="model", parts=[]), + finish_reason=types.FinishReason.SAFETY, + ) + ] + ) + + +def api_error(code, status, message): + """Builds an API error the way the client library reports one.""" + return errors.APIError( + code, {"error": {"code": code, "status": status, "message": message}} + ) + + +def replies_with(*outcomes): + """Returns a fake ``generate_content`` that replays the given outcomes. + + Each outcome is either a response to return or an exception to raise. The + returned coroutine function is defined here rather than at module level so + that it is serialized by value and can run inside a Spark worker. + + Args: + *outcomes: Responses or exceptions, consumed in order. The last one is + repeated once exhausted. + + Returns: + An awaitable callable with the same shape as the client method. + """ + remaining = list(outcomes) + + async def _generate_content(*, model, contents, config): + del model, contents, config + outcome = remaining.pop(0) if len(remaining) > 1 else remaining[0] + if isinstance(outcome, BaseException): + raise outcome + return outcome + + return _generate_content + + +def answers_the_prompt(): + """Returns a fake that echoes each prompt, so batching can be checked.""" + + async def _generate_content(*, model, contents, config): + del model, config + return response_with_text(f"answer to {contents}") + + return _generate_content + + +class SparkTestCase(unittest.TestCase): + """Base class holding a Spark session shared by every test module.""" + + @classmethod + def setUpClass(cls): + cls.spark = ( + SparkSession.builder.master("local[2]") + .appName("dataproc-ml-sql-tests") + .config("spark.sql.execution.arrow.maxRecordsPerBatch", "2") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + + +class TestPromptColumn(SparkTestCase): + + def test_a_string_names_a_column(self): + df = self.spark.createDataFrame([("hello",)], ["text"]) + column = ai_generate_module._to_prompt_column("text") + self.assertEqual(df.select(column).first()[0], "hello") + + def test_parts_are_concatenated_in_order(self): + df = self.spark.createDataFrame([("world",)], ["text"]) + column = ai_generate_module._to_prompt_column( + [sf.lit("hello "), sf.col("text"), sf.lit("!")] + ) + self.assertEqual(df.select(column).first()[0], "hello world!") + + def test_a_single_part_is_used_as_is(self): + df = self.spark.createDataFrame([("hello",)], ["text"]) + column = ai_generate_module._to_prompt_column([sf.col("text")]) + self.assertEqual(df.select(column).first()[0], "hello") + + def test_empty_prompt_is_rejected(self): + with self.assertRaises(ValueError): + ai_generate_module._to_prompt_column([]) + + def test_unsupported_prompt_is_rejected(self): + with self.assertRaises(ValueError): + ai_generate_module._to_prompt_column(42) + + def test_a_non_string_column_is_cast_to_string(self): + # Without the cast the worker would see an int, treat it as a null + # prompt and return nulls for the whole column without any error. + df = self.spark.createDataFrame([(7,)], ["n"]) + + for prompt in ("n", sf.col("n"), [sf.col("n")]): + with self.subTest(prompt=prompt): + column = ai_generate_module._to_prompt_column(prompt) + value = df.select(column).first()[0] + self.assertEqual(value, "7") + self.assertIsInstance(value, str) + + +class TestReturnType(SparkTestCase): + + def test_default_fields_match_bigquery(self): + struct = ai_generate_module._build_return_type(None) + self.assertEqual( + [field.name for field in struct.fields], + ["result", "full_response", "status"], + ) + self.assertIsInstance(struct["result"].dataType, StringType) + self.assertIsInstance(struct["status"].dataType, StringType) + + def test_output_schema_replaces_the_result_field(self): + struct = ai_generate_module._build_return_type( + StructType( + [ + StructField("sentiment", StringType()), + StructField("score", LongType()), + ] + ) + ) + self.assertEqual( + [field.name for field in struct.fields], + ["sentiment", "score", "full_response", "status"], + ) + + def test_reserved_field_names_are_rejected(self): + for reserved in ("full_response", "status"): + with self.subTest(field=reserved): + with self.assertRaises(ValueError): + ai_generate_module._build_return_type( + StructType([StructField(reserved, StringType())]) + ) + + def test_generated_fields_are_nullable(self): + # A failed row reports the failure in status and leaves the generated + # fields empty, so they must accept nulls. + struct = ai_generate_module._build_return_type( + StructType([StructField("a", StringType(), nullable=False)]) + ) + self.assertTrue(struct["a"].nullable) + + +class TestValidation(SparkTestCase): + """Configuration mistakes are reported on the driver, before any job.""" + + def test_empty_endpoint_is_rejected(self): + with self.assertRaises(ValueError): + ai_generate_udf(endpoint=" ") + + def test_invalid_model_params_are_rejected(self): + with self.assertRaises(ValueError): + ai_generate_udf(model_params={"contents": ["hi"]}) + + +class TestGenerateOneRow(SparkTestCase): + """The per row behavior, exercised without Spark for precision.""" + + def _run( + self, + generate_content, + prompt="a prompt", + *, + max_attempts=5, + **overrides, + ): + options = { + "generate_content": generate_content, + "prompt": prompt, + "endpoint": "gemini-3.6-flash", + "config": None, + "schema": None, + } + options.update(overrides) + # The retry settings are module constants. They are patched here so + # that the tests neither sleep nor depend on the shipped values. + with mock.patch.multiple( + ai_generate_module, + _MAX_ATTEMPTS=max_attempts, + _BASE_RETRY_DELAY_SECONDS=0.0, + _MAX_RETRY_DELAY_SECONDS=0.0, + ): + return asyncio.run(ai_generate_module._generate_one(**options)) + + def test_successful_row_has_empty_status(self): + row = self._run(replies_with(response_with_text("BLR"))) + self.assertEqual(row["result"], "BLR") + self.assertEqual(row["status"], "") + + def test_reasoning_parts_are_skipped(self): + row = self._run( + replies_with(response_with_text("BLR", thought="thinking...")) + ) + self.assertEqual(row["result"], "BLR") + + def test_response_without_text_produces_no_result(self): + # A filtered answer is not an error, so status stays empty. + row = self._run(replies_with(empty_response())) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "") + + def test_null_prompt_skips_the_model(self): + called = [] + + async def _never(*, contents, **_kwargs): + called.append(contents) + raise AssertionError("the model must not be called") + + for prompt in (None, float("nan")): + with self.subTest(prompt=prompt): + row = self._run(_never, prompt=prompt) + self.assertIsNone(row["result"]) + self.assertEqual(row["status"], "") + self.assertEqual(called, []) + + def test_a_non_string_prompt_is_rendered_as_text(self): + # Spark SQL can call the registered function on any column, and no + # cast is inserted there. Such a value must reach the model as text + # rather than be mistaken for a null. A whole number carried as a + # float, which is how Arrow delivers a nullable integer column, keeps + # its integer spelling. + cases = ((42, "42"), (2.0, "2"), (3.5, "3.5"), (True, "True")) + for prompt, expected in cases: + with self.subTest(prompt=prompt): + row = self._run(answers_the_prompt(), prompt=prompt) + self.assertEqual(row["result"], f"answer to {expected}") + + def test_transient_errors_are_retried_then_succeed(self): + row = self._run( + replies_with( + api_error(429, "RESOURCE_EXHAUSTED", "slow down"), + api_error(503, "UNAVAILABLE", "backend down"), + response_with_text("BLR"), + ) + ) + self.assertEqual(row["result"], "BLR") + self.assertEqual(row["status"], "") + + def test_exhausted_retries_report_the_error(self): + row = self._run( + replies_with(api_error(429, "RESOURCE_EXHAUSTED", "slow down")), + max_attempts=3, + ) + self.assertIsNone(row["result"]) + self.assertEqual( + row["status"], + "A retryable error occurred: RESOURCE_EXHAUSTED error from remote " + "service/endpoint.", + ) + + def test_request_errors_are_not_retried(self): + attempts = [] + + async def _fail(*, model, contents, config): + del model, contents, config + attempts.append(1) + raise api_error(404, "NOT_FOUND", "Publisher Model was not found") + + row = self._run(_fail) + self.assertEqual(len(attempts), 1) + self.assertEqual( + row["status"], "NOT_FOUND: Publisher Model was not found" + ) + + def test_library_errors_are_not_hidden_in_status(self): + async def _bug(*, model, contents, config): + del model, contents, config + raise KeyError("a bug in the library") + + with self.assertRaises(KeyError): + self._run(_bug) + + def test_structured_output_is_mapped_to_fields(self): + schema = StructType( + [ + StructField("score", LongType()), + StructField("sentiment", StringType()), + ] + ) + row = self._run( + replies_with( + response_with_text( + json.dumps({"sentiment": "positive", "score": 5}) + ) + ), + schema=schema, + ) + self.assertEqual(row["sentiment"], "positive") + self.assertEqual(row["score"], 5) + self.assertEqual(row["status"], "") + self.assertNotIn("result", row) + + def test_a_boolean_field_reads_the_answer(self): + schema = StructType([StructField("spam", BooleanType())]) + # A tuple rather than a dict: True and 1 are the same dictionary key. + cases = ( + (True, True), + (False, False), + # A model that answers with text rather than a JSON boolean. + ("true", True), + ("false", False), + ("False", False), + ("no", False), + (1, True), + (0, False), + # Nothing sensible to read, so the field is null. + ("perhaps", None), + ([], None), + ) + for answer, expected in cases: + with self.subTest(answer=answer): + row = self._run( + replies_with( + response_with_text(json.dumps({"spam": answer})) + ), + schema=schema, + ) + self.assertIs(row["spam"], expected) + + def test_unparsable_structured_output_is_reported(self): + schema = StructType([StructField("sentiment", StringType())]) + row = self._run( + replies_with(response_with_text("not json at all")), + schema=schema, + ) + self.assertIsNone(row["sentiment"]) + self.assertEqual( + row["status"], + "Failed to parse the response - please check the full_response " + "field.", + ) + self.assertIsNotNone(row["full_response"]) + + def test_missing_fields_become_null(self): + schema = StructType( + [ + StructField("sentiment", StringType()), + StructField("score", LongType()), + ] + ) + row = self._run( + replies_with( + response_with_text(json.dumps({"sentiment": "positive"})) + ), + schema=schema, + ) + self.assertEqual(row["sentiment"], "positive") + self.assertIsNone(row["score"]) + + +class TestSparkExecution(SparkTestCase): + """The function as users see it, running through real Spark tasks.""" + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.df = cls.spark.createDataFrame( + [("Bengaluru",), ("London",), ("Tokyo",), ("Paris",), ("Oslo",)], + ["city"], + ).repartition(2) + + def test_default_schema_and_dot_access(self): + result = self.df.withColumn( + "g", + ai_generate(sf.col("city"), _generate_content=answers_the_prompt()), + ) + self.assertEqual( + [field.name for field in result.schema["g"].dataType.fields], + ["result", "full_response", "status"], + ) + + rows = result.select("city", "g.result", "g.status").collect() + self.assertEqual(len(rows), 5) + for row in rows: + city = row["city"] + self.assertEqual(row["result"], f"answer to {city}") + self.assertEqual(row["status"], "") + + def test_prompt_parts_reach_the_model(self): + result = self.df.withColumn( + "g", + ai_generate( + [sf.lit("Airport code for "), sf.col("city")], + _generate_content=answers_the_prompt(), + ), + ) + rows = result.select("city", "g.result").collect() + for row in rows: + city = row["city"] + self.assertEqual( + row["result"], f"answer to Airport code for {city}" + ) + + def test_full_response_holds_the_whole_response(self): + result = self.df.withColumn( + "g", + ai_generate(sf.col("city"), _generate_content=answers_the_prompt()), + ) + # The column is VARIANT on Spark 4 and a JSON string otherwise, so it + # is read back as text either way. + as_text = result.selectExpr( + "CAST(to_json(g.full_response) AS STRING) AS payload" + if ai_generate_module._VARIANT_SUPPORTED + else "g.full_response AS payload" + ).first()["payload"] + payload = json.loads(as_text) + self.assertIn("candidates", payload) + + def test_null_rows_produce_null_results(self): + df = self.spark.createDataFrame([("Oslo",), (None,)], ["city"]) + rows = ( + df.withColumn( + "g", + ai_generate( + sf.col("city"), _generate_content=answers_the_prompt() + ), + ) + .select("city", "g.result", "g.status") + .collect() + ) + by_city = {row["city"]: row for row in rows} + self.assertEqual(by_city["Oslo"]["result"], "answer to Oslo") + self.assertIsNone(by_city[None]["result"]) + self.assertEqual(by_city[None]["status"], "") + + def test_structured_output_columns(self): + async def _structured(*, model, contents, config): + del model, config + return response_with_text( + json.dumps({"city": contents, "letters": len(contents)}) + ) + + result = self.df.withColumn( + "g", + ai_generate( + sf.col("city"), + output_schema="letters INT64, city STRING", + _generate_content=_structured, + ), + ) + self.assertEqual( + [field.name for field in result.schema["g"].dataType.fields], + ["city", "letters", "full_response", "status"], + ) + + rows = result.select("g.city", "g.letters").collect() + for row in rows: + self.assertEqual(row["letters"], len(row["city"])) + + def test_registered_for_spark_sql(self): + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=answers_the_prompt()), + ) + self.df.createOrReplaceTempView("cities") + + rows = self.spark.sql( + "SELECT city, ai_generate(city).result AS answer," + " ai_generate(city).status AS status FROM cities" + ).collect() + + self.assertEqual(len(rows), 5) + for row in rows: + city = row["city"] + self.assertEqual(row["answer"], f"answer to {city}") + self.assertEqual(row["status"], "") + + def test_a_non_string_column_works_in_spark_sql(self): + # Spark SQL applies the registered function to the column as declared, + # with no cast inserted, so the worker has to render the value. + self.spark.udf.register( + "ai_generate", + ai_generate_udf(_generate_content=answers_the_prompt()), + ) + numbers = self.spark.createDataFrame([(1,), (2,), (None,)], "n: bigint") + numbers.createOrReplaceTempView("numbers") + + rows = self.spark.sql( + "SELECT n, ai_generate(n).result AS answer FROM numbers" + ).collect() + + answers = {row["n"]: row["answer"] for row in rows} + self.assertEqual(answers[1], "answer to 1") + self.assertEqual(answers[2], "answer to 2") + self.assertIsNone(answers[None]) + + def test_failed_rows_do_not_fail_the_query(self): + result = self.df.withColumn( + "g", + ai_generate( + sf.col("city"), + _generate_content=replies_with( + api_error(404, "NOT_FOUND", "Publisher Model not found") + ), + ), + ) + + rows = result.select("g.result", "g.status").collect() + self.assertEqual(len(rows), 5) + for row in rows: + self.assertIsNone(row["result"]) + self.assertEqual( + row["status"], "NOT_FOUND: Publisher Model not found" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_ai_search.py b/tests/unit/sql/test_ai_search.py new file mode 100644 index 0000000..9bc90fd --- /dev/null +++ b/tests/unit/sql/test_ai_search.py @@ -0,0 +1,491 @@ +# 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. + +import math +import unittest + +from google.genai import errors, types +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf +from pyspark.sql.types import ArrayType, DoubleType, FloatType + +# pylint: disable=ungrouped-imports +from google.cloud.dataproc_ml.sql import _ai_search as ai_search_module +from google.cloud.dataproc_ml.sql import ai_embed, ai_search + +# pylint: enable=ungrouped-imports + +# A tiny stand-in for an embedding model. Each phrase has a hand-written +# vector, so the distance between any two rows can be worked out by hand and +# the expected ordering is obvious from reading the test. +VECTORS = { + "alpha": [1.0, 0.0, 0.0], + "almost alpha": [0.9, 0.1, 0.0], + "beta": [0.0, 1.0, 0.0], + "gamma": [0.0, 0.0, 1.0], +} + +# A phrase the fake model cannot embed, standing in for a row the real model +# fails on. +UNKNOWN = "not in the vocabulary" + + +def embeds_each_content(): + """Returns a fake ``embed_content`` answering from :data:`VECTORS`. + + Defined inside a function rather than at module level so that it is + serialized by value and can run inside a Spark worker. + """ + + async def _embed_content(*, model, contents, config): + del model, config + return types.EmbedContentResponse( + embeddings=[ + types.ContentEmbedding(values=VECTORS[text]) + if text in VECTORS + else types.ContentEmbedding(values=None) + for text in contents + ] + ) + + return _embed_content + + +def fails_with(error): + """Returns a fake ``embed_content`` that always raises.""" + + async def _embed_content(*, model, contents, config): + del model, contents, config + raise error + + return _embed_content + + +def api_error(code, status, message): + """Builds an API error the way the client library reports one.""" + return errors.APIError( + code, {"error": {"code": code, "status": status, "message": message}} + ) + + +class SparkTestCase(unittest.TestCase): + """Base class holding a Spark session shared by every test module.""" + + @classmethod + def setUpClass(cls): + cls.spark = ( + SparkSession.builder.master("local[2]") + .appName("dataproc-ml-sql-tests") + .config("spark.sql.execution.arrow.maxRecordsPerBatch", "2") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + + def tickets(self, *phrases): + """Builds a table of phrases the fake model knows, numbered from 1.""" + return self.spark.createDataFrame( + [(index, phrase) for index, phrase in enumerate(phrases, start=1)], + "id INT, body STRING", + ) + + def search(self, table, query="alpha", **kwargs): + """Searches with the fake model, so no credentials are needed.""" + kwargs.setdefault("_embed_content", embeds_each_content()) + return ai_search(table, "body", query, **kwargs) + + def ids(self, hits): + """Returns the ids of the rows found, in the order they came back.""" + return [row["base"]["id"] for row in hits.collect()] + + +class TestReturnedColumns(SparkTestCase): + """The output matches BigQuery's AI.SEARCH.""" + + def test_the_columns_are_base_and_distance(self): + hits = self.search(self.tickets("alpha")) + self.assertEqual(hits.columns, ["base", "distance"]) + + def test_base_holds_every_column_of_the_table(self): + hits = self.search(self.tickets("alpha")) + self.assertEqual( + hits.schema["base"].dataType.fieldNames(), ["id", "body"] + ) + + def test_the_distance_is_a_double(self): + # BigQuery returns FLOAT64, which is a double in Spark. + hits = self.search(self.tickets("alpha")) + self.assertEqual(hits.schema["distance"].dataType, DoubleType()) + + def test_an_embedding_column_is_not_returned(self): + # The column ai_search adds to hold each row's vector is an + # implementation detail and must not leak into the result. + hits = self.search(self.tickets("alpha")) + self.assertNotIn( + ai_search_module._VECTOR_COLUMN, + hits.schema["base"].dataType.fieldNames(), + ) + + +class TestOrdering(SparkTestCase): + """Rows come back nearest first, whichever distance is used.""" + + def setUp(self): + self.table = self.tickets("alpha", "beta", "gamma", "almost alpha") + + def test_the_nearest_row_comes_first(self): + self.assertEqual(self.ids(self.search(self.table))[:2], [1, 4]) + + def test_cosine_orders_the_same_way(self): + hits = self.search(self.table, distance_type="COSINE") + self.assertEqual(self.ids(hits)[:2], [1, 4]) + + def test_dot_product_orders_the_same_way(self): + # The inner product grows with similarity, so it is negated to keep + # one ordering for every distance. + hits = self.search(self.table, distance_type="DOT_PRODUCT") + self.assertEqual(self.ids(hits)[:2], [1, 4]) + + def test_the_distance_never_decreases(self): + distances = [ + row["distance"] for row in self.search(self.table).collect() + ] + self.assertEqual(distances, sorted(distances)) + + def test_a_lowercase_distance_type_is_accepted(self): + # A name written the way it reads in SQL should not be a mistake. + hits = self.search(self.table, distance_type="cosine") + self.assertEqual(self.ids(hits)[:2], [1, 4]) + + +class TestDistances(SparkTestCase): + """Each distance is the one it is named after.""" + + def distance_to_alpha(self, phrase, distance_type): + hits = self.search(self.tickets(phrase), distance_type=distance_type) + return hits.collect()[0]["distance"] + + def test_euclidean(self): + # |[0.9, 0.1, 0] - [1, 0, 0]| = sqrt(0.01 + 0.01). + self.assertAlmostEqual( + self.distance_to_alpha("almost alpha", "EUCLIDEAN"), + math.sqrt(0.02), + ) + + def test_cosine(self): + # 1 - 0.9 / (sqrt(0.82) * 1). + self.assertAlmostEqual( + self.distance_to_alpha("almost alpha", "COSINE"), + 1 - 0.9 / math.sqrt(0.82), + ) + + def test_dot_product_is_negated(self): + self.assertAlmostEqual( + self.distance_to_alpha("almost alpha", "DOT_PRODUCT"), -0.9 + ) + + def test_an_identical_row_is_at_no_distance(self): + self.assertAlmostEqual(self.distance_to_alpha("alpha", "EUCLIDEAN"), 0) + self.assertAlmostEqual(self.distance_to_alpha("alpha", "COSINE"), 0) + + +class TestTopK(SparkTestCase): + """Only the nearest rows are returned.""" + + def test_at_most_top_k_rows_come_back(self): + table = self.tickets("alpha", "beta", "gamma", "almost alpha") + self.assertEqual(self.search(table, top_k=2).count(), 2) + + def test_a_smaller_table_is_returned_whole(self): + table = self.tickets("alpha", "beta") + self.assertEqual(self.search(table, top_k=10).count(), 2) + + def test_the_default_matches_bigquery(self): + self.assertEqual(ai_search_module.DEFAULT_TOP_K, 10) + + +class TestExistingEmbeddings(SparkTestCase): + """The table's own embeddings are used when it has any.""" + + def embedded(self, *phrases, column="body_embedding"): + return self.tickets(*phrases).withColumn( + column, ai_embed("body", _embed_content=embeds_each_content()) + ) + + def test_a_column_named_for_the_one_it_embeds_is_found(self): + # BigQuery's own examples name the generated column this way. + hits = self.search(self.embedded("alpha", "beta"), top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_a_named_column_is_used(self): + table = self.embedded("alpha", "beta", column="vectors") + hits = self.search(table, embedding_column="vectors", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_a_plain_array_column_is_used(self): + table = self.embedded("alpha", "beta").withColumn( + "vectors", sf.col("body_embedding.result") + ) + hits = self.search(table, embedding_column="vectors", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_an_array_of_floats_is_accepted(self): + # An embedding read back from storage may be a float rather than a + # double, which is the same vector at less precision. + table = self.tickets("alpha", "beta").withColumn( + "vectors", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ) + .otherwise(sf.array(sf.lit(0.0), sf.lit(1.0), sf.lit(0.0))) + .cast(ArrayType(FloatType())), + ) + hits = self.search(table, embedding_column="vectors", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_the_named_column_wins_over_the_conventional_one(self): + # The named column embeds each row as if it were the next one along, + # so using it rather than body_embedding changes which row is nearest. + table = self.embedded("alpha", "beta").withColumn( + "shifted", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(0.0), sf.lit(1.0), sf.lit(0.0)), + ).otherwise(sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0))), + ) + self.assertEqual(self.ids(self.search(table, top_k=1)), [1]) + hits = self.search(table, embedding_column="shifted", top_k=1) + self.assertEqual(self.ids(hits), [2]) + + def test_the_embeddings_are_not_regenerated(self): + # The model is never reachable here, so the search can only succeed by + # using the embeddings the table already carries. + table = self.embedded("alpha", "beta") + hits = ai_search( + table, + "body", + "alpha", + top_k=1, + # The query still has to be embedded, and only the query. + _embed_content=embeds_each_content(), + ) + self.assertEqual(self.ids(hits), [1]) + + +class TestRowsWithoutAnEmbedding(SparkTestCase): + """A row that cannot be compared is skipped, not guessed at.""" + + def test_a_row_the_model_failed_on_is_skipped(self): + table = self.tickets("alpha", UNKNOWN) + self.assertEqual(self.ids(self.search(table, top_k=10)), [1]) + + def test_a_null_embedding_is_skipped(self): + table = self.tickets("alpha", "beta").withColumn( + "body_embedding", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ), + ) + hits = ai_search( + table, + "body", + "alpha", + _embed_content=embeds_each_content(), + ) + self.assertEqual(self.ids(hits), [1]) + + def test_an_embedding_of_a_different_width_is_skipped(self): + # Two vectors of different widths describe different spaces, so the + # row is dropped rather than compared on the values they share. + table = self.tickets("alpha", "beta").withColumn( + "body_embedding", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ).otherwise(sf.array(sf.lit(0.0), sf.lit(1.0))), + ) + hits = ai_search( + table, "body", "alpha", _embed_content=embeds_each_content() + ) + self.assertEqual(self.ids(hits), [1]) + + def test_a_zero_vector_is_skipped_by_cosine(self): + # A vector with no direction has no angle to the query. + table = self.tickets("alpha", "beta").withColumn( + "body_embedding", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ).otherwise(sf.array(sf.lit(0.0), sf.lit(0.0), sf.lit(0.0))), + ) + hits = ai_search( + table, + "body", + "alpha", + distance_type="COSINE", + _embed_content=embeds_each_content(), + ) + self.assertEqual(self.ids(hits), [1]) + + +class TestTheSearchTermFailsTheSearch(SparkTestCase): + """Unlike a row, a search term that cannot be embedded stops the search.""" + + def test_a_backend_failure_is_raised(self): + with self.assertRaises(ValueError) as caught: + ai_search( + self.tickets("alpha"), + "body", + "alpha", + _embed_content=fails_with( + api_error(403, "PERMISSION_DENIED", "Denied.") + ), + ) + self.assertIn("PERMISSION_DENIED", str(caught.exception)) + + def test_the_failure_names_the_search_term(self): + with self.assertRaises(ValueError) as caught: + ai_search( + self.tickets("alpha"), + "body", + "a missing term", + _embed_content=fails_with( + api_error(404, "NOT_FOUND", "No such model.") + ), + ) + self.assertIn("a missing term", str(caught.exception)) + + def test_an_empty_vector_is_raised(self): + with self.assertRaises(ValueError): + ai_search( + self.tickets("alpha"), + "body", + UNKNOWN, + _embed_content=embeds_each_content(), + ) + + def test_the_failure_happens_before_any_job(self): + # Nothing is scheduled, so the mistake surfaces at the call rather + # than at the first action taken on the result. + with self.assertRaises(ValueError): + ai_search( + self.tickets("alpha"), + "body", + "alpha", + _embed_content=fails_with( + api_error(400, "INVALID_ARGUMENT", "Bad request.") + ), + ) + + +class TestValidation(SparkTestCase): + """Mistakes are reported on the driver, before any job.""" + + def test_an_unknown_distance_type_is_rejected(self): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), distance_type="MANHATTAN") + + def test_top_k_must_be_positive(self): + for top_k in (0, -1): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), top_k=top_k) + + def test_top_k_must_be_a_whole_number(self): + for top_k in (1.5, "3", True): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), top_k=top_k) + + def test_an_empty_query_is_rejected(self): + for query in ("", " ", None): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), query=query) + + def test_a_missing_column_to_search_is_rejected(self): + with self.assertRaises(ValueError) as caught: + ai_search( + self.tickets("alpha"), + "subject", + "alpha", + _embed_content=embeds_each_content(), + ) + # The message lists what is there, so the typo is obvious. + self.assertIn("body", str(caught.exception)) + + def test_a_missing_embedding_column_is_rejected(self): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), embedding_column="vectors") + + def test_an_embedding_column_of_the_wrong_type_is_rejected(self): + with self.assertRaises(ValueError) as caught: + self.search(self.tickets("alpha"), embedding_column="body") + self.assertIn("array", str(caught.exception)) + + def test_a_struct_without_a_result_field_is_rejected(self): + table = self.tickets("alpha").withColumn( + "vectors", sf.struct(sf.lit(1.0).alias("value")) + ) + with self.assertRaises(ValueError) as caught: + self.search(table, embedding_column="vectors") + self.assertIn("result", str(caught.exception)) + + def test_a_base_table_of_the_wrong_type_is_rejected(self): + with self.assertRaises(ValueError): + self.search([("alpha",)]) + + def test_an_unknown_task_type_is_rejected(self): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), task_type="RETRIEVAL_EVERYTHING") + + def test_a_valid_configuration_is_accepted(self): + self.search( + self.tickets("alpha"), + task_type="retrieval_query", + top_k=1, + distance_type="dot_product", + ) + + +class TestNamingATable(SparkTestCase): + """A table can be named rather than passed, as in BigQuery.""" + + def test_a_registered_view_is_searched(self): + self.tickets("alpha", "beta").createOrReplaceTempView("ai_search_view") + hits = self.search("ai_search_view", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_an_empty_name_is_rejected(self): + with self.assertRaises(ValueError): + self.search(" ") + + +class TestTheModelIsCalledOncePerRow(SparkTestCase): + """Reading a row's vector twice must not embed it twice.""" + + def test_no_distance_evaluates_the_model_more_than_once(self): + # The cosine distance reads each vector twice, once for the dot + # product and once for its norm. Spark must not answer that by + # calling the model twice per row. + table = self.tickets("alpha", "beta") + for distance_type in ai_search_module.DISTANCE_TYPES: + with self.subTest(distance_type=distance_type): + hits = self.search(table, distance_type=distance_type) + plan = str( + hits._jdf.queryExecution().optimizedPlan() # pylint: disable=protected-access + ) + self.assertEqual(plan.count("_predict"), 1, plan) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_client.py b/tests/unit/sql/test_client.py new file mode 100644 index 0000000..b4bde32 --- /dev/null +++ b/tests/unit/sql/test_client.py @@ -0,0 +1,67 @@ +# 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. + +import unittest +from unittest import mock + +from google.cloud.dataproc_ml.sql import _client + + +class TestResolveProjectAndLocation(unittest.TestCase): + + def test_explicit_values_win(self): + env = { + "GOOGLE_CLOUD_PROJECT": "from-env", + "GOOGLE_CLOUD_LOCATION": "europe-west4", + } + with mock.patch.dict("os.environ", env, clear=True): + project, location = _client.resolve_project_and_location( + "explicit-project", "asia-northeast1" + ) + + self.assertEqual(project, "explicit-project") + self.assertEqual(location, "asia-northeast1") + + def test_project_falls_back_to_the_environment(self): + with mock.patch.dict( + "os.environ", {"GOOGLE_CLOUD_PROJECT": "from-env"}, clear=True + ): + project, _ = _client.resolve_project_and_location(None, None) + + self.assertEqual(project, "from-env") + + def test_location_defaults_to_global(self): + with mock.patch.dict("os.environ", {}, clear=True): + _, location = _client.resolve_project_and_location("p", None) + + self.assertEqual(location, "global") + self.assertEqual(_client.DEFAULT_LOCATION, "global") + + def test_location_ignores_the_environment(self): + # A Dataproc cluster sets these to the region it runs in. Honoring + # them would send the default model to a region that does not serve + # it, and the failure would be reported per row rather than raised. + env = { + "GOOGLE_CLOUD_PROJECT": "p", + "GOOGLE_CLOUD_LOCATION": "europe-west4", + "GOOGLE_CLOUD_REGION": "us-central1", + } + with mock.patch.dict("os.environ", env, clear=True): + _, location = _client.resolve_project_and_location(None, None) + + self.assertEqual(location, "global") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_model_params.py b/tests/unit/sql/test_model_params.py new file mode 100644 index 0000000..54ebd03 --- /dev/null +++ b/tests/unit/sql/test_model_params.py @@ -0,0 +1,211 @@ +# 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. + +import unittest + +from google.cloud.dataproc_ml.sql import _model_params + + +class TestNormalizeModelParams(unittest.TestCase): + + def test_no_params_produces_empty_config(self): + self.assertEqual(_model_params.normalize_model_params(None), {}) + + def test_generation_config_is_flattened(self): + flat = _model_params.normalize_model_params( + {"generation_config": {"temperature": 0.2, "top_p": 0.9}} + ) + self.assertEqual(flat, {"temperature": 0.2, "top_p": 0.9}) + + def test_camel_case_keys_are_accepted(self): + flat = _model_params.normalize_model_params( + {"generationConfig": {"maxOutputTokens": 128, "topK": 3}} + ) + self.assertEqual(flat, {"max_output_tokens": 128, "top_k": 3}) + + def test_nested_thinking_config_is_preserved(self): + flat = _model_params.normalize_model_params( + {"generation_config": {"thinking_config": {"thinking_budget": 0}}} + ) + self.assertEqual(flat, {"thinking_config": {"thinking_budget": 0}}) + + def test_top_level_request_fields_are_kept(self): + flat = _model_params.normalize_model_params( + { + "system_instruction": "Answer in one word.", + "generation_config": {"temperature": 0.0}, + } + ) + self.assertEqual( + flat, + { + "system_instruction": "Answer in one word.", + "temperature": 0.0, + }, + ) + + def test_contents_is_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params({"contents": ["hello"]}) + self.assertIn("should not contain contents", str(ctx.exception)) + + def test_invalid_parameters_are_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params( + {"generation_config": {"temperature": "warm"}} + ) + self.assertIn("model_params is invalid", str(ctx.exception)) + + def test_unknown_parameter_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_model_params({"not_a_real_field": 1}) + + def test_non_dictionary_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_model_params("temperature=0") + + def test_non_dictionary_generation_config_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_model_params({"generation_config": 1}) + + +class TestResponseSchemaHandling(unittest.TestCase): + + def test_output_schema_sets_json_mime_type(self): + schema = {"type": "OBJECT", "properties": {"a": {"type": "STRING"}}} + flat = _model_params.normalize_model_params(None, schema) + self.assertEqual(flat["response_mime_type"], "application/json") + self.assertEqual(flat["response_schema"], schema) + + def test_declaring_a_schema_twice_is_rejected(self): + schema = {"type": "OBJECT", "properties": {"a": {"type": "STRING"}}} + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_model_params( + {"generation_config": {"response_schema": schema}}, schema + ) + self.assertIn( + "can only be set either in the model_params or the output_schema", + str(ctx.exception), + ) + + def test_response_schema_alone_is_allowed(self): + schema = {"type": "OBJECT", "properties": {"a": {"type": "STRING"}}} + flat = _model_params.normalize_model_params( + {"generation_config": {"response_schema": schema}} + ) + self.assertEqual(flat["response_schema"], schema) + + +class TestBuildGenerationConfig(unittest.TestCase): + + def test_values_reach_the_generation_config(self): + config = _model_params.build_generation_config( + {"temperature": 0.25, "max_output_tokens": 64} + ) + self.assertEqual(config.temperature, 0.25) + self.assertEqual(config.max_output_tokens, 64) + + def test_thinking_config_is_understood(self): + config = _model_params.build_generation_config( + {"thinking_config": {"thinking_budget": 0}} + ) + self.assertEqual(config.thinking_config.thinking_budget, 0) + + +class TestNormalizeEmbedParams(unittest.TestCase): + + def test_no_params_produces_empty_config(self): + self.assertEqual(_model_params.normalize_embed_params(None), {}) + + def test_camel_case_keys_are_accepted(self): + flat = _model_params.normalize_embed_params( + {"outputDimensionality": 256} + ) + self.assertEqual(flat, {"output_dimensionality": 256}) + + def test_arguments_are_merged_into_the_params(self): + flat = _model_params.normalize_embed_params( + {"outputDimensionality": 256}, + "RETRIEVAL_DOCUMENT", + "A title", + ) + self.assertEqual( + flat, + { + "output_dimensionality": 256, + "task_type": "RETRIEVAL_DOCUMENT", + "title": "A title", + }, + ) + + def test_a_lowercase_task_type_is_normalized(self): + # A lowercase spelling reads naturally in SQL but the API only + # accepts the uppercase form. + flat = _model_params.normalize_embed_params(None, "clustering") + self.assertEqual(flat["task_type"], "CLUSTERING") + + def test_a_task_type_in_the_params_is_normalized_too(self): + flat = _model_params.normalize_embed_params({"taskType": "clustering"}) + self.assertEqual(flat["task_type"], "CLUSTERING") + + def test_unknown_task_type_is_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_embed_params(None, "RETRIEVAL_EVERYTHING") + self.assertIn("Unsupported task_type", str(ctx.exception)) + + def test_a_task_type_given_twice_is_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_embed_params( + {"taskType": "CLUSTERING"}, "CLUSTERING" + ) + self.assertIn("not both", str(ctx.exception)) + + def test_title_needs_a_document_task_type(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_embed_params( + None, "SEMANTIC_SIMILARITY", "A title" + ) + self.assertIn("RETRIEVAL_DOCUMENT", str(ctx.exception)) + + def test_contents_is_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.normalize_embed_params({"contents": ["hello"]}) + self.assertIn("should not contain contents", str(ctx.exception)) + + def test_unknown_parameter_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_embed_params({"not_a_real_field": 1}) + + def test_non_dictionary_is_rejected(self): + with self.assertRaises(ValueError): + _model_params.normalize_embed_params("outputDimensionality=256") + + +class TestBuildEmbedConfig(unittest.TestCase): + + def test_values_reach_the_embed_config(self): + config = _model_params.build_embed_config( + {"output_dimensionality": 256, "task_type": "CLUSTERING"} + ) + self.assertEqual(config.output_dimensionality, 256) + self.assertEqual(config.task_type, "CLUSTERING") + + def test_invalid_parameters_are_rejected(self): + with self.assertRaises(ValueError) as ctx: + _model_params.build_embed_config({"output_dimensionality": "many"}) + self.assertIn("model_params is invalid", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_output_schema.py b/tests/unit/sql/test_output_schema.py new file mode 100644 index 0000000..149fa90 --- /dev/null +++ b/tests/unit/sql/test_output_schema.py @@ -0,0 +1,241 @@ +# 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. + +import unittest + +from pyspark.sql import SparkSession +from pyspark.sql.types import ( + ArrayType, + BooleanType, + DoubleType, + LongType, + StringType, + StructField, + StructType, +) + +from google.cloud.dataproc_ml.sql import _output_schema + + +class OutputSchemaTestCase(unittest.TestCase): + """Parsing a schema string uses Spark's parser, so a session is needed.""" + + @classmethod + def setUpClass(cls): + cls.spark = ( + SparkSession.builder.master("local[2]") + .appName("dataproc-ml-sql-tests") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + + def field_names(self, struct): + return [field.name for field in struct.fields] + + def field_type(self, struct, name): + return struct[name].dataType + + +class TestParsing(OutputSchemaTestCase): + + def test_bigquery_type_names_are_accepted(self): + struct = _output_schema.parse_output_schema( + "a STRING, b INT64, c FLOAT64, d BOOL" + ) + self.assertIsInstance(self.field_type(struct, "a"), StringType) + self.assertIsInstance(self.field_type(struct, "b"), LongType) + self.assertIsInstance(self.field_type(struct, "c"), DoubleType) + self.assertIsInstance(self.field_type(struct, "d"), BooleanType) + + def test_spark_type_names_are_accepted(self): + struct = _output_schema.parse_output_schema( + "a STRING, b BIGINT, c DOUBLE, d BOOLEAN" + ) + self.assertIsInstance(self.field_type(struct, "b"), LongType) + self.assertIsInstance(self.field_type(struct, "c"), DoubleType) + self.assertIsInstance(self.field_type(struct, "d"), BooleanType) + + def test_a_field_named_after_a_type_keeps_its_name(self): + # Rewriting BigQuery type names must not touch the field names, or + # "int64 INT64" would silently declare a field called BIGINT. + struct = _output_schema.parse_output_schema( + "int64 INT64, bool BOOL, float64 FLOAT64" + ) + self.assertEqual(self.field_names(struct), ["bool", "float64", "int64"]) + self.assertIsInstance(self.field_type(struct, "int64"), LongType) + self.assertIsInstance(self.field_type(struct, "bool"), BooleanType) + self.assertIsInstance(self.field_type(struct, "float64"), DoubleType) + + def test_a_type_keeps_its_modifiers(self): + for schema in ( + "a INT64 NOT NULL", + "a INT64 COMMENT 'how many'", + "a INT64 NOT NULL, b BOOL", + ): + with self.subTest(schema=schema): + struct = _output_schema.parse_output_schema(schema) + self.assertIsInstance(self.field_type(struct, "a"), LongType) + + def test_arrays_and_structs_are_supported(self): + struct = _output_schema.parse_output_schema( + "tags ARRAY, location STRUCT" + ) + self.assertIsInstance(self.field_type(struct, "tags"), ArrayType) + nested = self.field_type(struct, "location") + self.assertEqual(self.field_names(nested), ["city", "zip"]) + self.assertIsInstance(nested["zip"].dataType, LongType) + + def test_a_struct_type_can_be_passed_directly(self): + struct = _output_schema.parse_output_schema( + StructType( + [ + StructField("b", StringType()), + StructField("a", LongType()), + ] + ) + ) + self.assertEqual(self.field_names(struct), ["a", "b"]) + + def test_unparsable_schema_is_rejected(self): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema("this is not a schema") + self.assertIn("Failed to parse", str(ctx.exception)) + + def test_empty_schema_is_rejected(self): + with self.assertRaises(ValueError): + _output_schema.parse_output_schema(StructType([])) + + def test_wrong_argument_type_is_rejected(self): + with self.assertRaises(ValueError): + _output_schema.parse_output_schema(42) + + +class TestFieldOrdering(OutputSchemaTestCase): + """BigQuery returns custom schema fields sorted by name.""" + + def test_fields_are_sorted_by_name(self): + struct = _output_schema.parse_output_schema( + "sentiment STRING, explanation STRING, score INT64" + ) + self.assertEqual( + self.field_names(struct), ["explanation", "score", "sentiment"] + ) + + def test_nested_struct_fields_are_sorted_too(self): + struct = _output_schema.parse_output_schema( + "location STRUCT" + ) + self.assertEqual( + self.field_names(self.field_type(struct, "location")), + ["city", "zip"], + ) + + def test_struct_inside_an_array_is_sorted(self): + struct = _output_schema.parse_output_schema( + "items ARRAY>" + ) + element = self.field_type(struct, "items").elementType + self.assertEqual(self.field_names(element), ["id", "name"]) + + +class TestValidation(OutputSchemaTestCase): + + def test_unsupported_types_are_rejected(self): + for declaration in ("d DATE", "t TIMESTAMP", "m MAP"): + with self.subTest(declaration=declaration): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema(declaration) + self.assertIn("Unsupported field type", str(ctx.exception)) + + def test_options_clause_is_reported_as_unsupported(self): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema( + "state STRING OPTIONS(description = 'the state')" + ) + self.assertIn("OPTIONS", str(ctx.exception)) + + def test_duplicate_field_names_are_rejected(self): + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema( + StructType( + [ + StructField("a", StringType()), + StructField("a", StringType()), + ] + ) + ) + self.assertIn("duplicate field names", str(ctx.exception)) + + def test_deep_nesting_is_rejected(self): + deep = "a STRUCT>>>" + with self.assertRaises(ValueError) as ctx: + _output_schema.parse_output_schema(deep) + self.assertIn("too many nested levels", str(ctx.exception)) + + def test_moderate_nesting_is_allowed(self): + struct = _output_schema.parse_output_schema( + "a STRUCT>" + ) + self.assertEqual(self.field_names(struct), ["a"]) + + +class TestResponseSchema(OutputSchemaTestCase): + + def test_scalars_map_to_gemini_types(self): + struct = _output_schema.parse_output_schema( + "a STRING, b INT64, c FLOAT64, d BOOL" + ) + schema = _output_schema.to_response_schema(struct) + self.assertEqual(schema["type"], "OBJECT") + self.assertEqual( + { + name: value["type"] + for name, value in schema["properties"].items() + }, + {"a": "STRING", "b": "INTEGER", "c": "NUMBER", "d": "BOOLEAN"}, + ) + + def test_arrays_declare_their_item_type(self): + struct = _output_schema.parse_output_schema("tags ARRAY") + schema = _output_schema.to_response_schema(struct) + self.assertEqual( + schema["properties"]["tags"], + {"type": "ARRAY", "items": {"type": "STRING"}}, + ) + + def test_nested_structs_become_nested_objects(self): + struct = _output_schema.parse_output_schema( + "location STRUCT" + ) + schema = _output_schema.to_response_schema(struct) + self.assertEqual( + schema["properties"]["location"], + {"type": "OBJECT", "properties": {"city": {"type": "STRING"}}}, + ) + + def test_non_nullable_fields_are_required(self): + struct = _output_schema.parse_output_schema( + StructType( + [ + StructField("a", StringType(), nullable=False), + StructField("b", StringType(), nullable=True), + ] + ) + ) + schema = _output_schema.to_response_schema(struct) + self.assertEqual(schema["required"], ["a"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_retry.py b/tests/unit/sql/test_retry.py new file mode 100644 index 0000000..7f11877 --- /dev/null +++ b/tests/unit/sql/test_retry.py @@ -0,0 +1,149 @@ +# 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. + +import random +import unittest + +import httpx +from google.genai import errors + +from google.cloud.dataproc_ml.sql import _retry + + +def _api_error(code, status, message): + """Builds an API error the way the client library reports one.""" + return errors.APIError( + code, {"error": {"code": code, "status": status, "message": message}} + ) + + +class TestErrorClassification(unittest.TestCase): + + def test_rate_limiting_and_transient_errors_are_retryable(self): + for code in (408, 429, 500, 503, 504): + with self.subTest(code=code): + error = _api_error(code, "UNAVAILABLE", "try again") + self.assertTrue(_retry.is_retryable(error)) + + def test_request_errors_are_not_retryable(self): + for code in (400, 401, 403, 404): + with self.subTest(code=code): + error = _api_error(code, "INVALID_ARGUMENT", "bad request") + self.assertFalse(_retry.is_retryable(error)) + + def test_transport_errors_are_retryable(self): + self.assertTrue(_retry.is_retryable(httpx.ConnectError("no route"))) + self.assertTrue(_retry.is_retryable(httpx.ReadTimeout("slow"))) + + def test_unrelated_exceptions_are_not_reportable(self): + # A bug in the library must fail the task rather than be hidden in the + # status column of an otherwise successful query. + self.assertFalse(_retry.is_reportable(KeyError("bug"))) + self.assertFalse(_retry.is_retryable(KeyError("bug"))) + + def test_backend_errors_are_reportable(self): + self.assertTrue( + _retry.is_reportable(_api_error(404, "NOT_FOUND", "no model")) + ) + self.assertTrue(_retry.is_reportable(httpx.ConnectError("no route"))) + + +class TestStatusMessages(unittest.TestCase): + """The wording matches BigQuery's AI functions.""" + + def test_request_errors_are_surfaced_verbatim(self): + error = _api_error( + 400, + "INVALID_ARGUMENT", + "HTTP links are not supported for requests restricted by VPCSC.", + ) + self.assertEqual( + _retry.to_status_message(error), + "INVALID_ARGUMENT: HTTP links are not supported for requests " + "restricted by VPCSC.", + ) + + def test_not_found_is_surfaced_verbatim(self): + error = _api_error(404, "NOT_FOUND", "Publisher Model was not found") + self.assertEqual( + _retry.to_status_message(error), + "NOT_FOUND: Publisher Model was not found", + ) + + def test_quota_messages_are_passed_through(self): + message = ( + "Quota exceeded for metric aiplatform.googleapis.com/" + "generate_content_requests_per_minute_per_project" + ) + error = _api_error(429, "RESOURCE_EXHAUSTED", message) + self.assertEqual(_retry.to_status_message(error), message) + + def test_other_errors_are_summarised_as_retryable(self): + error = _api_error(429, "RESOURCE_EXHAUSTED", "too many requests") + self.assertEqual( + _retry.to_status_message(error), + "A retryable error occurred: RESOURCE_EXHAUSTED error from remote " + "service/endpoint.", + ) + + def test_server_errors_are_summarised_as_retryable(self): + error = _api_error(503, "UNAVAILABLE", "backend down") + self.assertEqual( + _retry.to_status_message(error), + "A retryable error occurred: UNAVAILABLE error from remote " + "service/endpoint.", + ) + + def test_transport_errors_report_unavailable(self): + self.assertEqual( + _retry.to_status_message(httpx.ConnectError("no route")), + "A retryable error occurred: UNAVAILABLE error from remote " + "service/endpoint.", + ) + + def test_status_is_never_empty(self): + error = _api_error(418, "", "") + self.assertTrue(_retry.to_status_message(error)) + + +class TestFullJitterBackoff(unittest.TestCase): + + def test_delay_grows_exponentially_and_is_fully_randomized(self): + rng = random.Random(7) + for attempt in range(5): + with self.subTest(attempt=attempt): + ceiling = min(60.0, 1.0 * (2**attempt)) + samples = [ + _retry.full_jitter_delay(attempt, 1.0, 60.0, rng) + for _ in range(200) + ] + self.assertTrue(all(0.0 <= s <= ceiling for s in samples)) + # Full jitter draws from the whole interval, so samples should + # reach both ends rather than cluster near the ceiling. + self.assertLess(min(samples), ceiling * 0.2) + self.assertGreater(max(samples), ceiling * 0.8) + + def test_delay_is_capped(self): + rng = random.Random(1) + samples = [ + _retry.full_jitter_delay(20, 1.0, 30.0, rng) for _ in range(100) + ] + self.assertTrue(all(0.0 <= s <= 30.0 for s in samples)) + + def test_zero_base_delay_produces_no_wait(self): + self.assertEqual(_retry.full_jitter_delay(3, 0.0, 60.0), 0.0) + + +if __name__ == "__main__": + unittest.main()