From c8fd3e97d1bb35ad704ea442def79bf9188e8cbf Mon Sep 17 00:00:00 2001
From: Yang Cen <bubble-cal@outlook.com>
Date: Mon, 31 Aug 2026 22:34:34 +0800
Subject: [PATCH] test(python): cover stable main udf registration identity
 (#4094)

## Other changes

### What changed?

- Add a subprocess regression harness for an ordinary `@udf` function
defined in `__main__`.
- Verify the full registration request, artifact digest, and Function
signature stay identical across independent Python processes and
renamed/moved script paths.
- Verify body, referenced-global, and annotation changes still produce
distinct artifact identities, with annotation changes also producing a
distinct Function signature.

### Why is the change needed?


[ENT-2441](https://linear.app/lancedb/issue/ENT-2441/make-sure-function-defined-in-main-gets-stable-signature)
tracks the stability guarantee. Investigation on the exact `840e1d73`
main base found that LanceDB already packages canonical source instead
of cloudpickle bytes, so the unchanged `__main__` function is stable and
no production-code fix is needed. This change closes the missing
regression-test coverage.


[GEN-950](https://linear.app/lancedb/issue/GEN-950/class-based-udfs-defined-in-main-get-a-new-auto-version-on-every-run)
remains a separate Geneva checkpoint-version issue for class-based
callables. LanceDB's Function API continues to accept synchronous Python
functions only.

## Validation

- `cd python && uv run --extra tests pytest
python/tests/test_first_class_function_slice2.py -q` (`40 passed`)
- `uv run --project python --extra dev ruff format .`
- `uv run --project python --extra dev ruff check .` (`All checks
passed!`)
---
 .../tests/test_first_class_function_slice2.py | 76 +++++++++++++++++++
 1 file changed, 76 insertions(+)

diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py
index bab78316c..57b08e18d 100644
--- a/python/python/tests/test_first_class_function_slice2.py
+++ b/python/python/tests/test_first_class_function_slice2.py
@@ -12,6 +12,8 @@ from datetime import date
 import http.server
 import json
 from pathlib import Path
+import subprocess
+import sys
 import threading
 from typing import Optional
 
@@ -67,6 +69,80 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
     }
 
 
+def _main_udf_source(
+    *, threshold: int = 20, input_annotation: str = "int", comparison: str = ">="
+) -> str:
+    return (
+        "from __future__ import annotations\n"
+        "from lancedb.functions import udf\n"
+        f"THRESHOLD = {threshold}\n"
+        "\n"
+        "@udf\n"
+        f"def label(value: {input_annotation}) -> str:\n"
+        f"    return 'big' if value {comparison} THRESHOLD else 'small'\n"
+        "\n"
+        "assert label.__module__ == '__main__'\n"
+        "print(label.registration_request.to_canonical_json())\n"
+    )
+
+
+def _run_main_udf(path: Path, source: str) -> dict:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(source)
+    result = subprocess.run(
+        [sys.executable, str(path)],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return json.loads(result.stdout)
+
+
+def test_main_udf_registration_identity_is_stable_across_processes_and_paths(
+    tmp_path,
+):
+    source = _main_udf_source()
+    original_path = tmp_path / "original" / "job.py"
+    moved_path = tmp_path / "moved" / "renamed_job.py"
+
+    original_runs = [_run_main_udf(original_path, source) for _ in range(2)]
+    moved_run = _run_main_udf(moved_path, source)
+
+    assert len({run["artifact"]["digest"] for run in [*original_runs, moved_run]}) == 1
+    assert all(
+        run["signature"] == original_runs[0]["signature"]
+        for run in [original_runs[1], moved_run]
+    )
+    assert original_runs[0] == original_runs[1] == moved_run
+
+    body_change = _run_main_udf(
+        tmp_path / "changes" / "body.py", _main_udf_source(comparison=">")
+    )
+    global_change = _run_main_udf(
+        tmp_path / "changes" / "global.py", _main_udf_source(threshold=21)
+    )
+    annotation_change = _run_main_udf(
+        tmp_path / "changes" / "annotation.py",
+        _main_udf_source(input_annotation="float"),
+    )
+
+    baseline = original_runs[0]
+    assert baseline["signature"] == body_change["signature"]
+    assert baseline["signature"] == global_change["signature"]
+    assert baseline["signature"] != annotation_change["signature"]
+    assert (
+        len(
+            {
+                baseline["artifact"]["digest"],
+                body_change["artifact"]["digest"],
+                global_change["artifact"]["digest"],
+                annotation_change["artifact"]["digest"],
+            }
+        )
+        == 4
+    )
+
+
 def _run_packaged(definition, *args):
     """Execute the shipped artifact in a fresh namespace, as a worker would."""
     source = base64.b64decode(definition.registration_request.artifact.content.data)