diff --git a/src/main/java/org/apache/paimon/trino/FixedBucketTableShuffleFunction.java b/src/main/java/org/apache/paimon/trino/FixedBucketTableShuffleFunction.java
index 839ad0c..534917e 100644
--- a/src/main/java/org/apache/paimon/trino/FixedBucketTableShuffleFunction.java
+++ b/src/main/java/org/apache/paimon/trino/FixedBucketTableShuffleFunction.java
@@ -55,7 +55,7 @@ public FixedBucketTableShuffleFunction(
                 ThreadLocal.withInitial(
                         () ->
                                 CodeGenUtils.newProjection(
-                                        schema.logicalPrimaryKeysType(), schema.primaryKeys()));
+                                        schema.logicalBucketKeyType(), schema.bucketKeys()));
         this.bucketCount = new CoreOptions(schema.options()).bucket();
         this.workerCount = workerCount;
         this.isRowId =