Skip to content

Commit 2ea8a71

Browse files
johnedquinnjohnedquinn
committed
Splits out Lag/Lead into own classes
1 parent bf7b73c commit 2ea8a71

File tree

6 files changed

+186
-57
lines changed

6 files changed

+186
-57
lines changed

‎partiql-ast/api/partiql-ast.api‎

Lines changed: 44 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -443,8 +443,10 @@ public abstract class org/partiql/ast/AstRewriter : org/partiql/ast/AstVisitor {
443443
public fun visitWindowDefinition (Lorg/partiql/ast/WindowClause$WindowDefinition;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
444444
public synthetic fun visitWindowFunctionTypeCumeDist (Lorg/partiql/ast/WindowFunctionType$CumeDist;Ljava/lang/Object;)Ljava/lang/Object;
445445
public fun visitWindowFunctionTypeCumeDist (Lorg/partiql/ast/WindowFunctionType$CumeDist;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
446-
public synthetic fun visitWindowFunctionTypeLagOrLead (Lorg/partiql/ast/WindowFunctionType$LeadOrLag;Ljava/lang/Object;)Ljava/lang/Object;
447-
public fun visitWindowFunctionTypeLagOrLead (Lorg/partiql/ast/WindowFunctionType$LeadOrLag;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
446+
public synthetic fun visitWindowFunctionTypeLag (Lorg/partiql/ast/WindowFunctionType$Lag;Ljava/lang/Object;)Ljava/lang/Object;
447+
public fun visitWindowFunctionTypeLag (Lorg/partiql/ast/WindowFunctionType$Lag;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
448+
public synthetic fun visitWindowFunctionTypeLead (Lorg/partiql/ast/WindowFunctionType$Lead;Ljava/lang/Object;)Ljava/lang/Object;
449+
public fun visitWindowFunctionTypeLead (Lorg/partiql/ast/WindowFunctionType$Lead;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
448450
public synthetic fun visitWindowFunctionTypeRank (Lorg/partiql/ast/WindowFunctionType$Rank;Ljava/lang/Object;)Ljava/lang/Object;
449451
public fun visitWindowFunctionTypeRank (Lorg/partiql/ast/WindowFunctionType$Rank;Ljava/lang/Object;)Lorg/partiql/ast/AstNode;
450452
public synthetic fun visitWindowFunctionTypeRowNumber (Lorg/partiql/ast/WindowFunctionType$RowNumber;Ljava/lang/Object;)Ljava/lang/Object;
@@ -606,7 +608,8 @@ public abstract class org/partiql/ast/AstVisitor {
606608
public fun visitWindowDefinition (Lorg/partiql/ast/WindowClause$WindowDefinition;Ljava/lang/Object;)Ljava/lang/Object;
607609
public fun visitWindowFunctionType (Lorg/partiql/ast/WindowFunctionType;Ljava/lang/Object;)Ljava/lang/Object;
608610
public fun visitWindowFunctionTypeCumeDist (Lorg/partiql/ast/WindowFunctionType$CumeDist;Ljava/lang/Object;)Ljava/lang/Object;
609-
public fun visitWindowFunctionTypeLagOrLead (Lorg/partiql/ast/WindowFunctionType$LeadOrLag;Ljava/lang/Object;)Ljava/lang/Object;
611+
public fun visitWindowFunctionTypeLag (Lorg/partiql/ast/WindowFunctionType$Lag;Ljava/lang/Object;)Ljava/lang/Object;
612+
public fun visitWindowFunctionTypeLead (Lorg/partiql/ast/WindowFunctionType$Lead;Ljava/lang/Object;)Ljava/lang/Object;
610613
public fun visitWindowFunctionTypeRank (Lorg/partiql/ast/WindowFunctionType$Rank;Ljava/lang/Object;)Ljava/lang/Object;
611614
public fun visitWindowFunctionTypeRowNumber (Lorg/partiql/ast/WindowFunctionType$RowNumber;Ljava/lang/Object;)Ljava/lang/Object;
612615
public fun visitWindowPartition (Lorg/partiql/ast/WindowPartition;Ljava/lang/Object;)Ljava/lang/Object;
@@ -1620,31 +1623,57 @@ public class org/partiql/ast/WindowFunctionType$CumeDist$Builder {
16201623
public fun toString ()Ljava/lang/String;
16211624
}
16221625

1623-
public final class org/partiql/ast/WindowFunctionType$LeadOrLag : org/partiql/ast/WindowFunctionType {
1624-
public fun <init> (ZLorg/partiql/ast/expr/Expr;Ljava/lang/Long;Lorg/partiql/ast/expr/Expr;Lorg/partiql/ast/WindowFunctionNullTreatment;)V
1626+
public final class org/partiql/ast/WindowFunctionType$Lag : org/partiql/ast/WindowFunctionType, org/partiql/ast/WindowFunctionType$LeadOrLag {
1627+
public fun <init> (Lorg/partiql/ast/expr/Expr;Ljava/lang/Long;Lorg/partiql/ast/expr/Expr;Lorg/partiql/ast/WindowFunctionNullTreatment;)V
16251628
public fun accept (Lorg/partiql/ast/AstVisitor;Ljava/lang/Object;)Ljava/lang/Object;
1626-
public static fun builder ()Lorg/partiql/ast/WindowFunctionType$LeadOrLag$Builder;
1629+
public static fun builder ()Lorg/partiql/ast/WindowFunctionType$Lag$Builder;
16271630
public fun equals (Ljava/lang/Object;)Z
16281631
public fun getChildren ()Ljava/util/List;
16291632
public fun getDefaultValue ()Lorg/partiql/ast/expr/Expr;
16301633
public fun getExtent ()Lorg/partiql/ast/expr/Expr;
16311634
public fun getNullTreatment ()Lorg/partiql/ast/WindowFunctionNullTreatment;
16321635
public fun getOffset ()Ljava/lang/Long;
16331636
public fun hashCode ()I
1634-
public fun isLag ()Z
1635-
public fun isLead ()Z
16361637
}
16371638

1638-
public class org/partiql/ast/WindowFunctionType$LeadOrLag$Builder {
1639-
public fun build ()Lorg/partiql/ast/WindowFunctionType$LeadOrLag;
1640-
public fun defaultValue (Lorg/partiql/ast/expr/Expr;)Lorg/partiql/ast/WindowFunctionType$LeadOrLag$Builder;
1641-
public fun extent (Lorg/partiql/ast/expr/Expr;)Lorg/partiql/ast/WindowFunctionType$LeadOrLag$Builder;
1642-
public fun isLead (Z)Lorg/partiql/ast/WindowFunctionType$LeadOrLag$Builder;
1643-
public fun nullTreatment (Lorg/partiql/ast/WindowFunctionNullTreatment;)Lorg/partiql/ast/WindowFunctionType$LeadOrLag$Builder;
1644-
public fun offset (Ljava/lang/Long;)Lorg/partiql/ast/WindowFunctionType$LeadOrLag$Builder;
1639+
public class org/partiql/ast/WindowFunctionType$Lag$Builder {
1640+
public fun build ()Lorg/partiql/ast/WindowFunctionType$Lag;
1641+
public fun defaultValue (Lorg/partiql/ast/expr/Expr;)Lorg/partiql/ast/WindowFunctionType$Lag$Builder;
1642+
public fun extent (Lorg/partiql/ast/expr/Expr;)Lorg/partiql/ast/WindowFunctionType$Lag$Builder;
1643+
public fun nullTreatment (Lorg/partiql/ast/WindowFunctionNullTreatment;)Lorg/partiql/ast/WindowFunctionType$Lag$Builder;
1644+
public fun offset (Ljava/lang/Long;)Lorg/partiql/ast/WindowFunctionType$Lag$Builder;
16451645
public fun toString ()Ljava/lang/String;
16461646
}
16471647

1648+
public final class org/partiql/ast/WindowFunctionType$Lead : org/partiql/ast/WindowFunctionType, org/partiql/ast/WindowFunctionType$LeadOrLag {
1649+
public fun <init> (Lorg/partiql/ast/expr/Expr;Ljava/lang/Long;Lorg/partiql/ast/expr/Expr;Lorg/partiql/ast/WindowFunctionNullTreatment;)V
1650+
public fun accept (Lorg/partiql/ast/AstVisitor;Ljava/lang/Object;)Ljava/lang/Object;
1651+
public static fun builder ()Lorg/partiql/ast/WindowFunctionType$Lead$Builder;
1652+
public fun equals (Ljava/lang/Object;)Z
1653+
public fun getChildren ()Ljava/util/List;
1654+
public fun getDefaultValue ()Lorg/partiql/ast/expr/Expr;
1655+
public fun getExtent ()Lorg/partiql/ast/expr/Expr;
1656+
public fun getNullTreatment ()Lorg/partiql/ast/WindowFunctionNullTreatment;
1657+
public fun getOffset ()Ljava/lang/Long;
1658+
public fun hashCode ()I
1659+
}
1660+
1661+
public class org/partiql/ast/WindowFunctionType$Lead$Builder {
1662+
public fun build ()Lorg/partiql/ast/WindowFunctionType$Lead;
1663+
public fun defaultValue (Lorg/partiql/ast/expr/Expr;)Lorg/partiql/ast/WindowFunctionType$Lead$Builder;
1664+
public fun extent (Lorg/partiql/ast/expr/Expr;)Lorg/partiql/ast/WindowFunctionType$Lead$Builder;
1665+
public fun nullTreatment (Lorg/partiql/ast/WindowFunctionNullTreatment;)Lorg/partiql/ast/WindowFunctionType$Lead$Builder;
1666+
public fun offset (Ljava/lang/Long;)Lorg/partiql/ast/WindowFunctionType$Lead$Builder;
1667+
public fun toString ()Ljava/lang/String;
1668+
}
1669+
1670+
public abstract interface class org/partiql/ast/WindowFunctionType$LeadOrLag {
1671+
public abstract fun getDefaultValue ()Lorg/partiql/ast/expr/Expr;
1672+
public abstract fun getExtent ()Lorg/partiql/ast/expr/Expr;
1673+
public abstract fun getNullTreatment ()Lorg/partiql/ast/WindowFunctionNullTreatment;
1674+
public abstract fun getOffset ()Ljava/lang/Long;
1675+
}
1676+
16481677
public final class org/partiql/ast/WindowFunctionType$Rank : org/partiql/ast/WindowFunctionType {
16491678
public static field DENSE_RANK I
16501679
public static field PERCENT_RANK I

‎partiql-ast/src/main/java/org/partiql/ast/AstVisitor.java‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -400,7 +400,17 @@ public R visitWindowPartitionName(WindowPartition.Name node, C ctx) {
400400
* @param ctx TODO
401401
* @return TODO
402402
*/
403-
public R visitWindowFunctionTypeLagOrLead(WindowFunctionType.LeadOrLag node, C ctx) {
403+
public R visitWindowFunctionTypeLead(WindowFunctionType.Lead node, C ctx) {
404+
return defaultVisit(node, ctx);
405+
}
406+
407+
/**
408+
* TODO
409+
* @param node TODO
410+
* @param ctx TODO
411+
* @return TODO
412+
*/
413+
public R visitWindowFunctionTypeLag(WindowFunctionType.Lag node, C ctx) {
404414
return defaultVisit(node, ctx);
405415
}
406416

‎partiql-ast/src/main/java/org/partiql/ast/WindowFunctionType.java‎

Lines changed: 109 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -120,86 +120,160 @@ public <R, C> R accept(@NotNull AstVisitor<R, C> visitor, C ctx) {
120120
/**
121121
* Represents either the LAG or LEAD window function.
122122
* @see ExprWindowFunction#getFunctionType()
123+
* @see Lead
124+
* @see Lag
125+
*/
126+
public interface LeadOrLag {
127+
/**
128+
* Returns the extent of the window function.
129+
* @return the extent of the window function
130+
*/
131+
@NotNull
132+
Expr getExtent();
133+
134+
/**
135+
* Returns the offset of the window function.
136+
* @return the offset of the window function
137+
*/
138+
@Nullable
139+
Long getOffset();
140+
141+
/**
142+
* Returns the default value of the window function.
143+
* @return the default value of the window function
144+
*/
145+
@Nullable
146+
Expr getDefaultValue();
147+
148+
/**
149+
* Returns the null treatment of the window function.
150+
* @return the null treatment of the window function
151+
*/
152+
@Nullable
153+
WindowFunctionNullTreatment getNullTreatment();
154+
}
155+
156+
/**
157+
* Represents the LEAD window function.
158+
* @see ExprWindowFunction#getFunctionType()
123159
*/
124160
@Builder(builderClassName = "Builder")
125161
@EqualsAndHashCode(callSuper = false)
126-
public static final class LeadOrLag extends WindowFunctionType {
127-
private final boolean isLead;
162+
public static final class Lead extends WindowFunctionType implements LeadOrLag {
128163
private final Expr extent;
129164
private final Long offset;
130165
private final Expr defaultValue;
131166
private final WindowFunctionNullTreatment nullTreatment;
132-
133167
/**
134-
* Constructs a new LEAD or LAG window function type.
135-
* @param isLead whether this is a LEAD or LAG window function
168+
* Constructs a new LEAD window function type.
136169
* @param extent the extent of the window function
137170
* @param offset the offset of the window function
138171
* @param defaultValue the default value of the window function
139172
* @param nullTreatment the null treatment of the window function
140173
*/
141-
public LeadOrLag(
142-
boolean isLead,
174+
public Lead(
143175
@NotNull Expr extent,
144176
@Nullable Long offset,
145177
@Nullable Expr defaultValue,
146178
@Nullable WindowFunctionNullTreatment nullTreatment
147179
) {
148-
this.isLead = isLead;
180+
super();
149181
this.extent = extent;
150182
this.offset = offset;
151183
this.defaultValue = defaultValue;
152184
this.nullTreatment = nullTreatment;
153185
}
154186

155-
/**
156-
* Returns whether this is a LEAD or LAG window function.
157-
* @return whether this is a LEAD or LAG window function
158-
*/
159-
public boolean isLead() {
160-
return this.isLead;
187+
@NotNull
188+
@Override
189+
public Expr getExtent() {
190+
return this.extent;
161191
}
162192

163-
/**
164-
* Returns whether this is a LEAD or LAG window function.
165-
* @return whether this is a LEAD or LAG window function
166-
*/
167-
public boolean isLag() {
168-
return !this.isLead;
193+
@Nullable
194+
@Override
195+
public Long getOffset() {
196+
return this.offset;
169197
}
170198

199+
@Nullable
200+
@Override
201+
public Expr getDefaultValue() {
202+
return this.defaultValue;
203+
}
204+
205+
@Nullable
206+
@Override
207+
public WindowFunctionNullTreatment getNullTreatment() {
208+
return this.nullTreatment;
209+
}
210+
211+
@NotNull
212+
@Override
213+
public List<AstNode> getChildren() {
214+
List<AstNode> kids = new ArrayList<>();
215+
kids.add(extent);
216+
kids.add(defaultValue);
217+
return kids;
218+
}
219+
220+
@Override
221+
public <R, C> R accept(@NotNull AstVisitor<R, C> visitor, C ctx) {
222+
return visitor.visitWindowFunctionTypeLead(this, ctx);
223+
}
224+
}
225+
226+
/**
227+
* Represents the LAG window function.
228+
* @see ExprWindowFunction#getFunctionType()
229+
*/
230+
@Builder(builderClassName = "Builder")
231+
@EqualsAndHashCode(callSuper = false)
232+
public static final class Lag extends WindowFunctionType implements LeadOrLag {
233+
private final Expr extent;
234+
private final Long offset;
235+
private final Expr defaultValue;
236+
private final WindowFunctionNullTreatment nullTreatment;
171237
/**
172-
* Returns the extent of the window function.
173-
* @return the extent of the window function
238+
* Constructs a new LAG window function type.
239+
* @param extent the extent of the window function
240+
* @param offset the offset of the window function
241+
* @param defaultValue the default value of the window function
242+
* @param nullTreatment the null treatment of the window function
174243
*/
244+
public Lag(
245+
@NotNull Expr extent,
246+
@Nullable Long offset,
247+
@Nullable Expr defaultValue,
248+
@Nullable WindowFunctionNullTreatment nullTreatment
249+
) {
250+
super();
251+
this.extent = extent;
252+
this.offset = offset;
253+
this.defaultValue = defaultValue;
254+
this.nullTreatment = nullTreatment;
255+
}
256+
175257
@NotNull
258+
@Override
176259
public Expr getExtent() {
177260
return this.extent;
178261
}
179262

180-
/**
181-
* Returns the offset of the window function.
182-
* @return the offset of the window function
183-
*/
184263
@Nullable
264+
@Override
185265
public Long getOffset() {
186266
return this.offset;
187267
}
188268

189-
/**
190-
* Returns the default value of the window function.
191-
* @return the default value of the window function
192-
*/
193269
@Nullable
270+
@Override
194271
public Expr getDefaultValue() {
195272
return this.defaultValue;
196273
}
197274

198-
/**
199-
* Returns the null treatment of the window function.
200-
* @return the null treatment of the window function
201-
*/
202275
@Nullable
276+
@Override
203277
public WindowFunctionNullTreatment getNullTreatment() {
204278
return this.nullTreatment;
205279
}
@@ -215,7 +289,7 @@ public List<AstNode> getChildren() {
215289

216290
@Override
217291
public <R, C> R accept(@NotNull AstVisitor<R, C> visitor, C ctx) {
218-
return visitor.visitWindowFunctionTypeLagOrLead(this, ctx);
292+
return visitor.visitWindowFunctionTypeLag(this, ctx);
219293
}
220294
}
221295
}

‎partiql-ast/src/main/kotlin/org/partiql/ast/AstRewriter.kt‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -820,11 +820,21 @@ public abstract class AstRewriter<C> : AstVisitor<AstNode, C>() {
820820
}
821821
}
822822

823-
override fun visitWindowFunctionTypeLagOrLead(node: WindowFunctionType.LeadOrLag, ctx: C): AstNode {
823+
override fun visitWindowFunctionTypeLag(node: WindowFunctionType.Lag, ctx: C): AstNode {
824824
val extent = visitExpr(node.extent, ctx) as Expr
825825
val default = node.defaultValue?.let { visitExpr(it, ctx) as Expr }
826826
return if (extent !== node.extent || default !== node.defaultValue) {
827-
WindowFunctionType.LeadOrLag(node.isLead, extent, node.offset, default, node.nullTreatment)
827+
WindowFunctionType.Lag(extent, node.offset, default, node.nullTreatment)
828+
} else {
829+
node
830+
}
831+
}
832+
833+
override fun visitWindowFunctionTypeLead(node: WindowFunctionType.Lead, ctx: C): AstNode {
834+
val extent = visitExpr(node.extent, ctx) as Expr
835+
val default = node.defaultValue?.let { visitExpr(it, ctx) as Expr }
836+
return if (extent !== node.extent || default !== node.defaultValue) {
837+
WindowFunctionType.Lead(extent, node.offset, default, node.nullTreatment)
828838
} else {
829839
node
830840
}

‎partiql-parser/src/main/kotlin/org/partiql/parser/internal/PartiQLParserDefault.kt‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2031,7 +2031,10 @@ internal class PartiQLParserDefault : PartiQLParser {
20312031
val offset = ctx.offset?.text?.toLong()
20322032
val default = ctx.default_?.let { visitExpr(it) }
20332033
val nullTreatment = ctx.windowFunctionNullTreatment()?.let { visitWindowFunctionNullTreatment(it) }
2034-
WindowFunctionType.LeadOrLag(isLead, extent, offset, default, nullTreatment)
2034+
when (isLead) {
2035+
true -> WindowFunctionType.Lead(extent, offset, default, nullTreatment)
2036+
false -> WindowFunctionType.Lag(extent, offset, default, nullTreatment)
2037+
}
20352038
}
20362039

20372040
override fun visitWindowPartitionColumnReference(ctx: GeneratedParser.WindowPartitionColumnReferenceContext) = translate(ctx) {

‎partiql-planner/src/main/kotlin/org/partiql/planner/internal/transforms/RelConverter.kt‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ import org.partiql.ast.SetQuantifier
4848
import org.partiql.ast.Sort
4949
import org.partiql.ast.WindowFunctionNullTreatment
5050
import org.partiql.ast.WindowFunctionType
51+
import org.partiql.ast.WindowFunctionType.Lag
52+
import org.partiql.ast.WindowFunctionType.Lead
5153
import org.partiql.ast.WindowPartition
5254
import org.partiql.ast.WindowSpecification
5355
import org.partiql.ast.With
@@ -665,9 +667,10 @@ internal object RelConverter {
665667
relOpWindowWindowFunction("row_number", emptyList(), false, emptyList(), CompilerType(PType.dynamic()))
666668
}
667669
is WindowFunctionType.LeadOrLag -> {
668-
val name = when (windowType.isLag) {
669-
true -> "lag"
670-
false -> "lead"
670+
val name = when (windowType) {
671+
is Lag -> "lag"
672+
is Lead -> "lead"
673+
else -> error("Unsupported Lead/Lag type: ${windowType.javaClass.simpleName}.")
671674
}
672675
val isIgnoreNulls = when (windowType.nullTreatment?.code()) {
673676
WindowFunctionNullTreatment.RESPECT_NULLS, null -> false

0 commit comments

Comments
 (0)