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//!
15//! This hidden module is crate-owned example support for doctests and
16//! integration tests, not a stable extension API.
17
18use crate::extensions::args::{EnumValue, ExtensionArgs, ExtensionValue};
19use crate::extensions::registry::{Explainable, ExtensionError, ExtensionRegistry};
20
21// ---------------------------------------------------------------------------
22// PartitionStrategy enum
23// ---------------------------------------------------------------------------
24
25/// Partitioning strategy for a [`PartitionHint`] enhancement.
26#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, prost::Enumeration)]
27#[repr(i32)]
28pub enum PartitionStrategy {
29    /// No strategy specified.
30    Unspecified = 0,
31    /// Distribute rows by hashing one or more key columns.
32    Hash = 1,
33    /// Sort-based range partitioning.
34    Range = 2,
35    /// Broadcast the entire relation to every partition.
36    Broadcast = 3,
37}
38
39impl PartitionStrategy {
40    /// The identifier used in the text format (without the leading `&`).
41    pub fn as_str_name(self) -> &'static str {
42        match self {
43            PartitionStrategy::Unspecified => "UNSPECIFIED",
44            PartitionStrategy::Hash => "HASH",
45            PartitionStrategy::Range => "RANGE",
46            PartitionStrategy::Broadcast => "BROADCAST",
47        }
48    }
49
50    /// Parse from the text-format identifier (without the leading `&`).
51    pub fn from_str_name(s: &str) -> Option<Self> {
52        match s {
53            "UNSPECIFIED" => Some(PartitionStrategy::Unspecified),
54            "HASH" => Some(PartitionStrategy::Hash),
55            "RANGE" => Some(PartitionStrategy::Range),
56            "BROADCAST" => Some(PartitionStrategy::Broadcast),
57            _ => None,
58        }
59    }
60}
61
62// ---------------------------------------------------------------------------
63// PartitionHint
64// ---------------------------------------------------------------------------
65
66/// Enhancement that hints the executor how to partition a relation's output.
67///
68/// Attach this to any standard relation via `register_enhancement` to convey
69/// partitioning decisions made during planning.
70///
71/// # Text Format
72///
73/// ```rust
74/// # use substrait_explain::extensions::examples;
75/// # use substrait_explain::format_with_registry;
76/// # use substrait_explain::Parser;
77/// #
78/// # let registry = examples::registry();
79/// # let parser = Parser::new().with_extension_registry(registry.clone());
80/// #
81/// # let plan_text = r#"
82/// === Plan
83/// Root[result]
84///   Read[data => col:i64]
85///     + Enh:PartitionHint[&HASH, count=8]
86/// # "#;
87/// #
88/// # let plan = parser.parse_plan(plan_text).unwrap();
89/// # let (formatted, errors) = format_with_registry(&plan, &Default::default(), &registry);
90/// # assert!(errors.is_empty());
91/// # assert_eq!(formatted.trim(), plan_text.trim());
92/// ```
93#[derive(Clone, PartialEq, prost::Message)]
94pub struct PartitionHint {
95    /// The strategies to apply, in order of preference.  Each value is the
96    /// integer representation of [`PartitionStrategy`].
97    #[prost(enumeration = "PartitionStrategy", repeated, tag = "1")]
98    pub strategies: Vec<i32>,
99    /// Target number of partitions.  `0` means "let the executor decide".
100    #[prost(int64, tag = "2")]
101    pub count: i64,
102}
103
104impl prost::Name for PartitionHint {
105    const PACKAGE: &'static str = "example";
106    const NAME: &'static str = "PartitionHint";
107
108    fn full_name() -> String {
109        "example.PartitionHint".to_owned()
110    }
111
112    fn type_url() -> String {
113        "type.googleapis.com/example.PartitionHint".to_owned()
114    }
115}
116
117impl Explainable for PartitionHint {
118    fn name() -> &'static str {
119        "PartitionHint"
120    }
121
122    fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
123        // Positional arguments are PartitionStrategy enum values.
124        let strategies: Result<Vec<i32>, ExtensionError> = args
125            .positional
126            .iter()
127            .map(|val| {
128                let EnumValue(ident) = EnumValue::try_from(val)?;
129                PartitionStrategy::from_str_name(&ident)
130                    .map(|s| s as i32)
131                    .ok_or_else(|| {
132                        ExtensionError::InvalidArgument(format!(
133                            "Unknown PartitionStrategy variant '&{ident}'; \
134                             expected one of &UNSPECIFIED, &HASH, &RANGE, &BROADCAST"
135                        ))
136                    })
137            })
138            .collect();
139
140        let mut extractor = args.extractor();
141        let count: i64 = extractor.get_named_or("count", 0)?;
142        extractor.check_exhausted()?;
143
144        Ok(PartitionHint {
145            strategies: strategies?,
146            count,
147        })
148    }
149
150    fn to_args(&self) -> Result<ExtensionArgs, ExtensionError> {
151        let mut args = ExtensionArgs::default();
152        for &raw in &self.strategies {
153            let s = PartitionStrategy::try_from(raw).unwrap_or(PartitionStrategy::Unspecified);
154            args.positional
155                .push(ExtensionValue::Enum(s.as_str_name().to_owned()));
156        }
157        if self.count != 0 {
158            args.insert("count", self.count);
159        }
160        Ok(args)
161    }
162}
163
164// ---------------------------------------------------------------------------
165// PlanHint
166// ---------------------------------------------------------------------------
167
168/// Optimization hint that carries a planner directive as a string.
169///
170/// Attach this to any standard relation via `register_optimization` to convey
171/// planner choices without changing relation semantics.
172///
173/// # Text Format
174///
175/// ```rust
176/// # use substrait_explain::extensions::examples;
177/// # use substrait_explain::format_with_registry;
178/// # use substrait_explain::Parser;
179/// #
180/// # let registry = examples::registry();
181/// # let parser = Parser::new().with_extension_registry(registry.clone());
182/// #
183/// # let plan_text = r#"
184/// === Plan
185/// Root[result]
186///   Read[data => col:i64]
187///     + Opt:PlanHint[hint='use_index']
188/// # "#;
189/// #
190/// # let plan = parser.parse_plan(plan_text).unwrap();
191/// # let (formatted, errors) = format_with_registry(&plan, &Default::default(), &registry);
192/// # assert!(errors.is_empty());
193/// # assert_eq!(formatted.trim(), plan_text.trim());
194/// ```
195#[derive(Clone, PartialEq, prost::Message)]
196pub struct PlanHint {
197    /// Planner directive. The text format stores this as `hint='...'`.
198    #[prost(string, tag = "1")]
199    pub hint: String,
200}
201
202impl prost::Name for PlanHint {
203    const PACKAGE: &'static str = "example";
204    const NAME: &'static str = "PlanHint";
205
206    fn full_name() -> String {
207        "example.PlanHint".to_owned()
208    }
209
210    fn type_url() -> String {
211        "type.googleapis.com/example.PlanHint".to_owned()
212    }
213}
214
215impl Explainable for PlanHint {
216    fn name() -> &'static str {
217        "PlanHint"
218    }
219
220    fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
221        if !args.positional.is_empty() {
222            return Err(ExtensionError::InvalidArgument(
223                "PlanHint does not accept positional arguments".to_owned(),
224            ));
225        }
226
227        let mut extractor = args.extractor();
228        let hint: String = extractor.expect_named_arg::<&str>("hint")?.to_owned();
229        extractor.check_exhausted()?;
230        Ok(PlanHint { hint })
231    }
232
233    fn to_args(&self) -> Result<ExtensionArgs, ExtensionError> {
234        let mut args = ExtensionArgs::default();
235        args.named
236            .insert("hint".to_owned(), ExtensionValue::String(self.hint.clone()));
237        Ok(args)
238    }
239}
240
241// ---------------------------------------------------------------------------
242// BlobStoreRead
243// ---------------------------------------------------------------------------
244
245/// ExtensionTable detail for a simple blob-store backed read.
246///
247/// Attach this to `Read:Extension[...]` via `register_extension_table` to
248/// describe a custom table source whose output schema is carried by the
249/// surrounding `ReadRel.base_schema`.
250///
251/// # Text Format
252///
253/// ```rust
254/// # use substrait_explain::extensions::examples;
255/// # use substrait_explain::format_with_registry;
256/// # use substrait_explain::Parser;
257/// #
258/// # let registry = examples::registry();
259/// # let parser = Parser::new().with_extension_registry(registry.clone());
260/// #
261/// # let plan_text = r#"
262/// === Plan
263/// Root[id, payload]
264///   Read:Extension[id:i64, payload:string]
265///     + Ext:BlobStoreRead['path/to/file', limit=100]
266/// # "#;
267/// #
268/// # let plan = parser.parse_plan(plan_text).unwrap();
269/// # let (formatted, errors) = format_with_registry(&plan, &Default::default(), &registry);
270/// # assert!(errors.is_empty());
271/// # assert_eq!(formatted.trim(), plan_text.trim());
272/// ```
273#[derive(Clone, PartialEq, prost::Message)]
274pub struct BlobStoreRead {
275    /// Blob path or URI to read.
276    #[prost(string, tag = "1")]
277    pub path: String,
278    /// Optional row limit. `0` means no limit.
279    #[prost(int64, tag = "2")]
280    pub limit: i64,
281    /// Whether archived blobs should be included.
282    #[prost(bool, tag = "3")]
283    pub include_archived: bool,
284}
285
286impl prost::Name for BlobStoreRead {
287    const PACKAGE: &'static str = "example";
288    const NAME: &'static str = "BlobStoreRead";
289
290    fn full_name() -> String {
291        "example.BlobStoreRead".to_owned()
292    }
293
294    fn type_url() -> String {
295        "type.googleapis.com/example.BlobStoreRead".to_owned()
296    }
297}
298
299impl Explainable for BlobStoreRead {
300    fn name() -> &'static str {
301        "BlobStoreRead"
302    }
303
304    fn from_args(args: &ExtensionArgs) -> Result<Self, ExtensionError> {
305        if args.positional.len() != 1 {
306            return Err(ExtensionError::InvalidArgument(format!(
307                "BlobStoreRead expects exactly 1 positional path argument, got {}",
308                args.positional.len()
309            )));
310        }
311        if !args.output_columns.is_empty() {
312            return Err(ExtensionError::InvalidArgument(
313                "BlobStoreRead output columns belong in Read:Extension[...]".to_owned(),
314            ));
315        }
316
317        let path = <&str>::try_from(&args.positional[0])?.to_owned();
318        let mut extractor = args.extractor();
319        let limit: i64 = extractor.get_named_or("limit", 0)?;
320        let include_archived: bool = extractor.get_named_or("include_archived", false)?;
321        extractor.check_exhausted()?;
322
323        Ok(Self {
324            path,
325            limit,
326            include_archived,
327        })
328    }
329
330    fn to_args(&self) -> Result<ExtensionArgs, ExtensionError> {
331        let mut args = ExtensionArgs::default();
332        args.positional
333            .push(ExtensionValue::String(self.path.clone()));
334        if self.limit != 0 {
335            args.named
336                .insert("limit".to_owned(), ExtensionValue::Integer(self.limit));
337        }
338        if self.include_archived {
339            args.named
340                .insert("include_archived".to_owned(), ExtensionValue::Boolean(true));
341        }
342        Ok(args)
343    }
344}
345
346/// Create an [`ExtensionRegistry`] preloaded with the example extension types.
347pub fn registry() -> ExtensionRegistry {
348    let mut registry = ExtensionRegistry::new();
349    registry
350        .register_enhancement::<PartitionHint>()
351        .expect("register PartitionHint example enhancement");
352    registry
353        .register_optimization::<PlanHint>()
354        .expect("register PlanHint example optimization");
355    registry
356        .register_extension_table::<BlobStoreRead>()
357        .expect("register BlobStoreRead example extension table");
358    registry
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364    use crate::extensions::AnyConvertible;
365
366    fn make_hint(strategies: Vec<PartitionStrategy>, count: i64) -> PartitionHint {
367        PartitionHint {
368            strategies: strategies.into_iter().map(|s| s as i32).collect(),
369            count,
370        }
371    }
372
373    #[test]
374    fn round_trip_via_any() {
375        let original = make_hint(vec![PartitionStrategy::Hash, PartitionStrategy::Range], 4);
376        let any = original.to_any().expect("encode");
377        let decoded = PartitionHint::from_any(any.as_ref()).expect("decode");
378        assert_eq!(original, decoded);
379    }
380
381    #[test]
382    fn to_args_produces_enum_and_named() {
383        let hint = make_hint(vec![PartitionStrategy::Hash], 8);
384        let args = hint.to_args().unwrap();
385        assert_eq!(args.positional.len(), 1);
386        assert!(matches!(&args.positional[0], ExtensionValue::Enum(s) if s == "HASH"));
387        let count = args.named.get("count").expect("count arg");
388        assert_eq!(i64::try_from(count).unwrap(), 8);
389    }
390
391    #[test]
392    fn to_args_omits_zero_count() {
393        let hint = make_hint(vec![PartitionStrategy::Broadcast], 0);
394        let args = hint.to_args().unwrap();
395        assert!(args.named.is_empty(), "count=0 should be omitted");
396    }
397
398    #[test]
399    fn from_args_round_trip() {
400        let original = make_hint(vec![PartitionStrategy::Hash, PartitionStrategy::Range], 16);
401        let args = original.to_args().unwrap();
402        let decoded = PartitionHint::from_args(&args).unwrap();
403        assert_eq!(original, decoded);
404    }
405
406    #[test]
407    fn from_args_rejects_unknown_strategy() {
408        let mut args = ExtensionArgs::default();
409        args.positional
410            .push(ExtensionValue::Enum("BOGUS".to_owned()));
411        assert!(PartitionHint::from_args(&args).is_err());
412    }
413
414    #[test]
415    fn from_args_rejects_non_enum_positional() {
416        // An integer positional arg where an enum is expected should fail.
417        let mut args = ExtensionArgs::default();
418        args.push(1_i64);
419        let result = PartitionHint::from_args(&args);
420        assert!(
421            result.is_err(),
422            "expected error for non-enum positional arg, got {result:?}"
423        );
424    }
425
426    #[test]
427    fn from_args_rejects_extra_named_args() {
428        // check_exhausted should reject unknown named args.
429        let mut args = ExtensionArgs::default();
430        args.insert("unknown_key", 99_i64);
431        let result = PartitionHint::from_args(&args);
432        assert!(
433            result.is_err(),
434            "expected error for unknown named arg, got {result:?}"
435        );
436    }
437
438    #[test]
439    fn from_args_empty_strategies_roundtrip() {
440        let original = make_hint(vec![], 0);
441        let args = original.to_args().unwrap();
442        let decoded = PartitionHint::from_args(&args).unwrap();
443        assert_eq!(original, decoded);
444        assert!(decoded.strategies.is_empty());
445        assert_eq!(decoded.count, 0);
446    }
447
448    #[test]
449    fn registry_roundtrip() {
450        let registry = registry();
451
452        let original = make_hint(vec![PartitionStrategy::Hash], 4);
453        let any = original.to_any().unwrap();
454
455        let (name, args) = registry
456            .decode_enhancement(any.as_ref())
457            .expect("decode_enhancement");
458        assert_eq!(name, "PartitionHint");
459        assert_eq!(args.positional.len(), 1);
460
461        let any2 = registry
462            .parse_enhancement("PartitionHint", &args)
463            .expect("parse_enhancement");
464        let decoded = PartitionHint::from_any(any2.as_ref()).unwrap();
465        assert_eq!(original, decoded);
466    }
467
468    #[test]
469    fn plan_hint_args_round_trip() {
470        let original = PlanHint {
471            hint: "use_index".to_owned(),
472        };
473        let args = original.to_args().unwrap();
474        let decoded = PlanHint::from_args(&args).unwrap();
475        assert_eq!(original, decoded);
476    }
477
478    #[test]
479    fn plan_hint_registry_roundtrip() {
480        let registry = registry();
481
482        let original = PlanHint {
483            hint: "parallel".to_owned(),
484        };
485        let any = original.to_any().unwrap();
486
487        let (name, args) = registry
488            .decode_optimization(any.as_ref())
489            .expect("decode_optimization");
490        assert_eq!(name, "PlanHint");
491
492        let any2 = registry
493            .parse_optimization("PlanHint", &args)
494            .expect("parse_optimization");
495        let decoded = PlanHint::from_any(any2.as_ref()).unwrap();
496        assert_eq!(original, decoded);
497    }
498
499    #[test]
500    fn blob_store_read_args_round_trip() {
501        let original = BlobStoreRead {
502            path: "path/to/file".to_owned(),
503            limit: 100,
504            include_archived: true,
505        };
506
507        let args = original.to_args().unwrap();
508        let decoded = BlobStoreRead::from_args(&args).unwrap();
509
510        assert_eq!(original, decoded);
511    }
512
513    #[test]
514    fn blob_store_read_registry_roundtrip() {
515        let registry = registry();
516
517        let original = BlobStoreRead {
518            path: "path/to/file".to_owned(),
519            limit: 100,
520            include_archived: true,
521        };
522        let any = original.to_any().unwrap();
523
524        let (name, args) = registry
525            .decode_extension_table(any.as_ref())
526            .expect("decode_extension_table");
527        assert_eq!(name, "BlobStoreRead");
528
529        let any2 = registry
530            .parse_extension_table("BlobStoreRead", &args)
531            .expect("parse_extension_table");
532        let decoded = BlobStoreRead::from_any(any2.as_ref()).unwrap();
533        assert_eq!(original, decoded);
534    }
535}
substrait_explain/extensions/examples.rs

substrait_explain/extensions/
examples.rs
