Proposed API improvements

alamb · Dec 19, 2024 · f2d3029 · f2d3029
1 parent de2fbae
commit f2d3029
Showing 1 changed file with 12 additions and 6 deletions.
diff --git 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<Extensions>, state: Arc<SessionState>) -> Self {
-        DefaultSubstraitConsumer { extensions, state }
+    /// Create a new DefaultSubstraitConsumer
+    pub fn new(state: Arc<SessionState>) -> Self {
+        DefaultSubstraitConsumer {
+            extensions: Arc::new(Extensions::default()),
+            state,
+        }
+    }
+
+    pub fn with_extensions(mut self, extensions: Arc<Extensions>) -> 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()))
     }
 }