[Core]Support generatedClass cache in CodeGeneratorImpl

paimon-codegen/src/main/java/org/apache/paimon/codegen/CodeGeneratorImpl.java

@@ -22,43 +22,123 @@
 import org.apache.paimon.types.RowType;
 import org.apache.paimon.utils.TypeUtils;
 
+import org.apache.paimon.shade.guava30.com.google.common.cache.Cache;
+import org.apache.paimon.shade.guava30.com.google.common.cache.CacheBuilder;
+
+import javax.annotation.Nullable;
+
+import java.time.Duration;
+import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.ExecutionException;
 
 /** Default implementation of {@link CodeGenerator}. */
 public class CodeGeneratorImpl implements CodeGenerator {
+    static final Cache<ClassKey, GeneratedClass<?>> GENERATED_CLASS_CACHE =
+            CacheBuilder.newBuilder()
+                    // assume the table schema will stay the same for a period of time
+                    .expireAfterAccess(Duration.ofMinutes(30))
+                    // estimated cache size
+                    .maximumSize(300)
+                    .softValues()
+                    .build();
 
     @Override
+    @SuppressWarnings("unchecked")
     public GeneratedClass<Projection> generateProjection(
             String name, RowType inputType, int[] inputMapping) {
         RowType outputType = TypeUtils.project(inputType, inputMapping);
-        return ProjectionCodeGenerator.generateProjection(
-                new CodeGeneratorContext(), name, inputType, outputType, inputMapping);
+        ClassKey cacheKey = new ClassKey(name, inputType.getFieldTypes(), inputMapping);
+        GeneratedClass<Projection> generatedClass;
+        try {
+            generatedClass =
+                    (GeneratedClass<Projection>)
+                            GENERATED_CLASS_CACHE.get(
+                                    cacheKey,
+                                    () ->
+                                            ProjectionCodeGenerator.generateProjection(
+                                                    new CodeGeneratorContext(),
+                                                    name,
+                                                    inputType,
+                                                    outputType,
+                                                    inputMapping));
+        } catch (ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+        return generatedClass;
     }
 
     @Override
+    @SuppressWarnings("unchecked")
     public GeneratedClass<NormalizedKeyComputer> generateNormalizedKeyComputer(
             List<DataType> inputTypes, int[] sortFields, String name) {
-        return new SortCodeGenerator(
-                        RowType.builder().fields(inputTypes).build(),
-                        getAscendingSortSpec(sortFields))
-                .generateNormalizedKeyComputer(name);
+        ClassKey cacheKey = new ClassKey(name, inputTypes, sortFields);
+        GeneratedClass<NormalizedKeyComputer> generatedClass;
+        try {
+            generatedClass =
+                    (GeneratedClass<NormalizedKeyComputer>)
+                            GENERATED_CLASS_CACHE.get(
+                                    cacheKey,
+                                    () ->
+                                            new SortCodeGenerator(
+                                                            RowType.builder()
+                                                                    .fields(inputTypes)
+                                                                    .build(),
+                                                            getAscendingSortSpec(sortFields))
+                                                    .generateNormalizedKeyComputer(name));
+        } catch (ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+        return generatedClass;
     }
 
     @Override
+    @SuppressWarnings("unchecked")
     public GeneratedClass<RecordComparator> generateRecordComparator(
             List<DataType> inputTypes, int[] sortFields, String name) {
-        return ComparatorCodeGenerator.gen(
-                name,
-                RowType.builder().fields(inputTypes).build(),
-                getAscendingSortSpec(sortFields));
+        ClassKey cacheKey = new ClassKey(name, inputTypes, sortFields);
+        GeneratedClass<RecordComparator> generatedClass;
+
+        try {
+            generatedClass =
+                    (GeneratedClass<RecordComparator>)
+                            GENERATED_CLASS_CACHE.get(
+                                    cacheKey,
+                                    () ->
+                                            ComparatorCodeGenerator.gen(
+                                                    name,
+                                                    RowType.builder().fields(inputTypes).build(),
+                                                    getAscendingSortSpec(sortFields)));
+        } catch (ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+        return generatedClass;
     }
 
     /** Generate a {@link RecordEqualiser}. */
     @Override
+    @SuppressWarnings("unchecked")
     public GeneratedClass<RecordEqualiser> generateRecordEqualiser(
             List<DataType> fieldTypes, String name) {
-        return new EqualiserCodeGenerator(RowType.builder().fields(fieldTypes).build())
-                .generateRecordEqualiser(name);
+        ClassKey cacheKey = new ClassKey(name, fieldTypes, null);
+        GeneratedClass<RecordEqualiser> generatedClass;
+
+        try {
+            generatedClass =
+                    (GeneratedClass<RecordEqualiser>)
+                            GENERATED_CLASS_CACHE.get(
+                                    cacheKey,
+                                    () ->
+                                            new EqualiserCodeGenerator(
+                                                            RowType.builder()
+                                                                    .fields(fieldTypes)
+                                                                    .build())
+                                                    .generateRecordEqualiser(name));
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        return generatedClass;
     }
 
     private SortSpec getAscendingSortSpec(int[] sortFields) {
@@ -68,4 +148,40 @@ private SortSpec getAscendingSortSpec(int[] sortFields) {
         }
         return builder.build();
     }
+
+    /** Class to use as key for the {@link #GENERATED_CLASS_CACHE}. */
+    static class ClassKey {
+        private final String name;
+
+        private final List<DataType> fields;
+
+        private final int[] fieldsIndex;
+
+        private ClassKey(String name, List<DataType> fields, @Nullable int[] fieldsIndex) {
+            this.name = name;
+            this.fields = fields;
+            this.fieldsIndex = fieldsIndex;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            ClassKey classKey = (ClassKey) o;
+            return Objects.equals(name, classKey.name)
+                    && Objects.equals(fields, classKey.fields)
+                    && Arrays.equals(fieldsIndex, classKey.fieldsIndex);
+        }
+
+        @Override
+        public int hashCode() {
+            int result = Objects.hash(name, fields);
+            result = 31 * result + Arrays.hashCode(fieldsIndex);
+            return result;
+        }
+    }
 }

paimon-codegen/src/test/java/org/apache/paimon/codegen/CodeGeneratorCacheTest.java

@@ -0,0 +1,166 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.paimon.codegen;
+
+import org.apache.paimon.types.DataType;
+import org.apache.paimon.types.IntType;
+import org.apache.paimon.types.RowType;
+import org.apache.paimon.types.VarCharType;
+
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for codegen cache in {@link CodeGeneratorImpl}. */
+public class CodeGeneratorCacheTest {
+
+    private final CodeGeneratorImpl codeGenerator = new CodeGeneratorImpl();
+
+    @BeforeAll
+    public static void before() {
+        // cleanup cached class before tests
+        CodeGeneratorImpl.GENERATED_CLASS_CACHE.invalidateAll();
+    }
+
+    @Test
+    public void testProjectionClassCacheReuse() {
+        String name = "Projection";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        int[] fieldIndexes = new int[] {0, 1};
+        GeneratedClass<Projection> generatedClass1 =
+                codeGenerator.generateProjection(
+                        name, RowType.builder().fields(dataTypes).build(), fieldIndexes);
+
+        GeneratedClass<Projection> generatedClass2 =
+                codeGenerator.generateProjection(
+                        name, RowType.builder().fields(dataTypes).build(), fieldIndexes);
+
+        assertThat(generatedClass1 == generatedClass2).isTrue();
+    }
+
+    @Test
+    public void testNormalizedKeyComputerClassCacheReuse() {
+        String name = "NormalizedKeyComputer";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        int[] fieldIndexes = new int[] {0, 1};
+        GeneratedClass<NormalizedKeyComputer> generatedClass1 =
+                codeGenerator.generateNormalizedKeyComputer(dataTypes, fieldIndexes, name);
+
+        GeneratedClass<NormalizedKeyComputer> generatedClass2 =
+                codeGenerator.generateNormalizedKeyComputer(dataTypes, fieldIndexes, name);
+
+        assertThat(generatedClass1 == generatedClass2).isTrue();
+    }
+
+    @Test
+    public void testRecordComparatorClassCacheReuse() {
+        String name = "RecordComparator";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        int[] fieldIndexes = new int[] {0, 1};
+        GeneratedClass<RecordComparator> generatedClass1 =
+                codeGenerator.generateRecordComparator(dataTypes, fieldIndexes, name);
+
+        GeneratedClass<RecordComparator> generatedClass2 =
+                codeGenerator.generateRecordComparator(dataTypes, fieldIndexes, name);
+
+        assertThat(generatedClass1 == generatedClass2).isTrue();
+    }
+
+    @Test
+    public void testRecordEqualiserClassCacheReuse() {
+        String name = "RecordEqualiser";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        GeneratedClass<RecordEqualiser> generatedClass1 =
+                codeGenerator.generateRecordEqualiser(dataTypes, name);
+
+        GeneratedClass<RecordEqualiser> generatedClass2 =
+                codeGenerator.generateRecordEqualiser(dataTypes, name);
+
+        assertThat(generatedClass1 == generatedClass2).isTrue();
+    }
+
+    @Test
+    public void testProjectionClassWithoutCache() {
+        String name = "Projection";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        int[] fieldIndexes = new int[] {0, 1};
+        GeneratedClass<Projection> generatedClass1 =
+                codeGenerator.generateProjection(
+                        name, RowType.builder().fields(dataTypes).build(), fieldIndexes);
+
+        CodeGeneratorImpl.GENERATED_CLASS_CACHE.invalidateAll();
+
+        GeneratedClass<Projection> generatedClass2 =
+                codeGenerator.generateProjection(
+                        name, RowType.builder().fields(dataTypes).build(), fieldIndexes);
+
+        assertThat(generatedClass1 == generatedClass2).isFalse();
+    }
+
+    @Test
+    public void testNormalizedKeyComputerClassWithoutCache() {
+        String name = "NormalizedKeyComputer";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        int[] fieldIndexes = new int[] {0, 1};
+        GeneratedClass<NormalizedKeyComputer> generatedClass1 =
+                codeGenerator.generateNormalizedKeyComputer(dataTypes, fieldIndexes, name);
+
+        CodeGeneratorImpl.GENERATED_CLASS_CACHE.invalidateAll();
+
+        GeneratedClass<NormalizedKeyComputer> generatedClass2 =
+                codeGenerator.generateNormalizedKeyComputer(dataTypes, fieldIndexes, name);
+
+        assertThat(generatedClass1 == generatedClass2).isFalse();
+    }
+
+    @Test
+    public void testRecordComparatorClassWithoutCache() {
+        String name = "RecordComparator";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        int[] fieldIndexes = new int[] {0, 1};
+        GeneratedClass<RecordComparator> generatedClass1 =
+                codeGenerator.generateRecordComparator(dataTypes, fieldIndexes, name);
+
+        CodeGeneratorImpl.GENERATED_CLASS_CACHE.invalidateAll();
+
+        GeneratedClass<RecordComparator> generatedClass2 =
+                codeGenerator.generateRecordComparator(dataTypes, fieldIndexes, name);
+
+        assertThat(generatedClass1 == generatedClass2).isFalse();
+    }
+
+    @Test
+    public void testRecordEqualiserClassWithoutCache() {
+        String name = "RecordEqualiser";
+        List<DataType> dataTypes = Arrays.asList(new VarCharType(1), new IntType());
+        GeneratedClass<RecordEqualiser> generatedClass1 =
+                codeGenerator.generateRecordEqualiser(dataTypes, name);
+
+        CodeGeneratorImpl.GENERATED_CLASS_CACHE.invalidateAll();
+
+        GeneratedClass<RecordEqualiser> generatedClass2 =
+                codeGenerator.generateRecordEqualiser(dataTypes, name);
+
+        assertThat(generatedClass1 == generatedClass2).isFalse();
+    }
+}