diff --git a/docs/source/extension-guide/checklist.md b/docs/source/extension-guide/checklist.md index 6e204b3f5..832c97fb4 100644 --- a/docs/source/extension-guide/checklist.md +++ b/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, diff --git a/docs/source/extension-guide/codecs.md b/docs/source/extension-guide/codecs.md index 160439927..9bf1edb02 100644 --- a/docs/source/extension-guide/codecs.md +++ b/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)= diff --git a/examples/datafusion-ffi-example/README.md b/examples/datafusion-ffi-example/README.md index aadea909f..5d23dc417 100644 --- a/examples/datafusion-ffi-example/README.md +++ b/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. diff --git a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py b/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py index ac4c69d14..452a64b1d 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_logical_extension_codec.py +++ b/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. diff --git a/examples/datafusion-ffi-example/src/logical_extension_codec.rs b/examples/datafusion-ffi-example/src/logical_extension_codec.rs index 5660489d4..e38773441 100644 --- a/examples/datafusion-ffi-example/src/logical_extension_codec.rs +++ b/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>>> = 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>> { - 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) -> 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 { - 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 { + 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 = None; + let mut partitions: Vec> = 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::, _>>()?); + } + + 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 { + 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, /// Scalar function every table-provider decode must resolve from the /// `TaskContext` it is handed. See [`crate::required_udf`]. required_udf: Option, - /// 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> { 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, buf: &mut Vec, ) -> Result<()> { - if node.downcast_ref::().is_some() { + if let Some(table) = node.downcast_ref::() { 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, required_udf: Option, - 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(); diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index 810e090de..fc208b5dd 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -40,11 +40,14 @@ const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); static EXECUTION_PLANS: OnceLock>>> = OnceLock::new(); -/// Execution-plan counterpart of the logical codec's provider registry, with -/// the same lifecycle: encoding inserts, decoding removes, so a decode -/// consumes its token and an encode that is never decoded leaks. See -/// [`crate::logical_extension_codec`] for why that is acceptable here and not -/// in a real codec. +/// Hands an execution plan to another library in this process by token. +/// +/// Encoding inserts, decoding removes, so a decode consumes its token and an +/// encode that is never decoded leaks. That is acceptable for an example whose +/// job is to show that Rust type identity survives a trip through two other +/// libraries; a real codec should encode metadata sufficient to rebuild the +/// plan instead, the way [`crate::logical_extension_codec`] does for +/// `MemTable`. fn execution_plans() -> &'static Mutex>> { EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new())) }