From 64534a2a3b88c055b15cf006032ab5c5406eda86 Mon Sep 17 00:00:00 2001 From: Leonid Molotiievskyi Date: Mon, 29 Jun 2026 17:22:20 +0200 Subject: [PATCH 01/15] enable package build on main branch --- .github/workflows/publish-dev-rc.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/publish-dev-rc.yml b/.github/workflows/publish-dev-rc.yml index 1a79c868..3a314f1d 100644 --- a/.github/workflows/publish-dev-rc.yml +++ b/.github/workflows/publish-dev-rc.yml @@ -23,10 +23,10 @@ jobs: outputs: rc_version: ${{ steps.compute.outputs.rc_version }} steps: - - name: Guard – dev branch only - if: github.ref != 'refs/heads/dev' + - name: Guard – dev or main branch only + if: github.ref != 'refs/heads/dev' && github.ref != 'refs/heads/main' run: | - echo "::error::This workflow may only be triggered from the 'dev' branch. Current ref: ${{ github.ref }}" + echo "::error::This workflow may only be triggered from the 'dev' or 'main' branch. Current ref: ${{ github.ref }}" exit 1 - name: Validate version input format @@ -105,8 +105,8 @@ jobs: # ── 2. Pip-install smoke test ───────────────────────────────────────────── # NOTE: The full pytest suite is intentionally NOT run here. This workflow is - # guarded to the 'dev' branch, and every commit on 'dev' is already tested by - # publish-dev.yml + # guarded to the 'dev' and 'main' branches, and every commit on those branches + # is already tested by publish-dev.yml / publish-main.yml respectively. test-pip-install: name: Smoke Test – pip install runs-on: ubuntu-latest From b434280e60686b99bc8e4a7f4bd3cd2967581f1f Mon Sep 17 00:00:00 2001 From: Leonid Molotiievskyi Date: Mon, 29 Jun 2026 21:10:36 +0200 Subject: [PATCH 02/15] Add ability to trigger the deploy job directly from the wranglerspy repository (#1031) --- .github/workflows/publish-dev-rc.yml | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/.github/workflows/publish-dev-rc.yml b/.github/workflows/publish-dev-rc.yml index 3a314f1d..2cfaff17 100644 --- a/.github/workflows/publish-dev-rc.yml +++ b/.github/workflows/publish-dev-rc.yml @@ -108,7 +108,7 @@ jobs: # guarded to the 'dev' and 'main' branches, and every commit on those branches # is already tested by publish-dev.yml / publish-main.yml respectively. test-pip-install: - name: Smoke Test – pip install + name: smoke test runs-on: ubuntu-latest needs: [compute-version] permissions: @@ -130,7 +130,7 @@ jobs: # ── 3. Publish Python package to CodeArtifact ──────────────────────────── publish-codeartifact: - name: Publish RC to CodeArtifact + name: Publish RC Package runs-on: ubuntu-latest needs: [compute-version, test-pip-install] permissions: @@ -173,4 +173,26 @@ jobs: --repository ${{ vars.CODEARTIFACT_REPOSITORY }} - name: Publish to CodeArtifact - run: twine upload --repository codeartifact dist/* \ No newline at end of file + run: twine upload --repository codeartifact dist/* + + # ── 4. Trigger QA deployment in Lambda-Recipes ─────────────────────────── + # CROSS_REPO_PAT: a GitHub PAT with Actions read/write on Lambda-Recipes, + # stored as a repository or org secret. Required to dispatch workflows across repositories. + trigger-deploy-qa: + name: Trigger QA Deploy + runs-on: ubuntu-latest + needs: [compute-version, publish-codeartifact] + steps: + - name: Deploy QA + uses: convictional/trigger-workflow-and-wait@v1.6.5 + with: + owner: wrangleworks + repo: Lambda-Recipes + workflow_file_name: deploy-qa.yml + github_token: ${{ secrets.CROSS_REPO_PAT }} + ref: main + wait_interval: 15 + client_payload: '{"wrangles_version": "${{ needs.compute-version.outputs.rc_version }}"}' + propagate_failure: true + trigger_workflow: true + wait_workflow: true \ No newline at end of file From 74b94c8a9a15baa43df976ba789484eea1135c7f Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Wed, 1 Jul 2026 00:12:53 +0300 Subject: [PATCH 03/15] 825 add connectors for duckdb access (#1024) * Add DuckDB and Microsoft Access connectors Add DuckDB read/write/run connector with recipe schemas Add Microsoft Access read/write/run connector via ODBC Export new connectors from wrangles.connectors Document optional duckdb and pyodbc dependencies Add connector tests for DuckDB and Access * fix translate tests * fix duckdb tests --------- Co-authored-by: Eric Hills <53243273+ebhills@users.noreply.github.com> --- README.md | 2 + requirements-full.txt | 2 + tests/connectors/test_access.py | 109 ++++++++++++ tests/connectors/test_duckdb.py | 72 ++++++++ wrangles/connectors/__init__.py | 4 +- wrangles/connectors/access.py | 298 ++++++++++++++++++++++++++++++++ wrangles/connectors/duckdb.py | 205 ++++++++++++++++++++++ 7 files changed, 691 insertions(+), 1 deletion(-) create mode 100644 tests/connectors/test_access.py create mode 100644 tests/connectors/test_duckdb.py create mode 100644 wrangles/connectors/access.py create mode 100644 wrangles/connectors/duckdb.py diff --git a/README.md b/README.md index dab5f823..1469b7d1 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,8 @@ Connectors for databases, cloud storage, and external services require additiona | Capability | Install | |---|---| | Microsoft SQL Server | `pip install pymssql sqlalchemy` | +| Microsoft Access | `pip install pyodbc` | +| DuckDB | `pip install duckdb` | | PostgreSQL | `pip install psycopg2-binary sqlalchemy` | | MySQL | `pip install pymysql sqlalchemy` | | MongoDB | `pip install pymongo[srv]` | diff --git a/requirements-full.txt b/requirements-full.txt index 38b05081..15406ae9 100644 --- a/requirements-full.txt +++ b/requirements-full.txt @@ -8,5 +8,7 @@ # SQL databases sqlalchemy>=2.0,<3.0 +duckdb>=1.0.0 +pyodbc>=5.0.0 pymssql>=2.3.3 psycopg2-binary>=2.9.10 diff --git a/tests/connectors/test_access.py b/tests/connectors/test_access.py new file mode 100644 index 00000000..ec31891f --- /dev/null +++ b/tests/connectors/test_access.py @@ -0,0 +1,109 @@ +import pandas as pd +from types import SimpleNamespace +from unittest.mock import Mock + +from wrangles.connectors import access + + +class MockTables: + def __init__(self, exists): + self.exists = exists + + def fetchone(self): + return object() if self.exists else None + + +class MockCursor: + def __init__(self, table_exists=False): + self.table_exists = table_exists + self.executed = [] + self.executemany_calls = [] + + def tables(self, table=None, tableType=None): + return MockTables(self.table_exists) + + def execute(self, sql, params=()): + self.executed.append((sql, params)) + return self + + def executemany(self, sql, rows): + self.executemany_calls.append((sql, list(rows))) + return self + + +class MockConnection: + def __init__(self, cursor): + self.cursor_obj = cursor + self.committed = False + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def cursor(self): + return self.cursor_obj + + def commit(self): + self.committed = True + + +def test_connection_string_requires_database_or_connection_string(): + try: + access._connection_string() + assert False + except ValueError as e: + assert str(e) == 'database or connection_string must be provided' + + +def test_read_sql(monkeypatch): + data = pd.DataFrame({'Col1': ['Data1', 'Data2']}) + mock_connect = Mock(return_value=MockConnection(MockCursor())) + monkeypatch.setattr(access, "_pyodbc", SimpleNamespace(connect=mock_connect)) + monkeypatch.setattr(pd, "read_sql", Mock(return_value=data)) + + df = access.read( + database='database.accdb', + command='SELECT * from df_mock' + ) + + assert df.equals(data) + mock_connect.assert_called_once_with('DRIVER={Microsoft Access Driver (*.mdb, *.accdb)};DBQ=database.accdb;') + + +def test_write_sql_creates_table_and_inserts(monkeypatch): + cursor = MockCursor(table_exists=False) + monkeypatch.setattr(access, "_pyodbc", SimpleNamespace(connect=Mock(return_value=MockConnection(cursor)))) + + result = access.write( + df=pd.DataFrame({'Col1': ['Data1', 'Data2'], 'Col2': [1, 2]}), + database='database.accdb', + table='WrWx' + ) + + assert result is None + assert cursor.executed[0][0] == 'CREATE TABLE [WrWx] ([Col1] LONGTEXT, [Col2] INTEGER)' + assert cursor.executemany_calls[0] == ( + 'INSERT INTO [WrWx] ([Col1], [Col2]) VALUES (?, ?)', + [('Data1', 1), ('Data2', 2)] + ) + + +def test_run(monkeypatch): + cursor = MockCursor() + connection = MockConnection(cursor) + monkeypatch.setattr(access, "_pyodbc", SimpleNamespace(connect=Mock(return_value=connection))) + + result = access.run( + database='database.accdb', + command=['DELETE FROM WrWx WHERE Col1 = ?', 'UPDATE WrWx SET Col1 = ?'], + params=('Data1',) + ) + + assert result is None + assert cursor.executed == [ + ('DELETE FROM WrWx WHERE Col1 = ?', ('Data1',)), + ('UPDATE WrWx SET Col1 = ?', ('Data1',)) + ] + assert connection.committed diff --git a/tests/connectors/test_duckdb.py b/tests/connectors/test_duckdb.py new file mode 100644 index 00000000..fa7272f4 --- /dev/null +++ b/tests/connectors/test_duckdb.py @@ -0,0 +1,72 @@ +import importlib.util + +import pandas as pd +import pytest + +from wrangles.connectors import duckdb + + +pytestmark = pytest.mark.skipif( + importlib.util.find_spec('duckdb') is None, + reason='duckdb optional dependency is not installed' +) + + +def test_read_sql(tmp_path): + database = tmp_path / 'test.duckdb' + duckdb.run( + database=str(database), + command='CREATE TABLE df_mock AS SELECT \'Data1\' AS Col1, 1 AS Col2 UNION ALL SELECT \'Data2\', 2' + ) + + df = duckdb.read( + database=str(database), + command='SELECT * from df_mock ORDER BY Col2' + ) + + assert df.equals( + pd.DataFrame( + { + 'Col1': ['Data1', 'Data2'], + 'Col2': pd.array([1, 2], dtype='int32') + } + ) + ) + + +def test_write_sql(tmp_path): + database = tmp_path / 'write.duckdb' + df = pd.DataFrame({'Col1': ['Data1', 'Data2']}) + + duckdb.write( + df=df, + database=str(database), + table='temp_mock' + ) + + assert duckdb.read( + database=str(database), + command='SELECT * from temp_mock' + ).equals(df) + + +def test_run(tmp_path): + database = tmp_path / 'run.duckdb' + df = pd.DataFrame({'Col1': ['Data1', 'Data2']}) + duckdb.write( + df=df, + database=str(database), + table='test_table' + ) + + duckdb.run( + database=str(database), + command='CREATE TABLE test_table_copy AS SELECT * FROM test_table' + ) + + df_copy = duckdb.read( + database=str(database), + command='SELECT * from test_table_copy' + ) + + assert df.equals(df_copy) diff --git a/wrangles/connectors/__init__.py b/wrangles/connectors/__init__.py index e8039f7f..1c50c4c0 100644 --- a/wrangles/connectors/__init__.py +++ b/wrangles/connectors/__init__.py @@ -3,8 +3,10 @@ """ from . import akeneo +from . import access from . import ckan from . import concurrent +from . import duckdb from . import excel from . import file from . import http @@ -26,4 +28,4 @@ from . import train from . import jinja from . import _formatting -from . import input \ No newline at end of file +from . import input diff --git a/wrangles/connectors/access.py b/wrangles/connectors/access.py new file mode 100644 index 00000000..f7485f44 --- /dev/null +++ b/wrangles/connectors/access.py @@ -0,0 +1,298 @@ +""" +Connector to read/write from Microsoft Access databases using ODBC. +""" +import logging as _logging +from typing import Union as _Union + +import pandas as _pd +from pandas.api import types as _pd_types + +from ..utils import ( + LazyLoader as _LazyLoader, + wildcard_expansion as _wildcard_expansion, +) + + +_pyodbc = _LazyLoader('pyodbc') + +_schema = {} + + +def _quote_identifier(identifier: str) -> str: + return f"[{str(identifier).replace(']', ']]')}]" + + +def _connection_string( + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, +) -> str: + if connection_string: + return connection_string + if database is None: + raise ValueError('database or connection_string must be provided') + + conn = f"DRIVER={{{driver}}};DBQ={database};" + if password: + conn += f"PWD={password};" + return conn + + +def _access_type(dtype) -> str: + if _pd_types.is_bool_dtype(dtype): + return 'BIT' + if _pd_types.is_integer_dtype(dtype): + return 'INTEGER' + if _pd_types.is_float_dtype(dtype): + return 'DOUBLE' + if _pd_types.is_datetime64_any_dtype(dtype): + return 'DATETIME' + return 'LONGTEXT' + + +def _table_exists(cursor, table: str) -> bool: + return cursor.tables(table=table, tableType='TABLE').fetchone() is not None + + +def _create_table(cursor, table: str, df: _pd.DataFrame) -> None: + columns = ', '.join( + f"{_quote_identifier(column)} {_access_type(dtype)}" + for column, dtype in df.dtypes.items() + ) + cursor.execute(f"CREATE TABLE {_quote_identifier(table)} ({columns})") + + +def read( + command: str, + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, + columns: _Union[str, list] = None, + params: _Union[list, tuple] = None, + **kwargs +) -> _pd.DataFrame: + """ + Read data from a Microsoft Access database. + + >>> from wrangles.connectors import access + >>> df = access.read(database='database.accdb', command='SELECT * FROM table') + + :param command: SQL command to select data. + :param database: Access database file path. Not required if connection_string is supplied. + :param connection_string: Full ODBC connection string. If provided, database, driver, and password are ignored. + :param driver: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + :param password: Optional database password. + :param columns: (Optional) Subset of columns to be returned. This is less efficient than specifying in the SQL command. + :param params: (Optional) Variables to pass to a parameterized query. + """ + target = database or connection_string + _logging.info(f": Reading data from Microsoft Access :: {target}") + + conn_string = _connection_string(database, connection_string, driver, password) + with _pyodbc.connect(conn_string) as conn: + df = _pd.read_sql(command, conn, params=params, **kwargs) + + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + return df + + +_schema['read'] = r""" +type: object +description: Import data from a Microsoft Access Database +required: + - command +properties: + database: + type: string + description: Access database file path. Not required if connection_string is supplied. + connection_string: + type: string + description: Full ODBC connection string. If provided, database, driver, and password are ignored. + driver: + type: string + description: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + password: + type: string + description: Optional database password. + command: + type: string + description: |- + SQL command to select data. + Note - using variables here can make your recipe vulnerable + to sql injection. Use params if using variables from + untrusted sources. + columns: + type: + - string + - array + description: A list with a subset of the columns to import. This is less efficient than specifying in the command. + params: + type: array + description: Variables to pass to a parameterized query. +""" + + +def write( + df: _pd.DataFrame, + table: str, + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, + action: str = 'INSERT', + columns: _Union[str, list] = None, +) -> None: + """ + Write data to a Microsoft Access database. + + >>> from wrangles.connectors import access + >>> access.write(df, database='database.accdb', table='table') + + :param df: Pandas Dataframe to be written. + :param table: Table to be exported to. + :param database: Access database file path. Not required if connection_string is supplied. + :param connection_string: Full ODBC connection string. If provided, database, driver, and password are ignored. + :param driver: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + :param password: Optional database password. + :param action: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + :param columns: (Optional) Subset of the columns to be written. If not provided, all columns will be output. + """ + target = database or connection_string + _logging.info(f": Writing data to Microsoft Access :: {target} / {table}") + + action = action.upper() + if action not in ('INSERT', 'REPLACE', 'FAIL'): + raise ValueError('Invalid action. Expected INSERT, REPLACE, or FAIL.') + + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + conn_string = _connection_string(database, connection_string, driver, password) + with _pyodbc.connect(conn_string) as conn: + cursor = conn.cursor() + exists = _table_exists(cursor, table) + + if action == 'FAIL' and exists: + raise ValueError(f"Table already exists: {table}") + if action == 'REPLACE' and exists: + cursor.execute(f"DROP TABLE {_quote_identifier(table)}") + exists = False + if not exists: + _create_table(cursor, table, df) + + if not df.empty: + table_name = _quote_identifier(table) + column_names = ', '.join(_quote_identifier(column) for column in df.columns) + placeholders = ', '.join('?' for _ in df.columns) + sql = f"INSERT INTO {table_name} ({column_names}) VALUES ({placeholders})" + cursor.executemany(sql, df.where(_pd.notnull(df), None).itertuples(index=False, name=None)) + + conn.commit() + + +_schema['write'] = """ +type: object +description: Export data to a Microsoft Access Database +required: + - table +properties: + database: + type: string + description: Access database file path. Not required if connection_string is supplied. + connection_string: + type: string + description: Full ODBC connection string. If provided, database, driver, and password are ignored. + driver: + type: string + description: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + password: + type: string + description: Optional database password. + table: + type: string + description: The table to write to + action: + type: string + description: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + enum: + - INSERT + - REPLACE + - FAIL + columns: + type: + - string + - array + description: A list of the columns to write to the table. If omitted, all columns will be written. +""" + + +def run( + command: _Union[str, list], + database: str = None, + connection_string: str = None, + driver: str = 'Microsoft Access Driver (*.mdb, *.accdb)', + password: str = None, + params: _Union[list, tuple] = None, +) -> None: + """ + Run a command on a Microsoft Access database. + + >>> wrangles.connectors.access.run( + >>> database='database.accdb', + >>> command='' + >>> ) + + :param command: SQL command or a list of SQL commands to execute. + :param database: Access database file path. Not required if connection_string is supplied. + :param connection_string: Full ODBC connection string. If provided, database, driver, and password are ignored. + :param driver: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + :param password: Optional database password. + :param params: Variables to pass to a parameterized query. + """ + target = database or connection_string + _logging.info(f": Executing Microsoft Access Command :: {target}") + + if isinstance(command, str): + command = [command] + + conn_string = _connection_string(database, connection_string, driver, password) + with _pyodbc.connect(conn_string) as conn: + cursor = conn.cursor() + for sql in command: + cursor.execute(sql, params or ()) + conn.commit() + + +_schema['run'] = r""" +type: object +description: Run a command against a Microsoft Access Database +required: + - command +properties: + database: + type: string + description: Access database file path. Not required if connection_string is supplied. + connection_string: + type: string + description: Full ODBC connection string. If provided, database, driver, and password are ignored. + driver: + type: string + description: ODBC driver name. Defaults to Microsoft Access Driver (*.mdb, *.accdb). + password: + type: string + description: Optional database password. + command: + type: + - string + - array + description: SQL command or a list of SQL commands to execute + params: + type: array + description: Variables to pass to a parameterized query. +""" diff --git a/wrangles/connectors/duckdb.py b/wrangles/connectors/duckdb.py new file mode 100644 index 00000000..22d1115c --- /dev/null +++ b/wrangles/connectors/duckdb.py @@ -0,0 +1,205 @@ +""" +Connector to read/write from DuckDB database files. +""" +import logging as _logging +from typing import Union as _Union + +import pandas as _pd + +from ..utils import ( + LazyLoader as _LazyLoader, + wildcard_expansion as _wildcard_expansion, +) + + +_duckdb = _LazyLoader('duckdb') + +_schema = {} + + +def _quote_identifier(identifier: str) -> str: + return '.'.join( + f'"{part.replace(chr(34), chr(34) * 2)}"' + for part in str(identifier).split('.') + ) + + +def read( + database: str, + command: str, + columns: _Union[str, list] = None, + params: _Union[list, tuple, dict] = None, + **kwargs +) -> _pd.DataFrame: + """ + Read data from a DuckDB database. + + >>> from wrangles.connectors import duckdb + >>> df = duckdb.read(database='database.duckdb', command='SELECT * FROM table') + + :param database: The database to connect to including the file path. Use ':memory:' for an in-memory database. + :param command: SQL command to select data. + :param columns: (Optional) Subset of columns to be returned. This is less efficient than specifying in the SQL command. + :param params: (Optional) Variables to pass to a parameterized query. + """ + _logging.info(f": Reading data from DuckDB :: {database}") + + with _duckdb.connect(database=database, **kwargs) as conn: + df = conn.execute(command, params or ()).fetchdf() + + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + return df + + +_schema['read'] = r""" +type: object +description: Import data from a DuckDB Database +required: + - database + - command +properties: + database: + type: string + description: The database to connect to including the file path. Use ':memory:' for an in-memory database. + command: + type: string + description: |- + SQL command to select data. + Note - using variables here can make your recipe vulnerable + to sql injection. Use params if using variables from + untrusted sources. + columns: + type: + - string + - array + description: A list with a subset of the columns to import. This is less efficient than specifying in the command. + params: + type: + - array + - object + description: Variables to pass to a parameterized query. +""" + + +def write( + df: _pd.DataFrame, + database: str, + table: str, + action: str = 'INSERT', + columns: _Union[str, list] = None, + **kwargs +) -> None: + """ + Write data to a DuckDB database. + + >>> from wrangles.connectors import duckdb + >>> duckdb.write(df, database='database.duckdb', table='table') + + :param df: Pandas Dataframe to be written. + :param database: The database to connect to including the file path. Use ':memory:' for an in-memory database. + :param table: Table to be exported to. + :param action: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + :param columns: (Optional) Subset of the columns to be written. If not provided, all columns will be output. + """ + _logging.info(f": Writing data to DuckDB :: {database} / {table}") + + action = action.upper() + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + + table_name = _quote_identifier(table) + with _duckdb.connect(database=database, **kwargs) as conn: + conn.register('_wrangles_df', df) + + if action == 'FAIL': + conn.execute(f'CREATE TABLE {table_name} AS SELECT * FROM _wrangles_df') + elif action == 'REPLACE': + conn.execute(f'CREATE OR REPLACE TABLE {table_name} AS SELECT * FROM _wrangles_df') + elif action == 'INSERT': + conn.execute(f'CREATE TABLE IF NOT EXISTS {table_name} AS SELECT * FROM _wrangles_df WHERE false') + conn.execute(f'INSERT INTO {table_name} SELECT * FROM _wrangles_df') + else: + raise ValueError('Invalid action. Expected INSERT, REPLACE, or FAIL.') + + +_schema['write'] = """ +type: object +description: Export data to a DuckDB Database +required: + - database + - table +properties: + database: + type: string + description: The database to connect to including the file path. Use ':memory:' for an in-memory database. + table: + type: string + description: The table to write to + action: + type: string + description: INSERT appends, REPLACE recreates the table, FAIL errors if the table exists. Defaults to INSERT. + enum: + - INSERT + - REPLACE + - FAIL + columns: + type: + - string + - array + description: A list of the columns to write to the table. If omitted, all columns will be written. +""" + + +def run( + database: str, + command: _Union[str, list], + params: _Union[list, tuple, dict] = None, + **kwargs +) -> None: + """ + Run a command on a DuckDB database. + + >>> wrangles.connectors.duckdb.run( + >>> database='database.duckdb', + >>> command='' + >>> ) + + :param database: The database to connect to including the file path. Use ':memory:' for an in-memory database. + :param command: SQL command or a list of SQL commands to execute. + :param params: Variables to pass to a parameterized query. + """ + _logging.info(f": Executing DuckDB Command :: {database}") + + if isinstance(command, str): + command = [command] + + with _duckdb.connect(database=database, **kwargs) as conn: + for sql in command: + conn.execute(sql, params or ()) + + +_schema['run'] = r""" +type: object +description: Run a command against a DuckDB Database +required: + - database + - command +properties: + database: + type: string + description: The database to connect to including the file path. Use ':memory:' for an in-memory database. + command: + type: + - string + - array + description: SQL command or a list of SQL commands to execute + params: + type: + - array + - object + description: Variables to pass to a parameterized query. +""" From 369dfff887ab404e22d0bd68abbb68ddc2c03cd0 Mon Sep 17 00:00:00 2001 From: Leonid Molotiievskyi Date: Wed, 1 Jul 2026 13:55:00 +0200 Subject: [PATCH 04/15] build full version of the package (#1035) --- .github/workflows/publish-dev-rc.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/publish-dev-rc.yml b/.github/workflows/publish-dev-rc.yml index 2cfaff17..036cf046 100644 --- a/.github/workflows/publish-dev-rc.yml +++ b/.github/workflows/publish-dev-rc.yml @@ -162,6 +162,12 @@ jobs: - name: Install build tooling run: python -m pip install build twine --user + - name: Patch setup.py to use requirements-full.txt + run: | + sed -i "s/requirements\.txt/requirements-full.txt/" setup.py + echo "setup.py requirements line after patch:" + grep "requirements" setup.py + - name: Build sdist and wheel run: python -m build --sdist --wheel --outdir dist/ . From ab581495a07d9a75d2f38dfd1a485fd879af247a Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Thu, 2 Jul 2026 00:44:14 +0300 Subject: [PATCH 05/15] Add split.dictionary to_lists mode (#1021) * Add split.dictionary to_lists mode * fix tests * renamed parameters --- tests/recipes/wrangles/test_split.py | 169 +++++++++++++++++++++++++++ wrangles/recipe_wrangles/split.py | 50 ++++++-- 2 files changed, 210 insertions(+), 9 deletions(-) diff --git a/tests/recipes/wrangles/test_split.py b/tests/recipes/wrangles/test_split.py index c316f195..509f9df8 100644 --- a/tests/recipes/wrangles/test_split.py +++ b/tests/recipes/wrangles/test_split.py @@ -1281,6 +1281,175 @@ def test_split_dictionary_empty(self): ) assert df.empty and df.columns.to_list() == ['Col'] + def test_split_dictionary_to_lists(self): + """ + Test splitting dictionary keys and values to parallel lists + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [ + {"a": 123, "b": "kshdf"}, + {"b": 123, "c": "kshdf"} + ] + }) + ) + assert df['Keys'].to_list() == [["a", "b"], ["b", "c"]] + assert df['Values'].to_list() == [[123, "kshdf"], [123, "kshdf"]] + + def test_split_dictionary_to_lists_json(self): + """ + Test splitting JSON dictionary keys and values to lists + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': ['{"a": 123, "b": "kshdf"}'] + }) + ) + assert df['Keys'][0] == ["a", "b"] + assert df['Values'][0] == [123, "kshdf"] + + def test_split_dictionary_to_lists_multiple_inputs(self): + """ + Test splitting multiple dictionaries to keys and values lists + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: + - Dict1 + - Dict2 + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'Dict1': [{"a": 1, "b": 2}], + 'Dict2': [{"b": 3, "c": 4}] + }) + ) + assert df['Keys'][0] == ["a", "b", "c"] + assert df['Values'][0] == [1, 3, 4] + + def test_split_dictionary_to_lists_where(self): + """ + Test split.dictionary output_format to_lists using where + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + where: numbers > 3 + """, + dataframe=pd.DataFrame({ + 'My Dict': [ + {"a": 1}, + {"b": 2}, + {"c": 3} + ], + 'numbers': [3, 4, 5] + }) + ) + assert df['Keys'].to_list() == ["", ["b"], ["c"]] + assert df['Values'].to_list() == ["", [2], [3]] + + def test_split_dictionary_to_lists_default_output(self): + """ + Test split.dictionary output_format to_lists default output columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [{"a": 1}] + }) + ) + assert df['Keys'][0] == ["a"] + assert df['Values'][0] == [1] + + def test_split_dictionary_to_lists_output_error(self): + """ + Test split.dictionary output_format to_lists validates output column count + """ + with pytest.raises(ValueError, match="exactly two output columns"): + wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: Keys + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [{"a": 1}] + }) + ) + + def test_split_dictionary_invalid_output_format(self): + """ + Test split.dictionary validates output_format + """ + with pytest.raises(ValueError, match="output_format must be one of"): + wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output_format: invalid + """, + dataframe=pd.DataFrame({ + 'My Dict': [{"a": 1}] + }) + ) + + def test_split_dictionary_to_lists_empty(self): + """ + Test split.dictionary output_format to_lists with an empty column + """ + df = wrangles.recipe.run( + """ + wrangles: + - split.dictionary: + input: My Dict + output: + - Keys + - Values + output_format: to_lists + """, + dataframe=pd.DataFrame({ + 'My Dict': [] + }) + ) + assert df.empty and df.columns.to_list() == ['My Dict', 'Keys', 'Values'] + class TestTokenize: """ diff --git a/wrangles/recipe_wrangles/split.py b/wrangles/recipe_wrangles/split.py index 0d1a3468..01f18cd0 100644 --- a/wrangles/recipe_wrangles/split.py +++ b/wrangles/recipe_wrangles/split.py @@ -17,7 +17,8 @@ def dictionary( df: _pd.DataFrame, input: _Union[str, int, _list], output: _Union[str, _list] = None, - default: dict = None + default: dict = None, + output_format: str = "columns" ) -> _pd.DataFrame: """ type: object @@ -30,7 +31,7 @@ def dictionary( - input properties: input: - type: + type: - string - integer - array @@ -39,22 +40,36 @@ def dictionary( If providing multiple dictionaries and the dictionaries contain overlapping values, the last value will be returned. output: - type: + type: - string - array description: |- - (Optional) Subset of keys to extract from the dictionary. - If not provided, all keys will be returned. + In columns output_format, this is an optional subset of keys to extract + from the dictionary. If not provided, all keys will be returned. Columns can be renamed with the following syntax: output: - key1: new_column_name1 - key2: new_column_name2 + In to_lists output_format, this must be two output columns for the keys + and values lists. If not provided, Keys and Values will be used. default: type: object description: >- Provide a set of default headings and values if they are not found within the input - """ + output_format: + type: string + enum: + - columns + - to_lists + description: |- + How to split the dictionary. + columns creates one output column for each dictionary key. + to_lists creates two output columns containing lists of keys and values. + """ + if output_format not in ["columns", "to_lists"]: + raise ValueError("output_format must be one of: columns, to_lists") + if default is None: default = {} _logging.debug(f": Splitting dictionaries :: input :: {input}") @@ -73,11 +88,28 @@ def _parse_dict_or_json(val): raise ValueError(f'{val} is not a valid Dictionary') from None - # Generate new columns for each key in the dictionary - df_temp = _pd.DataFrame([ + # Merge each row's dictionaries so duplicate keys follow existing behavior: + # later input columns overwrite earlier input columns. + dicts = [ dict(_itertools.chain.from_iterable(_parse_dict_or_json(d) for d in ([default] + row.tolist()))) for row in df[input].values - ]) + ] + + if output_format == "to_lists": + if output is None: + output = ["Keys", "Values"] + elif not isinstance(output, _list): + output = [output] + + if len(output) != 2: + raise ValueError("split.dictionary with output_format to_lists requires exactly two output columns") + + df[output[0]] = [[key for key in item.keys()] for item in dicts] + df[output[1]] = [[value for value in item.values()] for item in dicts] + return df + + # Generate new columns for each key in the dictionary + df_temp = _pd.DataFrame(dicts) # If user has defined how they'd like the output columns if output is not None: From 624aa4f1d9878df8da9d2fee3fc4d569f1af8182 Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Thu, 2 Jul 2026 22:36:38 +0300 Subject: [PATCH 06/15] Add n parameter to lookup for multi-match retrieval (#1001) * Add n parameter to lookup for multi-match retrieval * n!=1 always returns a dict per match, allow * expansion for output * fix translate tests * Patch setup.py to use requirements-full.txt * implemented wildcard output expansion --- .github/workflows/publish-dev-rc.yml | 7 +- tests/recipes/wrangles/test_main.py | 144 +++++++++++++++++++++++++++ tests/test_wrangles.py | 29 ++++++ wrangles/lookup.py | 40 +++++--- wrangles/recipe_wrangles/main.py | 59 +++++++++-- 5 files changed, 254 insertions(+), 25 deletions(-) diff --git a/.github/workflows/publish-dev-rc.yml b/.github/workflows/publish-dev-rc.yml index 036cf046..60b4f987 100644 --- a/.github/workflows/publish-dev-rc.yml +++ b/.github/workflows/publish-dev-rc.yml @@ -125,6 +125,7 @@ jobs: - name: Test pip install run: | python -m pip install --upgrade pip + pip install -r requirements-full.txt pip install . wrangles.recipe tests/samples/generate-data.wrgl.yml @@ -162,12 +163,6 @@ jobs: - name: Install build tooling run: python -m pip install build twine --user - - name: Patch setup.py to use requirements-full.txt - run: | - sed -i "s/requirements\.txt/requirements-full.txt/" setup.py - echo "setup.py requirements line after patch:" - grep "requirements" setup.py - - name: Build sdist and wheel run: python -m build --sdist --wheel --outdir dist/ . diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index b7993006..ffea6e2a 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -6914,6 +6914,150 @@ def test_lookup_model_unrecognized_value_named_column(self): ) assert df['Value'][0] == "" + def test_lookup_n_single_output(self): + """ + Test lookup with n returns a list of n matches in a single output column + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: Matches + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert isinstance(df['Matches'].iloc[0], list) + assert len(df['Matches'].iloc[0]) == 2 + + def test_lookup_n_output_distribution(self): + """ + Test lookup with n where output list length equals n distributes matches across columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Match1 + - Match2 + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert 'Match1' in df.columns + assert 'Match2' in df.columns + + def test_lookup_n_output_wildcard_expansion(self): + """ + Test lookup with n where a single wildcard output name is expanded + into one column per match + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Top * + model_id: e8658a6f-c694-45d0 + n: 3 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert 'Top 1' in df.columns + assert 'Top 2' in df.columns + assert 'Top 3' in df.columns + assert df['Top 1'].iloc[0] != df['Top 2'].iloc[0] != df['Top 3'].iloc[0] + + def test_lookup_n_output_distribution_multiple_rows(self): + """ + Test lookup with n distributes matches correctly across multiple rows + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Match1 + - Match2 + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel', 'Dolores']}) + ) + assert len(df) == 2 + assert 'Match1' in df.columns + assert 'Match2' in df.columns + assert df['Match1'].iloc[0] != df['Match2'].iloc[0] + + def test_lookup_n_named_output_columns_distribution(self): + """ + Test lookup with n where the output columns match the model's + column names, distributing matches across those columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Value + - Score + model_id: e8658a6f-c694-45d0 + n: 2 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + assert 'Value' in df.columns + assert 'Score' in df.columns + assert df['Value'].iloc[0] != df['Score'].iloc[0] + + def test_lookup_n_output_mismatch_named_columns(self): + """ + Test that an error is raised when n does not match the number of + output columns that correspond to the model's column names + """ + with pytest.raises(ValueError, match="must equal n"): + wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Value + - Score + model_id: e8658a6f-c694-45d0 + n: 3 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + + def test_lookup_n_output_mismatch_unnamed_columns(self): + """ + Test that an error is raised when n does not match the number of + output columns that don't correspond to the model's column names + """ + with pytest.raises(ValueError, match="must equal n"): + wrangles.recipe.run( + """ + wrangles: + - lookup: + input: Col1 + output: + - Match1 + - Match2 + model_id: e8658a6f-c694-45d0 + n: 3 + """, + dataframe=pd.DataFrame({'Col1': ['Rachel']}) + ) + def test_lookup_wrong_model_id_type(self): """ Test the error message when passing through a model_id for a different wrangle type diff --git a/tests/test_wrangles.py b/tests/test_wrangles.py index 3f1d18de..e91f99c7 100644 --- a/tests/test_wrangles.py +++ b/tests/test_wrangles.py @@ -291,6 +291,35 @@ def test_lookup_list_value_list_column(): result = wrangles.lookup(["a"], "fe730444-1bda-4fcd", ["Value"]) assert result == [[1]] +def test_lookup_n_single_input(): + """ + Test n returns a list of n matches for a single input + """ + result = wrangles.lookup("Rachel", "e8658a6f-c694-45d0", n=2) + assert isinstance(result, list) + assert len(result) == 2 + +def test_lookup_n_list_input(): + """ + Test n returns a list of n-lists for a list of inputs + """ + result = wrangles.lookup(["Rachel", "Dolores"], "e8658a6f-c694-45d0", n=2) + assert isinstance(result, list) + assert len(result) == 2 + assert isinstance(result[0], list) and len(result[0]) == 2 + assert isinstance(result[1], list) and len(result[1]) == 2 + +def test_lookup_n_with_column(): + """ + Test n with a specific column still returns a list of n dicts - + requesting a single column does not collapse the dict to that + column's value when n > 1 + """ + result = wrangles.lookup("Rachel", "e8658a6f-c694-45d0", "Value", n=2) + assert isinstance(result, list) + assert len(result) == 2 + assert all(isinstance(match, dict) and "Value" in match for match in result) + def test_embedding_single(): """ Test generating an embedding from a single value diff --git a/wrangles/lookup.py b/wrangles/lookup.py index efd13dcd..49a6eb58 100644 --- a/wrangles/lookup.py +++ b/wrangles/lookup.py @@ -9,6 +9,7 @@ def lookup( input: _Union[str, list], model_id: str, columns: _Union[str, list] = None, + n: int = None, **kwargs ) -> _Union[str, list]: """ @@ -17,7 +18,9 @@ def lookup( :param input: A value or list of values to be looked up. :param model_id: The model to be used. :param columns: (Optional) The columns to be returned. If not provided, all columns will be returned as a dict. - """ + :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. + """ # Check if user has entered a single input or multiple inputs single_input = False if not isinstance(input, list): @@ -58,6 +61,10 @@ def lookup( ) _logging.info(f": Looking up {len(input)} values :: model_id :: {model_id}") + + if n: + kwargs['n'] = n + results = _batching.batch_api_calls( f'{_config.api_host}/wrangles/lookup', { @@ -69,20 +76,27 @@ def lookup( batch_size ) - if columns is None: - # If no columns specified, return as [{"col1": "val1", ...}, ...] - results = [ - {col: val for col, val in zip(results["columns"], row)} - for row in results["data"] - ] - elif single_columns: - # If single column specified, return as 1D array [val1, ...] - results = [r[0] for r in results["data"]] + if n and n > 1: + # API returns 1 row per input; row[0] is a list of n match dicts. + # When n > 1, every match is always returned as a dict, even if a + # single column was requested - naming an output column after a + # lookup column does not collapse the dict to that column's value. + results = [row[0] for row in results["data"]] else: - # If multiple columns specified, return as 2D array [[val1, ...], ...] - results = results["data"] + if columns is None: + # If no columns specified, return as [{"col1": "val1", ...}, ...] + results = [ + {col: val for col, val in zip(results["columns"], row)} + for row in results["data"] + ] + elif single_columns: + # If single column specified, return as 1D array [val1, ...] + results = [r[0] for r in results["data"]] + else: + # If multiple columns specified, return as 2D array [[val1, ...], ...] + results = results["data"] - # If input was a single value, return a single value + # If input was a single value, return a single value (or list of n matches) if single_input: results = results[0] diff --git a/wrangles/recipe_wrangles/main.py b/wrangles/recipe_wrangles/main.py index 2fa67d4b..820318c7 100644 --- a/wrangles/recipe_wrangles/main.py +++ b/wrangles/recipe_wrangles/main.py @@ -903,7 +903,8 @@ def lookup( input: str, output: _Union[str, list] = None, model_id: str = None, - lookup_mode: str = 'by_row', + lookup_mode: str = 'by_row', + n: int = None, **kwargs ) -> _pd.DataFrame: """ @@ -925,7 +926,17 @@ def lookup( type: - string - array - description: Name of the output column(s) + description: >- + Name of the output column(s). When n is provided and the output list + length equals n, each output column receives the corresponding match. + A single output containing a wildcard (*) is expanded into n columns, + e.g. "Top *" with n: 3 becomes "Top 1", "Top 2", "Top 3". + n: + type: integer + description: >- + Number of matches to return per input value. When the output list + length equals n, each output column receives the corresponding match. + Otherwise all n matches are stored as a list in each output column. lookup_mode: type: string description: >- @@ -950,6 +961,16 @@ def lookup( # Ensure output is a list if not isinstance(output, list): output = [output] + # Expand a single wildcard output name into one column per match + # e.g. output: "Top *" with n=3 -> ["Top 1", "Top 2", "Top 3"] + if ( + n and n > 1 and + len(output) == 1 and + isinstance(output[0], str) and + '*' in output[0] + ): + output = [output[0].replace('*', str(i)) for i in range(1, n + 1)] + # Return early on empty df if df.empty: # Add empty output columns to maintain expected structure @@ -986,6 +1007,12 @@ def _clean_kwargs(kwargs): kwargs.pop('matrix_variables') return kwargs + # Distribute the n matches for each row across the output columns, + # ranked match i goes to output column i + def _distribute_n_matches(data): + for i, out in enumerate(output): + df[out] = [row[i] if isinstance(row, list) and i < len(row) else None for row in data] + # Perform lookup based on lookup_mode if lookup_mode == 'by_row': # Current behavior - process all rows @@ -995,20 +1022,40 @@ def _clean_kwargs(kwargs): df[input].values.tolist(), model_id, columns=wrangle_output, + n=n, **_clean_kwargs(kwargs) ) - df[output] = data + if n and n > 1 and len(output) == n: + # Distribute: each output column gets the nth match + _distribute_n_matches(data) + elif n and n > 1 and len(output) > 1: + raise ValueError( + f'When n > 1 and multiple output columns are provided, the number ' + f'of output columns ({len(output)}) must equal n ({n}).' + ) + else: + df[output] = data elif not any([col in metadata["settings"]["columns"] for col in wrangle_output]): # User specified no columns from the wrangle data = _lookup( df[input].values.tolist(), model_id, + n=n, **_clean_kwargs(kwargs) ) - for out in output: - df[out] = data + if n and n > 1 and len(output) == n: + # Distribute: each output column gets the nth match + _distribute_n_matches(data) + elif n and n > 1 and len(output) > 1: + raise ValueError( + f'When n > 1 and multiple output columns are provided, the number ' + f'of output columns ({len(output)}) must equal n ({n}).' + ) + else: + for out in output: + df[out] = data else: - # User specified a mixture of unrecognized columns and columns from the wrangle + # User specified a mixture of unrecognized columns and columns from the wrangle raise ValueError('Lookup may only contain all named or unnamed columns.') elif lookup_mode == 'by_dataframe': From 8817e204cdfe8390e8c83b303b86e7e5a971f596 Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Thu, 2 Jul 2026 22:43:00 +0300 Subject: [PATCH 07/15] Create Column should NOT error if column exists (#957) * Create Column should NOT error if column exists * Added new value_if_exists and coalesce_value parameters to handle logic if column preexists --- tests/recipes/wrangles/test_create.py | 163 ++++++++++++++++++++++---- wrangles/recipe_wrangles/create.py | 84 ++++++++++--- 2 files changed, 212 insertions(+), 35 deletions(-) diff --git a/tests/recipes/wrangles/test_create.py b/tests/recipes/wrangles/test_create.py index b7404c2e..68487184 100644 --- a/tests/recipes/wrangles/test_create.py +++ b/tests/recipes/wrangles/test_create.py @@ -92,44 +92,165 @@ def test_create_columns_5(self): df = wrangles.recipe.run(recipe, dataframe=data) assert df.iloc[0]['column3'] in [True, False] - def test_column_exists(self): + def test_column_exists_no_error(self): """ - Check error if trying to create a column that already exists + Default behaviour: creating a column that already exists should not raise + an error — the existing column is left unchanged. """ - data = pd.DataFrame({ - 'col': ['data1'] - }) + data = pd.DataFrame({'col': ['data1']}) recipe = """ wrangles: - create.column: output: col + value: new_value """ - with pytest.raises(ValueError) as info: - wrangles.recipe.run(recipe, dataframe=data) - assert ( - info.typename == 'ValueError' and - '"col" column already exists in dataFrame.' in info.value.args[0] - ) + df = wrangles.recipe.run(recipe, dataframe=data) + assert df['col'][0] == 'data1' - def test_column_exists_list(self): + def test_column_exists_list_no_error(self): """ - Check error if trying to create a list of columns where one already exists + Default behaviour: creating a list of columns where one already exists + should not raise an error — the existing column is left unchanged. """ - data = pd.DataFrame({ - 'col': ['data1'] - }) + data = pd.DataFrame({'col': ['data1']}) recipe = """ wrangles: - create.column: output: - col + - col2 + value: new_value """ - with pytest.raises(ValueError) as info: + df = wrangles.recipe.run(recipe, dataframe=data) + assert df['col'][0] == 'data1' and df['col2'][0] == 'new_value' + + def test_column_exists_value_if_exists_existing(self): + """ + Explicit value_if_exists: existing leaves the column unchanged. + """ + data = pd.DataFrame({'col': ['data1']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: existing + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df['col'][0] == 'data1' + + def test_column_exists_value_if_exists_new(self): + """ + value_if_exists: new overwrites the entire existing column. + """ + data = pd.DataFrame({'col': ['data1', 'data2']}) + recipe = """ + wrangles: + - create.column: + output: col + value: replaced + value_if_exists: new + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['replaced', 'replaced'] + + def test_column_exists_value_if_exists_coalesce(self): + """ + value_if_exists: coalesce fills only null/empty cells, leaving non-null + cells intact. + """ + data = pd.DataFrame({'col': ['keep', None, '']}) + recipe = """ + wrangles: + - create.column: + output: col + value: filled + value_if_exists: coalesce + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['keep', 'filled', 'filled'] + + def test_column_exists_value_if_exists_coalesce_numeric(self): + """ + value_if_exists: coalesce on a numeric column only fills NaN cells, + without comparing values to an empty string. + """ + data = pd.DataFrame({'col': [1, None, 3]}) + recipe = """ + wrangles: + - create.column: + output: col + value: 99 + value_if_exists: coalesce + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == [1, 99, 3] + + def test_column_exists_coalesce_value_new(self): + """ + coalesce_value: new prefers the new value over a non-empty existing + value, only falling back to the existing value where the new value + is empty/null. + """ + data = pd.DataFrame({'col': ['keep1', 'keep2']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: coalesce + coalesce_value: new + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['new_value', 'new_value'] + + def test_column_exists_coalesce_value_new_fallback(self): + """ + coalesce_value: new falls back to the existing value when the new + value is empty/null. + """ + data = pd.DataFrame({'col': ['keep1', 'keep2']}) + recipe = """ + wrangles: + - create.column: + output: col + value: + value_if_exists: coalesce + coalesce_value: new + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert list(df['col']) == ['keep1', 'keep2'] + + def test_column_exists_coalesce_value_invalid(self): + """ + An invalid coalesce_value option raises a ValueError. + """ + data = pd.DataFrame({'col': ['data1']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: coalesce + coalesce_value: not_a_real_option + """ + with pytest.raises(ValueError, match="coalesce_value"): + wrangles.recipe.run(recipe, dataframe=data) + + def test_column_exists_value_if_exists_invalid(self): + """ + An invalid value_if_exists option raises a ValueError. + """ + data = pd.DataFrame({'col': ['data1']}) + recipe = """ + wrangles: + - create.column: + output: col + value: new_value + value_if_exists: existng + """ + with pytest.raises(ValueError, match="value_if_exists"): wrangles.recipe.run(recipe, dataframe=data) - assert ( - info.typename == 'ValueError' and - "['col'] column(s)" in info.value.args[0] - ) def test_create_column_value_number(self): """ diff --git a/wrangles/recipe_wrangles/create.py b/wrangles/recipe_wrangles/create.py index 3245d2f7..fdc755f3 100644 --- a/wrangles/recipe_wrangles/create.py +++ b/wrangles/recipe_wrangles/create.py @@ -85,7 +85,13 @@ def bins( return df -def column(df: _pd.DataFrame, output: _Union[str, list], value = None) -> _pd.DataFrame: +def column( + df: _pd.DataFrame, + output: _Union[str, list], + value = None, + value_if_exists: str = 'existing', + coalesce_value: str = 'existing' +) -> _pd.DataFrame: """ type: object description: Create column(s) with a user defined value. Defaults to None (empty). @@ -106,13 +112,45 @@ def column(df: _pd.DataFrame, output: _Union[str, list], value = None) -> _pd.Da - array - boolean description: (Optional) Value(s) to add in the new column(s). If using a dictionary in output, value can only be a string. + value_if_exists: + type: string + description: >- + Determines behaviour when the output column already exists. + existing (default): leave the column unchanged. + coalesce: fill empty/null cells with the new value, keeping non-null cells. + new: overwrite the entire column with the new value. + enum: + - existing + - coalesce + - new + coalesce_value: + type: string + description: >- + Only used when value_if_exists is coalesce. Determines which side + is preferred when both the existing and new values are non-empty. + existing (default): keep the existing value, only fill empty/null cells with the new value. + new: keep the new value, only fall back to the existing value where the new value is empty/null. + enum: + - existing + - new """ _logging.debug(f": Creating column(s) :: output :: {output}") + + valid_value_if_exists = ('existing', 'coalesce', 'new') + if value_if_exists not in valid_value_if_exists: + raise ValueError( + f'value_if_exists must be one of {valid_value_if_exists}, got "{value_if_exists}"' + ) + + valid_coalesce_value = ('existing', 'new') + if coalesce_value not in valid_coalesce_value: + raise ValueError( + f'coalesce_value must be one of {valid_coalesce_value}, got "{coalesce_value}"' + ) + # If a string provided, convert to list if isinstance(output, str): - if output in df.columns: - raise ValueError(f'"{output}" column already exists in dataFrame.') - output = [output] + output = [output] # gather the columns and values in a dictionary, if not a dict then use value as the value of dictionary output_dict = {} @@ -124,18 +162,36 @@ def column(df: _pd.DataFrame, output: _Union[str, list], value = None) -> _pd.Da else: output_dict.update({out: value}) - # Check if the list of outputs exist in dataFrame - check_list = [x for x in (output_dict.keys()) if x in df.columns] - if len(check_list) > 0: - raise ValueError(f'{check_list} column(s) already exists in the dataFrame') - for output_column, values_list in zip(output_dict.keys(), output_dict.values()): + column_exists = output_column in df.columns + + if column_exists and value_if_exists == 'existing': + continue + # Data to generate - data = _pd.DataFrame({ - output_column: _generate_cell_values(values_list, len(df)) - }).set_index(df.index) # use the same index as original to match rows - # Merging existing dataframe with values created - df = _pd.concat([df, data], axis=1) + new_values = _pd.Series( + _generate_cell_values(values_list, len(df)), + index=df.index + ) + + if column_exists and value_if_exists == 'coalesce': + if coalesce_value == 'existing': + primary, fallback = df[output_column], new_values + else: + primary, fallback = new_values, df[output_column] + + is_empty = primary.isna() + if primary.dtype == object: + is_empty = is_empty | (primary == '') + + df[output_column] = primary.where(~is_empty, fallback) + else: + # new or column doesn't exist — write/overwrite directly + if not column_exists: + data = _pd.DataFrame({output_column: new_values}) + df = _pd.concat([df, data], axis=1) + else: + df[output_column] = new_values return df From 673b07643fa958f16c6987fa3b764fa3f9282ed7 Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Thu, 2 Jul 2026 22:47:50 +0300 Subject: [PATCH 08/15] Fix where skipping wrangles when all rows are filtered ou (#1007) * 496 selectkeys (#1033) * Add split.dictionary to_lists mode * Fix where skipping wrangles when all rows are filtered out * 1005 bug where passes empty dataframe (#1034) * enable package build on main branch * 956 create column should not error if column exists (#1036) * Create Column should NOT error if column exists * install full package list * 990 enhancement add n to lookup params (#1037) * Add n parameter to lookup for multi-match retrieval * 985 enhancement allow list of onemultiple defaults (#1039) * convert.from_json and convert.from_yaml accept default as a list when input is also a list * 985-enhancement-allow-list-of-onemultiple-defaults * column-shift bug in wrangles/recipe.py * Revert "Merge branch 'dev' into 1005-bug-where-passes-empty-dataframe" This reverts commit a6b9b5721f8894f8641eb2aace204338a15b2f69, reversing changes made to af67032514b6aa9bd875949e98cd6c4ac130527e. * publesh dev rc * skip pre-creating output columns --- tests/recipes/wrangles/test_main.py | 55 ++++++++++++++++++++++++++++- wrangles/recipe.py | 36 +++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index ffea6e2a..e7321706 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -3969,7 +3969,32 @@ def test_recipe_where(self): }) ) assert df.values.tolist() == [['a', 'value1'], ['B', 'VALUE2']] - + + def test_recipe_where_empty_dataframe(self): + """ + Test that when where filters out all rows, the recipe wrangle is + skipped entirely and does not fail due to missing columns. Issue #1005. + """ + df = wrangles.recipe.run( + """ + wrangles: + - recipe: + where: successful_search == True + wrangles: + - split.dictionary: + input: scored_results + - split.dictionary: + input: summary + """, + dataframe=pd.DataFrame({ + 'scored_results': [{}], + 'successful_search': [False] + }) + ) + assert df['scored_results'].tolist() == [{}] + assert df['successful_search'].tolist() == [False] + assert 'summary' not in df.columns + def test_recipe_empty_column_preserved(self): data = [ ["col1", "", "col2"], @@ -5922,6 +5947,34 @@ def test_batch_where(self): ) assert df['output col'].to_list() == ["A","","C"] + def test_batch_size_one_where_no_column_shift(self): + """ + Test batch_size: 1 combined with a wrangle-level where. + Regression test - when a batch's single row does not match + the where clause, the output column must still be created + (as an empty value) rather than omitted entirely, otherwise + results become misaligned between batches. + """ + df = wrangles.recipe.run( + """ + wrangles: + - batch: + batch_size: 1 + wrangles: + - convert.case: + input: Desc + output: output_column_name + case: upper + where: WC = "A" + """, + dataframe=pd.DataFrame({ + "WC": ["A", "A", "B"], + "Desc": ["first A", "second A", "first B"] + }) + ) + assert df.columns.tolist() == ["WC", "Desc", "output_column_name"] + assert df["output_column_name"].to_list() == ["FIRST A", "SECOND A", ""] + def test_batch_variables(self): """ Test batch wrangle with a variable passed through diff --git a/wrangles/recipe.py b/wrangles/recipe.py index 089ccdf5..09149c4b 100644 --- a/wrangles/recipe.py +++ b/wrangles/recipe.py @@ -506,6 +506,42 @@ def _execute_wrangles( preserve_index=True ) + # If where filters out all rows, skip actually executing the + # wrangle (some wrangles error when given no rows), but if it + # explicitly declares output columns, still add them (empty) + # so the dataframe structure stays consistent with runs where + # at least one row matches - otherwise batching with where can + # produce inconsistent columns between batches + if len(df) == 0: + if 'output' in params: + if isinstance(params['output'], list): + output_columns = [ + list(col.values()) if isinstance(col, dict) else [col] + for col in params['output'] + ] + output_columns = [ + item + for sublist in output_columns + for item in sublist + ] + elif isinstance(params['output'], dict): + output_columns = list(params['output'].keys()) + else: + output_columns = [params['output']] + + for col in output_columns: + # Wildcard outputs (e.g. 'Col*') are expanded into + # concrete column names based on the actual data, + # which can't be determined with no rows to work + # with - skip adding those, only add named columns + if '*' in str(col): + continue + if col not in df_original.columns: + df_original[col] = '' + + df = df_original + continue + # Add to common_params dict and remove from params for key in ['where', 'where_params', 'if']: if key in params.keys(): From 6196673a2e70937e2991ff34e19796cb6993406a Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Tue, 21 Jul 2026 20:51:29 +0300 Subject: [PATCH 09/15] Add sort fallback for mixed-type columns (#1066) eliminate error when occasional strings get mixed in with mostly numbers --- tests/recipes/wrangles/test_main.py | 16 ++++++++ wrangles/recipe_wrangles/pandas.py | 62 +++++++++++++++++++++++++++-- 2 files changed, 74 insertions(+), 4 deletions(-) diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index e7321706..01073004 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -8984,6 +8984,22 @@ def test_pandas_sort_debug_log(self, caplog): ) assert any(': Sorting dataframe' in msg for msg in caplog.messages) + def test_pandas_sort_coerces_mixed_numeric_types(self): + df = wrangles.recipe.run( + """ + wrangles: + - sort: + by: score + """, + dataframe=pd.DataFrame({ + 'score': [10.5, '', 2.0, '3.5'], + 'item': ['ten', 'blank', 'two', 'three'], + }) + ) + + assert df['item'].tolist() == ['blank', 'two', 'three', 'ten'] + assert df['score'].tolist() == ['', 2.0, '3.5', 10.5] + def test_pandas_round_debug_log(self, caplog): import logging caplog.set_level(logging.DEBUG) diff --git a/wrangles/recipe_wrangles/pandas.py b/wrangles/recipe_wrangles/pandas.py index 552d8111..aadd33d3 100644 --- a/wrangles/recipe_wrangles/pandas.py +++ b/wrangles/recipe_wrangles/pandas.py @@ -4,6 +4,56 @@ import logging as _logging +def _sort_by_columns(by): + return [by] if isinstance(by, str) else list(by) + + +def _is_empty_sort_value(value): + if isinstance(value, str) and value.strip() == "": + return True + + try: + return bool(_pd.isna(value)) + except (TypeError, ValueError): + return False + + +def _coerce_sort_series(series: _pd.Series) -> _pd.Series: + non_empty = series[~series.map(_is_empty_sort_value)] + if non_empty.empty: + return series + + numeric_values = _pd.to_numeric(series, errors="coerce") + numeric_count = numeric_values.loc[non_empty.index].notna().sum() + if numeric_count >= len(non_empty) / 2: + return numeric_values.fillna(0) + + return series.map(lambda value: "" if _is_empty_sort_value(value) else str(value)) + + +def _coerced_sort_values(df: _pd.DataFrame, by_columns: list, ignore_index: bool, kwargs: dict) -> _pd.DataFrame: + sort_df = df.copy() + sort_kwargs = kwargs.copy() + temp_columns = [] + + for index, column in enumerate(by_columns): + if column not in sort_df.columns: + continue + + temp_column = f"__wrangles_sort_key_{index}" + while temp_column in sort_df.columns: + temp_column = f"_{temp_column}" + + sort_df[temp_column] = _coerce_sort_series(sort_df[column]) + temp_columns.append(temp_column) + + if not temp_columns: + return df.sort_values(ignore_index=ignore_index, **kwargs) + + sort_kwargs["by"] = temp_columns + return sort_df.sort_values(ignore_index=ignore_index, **sort_kwargs).drop(columns=temp_columns) + + def copy( df: _pd.DataFrame, input: _Union[str, int, list] = None, @@ -148,7 +198,7 @@ def sort(df: _pd.DataFrame, ignore_index=True, **kwargs) -> _pd.DataFrame: raise ValueError("'by' parameter is required for sorting") # Ensure 'by' is a list for consistent processing - by_columns = [by] if isinstance(by, str) else by + by_columns = _sort_by_columns(by) # Check if any columns need dtype conversion (float16 -> float32) # float16 can cause sorting issues in pandas @@ -166,8 +216,12 @@ def sort(df: _pd.DataFrame, ignore_index=True, **kwargs) -> _pd.DataFrame: if col in df.columns and df[col].dtype == "float16": df[col] = df[col].astype("float32") - # Perform the sort operation - return df.sort_values(ignore_index=ignore_index, **kwargs) + # Perform the sort operation. If mixed column types cannot be compared, + # retry using temporary sort keys coerced to the predominant compatible type. + try: + return df.sort_values(ignore_index=ignore_index, **kwargs) + except TypeError: + return _coerced_sort_values(df, by_columns, ignore_index, kwargs) def round(df: _pd.DataFrame, input: _Union[str, int, list], decimals: int = 0, output: _Union[str, list] = None) -> _pd.DataFrame: @@ -302,4 +356,4 @@ def explode( if drop_empty: df = df.dropna(subset=input, how='all') - return df \ No newline at end of file + return df From 45fba979cd9cdd00ec389cbe44e6154bda8c8f9f Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Tue, 21 Jul 2026 21:32:34 +0300 Subject: [PATCH 10/15] allow rename to skip missing inputs when output exists (#1068) * allow rename to skip missing inputs when output exists * Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Eric Hills <53243273+ebhills@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- tests/recipes/wrangles/test_main.py | 61 +++++++++++++++++++++++++++++ wrangles/recipe_wrangles/main.py | 54 ++++++++++++++++++------- 2 files changed, 101 insertions(+), 14 deletions(-) diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index 01073004..8fb0555d 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -2424,6 +2424,67 @@ def test_rename_optional_string_input_convert_case(self): # Should rename Col1 to COL1 assert 'COL1' in df.columns + def test_rename_missing_input_skips_when_output_exists_dict(self): + """ + Missing input should not error when the target output column already exists. + """ + data = pd.DataFrame({ + 'Description': ['already normalized'], + 'Part Number': ['PN-1'], + }) + recipe = """ + wrangles: + - rename: + desc: Description + """ + df = wrangles.recipe.run(recipe, dataframe=data) + + assert df.columns.tolist() == ['Description', 'Part Number'] + assert df.iloc[0]['Description'] == 'already normalized' + + def test_rename_multiple_possible_inputs_to_existing_output(self): + """ + Alternate input names can map to one output, or skip if output already exists. + """ + recipe = """ + wrangles: + - rename: + input: + - [input desc, desc] + output: + - Description + """ + + input_desc_df = wrangles.recipe.run( + recipe, + dataframe=pd.DataFrame({ + 'input desc': ['from input desc'], + 'Part Number': ['PN-1'], + }) + ) + assert input_desc_df.columns.tolist() == ['Description', 'Part Number'] + assert input_desc_df.iloc[0]['Description'] == 'from input desc' + + desc_df = wrangles.recipe.run( + recipe, + dataframe=pd.DataFrame({ + 'desc': ['from desc'], + 'Part Number': ['PN-2'], + }) + ) + assert desc_df.columns.tolist() == ['Description', 'Part Number'] + assert desc_df.iloc[0]['Description'] == 'from desc' + + existing_output_df = wrangles.recipe.run( + recipe, + dataframe=pd.DataFrame({ + 'Description': ['already normalized'], + 'Part Number': ['PN-3'], + }) + ) + assert existing_output_df.columns.tolist() == ['Description', 'Part Number'] + assert existing_output_df.iloc[0]['Description'] == 'already normalized' + class TestSimilarity: """ Test similarity diff --git a/wrangles/recipe_wrangles/main.py b/wrangles/recipe_wrangles/main.py index 820318c7..83426351 100644 --- a/wrangles/recipe_wrangles/main.py +++ b/wrangles/recipe_wrangles/main.py @@ -1640,6 +1640,39 @@ def resolve_input(candidate_list): ): del kwargs["functions"] + def output_exists(output_column): + return output_column in df.columns + + def resolve_rename_input(input_column, output_column): + candidates = input_column if isinstance(input_column, list) else [input_column] + optional_candidates = [] + required_candidates = [] + + for candidate in candidates: + optional = False + actual_col = candidate + if isinstance(candidate, str) and candidate.endswith("?"): + optional = True + actual_col = candidate[:-1] + + if actual_col in df.columns: + return actual_col + + if optional: + optional_candidates.append(actual_col) + continue + + required_candidates.append(actual_col) + + if optional_candidates and not required_candidates: + return None + + if output_exists(output_column): + return None + + missing = required_candidates[0] if required_candidates else candidates[0] + raise ValueError(f'Rename column "{missing}" not found.') + # If short form of paired names is provided, use that if input is None: @@ -1697,7 +1730,7 @@ def resolve_input(candidate_list): # Regular non-wildcard name if name not in cols: - if optional: + if optional or kwargs[x] in cols: continue else: raise ValueError(f'Rename column "{name}" not found.') @@ -1723,19 +1756,12 @@ def resolve_input(candidate_list): raise ValueError('The lists for input and output must be the same length.') for inp, out in zip(input, output): - if inp.endswith("?"): - actual_col = inp[:-1] - if actual_col not in list(df.columns): - # Skip this column if it doesn't exist - continue # This skips both input and output - else: - filtered_input.append(actual_col) - filtered_output.append(out) - elif inp not in list(df.columns): - raise ValueError(f'Rename column "{inp}" not found.') - else: - filtered_input.append(inp) - filtered_output.append(out) + actual_col = resolve_rename_input(inp, out) + if actual_col is None: + continue + + filtered_input.append(actual_col) + filtered_output.append(out) # Check that the output columns don't already exist if so drop them df = df.drop(columns=[x for x in filtered_output if x in df.columns and x not in filtered_input]) From 45dcdcf0d3402e00c5dda753e21e339af271cb9b Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Tue, 21 Jul 2026 22:09:36 +0300 Subject: [PATCH 11/15] perf: speed up convert_case function (#949) --- tests/recipes/wrangles/test_convert.py | 229 +++++++++++++++++++++++++ wrangles/recipe_wrangles/convert.py | 59 ++++--- 2 files changed, 268 insertions(+), 20 deletions(-) diff --git a/tests/recipes/wrangles/test_convert.py b/tests/recipes/wrangles/test_convert.py index 89d24d9e..0e329f88 100644 --- a/tests/recipes/wrangles/test_convert.py +++ b/tests/recipes/wrangles/test_convert.py @@ -252,6 +252,235 @@ def test_empty(self): ) assert df.empty and df.columns.to_list() == ['column', 'upper column'] + def test_sentence_tabs_after_punctuation(self): + """ + Sentence case with tab characters as whitespace between sentences. + The regex handles [ \t]* between punctuation and next word. + """ + data = pd.DataFrame({'Data': ["hello.\tthere! one more sentence."]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "Hello.\tThere! One more sentence." + + def test_sentence_leading_whitespace(self): + """ + Sentence case where the string starts with leading spaces/tabs. + The first non-whitespace character should be capitalised. + """ + data = pd.DataFrame({'Data': [" hello world. next sentence", "\t\thello again"]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == " Hello world. Next sentence" + assert df.iloc[1]['out'] == "\t\tHello again" + + def test_sentence_unicode_accented_chars(self): + """ + Sentence case preserves/capitalises Unicode accented characters. + """ + data = pd.DataFrame({'Data': ["héllo wörld. ñoño is here! über cool."]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "Héllo wörld. Ñoño is here! Über cool." + + def test_sentence_starts_with_number(self): + """ + Sentence case when the sentence starts with a digit. + The digit 'consumes' the capitalise flag; subsequent lowercase letters stay lower. + """ + data = pd.DataFrame({'Data': ["13 items found. that's all."]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "13 items found. That's all." + + def test_lower_non_string_preserved(self): + """ + Non-string values (int, list, None) in a lower-case column are returned unchanged. + The vectorised path must restore originals for non-string rows. + """ + data = pd.DataFrame({'Data': ["HELLO", 99, ["A", "B"], None]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "hello" + assert df.iloc[1]['out'] == 99 + assert df.iloc[2]['out'] == ["A", "B"] + + def test_title_non_string_preserved(self): + """ + Non-string values in a title-case column are returned unchanged. + """ + data = pd.DataFrame({'Data': ["hello world", {"key": "val"}]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: title + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == "Hello World" + assert df.iloc[1]['out'] == {"key": "val"} + + def test_purely_numeric_column_preserved(self): + """ + A column with an all-integer int64 dtype must not raise AttributeError + from the .str accessor and must return all values unchanged with a warning. + """ + data = pd.DataFrame({'Data': [1, 2, 3]}) + recipe = """ + wrangles: + - convert.case: + input: Data + output: out + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['out'] == 1 + assert df.iloc[1]['out'] == 2 + assert df.iloc[2]['out'] == 3 + + def test_lower_large_dataframe(self): + """ + Vectorised lower case on a large dataframe produces correct output. + """ + n = 100_000 + data = pd.DataFrame({'Data': ['Hello World MIXED Case'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'hello world mixed case' + assert df.iloc[-1]['Data'] == 'hello world mixed case' + + def test_upper_large_dataframe(self): + """ + Vectorised upper case on a large dataframe produces correct output. + """ + n = 100_000 + data = pd.DataFrame({'Data': ['hello world mixed case'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: upper + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'HELLO WORLD MIXED CASE' + + def test_title_large_dataframe(self): + """ + Vectorised title case on a large dataframe produces correct output. + """ + n = 100_000 + data = pd.DataFrame({'Data': ['hello world mixed case'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: title + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'Hello World Mixed Case' + + def test_sentence_large_dataframe(self): + """ + Regex-based sentence case on a large dataframe produces correct output. + """ + n = 50_000 + data = pd.DataFrame({'Data': ['first sentence. second one! third? yes.'] * n}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: sentence + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Data'] == 'First sentence. Second one! Third? Yes.' + + def test_lower_multiple_columns(self): + """ + Vectorised lower case across multiple input/output column pairs. + """ + n = 50_000 + data = pd.DataFrame({ + 'Col1': ['Hello World'] * n, + 'Col2': ['ANOTHER STRING'] * n, + 'Col3': ['YET ANOTHER'] * n, + }) + recipe = """ + wrangles: + - convert.case: + input: + - Col1 + - Col2 + - Col3 + output: + - Out1 + - Out2 + - Out3 + case: lower + """ + df = wrangles.recipe.run(recipe, dataframe=data) + assert df.iloc[0]['Out1'] == 'hello world' + assert df.iloc[0]['Out2'] == 'another string' + assert df.iloc[0]['Out3'] == 'yet another' + + def test_mixed_types_warning_logged_once(self, caplog): + """ + When a column contains non-string values, the invalid_data warning fires + exactly once regardless of how many non-string rows there are. + """ + import logging + data = pd.DataFrame({'Data': ["hello", 1, 2, 3, "world"]}) + recipe = """ + wrangles: + - convert.case: + input: Data + case: upper + """ + with caplog.at_level(logging.WARNING): + df = wrangles.recipe.run(recipe, dataframe=data) + warning_messages = [r.message for r in caplog.records if 'invalid' in r.message.lower() or 'non-string' in r.message.lower() or 'not a string' in r.message.lower()] + # Warning should fire at most once + assert len(warning_messages) <= 1 + # Strings are correctly transformed + assert df.iloc[0]['Data'] == 'HELLO' + assert df.iloc[4]['Data'] == 'WORLD' + # Non-strings are preserved + assert df.iloc[1]['Data'] == 1 + class TestConvertDataType: """ diff --git a/wrangles/recipe_wrangles/convert.py b/wrangles/recipe_wrangles/convert.py index 07fc1274..a2821216 100644 --- a/wrangles/recipe_wrangles/convert.py +++ b/wrangles/recipe_wrangles/convert.py @@ -16,6 +16,14 @@ except ImportError: from yaml import SafeLoader as _YAMLLoader, SafeDumper as _YAMLDumper +# Pre-compiled regex for sentence case: matches the first non-whitespace character +# at the start of the string or immediately after punctuation (. ! ?) + optional +# whitespace. Using (\S) rather than ([a-zA-Z]) preserves the original char-by-char +# behaviour where a digit after punctuation "consumes" the capitalize flag +# (e.g. "13.5mm" stays lowercase because the `5` consumes the flag before `m` is +# reached), while \s* preserves handling of newlines and other Unicode whitespace. +_SENTENCE_CASE_RE = _re.compile(r'(^\s*|[.!?]\s*)(\S)') + def case(df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, list] = None, case: str = 'lower') -> _pd.DataFrame: """ @@ -75,33 +83,44 @@ def case(df: _pd.DataFrame, input: _Union[str, int, list], output: _Union[str, l # Loop through and apply for all columns for input_column, output_column in zip(input, output): if desired_case != 'sentence': - df[output_column] = df[input_column].apply(lambda x: _safe_str_transform(x, desired_case, warnings)) + source = df[input_column] + # .str accessor raises AttributeError on non-object/non-string dtypes (e.g. int64). + # For those columns, preserve originals and warn once — matching prior behaviour. + if _pd.api.types.is_string_dtype(source.dtype) or _pd.api.types.is_object_dtype(source.dtype): + # Use vectorized pandas str methods (C-level, much faster than row-wise apply). + # Non-string values (lists, ints, etc.) become NaN after str operations; + # detect those and restore the originals so behaviour is unchanged. + transformed = getattr(source.str, desired_case)() + non_string_mask = source.notna() & transformed.isna() + if non_string_mask.any(): + if not warnings["invalid_data"]["logged"]: + _logging.warning(warnings['invalid_data']['message']) + warnings["invalid_data"]['logged'] = True + transformed = transformed.where(~non_string_mask, source) + df[output_column] = transformed + else: + if not warnings["invalid_data"]["logged"]: + _logging.warning(warnings['invalid_data']['message']) + warnings["invalid_data"]['logged'] = True + df[output_column] = source - elif desired_case == 'sentence': - def _getSentenceCase(source: str, warnings={}): + else: + # Sentence case: lowercase everything with the vectorized str method, then + # use a pre-compiled regex to re-capitalise sentence starts. This avoids + # the slow character-by-character Python loop of the previous implementation. + def _getSentenceCase(source, warnings=warnings): if isinstance(source, str): - output = [] - isFirstWord = True - - for character in source: - if isFirstWord and not character.isspace(): - output.append(character.upper()) - isFirstWord = False - elif not isFirstWord and character in ".!?": - isFirstWord = True - output.append(character.upper()) - else: - output.append(character.lower()) - - return ''.join(output) + return _SENTENCE_CASE_RE.sub( + lambda m: m.group(1) + m.group(2).upper(), + source.lower() + ) else: - # Only show this once to not spam the logs - if not warnings.get("invalid_data", {}).get('logged', False): + if not warnings["invalid_data"]["logged"]: _logging.warning(warnings['invalid_data']['message']) warnings["invalid_data"]['logged'] = True return source - df[output_column] = df[input_column].apply(lambda x: _getSentenceCase(x, warnings)) + df[output_column] = df[input_column].apply(_getSentenceCase) return df From 61457b468f20003660acbdc86b9b783d6b07b2a9 Mon Sep 17 00:00:00 2001 From: Eric Hills <53243273+ebhills@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:19:52 -0500 Subject: [PATCH 12/15] Add ignored output_files directory (#1028) --- .gitignore | 4 ++++ output_files/.gitkeep | 0 2 files changed, 4 insertions(+) create mode 100644 output_files/.gitkeep diff --git a/.gitignore b/.gitignore index 897ca8f8..29aa9f44 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,10 @@ tests/temp/* !tests/temp/README.md +# Local/generated output files +output_files/* +!output_files/.gitkeep + # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/output_files/.gitkeep b/output_files/.gitkeep new file mode 100644 index 00000000..e69de29b From 6f1bf4b7ef7f31f1adeb16aa7d2aeb132ec27092 Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Tue, 21 Jul 2026 22:29:41 +0300 Subject: [PATCH 13/15] Output empty string when extract use labels does not get a match mae (#880) - 817-output-empty-string-when-extract-use_labels-does-not-get-a-match-mae - add new param that can toggle off these empty labels - add param to format use_labels output as dict or columns --- tests/recipes/wrangles/test_extract.py | 57 +++++++++- wrangles/extract.py | 30 +++++- wrangles/recipe_wrangles/extract.py | 139 ++++++++++++++++++------- 3 files changed, 186 insertions(+), 40 deletions(-) diff --git a/tests/recipes/wrangles/test_extract.py b/tests/recipes/wrangles/test_extract.py index 854dac64..08c04dc6 100644 --- a/tests/recipes/wrangles/test_extract.py +++ b/tests/recipes/wrangles/test_extract.py @@ -1103,6 +1103,33 @@ def test_extract_custom_labels(self): df['col2'][0]['size'] == ['small'] ) + def test_extract_custom_labels_columns_format(self): + """ + Test use_labels option with output_format: columns to expand labels into dataframe columns + """ + df = wrangles.recipe.run( + """ + wrangles: + - extract.custom: + input: col1 + output: col2 + model_id: 829c1a73-1bfd-4ac0 + use_labels: true + output_format: columns + """, + dataframe = pd.DataFrame({ + 'col1': ['small blue cotton jacket'] + }) + ) + + # Expect new columns `colour` and `size` created and populated + assert ( + 'colour' in df.columns and + 'size' in df.columns and + df['colour'][0] == ['blue'] and + df['size'][0] == ['small'] + ) + def test_extract_custom_6(self): """ Incorrect model_id - forget to use ${} @@ -1734,6 +1761,7 @@ def test_unlabeled_only(self): model_id: 829c1a73-1bfd-4ac0 use_labels: true first_element: false + include_empty_labels: false """ df = wrangles.recipe.run(recipe, dataframe=data) assert df['out'][0] == {'Unlabeled': ['red']} @@ -1753,6 +1781,7 @@ def test_unlabeled_only_with_first_element_true(self): model_id: 829c1a73-1bfd-4ac0 use_labels: true first_element: true + include_empty_labels: false """ df = wrangles.recipe.run(recipe, dataframe=data) assert df['out'][0] == {'Unlabeled': 'red'} @@ -2044,7 +2073,33 @@ def test_extract_custom_sort_unexisting_sort_type(self): info.typename == 'ValueError' and 'Sort must be one of the following: training_order, input_order, longest, shortest, alphabetical, reverse_alphabetical, ascending, descending' in info.value.args[0] ) - + + def test_extract_custom_use_labels_empty_matches(self): + """ + Test that use_labels=True creates empty keys for all labels when no matches are found + """ + df = wrangles.recipe.run( + """ + wrangles: + - extract.custom: + input: col1 + output: out + model_id: 829c1a73-1bfd-4ac0 + use_labels: true + """, + dataframe=pd.DataFrame({ + 'col1': ['this text has no matching labels'] + }) + ) + result = df['out'][0] + + # Should contain empty keys for all possible labels from the model + expected_labels = ['colour', 'size'] # Based on the model's expected labels + + # Verify all expected labels exist with empty values + for label in expected_labels: + assert label in result, f"Missing label '{label}' in output" + assert result[label] == [], f"Label '{label}' should be empty list" class TestExtractRegex: diff --git a/wrangles/extract.py b/wrangles/extract.py index 734ddce6..bdf3cd1c 100644 --- a/wrangles/extract.py +++ b/wrangles/extract.py @@ -456,7 +456,9 @@ def custom( case_sensitive: bool = False, extract_raw: bool = False, use_spellcheck: bool = False, + include_empty_labels: bool = True, sort: str = 'training_order', + output_format: str = 'dict', **kwargs ) -> list: """ @@ -495,6 +497,15 @@ def custom( } model_properties = _data.model(model_id) + model_content = _data.model_content(model_id) + + model_labels = set() + for item in model_content['Data']: + if len(item) >= 2: + if ':' in item[1]: + label = item[1].split(':')[0] # Second column typically contains the label/type + model_labels.add(label.strip()) + # If model_id format is correct but no mode_id exists if model_properties.get('message', None) == 'error': raise ValueError('Incorrect model_id.\nmodel_id may be wrong or does not exists') @@ -524,13 +535,26 @@ def custom( {results["columns"][i]: row[i] for i in range(len(row))} for row in results["data"] ] - if isinstance(results, list): if first_element and not use_labels: results = [x[0] if len(x) >= 1 else "" for x in results] - if use_labels and first_element: - results = [{k:v[0] for (k, v) in zip(objs.keys(), objs.values())} for objs in results] + if use_labels: + if include_empty_labels: + # Ensure every label has a key, create empty keys if missing. + # Use both labels discovered from results and labels defined in the model. + all_labels = set(model_labels or []) + for objs in results: + all_labels.update([str(k).lower() for k in objs.keys()]) + + for objs in results: + # Normalize existing keys to lower-case while preserving original keys + existing = {str(k).lower(): k for k in objs.keys()} + for label in all_labels: + if label not in existing: + objs[label] = [] + if first_element: + results = [{k: v[0] if isinstance(v, list) and v else "" for k, v in objs.items()} for objs in results] else: raise ValueError(f'API Response did not return an expected format for model {model_id}') diff --git a/wrangles/recipe_wrangles/extract.py b/wrangles/recipe_wrangles/extract.py index eb4208d4..0e8236fd 100644 --- a/wrangles/recipe_wrangles/extract.py +++ b/wrangles/recipe_wrangles/extract.py @@ -553,7 +553,9 @@ def custom( case_sensitive: bool = False, extract_raw: bool = False, use_spellcheck: bool = False, + include_empty_labels: bool = True, sort: str = 'training_order', + output_format: str = 'dict', **kwargs ) -> _pd.DataFrame: """ @@ -606,6 +608,15 @@ def custom( - reverse_alphabetical - ascending - descending + include_empty_labels: + type: boolean + description: Include labels with no found values in the output when using use_labels=True + output_format: + type: string + description: Format of the output when using use_labels=True + enum: + - dict + - columns """ if output is None: output = input @@ -615,20 +626,39 @@ def custom( if not isinstance(model_id, list): model_id = [model_id] if len(input) == len(output) and len(model_id) == 1: - # if one model_id, then use that model for all columns inputs and outputs - model_id = [model_id[0] for _ in range(len(input))] - for in_col, out_col, model in zip(input, output, model_id): - df[out_col] = _extract.custom( - df[in_col].astype(str).tolist(), - model_id=model, - first_element=first_element, - use_labels=use_labels, - case_sensitive=case_sensitive, - extract_raw=extract_raw, - use_spellcheck=use_spellcheck, - sort=sort, - **kwargs - ) + # if one model_id, then use that model for all columns inputs and outputs + model_id = [model_id[0] for _ in range(len(input))] + for in_col, out_col, model in zip(input, output, model_id): + results = _extract.custom( + df[in_col].astype(str).tolist(), + model_id=model, + first_element=first_element, + use_labels=use_labels, + case_sensitive=case_sensitive, + extract_raw=extract_raw, + use_spellcheck=use_spellcheck, + include_empty_labels=include_empty_labels, + output_format=output_format, + sort=sort, + **kwargs + ) + + if use_labels and output_format == 'columns': + # Expand list-of-dicts into columns + try: + df_temp = _pd.DataFrame(results) + except Exception: + df[out_col] = results + continue + + df_temp.index = df.index + if not df_temp.empty: + df_temp = df_temp.fillna('') + df[df_temp.columns] = df_temp.values + else: + df[out_col] = results + else: + df[out_col] = results elif len(input) > 1 and len(output) == 1 and len(model_id) == 1: model_id = [model_id[0] for _ in range(len(input))] @@ -636,17 +666,36 @@ def custom( single_model_id = model_id[0] df_temp = _pd.DataFrame(index=range(len(df))) for i, in_col in enumerate(input): - df_temp[output + str(i)] = _extract.custom( - df[in_col].astype(str).tolist(), - model_id=single_model_id, - first_element=first_element, - use_labels=use_labels, - case_sensitive=case_sensitive, - extract_raw=extract_raw, - use_spellcheck=use_spellcheck, - sort=sort, - **kwargs - ) + results = _extract.custom( + df[in_col].astype(str).tolist(), + model_id=single_model_id, + first_element=first_element, + use_labels=use_labels, + case_sensitive=case_sensitive, + extract_raw=extract_raw, + use_spellcheck=use_spellcheck, + include_empty_labels=include_empty_labels, + output_format=output_format, + sort=sort, + **kwargs + ) + + # If requested as columns and use_labels, expand and add with suffix + if use_labels and output_format == 'columns': + try: + df_exp = _pd.DataFrame(results) + except Exception: + df_temp[output + str(i)] = results + continue + + # Prefixing to avoid collisions: use original output name + index + df_exp.index = df.index + df_exp = df_exp.fillna('') + # Insert each column with a suffix to keep unique names + for col in df_exp.columns: + df_temp[f"{col}{i}"] = df_exp[col].values + else: + df_temp[output + str(i)] = results # Concatenate the results into a single column df[output] = [list(dict.fromkeys(_format.concatenate([x for x in row if x], ' '))) for row in df_temp.values.tolist()] @@ -654,17 +703,35 @@ def custom( else: # Iterate through the inputs, outputs and model_ids for in_col, out_col, model in zip(input, output, model_id): - df[out_col] = _extract.custom( - df[in_col].astype(str).tolist(), - model_id=model, - first_element=first_element, - use_labels=use_labels, - case_sensitive=case_sensitive, - extract_raw=extract_raw, - use_spellcheck=use_spellcheck, - sort=sort, - **kwargs - ) + results = _extract.custom( + df[in_col].astype(str).tolist(), + model_id=model, + first_element=first_element, + use_labels=use_labels, + case_sensitive=case_sensitive, + extract_raw=extract_raw, + use_spellcheck=use_spellcheck, + include_empty_labels=include_empty_labels, + output_format=output_format, + sort=sort, + **kwargs + ) + + if use_labels and output_format == 'columns': + try: + df_temp = _pd.DataFrame(results) + except Exception: + df[out_col] = results + continue + + df_temp.index = df.index + if not df_temp.empty: + df_temp = df_temp.fillna('') + df[df_temp.columns] = df_temp.values + else: + df[out_col] = results + else: + df[out_col] = results return df From 8a28a56bf43d33c83b283f33c02626cdea13ad36 Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Wed, 22 Jul 2026 19:32:50 +0300 Subject: [PATCH 14/15] fix failing test (#1078) --- tests/connectors/test_train.py | 38 ++++++++++++++++++++++------------ 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/tests/connectors/test_train.py b/tests/connectors/test_train.py index 24d40a75..dff1534f 100644 --- a/tests/connectors/test_train.py +++ b/tests/connectors/test_train.py @@ -1361,21 +1361,33 @@ def test_upsert_key_only(self): assert 'Blade Runner Upsert' in df['City'].values assert 'New Value' in df['City'].values - def test_missing_columns_error_message(self): - """ - Verify that INSERT/UPSERT/UPDATE raise the expected error - when incoming columns are not present in the existing model. - """ + def test_missing_columns_error_message(self, mocker): + """ + Verify that INSERT/UPSERT/UPDATE raise the expected error + when incoming columns are not present in the existing model. + """ - df = pd.DataFrame({ - "Key": ["k3"], - "Value": ["v3"], - "ExtraCol": ["x"] # This column does not exist in the model - }) + df = pd.DataFrame({ + "Key": ["k3"], + "Value": ["v3"], + "ExtraCol": ["x"] # This column does not exist in the model + }) - - # Test each action that performs the column-alignment check - for action in ("insert", "upsert", "update"): + # Mock the existing model so the test does not depend on a real, live model + mocker.patch( + "wrangles.data.model_content", + return_value={ + "Columns": ["Key", "Value"], + "Data": [["k1", "v1"]] + } + ) + mocker.patch( + "wrangles.data.model", + return_value={"variant": "key"} + ) + + # Test each action that performs the column-alignment check + for action in ("insert", "upsert", "update"): recipe = f""" write: - train.lookup: From 161b96228f6444de6fac331bea049e2d48e4afc8 Mon Sep 17 00:00:00 2001 From: Mariia Borodii Date: Thu, 23 Jul 2026 07:11:04 +0000 Subject: [PATCH 15/15] 992 refactor lookup actions bug request key colm --- tests/connectors/test_train.py | 551 ++++++++++++++++++++++------ tests/recipes/wrangles/test_main.py | 79 ++++ wrangles/connectors/train.py | 84 +++-- wrangles/lookup.py | 2 +- wrangles/recipe_wrangles/main.py | 242 ++++++------ 5 files changed, 685 insertions(+), 273 deletions(-) diff --git a/tests/connectors/test_train.py b/tests/connectors/test_train.py index dff1534f..e3f2c734 100644 --- a/tests/connectors/test_train.py +++ b/tests/connectors/test_train.py @@ -1,4 +1,6 @@ import uuid +import time +import importlib from pytest_mock import mocker @@ -8,6 +10,28 @@ import logging import re + +def _wait_for_lookup(model_id, predicate, timeout=15, interval=0.5): + """ + Poll a live train.lookup model until predicate(df) is true. + + The lookup API is eventually consistent, so a fixed sleep after a write + can race with a subsequent read. Polling on the actual expected state + avoids that race instead of guessing a delay. + """ + deadline = time.time() + timeout + result = None + while time.time() < deadline: + result = wrangles.recipe.run(f"read:\n - train.lookup:\n model_id: {model_id}") + if predicate(result): + return result + time.sleep(interval) + raise AssertionError( + f"Lookup model {model_id} did not reach expected state within {timeout}s; " + f"last seen columns: {list(result.columns) if result is not None else None}" + ) + + class LogCapture(logging.Handler): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -641,9 +665,12 @@ def test_insert_success(self): recipe = f""" write: - train.lookup: - model_id: 3c8f6707-2de4-4be3 + model_id: 3c8f6707-2de4-4be3 action: INSERT variant: key + columns: + - Key + - Value """ data = pd.DataFrame({ 'Key': ['Rachel', 'Dolores', 'TARS'], @@ -651,23 +678,26 @@ def test_insert_success(self): }) df = wrangles.recipe.run(recipe, dataframe=data) assert df.iloc[0]['Key'] == 'Rachel' and df.iloc[0]['Value'] == 'Blade Runner' - - def test_insert_duplicate_keys(self): - """ - Test insert fails when DataFrame contains duplicate keys - """ - df = pd.DataFrame({ - 'Key': ['Rachel', 'Rachel', 'Dolores'], # Duplicate Rachel - 'Value': ['Blade Runner', 'Not Rachel', 'Westworld'] - }) + + def test_insert_duplicate_keys(self): + """ + Test insert fails when DataFrame contains duplicate keys + """ + df = pd.DataFrame({ + 'Key': ['Rachel', 'Rachel', 'Dolores'], # Duplicate Rachel + 'Value': ['Blade Runner', 'Not Rachel', 'Westworld'] + }) recipe = f""" write: - train.lookup: - model_id: 3c8f6707-2de4-4be3 + model_id: 3c8f6707-2de4-4be3 action: INSERT variant: key - """ - with pytest.raises(ValueError, match="Lookup: All Keys must be unique"): + columns: + - Key + - Value + """ + with pytest.raises(ValueError, match="Lookup: All Keys must be unique"): wrangles.recipe.run(recipe, dataframe=df) @@ -680,17 +710,20 @@ def test_update_model_not_found(self): 'Value': ['Blade Runner 2049', 'Westworld Updated'] }) - recipe = """ - write: - - train.lookup: - model_id: test-model-id - action: UPDATE - """ - - # This would test with an actual existing model - # For testing purposes, we'll catch the expected error - with pytest.raises(RuntimeError, match="Access denied to model test-model-id"): - wrangles.recipe.run(recipe, dataframe=df) + recipe = """ + write: + - train.lookup: + model_id: test-model-id + action: UPDATE + columns: + - Key + - Value + """ + + # This would test with an actual existing model + # For testing purposes, we'll catch the expected error + with pytest.raises(RuntimeError, match="Access denied to model test-model-id"): + wrangles.recipe.run(recipe, dataframe=df) def test_action_parameter_validation_recipe(self): """ @@ -720,57 +753,66 @@ def test_update_model(self): 'Value': ['Updated Rachel', 'Updated Dolores', 'Updated Phipi'] }) - recipe = """ - write: - - train.lookup: + recipe = """ + write: + - train.lookup: model_id: 3c8f6707-2de4-4be3 - action: UPDATE - """ - - df = wrangles.recipe.run(recipe, dataframe=df) + action: UPDATE + columns: + - Key + - Value + """ + + df = wrangles.recipe.run(recipe, dataframe=df) assert df.iloc[0]['Key'] == 'Rachel' and df.iloc[0]['Value'] == 'Updated Rachel' - - def test_upsert_new_model_recipe(self): - """ - Test upsert creates new model when model_id doesn't exist - """ - df = pd.DataFrame({ - 'Key': ['Rachel', 'NewCharacter'], - 'Value': ['Updated Rachel', 'New Movie'] - }) + + def test_upsert_new_model_recipe(self): + """ + Test upsert creates new model when model_id doesn't exist + """ + df = pd.DataFrame({ + 'Key': ['Rachel', 'NewCharacter'], + 'Value': ['Updated Rachel', 'New Movie'] + }) model_name = f"model {{ {uuid.uuid4().hex[:8]} }}" - - recipe = f""" - write: - - train.lookup: - name: {model_name} - action: UPSERT - variant: key - """ - - result = wrangles.recipe.run(recipe, dataframe=df) - assert len(result) == 2 - assert 'NewCharacter' in result['Key'].tolist() + + recipe = f""" + write: + - train.lookup: + name: {model_name} + action: UPSERT + variant: key + columns: + - Key + - Value + """ + + result = wrangles.recipe.run(recipe, dataframe=df) + assert len(result) == 2 + assert 'NewCharacter' in result['Key'].tolist() assert result['Value'].tolist() == ['Updated Rachel', 'New Movie'] models = wrangles.data.user.models('lookup') assert any(m['name'] == model_name for m in models) - - def test_action_parameter_upsert(self): - """ - Test write method with action='UPSERT' - """ - df = pd.DataFrame({ - 'Key': ['Rachel'], - 'Value': ['Updated Rachel'] - }) - recipe = f""" - write: - - train.lookup: + + def test_action_parameter_upsert(self): + """ + Test write method with action='UPSERT' + """ + df = pd.DataFrame({ + 'Key': ['Rachel'], + 'Value': ['Updated Rachel'] + }) + recipe = f""" + write: + - train.lookup: model_id: b2cd1a8a-4d99-4be1 - action: UPSERT - variant: key - """ - result = wrangles.recipe.run(recipe, dataframe=df) + action: UPSERT + variant: key + columns: + - Key + - Value + """ + result = wrangles.recipe.run(recipe, dataframe=df) assert len(result) == 1 assert result.iloc[0]['Key'] == 'Rachel' and result.iloc[0]['Value'] == 'Updated Rachel' @@ -857,24 +899,61 @@ def test_semantic_lookup_with_embeddings_columns_no_error(self): assert 'Interstellar' in result.values assert 'Value' in result.columns and 'Description' in result.columns - def test_upsert_mismatched_columns_existing_model(self): + def test_upsert_adds_new_column(self): """ - UPSERT should fail when new data includes columns not present - in the existing model schema. - """ - df = pd.DataFrame({ - 'Key': ['Rachel'], - 'Other': ['Blade Runner'] - }) - recipe = """ - write: - - train.lookup: - model_id: b2cd1a8a-4d99-4be1 - action: UPSERT - variant: key + UPSERT with a column that doesn't exist in the model yet must add + that column; existing rows receive '' for the new column. + """ - with pytest.raises(ValueError, match="The following columns are not present in the existing model: Other"): - wrangles.recipe.run(recipe, dataframe=df) + MODEL = 'c060a706-db4a-4564' + + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: overwrite + variant: key + """, + dataframe=pd.DataFrame({ + 'Key': ['apple', 'banana'], + 'Value': ['red', 'yellow'], + }), + ) + _wait_for_lookup( + MODEL, + lambda df: set(df.columns) == {'Key', 'Value'} and set(df['Key']) == {'apple', 'banana'}, + ) + + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: upsert + variant: key + columns: + - Key + - Value + - Weight + """, + dataframe=pd.DataFrame({ + 'Key': ['apple'], + 'Value': ['green'], + 'Weight': ['1.0'], + }), + ) + + result = _wait_for_lookup( + MODEL, + lambda df: 'Weight' in df.columns + and (df.loc[df['Key'] == 'apple', 'Value'] == 'green').all(), + ) + assert 'Weight' in result.columns, "New column must be present after upsert" + apple = result[result['Key'] == 'apple'].iloc[0] + banana = result[result['Key'] == 'banana'].iloc[0] + assert apple['Weight'] == '1.0', "Updated row must carry the new column value" + assert banana['Weight'] == '', "Unaffected rows get '' for the new column" def test_upsert_missing_key_for_key_variant(self): """ @@ -890,6 +969,8 @@ def test_upsert_missing_key_for_key_variant(self): model_id: b2cd1a8a-4d99-4be1 action: UPSERT variant: key + columns: + - Value """ with pytest.raises(ValueError, match="'Key' column must be provided for 'key' variant"): wrangles.recipe.run(recipe, dataframe=df) @@ -905,12 +986,15 @@ def test_lookup_update_semantic_allows_no_key(self): model_id: 89637e77-7214-49a0 action: UPDATE variant: semantic + columns: + - Not Key + - Not Value """ data = pd.DataFrame({ 'Not Key': ['A', 'B'], 'Not Value': ['X', 'Y'] }) - + with pytest.raises(ValueError, match="UPDATE requires 'Key' or 'MatchingColumns' for non-key variants"): wrangles.recipe.run(recipe, dataframe=data) @@ -924,6 +1008,9 @@ def test_lookup_update_key_variant_requires_key(self): model_id: 3c8f6707-2de4-4be3 action: UPDATE variant: key + columns: + - NotKey + - Value """ data = pd.DataFrame({ 'NotKey': ['A', 'B'], @@ -985,6 +1072,11 @@ def test_lookup_matchingcolumns_missing_raises_insert_existing_model(self): model_id: 3c8f6707-2de4-4be3 action: INSERT variant: key + columns: + - City + - Country + - Key + - Value settings: MatchingColumns: - Not City @@ -1028,6 +1120,9 @@ def test_upsert_new_matchingcolumns_missing(self): - train.lookup: name: 083ed6fe-a073-4b1a action: UPSERT + columns: + - City + - Country settings: MatchingColumns: - NotKey @@ -1049,6 +1144,9 @@ def test_upsert_existing_matchingcolumns_missing(self): - train.lookup: model_id: 083ed6fe-a073-4b1a action: UPSERT + columns: + - City + - Country settings: MatchingColumns: - NotKey @@ -1071,6 +1169,9 @@ def test_insert_matchingcolumns_missing(self): - train.lookup: model_id: 083ed6fe-a073-4b1a action: INSERT + columns: + - City + - Country settings: MatchingColumns: - NotKey @@ -1123,6 +1224,11 @@ def test_lookup_update_matchingcolumns_string(self, caplog): model_id: 4202c974-430a-46b9 action: update variant: semantic + columns: + - City + - Country + - Code + - Currency settings: MatchingColumns: City ''', @@ -1154,8 +1260,13 @@ def test_lookup_update_multiple_matchingcolumns_list(self, caplog): model_id: 4202c974-430a-46b9 action: update variant: semantic + columns: + - City + - Country + - Code + - Currency settings: - MatchingColumns: + MatchingColumns: - City - Country ''', @@ -1186,6 +1297,11 @@ def test_lookup_insert_matchingcolumns_single(self, caplog): - train.lookup: model_id: 4202c974-430a-46b9 action: insert + columns: + - City + - Country + - Code + - Currency settings: MatchingColumns: City ''', @@ -1217,8 +1333,13 @@ def test_lookup_insert_matchingcolumns_list(self, caplog): model_id: 4202c974-430a-46b9 action: insert variant: semantic + columns: + - City + - Country + - Code + - Currency settings: - MatchingColumns: + MatchingColumns: - City ''', dataframe=df @@ -1247,6 +1368,11 @@ def test_lookup_upsert_matchingcolumns_single(self, caplog): - train.lookup: model_id: 4202c974-430a-46b9 action: upsert + columns: + - City + - Country + - Code + - Currency settings: MatchingColumns: City ''', @@ -1277,8 +1403,13 @@ def test_lookup_upsert_matchingcolumns_list(self, caplog): model_id: 4202c974-430a-46b9 action: upsert variant: semantic + columns: + - City + - Country + - Code + - Currency settings: - MatchingColumns: + MatchingColumns: - City - Country ''', @@ -1288,7 +1419,7 @@ def test_lookup_upsert_matchingcolumns_list(self, caplog): messages = [record.message for record in caplog.records if record.levelname == "INFO"] assert any("Lookup UPSERT: 1 rows inserted, 2 rows updated. Total rows:" in msg for msg in messages), "Log should mention rows inserted (variant specified)" - def test_insert_key_only(self, caplog): + def test_insert_key_only(self): """ Test INSERT with only a Key column and no MatchingColumns/settings. """ @@ -1298,6 +1429,9 @@ def test_insert_key_only(self, caplog): model_id: 12b7ac66-7418-45b5 action: INSERT variant: key + columns: + - Key + - City """ data = pd.DataFrame({ 'Key': ['Rachel', 'Dolores'], @@ -1326,6 +1460,9 @@ def test_update_key_only(self): - train.lookup: model_id: 12b7ac66-7418-45b5 action: UPDATE + columns: + - Key + - City """ data = pd.DataFrame({ 'Key': ['Alice'], @@ -1352,6 +1489,9 @@ def test_upsert_key_only(self): model_id: 12b7ac66-7418-45b5 action: UPSERT variant: key + columns: + - Key + - City """ data = pd.DataFrame({ 'Key': ['Charlie', 'NewKey'], @@ -1360,47 +1500,236 @@ def test_upsert_key_only(self): df = wrangles.recipe.run(recipe, dataframe=data) assert 'Blade Runner Upsert' in df['City'].values assert 'New Value' in df['City'].values - - def test_missing_columns_error_message(self, mocker): + + def test_upsert_preserves_unspecified_columns(self): """ - Verify that INSERT/UPSERT/UPDATE raise the expected error - when incoming columns are not present in the existing model. + UPSERT with a partial set of columns must not delete columns that are + present in the model but absent from the incoming DataFrame. + Issue #992: unspecified columns were silently dropped. """ + MODEL = 'be1fcb1c-08ea-43bc' - df = pd.DataFrame({ - "Key": ["k3"], - "Value": ["v3"], - "ExtraCol": ["x"] # This column does not exist in the model - }) + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: overwrite + variant: key + """, + dataframe=pd.DataFrame({ + 'Key': ['apple', 'banana', 'cherry'], + 'Schema': ['fruit', 'fruit', 'fruit'], + 'Mapping': ['red', 'yellow', 'red'], + }), + ) + _wait_for_lookup( + MODEL, + lambda df: set(df.columns) == {'Key', 'Schema', 'Mapping'} + and set(df['Key']) == {'apple', 'banana', 'cherry'}, + ) + + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: upsert + variant: key + columns: + - Key + - Mapping + """, + dataframe=pd.DataFrame({'Key': ['apple'], 'Mapping': ['green']}), + ) + + result = _wait_for_lookup( + MODEL, + lambda df: 'Schema' in df.columns + and (df.loc[df['Key'] == 'apple', 'Mapping'] == 'green').all(), + ) + assert 'Schema' in result.columns, "Schema must be preserved after partial upsert" + apple = result[result['Key'] == 'apple'].iloc[0] + assert apple['Mapping'] == 'green', "Updated column must reflect new value" + assert apple['Schema'] == 'fruit', "Unspecified column must retain original value" + + def test_insert_preserves_unspecified_columns(self): + """ + INSERT must not drop columns that exist in the model but are absent + from the incoming DataFrame. New rows get '' for unspecified columns. + """ + MODEL = '29a87d05-0617-4283' + + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: overwrite + variant: key + """, + dataframe=pd.DataFrame({ + 'Key': ['apple', 'banana'], + 'Schema': ['fruit', 'fruit'], + 'Mapping': ['red', 'yellow'], + }), + ) + _wait_for_lookup( + MODEL, + lambda df: set(df.columns) == {'Key', 'Schema', 'Mapping'} + and set(df['Key']) == {'apple', 'banana'}, + ) - # Mock the existing model so the test does not depend on a real, live model + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: insert + variant: key + columns: + - Key + - Mapping + """, + dataframe=pd.DataFrame({'Key': ['cherry'], 'Mapping': ['red']}), + ) + + result = _wait_for_lookup( + MODEL, + lambda df: 'Schema' in df.columns and len(df) == 3, + ) + assert 'Schema' in result.columns, "Schema column must be preserved after insert" + assert len(result) == 3, "New row must be added" + cherry = result[result['Key'] == 'cherry'].iloc[0] + assert cherry['Mapping'] == 'red' + assert cherry['Schema'] == '', "Unspecified column for new row must be empty string" + + def test_update_preserves_unspecified_columns(self, mocker): + """ + UPDATE must only modify the columns present in the incoming DataFrame; + all other columns — on both the updated and untouched rows — must be + preserved exactly. + + Mocked rather than run against the live API: this previously polled + a real model for up to 15s and was flaky under eventual consistency. + """ + MODEL = 'dcac58c3-7da0-403f' + store = {} + + mocker.patch("wrangles.data.model", return_value={'variant': 'key'}) mocker.patch( "wrangles.data.model_content", - return_value={ - "Columns": ["Key", "Value"], - "Data": [["k1", "v1"]] - } + side_effect=lambda id, version_id=None: store[id] ) mocker.patch( - "wrangles.data.model", - return_value={"variant": "key"} + "wrangles.train.train.lookup", + side_effect=lambda data, name=None, model_id=None, settings=None: store.__setitem__(model_id, data) + ) + + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: overwrite + variant: key + """, + dataframe=pd.DataFrame({ + 'Key': ['apple', 'banana', 'cherry'], + 'Schema': ['fruit', 'fruit', 'fruit'], + 'Mapping': ['red', 'yellow', 'red'], + }), ) - # Test each action that performs the column-alignment check - for action in ("insert", "upsert", "update"): + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {MODEL} + action: update + variant: key + columns: + - Key + - Mapping + """, + dataframe=pd.DataFrame({'Key': ['apple'], 'Mapping': ['green']}), + ) + + result = wrangles.recipe.run(f"read:\n - train.lookup:\n model_id: {MODEL}") + assert 'Schema' in result.columns, "Schema must be preserved after update" + apple = result[result['Key'] == 'apple'].iloc[0] + banana = result[result['Key'] == 'banana'].iloc[0] + assert apple['Mapping'] == 'green', "Updated Mapping must reflect new value" + assert apple['Schema'] == 'fruit', "Unspecified Schema on updated row must be preserved" + assert banana['Mapping'] == 'yellow', "Untouched row must be unchanged" + assert banana['Schema'] == 'fruit', "Untouched row Schema must be unchanged" + + def test_missing_columns_error_message(self, monkeypatch): + """ + INSERT and UPDATE raise an error when incoming columns are not present + in the existing model schema. UPSERT is intentionally excluded: it + allows new columns (adds them to the model schema, filling existing + rows with ''). + """ + train_connector = importlib.import_module("wrangles.connectors.train") + monkeypatch.setattr( + train_connector._data, + "model", + lambda model_id: {"variant": "key"} + ) + monkeypatch.setattr( + train_connector._data, + "model_content", + lambda model_id: { + "Columns": ["Key", "Value"], + "Data": [["k1", "v1"], ["k2", "v2"]], + } + ) + + df = pd.DataFrame({ + "Key": ["k3"], + "Value": ["v3"], + "ExtraCol": ["x"] # does not exist in the model + }) + + for action in ("insert", "update"): recipe = f""" write: - train.lookup: model_id: bc3ee6a0-e104-4700 action: {action} variant: key - + columns: + - Key + - Value + - ExtraCol """ with pytest.raises(ValueError, match="Lookup: The following columns are not present in the existing model: ExtraCol"): wrangles.recipe.run(recipe, dataframe=df) - - -def test_lookup_write_logs_new_model_id(caplog): + + def test_columns_required_for_partial_update_actions(self): + """ + INSERT, UPDATE, and UPSERT must raise ValueError when 'columns' is not specified. + OVERWRITE does not require it. + """ + df = pd.DataFrame({'Key': ['apple'], 'Value': ['red']}) + for action in ('insert', 'update', 'upsert'): + with pytest.raises( + ValueError, + match=f"Lookup: 'columns' is required for action '{action}'" + ): + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: 3c8f6707-2de4-4be3 + action: {action} + """, + dataframe=df, + ) + + +def test_lookup_write_logs_new_model_id(caplog): """ Integration test for lookup model creation logging """ diff --git a/tests/recipes/wrangles/test_main.py b/tests/recipes/wrangles/test_main.py index 8fb0555d..a098ee09 100644 --- a/tests/recipes/wrangles/test_main.py +++ b/tests/recipes/wrangles/test_main.py @@ -6111,10 +6111,89 @@ def fail_on_2nd_batch(df): +def _seed_lookup_model(model_id, dataframe, timeout=15, interval=0.5): + """ + Overwrite a live train.lookup model and poll until the write is visible. + + The lookup API is eventually consistent, so callers must wait for the + write to land before reading it back rather than assuming it's immediate. + """ + wrangles.recipe.run( + f""" + write: + - train.lookup: + model_id: {model_id} + action: overwrite + variant: key + """, + dataframe=dataframe, + ) + expected_keys = set(dataframe['Key']) + deadline = time.time() + timeout + while time.time() < deadline: + result = wrangles.recipe.run(f"read:\n - train.lookup:\n model_id: {model_id}") + if set(result.columns) == set(dataframe.columns) and set(result['Key']) == expected_keys: + return + time.sleep(interval) + raise AssertionError(f"Lookup model {model_id} did not reach seeded state within {timeout}s") + + class TestLookup: """ Test lookup wrangle """ + def test_lookup_output_key_only(self): + """ + Specifying output: Key should return the looked-up key string, not a dict. + Issue #992: 'Key' was not in metadata["settings"]["columns"] so the unnamed- + columns path was hit and the full dict was returned instead. + """ + _seed_lookup_model( + '3f23acaf-a2e6-4327', + pd.DataFrame({ + 'Key': ['apple', 'banana', 'cherry'], + 'Schema': ['fruit', 'fruit', 'fruit'], + }), + ) + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: fruit + output: Key + model_id: 3f23acaf-a2e6-4327 + """, + dataframe=pd.DataFrame({'fruit': ['apple', 'banana', 'cherry']}), + ) + assert df['Key'].tolist() == ['apple', 'banana', 'cherry'] + + def test_lookup_output_key_and_value_column(self): + """ + Specifying output: [Key, Schema] must work without error. + Issue #992: mixing 'Key' with a real model column raised ValueError. + """ + _seed_lookup_model( + '33961b4e-92f5-4705', + pd.DataFrame({ + 'Key': ['apple', 'banana', 'cherry'], + 'Schema': ['fruit', 'fruit', 'fruit'], + }), + ) + df = wrangles.recipe.run( + """ + wrangles: + - lookup: + input: fruit + output: + - Key + - Schema + model_id: 33961b4e-92f5-4705 + """, + dataframe=pd.DataFrame({'fruit': ['apple', 'banana', 'cherry']}), + ) + assert df['Key'].tolist() == ['apple', 'banana', 'cherry'] + assert df['Schema'].tolist() == ['fruit', 'fruit', 'fruit'] + def test_lookup_mode_by_row_default(self): """ Test lookup with by_row mode (default behavior) diff --git a/wrangles/connectors/train.py b/wrangles/connectors/train.py index 6b5205a9..30b849b2 100644 --- a/wrangles/connectors/train.py +++ b/wrangles/connectors/train.py @@ -234,16 +234,17 @@ def read(model_id: str) -> _pd.DataFrame: description: Specific model to read """ - def write(df: _pd.DataFrame, name: str = None, model_id: str = None, settings: dict = None, variant: str = 'key', action: str = 'overwrite') -> None: + def write(df: _pd.DataFrame, columns: list = None, name: str = None, model_id: str = None, settings: dict = None, variant: str = 'key', action: str = 'overwrite') -> None: """ Train a new or existing lookup wrangle :param df: DataFrame to be written to a file + :param columns: Columns to include. Required for INSERT, UPDATE and UPSERT — only the listed columns will be added or updated. :param name: Name to give to a new Wrangle that will be created :param model_id: Model to be updated. Either this or name must be provided :param settings: Specific settings to apply to the wrangle :param variant: Variant of the Lookup Wrangle that will be created (key or semantic) - :param action: Action to take when training the lookup wrangle (insert, update, upsert) + :param action: Action to take when training the lookup wrangle (overwrite, insert, update, upsert) """ _logging.info(f": Training Lookup Wrangle") if settings is None: @@ -257,6 +258,18 @@ def write(df: _pd.DataFrame, name: str = None, model_id: str = None, settings: d act = (action or 'overwrite').upper() settings = dict(settings or {}) + # columns is required for partial-update actions + if columns is None and act in ('INSERT', 'UPDATE', 'UPSERT'): + raise ValueError( + f"Lookup: 'columns' is required for action '{action}'. " + "Specify the columns to add or update so that unrelated columns are not modified." + ) + + # Filter the DataFrame to only the requested columns + if columns is not None: + columns = _wildcard_expansion(df.columns, columns) + df = df[columns] + def _get_columns_from_payload(payload): # payload can be a DataFrame or the 'tight' dict produced by _to_tight if isinstance(payload, dict) and 'Columns' in payload: @@ -333,18 +346,8 @@ def _set_variant(mid: str, var: str) -> str: columns=existing_content['Columns'] ) - # Validate column compatibility with existing model - requested_cols = new_data['Columns'] - missing_in_existing = [c for c in requested_cols if c not in existing_df_all.columns] - if missing_in_existing: - raise ValueError( - "Lookup: The following columns are not present in the existing model: " - + ", ".join(missing_in_existing) - ) - - existing_df = existing_df_all[requested_cols] # Ensure same column order - # For key variant, ensure new data contains Key column + requested_cols = new_data['Columns'] normalized_variant = settings.get('variant', variant) if normalized_variant == 'key' and 'Key' not in df.columns: raise ValueError("Lookup: 'Key' column must be provided for 'key' variant") @@ -353,13 +356,17 @@ def _set_variant(mid: str, var: str) -> str: updated = 0 # Merge data - avoid duplicates based on Key column - if variant== 'key' and 'Key' in existing_df.columns and 'Key' in df.columns: + if variant == 'key' and 'Key' in existing_df_all.columns and 'Key' in df.columns: if df['Key'].duplicated().any(): raise ValueError("Lookup: All Keys must be unique") - existing_keys = set(existing_df['Key'].tolist()) + existing_keys = set(existing_df_all['Key'].tolist()) - # Start with current data - merged_df = existing_df.copy() + # Use ALL existing columns as the base so unspecified columns are preserved. + # New columns introduced by this upsert are added with '' for existing rows. + merged_df = existing_df_all.copy() + new_cols = [c for c in requested_cols if c not in merged_df.columns] + for col in new_cols: + merged_df[col] = '' # Apply updates for matching keys for _, row in df.iterrows(): @@ -367,13 +374,14 @@ def _set_variant(mid: str, var: str) -> str: if key in existing_keys: mask = merged_df['Key'] == key for col in df.columns: - if col != 'Key' and col in merged_df.columns: + if col != 'Key': merged_df.loc[mask, col] = row[col] updated += 1 else: - # Insert new rows for non-existing keys + # New key: fill columns absent from the incoming row with empty string + new_row = {col: (row[col] if col in df.columns else '') for col in merged_df.columns} merged_df = _pd.concat( - [merged_df, _pd.DataFrame([row[merged_df.columns].tolist()], columns=merged_df.columns)], + [merged_df, _pd.DataFrame([new_row])], ignore_index=True ) inserted += 1 @@ -422,7 +430,7 @@ def _set_variant(mid: str, var: str) -> str: merged_df = merged_df[existing_df_all.columns] else: # Without MatchingColumns, append all - merged_df = _pd.concat([existing_df, df], ignore_index=True) + merged_df = _pd.concat([existing_df_all, df], ignore_index=True) inserted = len(df) merged_data = { @@ -598,10 +606,11 @@ def _set_variant(mid: str, var: str) -> str: inserted = len(df) merged_df = _pd.concat([existing_df, df], ignore_index=True) - merged_data = { - 'Columns': merged_df.columns.tolist(), - 'Data': merged_df.values.tolist() - } + merged_df = merged_df.fillna('') + merged_data = { + 'Columns': merged_df.columns.tolist(), + 'Data': merged_df.values.tolist() + } _validate_matching_columns(merged_data, settings) total_rows = len(merged_df) _train.lookup(merged_data, None, model_id, settings) @@ -611,6 +620,8 @@ def _set_variant(mid: str, var: str) -> str: type: object description: Train a new or existing Lookup Wrangle additionalProperties: false + required: + - action properties: name: type: string @@ -620,21 +631,28 @@ def _set_variant(mid: str, var: str) -> str: description: Model to be updated. Either this or a name must be provided columns: type: array - description: Columns to submit - variant: - type: string - description: Variant of the Lookup Wrangle that will be created - enum: - - key - - semantic - action: + description: >- + Columns to include from the DataFrame. + Required for INSERT, UPDATE and UPSERT — only the listed columns + will be added or updated; all other columns in the model are left + unchanged. + items: + type: string + action: type: string description: Action to take when training the lookup wrangle + default: overwrite enum: - insert - update - upsert - overwrite + variant: + type: string + description: Variant of the Lookup Wrangle that will be created + enum: + - key + - semantic """ diff --git a/wrangles/lookup.py b/wrangles/lookup.py index 49a6eb58..f8117e01 100644 --- a/wrangles/lookup.py +++ b/wrangles/lookup.py @@ -69,7 +69,7 @@ def lookup( f'{_config.api_host}/wrangles/lookup', { "model_id": model_id, - "columns": _json.dumps(columns or metadata["settings"]["columns"]), + "columns": _json.dumps(columns if columns is not None else metadata.get("settings", {}).get("columns", [])), **kwargs }, input, diff --git a/wrangles/recipe_wrangles/main.py b/wrangles/recipe_wrangles/main.py index 83426351..108d9980 100644 --- a/wrangles/recipe_wrangles/main.py +++ b/wrangles/recipe_wrangles/main.py @@ -903,7 +903,7 @@ def lookup( input: str, output: _Union[str, list] = None, model_id: str = None, - lookup_mode: str = 'by_row', + lookup_mode: str = 'by_row', n: int = None, **kwargs ) -> _pd.DataFrame: @@ -999,7 +999,14 @@ def lookup( list(val.values())[0] if isinstance(val, dict) else val for val in output ] - + + # 'Key' is valid as output — it echoes the input key value, not a model value column. + # Strip it before routing so the named/unnamed check is not confused by it. + _key_indices = {i for i, col in enumerate(wrangle_output) if col == 'Key'} + _key_out_cols = [output[i] for i in sorted(_key_indices)] + _wrangle_cols = [col for i, col in enumerate(wrangle_output) if i not in _key_indices] + _out_cols = [col for i, col in enumerate(output) if i not in _key_indices] + # Remove matrix_variables from kwargs if present before passing to _lookup def _clean_kwargs(kwargs): if 'matrix_variables' in kwargs: @@ -1010,100 +1017,87 @@ def _clean_kwargs(kwargs): # Distribute the n matches for each row across the output columns, # ranked match i goes to output column i def _distribute_n_matches(data): - for i, out in enumerate(output): + for i, out in enumerate(_out_cols): df[out] = [row[i] if isinstance(row, list) and i < len(row) else None for row in data] # Perform lookup based on lookup_mode + model_columns = metadata.get("settings", {}).get("columns", []) if lookup_mode == 'by_row': - # Current behavior - process all rows - if all([col in metadata["settings"]["columns"] for col in wrangle_output]): - # User specified all columns from the wrangle - data = _lookup( - df[input].values.tolist(), - model_id, - columns=wrangle_output, - n=n, - **_clean_kwargs(kwargs) - ) - if n and n > 1 and len(output) == n: - # Distribute: each output column gets the nth match - _distribute_n_matches(data) - elif n and n > 1 and len(output) > 1: - raise ValueError( - f'When n > 1 and multiple output columns are provided, the number ' - f'of output columns ({len(output)}) must equal n ({n}).' + if _wrangle_cols: + if all([col in model_columns for col in _wrangle_cols]): + data = _lookup( + df[input].values.tolist(), + model_id, + columns=_wrangle_cols, + n=n, + **_clean_kwargs(kwargs) ) - else: - df[output] = data - elif not any([col in metadata["settings"]["columns"] for col in wrangle_output]): - # User specified no columns from the wrangle - data = _lookup( - df[input].values.tolist(), - model_id, - n=n, - **_clean_kwargs(kwargs) - ) - if n and n > 1 and len(output) == n: - # Distribute: each output column gets the nth match - _distribute_n_matches(data) - elif n and n > 1 and len(output) > 1: - raise ValueError( - f'When n > 1 and multiple output columns are provided, the number ' - f'of output columns ({len(output)}) must equal n ({n}).' + if n and n > 1 and len(_out_cols) == n: + # Distribute: each output column gets the nth match + _distribute_n_matches(data) + elif n and n > 1 and len(_out_cols) > 1: + raise ValueError( + f'When n > 1 and multiple output columns are provided, the number ' + f'of output columns ({len(_out_cols)}) must equal n ({n}).' + ) + else: + df[_out_cols] = data + + elif not any([col in model_columns for col in _wrangle_cols]): + data = _lookup( + df[input].values.tolist(), + model_id, + n=n, + **_clean_kwargs(kwargs) ) + if n and n > 1 and len(_out_cols) == n: + # Distribute: each output column gets the nth match + _distribute_n_matches(data) + elif n and n > 1 and len(_out_cols) > 1: + raise ValueError( + f'When n > 1 and multiple output columns are provided, the number ' + f'of output columns ({len(_out_cols)}) must equal n ({n}).' + ) + else: + for out in _out_cols: + df[out] = data else: - for out in output: - df[out] = data - else: - # User specified a mixture of unrecognized columns and columns from the wrangle - raise ValueError('Lookup may only contain all named or unnamed columns.') - + raise ValueError('Lookup may only contain all named or unnamed columns.') + elif lookup_mode == 'by_dataframe': - # Optimized - lookup unique values once, then map to all rows - unique_values = df[input].unique() - - if all([col in metadata["settings"]["columns"] for col in wrangle_output]): - # User specified all columns from the wrangle - unique_data = _lookup( - unique_values.tolist(), - model_id, - columns=wrangle_output, - **_clean_kwargs(kwargs) - ) - - # Create mapping from values to results - value_to_result = dict(zip(unique_values, unique_data)) - - # Map results back to all rows - extract correct values - if len(output) == 1: - df[output[0]] = df[input].map( - lambda x: value_to_result.get(x, [])[0] if x in value_to_result and value_to_result[x] else "" - ) - else: - for i, out_col in enumerate(output): - df[out_col] = df[input].map( - lambda x: value_to_result.get(x, [])[i] if x in value_to_result and len(value_to_result[x]) > i else "" - ) - - elif not any([col in metadata["settings"]["columns"] for col in wrangle_output]): - # User specified no columns from the wrangle - unique_data = _lookup( - unique_values.tolist(), - model_id, - **_clean_kwargs(kwargs) - ) - - # Create mapping from values to results (preserve full dict as in by_row) - value_to_result = dict(zip(unique_values, unique_data)) - - # Map results back to all rows - output is the full dict as in by_row - for out in output: - df[out] = df[input].map(lambda x: value_to_result.get(x, {})) - else: - # User specified a mixture of unrecognized columns and columns from the wrangle - raise ValueError('Lookup may only contain all named or unnamed columns.') - - elif lookup_mode == 'by_matrix': + unique_values = df[input].unique() + if _wrangle_cols: + if all([col in model_columns for col in _wrangle_cols]): + unique_data = _lookup( + unique_values.tolist(), + model_id, + columns=_wrangle_cols, + **_clean_kwargs(kwargs) + ) + value_to_result = dict(zip(unique_values, unique_data)) + if len(_out_cols) == 1: + df[_out_cols[0]] = df[input].map( + lambda x: value_to_result.get(x, [])[0] if x in value_to_result and value_to_result[x] else "" + ) + else: + for i, out_col in enumerate(_out_cols): + df[out_col] = df[input].map( + lambda x: value_to_result.get(x, [])[i] if x in value_to_result and len(value_to_result[x]) > i else "" + ) + elif not any([col in model_columns for col in _wrangle_cols]): + unique_data = _lookup( + unique_values.tolist(), + model_id, + **_clean_kwargs(kwargs) + ) + value_to_result = dict(zip(unique_values, unique_data)) + for out in _out_cols: + df[out] = df[input].map(lambda x: value_to_result.get(x, {})) + else: + # User specified a mixture of unrecognized columns and columns from the wrangle + raise ValueError('Lookup may only contain all named or unnamed columns.') + + elif lookup_mode == 'by_matrix': # Get matrix variables and permutations matrix_vars = kwargs.get('matrix_variables', []) if not matrix_vars: @@ -1130,48 +1124,40 @@ def _distribute_n_matches(data): # Get unique values for this permutation perm_values = df.loc[mask, input].unique() - if all([col in metadata["settings"]["columns"] for col in wrangle_output]): - # User specified all columns from the wrangle - perm_data = _lookup( - perm_values.tolist(), - model_id, - columns=wrangle_output, - **_clean_kwargs(kwargs) - ) - - # Create mapping for this permutation - value_to_result = dict(zip(perm_values, perm_data)) - - # Apply results to matching rows - extract correct values - if len(output) == 1: - df.loc[mask, output[0]] = df.loc[mask, input].map( - lambda x: value_to_result.get(x, [])[0] if x in value_to_result and value_to_result[x] else "" - ) - else: - for i, out_col in enumerate(output): - df.loc[mask, out_col] = df.loc[mask, input].map( - lambda x: value_to_result.get(x, [])[i] if x in value_to_result and len(value_to_result[x]) > i else "" - ) - - elif not any([col in metadata["settings"]["columns"] for col in wrangle_output]): - # User specified no columns from the wrangle - perm_data = _lookup( - perm_values.tolist(), - model_id, - **_clean_kwargs(kwargs) - ) - - # Create mapping for this permutation - value_to_result = dict(zip(perm_values, perm_data)) - - # Apply results to matching rows - extract correct values - for out in output: - df.loc[mask, out] = df.loc[mask, input].map(lambda x: value_to_result.get(x, {})) - else: - # User specified a mixture of unrecognized columns and columns from the wrangle - raise ValueError('Lookup may only contain all named or unnamed columns.') + if _wrangle_cols: + if all([col in model_columns for col in _wrangle_cols]): + perm_data = _lookup( + perm_values.tolist(), + model_id, + columns=_wrangle_cols, + **_clean_kwargs(kwargs) + ) + value_to_result = dict(zip(perm_values, perm_data)) + if len(_out_cols) == 1: + df.loc[mask, _out_cols[0]] = df.loc[mask, input].map( + lambda x: value_to_result.get(x, [])[0] if x in value_to_result and value_to_result[x] else "" + ) + else: + for i, out_col in enumerate(_out_cols): + df.loc[mask, out_col] = df.loc[mask, input].map( + lambda x: value_to_result.get(x, [])[i] if x in value_to_result and len(value_to_result[x]) > i else "" + ) + elif not any([col in model_columns for col in _wrangle_cols]): + perm_data = _lookup( + perm_values.tolist(), + model_id, + **_clean_kwargs(kwargs) + ) + value_to_result = dict(zip(perm_values, perm_data)) + for out in _out_cols: + df.loc[mask, out] = df.loc[mask, input].map(lambda x: value_to_result.get(x, {})) + else: + raise ValueError('Lookup may only contain all named or unnamed columns.') else: raise ValueError(f"Invalid lookup_mode: {lookup_mode}. Must be 'by_row', 'by_dataframe', or 'by_matrix'") + + for _key_col in _key_out_cols: + df[_key_col] = df[input].values else: raise ValueError('model_id is required for lookup')