datafusion-examples/examples/custom_data_source/adapter_serialization.rs

// 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.

//! See `main.rs` for how to run it.
//!
//! This example demonstrates how to use the `PhysicalProtoConverterExtension`
//! trait's interception methods (`execution_plan_to_proto` and
//! `proto_to_execution_plan`) to implement custom serialization logic.
//!
//! The key insight is that a non-default `FileScanConfig::expr_adapter_factory`
//! cannot be serialized by the default protobuf path. This example shows how to:
//! 1. Detect plans with custom adapters during serialization
//! 2. Wrap them as Extension nodes with JSON-serialized adapter metadata
//! 3. Explicitly remove the captured adapter from the inner `DataSourceExec`
//! 4. Unwrap and restore the adapter during deserialization
//!
//! This demonstrates nested serialization (protobuf outer, JSON inner) and the
//! power of `PhysicalProtoConverterExtension`. Both plan and expression
//! serialization route through converter hooks, enabling interception at every
//! node in the tree.

use std::fmt::Debug;
use std::sync::Arc;

use arrow::array::record_batch;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion::assert_batches_eq;
use datafusion::common::{Result, not_impl_err};
use datafusion::datasource::listing::{
    ListingTable, ListingTableConfig, ListingTableConfigExt, ListingTableUrl,
};
use datafusion::datasource::physical_plan::{FileScanConfig, FileScanConfigBuilder};
use datafusion::datasource::source::DataSourceExec;
use datafusion::execution::TaskContext;
use datafusion::execution::context::SessionContext;
use datafusion::execution::object_store::ObjectStoreUrl;
use datafusion::parquet::arrow::ArrowWriter;
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_expr_adapter::{
    DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory,
};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::SessionConfig;
use datafusion_proto::bytes::{
    physical_plan_from_bytes_with_proto_converter,
    physical_plan_to_bytes_with_proto_converter,
};
use datafusion_proto::physical_plan::from_proto::parse_physical_expr_with_converter;
use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter;
use datafusion_proto::physical_plan::{
    PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalPlanNodeExt,
    PhysicalProtoConverterExtension,
};
use datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType;
use datafusion_proto::protobuf::{
    PhysicalExprNode, PhysicalExtensionNode, PhysicalPlanNode,
};
use object_store::memory::InMemory;
use object_store::path::Path;
use object_store::{ObjectStore, ObjectStoreExt, PutPayload};
use serde::{Deserialize, Serialize};

