substrait_explain/extensions/
examples.rs

1//! [`PartitionHint`] demonstrates how to implement a custom enhancement:
2//! positional enum arguments combined with an optional named integer argument.
3//!
4//! # Text Format
5//!
6//! ```text
7//! Read[data => col:i64]
8//!   + Enh:PartitionHint[&HASH, count=8]
9//! ```
10//!
11//! Each positional argument is a [`PartitionStrategy`] variant rendered with
12//! the `&` enum prefix.  The optional named argument `count` gives the target
13//! number of partitions (`0` / absent means "let the executor decide").
14
15use crate::extensions::args::{EnumValue, ExtensionArgs, ExtensionValue};
16use crate::extensions::registry::{Explainable, ExtensionError, ExtensionRegistry};
17
18// ---------------------------------------------------------------------------
19// PartitionStrategy enum
20// ---------------------------------------------------------------------------
21
22/// Partitioning strategy for a [`PartitionHint`] enhancement.
23#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, prost::Enumeration)]
24#[repr(i32)]
25pub enum PartitionStrategy {
26    /// No strategy specified.
27    Unspecified = 0,
28    /// Distribute rows by hashing one or more key columns.
29    Hash = 1,
30    /// Sort-based range partitioning.
31    Range = 2,
32    /// Broadcast the entire relation to every partition.
33    Broadcast = 3,
34}
35
36impl PartitionStrategy {
37    /// The identifier used in the text format (without the leading `&`).
38    pub fn as_str_name(self) -> &'static str {
39        match self {
40            PartitionStrategy::Unspecified => "UNSPECIFIED",
41            PartitionStrategy::Hash => "HASH",
42            PartitionStrategy::Range => "RANGE",
43            PartitionStrategy::Broadcast => "BROADCAST",
44        }
45    }
46
47    /// Parse from the text-format identifier (without the leading `&`).
48    pub fn from_str_name(s: &str) -> Option<Self> {
49        match s {
50            "UNSPECIFIED" => Some(PartitionStrategy::Unspecified),
51            "HASH" => Some(PartitionStrategy::Hash),
52            "RANGE" => Some(PartitionStrategy::Range),
53            "BROADCAST" => Some(PartitionStrategy::Broadcast),
54            _ => None,
55        }
56    }
57}
58
59// ---------------------------------------------------------------------------
60// PartitionHint
61// ---------------------------------------------------------------------------
62
63/// Enhancement that hints the executor how to partition a relation's output.
64///
65/// Attach this to any standard relation via `register_enhancement` to convey
66/// partitioning decisions made during planning.
67///
68/// # Text Format
69///
70/// ```rust
71/// # use substrait_explain::extensions::examples;
72/// # use substrait_explain::format_with_registry;
73/// # use substrait_explain::parser::Parser;
74/// #
75/// # let registry = examples::registry();
76/// # let parser = Parser::new().with_extension_registry(registry.clone());
77/// #
78/// # let plan_text = r#"
79/// === Plan
80/// Root[result]
81///   Read[data => col:i64]
82///     + Enh:PartitionHint[&HASH, count=8]
83/// # "#;
84/// #
85/// # let plan = parser.parse_plan(plan_text).unwrap();
86/// # let (formatted, errors) = format_with_registry(&plan, &Default::default(), &registry);
87/// # assert!(errors.is_empty());
88/// # assert_eq!(formatted.trim(), plan_text.trim());
89/// ```
90#[derive(Clone, PartialEq, prost::Message)]
91pub struct PartitionHint {
92    /// The strategies to apply, in order of preference.  Each value is the
93    /// integer representation of [`PartitionStrategy`].
94    #[prost(enumeration = "PartitionStrategy", repeated, tag = "1")]
95    pub strategies: Vec<i32>,
96    /// Target number of partitions.  `0` means "let the executor decide".
97    #[prost(int64, tag = "2")]
98    pub count: i64,
99}
100
101impl prost::Name for PartitionHint {
102    const PACKAGE: &'static str = "example";
103    const NAME: &'static str = "PartitionHint";
104
105    fn full_name() -> String {
106        "example.PartitionHint".to_owned()
107    }
108
109    fn type_url() -> String {
110        "type.googleapis.com/example.PartitionHint".to_owned()
111    }
112}
113
114impl Explainable for PartitionHint {
115    fn name() -> &'static str {
116        "PartitionHint"
117    }
118
119    fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
120        // Positional arguments are PartitionStrategy enum values.
121        let strategies: Result<Vec<i32>, ExtensionError> = args
122            .positional
123            .iter()
124            .map(|val| {
125                let EnumValue(ident) = EnumValue::try_from(val)?;
126                PartitionStrategy::from_str_name(&ident)
127                    .map(|s| s as i32)
128                    .ok_or_else(|| {
129                        ExtensionError::InvalidArgument(format!(
130                            "Unknown PartitionStrategy variant '&{ident}'; \
131                             expected one of &UNSPECIFIED, &HASH, &RANGE, &BROADCAST"
132                        ))
133                    })
134            })
135            .collect();
136
137        let mut extractor = args.extractor();
138        let count: i64 = extractor.get_named_or("count", 0)?;
139        extractor.check_exhausted()?;
140
141        Ok(PartitionHint {
142            strategies: strategies?,
143            count,
144        })
145    }
146
147    fn to_args(&self) -> Result<ExtensionArgs, ExtensionError> {
148        let mut args = ExtensionArgs::default();
149        for &raw in &self.strategies {
150            let s = PartitionStrategy::try_from(raw).unwrap_or(PartitionStrategy::Unspecified);
151            args.positional
152                .push(ExtensionValue::Enum(s.as_str_name().to_owned()));
153        }
154        if self.count != 0 {
155            args.insert("count", self.count);
156        }
157        Ok(args)
158    }
159}
160
161// ---------------------------------------------------------------------------
162// PlanHint
163// ---------------------------------------------------------------------------
164
165/// Optimization hint that carries a planner directive as a string.
166///
167/// Attach this to any standard relation via `register_optimization` to convey
168/// planner choices without changing relation semantics.
169///
170/// # Text Format
171///
172/// ```rust
173/// # use substrait_explain::extensions::examples;
174/// # use substrait_explain::format_with_registry;
175/// # use substrait_explain::parser::Parser;
176/// #
177/// # let registry = examples::registry();
178/// # let parser = Parser::new().with_extension_registry(registry.clone());
179/// #
180/// # let plan_text = r#"
181/// === Plan
182/// Root[result]
183///   Read[data => col:i64]
184///     + Opt:PlanHint[hint='use_index']
185/// # "#;
186/// #
187/// # let plan = parser.parse_plan(plan_text).unwrap();
188/// # let (formatted, errors) = format_with_registry(&plan, &Default::default(), &registry);
189/// # assert!(errors.is_empty());
190/// # assert_eq!(formatted.trim(), plan_text.trim());
191/// ```
192#[derive(Clone, PartialEq, prost::Message)]
193pub struct PlanHint {
194    /// Planner directive. The text format stores this as `hint='...'`.
195    #[prost(string, tag = "1")]
196    pub hint: String,
197}
198
199impl prost::Name for PlanHint {
200    const PACKAGE: &'static str = "example";
201    const NAME: &'static str = "PlanHint";
202
203    fn full_name() -> String {
204        "example.PlanHint".to_owned()
205    }
206
207    fn type_url() -> String {
208        "type.googleapis.com/example.PlanHint".to_owned()
209    }
210}
211
212impl Explainable for PlanHint {
213    fn name() -> &'static str {
214        "PlanHint"
215    }
216
217    fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
218        if !args.positional.is_empty() {
219            return Err(ExtensionError::InvalidArgument(
220                "PlanHint does not accept positional arguments".to_owned(),
221            ));
222        }
223
224        let mut extractor = args.extractor();
225        let hint: String = extractor.expect_named_arg::<&str>("hint")?.to_owned();
226        extractor.check_exhausted()?;
227        Ok(PlanHint { hint })
228    }
229
230    fn to_args(&self) -> Result<ExtensionArgs, ExtensionError> {
231        let mut args = ExtensionArgs::default();
232        args.named
233            .insert("hint".to_owned(), ExtensionValue::String(self.hint.clone()));
234        Ok(args)
235    }
236}
237
238/// Create an [`ExtensionRegistry`] preloaded with the example extension types.
239pub fn registry() -> ExtensionRegistry {
240    let mut registry = ExtensionRegistry::new();
241    registry
242        .register_enhancement::<PartitionHint>()
243        .expect("register PartitionHint example enhancement");
244    registry
245        .register_optimization::<PlanHint>()
246        .expect("register PlanHint example optimization");
247    registry
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use crate::extensions::AnyConvertible;
254
255    fn make_hint(strategies: Vec<PartitionStrategy>, count: i64) -> PartitionHint {
256        PartitionHint {
257            strategies: strategies.into_iter().map(|s| s as i32).collect(),
258            count,
259        }
260    }
261
262    #[test]
263    fn round_trip_via_any() {
264        let original = make_hint(vec![PartitionStrategy::Hash, PartitionStrategy::Range], 4);
265        let any = original.to_any().expect("encode");
266        let decoded = PartitionHint::from_any(any.as_ref()).expect("decode");
267        assert_eq!(original, decoded);
268    }
269
270    #[test]
271    fn to_args_produces_enum_and_named() {
272        let hint = make_hint(vec![PartitionStrategy::Hash], 8);
273        let args = hint.to_args().unwrap();
274        assert_eq!(args.positional.len(), 1);
275        assert!(matches!(&args.positional[0], ExtensionValue::Enum(s) if s == "HASH"));
276        let count = args.named.get("count").expect("count arg");
277        assert_eq!(i64::try_from(count).unwrap(), 8);
278    }
279
280    #[test]
281    fn to_args_omits_zero_count() {
282        let hint = make_hint(vec![PartitionStrategy::Broadcast], 0);
283        let args = hint.to_args().unwrap();
284        assert!(args.named.is_empty(), "count=0 should be omitted");
285    }
286
287    #[test]
288    fn from_args_round_trip() {
289        let original = make_hint(vec![PartitionStrategy::Hash, PartitionStrategy::Range], 16);
290        let args = original.to_args().unwrap();
291        let decoded = PartitionHint::from_args(&args).unwrap();
292        assert_eq!(original, decoded);
293    }
294
295    #[test]
296    fn from_args_rejects_unknown_strategy() {
297        let mut args = ExtensionArgs::default();
298        args.positional
299            .push(ExtensionValue::Enum("BOGUS".to_owned()));
300        assert!(PartitionHint::from_args(&args).is_err());
301    }
302
303    #[test]
304    fn from_args_rejects_non_enum_positional() {
305        // An integer positional arg where an enum is expected should fail.
306        let mut args = ExtensionArgs::default();
307        args.push(1_i64);
308        let result = PartitionHint::from_args(&args);
309        assert!(
310            result.is_err(),
311            "expected error for non-enum positional arg, got {result:?}"
312        );
313    }
314
315    #[test]
316    fn from_args_rejects_extra_named_args() {
317        // check_exhausted should reject unknown named args.
318        let mut args = ExtensionArgs::default();
319        args.insert("unknown_key", 99_i64);
320        let result = PartitionHint::from_args(&args);
321        assert!(
322            result.is_err(),
323            "expected error for unknown named arg, got {result:?}"
324        );
325    }
326
327    #[test]
328    fn from_args_empty_strategies_roundtrip() {
329        let original = make_hint(vec![], 0);
330        let args = original.to_args().unwrap();
331        let decoded = PartitionHint::from_args(&args).unwrap();
332        assert_eq!(original, decoded);
333        assert!(decoded.strategies.is_empty());
334        assert_eq!(decoded.count, 0);
335    }
336
337    #[test]
338    fn registry_roundtrip() {
339        let registry = registry();
340
341        let original = make_hint(vec![PartitionStrategy::Hash], 4);
342        let any = original.to_any().unwrap();
343
344        let (name, args) = registry
345            .decode_enhancement(any.as_ref())
346            .expect("decode_enhancement");
347        assert_eq!(name, "PartitionHint");
348        assert_eq!(args.positional.len(), 1);
349
350        let any2 = registry
351            .parse_enhancement("PartitionHint", &args)
352            .expect("parse_enhancement");
353        let decoded = PartitionHint::from_any(any2.as_ref()).unwrap();
354        assert_eq!(original, decoded);
355    }
356
357    #[test]
358    fn plan_hint_args_round_trip() {
359        let original = PlanHint {
360            hint: "use_index".to_owned(),
361        };
362        let args = original.to_args().unwrap();
363        let decoded = PlanHint::from_args(&args).unwrap();
364        assert_eq!(original, decoded);
365    }
366
367    #[test]
368    fn plan_hint_registry_roundtrip() {
369        let registry = registry();
370
371        let original = PlanHint {
372            hint: "parallel".to_owned(),
373        };
374        let any = original.to_any().unwrap();
375
376        let (name, args) = registry
377            .decode_optimization(any.as_ref())
378            .expect("decode_optimization");
379        assert_eq!(name, "PlanHint");
380
381        let any2 = registry
382            .parse_optimization("PlanHint", &args)
383            .expect("parse_optimization");
384        let decoded = PlanHint::from_any(any2.as_ref()).unwrap();
385        assert_eq!(original, decoded);
386    }
387}
substrait_explain/extensions/examples.rs

substrait_explain/extensions/
examples.rs
