test: port codec Rust tests to pytest

The Rust tests added for the composable codec work never ran: CI invokes
`cargo fmt` and `cargo clippy --all-targets` but no `cargo test`, so the
tests compiled and were never executed. Rather than add a `cargo test`
job — which would also require feature-gating `pyo3/extension-module`,
since the test binary cannot link on Linux while it is unconditional —
move the coverage to pytest, matching this repository's practice of
treating the user-facing Python surface as the first line of defense.

Remove both `#[cfg(test)]` modules from crates/core/src/codec.rs and
replace them as follows:

- Four wire-header round-trip tests and the Python-minor-mismatch test
  were already covered by existing cases in test_pickle_expr.py.
- `strip_errors_on_too_old_version` asserted nothing: it returns early
  because WIRE_VERSION_MIN_SUPPORTED equals WIRE_VERSION_CURRENT.
- The unsupported-wire-version and Python-major-mismatch cases move to
  test_pickle_expr.py, patching the header in place inside the encoded
  protobuf. The patches preserve length so the outer message stays
  parseable and the bytes reach the codec.
- The three truncated-header cases are dropped. Truncation changes the
  payload length and breaks the protobuf framing, so they fail before
  reaching the header check and cannot be expressed from Python.
- The codec-chain tests move to the FFI example suite, which exercises
  the same chain through the real FFI boundary.

MyLogicalExtensionCodec gains an optional token overriding the byte
prefix it stamps on encoded table providers. Two instances with distinct
tokens own disjoint slices of the wire format, which is what makes chain
ordering and fall-through observable from Python.

The ported inlining test asserts encode and decode behavior rather than
the `python_udf_inlining()` getter the Rust test checked. This is a
stronger assertion: the getter is preserved even when a composed
Python-aware codec re-inlines a UDF that the outer strict codec declined
to inline, so the original test could not have caught that path.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: timsaucer

crates/core/src/codec.rs

@@ -1182,315 +1182,3 @@ fn decode_python_udaf(py: Python<'_>, payload: &[u8]) -> PyResult<PythonFunction
         volatility,
     ))
 }
