perf: Reuse compiled numba kernels across compile_components cache misses

[hoxbro]

Description

compile_components is memoized on the full schema and glyph, so a frame with different unused columns, dtypes or column names misses the cache and builds new functions, even when they are identical to ones already built. The glyphs memoize their kernels (extend_cpu, draw_segment, _bresenham, ...) on these functions, so the whole kernel chain is compiled again. The compiled code is never freed, so memory grows over the test suite.

This reuses the functions when their inputs are the same:

• The jitted append, which make_append builds from a source string with exec, is cached on its generated code and the globals it is executed with.
• make_info and make_antialias_stage_2_functions are memoized. They returned a new info closure and new antialias stage 2 functions on every cache miss, and the glyphs' memoized _build_extend takes both as arguments.

This came up while investigating the ubuntu unit jobs on #1539 being killed after hanging, where each of the 4 xdist workers ends at ~4.4 GB RSS (~17.5 GB total on a 16 GB runner). On main, the workers already end at ~3.5–3.7 GB each.

append cache, measured on the #1539 branch (which also passes columns to numba as read-only arrays), running the masked/dict/arrow tests in one process:

	RSS	numba dispatchers	compiled overloads	append functions built (unique)
before	1393 MB	1110	708	141 (44)
after	1308 MB	1013	620	44

make_info and make_antialias_stage_2_functions memoization, measured on this branch with test_pandas.py in one process:

	RSS	numba dispatchers	compiled overloads
before	3515 MB	2719	1959
after	3117 MB	2165	1696

The remaining duplicates mostly come from glyph kernels memoized on the glyph itself, so different column names still compile them again. That is left for a follow-up.

Scripts

Counts append functions built vs. unique code + globals (-p append_plugin):

import collections

from datashader import compiler

_keys = collections.Counter()
_orig_ngjit = compiler.ngjit


def _recording_ngjit(func):
    if func.__name__ == "append":
        ns = func.__globals__
        refs = tuple(sorted((k, v) for k, v in ns.items() if k in func.__code__.co_names))
        _keys[(func.__code__.co_code, func.__code__.co_consts, func.__code__.co_names, refs)] += 1
    return _orig_ngjit(func)


compiler.ngjit = _recording_ngjit


def pytest_sessionfinish(session):
    total = sum(_keys.values())
    print(f"\nappend functions built={total} unique code+namespace={len(_keys)}")

Records RSS, numba dispatchers and compiled overloads after each test (-p jit_plugin):

import gc
import os

import psutil
import pytest
from numba.core.dispatcher import Dispatcher

_proc = psutil.Process()
_out = open(f".scraith/ci-timeout/jit-{os.environ.get('JIT_TAG', 'run')}.tsv", "w")


def _jit_state():
    ds = [o for o in gc.get_objects() if isinstance(o, Dispatcher)]
    return len(ds), sum(len(d.overloads) for d in ds)


@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_protocol(item, nextitem):
    yield
    gc.collect()
    n_disp, n_over = _jit_state()
    _out.write(f"{item.nodeid}\t{_proc.memory_info().rss >> 20}\t{n_disp}\t{n_over}\n")
    _out.flush()

Run with:

PYTHONPATH=<plugin dir> pytest datashader/tests/test_pandas.py datashader/tests/test_utils.py \
  -k "masked or dict_source or arrow or polars" -p no:randomly -p jit_plugin -p append_plugin

AI Disclosure

Tool & Model: Claude Code + Opus 5.5
Usage: Investigated the CI memory growth, wrote the change and the measurement scripts.

• I have tested all AI-generated content in my PR.
• I take responsibility for all AI-generated content in my PR.

Checklist

• Tests added and are passing

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.70%. Comparing base (44c4beb) to head (c3472c5).
⚠️ Report is 18 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1540      +/-   ##
==========================================
+ Coverage   88.50%   88.70%   +0.19%     
==========================================
  Files          96       96              
  Lines       19536    19545       +9     
==========================================
+ Hits        17291    17338      +47     
+ Misses       2245     2207      -38     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[codspeed]

Merging this PR will not alter performance

✅ 52 untouched benchmarks

---

_{Comparing cache_append (c3472c5) with main (385fb0b)}