[Coral-Schema] Support ROW_NUMBER() OVER WINDOW (#437)

* [Coral-Schema] Support ROW_NUMBER() OVER WINDOW * fix inline comment
linkedin · Jul 20, 2023 · 8ebfdb1 · 8ebfdb1
1 parent 482fe22
commit 8ebfdb1
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 2 deletions.
diff --git a/coral-schema/src/main/java/com/linkedin/coral/schema/avro/RelToAvroSchemaConverter.java b/coral-schema/src/main/java/com/linkedin/coral/schema/avro/RelToAvroSchemaConverter.java
@@ -453,8 +453,14 @@ public RexNode visitCall(RexCall rexCall) {
 
     @Override
     public RexNode visitOver(RexOver rexOver) {
-      // TODO: implement this method
-      return super.visitOver(rexOver);
+      /**
+       * For aggregates over window, derive the expected field's schema using RexOver's return type alone.
+       * Example RexOver:
+       *  ROW_NUMBER() OVER (PARTITION BY colA, colB ORDER BY colC DESC)
+       */
+      appendField(rexOver.getType(), false,
+          SchemaUtilities.generateDocumentationForFunctionCall(rexOver, inputSchema, inputNode));
+      return rexOver;
     }
 
     @Override

diff --git a/...l-schema/src/test/java/com/linkedin/coral/schema/avro/ViewToAvroSchemaConverterTests.java b/...l-schema/src/test/java/com/linkedin/coral/schema/avro/ViewToAvroSchemaConverterTests.java
@@ -1015,6 +1015,19 @@ public void testProjectUdfReturnedStruct() {
         TestUtils.loadSchema("testProjectUdfReturnedStruct-expected.avsc"));
   }
 
+  @Test
+  public void testRowNumberOverWindow() {
+    String viewSql = "CREATE VIEW foo_with_rownumber_over_window AS SELECT *, ROW_NUMBER() OVER"
+        + " (PARTITION BY Id, Struct_Col.Bigint_Field" + " ORDER BY Id DESC) rank" + " FROM basecomplex bc";
+
+    TestUtils.executeCreateViewQuery("default", "foo_with_rownumber_over_window", viewSql);
+
+    ViewToAvroSchemaConverter viewToAvroSchemaConverter = ViewToAvroSchemaConverter.create(hiveMetastoreClient);
+    Schema actualSchema = viewToAvroSchemaConverter.toAvroSchema("default", "foo_with_rownumber_over_window");
+
+    Assert.assertEquals(actualSchema.toString(true), TestUtils.loadSchema("testRowNumberOverWindow-expected.avsc"));
+  }
+
   @Test
   public void testLowercaseSchema() {
     String viewSql = "CREATE VIEW v AS SELECT id as Id FROM basecomplex";

diff --git a/coral-schema/src/test/resources/testRowNumberOverWindow-expected.avsc b/coral-schema/src/test/resources/testRowNumberOverWindow-expected.avsc
@@ -0,0 +1,54 @@
+{
+  "type" : "record",
+  "name" : "foo_with_rownumber_over_window",
+  "namespace" : "default.foo_with_rownumber_over_window",
+  "fields" : [ {
+    "name" : "Id",
+    "type" : "int"
+  }, {
+    "name" : "Array_Col",
+    "type" : [ "null", {
+      "type" : "array",
+      "items" : [ "null", "string" ]
+    } ]
+  }, {
+    "name" : "Map_Col",
+    "type" : [ "null", {
+      "type" : "map",
+      "values" : [ "null", "string" ]
+    } ]
+  }, {
+    "name" : "Struct_Col",
+    "type" : [ "null", {
+      "type" : "record",
+      "name" : "Struct_col",
+      "namespace" : "default.foo_with_rownumber_over_window.foo_with_rownumber_over_window",
+      "fields" : [ {
+        "name" : "Bool_Field",
+        "type" : "boolean"
+      }, {
+        "name" : "Int_Field",
+        "type" : [ "null", "int" ]
+      }, {
+        "name" : "Bigint_Field",
+        "type" : "long"
+      }, {
+        "name" : "Float_Field",
+        "type" : [ "null", "float" ]
+      }, {
+        "name" : "Double_Field",
+        "type" : "double"
+      }, {
+        "name" : "Date_String_Field",
+        "type" : [ "null", "string" ]
+      }, {
+        "name" : "String_Field",
+        "type" : [ "null", "string" ]
+      } ]
+    } ]
+  }, {
+    "name" : "rank",
+    "type" : "long",
+    "doc" : "Field created in view by applying operator 'ROW_NUMBER'"
+  } ]
+}