Warn at build time when an expr pickles a module by reference that isn't packaged into the project wheel

[hussainsultan]

Problem

When a build script imports a helper module and puts one of its functions into an expr (UDF/UDAF body, Flight UDXF exchanger, etc.), cloudpickle serializes it by reference: the pickle records the module name and the module must be importable again at load time. The project wheel bundled into a catalog entry is the vehicle for that code — but nothing checks that the referenced modules actually made it into the wheel.

The gap is invisible at build time because the project source tree is on sys.path (cwd), so the import succeeds even inside the packaged env. It only surfaces when someone re-runs the entry outside the checkout:

Error: No module named 'cms_geo_join'

Real-world case: a project whose pyproject packaged only a vendored package (pi_xorq_verifier), while the analysis module cms_geo_join.py was a flat module never declared in the wheel targets. Every entry built fine, cataloged fine, and is unrunnable everywhere else — both in a foreign venv and in the entry's own isolated env (uv tool run --with <wheel> --with-requirements <pins> installs exactly the incomplete wheel). Because the failure is a bare ModuleNotFoundError at run time, it reads like an environment problem and sends users chasing venv flags instead of the packaging fix.

Proposal: warn when a referenced module is outside the package

There is a single choke point: serialize_callable (python/xorq/ibis_yaml/common.py), which all yaml-side pickles go through (translate.py pickled_fn, ibis_yaml/udf.py pickle and the Flight make_server/make_connection), plus one extra dumps site at expr/udf.py for model payloads.

A. Warn at pickle time (smallest fix). After cloudpickle.dumps, walk the bytes with stdlib pickletools.genops collecting STACK_GLOBAL/GLOBAL module operands — this catches transitively referenced modules (in the real case cms_geo_join sat nested inside a curry/Compose chain, so inspecting only the top function's __module__ would miss it). Classify each module with importlib.util.find_spec:

• stdlib or site-packages → provided by the env / pinned requirements → fine
• unresolvable, or resolving from the source tree → warn:

module 'cms_geo_join' is pickled by reference but resolves from the project source tree, not an installed package — include it in the project wheel or the entry won't load outside this checkout.

A prototype opcode scan against the real broken entry flags exactly cms_geo_join as the one OUTSIDE module (everything else resolves to stdlib/site-packages: pyarrow, toolz.functoolz, xorq.flight.exchanger, …). Requires a small stack simulation of the opcode stream for exact STACK_GLOBAL operands. No unpickling / code execution involved. Known caveat: editable installs also resolve outside site-packages, so the warning fires for them too (arguably correct, since the built wheel is what travels).

B. Wheel-membership check at packaging. Record A's module set into build_metadata.json, then in WheelPackager/PackagedBuilder (or xorq catalog add) compare against the built wheel's RECORD top-levels and name names: "expr references cms_geo_join; wheel cms_geo-0.1.0 provides only pi_xorq_verifier."

C. Clean-env load check (airtight gate). After building, reinvoke a deserialize-only load via uv tool run --with <wheel> --with-requirements <pins> from a temp cwd so the source tree is NOT on sys.path. Any ModuleNotFoundError there is byte-for-byte the failure users hit at re-run time. PackagedBuilder already has the reinvoke machinery; warm uv cache makes this seconds.

D. Fix instead of warn (opt-in). cloudpickle >=3.1.1 (already pinned) has register_pickle_by_value(module) — auto-register modules that classify OUTSIDE so their code embeds by value, making entries self-contained with no pyproject changes. Better as a flag (xorq build --pickle-local-by-value) than a default: it bloats expr.yaml and by-value pickles carry no cross-version compatibility guarantees.

Recommendation

Land A now (self-contained, exactly "warn when a module is outside the package"), fold its recorded module set into B's message at catalog add, keep C as a later strictness knob.

🤖 Generated with Claude Code