Skip to content

Capital extract labels #1127

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ebhills merged 7 commits into from
Aug 14, 2026
+65 −15
Merged

Capital extract labels #1127

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
2baef68
Added fix and test
thomasstvr Aug 12, 2026
e00584a
Updated test to include a row with no matches and check the entire di…
thomasstvr Aug 12, 2026
332413f
Updated to handle capitalized labels with no match
thomasstvr Aug 12, 2026
bc4bd15
Updated test and added a second
thomasstvr Aug 12, 2026
dbbc538
Added capital, and dropped or since sort order should not change
thomasstvr Aug 12, 2026
c79c5e8
Updated the test to include CAPITAL and fixed order
thomasstvr Aug 12, 2026
8556ccd
Potential fix for pull request finding
thomasstvr Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 58 additions & 11 deletions tests/recipes/wrangles/test_extract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1157,6 +1157,52 @@ def test_extract_custom_labels(self):
df['col2'][0]['size'] == ['small']
)

def test_extract_custom_preserves_capitalized_label(self):
"""
Test use_labels option to group output and preserve capitalized label
"""
df = wrangles.recipe.run(
"""
wrangles:
- extract.custom:
input: col1
output: col2
model_id: 829c1a73-1bfd-4ac0
use_labels: true
""",
dataframe = pd.DataFrame({
'col1': ['Washington D.C. is the capital of the United States', 'Austin is the capital of Texas', 'No Matches Here']
})
)
assert (
df['col2'][0] == {'CAPITAL': ['Washington D.C.'], 'size': [], 'colour': []} and
df['col2'][1] == {'CAPITAL': ['Austin'], 'size': [], 'colour': []} and
df['col2'][2] == {'CAPITAL': [], 'size': [], 'colour': []}
)

def test_extract_custom_preserves_capitalized_label_multi_matches(self):
"""
Test use_labels option to group output and preserve capitalized label
"""
df = wrangles.recipe.run(
"""
wrangles:
- extract.custom:
input: col1
output: col2
model_id: 829c1a73-1bfd-4ac0
use_labels: true
""",
dataframe = pd.DataFrame({
'col1': ['Washington D.C. green small', 'Austin black medium', 'No Matches Here']
})
)
assert (
df['col2'][0] == {'CAPITAL': ['Washington D.C.'], 'colour': ['green'], 'size': ['small']} and
df['col2'][1] == {'CAPITAL': ['Austin'], 'colour': ['black'], 'size': ['medium']} and
df['col2'][2] == {'CAPITAL': [], 'size': [], 'colour': []}
)

def test_extract_custom_labels_columns_format(self):
"""
Test use_labels option with output_format: columns to expand labels into dataframe columns
Expand Down Expand Up @@ -1826,7 +1872,7 @@ def test_use_labels_multiple(self):
Testing use labels with multiple same labels and other labels
"""
data = pd.DataFrame({
'col': ['colour: blue size: small colour: black']
'col': ['colour: blue size: small colour: black Washington D.C. Austin']
})
recipe = """
wrangles:
Expand All @@ -1838,7 +1884,7 @@ def test_use_labels_multiple(self):
first_element: false
"""
df = wrangles.recipe.run(recipe, dataframe=data)
assert df['out'][0] == {'size': ['small'], 'colour': ['blue', 'black']} or df['out'][0] == {'colour': ['black', 'blue'], 'size': ['small']}
assert df['out'][0] == {'CAPITAL': ['Washington D.C.', 'Austin'], 'colour': ['blue', 'black'], 'size': ['small']}
Comment thread
thomasstvr marked this conversation as resolved.

def test_use_labels_same_key(self):
"""
Expand Down Expand Up @@ -2102,18 +2148,18 @@ def test_extract_custom_sort_case_sensitive(self, sort_type, expected):


@pytest.mark.parametrize("sort_type, expected", [
("training_order", {'colour': ['blue', 'black'], 'size': ['small', 'medium']}),
("input_order", {'colour': ['blue', 'black'], 'size': ['medium', 'small']}),
("longest", {'colour': ['black', 'blue'], 'size': ['medium', 'small']}),
("shortest", {'colour': ['blue', 'black'], 'size': ['small', 'medium']}),
("alphabetical", {'colour': ['black', 'blue'], 'size': ['medium', 'small']}),
("reverse_alphabetical",{'colour': ['blue', 'black'], 'size': ['small', 'medium']}),
("descending", {'colour': ['black', 'blue'], 'size': ['medium', 'small']}),
("ascending", {'colour': ['blue', 'black'], 'size': ['small', 'medium']})
("training_order", {'CAPITAL': ['Washington D.C.', 'Austin'], 'colour': ['blue', 'black'], 'size': ['small', 'medium']}),
("input_order", {'CAPITAL': ['Austin', 'Washington D.C.'], 'colour': ['blue', 'black'], 'size': ['medium', 'small']}),
("longest", {'CAPITAL': ['Washington D.C.', 'Austin'], 'colour': ['black', 'blue'], 'size': ['medium', 'small']}),
("shortest", {'CAPITAL': ['Austin', 'Washington D.C.'], 'colour': ['blue', 'black'], 'size': ['small', 'medium']}),
("alphabetical", {'CAPITAL': ['Austin', 'Washington D.C.'], 'colour': ['black', 'blue'], 'size': ['medium', 'small']}),
("reverse_alphabetical",{'CAPITAL': ['Washington D.C.', 'Austin'], 'colour': ['blue', 'black'], 'size': ['small', 'medium']}),
("descending", {'CAPITAL': ['Washington D.C.', 'Austin'], 'colour': ['black', 'blue'], 'size': ['medium', 'small']}),
("ascending", {'CAPITAL': ['Austin', 'Washington D.C.'], 'colour': ['blue', 'black'], 'size': ['small', 'medium']})
])
def test_extract_custom_sort_use_labels(self, sort_type, expected):
df = pd.DataFrame({
"input_col": ['medium blue black small']
"input_col": ['medium blue black small austin washington d.c.']
})
recipe = f"""
wrangles:
Expand All @@ -2122,6 +2168,7 @@ def test_extract_custom_sort_use_labels(self, sort_type, expected):
output: result
sort: {sort_type}
use_labels: True
# include_empty_labels: False
model_id: 829c1a73-1bfd-4ac0
"""
result = wrangles.recipe.run(recipe, dataframe=df)
Expand Down
11 changes: 7 additions & 4 deletions wrangles/extract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,15 +681,18 @@ def _entities_to_list(value):
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 [])
all_labels = {}
for label in sorted(model_labels, key=lambda x: x.lower()):
all_labels.setdefault(label.lower(), label)
for objs in results:
all_labels.update([str(k).lower() for k in objs.keys()])
for label in objs:
all_labels.setdefault(str(label).lower(), label)

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:
for normalized_label, label in all_labels.items():
if normalized_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]
Expand Down