fix comment

Author: JingsongLi

paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/lookup/PrimaryKeyPartialLookupTable.java

@@ -49,6 +49,7 @@
 import java.util.Set;
 
 import static org.apache.paimon.table.BucketMode.POSTPONE_BUCKET;
+import static org.apache.paimon.utils.Preconditions.checkArgument;
 
 /** Lookup table for primary key which supports to read the LSM tree directly. */
 public class PrimaryKeyPartialLookupTable implements LookupTable {
@@ -209,7 +210,8 @@ static class LocalQueryExecutor implements QueryExecutor {
         private final StreamTableScan scan;
         private final String tableName;
 
-        private final Map<BinaryRow, Integer> totalBuckets;
+        private final Integer defaultNumBuckets;
+        private final Map<BinaryRow, Integer> numBuckets;
 
         private LocalQueryExecutor(
                 FileStoreTable table,
@@ -238,13 +240,14 @@ private LocalQueryExecutor(
                             .newStreamScan();
 
             this.tableName = table.name();
-            this.totalBuckets = new HashMap<>();
+            this.defaultNumBuckets = table.bucketSpec().getNumBuckets();
+            this.numBuckets = new HashMap<>();
         }
 
         @Override
         @Nullable
         public Integer numBuckets(BinaryRow partition) {
-            return totalBuckets.get(partition);
+            return numBuckets.get(partition);
         }
 
         @Override
@@ -264,18 +267,30 @@ public void refresh() {
                 }
 
                 for (Split split : splits) {
-                    DataSplit dataSplit = (DataSplit) split;
-                    BinaryRow partition = dataSplit.partition();
-                    int bucket = dataSplit.bucket();
-                    List<DataFileMeta> before = dataSplit.beforeFiles();
-                    List<DataFileMeta> after = dataSplit.dataFiles();
-
-                    tableQuery.refreshFiles(partition, bucket, before, after);
-                    totalBuckets.put(partition, dataSplit.totalBuckets());
+                    refreshSplit((DataSplit) split);
                 }
             }
         }
 
+        @VisibleForTesting
+        void refreshSplit(DataSplit split) {
+            BinaryRow partition = split.partition();
+            int bucket = split.bucket();
+            List<DataFileMeta> before = split.beforeFiles();
+            List<DataFileMeta> after = split.dataFiles();
+
+            tableQuery.refreshFiles(partition, bucket, before, after);
+            Integer totalBuckets = split.totalBuckets();
+            if (totalBuckets == null) {
+                // Just for compatibility with older versions
+                checkArgument(
+                        defaultNumBuckets > 0,
+                        "This is a bug, old version table numBuckets should be greater than 0.");
+                totalBuckets = defaultNumBuckets;
+            }
+            numBuckets.put(partition, totalBuckets);
+        }
+
         @Override
         public void close() throws IOException {
             tableQuery.close();

paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/lookup/FileStoreLookupFunctionTest.java

@@ -37,6 +37,7 @@
 import org.apache.paimon.table.sink.CommitMessage;
 import org.apache.paimon.table.sink.StreamTableWrite;
 import org.apache.paimon.table.sink.TableCommitImpl;
+import org.apache.paimon.table.source.DataSplit;
 import org.apache.paimon.types.DataType;
 import org.apache.paimon.types.DataTypes;
 import org.apache.paimon.types.RowType;
@@ -49,6 +50,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
+import java.lang.reflect.Field;
 import java.net.InetSocketAddress;
 import java.nio.file.Path;
 import java.time.Duration;
@@ -62,6 +64,7 @@
 import java.util.Random;
 import java.util.UUID;
 
+import static org.apache.paimon.data.BinaryRow.EMPTY_ROW;
 import static org.apache.paimon.flink.FlinkConnectorOptions.LOOKUP_REFRESH_TIME_PERIODS_BLACKLIST;
 import static org.apache.paimon.service.ServiceManager.PRIMARY_KEY_LOOKUP;
 import static org.apache.paimon.testutils.assertj.PaimonAssertions.anyCauseMatches;
@@ -148,6 +151,26 @@ public void close() throws Exception {
         }
     }
 
+    @Test
+    public void testCompatibilityForOldVersion() throws Exception {
+        createLookupFunction(false, true, false, false);
+        commit(writeCommit(1));
+        PrimaryKeyPartialLookupTable lookupTable =
+                (PrimaryKeyPartialLookupTable) lookupFunction.lookupTable();
+        LocalQueryExecutor queryExecutor = (LocalQueryExecutor) lookupTable.queryExecutor();
+
+        // set totalBuckets to null, for testing old version
+        DataSplit split = (DataSplit) table.newReadBuilder().newScan().plan().splits().get(0);
+        Field field = DataSplit.class.getDeclaredField("totalBuckets");
+        field.setAccessible(true);
+        field.set(split, null);
+        assertThat(split.totalBuckets()).isNull();
+
+        // assert num buckets should be 2
+        queryExecutor.refreshSplit(split);
+        assertThat(queryExecutor.numBuckets(EMPTY_ROW)).isEqualTo(2);
+    }
+
     @ParameterizedTest
     @ValueSource(booleans = {false, true})
     public void testDefaultLocalPartial(boolean refreshAsync) throws Exception {