diff --git a/README.md b/README.md index 2a5cab8..769d96c 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,122 @@ 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() +``` + ### 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..4c10a8c --- /dev/null +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -0,0 +1,54 @@ +# 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) +""" + +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 + +__all__ = ( + "ai_embed", + "ai_embed_udf", + "ai_generate", + "ai_generate_udf", +) 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..8699686 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_embed.py @@ -0,0 +1,413 @@ +# 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 + +_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) + + +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/_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/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_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()