-
-#[cfg(test)]
-mod wire_header_tests {
-    use super::*;
-
-    const TEST_PY: (u8, u8) = (3, 12);
-
-    #[test]
-    fn strip_returns_none_when_family_absent() {
-        let buf = b"OTHER_PAYLOAD";
-        assert!(matches!(
-            strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY),
-            Ok(None)
-        ));
-    }
-
-    #[test]
-    fn strip_errors_on_truncated_version_byte() {
-        let buf = PY_SCALAR_UDF_FAMILY;
-        let err = strip_wire_header(buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
-        assert!(format!("{err}").contains("missing wire-format version byte"));
-    }
-
-    #[test]
-    fn strip_errors_on_too_new_version() {
-        let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
-        buf.push(WIRE_VERSION_CURRENT.saturating_add(1));
-        buf.push(TEST_PY.0);
-        buf.push(TEST_PY.1);
-        buf.extend_from_slice(b"payload");
-        let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
-        let msg = format!("{err}");
-        assert!(msg.contains("wire-format version v"));
-        assert!(msg.contains("supports"));
-        assert!(msg.contains("Align datafusion-python versions"));
-    }
-
-    #[test]
-    fn strip_errors_on_too_old_version() {
-        if WIRE_VERSION_MIN_SUPPORTED == 0 {
-            return;
-        }
-        let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
-        buf.push(WIRE_VERSION_MIN_SUPPORTED - 1);
-        buf.push(TEST_PY.0);
-        buf.push(TEST_PY.1);
-        buf.extend_from_slice(b"payload");
-        assert!(strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).is_err());
-    }
-
-    #[test]
-    fn strip_errors_on_truncated_py_major() {
-        let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
-        buf.push(WIRE_VERSION_CURRENT);
-        let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
-        assert!(format!("{err}").contains("missing Python major version byte"));
-    }
-
-    #[test]
-    fn strip_errors_on_truncated_py_minor() {
-        let mut buf = PY_SCALAR_UDF_FAMILY.to_vec();
-        buf.push(WIRE_VERSION_CURRENT);
-        buf.push(TEST_PY.0);
-        let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY).unwrap_err();
-        assert!(format!("{err}").contains("missing Python minor version byte"));
-    }
-
-    #[test]
-    fn strip_errors_on_py_minor_mismatch() {
-        let mut buf = Vec::new();
-        write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, (3, 11));
-        buf.extend_from_slice(b"payload");
-        let err = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", (3, 12)).unwrap_err();
-        let msg = format!("{err}");
-        assert!(msg.contains("Python 3.11"));
-        assert!(msg.contains("Python 3.12"));
-        assert!(msg.contains("not portable across Python minor versions"));
-    }
-
-    #[test]
-    fn strip_errors_on_py_major_mismatch() {
-        let mut buf = Vec::new();
-        write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, (3, 12));
-        buf.extend_from_slice(b"payload");
-        assert!(strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", (4, 0)).is_err());
-    }
-
-    #[test]
-    fn write_then_strip_round_trips_scalar_payload() {
-        let mut buf = Vec::new();
-        write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, TEST_PY);
-        buf.extend_from_slice(b"scalar-payload");
-
-        let payload = strip_wire_header(&buf, PY_SCALAR_UDF_FAMILY, "scalar UDF", TEST_PY)
-            .unwrap()
-            .unwrap();
-        assert_eq!(payload, b"scalar-payload");
-    }
-
-    #[test]
-    fn write_then_strip_round_trips_agg_payload() {
-        let mut buf = Vec::new();
-        write_wire_header(&mut buf, PY_AGG_UDF_FAMILY, TEST_PY);
-        buf.extend_from_slice(b"agg-payload");
-
-        let payload = strip_wire_header(&buf, PY_AGG_UDF_FAMILY, "aggregate UDF", TEST_PY)
-            .unwrap()
-            .unwrap();
-        assert_eq!(payload, b"agg-payload");
-    }
-
-    #[test]
-    fn write_then_strip_round_trips_window_payload() {
-        let mut buf = Vec::new();
-        write_wire_header(&mut buf, PY_WINDOW_UDF_FAMILY, TEST_PY);
-        buf.extend_from_slice(b"window-payload");
-
-        let payload = strip_wire_header(&buf, PY_WINDOW_UDF_FAMILY, "window UDF", TEST_PY)
-            .unwrap()
-            .unwrap();
-        assert_eq!(payload, b"window-payload");
-    }
-
-    #[test]
-    fn strip_does_not_match_a_different_family() {
-        let mut buf = Vec::new();
-        write_wire_header(&mut buf, PY_SCALAR_UDF_FAMILY, TEST_PY);
-        buf.extend_from_slice(b"payload");
-        assert!(matches!(
-            strip_wire_header(&buf, PY_WINDOW_UDF_FAMILY, "window UDF", TEST_PY),
-            Ok(None)
-        ));
-    }
-}
-
-#[cfg(test)]
-mod codec_chain_tests {
-    use std::sync::atomic::{AtomicUsize, Ordering};
-
-    use datafusion::catalog::MemTable;
-    use datafusion::common::exec_err;
-
-    use super::*;
-
-    /// Codec that owns a single byte token for table providers and
-    /// errors on everything else, mirroring the family-prefix
-    /// discipline expected of downstream FFI codecs.
-    #[derive(Debug)]
-    struct TokenCodec {
-        token: &'static [u8],
-        /// Return `Ok` from `try_encode_table_provider` without
-        /// writing bytes, imitating a "no opinion" codec.
-        encode_by_name: bool,
-        decode_hits: AtomicUsize,
-        encode_hits: AtomicUsize,
-    }
-
-    impl TokenCodec {
-        fn new(token: &'static [u8]) -> Arc<Self> {
-            Arc::new(Self {
-                token,
-                encode_by_name: false,
-                decode_hits: AtomicUsize::new(0),
-                encode_hits: AtomicUsize::new(0),
-            })
-        }
-
-        fn new_by_name(token: &'static [u8]) -> Arc<Self> {
-            Arc::new(Self {
-                token,
-                encode_by_name: true,
-                decode_hits: AtomicUsize::new(0),
-                encode_hits: AtomicUsize::new(0),
-            })
-        }
-    }
-
-    impl LogicalExtensionCodec for TokenCodec {
-        fn try_decode(
-            &self,
-            _buf: &[u8],
-            _inputs: &[LogicalPlan],
-            _ctx: &TaskContext,
-        ) -> Result<Extension> {
-            exec_err!("TokenCodec does not decode extension nodes")
-        }
-
-        fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
-            exec_err!("TokenCodec does not encode extension nodes")
-        }
-
-        fn try_decode_table_provider(
-            &self,
-            buf: &[u8],
-            _table_ref: &TableReference,
-            schema: SchemaRef,
-            _ctx: &TaskContext,
-        ) -> Result<Arc<dyn TableProvider>> {
-            if buf != self.token {
-                return exec_err!("Unknown table provider token for TokenCodec");
-            }
-            self.decode_hits.fetch_add(1, Ordering::SeqCst);
-            Ok(Arc::new(MemTable::try_new(schema, vec![vec![]])?))
-        }
-
-        fn try_encode_table_provider(
-            &self,
-            _table_ref: &TableReference,
-            _node: Arc<dyn TableProvider>,
-            buf: &mut Vec<u8>,
-        ) -> Result<()> {
-            self.encode_hits.fetch_add(1, Ordering::SeqCst);
-            if !self.encode_by_name {
-                buf.extend_from_slice(self.token);
-            }
-            Ok(())
-        }
-    }
-
-    fn mem_table() -> Arc<dyn TableProvider> {
-        Arc::new(MemTable::try_new(Arc::new(Schema::empty()), vec![vec![]]).unwrap())
-    }
-
-    fn table_ref() -> TableReference {
-        TableReference::bare("t")
-    }
-
-    #[test]
-    fn decode_falls_through_to_earlier_installed_codec() {
-        let first = TokenCodec::new(b"AAAA");
-        let second = TokenCodec::new(b"BBBB");
-        let codec = PythonLogicalCodec::default()
-            .with_additional_codec(first.clone())
-            .with_additional_codec(second.clone());
-
-        let ctx = TaskContext::default();
-        codec
-            .try_decode_table_provider(b"AAAA", &table_ref(), Arc::new(Schema::empty()), &ctx)
-            .unwrap();
-
-        assert_eq!(first.decode_hits.load(Ordering::SeqCst), 1);
-        assert_eq!(second.decode_hits.load(Ordering::SeqCst), 0);
-    }
-
-    #[test]
-    fn most_recently_installed_codec_encodes_first() {
-        let first = TokenCodec::new(b"AAAA");
-        let second = TokenCodec::new(b"BBBB");
-        let codec = PythonLogicalCodec::default()
-            .with_additional_codec(first.clone())
-            .with_additional_codec(second.clone());
-
-        let mut buf = Vec::new();
-        codec
-            .try_encode_table_provider(&table_ref(), mem_table(), &mut buf)
-            .unwrap();
-
-        assert_eq!(buf, b"BBBB");
-        assert_eq!(first.encode_hits.load(Ordering::SeqCst), 0);
-    }
-
-    #[test]
-    fn empty_ok_encode_lets_later_codec_write_payload() {
-        let writer = TokenCodec::new(b"AAAA");
-        let by_name = TokenCodec::new_by_name(b"BBBB");
-        let codec = PythonLogicalCodec::default()
-            .with_additional_codec(writer.clone())
-            .with_additional_codec(by_name.clone());
-
-        let mut buf = Vec::new();
-        codec
-            .try_encode_table_provider(&table_ref(), mem_table(), &mut buf)
-            .unwrap();
-
-        assert_eq!(buf, b"AAAA");
-        assert_eq!(by_name.encode_hits.load(Ordering::SeqCst), 1);
-        assert_eq!(writer.encode_hits.load(Ordering::SeqCst), 1);
-    }
-
-    #[test]
-    fn decode_failure_aggregates_every_codec_error() {
-        let codec = PythonLogicalCodec::default()
-            .with_additional_codec(TokenCodec::new(b"AAAA"))
-            .with_additional_codec(TokenCodec::new(b"BBBB"));
-
-        let ctx = TaskContext::default();
-        let err = codec
-            .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx)
-            .unwrap_err();
-
-        let msg = err.to_string();
-        assert!(msg.contains("None of the 3 composed extension codecs"));
-        assert!(msg.contains("Unknown table provider token"));
-    }
-
-    #[test]
-    fn single_codec_chain_error_is_returned_verbatim() {
-        let codec = PythonLogicalCodec::default();
-        let ctx = TaskContext::default();
-        let err = codec
-            .try_decode_table_provider(b"????", &table_ref(), Arc::new(Schema::empty()), &ctx)
-            .unwrap_err();
-        assert!(!err.to_string().contains("composed extension codecs"));
-    }
-
-    #[test]
-    fn with_additional_codec_preserves_udf_inlining_setting() {
-        let strict = PythonLogicalCodec::default().with_python_udf_inlining(false);
-        let extended = strict.with_additional_codec(TokenCodec::new(b"AAAA"));
-        assert!(!extended.python_udf_inlining());
-    }
-}

examples/datafusion-ffi-example/README.md

@@ -37,6 +37,8 @@ The example codecs do not inspect the callback `TaskContext`. A production codec
 
 Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call prepends the codec to the session's codec chain, with the most recently installed codec consulted first and DataFusion's default codec as the terminal fallback. A codec signals "not mine" by returning an error, so several independent plugin libraries can install codecs on the same session as long as each only answers for payloads it owns (frame them with a distinct byte prefix). In this example the provider library is the only codec owner; the planner uses built-in physical nodes and receives the provider codecs from the host.
 
+`MyLogicalExtensionCodec` takes an optional token argument (`MyLogicalExtensionCodec("TOKENAAA")`) that overrides the byte prefix it stamps on encoded table providers. It exists so the tests can install two instances that own disjoint slices of the wire format, which is what makes chain ordering and fall-through observable from Python. Real plugin libraries should hard-code a prefix unique to the library rather than accept one from the caller.
+
 Register both provider codecs before installing the planner:
 
 ```python