Make the tests fail for the implementations they claim to catch

Each of these passed with the defect it existed to catch:

- The projection round trip never checked the projection was applied; a
  decoder that parsed it and passed None through produced the same values,
  just with extra columns. The schema is now asserted.
- Nothing exercised the limit half of the storage descriptor in either
  direction. A new test executes the restored scan node per partition,
  because the limit operator above the scan would mask a dropped one.
- The bundle-reuse test compared two session codec ids, which differ with
  no bundle installed at all. It now round-trips a plan through the second
  session and pins the decode counter.
- The decline test offered the codec nothing: a registered name is resolved
  by the registry before any codec is consulted. The foreign name is now
  manufactured by a length-preserving edit of the encoded plan, and the
  declined counter is asserted.
- aggregate_by_name and window_by_name, and the whole logical codec, ran in
  no test. Two new workers decode the aggregate and window function with
  nothing registered, and a receiver with only the logical codec installed
  rebuilds the scalar from a logical plan.
- The by-value pickle claim was pinned on one side only; the payload now has
  a size floor to match the by-reference cap.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: timsaucer

examples/distributed/engine-library/python/tests/_test_three_libraries.py

@@ -145,6 +145,11 @@ def bucket(prices: pa.Array) -> pa.Array:
     assert stage is not None
     # The callable itself is in the bytes, under the scalar-UDF family prefix.
     assert b"DFPYUDF" in stage.to_bytes(ctx)
+    # By value means the payload carries bytecode, not a two-word pointer --
+    # the other half of the size claim pinned at the by-reference test's
+    # `< 200`. If cloudpickle ever started resolving `<locals>` functions by
+    # reference, this is the line that would notice.
+    assert len(cloudpickle.dumps(bucket)) > 500
 
     result = run_distributed(sql, spec, extra_udfs=[price_bucket])
     # Prices run from one hundred to eight hundred, so four exceed four hundred.

examples/distributed/storage-library/python/tests/_test_portable_codec.py

@@ -131,13 +131,52 @@ def test_a_projection_survives_the_round_trip(readings_dir: pathlib.Path) -> Non
     plan = ctx.sql("select reading from readings").execution_plan()
     restored = ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx))
 
-    rows = [
-        value
+    batches = [
+        batch.to_pyarrow()
         for partition in range(restored.partition_count)
         for batch in ctx.execute(restored, partition)
-        for value in batch.to_pyarrow().column("reading").to_pylist()
     ]
-    assert sorted(rows) == [1.5, 1.5, 1.5, 2.5, 2.5, 2.5, 3.5, 3.5, 3.5]
+    # Exactly one column, not merely the right values in it: a decoder that
+    # parsed the projection and then dropped it would still produce every
+    # `reading` value below, just alongside the columns that were projected
+    # away. The schema is where that shows.
+    assert [batch.column_names for batch in batches] == [["reading"]] * 3
+    rows = sorted(
+        value for batch in batches for value in batch.column("reading").to_pylist()
+    )
+    assert rows == [1.5, 1.5, 1.5, 2.5, 2.5, 2.5, 3.5, 3.5, 3.5]
+
+
+def _scan_of(plan: ExecutionPlan) -> ExecutionPlan:
+    """Descend to the leaf, which for these plans is this library's scan."""
+    children = plan.children()
+    return _scan_of(children[0]) if children else plan
+
+
+def test_a_limit_survives_the_round_trip(readings_dir: pathlib.Path) -> None:
+    """The limit is part of the descriptor too, and it changes what is read.
+
+    Asserted on the scan node itself rather than on the whole plan: the
+    logical `PushDownLimit` rule keeps a limit operator *above* the scan, and
+    that operator would trim the result to the right count even if the decoder
+    dropped the scan's own limit. Executing the restored scan directly is what
+    makes a dropped limit visible -- three rows per file instead of two.
+    """
+    ctx, _ = _configured(readings_dir)
+    plan = ctx.sql("select reading from readings limit 2").execution_plan()
+    scan = _scan_of(plan)
+    blob = scan.to_bytes(ctx)
+    assert b'"limit":2' in blob
+
+    restored = ExecutionPlan.from_bytes(ctx, blob)
+    per_partition = [
+        sum(batch.to_pyarrow().num_rows for batch in ctx.execute(restored, partition))
+        for partition in range(restored.partition_count)
+    ]
+    # The limit is applied per partition -- each file stops after two of its
+    # three rows. The trim to two rows overall is the limit operator's job,
+    # and that operator is deliberately not part of what was serialized here.
+    assert per_partition == [2, 2, 2]
 
 
 def test_stock_nodes_never_reach_this_codec(readings_dir: pathlib.Path) -> None:
@@ -343,4 +382,13 @@ def test_the_bundle_is_reusable_across_sessions(readings_dir: pathlib.Path) -> N
             == 9
         )
 
-    assert first.__datafusion_codec_id__ != second.__datafusion_codec_id__
+    # The second install has to carry live codecs of its own, not a component
+    # cached from the first -- so the proof is a round trip through `second`
+    # specifically, not a property of the session objects.
+    assert bundle.decode_calls() == 0
+    blob = (
+        second.sql("select sensor_id from readings").execution_plan().to_bytes(second)
+    )
+    restored = ExecutionPlan.from_bytes(second, blob)
+    assert restored.partition_count == 3
+    assert bundle.decode_calls() == 1

examples/distributed/udf-library/python/tests/_test_udfs.py

@@ -27,7 +27,7 @@
 
 import pytest
 from datafusion import SessionConfig, SessionContext, udaf, udf, udwf
