diff --git a/datafusion/substrait/src/logical_plan/consumer.rs b/datafusion/substrait/src/logical_plan/consumer.rs index f3475219be3d3..a6941e4ddbe42 100644 --- a/datafusion/substrait/src/logical_plan/consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer.rs @@ -477,17 +477,23 @@ pub struct DefaultSubstraitConsumer { } impl DefaultSubstraitConsumer { - pub fn new(extensions: Arc, state: Arc) -> Self { - DefaultSubstraitConsumer { extensions, state } + /// Create a new DefaultSubstraitConsumer + pub fn new(state: Arc) -> Self { + DefaultSubstraitConsumer { + extensions: Arc::new(Extensions::default()), + state, + } + } + + pub fn with_extensions(mut self, extensions: Arc) -> Self { + self.extensions = extensions; + self } } impl Default for DefaultSubstraitConsumer { fn default() -> Self { - DefaultSubstraitConsumer { - extensions: Arc::new(Extensions::default()), - state: Arc::new(SessionContext::default().state()), - } + DefaultSubstraitConsumer::new(Arc::new(SessionContext::default().state())) } }