Skip to content

fix(memory): avoid identifier columns in aggregation seed queries - #2358

Merged
goldmedal merged 3 commits into
Canner:mainfrom
ttw225:fix/memory-seed-exclude-identifier
Jun 11, 2026
Merged

goldmedal merged 3 commits into
Canner:mainfrom
ttw225:fix/memory-seed-exclude-identifier

Conversation

@ttw225

@ttw225 ttw225 commented Jun 9, 2026 •

Copy link
Copy Markdown
Contributor

Problem

wren memory index generates seed NL-SQL pairs from the MDL manifest. The
aggregation template picks the first numeric column and applies SUM(),
previously excluding only a string primaryKey. Identifier columns that are
numeric — foreign keys, join keys, composite-PK members — were treated as
measures, producing semantically meaningless seeds such as:

SELECT SUM(customer_id) FROM orders

These pollute the memory index and dilute wren memory recall quality.

Fix

In seed_queries.py, exclude identifier columns from the aggregation target.
A column is an identifier if it is any of:

  • a primary key — now supports both string and composite-list primaryKey
    (aligned with the schema change in feat(wren-core): support composite (multi-column) primary keys #2345).
  • a relationship key — parsed from relationships[].condition; both sides
    of the join (and every column of an AND-joined composite condition) are
    excluded. This catches FKs whose name does not end in _id (e.g. created_by).
  • id / *_id by name — a case-insensitive fallback for identifiers that
    are not declared via a relationship (covers upper-case warehouses like
    Snowflake/Oracle).

Scope is intentionally narrow:

  • Only the aggregation target changes; group-column selection is untouched
    to avoid widening behavior.
  • The *_id heuristic also drops a would-be metric literally named *_id
    (e.g. household_id) — a deliberate, documented trade-off; real metrics
    should avoid the _id suffix.

No new MDL fields required — all signals already exist in the manifest.

Smoke test

cd core/wren && uv run python - <<'PY'
from wren.memory.seed_queries import generate_seed_queries
manifest = {
    "models": [
        {"name": "orders", "primaryKey": "order_id", "columns": [
            {"name": "order_id", "type": "int"}, {"name": "customer_id", "type": "int"},
            {"name": "amount", "type": "double"}, {"name": "status", "type": "varchar"}]},
        {"name": "customers", "primaryKey": "customer_id", "columns": [
            {"name": "customer_id", "type": "int"}, {"name": "name", "type": "varchar"}]},
        {"name": "store_sales", "primaryKey": ["ss_item_sk", "ss_ticket_number"], "columns": [
            {"name": "ss_item_sk", "type": "int"}, {"name": "ss_ticket_number", "type": "int"},
            {"name": "sales_price", "type": "decimal"}, {"name": "status", "type": "varchar"}]},
    ],
    "relationships": [{"name": "orders_customers", "models": ["orders", "customers"],
                       "condition": "orders.customer_id = customers.customer_id"}],
}
sqls = [p["sql"] for p in generate_seed_queries(manifest)]
print("\n".join(sqls))
noise = [c for c in ("SUM(customer_id)", "SUM(order_id)", "SUM(ss_item_sk)", "SUM(ss_ticket_number)")
         if any(c in s for s in sqls)]
assert not noise, f"identifier noise leaked: {noise}"
assert "SELECT SUM(amount) FROM orders" in sqls
assert "SELECT SUM(sales_price) FROM store_sales" in sqls
print("\nsmoke test passed — no identifier SUM noise; real metrics preserved")
PY

Expected: prints the seeds, ends with smoke test passed.
On main it raises AssertionError: identifier noise leaked: ['SUM(customer_id)', ...].

Validation

cd core/wren
uv run pytest tests/unit/test_seed_queries.py                                        # 32 passed
uv run pytest tests/unit/test_memory.py tests/unit/test_seed_queries.py              # 104 passed
uv run ruff check  src/wren/memory/seed_queries.py tests/unit/test_seed_queries.py   # clean
uv run ruff format --check src/wren/memory/seed_queries.py tests/unit/test_seed_queries.py  # clean

Known limitations (out of scope)

  • Group-column selection is unchanged, so a string FK can still appear in a
    grouped seed (amount by customer_uuid). Left for a follow-up.
  • Relationship-condition parsing assumes two-part model.column refs (what
    Wren emits); three-part schema-qualified refs are not parsed.

Summary by CodeRabbit

  • Bug Fixes

    • Seed generation now detects identifier-like numeric columns (primary keys, relationship join keys, and id-named fields, including composite and quoted cases) and excludes them from aggregation templates to avoid summing IDs.
  • Tests

    • Added comprehensive tests covering composite PKs, relationship-derived join keys, complex/quoted conditions, case-insensitivity, invalid conditions, and ensuring real numeric metrics still get aggregated.

@github-actions github-actions Bot added python Pull requests that update Python code core labels Jun 9, 2026
@coderabbitai

coderabbitai Bot commented Jun 9, 2026 •

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 0da7b58b-df42-45f1-9202-1cfafc4b5efa

📥 Commits

Reviewing files that changed from the base of the PR and between adf280b and 0fd5438.

📒 Files selected for processing (2)
  • core/wren/src/wren/memory/seed_queries.py
  • core/wren/tests/unit/test_seed_queries.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • core/wren/tests/unit/test_seed_queries.py

Walkthrough

This PR prevents SUM aggregation over identifier columns by parsing relationship.condition expressions for join keys, normalizing model primary keys (including composites), applying id-name heuristics, and excluding those columns when selecting numeric aggregation candidates during seed query generation.

Changes

Identifier-Aware Seed Query Generation

Layer / File(s) Summary
Identifier Classification Helpers
core/wren/src/wren/memory/seed_queries.py
Adds sqlglot imports and helper functions: _relationship_key_columns(manifest) parses relationship condition strings for join keys; _primary_key_columns(model) normalizes PKs including composite keys; _norm_ident and _is_id_like perform case-insensitive identifier normalization and id/*_id-style detection.
Integration into Seed Generation
core/wren/src/wren/memory/seed_queries.py
generate_seed_queries() precomputes relationship key sets and passes them into _model_seeds(model, relationship_keys). _model_seeds now computes primary_keys and excludes identifier-classified numeric columns (primary keys, relationship keys, or _is_id_like matches) from aggregation candidate selection and grouping.
Identifier Exclusion Test Suite
core/wren/tests/unit/test_seed_queries.py
Test helper _model accepts pk as string or list for composite keys. Adds TestIdentifierExclusion tests verifying exclusion of identifier-like numeric columns (relationship-derived keys, composite PK members, case-insensitive id/*_id heuristics, quoted/AND-composed relationship conditions) from SUM aggregations while genuine numeric metrics remain aggregated.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested labels: dependencies

Suggested reviewers:

  • douenergy

"I nibbled through the schema trees,
Found keys in joins and hiding slews,
No more summing IDs by mistake,
Seeds now know which fields to take,
A rabbit's hop toward tidier views." 🐇✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 22.73% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and specifically describes the main change: preventing identifier columns from being included in aggregation seed queries, which is the core purpose of this PR.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@core/wren/src/wren/memory/seed_queries.py`:
- Around line 45-46: The membership checks against relationship_keys are
case-sensitive and allow mismatches; update the code that builds and queries
relationship_keys (refer to the use of relationship_keys and the
_model_seeds(...) call, and any checks like model["name"] or column identifiers)
to normalize identifiers (e.g., .lower() or a consistent canonical form) both
when inserting keys into relationship_keys and when performing membership checks
inside _model_seeds and related functions so comparisons are case-insensitive
and consistent across model and relationship names.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 5648bf84-9e92-43b6-a4c7-d6336d3990d5

📥 Commits

Reviewing files that changed from the base of the PR and between d33917f and e1d6cd0.

📒 Files selected for processing (2)
  • core/wren/src/wren/memory/seed_queries.py
  • core/wren/tests/unit/test_seed_queries.py

Comment thread core/wren/src/wren/memory/seed_queries.py Outdated
Comment thread core/wren/src/wren/memory/seed_queries.py Outdated
@ttw225

ttw225 commented Jun 10, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the suggestion and example. I agree this is a better fit since the relationship condition is an SQL expression.

I’ve updated the implementation to use sqlglot to parse the condition and collect qualified columns from the AST, instead of relying on regex. I also added regression coverage for quoted identifiers and invalid relationship conditions.

@goldmedal

Copy link
Copy Markdown
Collaborator

Code review

No issues found. Checked for bugs and CLAUDE.md compliance.

🤖 Generated with Claude Code

- If this code review was useful, please react with 👍. Otherwise, react with 👎.

@goldmedal goldmedal left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @ttw225 Look great 👍

@goldmedal
goldmedal merged commit d5e7879 into Canner:main Jun 11, 2026
10 checks passed
@ttw225
ttw225 deleted the fix/memory-seed-exclude-identifier branch June 11, 2026 02:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core python Pull requests that update Python code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants