Encode MemTable as durable metadata in the FFI example logical codec

Author: s5dsn-eqee

docs/source/extension-guide/checklist.md

@@ -91,8 +91,9 @@ publish. Each links to the page that explains it.
       derived from it is in use. This is the rule most likely to arrive as a
       bug report against your library. → {ref}`extension_sessions`
 - [ ] **Your production codec serializes durable metadata**, not a
-      process-local token. The examples in this repository use tokens to make
-      ownership observable; that is a demonstration, not a pattern.
+      process-local token. The example logical codec in this repository does
+      this; the example physical codec uses a token to make ownership
+      observable, which is a demonstration, not a pattern.
       → {ref}`extension_codec_durable_metadata`
 - [ ] **You have integration tests across a real FFI boundary.** The two
       example crates in this repository are the pattern: build the cdylib,

docs/source/extension-guide/codecs.md

@@ -58,15 +58,20 @@ Your payload has to be enough to rebuild the object somewhere your process is
 not. Write the metadata a fresh instance can be constructed from — a path, a
 connection string, a schema, the options the object was created with.
 
-The example codecs in this repository do not do this, and it is worth knowing
-before copying them. They keep a process-local `HashMap` of live providers and
-encode an integer token into it: encoding inserts, decoding removes. That makes
-Rust type identity observable across three separately loaded libraries in one
-test, which is what the examples exist to show. It also means a decode consumes
-its token, so the same bytes cannot be decoded twice, one encoded plan cannot
-fan out to several readers, and a plan that never reaches a decoder keeps its
-provider alive for the life of the process. A real codec has none of those
-properties because it does not park the object anywhere.
+The logical codec in `examples/datafusion-ffi-example` is the pattern to copy.
+It encodes a `MemTable` as its schema and batches, one Arrow IPC stream per
+partition, and decodes by building a new `MemTable` from those streams. Nothing
+is kept between encode and decode, so the same bytes decode any number of times
+and an encoded plan can fan out to several readers.
+
+The physical codec in the same example does not do this, and it is worth
+knowing before copying it. It keeps a process-local `HashMap` of live execution
+plans and encodes an integer token into it: encoding inserts, decoding removes.
+That makes Rust type identity observable across three separately loaded
+libraries in one test, which is what it exists to show. It also means a decode
+consumes its token, so the same bytes cannot be decoded twice, and a plan that
+never reaches a decoder stays alive for the life of the process. A real codec
+has none of those properties because it does not park the object anywhere.
 
 (extension_codec_ids)=
 

examples/datafusion-ffi-example/README.md

@@ -31,7 +31,9 @@ Separate shared libraries guarantee distinct DataFusion library markers. This ca
 
 ## Codec behavior
 
-`MyLogicalExtensionCodec` serializes this example's in-memory table providers, and `MyPhysicalExtensionCodec` serializes provider-owned memory scans and opaque FFI wrappers around them. Both use documented, process-local, one-shot token registries. The registries make ownership and callback routing visible without pretending to be a portable format. They assume trusted in-process payloads and consume each token during decoding. A production provider should instead encode durable metadata from which its provider and plans can be reconstructed.
+`MyLogicalExtensionCodec` serializes this example's in-memory table providers as durable metadata: a `MemTable` is written as its schema and batches, one Arrow IPC stream per partition, and decoding builds a new `MemTable` from those streams. Nothing is retained between encode and decode, so one encoded plan can be decoded any number of times, in any process. This is the pattern a production provider should follow.
+
+`MyPhysicalExtensionCodec` serializes provider-owned memory scans and opaque FFI wrappers around them through a documented, process-local, one-shot token registry. The registry makes ownership and callback routing visible without pretending to be a portable format. It assumes trusted in-process payloads and consumes each token during decoding.
 
 Both codec getters take the `SessionContext` they are being installed on and pull the `TaskContextProvider` off it, so decode callbacks resolve session configuration and registered functions against the session that is running the query. Passing `require_udf_on_decode` to either constructor makes every decode call resolve a named scalar function out of that context, which is how the tests check where the registry came from.
 

examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py

@@ -203,12 +203,13 @@ def test_installing_a_codec_cannot_hijack_an_earlier_codecs_objects():
     ctx = ctx.with_logical_extension_codec(later, codec_id="TOKENBBB")
     after = ctx.sql('SELECT "A" FROM numbers').logical_plan().to_bytes(ctx)
 
-    # Provider tokens are minted per encode, so the payloads differ in the
-    # token id. What must not change is which codec claimed the provider.
+    # The provider is encoded as its schema and batches, so re-encoding the
+    # same table yields the same bytes. Which codec claimed it must not
+    # change either.
     assert b"TOKENAAA" in after
     assert b"TOKENBBB" not in after
     assert later.table_provider_encode_calls() == 0
-    assert len(before) == len(after)
+    assert before == after
 
 
 def test_decode_dispatches_to_the_codec_that_encoded():
@@ -233,6 +234,39 @@ def test_decode_dispatches_to_the_codec_that_encoded():
     assert second.table_provider_decode_calls() == 0
 
 
+def test_one_encoded_plan_decodes_more_than_once():
+    """A table provider payload is durable metadata -- the table's schema
+    and batches -- not a handle into the encoding process, so the same
+    bytes can be decoded again and again, on the session that wrote them
+    or on one that never saw the original provider.
+
+    Every decode rebuilds an equivalent table: the rows come back, and
+    they are the rows the provider was created with.
+    """
+    blob, owner = _encode_provider_plan("TOKENAAA")
+    expected = [[0, 1, 2, 3]]
+
+    def rows(ctx: SessionContext) -> list[list[int]]:
+        restored = LogicalPlan.from_bytes(ctx, blob)
+        batches = ctx.create_dataframe_from_logical_plan(restored).collect()
+        return [batch.column(0).to_pylist() for batch in batches]
+
+    same_session = SessionContext().with_logical_extension_codec(
+        owner, codec_id="TOKENAAA"
+    )
+    assert rows(same_session) == expected
+    assert rows(same_session) == expected
+
+    # A fresh session with a fresh codec instance has no access to anything
+    # the encoding side might have kept; the bytes alone must suffice.
+    elsewhere = SessionContext().with_logical_extension_codec(
+        MyLogicalExtensionCodec(provider_prefix="TOKENAAA"), codec_id="TOKENAAA"
+    )
+    assert rows(elsewhere) == expected
+
+    assert owner.table_provider_decode_calls() == 2
+
+
 def test_decode_survives_a_different_install_order():
     """Dispatch keys off codec identity, not chain position, so the
     decoding session may install the same codecs in any order.

examples/datafusion-ffi-example/src/logical_extension_codec.rs

@@ -15,12 +15,15 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::collections::HashMap;
 use std::fmt;
-use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
-use std::sync::{Arc, Mutex, OnceLock};
+use std::io::Cursor;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
 
 use arrow::datatypes::SchemaRef;
+use arrow::ipc::reader::StreamReader;
+use arrow::ipc::writer::StreamWriter;
+use arrow::record_batch::RecordBatch;
 use datafusion::catalog::MemTable;
 use datafusion::common::{DataFusionError, Result, TableReference};
 use datafusion::datasource::TableProvider;
@@ -34,35 +37,126 @@ use pyo3::types::PyCapsule;
 
 use crate::required_udf::{TaskContextProbe, resolve_required_udf};
 
-const TABLE_PROVIDER_TOKEN: &[u8] = b"DFPYEXTP";
-static NEXT_TABLE_PROVIDER_ID: AtomicU64 = AtomicU64::new(1);
-static TABLE_PROVIDERS: OnceLock<Mutex<HashMap<u64, Arc<dyn TableProvider>>>> = OnceLock::new();
+/// Default byte prefix stamped on every table provider this codec encodes.
+const TABLE_PROVIDER_PREFIX: &[u8] = b"DFPYEXTP";
 
-/// Hands a provider to another library in this process by token.
+/// Format tag that follows the prefix. Bump it if the layout below changes.
+const MEM_TABLE_FORMAT: &[u8] = b"MEMTBL1";
+
+/// Write a [`MemTable`] as durable metadata: its schema and every batch of
+/// every partition, so that a decoder anywhere can rebuild an equivalent
+/// table from the bytes alone.
 ///
-/// Encoding inserts, decoding removes. Two consequences worth knowing before
-/// copying this:
+/// Layout, after the caller's provider prefix:
 ///
-/// - **Decode consumes the token.** Decoding the same encoded bytes twice
-///   fails the second time with `Unknown ... table provider token`. That is
-///   fine here because every plan is encoded immediately before the single
-///   decode that consumes it, but it rules out anything that replays a stored
-///   plan, retries a decode, or fans one encoded plan out to several readers.
-/// - **An encode that is never decoded leaks.** Nothing expires entries, so a
-///   plan that fails to reach its decoder keeps its provider alive for the
-///   life of the process.
+/// ```text
+/// b"MEMTBL1" | u32 LE n_partitions | { u32 LE ipc_len | ipc stream }*
+/// ```
 ///
-/// Both are acceptable for an example whose job is to show that Rust type
-/// identity survives a trip through two other libraries. Neither is acceptable
-/// in a real codec, which should encode metadata sufficient to rebuild the
-/// provider rather than parking the object here.
-fn table_providers() -> &'static Mutex<HashMap<u64, Arc<dyn TableProvider>>> {
-    TABLE_PROVIDERS.get_or_init(|| Mutex::new(HashMap::new()))
+/// Each partition is one Arrow IPC stream. The stream carries the schema, so
+/// the decoder never has to trust a schema handed to it out of band.
+fn encode_mem_table(table: &MemTable, buf: &mut Vec<u8>) -> Result<()> {
+    let schema = table.schema();
+    buf.extend_from_slice(MEM_TABLE_FORMAT);
+    buf.extend_from_slice(&length_prefix(table.batches.len())?);
+
+    for partition in &table.batches {
+        // `MemTable` guards each partition with a tokio `RwLock`. This encode
+        // runs on a tokio worker thread, where `blocking_read` panics, so
+        // take the lock only if it is free. A partition that is mid-insert
+        // is reported rather than waited for.
+        let batches = partition.try_read().map_err(|_| {
+            DataFusionError::Internal(
+                "datafusion-ffi-example cannot encode a MemTable while a partition is locked"
+                    .to_string(),
+            )
+        })?;
+
+        let mut ipc = Vec::new();
+        let mut writer = StreamWriter::try_new(&mut ipc, schema.as_ref())?;
+        for batch in batches.iter() {
+            writer.write(batch)?;
+        }
+        writer.finish()?;
+        drop(writer);
+
+        buf.extend_from_slice(&length_prefix(ipc.len())?);
+        buf.extend_from_slice(&ipc);
+    }
+    Ok(())
 }
 
-fn token_id(buf: &[u8], prefix: &[u8]) -> Option<u64> {
-    let id: [u8; 8] = buf.strip_prefix(prefix)?.try_into().ok()?;
-    Some(u64::from_le_bytes(id))
+/// Rebuild a [`MemTable`] from bytes written by [`encode_mem_table`].
+///
+/// The table's schema is the one carried inside the IPC streams, not the
+/// `schema` argument DataFusion passes to `try_decode_table_provider`. A
+/// payload that does not describe itself consistently is rejected here
+/// instead of producing a table whose batches disagree with its schema.
+fn decode_mem_table(payload: &[u8]) -> Result<MemTable> {
+    let mut rest = payload.strip_prefix(MEM_TABLE_FORMAT).ok_or_else(|| {
+        DataFusionError::Internal(
+            "datafusion-ffi-example table provider payload has an unknown format tag".to_string(),
+        )
+    })?;
+
+    let n_partitions = read_length_prefix(&mut rest)?;
+    let mut schema: Option<SchemaRef> = None;
+    let mut partitions: Vec<Vec<RecordBatch>> = Vec::with_capacity(n_partitions);
+
+    for _ in 0..n_partitions {
+        let ipc_len = read_length_prefix(&mut rest)?;
+        if rest.len() < ipc_len {
+            return Err(DataFusionError::Internal(
+                "datafusion-ffi-example table provider payload is truncated".to_string(),
+            ));
+        }
+        let (ipc, tail) = rest.split_at(ipc_len);
+        rest = tail;
+
+        let reader = StreamReader::try_new(Cursor::new(ipc), None)?;
+        let ipc_schema = reader.schema();
+        match &schema {
+            None => schema = Some(ipc_schema),
+            Some(first) if *first != ipc_schema => {
+                return Err(DataFusionError::Internal(
+                    "datafusion-ffi-example table provider partitions disagree on schema"
+                        .to_string(),
+                ));
+            }
+            Some(_) => {}
+        }
+        partitions.push(reader.collect::<std::result::Result<Vec<_>, _>>()?);
+    }
+
+    if !rest.is_empty() {
+        return Err(DataFusionError::Internal(
+            "datafusion-ffi-example table provider payload has trailing bytes".to_string(),
+        ));
+    }
+
+    let schema = schema.ok_or_else(|| {
+        DataFusionError::Internal(
+            "datafusion-ffi-example table provider payload has no partitions".to_string(),
+        )
+    })?;
+    MemTable::try_new(schema, partitions)
+}
+
+fn length_prefix(len: usize) -> Result<[u8; 4]> {
+    u32::try_from(len)
+        .map(u32::to_le_bytes)
+        .map_err(|_| DataFusionError::Internal(format!("length {len} does not fit in u32")))
+}
+
+fn read_length_prefix(rest: &mut &[u8]) -> Result<usize> {
+    let (head, tail) = rest.split_at_checked(4).ok_or_else(|| {
+        DataFusionError::Internal(
+            "datafusion-ffi-example table provider payload is truncated".to_string(),
+        )
+    })?;
+    *rest = tail;
+    let bytes: [u8; 4] = head.try_into().expect("split_at_checked returned 4 bytes");
+    Ok(u32::from_le_bytes(bytes) as usize)
 }
 
 #[derive(Debug, Default)]
@@ -76,23 +170,20 @@ pub(crate) struct CallCounters {
 
 /// Example codec for objects owned by this extension library.
 ///
-/// The table-provider token registry is intentionally process-local. It is a compact
-/// example of preserving Rust type identity across three loaded libraries, not a
-/// network serialization format. Production libraries should encode reconstructible
-/// provider metadata rather than retaining objects in a global registry.
-///
-/// See [`table_providers`] for the token lifecycle, which is narrower than it
-/// looks: a decode consumes its token, so the same encoded plan cannot be
-/// decoded twice.
+/// Table providers are encoded as durable metadata, see [`encode_mem_table`].
+/// Nothing is retained between encode and decode, so the same bytes decode
+/// any number of times, in any process, and an encoded plan that never
+/// reaches a decoder costs nothing.
 struct CountingLogicalExtensionCodec {
     inner: DefaultLogicalExtensionCodec,
     counters: Arc<CallCounters>,
     /// Scalar function every table-provider decode must resolve from the
     /// `TaskContext` it is handed. See [`crate::required_udf`].
     required_udf: Option<String>,
-    /// Byte prefix identifying providers this codec owns. Distinct tokens let a
-    /// test install several instances and observe which one the chain picks.
-    token: Arc<[u8]>,
+    /// Byte prefix identifying providers this codec owns. Distinct prefixes
+    /// let a test install several instances and observe which one the chain
+    /// picks.
+    provider_prefix: Arc<[u8]>,
 }
 
 impl fmt::Debug for CountingLogicalExtensionCodec {
@@ -127,19 +218,11 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec {
         ctx: &TaskContext,
     ) -> Result<Arc<dyn TableProvider>> {
         resolve_required_udf(self.required_udf.as_deref(), ctx, &self.counters.task_ctx)?;
-        if let Some(id) = token_id(buf, &self.token) {
+        if let Some(payload) = buf.strip_prefix(self.provider_prefix.as_ref()) {
             self.counters
                 .decode_table_provider
                 .fetch_add(1, Ordering::SeqCst);
-            return table_providers()
-                .lock()
-                .map_err(|err| DataFusionError::Internal(err.to_string()))?
-                .remove(&id)
-                .ok_or_else(|| {
-                    DataFusionError::Internal(format!(
-                        "Unknown datafusion-ffi-example table provider token {id}"
-                    ))
-                });
+            return Ok(Arc::new(decode_mem_table(payload)?));
         }
         self.inner
             .try_decode_table_provider(buf, table_ref, schema, ctx)
@@ -151,18 +234,12 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec {
         node: Arc<dyn TableProvider>,
         buf: &mut Vec<u8>,
     ) -> Result<()> {
-        if node.downcast_ref::<MemTable>().is_some() {
+        if let Some(table) = node.downcast_ref::<MemTable>() {
             self.counters
                 .encode_table_provider
                 .fetch_add(1, Ordering::SeqCst);
-            let id = NEXT_TABLE_PROVIDER_ID.fetch_add(1, Ordering::SeqCst);
-            table_providers()
-                .lock()
-                .map_err(|err| DataFusionError::Internal(err.to_string()))?
-                .insert(id, node);
-            buf.extend_from_slice(&self.token);
-            buf.extend_from_slice(&id.to_le_bytes());
-            return Ok(());
+            buf.extend_from_slice(&self.provider_prefix);
+            return encode_mem_table(table, buf);
         }
         self.inner.try_encode_table_provider(table_ref, node, buf)
     }
@@ -188,7 +265,7 @@ impl LogicalExtensionCodec for CountingLogicalExtensionCodec {
 pub(crate) struct MyLogicalExtensionCodec {
     counters: Arc<CallCounters>,
     required_udf: Option<String>,
-    token: Arc<[u8]>,
+    provider_prefix: Arc<[u8]>,
 }
 
 #[pymethods]
@@ -200,7 +277,7 @@ impl MyLogicalExtensionCodec {
     /// unset for the ordinary behaviour; set it to observe *which* session's
     /// registry the FFI decode callback actually receives.
     ///
-    /// `provider_prefix` overrides [`TABLE_PROVIDER_TOKEN`], the byte prefix
+    /// `provider_prefix` overrides [`TABLE_PROVIDER_PREFIX`], the byte prefix
     /// stamped on encoded table providers. Two instances built with different
     /// prefixes each own a disjoint slice of the wire format, which is what
     /// lets a test install both and tell from the decoded bytes which one the
@@ -211,8 +288,8 @@ impl MyLogicalExtensionCodec {
         Self {
             counters: Arc::new(CallCounters::default()),
             required_udf: require_udf_on_decode,
-            token: provider_prefix.map_or_else(
-                || Arc::from(TABLE_PROVIDER_TOKEN),
+            provider_prefix: provider_prefix.map_or_else(
+                || Arc::from(TABLE_PROVIDER_PREFIX),
                 |prefix| Arc::from(prefix.as_bytes()),
             ),
         }
@@ -259,7 +336,7 @@ impl MyLogicalExtensionCodec {
             inner: DefaultLogicalExtensionCodec {},
             counters: Arc::clone(&self.counters),
             required_udf: self.required_udf.clone(),
-            token: Arc::clone(&self.token),
+            provider_prefix: Arc::clone(&self.provider_prefix),
         });
 
         let runtime = get_tokio_runtime().handle().clone();