feat: support window function lookup

substrait-io · Aug 10, 2023 · 418ad58 · 418ad58
1 parent c93f826
commit 418ad58
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/core/src/main/java/io/substrait/extension/AbstractExtensionLookup.java b/core/src/main/java/io/substrait/extension/AbstractExtensionLookup.java
@@ -24,6 +24,17 @@ public SimpleExtension.ScalarFunctionVariant getScalarFunction(
     return extensions.getScalarFunction(anchor);
   }
 
+  public SimpleExtension.WindowFunctionVariant getWindowFunction(
+      int reference, SimpleExtension.ExtensionCollection extensions) {
+    var anchor = functionAnchorMap.get(reference);
+    if (anchor == null) {
+      throw new IllegalArgumentException(
+          "Unknown function id. Make sure that the function id provided was shared in the extensions section of the plan.");
+    }
+
+    return extensions.getWindowFunction(anchor);
+  }
+
   public SimpleExtension.AggregateFunctionVariant getAggregateFunction(
       int reference, SimpleExtension.ExtensionCollection extensions) {
     var anchor = functionAnchorMap.get(reference);

diff --git a/core/src/main/java/io/substrait/extension/ExtensionLookup.java b/core/src/main/java/io/substrait/extension/ExtensionLookup.java
@@ -9,6 +9,9 @@ public interface ExtensionLookup {
   SimpleExtension.ScalarFunctionVariant getScalarFunction(
       int reference, SimpleExtension.ExtensionCollection extensions);
 
+  SimpleExtension.WindowFunctionVariant getWindowFunction(
+      int reference, SimpleExtension.ExtensionCollection extensions);
+
   SimpleExtension.AggregateFunctionVariant getAggregateFunction(
       int reference, SimpleExtension.ExtensionCollection extensions);