fix: add VecType to codegen

ColdL · ColdL · commit 1f4ea3f62fb2 · 2026-01-13T16:23:06.000+08:00
diff --git a/paimon-codegen/src/main/scala/org/apache/paimon/codegen/GenerateUtils.scala b/paimon-codegen/src/main/scala/org/apache/paimon/codegen/GenerateUtils.scala
@@ -127,10 +127,13 @@ object GenerateUtils {
       s"$sortUtil.compareBinary($leftTerm, $rightTerm)"
     case TINYINT | SMALLINT | INTEGER | BIGINT | FLOAT | DOUBLE | DATE | TIME_WITHOUT_TIME_ZONE =>
       s"($leftTerm > $rightTerm ? 1 : $leftTerm < $rightTerm ? -1 : 0)"
-    case ARRAY =>
-      val at = t.asInstanceOf[ArrayType]
+    case ARRAY | VECTOR =>
+      val elementType = t.getTypeRoot match {
+        case ARRAY => t.asInstanceOf[ArrayType].getElementType
+        case VECTOR => t.asInstanceOf[VecType].getElementType
+      }
       val compareFunc = newName("compareArray")
-      val compareCode = generateArrayCompare(ctx, nullsIsLast = false, at, "a", "b")
+      val compareCode = generateArrayCompare(ctx, nullsIsLast = false, elementType, "a", "b")
       val funcCode: String =
         s"""
           public int $compareFunc($ARRAY_DATA a, $ARRAY_DATA b) {
@@ -188,11 +191,10 @@ object GenerateUtils {
   def generateArrayCompare(
       ctx: CodeGeneratorContext,
       nullsIsLast: Boolean,
-      arrayType: ArrayType,
+      elementType: DataType,
       leftTerm: String,
       rightTerm: String): String = {
     val nullIsLastRet = if (nullsIsLast) 1 else -1
-    val elementType = arrayType.getElementType
     val fieldA = newName("fieldA")
     val isNullA = newName("isNullA")
     val lengthA = newName("lengthA")
@@ -379,6 +381,7 @@ object GenerateUtils {
     case DOUBLE => className[JDouble]
     case TIMESTAMP_WITHOUT_TIME_ZONE | TIMESTAMP_WITH_LOCAL_TIME_ZONE => className[Timestamp]
     case ARRAY => className[InternalArray]
+    case VECTOR => className[InternalVec]
     case MULTISET | MAP => className[InternalMap]
     case ROW => className[InternalRow]
     case VARIANT => className[Variant]
@@ -417,6 +420,8 @@ object GenerateUtils {
         s"$rowTerm.getTimestamp($indexTerm, ${getPrecision(t)})"
       case ARRAY =>
         s"$rowTerm.getArray($indexTerm)"
+      case VECTOR =>
+        s"$rowTerm.getVec($indexTerm)"
       case MULTISET | MAP =>
         s"$rowTerm.getMap($indexTerm)"
       case ROW =>
diff --git a/paimon-codegen/src/main/scala/org/apache/paimon/codegen/ScalarOperatorGens.scala b/paimon-codegen/src/main/scala/org/apache/paimon/codegen/ScalarOperatorGens.scala
@@ -60,7 +60,13 @@ object ScalarOperatorGens {
     }
     // array types
     else if (isArray(left.resultType) && canEqual) {
-      generateArrayComparison(ctx, left, right, resultType)
+      val elementType = left.resultType.asInstanceOf[ArrayType].getElementType
+      generateArrayComparison(ctx, left, right, elementType, resultType)
+    }
+    // vector type
+    else if (isVec(left.resultType) && canEqual) {
+      val elementType = left.resultType.asInstanceOf[VecType].getElementType
+      generateArrayComparison(ctx, left, right, elementType, resultType)
     }
     // map types
     else if (isMap(left.resultType) && canEqual) {
@@ -196,6 +202,7 @@ object ScalarOperatorGens {
       ctx: CodeGeneratorContext,
       left: GeneratedExpression,
       right: GeneratedExpression,
+      elementType: DataType,
       resultType: DataType): GeneratedExpression = {
     generateCallWithStmtIfArgsNotNull(ctx, resultType, Seq(left, right)) {
       args =>
@@ -204,7 +211,6 @@ object ScalarOperatorGens {
 
         val resultTerm = newName("compareResult")
 
-        val elementType = left.resultType.asInstanceOf[ArrayType].getElementType
         val elementCls = primitiveTypeTermForType(elementType)
         val elementDefault = primitiveDefaultValue(elementType)
 
@@ -225,6 +231,7 @@ object ScalarOperatorGens {
           rightElementExpr,
           new BooleanType(elementType.isNullable))
 
+        // TODO: With BinaryVec available, we can use it here.
         val stmt =
           s"""
              |boolean $resultTerm;
diff --git a/paimon-codegen/src/test/java/org/apache/paimon/codegen/EqualiserCodeGeneratorTest.java b/paimon-codegen/src/test/java/org/apache/paimon/codegen/EqualiserCodeGeneratorTest.java
@@ -30,6 +30,7 @@
 import org.apache.paimon.data.serializer.InternalArraySerializer;
 import org.apache.paimon.data.serializer.InternalMapSerializer;
 import org.apache.paimon.data.serializer.InternalRowSerializer;
+import org.apache.paimon.data.serializer.InternalVecSerializer;
 import org.apache.paimon.data.serializer.Serializer;
 import org.apache.paimon.data.variant.GenericVariant;
 import org.apache.paimon.types.DataType;
@@ -133,6 +134,16 @@ public class EqualiserCodeGeneratorTest {
                                 castFromString("[1,2,3]", DataTypes.ARRAY(new VarCharType())),
                                 castFromString("[4,5,6]", DataTypes.ARRAY(new VarCharType()))),
                         new InternalArraySerializer(DataTypes.VARCHAR(1))));