/// Example showing how to preserve custom adapter information during plan serialization.
///
/// This demonstrates:
/// 1. Creating a custom PhysicalExprAdapter with metadata
/// 2. Using PhysicalExtensionCodec to intercept serialization
/// 3. Wrapping adapter info as Extension nodes
/// 4. Restoring adapters during deserialization
pub async fn adapter_serialization() -> Result<()> {
    println!("=== PhysicalExprAdapter Serialization Example ===\n");

    // Step 1: Create sample Parquet data in memory
    println!("Step 1: Creating sample Parquet data...");
    let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
    let batch = record_batch!(("id", Int32, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))?;
    let path = Path::from("data.parquet");
    write_parquet(&store, &path, &batch).await?;

    // Step 2: Set up session with custom adapter
    println!("Step 2: Setting up session with custom adapter...");
    let logical_schema =
        Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));

    let mut cfg = SessionConfig::new();
    cfg.options_mut().execution.parquet.pushdown_filters = true;
    let ctx = SessionContext::new_with_config(cfg);
    ctx.runtime_env().register_object_store(
        ObjectStoreUrl::parse("memory://")?.as_ref(),
        Arc::clone(&store),
    );

    // Create a table with our custom MetadataAdapterFactory
    let adapter_factory = Arc::new(MetadataAdapterFactory::new("v1"));
    let listing_config =
        ListingTableConfig::new(ListingTableUrl::parse("memory:///data.parquet")?)
            .infer_options(&ctx.state())
            .await?
            .with_schema(logical_schema)
            .with_expr_adapter_factory(
                Arc::clone(&adapter_factory) as Arc<dyn PhysicalExprAdapterFactory>
            );
    let table = ListingTable::try_new(listing_config)?;
    ctx.register_table("my_table", Arc::new(table))?;

    // Step 3: Create physical plan with filter
    println!("Step 3: Creating physical plan with filter...");
    let df = ctx.sql("SELECT * FROM my_table WHERE id > 5").await?;
    let original_plan = df.create_physical_plan().await?;

    // Verify adapter is present in original plan
    let has_adapter_before = verify_adapter_in_plan(&original_plan, "original");
    println!("  Original plan has adapter: {has_adapter_before}");

    // Step 4: Serialize with our custom codec
    println!("\nStep 4: Serializing plan with AdapterPreservingCodec...");
    let codec = AdapterPreservingCodec;
    let bytes = physical_plan_to_bytes_with_proto_converter(
        Arc::clone(&original_plan),
        &codec,
        &codec,
    )?;
    println!("  Serialized {} bytes", bytes.len());
    println!("  (DataSourceExec with adapter was wrapped as PhysicalExtensionNode)");

    // Step 5: Deserialize with our custom codec
    println!("\nStep 5: Deserializing plan with AdapterPreservingCodec...");
    let task_ctx = ctx.task_ctx();
    let restored_plan =
        physical_plan_from_bytes_with_proto_converter(&bytes, &task_ctx, &codec, &codec)?;

    // Verify adapter is restored
    let has_adapter_after = verify_adapter_in_plan(&restored_plan, "restored");
    println!("  Restored plan has adapter: {has_adapter_after}");

    // Step 6: Execute and compare results
    println!("\nStep 6: Executing plans and comparing results...");
    let original_results =
        datafusion::physical_plan::collect(Arc::clone(&original_plan), task_ctx.clone())
            .await?;
    let restored_results =
        datafusion::physical_plan::collect(restored_plan, task_ctx).await?;

    #[rustfmt::skip]
    let expected = [
        "+----+",
        "| id |",
        "+----+",
        "| 6  |",
        "| 7  |",
        "| 8  |",
        "| 9  |",
        "| 10 |",
        "+----+",
    ];

    println!("\n  Original plan results:");
    arrow::util::pretty::print_batches(&original_results)?;
    assert_batches_eq!(expected, &original_results);

    println!("\n  Restored plan results:");
    arrow::util::pretty::print_batches(&restored_results)?;
    assert_batches_eq!(expected, &restored_results);

    println!("\n=== Example Complete! ===");
    println!("Key takeaways:");
    println!(
        "  1. PhysicalProtoConverterExtension provides execution_plan_to_proto/proto_to_execution_plan hooks"
    );
    println!("  2. Custom metadata can be wrapped as PhysicalExtensionNode");
    println!("  3. Nested serialization (protobuf + JSON) works seamlessly");
    println!(
        "  4. Both plans produce identical results despite serialization round-trip"
    );
    println!("  5. Adapters are fully preserved through the serialization round-trip");

    Ok(())
}

// ============================================================================
// MetadataAdapter - A simple custom adapter with a tag
// ============================================================================

/// A custom PhysicalExprAdapter that wraps another adapter.
/// The tag metadata is stored in the factory, not the adapter itself.
#[derive(Debug)]
struct MetadataAdapter {
    inner: Arc<dyn PhysicalExprAdapter>,
}

impl PhysicalExprAdapter for MetadataAdapter {
    fn rewrite(&self, expr: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
        // Simply delegate to inner adapter
        self.inner.rewrite(expr)
    }
}

// ============================================================================
// MetadataAdapterFactory - Factory for creating MetadataAdapter instances
// ============================================================================

/// Factory for creating MetadataAdapter instances.
/// The tag is stored in the factory and extracted via Debug formatting in `extract_adapter_tag`.
#[derive(Debug)]
struct MetadataAdapterFactory {
    // Note: This field is read via Debug formatting in `extract_adapter_tag`.
    // Rust's dead code analysis doesn't recognize Debug-based field access.
    // In PR #19234, this field is used by `with_partition_values`, but that method
    // doesn't exist in upstream DataFusion's PhysicalExprAdapter trait.
    #[expect(dead_code)]
    tag: String,
}

impl MetadataAdapterFactory {
    fn new(tag: impl Into<String>) -> Self {
        Self { tag: tag.into() }
    }
}

impl PhysicalExprAdapterFactory for MetadataAdapterFactory {
    fn create(
        &self,
        logical_file_schema: SchemaRef,
        physical_file_schema: SchemaRef,
    ) -> Result<Arc<dyn PhysicalExprAdapter>> {
        let inner = DefaultPhysicalExprAdapterFactory
            .create(logical_file_schema, physical_file_schema)?;
        Ok(Arc::new(MetadataAdapter { inner }))
    }
}

// ============================================================================
// AdapterPreservingCodec - Custom codec that preserves adapters
// ============================================================================

/// Extension payload structure for serializing adapter info
#[derive(Serialize, Deserialize)]
struct ExtensionPayload {
    /// Marker to identify this is our custom extension
    marker: String,
    /// JSON-serialized adapter metadata
    adapter_metadata: AdapterMetadata,
}

/// Metadata about the adapter to recreate it during deserialization
#[derive(Serialize, Deserialize)]
struct AdapterMetadata {
    /// The adapter tag (e.g., "v1")
    tag: String,
}

