diff --git a/tests/test_lookup_variants.py b/tests/test_lookup_variants.py new file mode 100644 index 00000000..c7afe297 --- /dev/null +++ b/tests/test_lookup_variants.py @@ -0,0 +1,291 @@ +"""Offline execution and contract checks for the explicit lookup operations.""" +import importlib +import inspect +import json +from pathlib import Path +import runpy +import shutil +from types import SimpleNamespace + +import jsonschema +import pandas as pd +import pytest +import requests +import yaml + +import wrangles + + +MODEL_ID = '12345678-1234-1234' +VARIANTS = [('key', 'key', 1000), ('semantic', 'embedding', 20)] +lookup_module = importlib.import_module('wrangles.lookup') +recipe_main = importlib.import_module('wrangles.recipe_wrangles.main') + + +@pytest.fixture +def backend(monkeypatch): + metadata = {'purpose': 'lookup', 'variant': 'key', 'settings': {'columns': ['Value']}} + metadata_calls = [] + calls = [] + + def model(model_id): + metadata_calls.append(model_id) + return metadata + + def batch(url, params, values, batch_size): + calls.append((url, params, values, batch_size)) + count = params.get('n') or 1 + if count > 1: + data = [[[{'Value': f'{value}:{i}'} for i in range(1, count + 1)]] for value in values] + else: + data = [[f'found:{value}'] for value in values] + return {'columns': ['Value'], 'data': data} + + def no_network(*args, **kwargs): + pytest.fail('Lookup tests must not call a live service') + + monkeypatch.setattr(requests.sessions.Session, 'request', no_network) + monkeypatch.setattr(lookup_module._data, 'model', model) + monkeypatch.setattr(recipe_main, '_model', model) + monkeypatch.setattr(lookup_module._batching, 'batch_api_calls', batch) + return SimpleNamespace(metadata=metadata, metadata_calls=metadata_calls, calls=calls) + + +@pytest.mark.parametrize('operation,variant,batch_size', VARIANTS) +@pytest.mark.parametrize('values,columns,expected', [ + ('one', None, {'Value': 'found:one'}), + ('one', 'Value', 'found:one'), + (['one', 'two'], ['Value'], [['found:one'], ['found:two']]), + (['one', 'two'], None, [{'Value': 'found:one'}, {'Value': 'found:two'}]), + ([], None, []), +]) +def test_python_variants_preserve_shapes_and_arguments(backend, operation, variant, batch_size, values, columns, expected): + backend.metadata['variant'] = variant + result = getattr(wrangles.lookup, operation)( + values, MODEL_ID, columns=columns, custom_option='kept' + ) + assert result == expected + assert backend.metadata_calls == [MODEL_ID] + assert backend.calls == [( + f'{wrangles.config.api_host}/wrangles/lookup', + {'model_id': MODEL_ID, 'columns': '["Value"]', 'custom_option': 'kept'}, + values if isinstance(values, list) else [values], batch_size, + )] + + +@pytest.mark.parametrize('operation,variant,batch_size', VARIANTS) +def test_python_n_keeps_match_dicts_and_batch_override(backend, operation, variant, batch_size): + backend.metadata.update(variant=variant, batch_size=7) + result = getattr(wrangles.lookup, operation)('one', MODEL_ID, columns='Value', n=2) + assert result == [{'Value': 'one:1'}, {'Value': 'one:2'}] + assert backend.calls[0][1] == {'model_id': MODEL_ID, 'columns': '["Value"]', 'n': 2} + assert backend.calls[0][3] == 7 + + +@pytest.mark.parametrize('layer', ['python', 'recipe']) +@pytest.mark.parametrize('operation,variant', [ + ('key', 'embedding'), ('semantic', 'key'), + ('key', None), ('semantic', None), + ('key', ''), ('semantic', ''), + ('key', 'recipe'), ('semantic', 'semantic'), +]) +def test_explicit_names_reject_incompatible_variants(backend, layer, operation, variant): + backend.metadata['variant'] = variant + with pytest.raises(ValueError, match=rf'lookup\.{operation} requires a model with variant'): + if layer == 'python': + getattr(wrangles.lookup, operation)('one', MODEL_ID) + else: + getattr(wrangles.recipe._recipe_wrangles.lookup, operation)( + pd.DataFrame({'Raw': ['one']}), 'Raw', MODEL_ID, output='Result' + ) + assert backend.metadata_calls == [MODEL_ID] + assert backend.calls == [] + + +@pytest.mark.parametrize('operation', ['key', 'semantic']) +def test_missing_variant_is_not_inferred(backend, operation): + del backend.metadata['variant'] + with pytest.raises(ValueError, match='got None'): + getattr(wrangles.lookup, operation)('one', MODEL_ID) + assert backend.calls == [] + + +@pytest.mark.parametrize('operation,variant,batch_size', VARIANTS) +@pytest.mark.parametrize('layer', ['python', 'recipe']) +def test_wrong_purpose_and_missing_model_fail_before_execution(backend, operation, variant, batch_size, layer): + backend.metadata.update(purpose='extract', variant=variant) + + def call(): + if layer == 'python': + return getattr(wrangles.lookup, operation)('one', MODEL_ID) + return getattr(wrangles.recipe._recipe_wrangles.lookup, operation)( + pd.DataFrame({'Raw': ['one']}), 'Raw', MODEL_ID + ) + + with pytest.raises(ValueError, match='Using extract model_id'): + call() + backend.metadata.clear() + backend.metadata['message'] = 'error' + with pytest.raises(ValueError, match='Incorrect model_id'): + call() + assert backend.calls == [] + + +@pytest.mark.parametrize('variant,batch_size', [('key', 1000), ('embedding', 20), (None, 1000), ('recipe', 1000)]) +def test_legacy_lookup_keeps_model_dispatch(backend, variant, batch_size): + backend.metadata['variant'] = variant + assert wrangles.lookup('one', MODEL_ID, columns='Value') == 'found:one' + result = wrangles.recipe.run( + {'wrangles': [{'lookup': {'input': 'Raw', 'output': 'Value', 'model_id': MODEL_ID}}]}, + dataframe=pd.DataFrame({'Raw': ['one', 'two']}), + ) + assert result.to_dict('list') == {'Raw': ['one', 'two'], 'Value': ['found:one', 'found:two']} + assert [call[3] for call in backend.calls] == [batch_size, batch_size] + + +def test_legacy_n_and_passthrough_stay_supported(backend): + expected = [{'Value': 'one:1'}, {'Value': 'one:2'}] + assert wrangles.lookup('one', MODEL_ID, n=2, custom_option='kept') == expected + result = wrangles.recipe.run( + {'wrangles': [{'lookup': { + 'input': 'Raw', 'output': 'Result', 'model_id': MODEL_ID, + 'n': 2, 'custom_option': 'kept', + }}]}, + dataframe=pd.DataFrame({'Raw': ['one']}), + ) + assert result.to_dict('list') == {'Raw': ['one'], 'Result': [expected]} + assert [call[1] for call in backend.calls] == [ + {'model_id': MODEL_ID, 'columns': '["Value"]', 'n': 2, 'custom_option': 'kept'}, + {'model_id': MODEL_ID, 'columns': '["Value"]', 'n': 2, 'custom_option': 'kept'}, + ] + + +@pytest.mark.parametrize('operation,variant,batch_size', VARIANTS) +@pytest.mark.parametrize('mode,expected_inputs', [ + ('by_row', [['one', 'one', 'two']]), + ('by_dataframe', [['one', 'two']]), + ('by_matrix', [['one'], ['two']]), +]) +def test_recipe_variants_execute_existing_modes(backend, operation, variant, batch_size, mode, expected_inputs): + backend.metadata['variant'] = variant + result = wrangles.recipe.run( + {'wrangles': [{f'lookup.{operation}': { + 'input': 'Raw', 'output': {'Value': 'Result'}, 'model_id': MODEL_ID, + 'lookup_mode': mode, 'matrix_variables': ['Group'], 'custom_option': 'kept', + }}]}, + dataframe=pd.DataFrame({'Raw': ['one', 'one', 'two'], 'Group': ['a', 'a', 'b']}), + ) + assert result.to_dict('list') == { + 'Raw': ['one', 'one', 'two'], 'Group': ['a', 'a', 'b'], + 'Result': ['found:one', 'found:one', 'found:two'], + } + assert [call[2] for call in backend.calls] == expected_inputs + assert all(call[1] == { + 'model_id': MODEL_ID, 'columns': '["Value"]', 'custom_option': 'kept' + } for call in backend.calls) + assert [call[3] for call in backend.calls] == [batch_size] * len(expected_inputs) + + +@pytest.mark.parametrize('operation,variant,batch_size', VARIANTS) +def test_recipe_n_and_empty_behavior(backend, operation, variant, batch_size): + backend.metadata['variant'] = variant + params = {'input': 'Raw', 'output': 'Match *', 'model_id': MODEL_ID, 'n': 2} + result = wrangles.recipe.run( + {'wrangles': [{f'lookup.{operation}': params}]}, + dataframe=pd.DataFrame({'Raw': ['one', 'two']}), + ) + assert result.to_dict('list') == { + 'Raw': ['one', 'two'], + 'Match 1': [{'Value': 'one:1'}, {'Value': 'two:1'}], + 'Match 2': [{'Value': 'one:2'}, {'Value': 'two:2'}], + } + assert backend.calls[0][1]['n'] == 2 + backend.calls.clear() + backend.metadata_calls.clear() + empty = getattr(wrangles.recipe._recipe_wrangles.lookup, operation)( + pd.DataFrame({'Raw': []}), **params + ) + assert empty.to_dict('list') == {'Raw': [], 'Match 1': [], 'Match 2': []} + assert backend.metadata_calls == [] + assert backend.calls == [] + + +@pytest.mark.parametrize('operation', ['key', 'semantic']) +@pytest.mark.parametrize('mode', ['by_dataframe', 'by_matrix']) +def test_explicit_recipe_rejects_silently_ignored_n(backend, operation, mode): + with pytest.raises(ValueError, match='n is only supported with lookup_mode=by_row'): + getattr(wrangles.recipe._recipe_wrangles.lookup, operation)( + pd.DataFrame({'Raw': ['one']}), 'Raw', MODEL_ID, n=2, lookup_mode=mode + ) + assert backend.metadata_calls == [] + assert backend.calls == [] + + +@pytest.mark.parametrize('operation,variant,batch_size', VARIANTS) +def test_dataframe_accessor_preserves_source(backend, operation, variant, batch_size): + backend.metadata['variant'] = variant + source = wrangles.DataFrame({'Raw': ['one', 'two']}) + result = getattr(source.wrangles.lookup, operation)('Raw', MODEL_ID, output='Value') + legacy = source.wrangles.lookup(input='Raw', model_id=MODEL_ID, output='Value') + assert result.to_dict('list') == {'Raw': ['one', 'two'], 'Value': ['found:one', 'found:two']} + assert legacy.equals(result) + assert source.to_dict('list') == {'Raw': ['one', 'two']} + + +@pytest.mark.parametrize('operation,variant,batch_size', VARIANTS) +def test_recipe_where_only_executes_selected_rows(backend, operation, variant, batch_size): + backend.metadata['variant'] = variant + result = wrangles.recipe.run( + {'wrangles': [{f'lookup.{operation}': { + 'input': 'Raw', 'output': 'Value', 'model_id': MODEL_ID, + 'where': 'Apply = true', + }}]}, + dataframe=pd.DataFrame({'Raw': ['one', 'two'], 'Apply': [True, False]}), + ) + assert result.to_dict('list') == { + 'Raw': ['one', 'two'], 'Apply': [True, False], 'Value': ['found:one', ''], + } + assert [call[2] for call in backend.calls] == [['one']] + + +def test_generated_schema_preserves_legacy_and_explicit_contracts(tmp_path, monkeypatch, backend): + root = Path(__file__).resolve().parents[1] + shutil.copyfile(root / 'schema/recipe_base_schema.json', tmp_path / 'recipe_base_schema.json') + monkeypatch.chdir(tmp_path) + + def schema_response(url): + assert url == 'http://json-schema.org/draft-07/schema#' + return SimpleNamespace(json=lambda: jsonschema.Draft7Validator.META_SCHEMA) + + monkeypatch.setattr(requests, 'get', schema_response) + runpy.run_path(str(root / 'schema/generate_recipe_schema.py')) + generated = json.loads((tmp_path / 'schema.json').read_text()) + contracts = generated['$defs']['wrangles']['items']['properties'] + assert {'lookup', 'lookup.key', 'lookup.semantic'} <= contracts.keys() + assert contracts['lookup']['description'] == 'Lookup values from a saved lookup wrangle' + for operation, variant, _ in VARIANTS: + method = getattr(wrangles.recipe._recipe_wrangles.lookup, operation) + contract = contracts[f'lookup.{operation}'] + raw = yaml.safe_load(method.__doc__) + assert f'stored variant {variant}' in contract['description'] + assert contract['required'] == ['input', 'model_id'] + assert set(inspect.signature(method).parameters) - {'df', 'kwargs'} <= contract['properties'].keys() + for name in ['output', 'lookup_mode', 'n']: + assert contract['properties'][name]['default'] == inspect.signature(method).parameters[name].default + jsonschema.Draft7Validator.check_schema(raw) + validator = jsonschema.Draft7Validator(raw) + validator.validate(raw['examples'][0]) + with pytest.raises(jsonschema.ValidationError, match='model_id'): + validator.validate({'input': 'Raw'}) + with pytest.raises(jsonschema.ValidationError): + validator.validate({'input': ['First', 'Second'], 'model_id': MODEL_ID}) + with pytest.raises(jsonschema.ValidationError): + validator.validate({'input': 'Raw', 'model_id': MODEL_ID, 'n': 2, 'lookup_mode': 'by_dataframe'}) + with pytest.raises(jsonschema.ValidationError, match='matrix_variables'): + validator.validate({'input': 'Raw', 'model_id': MODEL_ID, 'lookup_mode': 'by_matrix'}) + validator.validate({ + 'input': 'Raw', 'model_id': MODEL_ID, 'lookup_mode': 'by_matrix', + 'matrix_variables': ['Group'], 'n': None, + }) + assert {'if', 'where', 'where_params'} <= contract['properties'].keys() + assert backend.calls == [] diff --git a/wrangles/lookup.py b/wrangles/lookup.py index 49a6eb58..374509fb 100644 --- a/wrangles/lookup.py +++ b/wrangles/lookup.py @@ -21,6 +21,21 @@ def lookup( :param n: (Optional) Number of matches to return per input. When > 1, returns a list of n dicts per input - each match is always a dict, even if a single column is requested. """ + return _lookup(None, input, model_id, columns, n, **kwargs) + + +def _validate_lookup_variant(metadata, operation): + """Check explicit operations against the currently stored model variants.""" + expected = 'key' if operation == 'key' else 'embedding' + actual = metadata.get('variant') + if actual != expected: + raise ValueError( + f"lookup.{operation} requires a model with variant '{expected}'; " + f"got {actual!r}. Use lookup for legacy or unknown variants." + ) + + +def _lookup(expected_variant, input, model_id, columns=None, n=None, **kwargs): # Check if user has entered a single input or multiple inputs single_input = False if not isinstance(input, list): @@ -60,6 +75,9 @@ def lookup( f'Using {purpose} model_id {model_id} in a lookup wrangle.' ) + if expected_variant is not None: + _validate_lookup_variant(metadata, expected_variant) + _logging.info(f": Looking up {len(input)} values :: model_id :: {model_id}") if n: @@ -101,3 +119,58 @@ def lookup( results = results[0] return results + + +def key( + input: _Union[str, list], + model_id: str, + columns: _Union[str, list] = None, + n: int = None, + **kwargs +) -> _Union[str, list]: + """ + Look up values using a saved key lookup model. Requires WrangleWorks Account. + + The model must have purpose ``lookup`` and stored variant ``key``. + Missing or unknown variants are rejected; legacy callers can use ``lookup``. + + :param input: A value or list of values to look up. + :param model_id: The saved key lookup model to execute. + :param columns: Return one named column, a list of columns, or all columns + as a dictionary when omitted. + :param n: Optional number of matches, using the existing lookup behavior. + When greater than one, each input produces a list of match dictionaries. + :param kwargs: Additional options forwarded to the existing lookup API. + :return: One result for scalar input, or a list of results for list input. + """ + return _lookup('key', input, model_id, columns, n, **kwargs) + + +def semantic( + input: _Union[str, list], + model_id: str, + columns: _Union[str, list] = None, + n: int = None, + **kwargs +) -> _Union[str, list]: + """ + Look up values using a saved semantic model. Requires WrangleWorks Account. + + The model must have purpose ``lookup`` and stored variant ``embedding``. + Missing or unknown variants are rejected; legacy callers can use ``lookup``. + + :param input: A value or list of values to look up semantically. + :param model_id: The saved semantic lookup model to execute. + :param columns: Return one named column, a list of columns, or all columns + as a dictionary when omitted. + :param n: Optional number of matches. When greater than one, each input + produces a list of match dictionaries, even for a single named column. + :param kwargs: Additional options forwarded to the existing lookup API. + :return: One result for scalar input, or a list of results for list input. + """ + return _lookup('semantic', input, model_id, columns, n, **kwargs) + + +# Preserve wrangles.lookup(...) while exposing the explicit dotted operations. +lookup.key = key +lookup.semantic = semantic diff --git a/wrangles/recipe_wrangles/__init__.py b/wrangles/recipe_wrangles/__init__.py index 46f1e5e2..b8ab77af 100644 --- a/wrangles/recipe_wrangles/__init__.py +++ b/wrangles/recipe_wrangles/__init__.py @@ -17,6 +17,8 @@ """ from .standardize import clean as _standardize_clean from .standardize import custom as _standardize_custom +from .lookup import key as _lookup_key +from .lookup import semantic as _lookup_semantic from .main import * from .pandas import * from . import convert @@ -36,3 +38,6 @@ # standardize namespace to recipe resolution and schema discovery. standardize.clean = _standardize_clean standardize.custom = _standardize_custom + +lookup.key = _lookup_key +lookup.semantic = _lookup_semantic diff --git a/wrangles/recipe_wrangles/lookup.py b/wrangles/recipe_wrangles/lookup.py new file mode 100644 index 00000000..d02f6a8b --- /dev/null +++ b/wrangles/recipe_wrangles/lookup.py @@ -0,0 +1,165 @@ +"""Explicit saved-model lookup operations for recipes.""" +from typing import Union as _Union + +import pandas as _pd + +from .main import _lookup_with_variant + + +def key( + df: _pd.DataFrame, + input: _Union[str, int, list], + model_id: str, + output: _Union[str, list] = None, + lookup_mode: str = 'by_row', + n: int = None, + **kwargs +) -> _pd.DataFrame: + """ + type: object + description: >- + Look up values using a saved key lookup model. Requires a WrangleWorks + account and a model with purpose lookup and stored variant key. + Missing or unknown variants are rejected; use lookup for legacy models. + additionalProperties: true + required: [input, model_id] + allOf: + - if: + required: [n] + properties: + n: {type: integer} + then: + properties: + lookup_mode: {const: by_row} + - if: + required: [lookup_mode] + properties: + lookup_mode: {const: by_matrix} + then: + required: [matrix_variables] + properties: + input: + type: [string, integer, array] + items: + type: [string, integer] + minItems: 1 + maxItems: 1 + description: One input column, optionally expressed as a single-item list. + model_id: + type: string + description: Existing saved key lookup model ID, not a catalog ID. + output: + type: [string, array, 'null'] + default: null + description: >- + Output column or columns; defaults to the input column. Named lookup + columns return their values; other names receive match dictionaries. + With n greater than one, n output columns or a wildcard name such as + 'Match *' distribute match dictionaries across columns. + lookup_mode: + type: string + default: by_row + enum: [by_row, by_dataframe, by_matrix] + description: >- + Look up each row, each distinct input, or each matrix permutation. + by_matrix requires matrix_variables in the forwarded options. + matrix_variables: + type: array + minItems: 1 + items: {type: string} + description: Columns that define the permutations for by_matrix mode. + n: + type: [integer, 'null'] + default: null + description: >- + Optional number of matches using the existing lookup behavior. + Supported only with by_row. When greater than one, a single output + contains a list of match dictionaries unless expanded with a wildcard. + examples: + - input: ProductCode + output: KeyResult + model_id: '${KEY_LOOKUP_MODEL_ID}' + """ + return _lookup_with_variant( + 'key', df, input, output, model_id, lookup_mode, n, **kwargs + ) + + +def semantic( + df: _pd.DataFrame, + input: _Union[str, int, list], + model_id: str, + output: _Union[str, list] = None, + lookup_mode: str = 'by_row', + n: int = None, + **kwargs +) -> _pd.DataFrame: + """ + type: object + description: >- + Look up values using a saved semantic lookup model. Requires a WrangleWorks + account and a model with purpose lookup and stored variant embedding. + Missing or unknown variants are rejected; use lookup for legacy models. + additionalProperties: true + required: [input, model_id] + allOf: + - if: + required: [n] + properties: + n: {type: integer} + then: + properties: + lookup_mode: {const: by_row} + - if: + required: [lookup_mode] + properties: + lookup_mode: {const: by_matrix} + then: + required: [matrix_variables] + properties: + input: + type: [string, integer, array] + items: + type: [string, integer] + minItems: 1 + maxItems: 1 + description: One input column, optionally expressed as a single-item list. + model_id: + type: string + description: Existing saved semantic lookup model ID, not a catalog ID. + output: + type: [string, array, 'null'] + default: null + description: >- + Output column or columns; defaults to the input column. Named lookup + columns return their values; other names receive match dictionaries. + With n greater than one, n output columns or a wildcard name such as + 'Match *' distribute match dictionaries across columns. + lookup_mode: + type: string + default: by_row + enum: [by_row, by_dataframe, by_matrix] + description: >- + Look up each row, each distinct input, or each matrix permutation. + by_matrix requires matrix_variables in the forwarded options. + matrix_variables: + type: array + minItems: 1 + items: {type: string} + description: Columns that define the permutations for by_matrix mode. + n: + type: [integer, 'null'] + default: null + description: >- + Optional number of matches. Supported only with by_row. When greater + than one, a single output contains a list of match dictionaries unless + expanded with a wildcard. Each of n output columns receives one match. + examples: + - input: Description + output: 'Match *' + model_id: '${SEMANTIC_LOOKUP_MODEL_ID}' + n: 3 + """ + return _lookup_with_variant( + 'semantic', df, input, output, model_id, lookup_mode, n, **kwargs + ) diff --git a/wrangles/recipe_wrangles/main.py b/wrangles/recipe_wrangles/main.py index 28cfbc9a..df9df2bd 100644 --- a/wrangles/recipe_wrangles/main.py +++ b/wrangles/recipe_wrangles/main.py @@ -25,6 +25,7 @@ from ..translate import translate as _translate from ..data import model as _model from ..lookup import lookup as _lookup +from ..lookup import _validate_lookup_variant from .. import extract as _extract from .. import recipe as _recipe from .convert import to_json as _to_json @@ -987,6 +988,20 @@ def lookup( - by_matrix - by_dataframe """ + return _lookup_with_variant( + None, df, input, output, model_id, lookup_mode, n, **kwargs + ) + + +def _lookup_with_variant( + expected_variant, df, input, output=None, model_id=None, + lookup_mode='by_row', n=None, **kwargs +): + # The legacy implementation only forwards n in by_row mode. Explicit + # operations reject unsupported combinations instead of silently ignoring n. + if expected_variant is not None and n is not None and lookup_mode != 'by_row': + raise ValueError('n is only supported with lookup_mode=by_row') + # Ensure input is only 1 value if isinstance(input, list): if len(input) == 1: @@ -1027,6 +1042,9 @@ def lookup( purpose = metadata['purpose'] if purpose != 'lookup': raise ValueError(f'Using {purpose} model_id {model_id} in a lookup function.') + + if expected_variant is not None: + _validate_lookup_variant(metadata, expected_variant) # Split input/output if user differentiated e.g. "wrangle_column: output_column" wrangle_output = [