substrait-explain 0.7.0

//! Example demonstrating typed extension handlers using protobuf messages.
//!
//! This example shows how to:
//! 1. Define a custom protobuf message for extension configuration
//! 2. Implement `Explainable` for Substrait text round-tripping
//! 3. Parse Substrait text into the protobuf plan with a registry
//! 4. Recover the typed payload and format the plan back to text
//!
//! `AnyConvertible` is implemented automatically for any type that implements
//! `prost::Message` and `prost::Name`. If you are not using prost—or you need a
//! bespoke payload format—you can provide a manual `AnyConvertible`
//! implementation following the same pattern shown for `Explainable`.

use prost::{Message, Name};
use substrait::proto::{self, plan_rel, rel};
use substrait_explain::extensions::any::AnyRef;
use substrait_explain::extensions::{
    AnyConvertible, Explainable, ExtensionArgs, ExtensionColumn, ExtensionError, ExtensionRegistry,
};
use substrait_explain::{OutputOptions, Parser, format_with_registry};

/// Custom protobuf message for a ParquetScan extension configuration. In a real
/// implementation, this would be generated from a .proto file.
///
/// Also note: this would probably make more sense as a `ReadRel` with an
/// `ExtensionTable` rather than a full `ExtensionLeafRel`, but the registry
/// does not yet support `ExtensionTable` — see #114.
#[derive(Clone, PartialEq, Message)]
pub struct ParquetScanConfig {
    #[prost(string, tag = "1")]
    pub path: String,
    #[prost(int64, tag = "2")]
    pub batch_size: i64,
    #[prost(bool, tag = "3")]
    pub use_dictionary: bool,
    // A collection of (name, type) pairs, one per selected column. Conceptually
    // mirrors Substrait's `NamedStruct` (names + types bundled together)
    #[prost(message, repeated, tag = "4")]
    pub selected_columns: Vec<ParquetColumn>,
}

/// One selected column in a `ParquetScanConfig`: a (name, type) pair.
#[derive(Clone, PartialEq, Message)]
pub struct ParquetColumn {
    #[prost(string, tag = "1")]
    pub name: String,
    #[prost(message, optional, tag = "2")]
    pub r#type: Option<proto::Type>,
}

// Manually implement Name trait for our custom message
impl Name for ParquetScanConfig {
    const NAME: &'static str = "ParquetScanConfig";
    const PACKAGE: &'static str = "example";

    fn full_name() -> String {
        "example.ParquetScanConfig".to_string()
    }

    fn type_url() -> String {
        "type.googleapis.com/example.ParquetScanConfig".to_string()
    }
}

// Implement Explainable for our custom message
impl Explainable for ParquetScanConfig {
    fn name() -> &'static str {
        "TypedParquetScan"
    }

    fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
        let mut extractor = args.extractor();
        // path is required
        let path: &str = extractor.expect_named_arg("path")?;
        // batch_size and use_dictionary are optional, with default values
        let batch_size: i64 = extractor.get_named_or("batch_size", 1024)?;
        let use_dictionary: bool = extractor.get_named_or("use_dictionary", true)?;

        // Validate there are no other named arguments
        extractor.check_exhausted()?;

        let selected_columns = args
            .output_columns
            .iter()
            .map(|column| match column {
                ExtensionColumn::Named { name, r#type } => Ok(ParquetColumn {
                    name: name.clone(),
                    r#type: Some(r#type.clone()),
                }),
                _ => Err(ExtensionError::InvalidArgument(format!(
                    "expected named columns only: {column:?}"
                ))),
            })
            .collect::<Result<Vec<ParquetColumn>, ExtensionError>>()?;

        Ok(ParquetScanConfig {
            path: path.to_string(),
            // batch_size and use_dictionary are optional, with default values
            batch_size,
            use_dictionary,
            selected_columns,
        })
    }

    fn to_args(&self) -> Result<ExtensionArgs, ExtensionError> {
        let mut args = ExtensionArgs::default();

        // Add named arguments from the message
        args.insert("path", self.path.clone());
        args.insert("batch_size", self.batch_size);
        args.insert("use_dictionary", self.use_dictionary);

        for column in &self.selected_columns {
            let column_type = column.r#type.clone().ok_or_else(|| {
                ExtensionError::InvalidArgument(format!(
                    "column '{}' is missing a type",
                    column.name
                ))
            })?;
            args.output_columns.push(ExtensionColumn::Named {
                name: column.name.clone(),
                r#type: column_type,
            });
        }

        Ok(args)
    }
}

fn main() -> anyhow::Result<()> {
    let mut registry = ExtensionRegistry::new();
    registry.register_relation::<ParquetScanConfig>()?;

    let plan_text = r#"=== Plan
Root[customer_id, amount]
  ExtensionLeaf:TypedParquetScan[path='data/sales.parquet', batch_size=2048, use_dictionary=true => customer_id:i64, amount:fp64]"#;

    println!("Parsing plan text:\n---\n{plan_text}\n---");
    let plan = Parser::new()
        .with_extension_registry(registry.clone())
        .parse_plan(plan_text)?;
    println!("Parsed plan with {} relation(s).", plan.relations.len());

    let detail = extension_detail(&plan)?;
    let config = <ParquetScanConfig as AnyConvertible>::from_any(detail)?;
    println!("Decoded: {config:?}");

    let (roundtrip, errors) = format_with_registry(&plan, &OutputOptions::default(), &registry);
    if !errors.is_empty() {
        return Err(anyhow::anyhow!("formatting produced errors: {errors:?}"));
    }
    println!("Round-tripped text:\n{roundtrip}");

    let reparsed_plan = Parser::new()
        .with_extension_registry(registry)
        .parse_plan(&roundtrip)?;
    let reparsed_detail = extension_detail(&reparsed_plan)?;
    let reparsed_config = <ParquetScanConfig as AnyConvertible>::from_any(reparsed_detail)?;
    if reparsed_config != config {
        return Err(anyhow::anyhow!(
            "reparsed config does not match initial payload: {reparsed_config:?}"
        ));
    }

    if Parser::parse(plan_text).is_ok() {
        return Err(anyhow::anyhow!(
            "parsing without registry unexpectedly succeeded"
        ));
    }
    // else: Parsing without the registry fails, as expected.

    Ok(())
}

/// Extract the extension detail from the plan, for validation. Assumes Root -> ExtensionLeaf plan.
fn extension_detail<'a>(plan: &'a substrait::proto::Plan) -> Result<AnyRef<'a>, anyhow::Error> {
    assert!(plan.relations.len() == 1);
    let root = match plan.relations.first().unwrap() {
        substrait::proto::PlanRel {
            rel_type: Some(plan_rel::RelType::Root(root)),
        } => root,
        rel => return Err(anyhow::anyhow!("expected Root relation, got {rel:?}")),
    };

    let rel = root.input.as_ref().unwrap();

    match root.input.as_ref() {
        Some(substrait::proto::Rel {
            rel_type: Some(rel::RelType::ExtensionLeaf(leaf)),
        }) => leaf
            .detail
            .as_ref()
            .map(AnyRef::from)
            .ok_or_else(|| anyhow::anyhow!("expected Extension detail in parsed plan")),
        _ => Err(anyhow::anyhow!(
            "expected ExtensionLeaf relation, got {rel:?}"
        )),
    }
}