[apache#2921] fix(catalog-lakehouse-iceberg): Add width param to buck…

…et and truncate functions in Gravitino SortOrder
caican00 · Apr 13, 2024 · 3e52bc0 · 3e52bc0
1 parent 72502e0
commit 3e52bc0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/...st/java/com/datastrato/gravitino/catalog/lakehouse/iceberg/converter/TestBaseConvert.java b/...st/java/com/datastrato/gravitino/catalog/lakehouse/iceberg/converter/TestBaseConvert.java
@@ -9,6 +9,7 @@
 import com.datastrato.gravitino.catalog.lakehouse.iceberg.IcebergColumn;
 import com.datastrato.gravitino.rel.Column;
 import com.datastrato.gravitino.rel.expressions.FunctionExpression;
+import com.datastrato.gravitino.rel.expressions.literals.Literals;
 import com.datastrato.gravitino.rel.expressions.sorts.NullOrdering;
 import com.datastrato.gravitino.rel.expressions.sorts.SortDirection;
 import com.datastrato.gravitino.rel.expressions.sorts.SortOrder;
@@ -105,6 +106,13 @@ protected static SortOrder createFunctionSortOrder(String name, String colName)
         RandomUtils.nextBoolean() ? NullOrdering.NULLS_FIRST : NullOrdering.NULLS_LAST);
   }
 
+  protected static SortOrder createFunctionSortOrder(String name, int width, String colName) {
+    return SortOrders.of(
+        FunctionExpression.of(name, Literals.integerLiteral(width), field(colName)),
+        RandomUtils.nextBoolean() ? SortDirection.DESCENDING : SortDirection.ASCENDING,
+        RandomUtils.nextBoolean() ? NullOrdering.NULLS_FIRST : NullOrdering.NULLS_LAST);
+  }
+
   protected static Types.NestedField createNestedField(
       int id, String name, org.apache.iceberg.types.Type type) {
     return Types.NestedField.optional(id, name, type, TEST_COMMENT);

diff --git a/.../com/datastrato/gravitino/catalog/lakehouse/iceberg/converter/TestToIcebergSortOrder.java b/.../com/datastrato/gravitino/catalog/lakehouse/iceberg/converter/TestToIcebergSortOrder.java
@@ -29,6 +29,8 @@ public void testToSortOrder() {
     sortOrders = ArrayUtils.add(sortOrders, createFunctionSortOrder("hour", "col_7"));
     sortOrders = ArrayUtils.add(sortOrders, createFunctionSortOrder("month", "col_8"));
     sortOrders = ArrayUtils.add(sortOrders, createFunctionSortOrder("year", "col_9"));
+    sortOrders = ArrayUtils.add(sortOrders, createFunctionSortOrder("bucket", 10, "col_10"));
+    sortOrders = ArrayUtils.add(sortOrders, createFunctionSortOrder("truncate", 2, "col_11"));
 
     Types.NestedField[] nestedFields =
         createNestedField("col_1", "col_2", "col_3", "col_4", "col_5");
@@ -40,6 +42,10 @@ public void testToSortOrder() {
         ArrayUtils.add(nestedFields, createNestedField(8, "col_8", Types.DateType.get()));
     nestedFields =
         ArrayUtils.add(nestedFields, createNestedField(9, "col_9", Types.DateType.get()));
+    nestedFields =
+        ArrayUtils.add(nestedFields, createNestedField(10, "col_10", Types.IntegerType.get()));
+    nestedFields =
+        ArrayUtils.add(nestedFields, createNestedField(11, "col_11", Types.StringType.get()));
     Schema schema = new Schema(nestedFields);
     org.apache.iceberg.SortOrder icebergSortOrder =
         ToIcebergSortOrder.toSortOrder(schema, sortOrders);
@@ -61,7 +67,9 @@ public void testToSortOrder() {
       if (colName.equals("col_6")
           || colName.equals("col_7")
           || colName.equals("col_8")
-          || colName.equals("col_9")) {
+          || colName.equals("col_9")
+          || colName.equals("col_10")
+          || colName.equals("col_11")) {
         Assertions.assertFalse(sortField.transform().isIdentity());
       } else {
         Assertions.assertTrue(sortField.transform().isIdentity());