const EXTENSION_MARKER: &str = "adapter_preserving_extension_v1";

/// A codec that intercepts serialization to preserve adapter information.
#[derive(Debug)]
struct AdapterPreservingCodec;

impl PhysicalExtensionCodec for AdapterPreservingCodec {
    // Required method: decode custom extension nodes
    fn try_decode(
        &self,
        buf: &[u8],
        inputs: &[Arc<dyn ExecutionPlan>],
        _ctx: &TaskContext,
        _proto_converter: &dyn PhysicalProtoConverterExtension,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        // Try to parse as our extension payload
        if let Ok(payload) = serde_json::from_slice::<ExtensionPayload>(buf)
            && payload.marker == EXTENSION_MARKER
        {
            if inputs.len() != 1 {
                return Err(datafusion::error::DataFusionError::Plan(format!(
                    "Extension node expected exactly 1 child, got {}",
                    inputs.len()
                )));
            }
            let inner_plan = inputs[0].clone();

            // Recreate the adapter factory
            let adapter_factory = create_adapter_factory(&payload.adapter_metadata.tag);

            // Inject adapter into the plan
            return inject_adapter_into_plan(inner_plan, adapter_factory);
        }

        not_impl_err!("Unknown extension type")
    }

    // Required method: encode custom execution plans
    fn try_encode(
        &self,
        _node: Arc<dyn ExecutionPlan>,
        _buf: &mut Vec<u8>,
        _proto_converter: &dyn PhysicalProtoConverterExtension,
    ) -> Result<()> {
        // We don't need this for the example - adapter wrapping happens in
        // `execution_plan_to_proto` instead.
        not_impl_err!(
            "try_encode not used - adapter wrapping happens in execution_plan_to_proto"
        )
    }
}

impl PhysicalProtoConverterExtension for AdapterPreservingCodec {
    fn execution_plan_to_proto(
        &self,
        plan: &Arc<dyn ExecutionPlan>,
        extension_codec: &dyn PhysicalExtensionCodec,
    ) -> Result<PhysicalPlanNode> {
        // Check if this is a DataSourceExec with adapter
        if let Some(exec) = plan.downcast_ref::<DataSourceExec>()
            && let Some(config) = exec.data_source().downcast_ref::<FileScanConfig>()
            && let Some(adapter_factory) = &config.expr_adapter_factory
            && let Some(tag) = extract_adapter_tag(adapter_factory.as_ref())
        {
            // Try to extract our MetadataAdapterFactory's tag
            println!("    [Serialize] Found DataSourceExec with adapter tag: {tag}");

            // 1. Create adapter metadata
            let adapter_metadata = AdapterMetadata { tag };

            // 2. The metadata is captured above, so explicitly remove the custom
            //    adapter before using the guarded default serializer for the child.
            let mut inner_config = config.clone();
            inner_config.expr_adapter_factory = None;
            let inner_plan: Arc<dyn ExecutionPlan> =
                DataSourceExec::from_data_source(inner_config);
            let inner_proto = PhysicalPlanNode::try_from_physical_plan_with_converter(
                inner_plan,
                extension_codec,
                self,
            )?;

            // 3. Create extension payload to wrap the plan
            //    so that the custom adapter gets re-attached during deserialization
            //    The choice of JSON is arbitrary; other formats could be used.
            let payload = ExtensionPayload {
                marker: EXTENSION_MARKER.to_string(),
                adapter_metadata,
            };
            let payload_bytes = serde_json::to_vec(&payload).map_err(|e| {
                datafusion::error::DataFusionError::Plan(format!(
                    "Failed to serialize payload: {e}"
                ))
            })?;

            // 4. Return as PhysicalExtensionNode with child plan in inputs
            return Ok(PhysicalPlanNode {
                physical_plan_type: Some(PhysicalPlanType::Extension(
                    PhysicalExtensionNode {
                        node: payload_bytes,
                        inputs: vec![inner_proto],
                        // This converter recognizes its own nodes by the marker
                        // inside the payload, so it needs no `plan_name`
                        // discriminator. Plans that decode through the session
                        // registry set it to their `PLAN_NAME` instead.
                        plan_name: None,
                    },
                )),
            });
        }

        // No adapter found, not a DataSourceExec, etc. - use default serialization
        PhysicalPlanNode::try_from_physical_plan_with_converter(
            Arc::clone(plan),
            extension_codec,
            self,
        )
    }

