diff --git a/src/context.rs b/src/context.rs
index bd31bf1..ab6a5b6 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -67,7 +67,7 @@ pub(crate) fn execution_plan_from_pyany(
         ))
     })?;
 
-    let codec = DefaultPhysicalExtensionCodec {};
+    let codec = ShuffleCodec {};
     let runtime = RuntimeEnv::default();
     let registry = SessionContext::new();
     plan_node