+        TEST_DATA.put(
+                DataTypeRoot.VECTOR,
+                new GeneratedData(
+                        DataTypes.VECTOR(3, DataTypes.FLOAT()),
+                        Pair.of(
+                                castFromString(
+                                        "[1.1,2.2,3.3]", DataTypes.VECTOR(3, DataTypes.FLOAT())),
+                                castFromString(
+                                        "[4.4,5.5,6.6]", DataTypes.VECTOR(3, DataTypes.FLOAT()))),
+                        new InternalVecSerializer(DataTypes.FLOAT(), 3)));
         TEST_DATA.put(
                 DataTypeRoot.MULTISET,
                 new GeneratedData(
diff --git a/paimon-common/src/main/java/org/apache/paimon/utils/TypeCheckUtils.java b/paimon-common/src/main/java/org/apache/paimon/utils/TypeCheckUtils.java
@@ -33,6 +33,7 @@
 import static org.apache.paimon.types.DataTypeRoot.TIMESTAMP_WITHOUT_TIME_ZONE;
 import static org.apache.paimon.types.DataTypeRoot.TIMESTAMP_WITH_LOCAL_TIME_ZONE;
 import static org.apache.paimon.types.DataTypeRoot.VARIANT;
+import static org.apache.paimon.types.DataTypeRoot.VECTOR;
 
 /** Utils for type. */
 public class TypeCheckUtils {
@@ -85,6 +86,10 @@ public static boolean isArray(DataType type) {
         return type.getTypeRoot() == ARRAY;
     }
 
+    public static boolean isVec(DataType type) {
+        return type.getTypeRoot() == VECTOR;
+    }
+
     public static boolean isMap(DataType type) {
         return type.getTypeRoot() == MAP;
     }
@@ -110,6 +115,7 @@ public static boolean isComparable(DataType type) {
                 && !isMultiset(type)
                 && !isRow(type)
                 && !isArray(type)
+                && !isVec(type)
                 && !isVariant(type)
                 && !isBlob(type);
     }
@@ -120,6 +126,7 @@ public static boolean isMutable(DataType type) {
             case CHAR:
             case VARCHAR: // the internal representation of String is BinaryString which is mutable
             case ARRAY:
+            case VECTOR:
             case MULTISET:
             case MAP:
             case ROW:
diff --git a/paimon-common/src/main/java/org/apache/paimon/utils/TypeUtils.java b/paimon-common/src/main/java/org/apache/paimon/utils/TypeUtils.java
@@ -18,22 +18,26 @@
 
 package org.apache.paimon.utils;
 
+import org.apache.paimon.data.ArrayBasedVec;
 import org.apache.paimon.data.BinaryString;
 import org.apache.paimon.data.Decimal;
 import org.apache.paimon.data.GenericArray;
 import org.apache.paimon.data.GenericMap;
 import org.apache.paimon.data.GenericRow;
+import org.apache.paimon.data.InternalArray;
 import org.apache.paimon.types.ArrayType;
 import org.apache.paimon.types.DataField;
 import org.apache.paimon.types.DataType;
 import org.apache.paimon.types.DataTypeChecks;
 import org.apache.paimon.types.DataTypeRoot;
+import org.apache.paimon.types.DataTypes;
 import org.apache.paimon.types.DecimalType;
 import org.apache.paimon.types.LocalZonedTimestampType;
 import org.apache.paimon.types.MapType;
 import org.apache.paimon.types.RowType;
 import org.apache.paimon.types.TimestampType;
 import org.apache.paimon.types.VarCharType;
+import org.apache.paimon.types.VecType;
 
 import org.apache.paimon.shade.jackson2.com.fasterxml.jackson.core.JsonProcessingException;
 import org.apache.paimon.shade.jackson2.com.fasterxml.jackson.databind.JsonNode;
@@ -212,6 +216,16 @@ public static Object castFromStringInternal(String s, DataType type, boolean isC
                     throw new RuntimeException(
                             String.format("Failed to parse Json String %s", s), e);
                 }
+            case VECTOR:
+                VecType vecType = (VecType) type;
+                DataType vecElementType = vecType.getElementType();
+                Object vecBaseArr =
+                        castFromStringInternal(s, DataTypes.ARRAY(vecElementType), isCdcValue);
+                if (vecBaseArr instanceof InternalArray) {
+                    return ArrayBasedVec.from((InternalArray) vecBaseArr);
+                } else {
+                    throw new RuntimeException("Failed to make array during building a vector");
+                }
             case MAP:
                 MapType mapType = (MapType) type;
                 DataType keyType = mapType.getKeyType();
@@ -333,6 +347,7 @@ public static boolean isInteroperable(DataType t1, DataType t2) {
 
         switch (t1.getTypeRoot()) {
             case ARRAY:
+            case VECTOR:
             case MAP:
             case MULTISET:
             case ROW:
diff --git a/paimon-core/src/test/java/org/apache/paimon/codegen/CodeGenUtilsTest.java b/paimon-core/src/test/java/org/apache/paimon/codegen/CodeGenUtilsTest.java
@@ -32,6 +32,7 @@
 import static org.apache.paimon.types.DataTypes.DOUBLE;
 import static org.apache.paimon.types.DataTypes.INT;
 import static org.apache.paimon.types.DataTypes.STRING;
+import static org.apache.paimon.types.DataTypes.VECTOR;
 import static org.assertj.core.api.Assertions.assertThat;
 
 class CodeGenUtilsTest {
@@ -74,6 +75,14 @@ public void testRecordComparatorCodegenCache() {
                 () -> newRecordComparator(Arrays.asList(STRING(), INT()), new int[] {0, 1}, true));
     }
 
+    @Test
+    public void testRecordComparatorCodegenCacheWithVec() {
+        assertClassEquals(
+                () ->
+                        newRecordComparator(
+                                Arrays.asList(STRING(), VECTOR(3, INT())), new int[] {0, 1}, true));
+    }
+
     @Test
     public void testRecordComparatorCodegenCacheMiss() {
         assertClassNotEquals(