    // Interception point: override deserialization to unwrap adapters
    fn proto_to_execution_plan(
        &self,
        proto: &PhysicalPlanNode,
        ctx: &PhysicalPlanDecodeContext<'_>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        // Check if this is our custom extension wrapper
        if let Some(PhysicalPlanType::Extension(extension)) = &proto.physical_plan_type
            && let Ok(payload) =
                serde_json::from_slice::<ExtensionPayload>(&extension.node)
            && payload.marker == EXTENSION_MARKER
        {
            println!(
                "    [Deserialize] Found adapter extension with tag: {}",
                payload.adapter_metadata.tag
            );

            // Get the inner plan proto from inputs field
            if extension.inputs.is_empty() {
                return Err(datafusion::error::DataFusionError::Plan(
                    "Extension node missing child plan in inputs".to_string(),
                ));
            }
            let inner_proto = &extension.inputs[0];

            // Deserialize the inner plan
            let inner_plan = self.default_proto_to_execution_plan(inner_proto, ctx)?;

            // Recreate the adapter factory
            let adapter_factory = create_adapter_factory(&payload.adapter_metadata.tag);

            // Inject adapter into the plan
            return inject_adapter_into_plan(inner_plan, adapter_factory);
        }

        // Not our extension - use default deserialization
        self.default_proto_to_execution_plan(proto, ctx)
    }

    fn proto_to_physical_expr(
        &self,
        proto: &PhysicalExprNode,
        input_schema: &Schema,
        ctx: &PhysicalPlanDecodeContext<'_>,
    ) -> Result<Arc<dyn PhysicalExpr>> {
        parse_physical_expr_with_converter(proto, input_schema, ctx, self)
    }

    fn physical_expr_to_proto(
        &self,
        expr: &Arc<dyn PhysicalExpr>,
        codec: &dyn PhysicalExtensionCodec,
    ) -> Result<PhysicalExprNode> {
        serialize_physical_expr_with_converter(expr, codec, self)
    }
}

// ============================================================================
// Helper functions
// ============================================================================

/// Write a RecordBatch to Parquet in the object store
async fn write_parquet(
    store: &dyn ObjectStore,
    path: &Path,
    batch: &arrow::record_batch::RecordBatch,
) -> Result<()> {
    let mut buf = vec![];
    let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None)?;
    writer.write(batch)?;
    writer.close()?;

    let payload = PutPayload::from_bytes(buf.into());
    store.put(path, payload).await?;
    Ok(())
}

/// Extract the tag from a MetadataAdapterFactory.
///
/// Note: Since `PhysicalExprAdapterFactory` doesn't provide `as_any()` for downcasting,
/// we parse the Debug output. In a production system, you might add a dedicated trait
/// method for metadata extraction.
fn extract_adapter_tag(factory: &dyn PhysicalExprAdapterFactory) -> Option<String> {
    let debug_str = format!("{factory:?}");
    if debug_str.contains("MetadataAdapterFactory") {
        // Extract tag from debug output: MetadataAdapterFactory { tag: "v1" }
        if let Some(start) = debug_str.find("tag: \"") {
            let after_tag = &debug_str[start + 6..];
            if let Some(end) = after_tag.find('"') {
                return Some(after_tag[..end].to_string());
            }
        }
    }
    None
}

/// Create an adapter factory from a tag
fn create_adapter_factory(tag: &str) -> Arc<dyn PhysicalExprAdapterFactory> {
    Arc::new(MetadataAdapterFactory::new(tag))
}

/// Inject an adapter into a plan (assumes plan is a DataSourceExec with FileScanConfig)
fn inject_adapter_into_plan(
    plan: Arc<dyn ExecutionPlan>,
    adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
) -> Result<Arc<dyn ExecutionPlan>> {
    if let Some(exec) = plan.downcast_ref::<DataSourceExec>()
        && let Some(config) = exec.data_source().downcast_ref::<FileScanConfig>()
    {
        let new_config = FileScanConfigBuilder::from(config.clone())
            .with_expr_adapter(Some(adapter_factory))
            .build();
        return Ok(DataSourceExec::from_data_source(new_config));
    }
    // If not a DataSourceExec with FileScanConfig, return as-is
    Ok(plan)
}

/// Helper to verify if a plan has an adapter (for testing/validation)
fn verify_adapter_in_plan(plan: &Arc<dyn ExecutionPlan>, label: &str) -> bool {
    // Walk the plan tree to find DataSourceExec with adapter
    fn check_plan(plan: &dyn ExecutionPlan) -> bool {
        if let Some(exec) = plan.downcast_ref::<DataSourceExec>()
            && let Some(config) = exec.data_source().downcast_ref::<FileScanConfig>()
            && config.expr_adapter_factory.is_some()
        {
            return true;
        }
        // Check children
        for child in plan.children() {
            if check_plan(child.as_ref()) {
                return true;
            }
        }
        false
    }

    let has_adapter = check_plan(plan.as_ref());
    println!("    [Verify] {label} plan adapter check: {has_adapter}");
    has_adapter
}