milenkovicm commented on code in PR #1678:
URL:
https://github.com/apache/datafusion-python/pull/1678#discussion_r3918593844
##########
examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py:
##########
@@ -80,3 +118,430 @@ def test_ffi_logical_codec_roundtrip():
restored = LogicalPlan.from_bytes(ctx, blob)
df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
assert df.collect() == df_round_trip.collect()
+
+
+def test_ffi_logical_codec_composes_with_later_install():
+ """Codecs compose: installing a second codec appends it to the
+ session's codec chain instead of replacing the first. The second
+ codec here (a default-backed codec exported from a fresh session)
+ cannot encode this library's table provider, so the first codec
+ still claims it. Under replace semantics this test fails with
+ `LogicalExtensionCodec is not provided`."""
+ ctx, codec = _setup_session_with_codec()
+ ctx = ctx.with_logical_extension_codec(
+ SessionContext().__datafusion_logical_extension_codec__()
+ )
+
+ ctx.register_table("numbers", MyTableProvider(1, 4, 1))
+ df = ctx.sql('SELECT "A" FROM numbers')
+ plan = df.logical_plan()
+
+ before = codec.table_provider_encode_calls()
+ blob = plan.to_bytes(ctx)
+ assert codec.table_provider_encode_calls() > before
+
+ restored = LogicalPlan.from_bytes(ctx, blob)
+ df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
+ assert df.collect() == df_round_trip.collect()
Review Comment:
nit, should you assert on decode call number as well?
##########
examples/datafusion-ffi-example/src/name_only_codec.rs:
##########
@@ -0,0 +1,268 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A codec whose functions need no payload at all.
+//!
+//! Most extension codecs answer with bytes. This one owns a fixed catalog of
+//! functions that are fully described by their names, so `try_encode_udf`
+//! writes nothing and `try_decode_udf` rebuilds the function from `name`
+//! alone. DataFusion supports that shape directly: an encoder that writes no
+//! bytes leaves `fun_definition` unset, and the decoder then tries the
+//! `FunctionRegistry` first and the codec second — see the
+//! `None => ctx.udf(..).or_else(|_| codec.try_decode_udf(name, &[]))` arm in
+//! `datafusion-proto`'s `from_proto.rs`.
+//!
+//! It exists here to pin that arm. Because there are no bytes, there is
+//! nothing to tag with the codec's identity, so this is the one path where
+//! `PythonLogicalCodec` still offers a payload to every installed codec in
+//! turn. A change that wrapped empty encodings in an envelope would set
+//! `fun_definition`, skip the registry lookup permanently, and break both this
+//! codec and plain by-name round trips — with no other test noticing.
+
+use std::fmt;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use arrow_schema::DataType;
+use datafusion::common::error::Result;
+use datafusion::common::not_impl_err;
+use datafusion::logical_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
TypeSignature,
+ Volatility,
+};
+use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec,
LogicalExtensionCodec};
+use datafusion_python_util::{ffi_task_context_provider_from_pycapsule,
get_tokio_runtime};
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+/// Prefix marking the functions this library owns. A name is the entire
+/// encoding, so the prefix is the whole ownership test.
+const NAME_PREFIX: &str = "name_only_";
+
+/// Scalar function reconstructed purely from its name.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct NameOnlyUdf {
+ name: String,
+ signature: Signature,
+}
+
+impl NameOnlyUdf {
+ fn new(name: impl Into<String>) -> Self {
+ Self {
+ name: name.into(),
+ signature: Signature::new(TypeSignature::Any(1),
Volatility::Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for NameOnlyUdf {
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
+ Ok(args.args[0].clone())
+ }
+}
+
+#[derive(Default)]
+struct Counters {
+ encode_udf: AtomicUsize,
+ decode_udf: AtomicUsize,
+}
+
+impl fmt::Debug for Counters {
+ fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+ formatter.debug_struct("Counters").finish_non_exhaustive()
+ }
+}
+
+struct NameOnlyLogicalExtensionCodec {
+ inner: DefaultLogicalExtensionCodec,
+ counters: Arc<Counters>,
+}
+
+impl fmt::Debug for NameOnlyLogicalExtensionCodec {
+ fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+ formatter
+ .debug_struct("NameOnlyLogicalExtensionCodec")
+ .finish_non_exhaustive()
+ }
+}
+
+impl LogicalExtensionCodec for NameOnlyLogicalExtensionCodec {
+ fn try_decode(
+ &self,
+ buf: &[u8],
+ inputs: &[datafusion::logical_expr::LogicalPlan],
+ ctx: &datafusion::execution::TaskContext,
+ ) -> Result<datafusion::logical_expr::Extension> {
+ self.inner.try_decode(buf, inputs, ctx)
+ }
+
+ fn try_encode(
+ &self,
+ node: &datafusion::logical_expr::Extension,
+ buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ self.inner.try_encode(node, buf)
+ }
+
+ fn try_decode_table_provider(
+ &self,
+ buf: &[u8],
+ table_ref: &datafusion::common::TableReference,
+ schema: arrow_schema::SchemaRef,
+ ctx: &datafusion::execution::TaskContext,
+ ) -> Result<Arc<dyn datafusion::datasource::TableProvider>> {
+ self.inner
+ .try_decode_table_provider(buf, table_ref, schema, ctx)
+ }
+
+ fn try_encode_table_provider(
+ &self,
+ table_ref: &datafusion::common::TableReference,
+ node: Arc<dyn datafusion::datasource::TableProvider>,
+ buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ self.inner.try_encode_table_provider(table_ref, node, buf)
+ }
+
+ /// Writes nothing on purpose. The name is the whole encoding, so there is
+ /// no payload to emit, and returning `Ok` with an empty buffer is how a
+ /// codec says "encoded by name" to DataFusion.
+ fn try_encode_udf(&self, node: &ScalarUDF, _buf: &mut Vec<u8>) ->
Result<()> {
+ if node.name().starts_with(NAME_PREFIX) {
+ self.counters.encode_udf.fetch_add(1, Ordering::SeqCst);
+ }
+ Ok(())
+ }
+
+ /// Rebuilds the function from `name`, with no registry entry and no bytes.
+ fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>>
{
+ if !name.starts_with(NAME_PREFIX) {
+ return not_impl_err!("Not a name-only function: {name}");
+ }
+ if !buf.is_empty() {
+ return not_impl_err!(
+ "name-only functions carry no payload, but {} bytes were
supplied for {name}",
+ buf.len()
+ );
+ }
+ self.counters.decode_udf.fetch_add(1, Ordering::SeqCst);
+ Ok(Arc::new(ScalarUDF::from(NameOnlyUdf::new(name))))
+ }
+}
+
+/// The function [`NameOnlyUdfCodec`] owns, exported so a session can register
+/// it and build a plan that references it.
+///
+/// Only the *encoding* session needs it registered. The decoding session
+/// deliberately does not, which is what forces the codec's name-only decode
+/// path to run.
+#[pyclass(
+ from_py_object,
+ name = "NameOnlyFunction",
+ module = "datafusion_ffi_example",
+ subclass
+)]
+#[derive(Debug, Clone)]
+pub(crate) struct NameOnlyFunction;
+
+#[pymethods]
+impl NameOnlyFunction {
+ #[new]
+ fn new() -> Self {
+ Self
+ }
+
+ fn __datafusion_scalar_udf__<'py>(&self, py: Python<'py>) ->
PyResult<Bound<'py, PyCapsule>> {
+ let func = Arc::new(ScalarUDF::from(NameOnlyUdf::new(format!(
+ "{NAME_PREFIX}identity"
+ ))));
+ PyCapsule::new_with_value(
+ py,
+ datafusion_ffi::udf::FFI_ScalarUDF::from(func),
+ cr"datafusion_scalar_udf",
+ )
+ }
+}
+
+/// Codec owning functions that are reconstructible from their names alone.
+///
+/// A real library shaped like this would be one shipping a fixed catalog of
+/// built-ins: nothing about a call site varies, so there is nothing to encode.
+#[pyclass(
+ from_py_object,
+ name = "NameOnlyUdfCodec",
+ module = "datafusion_ffi_example",
+ subclass
+)]
+#[derive(Clone)]
+pub(crate) struct NameOnlyUdfCodec {
+ counters: Arc<Counters>,
+}
+
+#[pymethods]
+impl NameOnlyUdfCodec {
+ #[new]
+ fn new() -> Self {
+ Self {
+ counters: Arc::new(Counters::default()),
+ }
+ }
+
+ /// Name of the function this codec can rebuild, for use in a query.
+ #[staticmethod]
+ fn function_name() -> String {
+ format!("{NAME_PREFIX}identity")
+ }
+
+ fn encode_udf_calls(&self) -> usize {
+ self.counters.encode_udf.load(Ordering::SeqCst)
+ }
+
+ fn decode_udf_calls(&self) -> usize {
+ self.counters.decode_udf.load(Ordering::SeqCst)
+ }
+
+ fn __datafusion_logical_extension_codec__<'py>(
Review Comment:
thanks for example
##########
examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py:
##########
@@ -667,3 +686,26 @@ def test_query_planner_rejects_invalid_config(max_rows:
str):
with pytest.raises(Exception, match=r"max_rows|Invalid value"):
ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect()
+
+
+def test_composed_codecs_with_query_planner():
+ """A second pair of codecs installed on top of the provider codecs
+ composes with them instead of replacing them. The extra codecs
+ (default-backed exports from a fresh session) decline everything,
+ so planner-driven encode/decode falls through to the provider
Review Comment:
decline everything but should be called? Am I correct?
##########
crates/core/src/codec.rs:
##########
@@ -223,12 +233,339 @@ fn strip_wire_header<'a>(
Ok(Some(&buf[py_minor_idx + 1..]))
}
+/// Family prefix for the envelope wrapping a chained codec's payload.
+///
+/// A distinct magic is what makes "is this framed?" a definite test
+/// rather than a speculative decode. Probing by attempting to parse the
+/// envelope would reintroduce exactly the protobuf ambiguity this
+/// framing exists to remove: prost skips unknown fields and defaults
+/// missing ones, so a foreign payload can parse cleanly as an envelope.
+pub(crate) const CHAINED_PAYLOAD_FAMILY: &[u8] = b"DFPYCHN";
+
+/// Wire-format version for the chained-payload envelope. Independent of
+/// [`WIRE_VERSION_CURRENT`], which versions the cloudpickle framing.
+pub(crate) const CHAIN_WIRE_VERSION_CURRENT: u8 = 1;
+
+/// Oldest chained-payload envelope version this build decodes.
+pub(crate) const CHAIN_WIRE_VERSION_MIN_SUPPORTED: u8 = 1;
+
+/// Prefix for the synthetic id given to a codec installed from a bare
+/// PyCapsule, which exposes nothing stable to derive an identity from.
+/// The rest of the id is random per install, so no other session can
+/// mint it: a payload carrying one decodes within the installing
+/// session's lineage, which clones the id along with the chain, and
+/// fails with a pointed error anywhere else rather than resolving to a
+/// different codec. A counter or a chain position would not do — every
+/// session numbers from the same end, so the first bare capsule
+/// installed anywhere would answer for every other session's first.
Review Comment:
its hard for me to make sense of this comment, it has quite a "legal tone"
for me (non english native)
##########
crates/core/src/codec.rs:
##########
@@ -321,30 +724,54 @@ impl LogicalExtensionCodec for PythonLogicalCodec {
node: Arc<dyn TableProvider>,
buf: &mut Vec<u8>,
) -> Result<()> {
- self.inner.try_encode_table_provider(table_ref, node, buf)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a table provider",
+ |codec, buf| codec.try_encode_table_provider(table_ref,
Arc::clone(&node), buf),
+ )
}
fn try_decode_file_format(
&self,
buf: &[u8],
ctx: &TaskContext,
) -> Result<Arc<dyn FileFormatFactory>> {
- self.inner.try_decode_file_format(buf, ctx)
+ chain_decode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a file format",
+ |codec, buf| codec.try_decode_file_format(buf, ctx),
+ )
}
fn try_encode_file_format(
&self,
buf: &mut Vec<u8>,
node: Arc<dyn FileFormatFactory>,
) -> Result<()> {
- self.inner.try_encode_file_format(buf, node)
+ chain_encode(
+ &self.chain,
+ &self.terminal,
+ buf,
+ "a file format",
+ |codec, buf| codec.try_encode_file_format(buf, Arc::clone(&node)),
+ )
}
fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec<u8>) ->
Result<()> {
if self.python_udf_inlining && try_encode_python_scalar_udf(node,
buf)? {
Review Comment:
i believe python udf inlining could be extracted to a separate codec, and
then codec could be added to chain of codecs, when inlining is enabled.
no need to change it now, just thinking aloud
##########
docs/source/contributor-guide/ffi.md:
##########
@@ -248,10 +248,96 @@ 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` appends the codec to the session's codec chain
+rather than replacing prior codecs.
+
+**Nothing is asked of the codec itself.** Implement `LogicalExtensionCodec` or
+`PhysicalExtensionCodec` exactly as you would for a session that installs only
+yours. When your codec writes bytes into a serialized plan, datafusion-python
+records which codec wrote them, and strips that record off again before
handing the
+bytes back. So your codec receives, byte for byte, the payload it wrote, and is
+never offered a payload another codec wrote.
+
+A codec that also ships to hosts which dispatch differently may still want its
own
+guard against foreign payloads. Keeping one is fine; it is simply not needed
for the
+datafusion-python path.
+
+That record is the codec's **id**: a short string stored inside the plan,
naming the
+codec that wrote each payload. Because plans are decoded in another process —
or
+another program — the id has to name the same codec there as it did where the
plan
+was written.
+
+Ids are assigned for you. A codec's id is normally its exporting class's import
+path, such as `my_library.Codec`, which is what you will see in
+`logical_extension_codec_ids()` and in decode errors. You choose one yourself
in
+three cases:
+
+- **Two instances of one class.** Both get the same id, so the second install
+ raises `ValueError`. Pass `codec_id=` to tell them apart.
+- **A bare `PyCapsule`.** A capsule has no class to take a name from, so it
gets an
+ id private to the session that installed it. Plans it encodes fail with a
clear
+ error on any other session, rather than being decoded by the wrong codec.
Pass
+ `codec_id=` if those plans have to cross sessions.
+- **A class you intend to rename.** The id follows the class name, so renaming
stops
+ older plans from decoding. Declare `__datafusion_codec_id__` on the exporting
+ object to pin an id that survives the rename.
+
+`SessionContext.logical_extension_codec_ids()` and its physical counterpart
list the
+ids installed on a session, which is also what a decode failure names.
+
+Installing one context's codec stack on another session composes the two
sessions
+rather than copying codecs out of one: the imported codecs resolve their task
context
+against the original and stop working when it is dropped — see
+[One session, one `Arc<SessionContext>`](#one-session-one-arcsessioncontext).
Pass
+the context itself rather than the capsule it exports, so its codecs get an id
that
+other sessions can decode.
+
+Because decoding keys off the id rather than install position, registration
order
+between independent libraries does not affect decoding at all. It is visible
only
+on encoding, where codecs are consulted in install order and the first to
claim an
+object wins — so installing a library can claim objects nothing else claimed,
but
+never takes over an object an earlier codec was already encoding. Two libraries
+that each own tables, functions, and a planner register like this:
+
+```python
+ctx = SessionContext(config)
+
+# Codecs from both libraries. Order between libraries does not matter.
Review Comment:
Most of the codecs do lot of `downcast_ref` which under to hood compares
`TypeId` before it tries to access structure properties. With chained encoders,
Is there a chance that we have `TypeId` collision, with two different
libraries have a different structures with same `TypeId`, triggering encoding
in the wrong encoder?
I guess even its possible it is low probability, this is a bit of
philosophical question
--
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]