-from datafusion.plan import ExecutionPlan
+from datafusion.plan import ExecutionPlan, LogicalPlan
 from dfx_udfs import (
     CodecObservations,
     NetRevenueUDF,
@@ -184,6 +184,10 @@ def test_the_payload_carries_no_bytes(lineitem: pathlib.Path) -> None:
 )
 
 
+def _decoded(stdout: str) -> int:
+    return int(re.search(r"decoded=(\d+)", stdout).group(1))
+
+
 def _run_worker(
     tmp_path: pathlib.Path, blob: bytes, mode: str, data: pathlib.Path
 ) -> subprocess.CompletedProcess[str]:
@@ -252,23 +256,115 @@ def test_a_worker_with_neither_names_the_function_it_cannot_find(
     assert "dfx_net_revenue" in result.stderr
 
 
+def test_a_worker_with_only_the_codec_rebuilds_the_aggregate(
+    lineitem: pathlib.Path, tmp_path: pathlib.Path
+) -> None:
+    """The aggregate travels by name too, through `aggregate_by_name`.
+
+    Worth its own worker: the scalar tests exercise `scalar_by_name` and
+    nothing else, so a typo in the aggregate's decode arm would only ever
+    surface here.
+    """
+    ctx, _ = _session(lineitem)
+    blob = (
+        ctx.sql("select dfx_weighted_avg(l_extendedprice, l_quantity) from lineitem")
+        .execution_plan()
+        .to_bytes(ctx)
+    )
+
+    result = _run_worker(tmp_path, blob, "codec", lineitem)
+
+    assert result.returncode == 0, result.stderr
+    # (100*1 + 200*3 + 400*4) / (1 + 3 + 4), same as the in-process test.
+    assert "total=287.5" in result.stdout
+    # More than one decode is fine -- the plan holds the function once per
+    # aggregate phase -- but zero would mean the registry answered, and this
+    # worker has nothing registered.
+    assert _decoded(result.stdout) >= 1
+
+
+def test_a_worker_with_only_the_codec_rebuilds_the_window_function(
+    lineitem: pathlib.Path, tmp_path: pathlib.Path
+) -> None:
+    """And the window function, through `window_by_name`."""
+    ctx, _ = _session(lineitem)
+    blob = (
+        ctx.sql(
+            "select dfx_revenue_rank() over (order by l_extendedprice desc) as rnk "
+            "from lineitem"
+        )
+        .execution_plan()
+        .to_bytes(ctx)
+    )
+
+    result = _run_worker(tmp_path, blob, "codec", lineitem)
+
+    assert result.returncode == 0, result.stderr
+    # Ranks 1 + 2 + 3 over the three rows.
+    assert "total=6.0" in result.stdout
+    assert _decoded(result.stdout) >= 1
+
+
 def test_the_codec_declines_names_it_does_not_own(lineitem: pathlib.Path) -> None:
     """A name-only payload reaches every codec, so declining matters.
 
     With no bytes there is no codec id to route on. A codec that answered for
     any name it was handed would hijack another library's functions.
+
+    The foreign name is manufactured by editing the encoded plan: the session's
+    own registry is consulted before any codec, so a name that is actually
+    registered anywhere -- a built-in, say -- never reaches a codec at all.
+    The replacement keeps the byte length, because the name sits inside
+    length-delimited protobuf fields.
     """
     ctx, observations = _session(lineitem, with_codecs=True)
-    # `abs` is a built-in, so the plan references a name this library does not
-    # own; decoding offers it around.
     blob = (
-        ctx.sql("select abs(l_discount) from lineitem").execution_plan().to_bytes(ctx)
+        ctx.sql(
+            "select dfx_net_revenue(l_extendedprice, l_discount, l_tax) from lineitem"
+        )
+        .execution_plan()
+        .to_bytes(ctx)
     )
-    ExecutionPlan.from_bytes(ctx, blob)
+    mangled = blob.replace(b"dfx_net_revenue", b"dfx_not_revenue")
+    assert len(mangled) == len(blob), "the edit has to preserve the protobuf framing"
 
+    with pytest.raises(Exception, match="dfx_not_revenue"):
+        ExecutionPlan.from_bytes(ctx, mangled)
+
+    # The codec was offered the foreign name and said no; it never rebuilt
+    # anything. A codec that hijacked the name would show up as the mirror
+    # image -- a decode, no decline, and a query that "works".
+    assert observations.declined_calls() >= 1
     assert observations.decode_calls() == 0
 
 
+def test_the_logical_codec_rebuilds_functions_by_name(lineitem: pathlib.Path) -> None:
+    """The logical half is a peer, not a passenger.
+
+    `LogicalPlan.to_bytes` is the layer an engine that ships *logical* plans
+    exercises, and nothing about the physical tests touches it. The receiving
+    session installs only the logical codec and registers nothing, so the
+    rebuild below can only have come through `try_decode_udf` on this codec.
+    """
+    ctx, _ = _session(lineitem)
+    blob = (
+        ctx.sql(
+            "select dfx_net_revenue(l_extendedprice, l_discount, l_tax) from lineitem"
+        )
+        .logical_plan()
+        .to_bytes(ctx)
+    )
+
+    receiver_observations = CodecObservations()
+    receiver = SessionContext().with_logical_extension_codec(
+        receiver_observations.logical_codec()
+    )
+    restored = LogicalPlan.from_bytes(receiver, blob)
+
+    assert receiver_observations.decode_calls() == 1
+    assert "dfx_net_revenue" in restored.display_indent()
+
+
 def test_the_codec_ids_are_pinned() -> None:
     """Renaming the exporting class must not invalidate written plans."""
     observations = CodecObservations()