992-refactor-lookup-actions-bug-request-key-colm

[mborodii-prog]

Summary

Two bugs in train.lookup write actions and the lookup wrangle are resolved.

---

Bug 1 — upsert / insert / update destroyed columns not in the incoming DataFrame

Problem

When running any non-overwrite action with a DataFrame that only contains a subset of the model's columns, the columns absent from the DataFrame were silently deleted from the model.

# Model has: Key, Schema, Mapping
# User only wants to update Mapping for one row — Schema is intentionally omitted

wrangles.recipe.run("""
write:
  - train.lookup:
      model_id: <model_id>
      action: upsert
      variant: key
""", dataframe=pd.DataFrame({'Key': ['apple'], 'Mapping': ['green']}))

# BEFORE (bug): model now has Key, Mapping only — Schema was dropped
# AFTER (fix):  model still has Key, Schema, Mapping — Schema preserved

The same issue existed for insert (new rows had NaN in unspecified columns, which broke JSON serialisation) and update.

Root cause

wrangles/connectors/train.py — UPSERT branch (~line 345):

# BEFORE — slices the existing data to only the incoming columns
existing_df = existing_df_all[requested_cols]   # ← BUG
merged_df   = existing_df.copy()                # merge base is now partial

# AFTER — uses ALL existing columns as the merge base
merged_df = existing_df_all.copy()
# New columns introduced by the upsert are added with '' for existing rows
new_cols = [c for c in requested_cols if c not in merged_df.columns]
for col in new_cols:
    merged_df[col] = ''

The same existing_df reference was also used in the non-key no-MatchingColumns path (_pd.concat([existing_df, df])), which is now existing_df_all.

For INSERT, _pd.concat can produce NaN for columns absent from new rows. Added merged_df = merged_df.fillna('') before building the JSON payload.

upsert is the only action that accepts new columns — it treats them as an additive schema extension.
insert and update still raise ValueError when a column that doesn't exist in the model is supplied.

---

Bug 2 — Key column could not be requested as output in a lookup wrangle

Problem

wrangles:
  - lookup:
      input: fruit
      output: Key           # wanted the key string back
      model_id: <model_id>

Case A — output: Key only
Key is not in metadata["settings"]["columns"] (which lists only value columns).
The routing fell into the "no named columns" branch, treating Key as an arbitrary output name and returning a full dict of all columns instead of the key string.

Case B — output: [Key, Schema]
Key not in settings cols, Schema is → routing hit the "mixture" branch → ValueError: Lookup may only contain all named or unnamed columns.

Root cause

wrangles/recipe_wrangles/main.py — lookup() (~line 980):

# routing check used wrangle_output directly, which included 'Key'
if all([col in metadata["settings"]["columns"] for col in wrangle_output]):
    ...
elif not any([col in metadata["settings"]["columns"] for col in wrangle_output]):
    ...
else:
    raise ValueError('Lookup may only contain all named or unnamed columns.')

Fix

Before the routing check, Key entries are extracted from wrangle_output. The routing and API call use only the remaining non-Key columns. After the lookup completes, any requested Key output column is populated directly from the input values.

# Strip 'Key' from routing — it echoes the input, not a model value column
_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]

# ... (routing + API call using _wrangle_cols / _out_cols) ...

# Populate Key output column(s) with the input key values
for _key_col in _key_out_cols:
    df[_key_col] = df[input].values

Applied consistently across by_row, by_dataframe, and by_matrix modes.

Examples after fix

import wrangles, pandas as pd

df = pd.DataFrame({'fruit': ['apple', 'banana', 'cherry']})

# Case A — Key only
result = wrangles.recipe.run("""
wrangles:
  - lookup:
      input: fruit
      output: Key
      model_id: <model_id>
""", dataframe=df)
# result['Key'] → ['apple', 'banana', 'cherry']  ✓

# Case B — Key + value column together
result = wrangles.recipe.run("""
wrangles:
  - lookup:
      input: fruit
      output:
        - Key
        - Schema
      model_id: <model_id>
""", dataframe=df)
# result['Key']    → ['apple', 'banana', 'cherry']  ✓
# result['Schema'] → ['fruit', 'fruit',  'fruit']   ✓

---

[mborodii-prog]

@thomasstvr @ebhills as we discussed yesterday I made columns parameter mandatory ## Breaking change: columns is now required for insert, update, and upsert

train.lookup (Python API and recipe) now requires an explicit columns list when using any partial-update action (insert, update, upsert). Calling one of these actions without columns raises a ValueError.

ValueError: Lookup: 'columns' is required for action 'upsert'.
Specify the columns to add or update so that unrelated columns are not modified.

overwrite is unaffected — columns remains optional there.

---

Migration

Before (will now raise ValueError):

# Python API
wrangles.connectors.train.lookup.write(
    df,
    model_id='<model-id>',
    action='upsert',
    variant='key',
)

# Recipe
write:
  - train.lookup:
      model_id: <model-id>
      action: upsert
      variant: key

After (explicit columns required):

# Python API
wrangles.connectors.train.lookup.write(
    df,
    model_id='<model-id>',
    columns=['Key', 'Mapping', 'Weight'],
    action='upsert',
    variant='key',
)

# Recipe
write:
  - train.lookup:
      model_id: <model-id>
      action: upsert
      columns:
        - Key
        - Mapping
        - Weight
      variant: key

---

Behaviour by action

Action	columns required?	Effect
overwrite	No	Replaces the entire model with the DataFrame
insert	Yes	Appends rows whose Key doesn't exist yet; only the listed columns are written
update	Yes	Updates only existing keys; only the listed columns change
upsert	Yes	Updates existing keys, inserts new ones; only the listed columns are touched

---

Wildcard support

columns accepts wildcard patterns via the existing _wildcard_expansion helper, so patterns like * or Mapping* work as expected:

columns:
  - Key
  - Mapping*   # matches Mapping, MappingDetail, etc.

---

Example: upsert preserving unspecified columns

Model state before upsert — columns: Key, Schema, Mapping, Mandatory

update_df = pd.DataFrame({
    'Key':     ['apple'],
    'Mapping': ['green'],
    'Weight':  [1.0],    # new column not yet in the model
})

wrangles.connectors.train.lookup.write(
    update_df,
    model_id='<model-id>',
    columns=['Key', 'Mapping', 'Weight'],
    action='upsert',
    variant='key',
)

Model state after upsert:

Key	Schema	Mapping	Mandatory	Weight
apple	fruit	green	true	1.0
banana	fruit	yellow	false
cherry	fruit	red	true

• Schema and Mandatory on the apple row are preserved (not in columns, not touched).
• Weight is a new column — existing rows get an empty string.

[thomasstvr]

@mborodii-prog is there a series of commits that did not get pushed or am I missing something? I can't find anything about columns being required, even the tests are passing without columns (which none seem to have).

[mborodii-prog]

@thomasstvr It's because it will be a breaking change and on one of 1:1 call with Eric we decided to postpone that change

[mborodii-prog]

@ebhills @thomasstvr you can use recipe 'test lookup 922' for testing in QA

[ebhills]

Queue triage (2026-07-27)

• Disposition: Draft — delivery-owner action
• Delivery owner: @mborodii-prog
• Next action: This is the canonical PR for the branch (duplicate 992 refactor lookup actions bug request key colm #1038 is closed). Resolve the open thread, rerun the cancelled checks, and re-request one reviewer for v1.20.

GitHub is the status record; update this PR rather than the external spreadsheet.