refactor: de-duplicate small helpers in the DataFrame/Expr API

[nielspardon]

Follow-up cleanup from the review of #204 (DRY, no behavior change). Four small helpers in the substrait.dataframe / substrait.narwhals layer re-implemented logic that already existed elsewhere; this consolidates each onto a single source of truth so the copies cannot drift apart.

Changes

• Extension merge — replace dataframe.expr._merge_extensions_into (which deduped on strict proto equality) with a shared utils.merge_extensions_into that reuses merge_extension_urns / merge_extension_declarations, keying on the same anchor/name identity as builders.plan._merge_extensions. Used by Expr.order_by and Expr.over.
• Sort direction — drop the duplicate (descending, nulls_last) -> SortDirection mapping. A single _SORT_DIRECTIONS + sort_direction() now lives in dataframe.expr and is used by both dataframe.expr and dataframe.frame.
• Multi-URN lookup — add ExtensionRegistry.find_function(name, signature, urns=None), a generalization of lookup_function (one URN) and list_functions_across_urns (all URNs) over an ordered URN subset. _resolve_over_urns now resolves in one registry call and recovers the winning extension via FunctionEntry.urn instead of looping lookup_function per URN.
• Alias — extract the referred_expr[0].output_names[0] rewrite into a shared builders.extended_expression.alias, used by both Expr.alias and the Narwhals expression wrapper.

Also replaces the empty Exception("") in merge_extension_declarations with an informative NotImplementedError naming the unsupported mapping type.

Notes

The extension-merge and multi-URN changes are behavior-preserving: extension anchors are assigned deterministically per URN/overload and declaration names embed the full overload signature, so the identity-based dedup and the previous proto-equality dedup agree on all real inputs, and find_function selects the same winning URN as the previous per-URN loop.

Testing

• pytest: 528 passed, 30 skipped
• ruff check and ruff format --check: clean
• New unit tests for merge_extensions_into (identity dedup, multiple sources) and find_function (single-URN parity with lookup_function, candidate-order priority, no-match, all-URNs).

Closes #208

🤖 Generated with AI

[andrew-coleman]

PR review by Claude:

CRITICAL FINDINGS

1. ❌ EXTENSION URN DEDUP: Proto Equality → String-Only Comparison (BEHAVIOR CHANGE)

Severity: CRITICAL
Category: Correctness
Files: src/substrait/utils/__init__.py:merge_extension_urns()

The Issue

The PR consolidates extension merging from two places:

Old code (expr.py:_merge_extensions_into):

for urn in source.extension_urns:
    if urn not in target.extension_urns:  # Uses proto equality (all fields)
        target.extension_urns.append(urn)

New code (utils/__init__.py:merge_extension_urns):

seen_urns = set()
for urn in urns:
    if urn.urn not in seen_urns:  # Only checks URN string!
        seen_urns.add(urn.urn)
        ret.append(urn)

Why This Breaks

Proto equality checks all fields: extension_urn_anchor AND urn string.
String dedup only checks the urn field, ignoring extension_urn_anchor.

Concrete Attack Scenario:

target.extension_urns = [SimpleExtensionURN(anchor=1, urn="extension:io.substrait:functions_arithmetic")]
source.extension_urns = [SimpleExtensionURN(anchor=2, urn="extension:io.substrait:functions_arithmetic")]

# Old behavior: Add source (different protos due to different anchors)
# New behavior: Skip source (same urn string)

Evidence

Verified with Python test (see verification section) — proto equality and string equality diverge.

PR's Counterargument

The PR states: "extension anchors are assigned deterministically per URN/overload... the identity-based dedup and the previous proto-equality dedup agree on all real inputs"

Rebuttal:

• No evidence provided that anchors are deterministically assigned
• No invariant enforced in code that anchor ≠ anchor for same URN is impossible
• The behavior still changes, even if "real inputs" never trigger it

Impact

• Low probability: Anchors may indeed be deterministic in practice
• High consequence: Silent data loss (extensions dropped) if anchors differ
• Verdict: Requires explicit confirmation or code change

Recommendation:

• Add a test that verifies anchor mismatch is impossible, OR
• Revert to proto equality check with explicit comment explaining why, OR
• Document this assumption clearly

