Skip to content

Commit

Permalink
feat: add data field in define_plan_node macro (#129)
Browse files Browse the repository at this point in the history
Add a `data_name` field to `define_plan_node` macro. Plan nodes with
data field must implement a method `explain_data`.

**An example**

```rust
#[derive(Clone, Debug, Serialize, Deserialize)]
struct ComplexData {
    a: i32,
    b: String,
}

#[derive(Clone, Debug)]
struct PhysicalComplexDummy(PlanNode);

impl PhysicalComplexDummy {
    pub fn explain_data(data: &Value) -> Vec<(&'static str, Pretty<'static>)> {
        if let Value::Serialized(serialized_data) = data {
            let data: ComplexData = bincode::deserialize(serialized_data).unwrap();
            vec![
                ("a", data.a.to_string().into()),
                ("b", data.b.to_string().into()),
            ]
        } else {
            unreachable!()
        }
    }
}

define_plan_node!(
    PhysicalComplexDummy: PlanNode,
    PhysicalScan, [
        { 0, child: PlanNode }
    ], [
    ],
    complex_data
);
```

I could just print data with `format!("{:?}, data)`, but the downside of
this is that for a struct

```rust
struct Foo {
    a: i32,
    b: i32,
}
```

The debug output would be `Foo { a: 1, b: 2}`. Here, the struct name is
implementation detail that should be hidden from the user, and the
bracket surrounding the fields makes these fields seem related in some
way, when they are actually not. So just letting user implement
`explain_data` gives maximum flexibility.
  • Loading branch information
Gun9niR authored Mar 24, 2024
1 parent 76b98c6 commit 026a01e
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 2 deletions.
6 changes: 5 additions & 1 deletion optd-datafusion-repr/src/plan_nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use std::sync::Arc;
use arrow_schema::DataType;
use optd_core::{
cascades::{CascadesOptimizer, GroupId},
rel_node::{RelNode, RelNodeMeta, RelNodeMetaMap, RelNodeRef, RelNodeTyp},
rel_node::{RelNode, RelNodeMeta, RelNodeMetaMap, RelNodeRef, RelNodeTyp, Value},
};

pub use agg::{LogicalAgg, PhysicalAgg};
Expand Down Expand Up @@ -208,6 +208,10 @@ pub trait OptRelNode: 'static + Clone {
}
}

pub trait ExplainData: OptRelNode {
fn explain_data(data: &Value) -> Vec<(&'static str, Pretty<'static>)>;
}

#[derive(Clone, Debug)]
pub struct PlanNode(pub(crate) OptRelNodeRef);

Expand Down
106 changes: 105 additions & 1 deletion optd-datafusion-repr/src/plan_nodes/macros.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,29 @@
/// Plan nodes with data fields must implement `ExplainData` trait. An example:
///
/// ```ignore
/// #[derive(Clone, Debug)]
/// struct PhysicalDummy(PlanNode);
///
/// // Implement `OptRelNode` using `define_plan_node!`...
///
/// impl ExplainData for PhysicalDummy {
/// fn explain_data(data: &Value) -> Vec<(&'static str, Pretty<'static>)> {
/// if let Value::Int32(i) = data {
/// vec![("primitive_data", i.to_string().into())]
/// } else {
/// unreachable!()
/// }
/// }
/// }
/// ```
macro_rules! define_plan_node {
(
$struct_name:ident : $meta_typ:tt,
$variant:ident,
[ $({ $child_id:literal, $child_name:ident : $child_meta_typ:ty }),* ] ,
[ $({ $attr_id:literal, $attr_name:ident : $attr_meta_typ:ty }),* ]
$(, { $inner_name:ident : $inner_typ:ty })?
$(, $data_name:ident)?
) => {
impl OptRelNode for $struct_name {
fn into_rel_node(self) -> OptRelNodeRef {
Expand All @@ -30,6 +49,8 @@ macro_rules! define_plan_node {
if let Some(meta_map) = meta_map {
fields = fields.with_meta(self.0.get_meta(meta_map));
};
define_plan_node!(@expand_fields self, $struct_name, fields $(, $data_name)?);

pretty_xmlish::Pretty::simple_record(
stringify!($struct_name),
fields,
Expand All @@ -44,16 +65,22 @@ macro_rules! define_plan_node {
pub fn new(
$($child_name : $child_meta_typ,)*
$($attr_name : $attr_meta_typ),*
$($data_name: Value)?
$(, $inner_name : $inner_typ)?
) -> $struct_name {
#[allow(unused_mut, unused)]
let mut data = None;
$(
data = Some($data_name);
)*
$struct_name($meta_typ(
optd_core::rel_node::RelNode {
typ: OptRelNodeTyp::$variant $( ($inner_name) )?,
children: vec![
$($child_name.into_rel_node(),)*
$($attr_name.into_rel_node()),*
],
data: None,
data,
}
.into(),
))
Expand Down Expand Up @@ -83,6 +110,83 @@ macro_rules! define_plan_node {
)?
}
};
// Dummy branch that does nothing when data is `None`.
(@expand_fields $self:ident, $struct_name:ident, $fields:ident) => {};
// Expand explain fields with data.
(@expand_fields $self:ident, $struct_name:ident, $fields:ident, $data_name:ident) => {
let data = $self.0 .0.data.as_ref().unwrap();
$fields.extend($struct_name::explain_data(data));
};
}

pub(crate) use define_plan_node;

#[cfg(test)]
mod test {
use crate::plan_nodes::*;
use optd_core::rel_node::Value;
use serde::{Deserialize, Serialize};

fn get_explain_str(pretty: &Pretty) -> String {
let mut config = PrettyConfig {
need_boundaries: false,
reduced_spaces: false,
width: 300,
..Default::default()
};
let mut out = String::new();
config.unicode(&mut out, pretty);
out
}

/// Ensure `define_plan_node` works with data field.
#[test]
fn test_explain_complex_data() {
#[derive(Clone, Debug, Serialize, Deserialize)]
struct ComplexData {
a: i32,
b: String,
}

#[derive(Clone, Debug)]
struct PhysicalComplexDummy(PlanNode);

impl ExplainData for PhysicalComplexDummy {
fn explain_data(data: &Value) -> Vec<(&'static str, Pretty<'static>)> {
if let Value::Serialized(serialized_data) = data {
let data: ComplexData = bincode::deserialize(serialized_data).unwrap();
vec![
("a", data.a.to_string().into()),
("b", data.b.to_string().into()),
]
} else {
unreachable!()
}
}
}

define_plan_node!(
PhysicalComplexDummy: PlanNode,
PhysicalScan, [
{ 0, child: PlanNode }
], [
],
complex_data
);

let node = PhysicalComplexDummy::new(
LogicalScan::new("a".to_string()).0,
Value::Serialized(
bincode::serialize(&ComplexData {
a: 1,
b: "a".to_string(),
})
.unwrap()
.into_iter()
.collect(),
),
);
let pretty = node.dispatch_explain(None);
println!("{}", get_explain_str(&pretty));
}
}

0 comments on commit 026a01e

Please sign in to comment.