cleanup

alamb · Dec 19, 2024 · 93340dc · 93340dc
1 parent e8000d5
commit 93340dc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs
@@ -991,7 +991,8 @@ async fn extension_logical_plan() -> Result<()> {
         }),
     });
 
-    let consumer = DefaultSubstraitConsumer::new(Arc::new(ctx.state()));
+    let state = ctx.state();
+    let consumer = DefaultSubstraitConsumer::new(&state);
 
     let proto = to_substrait_plan(&ext_plan, &ctx.state())?;
     let plan2 = from_substrait_plan_with_consumer(&consumer, &proto).await?;

diff --git a/datafusion/substrait/tests/utils.rs b/datafusion/substrait/tests/utils.rs
@@ -51,7 +51,8 @@ pub mod test {
         ctx: SessionContext,
         plan: &Plan,
     ) -> Result<SessionContext> {
-        let consumer = Arc::new(DefaultSubstraitConsumer::default());
+        let state = ctx.state();
+        let consumer = Arc::new(DefaultSubstraitConsumer::new(&state));
         let schemas = TestSchemaCollector::collect_schemas(consumer, plan)?;
         let mut schema_map: HashMap<TableReference, Arc<dyn TableProvider>> =
             HashMap::new();