---

2. ❌ PROTOBUF REFERENCE INVALIDATION via ClearField (MEMORY SAFETY)

Severity: CRITICAL
Category: Memory Safety / Reference Semantics
Files: src/substrait/utils/__init__.py:merge_extensions_into() lines 338-341

The Issue

def merge_extensions_into(target, *sources):
    merged_urns = merge_extension_urns(...)
    merged_extensions = merge_extension_declarations(...)
    target.ClearField("extension_urns")      # ← Replaces list object
    target.extension_urns.extend(merged_urns)  # ← Different list object
    target.ClearField("extensions")          # ← Replaces list object
    target.extensions.extend(merged_extensions)  # ← Different list object

Why This Is Unsafe

In Python protobuf, ClearField() on a repeated field replaces the underlying list object. External references become stale:

bound = <ExtendedExpression>
urns_ref = bound.extension_urns  # Reference to list object A

merge_extensions_into(bound, other_bound)
# Inside merge_extensions_into: ClearField replaces bound.extension_urns with list object B

# Now bound.extension_urns points to object B, but urns_ref still points to object A
for urn in urns_ref:  # Iterates over STALE data
    ...

Verification

Python mock test shows:

• Old code (in-place append): Same list object, no reference invalidation ✓
• New code (ClearField + extend): Different list object, references invalidated ✗

Current Risk Assessment

Checked dataframe/expr.py for external references to extension_urns — none found in immediate callers. However:

1. Futures risk: Code could be added in Narwhals or other layers that holds references
2. Fragility: Making the implementation fragile to future changes
3. Semantics: Violates expectation that mutating target in-place keeps it "in-place"

Recommendation

Change merge_extensions_into to extend in-place instead of ClearField:

def merge_extensions_into(target, *sources):
    merged_urns = merge_extension_urns(
        target.extension_urns, *(s.extension_urns for s in sources)
    )
    merged_extensions = merge_extension_declarations(
        target.extensions, *(s.extensions for s in sources)
    )
    # Remove old entries that will be replaced
    del target.extension_urns[:]  # Clear in-place (same object)
    del target.extensions[:]      # Clear in-place (same object)
    # Extend with new values (same objects)
    target.extension_urns.extend(merged_urns)
    target.extensions.extend(merged_extensions)

---

3. ❌ SORT_DIRECTION: Silent Type Coercion Masks Invalid Inputs

Severity: HIGH
Category: Robustness / Input Validation
Files: src/substrait/dataframe/expr.py:119

The Issue

def sort_direction(descending: bool, nulls_last: bool):
    """The ``SortField.SortDirection`` for a ``(descending, nulls_last)`` pair."""
    return _SORT_DIRECTIONS[(bool(descending), bool(nulls_last))]

Explicit bool() coercion: Converts any input to boolean.

Why This Is Problematic

Python truthiness semantics are permissive:

• bool("false") → True (non-empty string is truthy!)
• bool("") → False (empty string is falsy)
• bool(None) → False (None is falsy)
• bool(0) → False (zero is falsy)

Concrete Bug Scenario

expr.sort(..., descending="false")  # Intends False, but gets True!
# Result: DESC_NULLS_FIRST (wrong!)

expr.sort(..., descending=None)  # Intends error, but silently coerces to False
# Result: ASC_NULLS_FIRST (unexpected!)

Verified with Test

Python test output:

sort_direction('true' , 'false') -> DESC_NULLS_LAST  # Should error!
sort_direction(None   , '')     -> ASC_NULLS_FIRST   # Should error!

Test Gap

The new test file (test_registry_lookup.py) only tests find_function(), not sort_direction(). No tests exist for invalid inputs.

Recommendation

Either:

1. Add explicit validation: Check isinstance(descending, bool) and raise TypeError if not
2. Document behavior: Add docstring saying "truthy/falsy values accepted" and add tests
3. Revert to if/else: Use the original pattern which is less deceptive about coercion

---

HIGH FINDINGS

4. ⚠️ EXTENSION DECLARATION DEDUP: Identity vs Proto Equality (DOCUMENTED BUT CONCERNING)

