From 16b48897f3264f546db0fa3a4ca7b27f582dd110 Mon Sep 17 00:00:00 2001 From: Saurabh Maurya Date: Fri, 11 Sep 2026 15:08:52 +0000 Subject: [PATCH 1/4] feat(sql): add ai_generate, a Spark column function for Gemini Introduces google.cloud.dataproc_ml.sql.ai_generate, the Spark counterpart of BigQuery's AI.GENERATE. It calls a Gemini model once per row and returns a struct of result, full_response and status, so queries can move between the two engines with little change. The function is a pandas UDF, which means it works both as a DataFrame column expression and, through ai_generate_udf(), as a function registered with spark.udf.register for use from Spark SQL. The argument list is deliberately BigQuery's: prompt, endpoint, model_params and output_schema, with project and location standing in for connection_id. Concurrency and retry timing are implemented but kept as internal constants, as they are in BigQuery, rather than exposed as arguments that would be hard to withdraw later. Notable behaviour, all matching BigQuery: * A row that fails never fails the stage. Its error is reported in status and result is NULL. status is empty on success. * output_schema replaces result with one field per declared field, sorted by name. BigQuery type names (INT64, FLOAT64, BOOL) and their Spark spellings are both accepted. * full_response holds the complete model response, typed as VARIANT on Spark 4 and as a JSON string on older versions. Two things differ from BigQuery by necessity. There is no connection_id, because a BigQuery connection's service account cannot be impersonated from outside BigQuery; calls use Application Default Credentials instead. And the default model is gemini-3.6-flash, which is served only from the global endpoint, so the default location is global rather than a region. Tested with unit tests that exercise the full Spark path against an injected generate_content, and with integration tests that call Vertex AI for real. --- README.md | 55 ++ docs/api.rst | 4 + google/cloud/dataproc_ml/sql/__init__.py | 44 ++ google/cloud/dataproc_ml/sql/_ai_generate.py | 519 ++++++++++++++++++ google/cloud/dataproc_ml/sql/_client.py | 190 +++++++ google/cloud/dataproc_ml/sql/_model_params.py | 130 +++++ .../cloud/dataproc_ml/sql/_output_schema.py | 228 ++++++++ google/cloud/dataproc_ml/sql/_retry.py | 187 +++++++ pyproject.toml | 1 + tests/integration/sql/__init__.py | 13 + tests/integration/sql/test_ai_generate.py | 196 +++++++ tests/unit/sql/__init__.py | 13 + tests/unit/sql/test_ai_generate.py | 502 +++++++++++++++++ tests/unit/sql/test_client.py | 67 +++ tests/unit/sql/test_model_params.py | 127 +++++ tests/unit/sql/test_output_schema.py | 220 ++++++++ tests/unit/sql/test_retry.py | 149 +++++ 17 files changed, 2645 insertions(+) create mode 100644 google/cloud/dataproc_ml/sql/__init__.py create mode 100644 google/cloud/dataproc_ml/sql/_ai_generate.py create mode 100644 google/cloud/dataproc_ml/sql/_client.py create mode 100644 google/cloud/dataproc_ml/sql/_model_params.py create mode 100644 google/cloud/dataproc_ml/sql/_output_schema.py create mode 100644 google/cloud/dataproc_ml/sql/_retry.py create mode 100644 tests/integration/sql/__init__.py create mode 100644 tests/integration/sql/test_ai_generate.py create mode 100644 tests/unit/sql/__init__.py create mode 100644 tests/unit/sql/test_ai_generate.py create mode 100644 tests/unit/sql/test_client.py create mode 100644 tests/unit/sql/test_model_params.py create mode 100644 tests/unit/sql/test_output_schema.py create mode 100644 tests/unit/sql/test_retry.py diff --git a/README.md b/README.md index 2a5cab8..fdeb934 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,61 @@ pip install dataproc-ml Here are a couple of examples demonstrating how to use the handlers for distributed inference on a Spark DataFrame. +### AI Functions: `ai_generate` + +> **Note:** `ai_generate` makes API calls to Vertex AI, which will incur costs. +> Please review the [Vertex AI Generative +> AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing). + +`ai_generate` is a Spark column function that calls a Gemini model on every +row, mirroring the semantics of BigQuery's +[`AI.GENERATE`](https://cloud.google.com/bigquery/docs/reference/standard-sql/AI-functions#ai.generate). +It returns a `STRUCT` with a `result`, the raw `full_response`, and a `status` +that is empty on success. + +```python +from pyspark.sql import SparkSession, functions as F +from google.cloud.dataproc_ml.sql import ai_generate + +spark = SparkSession.builder.getOrCreate() + +df = spark.createDataFrame([("Paris",), ("Tokyo",)], ["city"]) + +result_df = df.withColumn( + "generated", + ai_generate(F.concat(F.lit("Airport code for "), F.col("city"))), +).select("city", "generated.result", "generated.status") + +result_df.show() +# +-----+------+------+ +# | city|result|status| +# +-----+------+------+ +# |Paris| CDG| | +# |Tokyo| HND| | +# +-----+------+------+ +``` + +Pass `output_schema` to get typed columns back instead of a single string: + +```python +df.withColumn( + "review", + ai_generate( + F.col("text"), + output_schema="sentiment STRING, score FLOAT64", + ), +) +``` + +The same function can be registered for use from Spark SQL: + +```python +from google.cloud.dataproc_ml.sql import ai_generate_udf + +spark.udf.register("ai_generate", ai_generate_udf()) +spark.sql("SELECT ai_generate('Airport code for Paris').result").show() +``` + ### Generative AI (Gemini) Model Inference > **Note:** Using the `GenAiModelHandler` involves making API calls to 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..98b7e6a --- /dev/null +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -0,0 +1,44 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Declarative AI functions for Spark DataFrames and Spark SQL. + +These functions apply Google's Gemini models to a DataFrame column without any +of the usual plumbing: no client setup, no retry loops and no response +parsing. They mirror the BigQuery ``AI.*`` functions, so a query written for +one engine reads the same on the other. + +Example: + >>> from pyspark.sql.functions import col + >>> from google.cloud.dataproc_ml.sql import ai_generate + >>> + >>> summarized = df.withColumn( + ... "summary", ai_generate(col("feedback")).result + ... ) + +The same function can be registered for use from Spark SQL: + + >>> from google.cloud.dataproc_ml.sql import ai_generate_udf + >>> + >>> spark.udf.register("ai_generate", ai_generate_udf()) + >>> spark.sql("SELECT ai_generate(feedback).result FROM feedback_table") +""" + +from ._ai_generate import ai_generate +from ._ai_generate import ai_generate_udf + +__all__ = ( + "ai_generate", + "ai_generate_udf", +) 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..6a59139 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -0,0 +1,519 @@ +# 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, Sequence, Union + +import pandas as pd +from pyspark.sql import Column +from pyspark.sql import functions as sf +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 _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" + +PromptArg = Union[Column, str, Sequence[Union[Column, str]]] + + +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. + """ + if isinstance(prompt, Column): + return prompt.cast(StringType()) + if isinstance(prompt, str): + return sf.col(prompt).cast(StringType()) + if isinstance(prompt, Sequence): + parts = list(prompt) + if not parts: + raise ValueError("prompt 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 prompt + # are combined in the order they are given. + return sf.concat(*[column.cast(StringType()) for column in columns]) + raise ValueError( + "prompt must be a Column, a column name, or a sequence of those, got " + f"{type(prompt).__name__}." + ) + + +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="") + + for attempt in range(_MAX_ATTEMPTS): + try: + response = await generate_content( + model=endpoint, contents=prompt, 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`` or, in a numeric batch, as NaN. + """ + if prompt is None: + return True + if isinstance(prompt, float) and math.isnan(prompt): + return True + return not isinstance(prompt, str) + + +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 bool(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 diff --git a/google/cloud/dataproc_ml/sql/_client.py b/google/cloud/dataproc_ml/sql/_client.py new file mode 100644 index 0000000..7b1592b --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_client.py @@ -0,0 +1,190 @@ +# 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 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/_model_params.py b/google/cloud/dataproc_ml/sql/_model_params.py new file mode 100644 index 0000000..8ad8470 --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_model_params.py @@ -0,0 +1,130 @@ +# 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 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..ba68a1f --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_output_schema.py @@ -0,0 +1,228 @@ +# 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", +} + +# 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 bigquery_name, spark_name in _TYPE_ALIASES.items(): + normalized = re.sub( + rf"\b{bigquery_name}\b", + spark_name, + normalized, + flags=re.IGNORECASE, + ) + + 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_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_generate.py b/tests/unit/sql/test_ai_generate.py new file mode 100644 index 0000000..bb940d6 --- /dev/null +++ b/tests/unit/sql/test_ai_generate.py @@ -0,0 +1,502 @@ +# 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 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_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_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_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..27fdeaa --- /dev/null +++ b/tests/unit/sql/test_model_params.py @@ -0,0 +1,127 @@ +# 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) + + +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..f6243c7 --- /dev/null +++ b/tests/unit/sql/test_output_schema.py @@ -0,0 +1,220 @@ +# 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_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() From a20e89c5bf5d3977c8e448c0f5855d104825c916 Mon Sep 17 00:00:00 2001 From: Saurabh Maurya Date: Fri, 11 Sep 2026 16:34:43 +0000 Subject: [PATCH 2/4] fix(sql): correct type and prompt handling raised in review Three fixes from the review bot on #8, each with a regression test: - Only rewrite a BigQuery type name where a type may appear. Matching the bare name also rewrote a field named after a type, so "int64 INT64" parsed as a field actually called BIGINT. The narrower match still accepts the NOT NULL and COMMENT modifiers that Spark allows after a type. - Treat only nulls as a missing prompt. Any other value is now rendered as text, so a non-string column no longer arrives as a whole column of nulls when the function is registered and called from Spark SQL, where no cast to STRING is inserted. A whole number keeps its integer spelling: Arrow delivers a nullable BIGINT column as floats, so otherwise a row's prompt would depend on whether an unrelated row in the same batch happened to be null. - Read a boolean field explicitly rather than through Python truthiness, which recorded the answer "false" as True. --- google/cloud/dataproc_ml/sql/_ai_generate.py | 65 ++++++++++++++++-- .../cloud/dataproc_ml/sql/_output_schema.py | 27 ++++++-- tests/unit/sql/test_ai_generate.py | 66 ++++++++++++++++++- tests/unit/sql/test_output_schema.py | 21 ++++++ 4 files changed, 164 insertions(+), 15 deletions(-) diff --git a/google/cloud/dataproc_ml/sql/_ai_generate.py b/google/cloud/dataproc_ml/sql/_ai_generate.py index 6a59139..202a799 100644 --- a/google/cloud/dataproc_ml/sql/_ai_generate.py +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -69,6 +69,12 @@ _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 = Union[Column, str, Sequence[Union[Column, str]]] @@ -335,10 +341,17 @@ async def _generate_one( # 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=prompt, config=config + model=endpoint, contents=text, config=config ) return _to_row(response, schema) except Exception as e: # pylint: disable=broad-except @@ -376,13 +389,28 @@ async def _generate_one( def _is_missing(prompt: Any) -> bool: """Reports whether a prompt is null. - Spark nulls arrive as ``None`` or, in a numeric batch, as NaN. + 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: + if prompt is None or prompt is pd.NA: return True - if isinstance(prompt, float) and math.isnan(prompt): - return True - return not isinstance(prompt, str) + 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]: @@ -498,7 +526,7 @@ def _coerce(value: Any, data_type: DataType) -> Any: if isinstance(data_type, StringType): return value if isinstance(value, str) else json.dumps(value) if isinstance(data_type, BooleanType): - return bool(value) + return _to_boolean(value) if isinstance(data_type, (LongType, IntegerType)): return int(value) if isinstance(data_type, (DoubleType, FloatType)): @@ -517,3 +545,26 @@ def _coerce(value: Any, data_type: DataType) -> Any: 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/_output_schema.py b/google/cloud/dataproc_ml/sql/_output_schema.py index ba68a1f..d6c8f64 100644 --- a/google/cloud/dataproc_ml/sql/_output_schema.py +++ b/google/cloud/dataproc_ml/sql/_output_schema.py @@ -43,6 +43,24 @@ "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) @@ -105,13 +123,8 @@ def _parse_ddl(output_schema: str) -> StructType: ) normalized = output_schema - for bigquery_name, spark_name in _TYPE_ALIASES.items(): - normalized = re.sub( - rf"\b{bigquery_name}\b", - spark_name, - normalized, - flags=re.IGNORECASE, - ) + for pattern, spark_name in _ALIAS_PATTERNS: + normalized = pattern.sub(spark_name, normalized) try: parsed = StructType.fromDDL(normalized) diff --git a/tests/unit/sql/test_ai_generate.py b/tests/unit/sql/test_ai_generate.py index bb940d6..52ad61c 100644 --- a/tests/unit/sql/test_ai_generate.py +++ b/tests/unit/sql/test_ai_generate.py @@ -20,7 +20,13 @@ from google.genai import errors, types from pyspark.sql import SparkSession from pyspark.sql import functions as sf -from pyspark.sql.types import LongType, StringType, StructField, StructType +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 @@ -265,6 +271,18 @@ async def _never(*, contents, **_kwargs): 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( @@ -330,6 +348,33 @@ def test_structured_output_is_mapped_to_fields(self): 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( @@ -478,6 +523,25 @@ def test_registered_for_spark_sql(self): 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", diff --git a/tests/unit/sql/test_output_schema.py b/tests/unit/sql/test_output_schema.py index f6243c7..149fa90 100644 --- a/tests/unit/sql/test_output_schema.py +++ b/tests/unit/sql/test_output_schema.py @@ -66,6 +66,27 @@ def test_spark_type_names_are_accepted(self): 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" From e5b819a299bcab3b17c75fb537f269fcd7e3e50c Mon Sep 17 00:00:00 2001 From: Grace Qi Date: Tue, 15 Sep 2026 03:24:04 +0000 Subject: [PATCH 3/4] feat(sql): add ai_embed, the Spark counterpart of BigQuery AI.EMBED Adds `ai_embed` and `ai_embed_udf`, which turn a text column into dense embedding vectors and return the same `STRUCT, status STRING>` that BigQuery's `AI.EMBED` does, so a query can be moved between the two engines with little change. Authentication follows `ai_generate`: Application Default Credentials, with `project` and `location` in place of a BigQuery connection. Unlike `ai_generate`, which calls the model once per row, the rows of an Arrow batch are embedded together. The endpoint accepts at most 250 instances per request while a batch holds 10,000 rows by default, so the splitting is a correctness requirement rather than a tuning knob. A request either succeeds or fails as a whole, so a terminal failure is recorded against every row it carried, which keeps the per-row status reporting. Empty content is answered locally and never sent. The endpoint rejects a request whose only instance is empty and, worse, silently drops an empty instance from a request carrying others, which would shift every later row onto its neighbour's vector. For the same reason a response whose length does not match the request fails the task instead of being reported in `status`: a wrong embedding is worse than a failed query. `_to_prompt_column` moves to a shared `_columns.to_string_column` so both functions build their input column the same way. --- README.md | 61 +++ google/cloud/dataproc_ml/sql/__init__.py | 10 + google/cloud/dataproc_ml/sql/_ai_embed.py | 413 +++++++++++++++ google/cloud/dataproc_ml/sql/_ai_generate.py | 27 +- google/cloud/dataproc_ml/sql/_client.py | 17 + google/cloud/dataproc_ml/sql/_columns.py | 70 +++ google/cloud/dataproc_ml/sql/_model_params.py | 131 +++++ tests/integration/sql/test_ai_embed.py | 201 ++++++++ tests/unit/sql/test_ai_embed.py | 484 ++++++++++++++++++ tests/unit/sql/test_model_params.py | 84 +++ 10 files changed, 1475 insertions(+), 23 deletions(-) create mode 100644 google/cloud/dataproc_ml/sql/_ai_embed.py create mode 100644 google/cloud/dataproc_ml/sql/_columns.py create mode 100644 tests/integration/sql/test_ai_embed.py create mode 100644 tests/unit/sql/test_ai_embed.py diff --git a/README.md b/README.md index fdeb934..769d96c 100644 --- a/README.md +++ b/README.md @@ -79,6 +79,67 @@ 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/google/cloud/dataproc_ml/sql/__init__.py b/google/cloud/dataproc_ml/sql/__init__.py index 98b7e6a..4c10a8c 100644 --- a/google/cloud/dataproc_ml/sql/__init__.py +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -33,12 +33,22 @@ >>> >>> 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 index 202a799..ef80a50 100644 --- a/google/cloud/dataproc_ml/sql/_ai_generate.py +++ b/google/cloud/dataproc_ml/sql/_ai_generate.py @@ -18,11 +18,10 @@ import json import logging import math -from typing import Any, Dict, Iterator, List, Optional, Sequence, Union +from typing import Any, Dict, Iterator, List, Optional, Union import pandas as pd from pyspark.sql import Column -from pyspark.sql import functions as sf from pyspark.sql.pandas.functions import pandas_udf from pyspark.sql.types import ( ArrayType, @@ -38,6 +37,7 @@ ) 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 @@ -75,7 +75,7 @@ _TRUE_WORDS = frozenset({"true", "t", "yes", "y", "1"}) _FALSE_WORDS = frozenset({"false", "f", "no", "n", "0"}) -PromptArg = Union[Column, str, Sequence[Union[Column, str]]] +PromptArg = _columns.ColumnArg def ai_generate( @@ -248,26 +248,7 @@ def _to_prompt_column(prompt: PromptArg) -> Column: 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. """ - if isinstance(prompt, Column): - return prompt.cast(StringType()) - if isinstance(prompt, str): - return sf.col(prompt).cast(StringType()) - if isinstance(prompt, Sequence): - parts = list(prompt) - if not parts: - raise ValueError("prompt 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 prompt - # are combined in the order they are given. - return sf.concat(*[column.cast(StringType()) for column in columns]) - raise ValueError( - "prompt must be a Column, a column name, or a sequence of those, got " - f"{type(prompt).__name__}." - ) + return _columns.to_string_column(prompt, argument_name="prompt") def _build_return_type(parsed_schema: Optional[StructType]) -> StructType: diff --git a/google/cloud/dataproc_ml/sql/_client.py b/google/cloud/dataproc_ml/sql/_client.py index 7b1592b..427d811 100644 --- a/google/cloud/dataproc_ml/sql/_client.py +++ b/google/cloud/dataproc_ml/sql/_client.py @@ -182,6 +182,23 @@ def get_generate_content(project: Optional[str], location: Optional[str]): 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 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 index 8ad8470..cdcbec0 100644 --- a/google/cloud/dataproc_ml/sql/_model_params.py +++ b/google/cloud/dataproc_ml/sql/_model_params.py @@ -39,6 +39,26 @@ "output_schema, not both." ) +_EMBED_CONTENTS_ERROR = ( + "The model_params is invalid and should not contain contents. The content " + "argument supplies the contents of the request." +) + +#: The ``task_type`` values BigQuery's ``AI.EMBED`` accepts. Validated on the +#: driver so that a typo fails the query rather than every row. +TASK_TYPES = frozenset( + { + "RETRIEVAL_QUERY", + "RETRIEVAL_DOCUMENT", + "SEMANTIC_SIMILARITY", + "CLASSIFICATION", + "CLUSTERING", + "QUESTION_ANSWERING", + "FACT_VERIFICATION", + "CODE_RETRIEVAL_QUERY", + } +) + def _to_snake_case(key: str) -> str: """Converts a camelCase request body key to its snake_case equivalent.""" @@ -128,3 +148,114 @@ def build_generation_config(flat_params: Dict[str, Any]): "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/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/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_model_params.py b/tests/unit/sql/test_model_params.py index 27fdeaa..54ebd03 100644 --- a/tests/unit/sql/test_model_params.py +++ b/tests/unit/sql/test_model_params.py @@ -123,5 +123,89 @@ def test_thinking_config_is_understood(self): 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() From 5474901693165cbd1b48b469ce5e086f169ca2d8 Mon Sep 17 00:00:00 2001 From: Grace Qi Date: Wed, 16 Sep 2026 23:54:58 +0000 Subject: [PATCH 4/4] feat(sql): add ai_search, the Spark counterpart of BigQuery AI.SEARCH Adds `ai_search`, which ranks the rows of a table by semantic similarity to a query string, mirroring BigQuery's `AI.SEARCH`. ai_search(base_table, column_to_search, query, *, top_k=10, distance_type="EUCLIDEAN", embedding_column=None, endpoint="text-embedding-005", task_type=None, model_params=None, project=None, location=None) -> DataFrame The result has two columns: `base`, a struct of the input row, and `distance`. Smaller is always nearer, including for DOT_PRODUCT, which is negated so the three distance types order consistently. Design notes: * It is a DataFrame function rather than a SQL table function. BigQuery's AI.SEARCH is a TVF over indexes whose embeddings are generated and maintained automatically; Spark has neither concept, so the table and its embedding column are passed explicitly. * `base_table` accepts a DataFrame or a table/view name. * Embeddings are resolved in the order: an explicit `embedding_column`, then a column named `f"{column_to_search}_embedding"` (matching BigQuery's generated-column convention), then falling back to embedding the column inline. An existing column may be an ARRAY of any numeric type or the STRUCT that `ai_embed` returns. * When an embedding column is supplied the whole search stays in the JVM: the query is embedded once on the driver and the distance is a Catalyst expression, so no rows cross into Python. Verified against the optimized plan (zero ArrowEvalPython/BatchEvalPython nodes). * The query embedding raises on failure, matching BigQuery, while per-row failures stay soft: the row gets a null distance and drops out of the ranking rather than failing the query. Two bugs found while testing against Vertex AI are fixed here: * The model was called twice per row, because Catalyst pushes a Filter below a projection containing the embedding UDF. Ordering the pipeline as orderBy -> limit -> where avoids it, since a Filter is not pushed below a Limit. * Cosine distance raised DIVIDE_BY_ZERO under Spark 4's ANSI mode for zero-norm vectors. Now uses try_divide. Supporting change to _ai_embed: RESULT_FIELD and STATUS_FIELD become public so ai_search can read the struct it produces, and a new `embed_one` helper embeds a single driver-side value, raising rather than soft-failing. --- README.md | 72 +++ google/cloud/dataproc_ml/sql/__init__.py | 10 + google/cloud/dataproc_ml/sql/_ai_embed.py | 60 ++- google/cloud/dataproc_ml/sql/_ai_search.py | 428 ++++++++++++++++++ tests/integration/sql/test_ai_search.py | 174 ++++++++ tests/unit/sql/test_ai_search.py | 491 +++++++++++++++++++++ 6 files changed, 1230 insertions(+), 5 deletions(-) create mode 100644 google/cloud/dataproc_ml/sql/_ai_search.py create mode 100644 tests/integration/sql/test_ai_search.py create mode 100644 tests/unit/sql/test_ai_search.py diff --git a/README.md b/README.md index 769d96c..afb8d10 100644 --- a/README.md +++ b/README.md @@ -140,6 +140,78 @@ spark.udf.register("ai_embed", ai_embed_udf()) spark.sql("SELECT ai_embed(text).result AS vector FROM documents").show() ``` +### AI Functions: `ai_search` + +> **Note:** `ai_search` makes API calls to Vertex AI, which will incur costs. +> Please review the [Vertex AI Generative +> AI pricing](https://cloud.google.com/vertex-ai/generative-ai/pricing). + +`ai_search` finds the rows closest in meaning to a search term, mirroring the +semantics of BigQuery's +[`AI.SEARCH`](https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-search). +It takes the whole DataFrame rather than a column and returns the same two +columns BigQuery does: `base`, a struct of every column of the table searched, +and `distance`. A smaller distance always means a closer match. + +```python +from pyspark.sql import SparkSession +from google.cloud.dataproc_ml.sql import ai_search + +spark = SparkSession.builder.getOrCreate() + +products = spark.createDataFrame( + [ + ("Rocket kite", "A huge stunt kite that loops in a light breeze."), + ("Desk lamp", "An adjustable lamp with a warm dimmable bulb."), + ("Bath ducks", "A pack of six squeaky ducks for the tub."), + ], + ["name", "description"], +) + +ai_search(products, "description", "A really fun toy", top_k=2).select( + "base.name", "distance" +).show() +# +-----------+------------------+ +# | name| distance| +# +-----------+------------------+ +# | Bath ducks|0.7721... | +# |Rocket kite|0.8034... | +# +-----------+------------------+ +``` + +BigQuery searches a table whose embeddings are maintained for it in a +generated column. Spark has no such column, so `ai_search` looks for the +embeddings in three places, in order: + +1. the column named by `embedding_column`; +2. a column named `_embedding`, which is the name + BigQuery's own examples give the generated column; +3. failing both, it generates them by calling `ai_embed` on + `column_to_search`. + +An existing embedding column may hold either a plain `ARRAY` or the +struct that `ai_embed` returns. Prefer to embed the table once and keep the +result, as in the third case every search re-embeds every row: + +```python +from google.cloud.dataproc_ml.sql import ai_embed + +embedded = products.withColumn( + "description_embedding", ai_embed("description") +).cache() + +ai_search(embedded, "description", "Something to read by", top_k=1) +``` + +Pass `distance_type` to choose how distance is measured, one of `EUCLIDEAN` +(the default, as in BigQuery), `COSINE` or `DOT_PRODUCT`. Because a smaller +distance always means a closer match, `DOT_PRODUCT`, which grows with +similarity, is returned negated. + +Rows whose own embedding is missing are skipped. A search term that cannot be +embedded raises instead, since there is no result to return without it. This +also matches BigQuery. + ### Generative AI (Gemini) Model Inference > **Note:** Using the `GenAiModelHandler` involves making API calls to diff --git a/google/cloud/dataproc_ml/sql/__init__.py b/google/cloud/dataproc_ml/sql/__init__.py index 4c10a8c..59d4412 100644 --- a/google/cloud/dataproc_ml/sql/__init__.py +++ b/google/cloud/dataproc_ml/sql/__init__.py @@ -39,16 +39,26 @@ >>> from google.cloud.dataproc_ml.sql import ai_embed >>> >>> embedded = df.withColumn("vector", ai_embed(col("body")).result) + +Semantic search takes the whole DataFrame rather than a column, and returns +the matching rows: + + >>> from google.cloud.dataproc_ml.sql import ai_search + >>> + >>> hits = ai_search(df, "body", "I cannot sign in", top_k=3) + >>> hits.select("base.id", "base.body", "distance") """ from ._ai_embed import ai_embed from ._ai_embed import ai_embed_udf from ._ai_generate import ai_generate from ._ai_generate import ai_generate_udf +from ._ai_search import ai_search __all__ = ( "ai_embed", "ai_embed_udf", "ai_generate", "ai_generate_udf", + "ai_search", ) diff --git a/google/cloud/dataproc_ml/sql/_ai_embed.py b/google/cloud/dataproc_ml/sql/_ai_embed.py index 8699686..095fd57 100644 --- a/google/cloud/dataproc_ml/sql/_ai_embed.py +++ b/google/cloud/dataproc_ml/sql/_ai_embed.py @@ -53,8 +53,10 @@ # 10,000 rows by default, so splitting is required rather than optional. _MAX_INSTANCES_PER_REQUEST = 250 -_RESULT_FIELD = "result" -_STATUS_FIELD = "status" +#: Field names of :data:`RETURN_TYPE`. Not private to the module because +#: ``ai_search`` reads the embeddings back out of a column of this shape. +RESULT_FIELD = "result" +STATUS_FIELD = "status" # The endpoint rejects a request whose only instance is the empty string, and, # worse, silently drops an empty instance from a request that has other @@ -73,8 +75,8 @@ #: 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), + StructField(RESULT_FIELD, ArrayType(DoubleType()), nullable=True), + StructField(STATUS_FIELD, StringType(), nullable=True), ] ) @@ -235,6 +237,54 @@ def _predict( return pandas_udf(_predict, returnType=RETURN_TYPE) +def embed_one( + *, + embed_content, + content: str, + endpoint: str, + config, +) -> List[float]: + """Embeds a single value on the driver, raising if it cannot be embedded. + + The column functions report a failed row in its ``status`` field so that + one bad row never fails a stage. A single value embedded on the driver has + no such row to report against, and a caller that asked for one embedding + and got none cannot carry on, so this raises instead. + + Retries and error wording are shared with the column function. + + Args: + embed_content: The awaitable Vertex AI call. + content: The text to embed. + endpoint: The embedding model to call. + config: The embedding config. + + Returns: + The embedding vector. + + Raises: + ValueError: If the content is empty or the model reported an error. + """ + rows = _client.run_coroutine( + _embed_batch( + embed_content=embed_content, + contents=[content], + endpoint=endpoint, + config=config, + ) + ) + row = rows[0] + if row[STATUS_FIELD]: + raise ValueError(f"Could not embed {content!r}: {row[STATUS_FIELD]}") + vector = row[RESULT_FIELD] + if not vector: + raise ValueError( + f"The embedding model {endpoint} returned no vector for " + f"{content!r}." + ) + return vector + + async def _embed_batch( *, embed_content, @@ -383,7 +433,7 @@ def _record( def _row(*, result: Optional[List[float]], status: str) -> Dict[str, Any]: """Builds one output row.""" - return {_RESULT_FIELD: result, _STATUS_FIELD: status} + return {RESULT_FIELD: result, STATUS_FIELD: status} def _is_missing(content: Any) -> bool: diff --git a/google/cloud/dataproc_ml/sql/_ai_search.py b/google/cloud/dataproc_ml/sql/_ai_search.py new file mode 100644 index 0000000..c8b3bcd --- /dev/null +++ b/google/cloud/dataproc_ml/sql/_ai_search.py @@ -0,0 +1,428 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``AI.SEARCH``: semantic search over a DataFrame.""" + +import logging +from typing import Any, Dict, List, Optional, Union + +from pyspark.sql import Column +from pyspark.sql import DataFrame +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf +from pyspark.sql.types import ArrayType +from pyspark.sql.types import DoubleType +from pyspark.sql.types import StructType + +from google.cloud.dataproc_ml.sql import _ai_embed +from google.cloud.dataproc_ml.sql import _client +from google.cloud.dataproc_ml.sql import _model_params + +logger = logging.getLogger(__name__) + +#: The number of results returned when ``top_k`` is not given, matching +#: BigQuery. +DEFAULT_TOP_K = 10 + +#: The distance used when ``distance_type`` is not given, matching BigQuery. +DEFAULT_DISTANCE_TYPE = "EUCLIDEAN" + +#: The distances BigQuery's ``AI.SEARCH`` accepts. +DISTANCE_TYPES = ("EUCLIDEAN", "COSINE", "DOT_PRODUCT") + +#: The columns of the returned DataFrame, matching BigQuery's ``AI.SEARCH``. +BASE_COLUMN = "base" +DISTANCE_COLUMN = "distance" + +#: BigQuery searches a table whose embeddings are maintained for it in a +#: generated column conventionally named after the column it embeds. The same +#: name is looked for here, so a DataFrame that carries its embeddings needs no +#: extra argument. +EMBEDDING_COLUMN_SUFFIX = "_embedding" + +# Holds each row's vector between the two projections below. Prefixed so that +# it cannot collide with a column of the table being searched. +_VECTOR_COLUMN = "__dataproc_ml_ai_search_vector" + +_VECTOR_TYPE = ArrayType(DoubleType()) + +BaseTable = Union[DataFrame, str] + + +def ai_search( + base_table: BaseTable, + column_to_search: str, + query: str, + *, + top_k: int = DEFAULT_TOP_K, + distance_type: str = DEFAULT_DISTANCE_TYPE, + embedding_column: Optional[str] = None, + endpoint: str = _ai_embed.DEFAULT_ENDPOINT, + task_type: Optional[str] = None, + model_params: Optional[Dict[str, Any]] = None, + project: Optional[str] = None, + location: Optional[str] = None, + _embed_content=None, +) -> DataFrame: + """Returns the rows closest in meaning to a search term. + + This is the Spark counterpart of BigQuery's ``AI.SEARCH`` function and + returns the same two columns, ``base`` and ``distance``, so a query can be + moved between the two engines with little change. + + BigQuery searches a table whose embeddings are maintained for it, in a + generated column. Spark has no such column, so the embeddings are found in + one of three ways, in order: + + 1. the column named by ``embedding_column``; + 2. a column named ``_embedding``, which is the name + BigQuery's own examples give the generated column; + 3. failing both, they are generated here by calling + :func:`~google.cloud.dataproc_ml.sql.ai_embed` on + ``column_to_search``. + + An existing embedding column may hold either a plain ``ARRAY`` or + the struct that ``ai_embed`` returns. Prefer to embed the table once and + keep the result, as in the third case every search re-embeds every row. + + The search term itself is embedded once, on the driver. Unlike the column + functions, which report a failed row in its ``status`` field, a search term + that cannot be embedded raises: there is no result to return without it. + This also matches BigQuery. Rows whose own embedding is missing are skipped + rather than failing the search. + + .. note:: + Using this function incurs Vertex AI charges. See the + `Vertex AI Generative AI pricing + `_ page. + + Args: + base_table: The DataFrame to search, or the name of a table or view to + read it from. + column_to_search: The name of the column holding the text to search. + As in BigQuery this names the source column, not the embeddings. + query: The search term. + top_k: The number of rows to return, nearest first. Defaults to 10. + distance_type: How to measure the distance between two embeddings, one + of ``'EUCLIDEAN'``, ``'COSINE'`` or ``'DOT_PRODUCT'``. Defaults to + ``'EUCLIDEAN'``. + embedding_column: The name of a column holding the embeddings of + ``column_to_search``, when it is not named for the column it + embeds. + endpoint: The embedding model to call. Defaults to + ``text-embedding-005``. It must be the model the table's own + embeddings were generated with, since embeddings from two models + cannot be compared. + task_type: The ``task_type`` of the search term, and of the table's + embeddings when they are generated here. See + :func:`~google.cloud.dataproc_ml.sql.ai_embed`. + model_params: Additional request settings. See + :func:`~google.cloud.dataproc_ml.sql.ai_embed`. + project: The Google Cloud project billed for the calls. Inferred from + the environment when not set. + location: The Vertex AI location. Defaults to ``global``. + _embed_content: Internal. Replaces the call to Vertex AI with an + awaitable of the same shape, so that tests can run the whole Spark + path without credentials or network access. + + Returns: + A DataFrame of at most ``top_k`` rows ordered by ``distance``, holding + ``base``, a struct of every column of the table searched, and + ``distance``, the distance between the row and the search term. A + smaller distance always means a closer match, so ``DOT_PRODUCT`` + returns the negated inner product. + + Raises: + ValueError: If any argument is invalid, or if the search term cannot + be embedded. + + Example: + >>> hits = ai_search(products, "description", "A really fun toy") + >>> hits.select("base.name", "base.description", "distance").show() + + Searching a table that already carries its embeddings, which is both + cheaper and what BigQuery does: + + >>> embedded = products.withColumn( + ... "description_embedding", ai_embed("description") + ... ).cache() + >>> ai_search(embedded, "description", "A really fun toy", top_k=3) + """ + table = _resolve_base_table(base_table) + distance_type = _validate_distance_type(distance_type) + _validate_top_k(top_k) + _validate_query(query) + _validate_column_to_search(table, column_to_search) + + flat_params = _model_params.normalize_embed_params(model_params, task_type) + + query_vector = _ai_embed.embed_one( + embed_content=( + _embed_content or _client.get_embed_content(project, location) + ), + content=query, + endpoint=endpoint, + config=_model_params.build_embed_config(flat_params), + ) + + vector = _row_vector( + table, + column_to_search=column_to_search, + embedding_column=embedding_column, + endpoint=endpoint, + task_type=task_type, + model_params=model_params, + project=project, + location=location, + _embed_content=_embed_content, + ) + + base_columns = list(table.columns) + if _VECTOR_COLUMN in base_columns: + raise ValueError( + f"{_VECTOR_COLUMN} is reserved by ai_search and cannot be a " + "column of the table being searched." + ) + + # The vector is attached in its own projection and read back by name. When + # it is generated here that expression is a call to the model, and Spark + # will not fold a projection whose output is read more than once into the + # one that reads it, so the distance below costs one call per row rather + # than one per mention. + scored = table.withColumn(_VECTOR_COLUMN, vector).select( + sf.struct(*[sf.col(name) for name in base_columns]).alias(BASE_COLUMN), + _distance(sf.col(_VECTOR_COLUMN), query_vector, distance_type).alias( + DISTANCE_COLUMN + ), + ) + + # A row with no embedding, or one of a different width, scores null and is + # dropped. BigQuery likewise skips rows whose embedding is missing. + # + # Those rows are sorted to the end and dropped after the limit rather than + # before it, which returns the same rows: the limit keeps the nearest + # matches, and only reaches a null once it has run out of real ones. The + # order matters because a filter placed before the limit is pushed down + # past the projection above, which would evaluate the embeddings a second + # time for every row. Nothing is pushed past a limit. + return ( + scored.orderBy(sf.col(DISTANCE_COLUMN).asc_nulls_last()) + .limit(top_k) + .where(sf.col(DISTANCE_COLUMN).isNotNull()) + ) + + +def _resolve_base_table(base_table: BaseTable) -> DataFrame: + """Returns the DataFrame to search. + + Raises: + ValueError: If the argument is neither a DataFrame nor a table name, + or if no Spark session is available to resolve a name. + """ + if isinstance(base_table, DataFrame): + return base_table + if isinstance(base_table, str): + if not base_table.strip(): + raise ValueError("base_table must be a non-empty table name.") + spark = SparkSession.getActiveSession() + if spark is None: + raise ValueError( + f"Cannot read the table {base_table!r} because there is no " + "active Spark session. Pass a DataFrame instead." + ) + return spark.table(base_table) + raise ValueError( + "base_table must be a DataFrame or the name of a table, got " + f"{type(base_table).__name__}." + ) + + +def _validate_distance_type(distance_type: str) -> str: + """Returns the distance type, normalized to upper case. + + Raises: + ValueError: If the distance type is not one of :data:`DISTANCE_TYPES`. + """ + if ( + not isinstance(distance_type, str) + or distance_type.upper() not in DISTANCE_TYPES + ): + supported = ", ".join(DISTANCE_TYPES) + raise ValueError( + f"Unsupported distance_type: {distance_type!r}. Supported values " + f"are {supported}." + ) + return distance_type.upper() + + +def _validate_top_k(top_k: int) -> None: + """Checks ``top_k``. + + Raises: + ValueError: If it is not a positive whole number. + """ + # A bool is an int in Python, and top_k=True is a mistake, not a request + # for one row. + if isinstance(top_k, bool) or not isinstance(top_k, int) or top_k < 1: + raise ValueError(f"top_k must be a positive integer, got {top_k!r}.") + + +def _validate_query(query: str) -> None: + """Checks the search term. + + Raises: + ValueError: If it is not a non-empty string. + """ + if not isinstance(query, str) or not query.strip(): + raise ValueError( + f"query must be a non-empty search term, got {query!r}." + ) + + +def _validate_column_to_search(table: DataFrame, name: str) -> None: + """Checks that the column to search exists. + + Raises: + ValueError: If the name is empty or names no column of the table. + """ + if not isinstance(name, str) or not name.strip(): + raise ValueError( + "column_to_search must be the name of a column, got " f"{name!r}." + ) + if name not in table.columns: + available = ", ".join(table.columns) + raise ValueError( + f"column_to_search {name!r} is not a column of the table being " + f"searched. Its columns are {available}." + ) + + +def _row_vector( + table: DataFrame, + *, + column_to_search: str, + embedding_column: Optional[str], + endpoint: str, + task_type: Optional[str], + model_params: Optional[Dict[str, Any]], + project: Optional[str], + location: Optional[str], + _embed_content, +) -> Column: + """Returns the column holding each row's embedding. + + Raises: + ValueError: If a named embedding column is missing or holds neither an + array nor the struct returned by ``ai_embed``. + """ + if embedding_column is not None: + if embedding_column not in table.columns: + available = ", ".join(table.columns) + raise ValueError( + f"embedding_column {embedding_column!r} is not a column of " + f"the table being searched. Its columns are {available}." + ) + return _existing_vector(table, embedding_column) + + conventional = f"{column_to_search}{EMBEDDING_COLUMN_SUFFIX}" + if conventional in table.columns: + logger.info( + "Searching the embeddings in %s. Pass embedding_column to use a " + "different column.", + conventional, + ) + return _existing_vector(table, conventional) + + logger.info( + "The table carries no embeddings of %s, so they will be generated " + "for every row of this search. Add a %s column to avoid paying for " + "them again on the next search.", + column_to_search, + conventional, + ) + return _ai_embed.ai_embed( + column_to_search, + endpoint=endpoint, + task_type=task_type, + model_params=model_params, + project=project, + location=location, + _embed_content=_embed_content, + ).getField(_ai_embed.RESULT_FIELD) + + +def _existing_vector(table: DataFrame, name: str) -> Column: + """Returns the embeddings held in an existing column. + + Raises: + ValueError: If the column holds neither an array of numbers nor the + struct returned by ``ai_embed``. + """ + data_type = table.schema[name].dataType + if isinstance(data_type, StructType): + if _ai_embed.RESULT_FIELD not in data_type.fieldNames(): + raise ValueError( + f"The embedding column {name} is a struct without a " + f"{_ai_embed.RESULT_FIELD} field, so it does not hold the " + "output of ai_embed." + ) + return sf.col(name).getField(_ai_embed.RESULT_FIELD).cast(_VECTOR_TYPE) + if isinstance(data_type, ArrayType): + return sf.col(name).cast(_VECTOR_TYPE) + raise ValueError( + f"The embedding column {name} is a {data_type.simpleString()}. It " + "must be an array of numbers, or the struct returned by ai_embed." + ) + + +def _distance( + vector: Column, query_vector: List[float], distance_type: str +) -> Column: + """Returns the distance between each row's embedding and the search term. + + Smaller always means nearer, so ``DOT_PRODUCT``, which grows with + similarity, is negated. That keeps one ordering for every distance and + lets ``top_k`` mean the same thing throughout. + + Rows whose embedding is null, or of a different width from the search + term, score null: ``zip_with`` pads the shorter array with nulls, which + carry through the sum. They are dropped rather than compared against a + vector they cannot meaningfully be compared with. + """ + query = sf.array(*[sf.lit(float(value)) for value in query_vector]) + + if distance_type == "EUCLIDEAN": + squares = sf.zip_with(vector, query, lambda x, y: (x - y) * (x - y)) + return sf.sqrt(_total(squares)) + + products = sf.zip_with(vector, query, lambda x, y: x * y) + if distance_type == "DOT_PRODUCT": + return -_total(products) + + # COSINE. The search term's norm is a constant, so only the row's own norm + # is computed in Spark. A vector of zeroes has no direction to compare, so + # it scores null and the row is dropped. The division is the tolerant one + # because Spark raises on a division by zero under ANSI mode, which is the + # default from Spark 4, and one such row would otherwise fail the search. + query_norm = sum(value * value for value in query_vector) ** 0.5 + row_norm = sf.sqrt(_total(sf.transform(vector, lambda x: x * x))) + similarity = sf.try_divide(_total(products), row_norm * sf.lit(query_norm)) + return sf.lit(1.0) - similarity + + +def _total(values: Column) -> Column: + """Sums an array of doubles, yielding null if any element is null.""" + return sf.aggregate( + values, sf.lit(0.0), lambda accumulator, value: accumulator + value + ) diff --git a/tests/integration/sql/test_ai_search.py b/tests/integration/sql/test_ai_search.py new file mode 100644 index 0000000..5feba7c --- /dev/null +++ b/tests/integration/sql/test_ai_search.py @@ -0,0 +1,174 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration tests that call Vertex AI for real. + +They rely on Application Default Credentials and, unless GOOGLE_CLOUD_PROJECT +is set, on the project discovered from the environment. The location is not +taken from the environment: these tests exercise the default global endpoint. +Running them incurs Vertex AI charges. +""" + +import unittest + +from pyspark.sql import SparkSession + +from google.cloud.dataproc_ml.sql import _ai_search as ai_search_module +from google.cloud.dataproc_ml.sql import ai_embed, ai_search + +# Products that a reader can rank by hand. Each pair of rows is close on one +# subject and far apart on every other, so a search for one subject has an +# unambiguous answer. +PRODUCTS = [ + (1, "Stunt kite", "A huge kite that loops and dives in a light breeze."), + (2, "Bath ducks", "Six squeaky rubber ducks for the tub."), + (3, "Desk lamp", "An adjustable lamp with a warm dimmable bulb."), + (4, "Floor lamp", "A tall standing light for a reading corner."), + (5, "Chef's knife", "A forged carbon steel blade for the kitchen."), + (6, "Cutting board", "An end grain maple block for chopping."), +] + +# The rows a reader would pick for each search, in no particular order. +EXPECTED = { + "a really fun toy": {1, 2}, + "something to read by in the evening": {3, 4}, + "preparing dinner": {5, 6}, +} + + +class TestAiSearch(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.spark = SparkSession.builder.getOrCreate() + cls.products = cls.spark.createDataFrame( + PRODUCTS, ["id", "name", "description"] + ).repartition(3) + # Embedded once and kept, both because it is what the function is + # meant to be used with and because it keeps these tests cheap. + cls.embedded = cls.products.withColumn( + "description_embedding", ai_embed("description") + ).cache() + cls.embedded.count() + + @classmethod + def tearDownClass(cls): + cls.spark.stop() + + def ids(self, hits): + """Returns the ids of the rows found, in the order they came back.""" + return [row["base"]["id"] for row in hits.collect()] + + def test_finds_the_rows_a_reader_would_pick(self): + for query, expected in EXPECTED.items(): + with self.subTest(query=query): + hits = ai_search(self.embedded, "description", query, top_k=2) + self.assertEqual(set(self.ids(hits)), expected) + + def test_every_distance_agrees_on_the_nearest_row(self): + query = "preparing dinner" + nearest = set() + for distance_type in ai_search_module.DISTANCE_TYPES: + with self.subTest(distance_type=distance_type): + hits = ai_search( + self.embedded, + "description", + query, + top_k=2, + distance_type=distance_type, + ) + found = set(self.ids(hits)) + self.assertEqual(found, EXPECTED[query]) + nearest.add(tuple(self.ids(hits))) + # The three distances rank the same two rows. They may disagree on + # the order within the pair, which is not something to assert on. + + def test_the_returned_columns_match_bigquery(self): + hits = ai_search(self.embedded, "description", "a lamp", top_k=1) + self.assertEqual(hits.columns, ["base", "distance"]) + row = hits.collect()[0] + self.assertEqual( + set(row["base"].asDict()), + {"id", "name", "description", "description_embedding"}, + ) + self.assertIsInstance(row["distance"], float) + + def test_top_k_bounds_the_result(self): + hits = ai_search(self.embedded, "description", "a lamp", top_k=3) + self.assertEqual(hits.count(), 3) + + def test_the_distance_never_decreases(self): + distances = [ + row["distance"] + for row in ai_search( + self.embedded, "description", "a lamp", top_k=6 + ).collect() + ] + self.assertEqual(distances, sorted(distances)) + + def test_embeds_the_table_when_it_carries_no_embeddings(self): + # The same answer, paid for with a call per row instead of one call. + hits = ai_search( + self.products, "description", "a really fun toy", top_k=2 + ) + self.assertEqual(set(self.ids(hits)), EXPECTED["a really fun toy"]) + + def test_a_named_embedding_column_is_used(self): + renamed = self.embedded.withColumnRenamed( + "description_embedding", "vectors" + ) + hits = ai_search( + renamed, + "description", + "preparing dinner", + top_k=2, + embedding_column="vectors", + ) + self.assertEqual(set(self.ids(hits)), EXPECTED["preparing dinner"]) + + def test_a_named_table_is_searched(self): + self.embedded.createOrReplaceTempView("products_for_search") + hits = ai_search( + "products_for_search", "description", "a lamp", top_k=2 + ) + self.assertEqual(set(self.ids(hits)), {3, 4}) + + def test_a_shorter_vector_still_searches(self): + # outputDimensionality has to match on both sides, so the table is + # re-embedded at the same width as the search term. + params = {"outputDimensionality": 256} + hits = ai_search( + self.products, + "description", + "preparing dinner", + top_k=2, + model_params=params, + ) + self.assertEqual(set(self.ids(hits)), EXPECTED["preparing dinner"]) + + def test_an_unknown_endpoint_fails_the_search(self): + # Unlike a row, a search term that cannot be embedded has nowhere to + # report a failure, so it raises. + with self.assertRaises(ValueError) as caught: + ai_search( + self.embedded, + "description", + "a lamp", + endpoint="text-embedding-does-not-exist", + ) + self.assertIn("a lamp", str(caught.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sql/test_ai_search.py b/tests/unit/sql/test_ai_search.py new file mode 100644 index 0000000..9bc90fd --- /dev/null +++ b/tests/unit/sql/test_ai_search.py @@ -0,0 +1,491 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import unittest + +from google.genai import errors, types +from pyspark.sql import SparkSession +from pyspark.sql import functions as sf +from pyspark.sql.types import ArrayType, DoubleType, FloatType + +# pylint: disable=ungrouped-imports +from google.cloud.dataproc_ml.sql import _ai_search as ai_search_module +from google.cloud.dataproc_ml.sql import ai_embed, ai_search + +# pylint: enable=ungrouped-imports + +# A tiny stand-in for an embedding model. Each phrase has a hand-written +# vector, so the distance between any two rows can be worked out by hand and +# the expected ordering is obvious from reading the test. +VECTORS = { + "alpha": [1.0, 0.0, 0.0], + "almost alpha": [0.9, 0.1, 0.0], + "beta": [0.0, 1.0, 0.0], + "gamma": [0.0, 0.0, 1.0], +} + +# A phrase the fake model cannot embed, standing in for a row the real model +# fails on. +UNKNOWN = "not in the vocabulary" + + +def embeds_each_content(): + """Returns a fake ``embed_content`` answering from :data:`VECTORS`. + + Defined inside a function rather than at module level so that it is + serialized by value and can run inside a Spark worker. + """ + + async def _embed_content(*, model, contents, config): + del model, config + return types.EmbedContentResponse( + embeddings=[ + types.ContentEmbedding(values=VECTORS[text]) + if text in VECTORS + else types.ContentEmbedding(values=None) + for text in contents + ] + ) + + return _embed_content + + +def fails_with(error): + """Returns a fake ``embed_content`` that always raises.""" + + async def _embed_content(*, model, contents, config): + del model, contents, config + raise error + + return _embed_content + + +def api_error(code, status, message): + """Builds an API error the way the client library reports one.""" + return errors.APIError( + code, {"error": {"code": code, "status": status, "message": message}} + ) + + +class SparkTestCase(unittest.TestCase): + """Base class holding a Spark session shared by every test module.""" + + @classmethod + def setUpClass(cls): + cls.spark = ( + SparkSession.builder.master("local[2]") + .appName("dataproc-ml-sql-tests") + .config("spark.sql.execution.arrow.maxRecordsPerBatch", "2") + .config("spark.ui.enabled", "false") + .getOrCreate() + ) + + def tickets(self, *phrases): + """Builds a table of phrases the fake model knows, numbered from 1.""" + return self.spark.createDataFrame( + [(index, phrase) for index, phrase in enumerate(phrases, start=1)], + "id INT, body STRING", + ) + + def search(self, table, query="alpha", **kwargs): + """Searches with the fake model, so no credentials are needed.""" + kwargs.setdefault("_embed_content", embeds_each_content()) + return ai_search(table, "body", query, **kwargs) + + def ids(self, hits): + """Returns the ids of the rows found, in the order they came back.""" + return [row["base"]["id"] for row in hits.collect()] + + +class TestReturnedColumns(SparkTestCase): + """The output matches BigQuery's AI.SEARCH.""" + + def test_the_columns_are_base_and_distance(self): + hits = self.search(self.tickets("alpha")) + self.assertEqual(hits.columns, ["base", "distance"]) + + def test_base_holds_every_column_of_the_table(self): + hits = self.search(self.tickets("alpha")) + self.assertEqual( + hits.schema["base"].dataType.fieldNames(), ["id", "body"] + ) + + def test_the_distance_is_a_double(self): + # BigQuery returns FLOAT64, which is a double in Spark. + hits = self.search(self.tickets("alpha")) + self.assertEqual(hits.schema["distance"].dataType, DoubleType()) + + def test_an_embedding_column_is_not_returned(self): + # The column ai_search adds to hold each row's vector is an + # implementation detail and must not leak into the result. + hits = self.search(self.tickets("alpha")) + self.assertNotIn( + ai_search_module._VECTOR_COLUMN, + hits.schema["base"].dataType.fieldNames(), + ) + + +class TestOrdering(SparkTestCase): + """Rows come back nearest first, whichever distance is used.""" + + def setUp(self): + self.table = self.tickets("alpha", "beta", "gamma", "almost alpha") + + def test_the_nearest_row_comes_first(self): + self.assertEqual(self.ids(self.search(self.table))[:2], [1, 4]) + + def test_cosine_orders_the_same_way(self): + hits = self.search(self.table, distance_type="COSINE") + self.assertEqual(self.ids(hits)[:2], [1, 4]) + + def test_dot_product_orders_the_same_way(self): + # The inner product grows with similarity, so it is negated to keep + # one ordering for every distance. + hits = self.search(self.table, distance_type="DOT_PRODUCT") + self.assertEqual(self.ids(hits)[:2], [1, 4]) + + def test_the_distance_never_decreases(self): + distances = [ + row["distance"] for row in self.search(self.table).collect() + ] + self.assertEqual(distances, sorted(distances)) + + def test_a_lowercase_distance_type_is_accepted(self): + # A name written the way it reads in SQL should not be a mistake. + hits = self.search(self.table, distance_type="cosine") + self.assertEqual(self.ids(hits)[:2], [1, 4]) + + +class TestDistances(SparkTestCase): + """Each distance is the one it is named after.""" + + def distance_to_alpha(self, phrase, distance_type): + hits = self.search(self.tickets(phrase), distance_type=distance_type) + return hits.collect()[0]["distance"] + + def test_euclidean(self): + # |[0.9, 0.1, 0] - [1, 0, 0]| = sqrt(0.01 + 0.01). + self.assertAlmostEqual( + self.distance_to_alpha("almost alpha", "EUCLIDEAN"), + math.sqrt(0.02), + ) + + def test_cosine(self): + # 1 - 0.9 / (sqrt(0.82) * 1). + self.assertAlmostEqual( + self.distance_to_alpha("almost alpha", "COSINE"), + 1 - 0.9 / math.sqrt(0.82), + ) + + def test_dot_product_is_negated(self): + self.assertAlmostEqual( + self.distance_to_alpha("almost alpha", "DOT_PRODUCT"), -0.9 + ) + + def test_an_identical_row_is_at_no_distance(self): + self.assertAlmostEqual(self.distance_to_alpha("alpha", "EUCLIDEAN"), 0) + self.assertAlmostEqual(self.distance_to_alpha("alpha", "COSINE"), 0) + + +class TestTopK(SparkTestCase): + """Only the nearest rows are returned.""" + + def test_at_most_top_k_rows_come_back(self): + table = self.tickets("alpha", "beta", "gamma", "almost alpha") + self.assertEqual(self.search(table, top_k=2).count(), 2) + + def test_a_smaller_table_is_returned_whole(self): + table = self.tickets("alpha", "beta") + self.assertEqual(self.search(table, top_k=10).count(), 2) + + def test_the_default_matches_bigquery(self): + self.assertEqual(ai_search_module.DEFAULT_TOP_K, 10) + + +class TestExistingEmbeddings(SparkTestCase): + """The table's own embeddings are used when it has any.""" + + def embedded(self, *phrases, column="body_embedding"): + return self.tickets(*phrases).withColumn( + column, ai_embed("body", _embed_content=embeds_each_content()) + ) + + def test_a_column_named_for_the_one_it_embeds_is_found(self): + # BigQuery's own examples name the generated column this way. + hits = self.search(self.embedded("alpha", "beta"), top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_a_named_column_is_used(self): + table = self.embedded("alpha", "beta", column="vectors") + hits = self.search(table, embedding_column="vectors", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_a_plain_array_column_is_used(self): + table = self.embedded("alpha", "beta").withColumn( + "vectors", sf.col("body_embedding.result") + ) + hits = self.search(table, embedding_column="vectors", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_an_array_of_floats_is_accepted(self): + # An embedding read back from storage may be a float rather than a + # double, which is the same vector at less precision. + table = self.tickets("alpha", "beta").withColumn( + "vectors", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ) + .otherwise(sf.array(sf.lit(0.0), sf.lit(1.0), sf.lit(0.0))) + .cast(ArrayType(FloatType())), + ) + hits = self.search(table, embedding_column="vectors", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_the_named_column_wins_over_the_conventional_one(self): + # The named column embeds each row as if it were the next one along, + # so using it rather than body_embedding changes which row is nearest. + table = self.embedded("alpha", "beta").withColumn( + "shifted", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(0.0), sf.lit(1.0), sf.lit(0.0)), + ).otherwise(sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0))), + ) + self.assertEqual(self.ids(self.search(table, top_k=1)), [1]) + hits = self.search(table, embedding_column="shifted", top_k=1) + self.assertEqual(self.ids(hits), [2]) + + def test_the_embeddings_are_not_regenerated(self): + # The model is never reachable here, so the search can only succeed by + # using the embeddings the table already carries. + table = self.embedded("alpha", "beta") + hits = ai_search( + table, + "body", + "alpha", + top_k=1, + # The query still has to be embedded, and only the query. + _embed_content=embeds_each_content(), + ) + self.assertEqual(self.ids(hits), [1]) + + +class TestRowsWithoutAnEmbedding(SparkTestCase): + """A row that cannot be compared is skipped, not guessed at.""" + + def test_a_row_the_model_failed_on_is_skipped(self): + table = self.tickets("alpha", UNKNOWN) + self.assertEqual(self.ids(self.search(table, top_k=10)), [1]) + + def test_a_null_embedding_is_skipped(self): + table = self.tickets("alpha", "beta").withColumn( + "body_embedding", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ), + ) + hits = ai_search( + table, + "body", + "alpha", + _embed_content=embeds_each_content(), + ) + self.assertEqual(self.ids(hits), [1]) + + def test_an_embedding_of_a_different_width_is_skipped(self): + # Two vectors of different widths describe different spaces, so the + # row is dropped rather than compared on the values they share. + table = self.tickets("alpha", "beta").withColumn( + "body_embedding", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ).otherwise(sf.array(sf.lit(0.0), sf.lit(1.0))), + ) + hits = ai_search( + table, "body", "alpha", _embed_content=embeds_each_content() + ) + self.assertEqual(self.ids(hits), [1]) + + def test_a_zero_vector_is_skipped_by_cosine(self): + # A vector with no direction has no angle to the query. + table = self.tickets("alpha", "beta").withColumn( + "body_embedding", + sf.when( + sf.col("id") == 1, + sf.array(sf.lit(1.0), sf.lit(0.0), sf.lit(0.0)), + ).otherwise(sf.array(sf.lit(0.0), sf.lit(0.0), sf.lit(0.0))), + ) + hits = ai_search( + table, + "body", + "alpha", + distance_type="COSINE", + _embed_content=embeds_each_content(), + ) + self.assertEqual(self.ids(hits), [1]) + + +class TestTheSearchTermFailsTheSearch(SparkTestCase): + """Unlike a row, a search term that cannot be embedded stops the search.""" + + def test_a_backend_failure_is_raised(self): + with self.assertRaises(ValueError) as caught: + ai_search( + self.tickets("alpha"), + "body", + "alpha", + _embed_content=fails_with( + api_error(403, "PERMISSION_DENIED", "Denied.") + ), + ) + self.assertIn("PERMISSION_DENIED", str(caught.exception)) + + def test_the_failure_names_the_search_term(self): + with self.assertRaises(ValueError) as caught: + ai_search( + self.tickets("alpha"), + "body", + "a missing term", + _embed_content=fails_with( + api_error(404, "NOT_FOUND", "No such model.") + ), + ) + self.assertIn("a missing term", str(caught.exception)) + + def test_an_empty_vector_is_raised(self): + with self.assertRaises(ValueError): + ai_search( + self.tickets("alpha"), + "body", + UNKNOWN, + _embed_content=embeds_each_content(), + ) + + def test_the_failure_happens_before_any_job(self): + # Nothing is scheduled, so the mistake surfaces at the call rather + # than at the first action taken on the result. + with self.assertRaises(ValueError): + ai_search( + self.tickets("alpha"), + "body", + "alpha", + _embed_content=fails_with( + api_error(400, "INVALID_ARGUMENT", "Bad request.") + ), + ) + + +class TestValidation(SparkTestCase): + """Mistakes are reported on the driver, before any job.""" + + def test_an_unknown_distance_type_is_rejected(self): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), distance_type="MANHATTAN") + + def test_top_k_must_be_positive(self): + for top_k in (0, -1): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), top_k=top_k) + + def test_top_k_must_be_a_whole_number(self): + for top_k in (1.5, "3", True): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), top_k=top_k) + + def test_an_empty_query_is_rejected(self): + for query in ("", " ", None): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), query=query) + + def test_a_missing_column_to_search_is_rejected(self): + with self.assertRaises(ValueError) as caught: + ai_search( + self.tickets("alpha"), + "subject", + "alpha", + _embed_content=embeds_each_content(), + ) + # The message lists what is there, so the typo is obvious. + self.assertIn("body", str(caught.exception)) + + def test_a_missing_embedding_column_is_rejected(self): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), embedding_column="vectors") + + def test_an_embedding_column_of_the_wrong_type_is_rejected(self): + with self.assertRaises(ValueError) as caught: + self.search(self.tickets("alpha"), embedding_column="body") + self.assertIn("array", str(caught.exception)) + + def test_a_struct_without_a_result_field_is_rejected(self): + table = self.tickets("alpha").withColumn( + "vectors", sf.struct(sf.lit(1.0).alias("value")) + ) + with self.assertRaises(ValueError) as caught: + self.search(table, embedding_column="vectors") + self.assertIn("result", str(caught.exception)) + + def test_a_base_table_of_the_wrong_type_is_rejected(self): + with self.assertRaises(ValueError): + self.search([("alpha",)]) + + def test_an_unknown_task_type_is_rejected(self): + with self.assertRaises(ValueError): + self.search(self.tickets("alpha"), task_type="RETRIEVAL_EVERYTHING") + + def test_a_valid_configuration_is_accepted(self): + self.search( + self.tickets("alpha"), + task_type="retrieval_query", + top_k=1, + distance_type="dot_product", + ) + + +class TestNamingATable(SparkTestCase): + """A table can be named rather than passed, as in BigQuery.""" + + def test_a_registered_view_is_searched(self): + self.tickets("alpha", "beta").createOrReplaceTempView("ai_search_view") + hits = self.search("ai_search_view", top_k=1) + self.assertEqual(self.ids(hits), [1]) + + def test_an_empty_name_is_rejected(self): + with self.assertRaises(ValueError): + self.search(" ") + + +class TestTheModelIsCalledOncePerRow(SparkTestCase): + """Reading a row's vector twice must not embed it twice.""" + + def test_no_distance_evaluates_the_model_more_than_once(self): + # The cosine distance reads each vector twice, once for the dot + # product and once for its norm. Spark must not answer that by + # calling the model twice per row. + table = self.tickets("alpha", "beta") + for distance_type in ai_search_module.DISTANCE_TYPES: + with self.subTest(distance_type=distance_type): + hits = self.search(table, distance_type=distance_type) + plan = str( + hits._jdf.queryExecution().optimizedPlan() # pylint: disable=protected-access + ) + self.assertEqual(plan.count("_predict"), 1, plan) + + +if __name__ == "__main__": + unittest.main()