refactor: move WindowFunctionInvocation into Expression

substrait-io · Aug 30, 2023 · 082139c · 082139c
1 parent 00d52fc
commit 082139c
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 57 deletions.
diff --git a/core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java b/core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java
@@ -140,7 +140,7 @@ public OUTPUT visit(Expression.ScalarFunctionInvocation expr) throws EXCEPTION {
   }
 
   @Override
-  public OUTPUT visit(WindowFunctionInvocation expr) throws EXCEPTION {
+  public OUTPUT visit(Expression.WindowFunctionInvocation expr) throws EXCEPTION {
     return visitFallback(expr);
   }
 

diff --git a/core/src/main/java/io/substrait/expression/Expression.java b/core/src/main/java/io/substrait/expression/Expression.java
@@ -570,6 +570,42 @@ public <R, E extends Throwable> R accept(ExpressionVisitor<R, E> visitor) throws
     }
   }
 
+  @Value.Immutable
+  abstract class WindowFunctionInvocation implements Expression {
+
+    public abstract SimpleExtension.WindowFunctionVariant declaration();
+
+    public abstract List<FunctionArg> arguments();
+
+    public abstract Map<String, FunctionOption> options();
+
+    public abstract AggregationPhase aggregationPhase();
+
+    public abstract List<Expression> partitionBy();
+
+    public abstract List<SortField> sort();
+
+    public abstract WindowBound lowerBound();
+
+    public abstract WindowBound upperBound();
+
+    public abstract Type outputType();
+
+    public Type getType() {
+      return outputType();
+    }
+
+    public abstract AggregationInvocation invocation();
+
+    public static ImmutableExpression.WindowFunctionInvocation.Builder builder() {
+      return ImmutableExpression.WindowFunctionInvocation.builder();
+    }
+
+    public <R, E extends Throwable> R accept(ExpressionVisitor<R, E> visitor) throws E {
+      return visitor.visit(this);
+    }
+  }
+
   @Value.Immutable
   abstract static class SingleOrList implements Expression {
     public abstract Expression condition();

diff --git a/core/src/main/java/io/substrait/expression/ExpressionCreator.java b/core/src/main/java/io/substrait/expression/ExpressionCreator.java
@@ -313,7 +313,7 @@ public static AggregateFunctionInvocation aggregateFunction(
         .build();
   }
 
-  public static WindowFunctionInvocation windowFunction(
+  public static Expression.WindowFunctionInvocation windowFunction(
       SimpleExtension.WindowFunctionVariant declaration,
       Type outputType,
       Expression.AggregationPhase phase,
@@ -323,7 +323,7 @@ public static WindowFunctionInvocation windowFunction(
       WindowBound lowerBound,
       WindowBound upperBound,
       Iterable<? extends FunctionArg> arguments) {
-    return WindowFunctionInvocation.builder()
+    return Expression.WindowFunctionInvocation.builder()
         .declaration(declaration)
         .outputType(outputType)
         .aggregationPhase(phase)
@@ -336,7 +336,7 @@ public static WindowFunctionInvocation windowFunction(
         .build();
   }
 
-  public static WindowFunctionInvocation windowFunction(
+  public static Expression.WindowFunctionInvocation windowFunction(
       SimpleExtension.WindowFunctionVariant declaration,
       Type outputType,
       Expression.AggregationPhase phase,
@@ -346,7 +346,7 @@ public static WindowFunctionInvocation windowFunction(
       WindowBound lowerBound,
       WindowBound upperBound,
       FunctionArg... arguments) {
-    return WindowFunctionInvocation.builder()
+    return Expression.WindowFunctionInvocation.builder()
         .declaration(declaration)
         .outputType(outputType)
         .aggregationPhase(phase)

diff --git a/core/src/main/java/io/substrait/expression/ExpressionVisitor.java b/core/src/main/java/io/substrait/expression/ExpressionVisitor.java
@@ -57,7 +57,7 @@ public interface ExpressionVisitor<R, E extends Throwable> {
 
   R visit(Expression.ScalarFunctionInvocation expr) throws E;
 
-  R visit(WindowFunctionInvocation expr) throws E;
+  R visit(Expression.WindowFunctionInvocation expr) throws E;
 
   R visit(Expression.Cast expr) throws E;
 

diff --git a/core/src/main/java/io/substrait/expression/WindowFunctionInvocation.java b/core/src/main/java/io/substrait/expression/WindowFunctionInvocation.java
diff --git a/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java b/core/src/main/java/io/substrait/expression/proto/ExpressionProtoConverter.java
@@ -416,7 +416,7 @@ public Expression visit(io.substrait.expression.Expression.InPredicate expr)
         .build();
   }
 
-  public Expression visit(io.substrait.expression.WindowFunctionInvocation expr)
+  public Expression visit(io.substrait.expression.Expression.WindowFunctionInvocation expr)
       throws RuntimeException {
     var argVisitor = FunctionArg.toProto(typeProtoConverter, this);
     List<FunctionArgument> args =

diff --git a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java
@@ -6,7 +6,6 @@
 import io.substrait.expression.FunctionArg;
 import io.substrait.expression.ImmutableExpression;
 import io.substrait.expression.WindowBound;
-import io.substrait.expression.WindowFunctionInvocation;
 import io.substrait.extension.ExtensionLookup;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.relation.ProtoRelConverter;
@@ -142,7 +141,7 @@ public Expression from(io.substrait.proto.Expression expr) {
         WindowBound lowerBound = toLowerBound(windowFunction.getLowerBound());
         WindowBound upperBound = toUpperBound(windowFunction.getUpperBound());
 
-        yield WindowFunctionInvocation.builder()
+        yield Expression.WindowFunctionInvocation.builder()
             .arguments(args)
             .declaration(declaration)
             .outputType(protoTypeConverter.from(windowFunction.getOutputType()))

diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java
@@ -244,7 +244,7 @@ public RexNode visit(Expression.ScalarFunctionInvocation expr) throws RuntimeExc
   }
 
   @Override
-  public RexNode visit(WindowFunctionInvocation expr) throws RuntimeException {
+  public RexNode visit(Expression.WindowFunctionInvocation expr) throws RuntimeException {
     // todo:to construct the RexOver
     return visitFallback(expr);
   }

diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/WindowFunctionConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/WindowFunctionConverter.java
@@ -6,7 +6,6 @@
 import io.substrait.expression.FunctionArg;
 import io.substrait.expression.ImmutableWindowBound;
 import io.substrait.expression.WindowBound;
-import io.substrait.expression.WindowFunctionInvocation;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.type.Type;
 import java.math.BigDecimal;
@@ -29,7 +28,7 @@
 public class WindowFunctionConverter
     extends FunctionConverter<
         SimpleExtension.WindowFunctionVariant,
-        WindowFunctionInvocation,
+        Expression.WindowFunctionInvocation,
         WindowFunctionConverter.WrappedWindowCall> {
 
   @Override
@@ -43,7 +42,7 @@ public WindowFunctionConverter(
   }
 
   @Override
-  protected WindowFunctionInvocation generateBinding(
+  protected Expression.WindowFunctionInvocation generateBinding(
       WrappedWindowCall call,
       SimpleExtension.WindowFunctionVariant function,
       List<FunctionArg> arguments,
@@ -82,7 +81,7 @@ protected WindowFunctionInvocation generateBinding(
         arguments);
   }
 
-  public Optional<WindowFunctionInvocation> convert(
+  public Optional<Expression.WindowFunctionInvocation> convert(
       RexOver over,
       Function<RexNode, Expression> topLevelConverter,
       RexExpressionConverter rexExpressionConverter) {