Severity: MEDIUM (documented as intentional, but still a change)
Category: Correctness
Files: src/substrait/utils/__init__.py:merge_extension_declarations()

The Issue

Similar to Finding #1, but for declarations:

# Old code: proto equality check
for decl in source.extensions:
    if decl not in target.extensions:
        target.extensions.append(decl)

# New code: identity-based (urn_ref, name) tuple
ident = (ext_func.extension_urn_reference, ext_func.name)
if ident not in seen_extension_functions:
    seen_extension_functions.add(ident)
    ret.append(declaration)

Evidence

The PR's own test (tests/test_utils.py:test_merge_extensions_into_appends_new_and_dedupes_on_identity()) explicitly documents this:

# Same (urn reference, name) as target's -- a duplicate by identity even
# though the function anchor differs, so it must not be appended (this is
# what distinguishes identity dedup from strict proto equality).
_extension_function(1, 99, "f:i8"),  # ← Different anchor (99 vs 10)

Mitigation Status

Unlike Finding #1, this is explicitly tested and the behavior is documented. However:

1. It still violates "behavior-preserving" claim (it's documented but changed)
2. The test only has function_anchor differ; what if other fields differ?
3. No guarantee this dedup strategy is correct for all fields

Recommendation

• Accept this if anchors are proven deterministic
• Add comment explaining why other declaration fields are irrelevant
• Consider whether function_anchor should matter (it shouldn't, but verify)

---

5. ⚠️ ALIAS MUTATION OF SHARED PROTO STATE

Severity: MEDIUM
Category: Correctness / State Mutation
Files: src/substrait/builders/extended_expression.py:20-26

The Issue

def alias(expression: ExtendedExpressionOrUnbound, name: str) -> UnboundExtendedExpression:
    def resolve(base_schema, registry) -> stee.ExtendedExpression:
        bound_expression = resolve_expression(expression, base_schema, registry)
        bound_expression.referred_expr[0].output_names[0] = name  # ← MUTATES
        return bound_expression
    return resolve

The function mutates bound_expression in-place. If resolve_expression() returns a shared proto object (not a copy), this breaks:

expr1 = Expr(...).alias("col1")  # Calls alias(expr, "col1")
expr2 = Expr(...).alias("col2")  # Calls alias(expr, "col2")

# If both resolve() calls get the SAME proto object:
# Both expr1 and expr2 end up with output_name = "col2" (last write wins)

Root Cause

resolve_expression() at line 53:

def resolve_expression(expression, base_schema, registry):
    return (
        expression
        if isinstance(expression, stee.ExtendedExpression)
        else expression(base_schema, registry)
    )

If expression is already a bound ExtendedExpression, it's returned as-is (no copy). Subsequent mutations will affect the shared object.

Likelihood

Depends on calling patterns:

• If Expr.alias() always gets a fresh unbound expression, safe ✓
• If unbound expressions cache results, unsafe ✗
• If pre-bound expressions are passed to alias(), unsafe ✗

Investigation

The old code did the same thing, so this is not introduced by the PR. But the consolidation makes it more likely to be exposed if the calling pattern changes.

Recommendation

• Add documentation: "alias() assumes resolve_expression() returns a fresh proto"
• Consider adding .CopyFrom() in resolve() if safety is needed: bound_expression = resolve_expression(...).CopyFrom(...)
• Actually, that doesn't work. Better: bound = resolve_expression(...); result = stee.ExtendedExpression(); result.CopyFrom(bound); result.referred_expr[0].output_names[0] = name

---

6. ⚠️ FIND_FUNCTION URN RECOVERY: Fragile Assumption

Severity: MEDIUM
Category: Robustness / Implicit Assumptions
Files: src/substrait/dataframe/expr.py:74-79

The Issue

match = registry.find_function(name, signature, urns)
if match is not None:
    winning_urn = match[0].urn  # ← Assumes match[0].urn is in urns list

The Risk

FunctionEntry.urn is set at object creation time. If:

1. The registry is mutated between creation and lookup
2. Function entries are shared/cached incorrectly
3. URN is stale for any reason

Then winning_urn could be incorrect or out of sync with the actual winning URN in the registry.

Current Code Assessment

Looking at registry.py, find_function() calls _find_matching_functions() which iterates over urns_to_search and collects entries:

for urn in urns_to_search:
    if urn not in self._function_mapping:
        continue
    functions = self._function_mapping[urn][function_name]
    for f in functions:
        rtn = f.satisfies_signature(signature)
        if rtn is not None:
            matches.append((f, rtn))  # ← f.urn should be from current urn
return matches[0] if matches else None

Since f comes from self._function_mapping[urn][function_name], and f.urn is set at creation, the assumption appears safe. However:

Why It's Still Fragile

1. Implicit contract: Code relies on f.urn being from the urn that contains it in the mapping
2. No validation: Nothing checks that f.urn in urns
3. No documentation: Doesn't explicitly state winning URN is recovered from entry

Better Approach

Return the winning URN explicitly:

def find_function(...):
    matches = self._find_matching_functions(function_name, signature, urns)
    if matches:
        entry, rtn_type = matches[0]
        return entry, rtn_type, <winning_urn_explicitly>  # ← Return URN
    return None

Or modify find_function() to return a triple: (entry, output_type, winning_urn)

Recommendation

• Add assertion: assert winning_urn in urns after line 76
• Or modify find_function() signature to return winning_urn explicitly

---

7. ⚠️ MISSING NARWHALS INTEGRATION TESTS

Severity: MEDIUM
Category: Test Coverage
Files: Tests (notably narwhals/expression.py changes)

The Issue

The PR refactors narwhals/expression.py to use the shared alias() builder:

from substrait.builders.extended_expression import (
    alias as _alias,
)

But the new tests only verify:

• find_function() with registry lookup tests
• merge_extensions_into() with extension merge tests

No tests verify that Narwhals expressions still work correctly with the shared builder.

Test Gap

The PR runs "528 passed, 30 skipped" but doesn't report breakdown:

• How many are Narwhals tests?
• Did Narwhals test count stay the same?
• Are there integration tests between Expr and Narwhals?

Recommendation

• Add explicit Narwhals-layer test: expr.alias("new_name") produces correct output_name
• Add integration test combining Expr and Narwhals with aliasing
• Report test breakdown by module in PR

---

MEDIUM FINDINGS

8. ✅ ERROR MESSAGE IMPROVEMENT (POSITIVE, but adds feature)

Severity: LOW (improvement)
Category: Documentation / Error Handling
Files: src/substrait/utils/__init__.py:306-312

The Change

Before:

else:
    raise Exception("")  # TODO handle extension types

After:

else:
    mapping_type = declaration.WhichOneof("mapping_type")
    raise NotImplementedError(
        f"cannot merge extension declaration of type {mapping_type!r}; "
        f"only 'extension_function' declarations are supported so far"
    )

Why It's Positive

• Provides context (actual mapping type discovered)
• Changes exception type to NotImplementedError (more accurate)
• Adds informative message

Why It's Concerning

• Calls WhichOneof() which is new code execution
• If error path was rarely hit, adding logic here could introduce new bugs
• No test verifies this error path still works

Recommendation

• Add test that triggers this error path (with non-function declaration)
• Verify NotImplementedError is correct vs TypeError or other

---

SUMMARY TABLE

ID	Finding	Severity	Category	Testable	Status
1	Extension URN dedup: proto equality → string	CRITICAL	Correctness	Yes	Needs verification
2	ClearField reference invalidation	CRITICAL	Memory Safety	Yes	Needs fix
3	sort_direction silent coercion	HIGH	Robustness	Yes	Needs tests/fix
4	Extension declaration dedup identity	MEDIUM	Correctness	Yes	Documented but concerning
5	Alias mutation of shared state	MEDIUM	Correctness	Hard	Inherited from old code
6	find_function URN recovery fragile	MEDIUM	Robustness	Yes	Needs assertion/refactor
7	Narwhals integration tests missing	MEDIUM	Test Coverage	Yes	Needs new tests
8	Error message improvement	LOW	Documentation	Yes	Minor enhancement

---

VERIFICATION CHECKLIST

• Proto equality vs string dedup behavior differs (verified with Python test)
• ClearField invalidates references (verified with mock protobuf test)
• sort_direction coerces silently (verified: "false" → True)
• Extension anchors are deterministically assigned (needs investigation)
• Narwhals tests still pass with shared alias (needs run)
• No code holds references to extension_urns (grep shows none, but fragile)
• resolve_expression always returns fresh protos (needs verification)
• find_function always returns correct URN (needs assertion)

---

RECOMMENDATIONS PRIORITY

MUST FIX (blocking approval)

1. ClearField reference invalidation — Change to in-place clear (safety issue)
2. Verify URN dedup assumption — Add test or document why anchors can't differ

SHOULD FIX (before merge)

3. sort_direction validation — Add type checking or tests for coercion
4. Narwhals integration test — Verify alias() works across layers

NICE TO HAVE (follow-up)

5. find_function URN assertion — Add safety check
6. Error message path test — Ensure new error handling works

---

CONCLUSION

The PR consolidates duplicate code with good intent, but introduces two critical issues:

1. Behavior changes that violate "behavior-preserving" claim (URN dedup by string instead of proto equality)
2. Memory safety issue with ClearField invalidating external references

The first is likely safe in practice (anchors probably deterministic), but undocumented and unverified. The second is a real fragility that should be fixed before merge.

Recommended Action: Request fixes for findings #1, #2, and #3 before approval. Findings #4-7 should be addressed in follow-up or documented.

[nielspardon]

Thanks for the review. I went through each finding against the actual code and verified the behavioral claims empirically (real protobuf upb runtime + the default registry). Summary: findings 1–6 don't hold up, 7 is already covered, and 8 was a genuine (minor) gap that I've now fixed in 225448c.

Finding 1 — URN dedup by string vs proto equality — not a behavior change

extension_urn_anchor is assigned as registry.lookup_urn(urn), which is deterministic per registry: the same URN string always maps to the same anchor. So two SimpleExtensionURN with the same urn string can never carry different anchors on real inputs, and the identity (string) dedup and the previous proto-equality dedup agree. This is also not new behavior — merge_extension_urns is already what builders.plan._merge_extensions and the scalar_function builder itself use; the PR just makes the DataFrame/Expr layer consistent with them.

Finding 2 — ClearField reference invalidation — not a bug

Tested against the upb implementation: ClearField + extend produces the correct merged list, and the objects referenced by the merged list stay valid (protobuf detaches them). No caller holds an external reference to target.extension_urns across the call, so there's no invalidation in practice.

Finding 3 — sort_direction bool() coercion — not a regression

The pre-PR expr._sort_direction used truthiness (if descending:), so bool(descending) is behavior-preserving. The parameters are typed bool with bool defaults; passing "false" was already truthy before this PR.

Finding 4 — declaration dedup by (urn_ref, name) — no collision on real inputs

The declaration name is str(func[0]), which embeds the full overload signature (e.g. add:i64_i64 vs add:i32_i32). Different overloads therefore have different names, and identical (urn_ref, name) implies the same overload ⇒ same function_anchor. The test with a differing anchor exists only to prove the dedup keys on identity rather than anchor; it can't arise from the builder.

Finding 5 — alias mutation of shared proto — inherited, and safe here

This is unchanged from the old code. Expr.alias/the Narwhals wrapper always wrap an unbound callable, which builds a fresh ExtendedExpression on each resolve, so there's no shared proto to clobber.

Finding 6 — find_function URN recovery — correct by construction

Entries live in _function_mapping[urn][name] and each FunctionEntry.urn is set to that same urn at registration, so match[0].urn is exactly the URN the match was found under. Verified.

Finding 7 — Narwhals alias coverage — already exercised

narwhals.DataFrame.select(name=expr) routes through Expression.alias → the shared builder, and tests/narwhals/test_df_project.py::test_project asserts the resulting names=["id"].

Finding 8 — untested NotImplementedError branch — fixed

Good catch — the non-function-declaration error path had no test. Added test_merge_extension_declarations_rejects_non_function_mapping in 225448c asserting the error names the unsupported mapping type.

Full suite (550 passed, 30 skipped), ruff check, and ruff format --check all clean.