ntjohnson1 commented on code in PR #1678:
URL:
https://github.com/apache/datafusion-python/pull/1678#discussion_r3786206600
##########
crates/core/src/codec.rs:
##########
@@ -223,32 +224,129 @@ fn strip_wire_header<'a>(
Ok(Some(&buf[py_minor_idx + 1..]))
}
+/// Run `f` against each codec in `chain`, returning the first `Ok`.
+///
+/// A codec signals "not mine" by returning an error, so the chain
+/// keeps trying until a codec succeeds. When every codec fails and the
+/// chain has more than one entry, the errors are aggregated into a
+/// single message — returning only the last error would surface the
+/// terminal `Default*ExtensionCodec` "not provided" message and mask
+/// the more specific diagnostic from an installed codec (e.g. a
+/// corrupt-token error from the codec that owns the payload family).
+fn chain_try<C: ?Sized, R>(chain: &[Arc<C>], what: &str, f: impl Fn(&C) ->
Result<R>) -> Result<R> {
+ let mut errors: Vec<datafusion::error::DataFusionError> = Vec::new();
+ for codec in chain {
+ match f(codec) {
+ Ok(value) => return Ok(value),
+ Err(err) => errors.push(err),
+ }
+ }
+ Err(aggregate_chain_errors(what, errors))
+}
+
+/// Collapse per-codec failures into one error. A single failure is
+/// returned as-is so the one-codec (default-only) chain behaves
+/// exactly like the pre-chain implementation.
+fn aggregate_chain_errors(
+ what: &str,
+ mut errors: Vec<datafusion::error::DataFusionError>,
+) -> datafusion::error::DataFusionError {
+ match errors.len() {
+ 0 => datafusion::error::DataFusionError::Internal(format!(
+ "Empty extension codec chain while handling {what}"
+ )),
+ 1 => errors.swap_remove(0),
+ _ => {
+ let joined = errors
+ .iter()
+ .map(|err| err.to_string())
+ .collect::<Vec<_>>()
+ .join("; ");
+ datafusion::error::DataFusionError::Execution(format!(
+ "None of the {} composed extension codecs handled {what}:
{joined}",
+ errors.len()
+ ))
+ }
+ }
+}
+
+/// Encode variant of [`chain_try`] for methods that write into a
+/// caller-provided buffer.
+///
+/// Each codec encodes into a scratch buffer so a failed attempt cannot
+/// leave partial bytes behind. `Ok` with bytes written commits those
+/// bytes and ends the chain. `Ok` with an empty buffer is treated as
+/// "no opinion" — the standard `Default*ExtensionCodec` behavior of
+/// encoding a UDF by name writes nothing — so later codecs still get a
+/// chance to emit a richer payload. If no codec writes bytes but at
+/// least one returned `Ok`, the overall result is `Ok` with nothing
+/// written (encode by name).
+fn chain_encode<C: ?Sized>(
+ chain: &[Arc<C>],
+ buf: &mut Vec<u8>,
+ what: &str,
+ f: impl Fn(&C, &mut Vec<u8>) -> Result<()>,
+) -> Result<()> {
+ let mut saw_empty_ok = false;
+ let mut errors: Vec<datafusion::error::DataFusionError> = Vec::new();
+ for codec in chain {
+ let mut scratch = Vec::new();
+ match f(codec, &mut scratch) {
+ Ok(()) if !scratch.is_empty() => {
+ buf.extend_from_slice(&scratch);
+ return Ok(());
+ }
+ Ok(()) => saw_empty_ok = true,
+ Err(err) => errors.push(err),
+ }
+ }
+ if saw_empty_ok {
+ return Ok(());
+ }
+ Err(aggregate_chain_errors(what, errors))
+}
+
/// `LogicalExtensionCodec` parked on every `SessionContext`. Holds
/// the Python-aware encoding hooks for logical-layer types
/// (`LogicalPlan`, `Expr`) and delegates everything it does not
-/// handle to the composable `inner` codec — typically
-/// `DefaultLogicalExtensionCodec`, or a downstream FFI codec
-/// installed via `SessionContext.with_logical_extension_codec(...)`.
+/// handle to a chain of composable codecs. The chain starts as just
+/// `DefaultLogicalExtensionCodec`; each downstream FFI codec installed
+/// via `SessionContext.with_logical_extension_codec(...)` is prepended,
+/// so the most recently installed codec is consulted first and the
+/// default codec always runs last.
+///
+/// Chain dispatch relies on each codec recognizing its own payloads
+/// (distinct family prefixes — see the module docs) and returning an
+/// error for everything else so the next codec gets a chance.
///
/// Sitting at the top of the session's logical codec stack means
/// every serializer that reads `session.logical_codec()` automatically
/// picks up Python-aware encoding for free.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
pub struct PythonLogicalCodec {
- inner: Arc<dyn LogicalExtensionCodec>,
+ chain: Vec<Arc<dyn LogicalExtensionCodec>>,
python_udf_inlining: bool,
}
impl PythonLogicalCodec {
pub fn new(inner: Arc<dyn LogicalExtensionCodec>) -> Self {
Self {
- inner,
+ chain: vec![inner],
python_udf_inlining: true,
}
}
- pub fn inner(&self) -> &Arc<dyn LogicalExtensionCodec> {
Review Comment:
This was pub before. Do people care about inspecting the codecs directly?
##########
docs/source/contributor-guide/ffi.md:
##########
@@ -248,15 +248,65 @@ foreign planner. This lets the planner decode
provider-owned objects and lets
process-local tokens to demonstrate ownership; production codecs should
serialize
durable metadata instead.
-The current Python API has one external logical codec and one external
physical codec.
-Installing another codec replaces the prior codec rather than composing a
registry.
-The example therefore has one external codec owner, and the planner uses
built-in
-physical nodes. Install the provider codecs before the planner where possible.
+### Composable codecs
+
+Extension codecs compose. Each call to `with_logical_extension_codec` or
+`with_physical_extension_codec` adds the codec to the front of the session's
codec
+chain rather than replacing prior codecs. During encoding and decoding, the
most
+recently installed codec is consulted first, falling through codec by codec to
+DataFusion's default codec. A codec signals "not mine" by returning an error,
which
+sends the chain on to the next codec. Two conventions keep this dispatch sound:
+
+- Frame your payloads with a distinct byte prefix (pick a `DF` namespace plus a
+ crate-specific suffix) and only decode payloads carrying your prefix.
+- Return an error for objects and payloads you do not own. A codec that answers
+ success for objects outside its family shadows every codec installed before
it.
+
+Because dispatch keys off payload prefixes rather than install position, codec
+registration order between independent libraries does not matter.
The current FFI logical codec supports providers and UDFs but not arbitrary
custom
`LogicalPlan::Extension` nodes. See both example READMEs for the supported
flow and
local build commands.
+### One planner per session, with explicit fallback
+
+Unlike codecs, a `SessionState` holds exactly one query planner — installing
another
+replaces it. Planner layering is therefore explicit: a planner that wants to
handle
+only some queries should accept a fallback planner and delegate the rest to
it. The
+current planner can be exported for that purpose with
+`ctx.__datafusion_query_planner__()`.
+
+One ordering rule applies: a planner capsule captures the session's codecs at
export
+time and cannot be rebound afterward. Codec changes made after installing a
single
+planner are rebound automatically, but a planner wrapped inside another
planner as a
+fallback is opaque and keeps the codecs it was exported with. **Install all
extension
+codecs before exporting or chaining planners.**
+
+Putting it together for a session using two extension libraries that each
provide
+tables, functions, and a query planner:
+
+```python
+ctx = SessionContext(config)
+
+# 1. Codecs from both libraries. Order between libraries does not matter.
+ctx = ctx.with_logical_extension_codec(lib_a.codec())
+ctx = ctx.with_logical_extension_codec(lib_b.codec())
Review Comment:
So is the resulting stage at this point `[default, lib_a, lib_b]`? I could
read it or ask claude if some assumed based codec for standard datafusion
python gets installed always as a fallback. Mostly I'm curious if it's a
reasonable workflow for someone to set things up to only get `[lib_a, lib_b]`
and if anything unsupported by their custom codecs barfs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
