[core] Fix keyRearrange in PrimaryKeyLookupTable

lmagic233 · Jan 10, 2024 · af25cd2 · af25cd2
1 parent b57dc35
commit af25cd2
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 19 deletions.
diff --git a/...flink-common/src/main/java/org/apache/paimon/flink/lookup/FixedBucketFromPkExtractor.java b/...flink-common/src/main/java/org/apache/paimon/flink/lookup/FixedBucketFromPkExtractor.java
@@ -33,8 +33,6 @@ public class FixedBucketFromPkExtractor implements KeyAndBucketExtractor<Interna
 
     private transient InternalRow primaryKey;
 
-    private BinaryRow bucketKey;
-
     private final boolean sameBucketKeyAndTrimmedPrimaryKey;
 
     private final int numBuckets;
@@ -90,8 +88,7 @@ private BinaryRow bucketKey() {
             return trimmedPrimaryKey();
         }
 
-        bucketKey = bucketKeyProjection.apply(primaryKey);
-        return bucketKey;
+        return bucketKeyProjection.apply(primaryKey);
     }
 
     @Override

diff --git a/...aimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FullCacheLookupTable.java b/...aimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/FullCacheLookupTable.java
@@ -171,7 +171,7 @@ static FullCacheLookupTable create(Context context, long lruCacheSize) throws IO
             return new NoPrimaryKeyLookupTable(context, lruCacheSize);
         } else {
             if (new HashSet<>(primaryKeys).equals(new HashSet<>(context.joinKey))) {
-                return new PrimaryKeyLookupTable(context, lruCacheSize);
+                return new PrimaryKeyLookupTable(context, lruCacheSize, context.joinKey);
             } else {
                 return new SecondaryIndexLookupTable(context, lruCacheSize);
             }

diff --git a/...imon-flink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyLookupTable.java b/...imon-flink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyLookupTable.java
@@ -22,10 +22,14 @@
 import org.apache.paimon.data.serializer.InternalSerializers;
 import org.apache.paimon.lookup.BulkLoader;
 import org.apache.paimon.lookup.RocksDBValueState;
+import org.apache.paimon.table.FileStoreTable;
 import org.apache.paimon.types.RowKind;
 import org.apache.paimon.utils.KeyProjectedRow;
+import org.apache.paimon.utils.ProjectedRow;
 import org.apache.paimon.utils.TypeUtils;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 import java.util.Collections;
 import java.util.Iterator;
@@ -40,11 +44,15 @@ public class PrimaryKeyLookupTable extends FullCacheLookupTable {
 
     protected final KeyProjectedRow primaryKeyRow;
 
-    public PrimaryKeyLookupTable(Context context, long lruCacheSize) throws IOException {
+    @Nullable private final ProjectedRow keyRearrange;
+
+    public PrimaryKeyLookupTable(Context context, long lruCacheSize, List<String> joinKey)
+            throws IOException {
         super(context);
         List<String> fieldNames = projectedType.getFieldNames();
+        FileStoreTable table = context.table;
         this.primaryKeyMapping =
-                context.table.primaryKeys().stream().mapToInt(fieldNames::indexOf).toArray();
+                table.primaryKeys().stream().mapToInt(fieldNames::indexOf).toArray();
         this.primaryKeyRow = new KeyProjectedRow(primaryKeyMapping);
 
         this.tableState =
@@ -54,10 +62,24 @@ public PrimaryKeyLookupTable(Context context, long lruCacheSize) throws IOExcept
                                 TypeUtils.project(projectedType, primaryKeyMapping)),
                         InternalSerializers.create(projectedType),
                         lruCacheSize);
+
+        ProjectedRow keyRearrange = null;
+        if (!table.primaryKeys().equals(joinKey)) {
+            keyRearrange =
+                    ProjectedRow.from(
+                            table.primaryKeys().stream()
+                                    .map(joinKey::indexOf)
+                                    .mapToInt(value -> value)
+                                    .toArray());
+        }
+        this.keyRearrange = keyRearrange;
     }
 
     @Override
     public List<InternalRow> innerGet(InternalRow key) throws IOException {
+        if (keyRearrange != null) {
+            key = keyRearrange.replaceRow(key);
+        }
         InternalRow value = tableState.get(key);
         return value == null ? Collections.emptyList() : Collections.singletonList(value);
     }

diff --git a/...ink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyPartialLookupTable.java b/...ink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyPartialLookupTable.java
@@ -47,7 +47,7 @@ public class PrimaryKeyPartialLookupTable implements LookupTable {
 
     private final TableFileMonitor fileMonitor;
 
-    private final ProjectedRow projectedKey;
+    @Nullable private final ProjectedRow keyRearrange;
 
     public PrimaryKeyPartialLookupTable(
             FileStoreTable table,
@@ -71,12 +71,17 @@ public PrimaryKeyPartialLookupTable(
                         .withIOManager(new IOManagerImpl(tempPath.toString()));
         this.extractor = new FixedBucketFromPkExtractor(table.schema());
         this.fileMonitor = new TableFileMonitor(table, predicate);
-        this.projectedKey =
-                ProjectedRow.from(
-                        table.primaryKeys().stream()
-                                .map(joinKey::indexOf)
-                                .mapToInt(value -> value)
-                                .toArray());
+
+        ProjectedRow keyRearrange = null;
+        if (!table.primaryKeys().equals(joinKey)) {
+            keyRearrange =
+                    ProjectedRow.from(
+                            table.primaryKeys().stream()
+                                    .map(joinKey::indexOf)
+                                    .mapToInt(value -> value)
+                                    .toArray());
+        }
+        this.keyRearrange = keyRearrange;
     }
 
     @Override
@@ -86,7 +91,9 @@ public void open() throws Exception {
 
     @Override
     public List<InternalRow> get(InternalRow key) throws IOException {
-        key = projectedKey.replaceRow(key);
+        if (keyRearrange != null) {
+            key = keyRearrange.replaceRow(key);
+        }
         extractor.setRecord(key);
         int bucket = extractor.bucket();
         BinaryRow partition = extractor.partition();

diff --git a/...-flink-common/src/main/java/org/apache/paimon/flink/lookup/SecondaryIndexLookupTable.java b/...-flink-common/src/main/java/org/apache/paimon/flink/lookup/SecondaryIndexLookupTable.java
@@ -38,7 +38,7 @@ public class SecondaryIndexLookupTable extends PrimaryKeyLookupTable {
     private final KeyProjectedRow secKeyRow;
 
     public SecondaryIndexLookupTable(Context context, long lruCacheSize) throws IOException {
-        super(context, lruCacheSize / 2);
+        super(context, lruCacheSize / 2, context.table.primaryKeys());
         List<String> fieldNames = projectedType.getFieldNames();
         int[] secKeyMapping = context.joinKey.stream().mapToInt(fieldNames::indexOf).toArray();
         this.secKeyRow = new KeyProjectedRow(secKeyMapping);

diff --git a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/LookupJoinITCase.java b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/LookupJoinITCase.java
@@ -746,11 +746,13 @@ public void testAsyncRetryLookupSecKeyWithSequenceField() throws Exception {
         iterator.close();
     }
 
-    @Test
-    public void testPartialCacheBucketKeyOrder() throws Exception {
+    @ParameterizedTest
+    @EnumSource(LookupCacheMode.class)
+    public void testPartialCacheBucketKeyOrder(LookupCacheMode mode) throws Exception {
         sql(
                 "CREATE TABLE DIM (k2 INT, k1 INT, j INT , i INT, PRIMARY KEY(i, j) NOT ENFORCED) WITH"
-                        + " ('continuous.discovery-interval'='1 ms', 'lookup.cache'='auto', 'bucket' = '2', 'bucket-key' = 'j')");
+                        + " ('continuous.discovery-interval'='1 ms', 'lookup.cache'='%s', 'bucket' = '2', 'bucket-key' = 'j')",
+                mode);
 
         sql("CREATE TABLE T2 (j INT, i INT, `proctime` AS PROCTIME())");