Skip to content

Commit 19f0aa4

Browse files
manticore-projectsmanticore-projects
committed
feat: improve TableSample Operator
- support DuckDB - support Google BigQuery - support `ParenthesedSelect` `FromItem` Signed-off-by: Andreas Reichel <[email protected]>
1 parent b1b86f5 commit 19f0aa4

File tree

15 files changed

+186
-16
lines changed

15 files changed

+186
-16
lines changed

‎src/main/java/net/sf/jsqlparser/statement/piped/FromQuery.java

+11
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import net.sf.jsqlparser.statement.select.FromItem;
66
import net.sf.jsqlparser.statement.select.FromItemVisitor;
77
import net.sf.jsqlparser.statement.select.Pivot;
8+
import net.sf.jsqlparser.statement.select.SampleClause;
89
import net.sf.jsqlparser.statement.select.Select;
910
import net.sf.jsqlparser.statement.select.SelectVisitor;
1011
import net.sf.jsqlparser.statement.select.UnPivot;
@@ -188,6 +189,16 @@ public void setUnPivot(UnPivot unpivot) {
188189

189190
}
190191

192+
@Override
193+
public SampleClause getSampleClause() {
194+
return null;
195+
}
196+
197+
@Override
198+
public FromItem setSampleClause(SampleClause sampleClause) {
199+
return null;
200+
}
201+
191202
@Override
192203
public <T, S> T accept(ExpressionVisitor<T> expressionVisitor, S context) {
193204
return expressionVisitor.visit(this, context);

‎src/main/java/net/sf/jsqlparser/statement/select/FromItem.java

+11-2
Original file line numberDiff line numberDiff line change
@@ -47,16 +47,25 @@ default FromItem withUnPivot(UnPivot unpivot) {
4747
return this;
4848
}
4949

50+
SampleClause getSampleClause();
51+
52+
FromItem setSampleClause(SampleClause sampleClause);
53+
5054
default StringBuilder appendTo(StringBuilder builder, Alias alias) {
51-
return appendTo(builder, alias, null, null);
55+
return appendTo(builder, alias, null, null, null);
5256
}
5357

54-
default StringBuilder appendTo(StringBuilder builder, Alias alias, Pivot pivot,
58+
default StringBuilder appendTo(StringBuilder builder, Alias alias, SampleClause sampleClause,
59+
Pivot pivot,
5560
UnPivot unPivot) {
5661
if (alias != null) {
5762
builder.append(alias);
5863
}
5964

65+
if (sampleClause != null) {
66+
builder.append(sampleClause);
67+
}
68+
6069
if (pivot != null) {
6170
builder.append(" ").append(pivot);
6271
}

‎src/main/java/net/sf/jsqlparser/statement/select/ParenthesedFromItem.java

+17
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ public class ParenthesedFromItem extends ASTNodeAccessImpl implements FromItem {
2323
private Alias alias;
2424
private Pivot pivot;
2525
private UnPivot unPivot;
26+
private SampleClause sampleClause;
2627

2728
public ParenthesedFromItem() {}
2829

@@ -130,6 +131,22 @@ public void setUnPivot(UnPivot unpivot) {
130131
this.unPivot = unpivot;
131132
}
132133

134+
@Override
135+
public SampleClause getSampleClause() {
136+
return sampleClause;
137+
}
138+
139+
@Override
140+
public FromItem setSampleClause(SampleClause sampleClause) {
141+
this.sampleClause = sampleClause;
142+
return this;
143+
}
144+
145+
public ParenthesedFromItem withSampleClause(SampleClause sampleClause) {
146+
this.sampleClause = sampleClause;
147+
return this;
148+
}
149+
133150
public ParenthesedFromItem withFromItem(FromItem fromItem) {
134151
this.setFromItem(fromItem);
135152
return this;

‎src/main/java/net/sf/jsqlparser/statement/select/ParenthesedSelect.java

+20-2
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ public class ParenthesedSelect extends Select implements FromItem, ParenthesedSt
2323
Pivot pivot;
2424
UnPivot unPivot;
2525
Select select;
26+
SampleClause sampleClause = null;
2627

2728
public ParenthesedSelect() {}
2829

@@ -79,7 +80,8 @@ private static Alias getAliasFromItem(FromItem fromItem) {
7980
TableFunction t = (TableFunction) fromItem;
8081
return new Alias(t.getName(), true);
8182
} else {
82-
return fromItem.getAlias()!=null ? new Alias(fromItem.getAlias().getName(), true) : null;
83+
return fromItem.getAlias() != null ? new Alias(fromItem.getAlias().getName(), true)
84+
: null;
8385
}
8486
}
8587

@@ -116,6 +118,22 @@ public void setUnPivot(UnPivot unPivot) {
116118
this.unPivot = unPivot;
117119
}
118120

121+
@Override
122+
public SampleClause getSampleClause() {
123+
return sampleClause;
124+
}
125+
126+
@Override
127+
public FromItem setSampleClause(SampleClause sampleClause) {
128+
this.sampleClause = sampleClause;
129+
return this;
130+
}
131+
132+
public ParenthesedSelect withSampleClause(SampleClause sampleClause) {
133+
this.sampleClause = sampleClause;
134+
return this;
135+
}
136+
119137
public Select getSelect() {
120138
return select;
121139
}
@@ -163,7 +181,7 @@ public <T, S> T accept(StatementVisitor<T> statementVisitor, S context) {
163181

164182
public StringBuilder appendSelectBodyTo(StringBuilder builder) {
165183
builder.append("(").append(select).append(")");
166-
appendTo(builder, alias, pivot, unPivot);
184+
appendTo(builder, alias, sampleClause, pivot, unPivot);
167185
return builder;
168186
}
169187
}

‎src/main/java/net/sf/jsqlparser/statement/select/PlainSelect.java

+10
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,16 @@ public <T, S> T accept(FromItemVisitor<T> fromItemVisitor, S context) {
324324
return fromItemVisitor.visit(this, context);
325325
}
326326

327+
@Override
328+
public SampleClause getSampleClause() {
329+
return null;
330+
}
331+
332+
@Override
333+
public FromItem setSampleClause(SampleClause sampleClause) {
334+
return null;
335+
}
336+
327337
public OptimizeFor getOptimizeFor() {
328338
return optimizeFor;
329339
}

‎src/main/java/net/sf/jsqlparser/statement/select/SampleClause.java

+32-6
Original file line numberDiff line numberDiff line change
@@ -13,25 +13,28 @@ public class SampleClause {
1313
private SampleKeyword keyword;
1414
private SampleMethod method;
1515
private Number percentageArgument;
16+
private String percentageUnit;
1617
private Number repeatArgument;
1718
// Oracle Specific
1819
private Number seedArgument;
1920

2021
public SampleClause(String keyword, String method, Number percentageArgument,
22+
String percentageUnit,
2123
Number repeatArgument, Number seedArgument) {
2224
this.keyword = SampleKeyword.from(keyword);
2325
this.method = method == null || method.length() == 0 ? null : SampleMethod.from(method);
2426
this.percentageArgument = percentageArgument;
27+
this.percentageUnit = percentageUnit;
2528
this.repeatArgument = repeatArgument;
2629
this.seedArgument = seedArgument;
2730
}
2831

2932
public SampleClause() {
30-
this(SampleKeyword.TABLESAMPLE.toString(), null, null, null, null);
33+
this(SampleKeyword.TABLESAMPLE.toString(), null, null, null, null, null);
3134
}
3235

3336
public SampleClause(String keyword) {
34-
this(keyword, null, null, null, null);
37+
this(keyword, null, null, null, null, null);
3538
}
3639

3740
public SampleKeyword getKeyword() {
@@ -56,6 +59,15 @@ public Number getRepeatArgument() {
5659
return repeatArgument;
5760
}
5861

62+
public String getPercentageUnit() {
63+
return percentageUnit;
64+
}
65+
66+
public SampleClause setPercentageUnit(String percentageUnit) {
67+
this.percentageUnit = percentageUnit;
68+
return this;
69+
}
70+
5971
public SampleClause setRepeatArgument(Number repeatArgument) {
6072
this.repeatArgument = repeatArgument;
6173
return this;
@@ -92,7 +104,8 @@ public StringBuilder appendTo(StringBuilder builder) {
92104
}
93105

94106
if (percentageArgument != null) {
95-
builder.append(" (").append(percentageArgument).append(")");
107+
builder.append(" (").append(percentageArgument)
108+
.append(percentageUnit != null ? " " + percentageUnit : "").append(")");
96109
}
97110

98111
if (repeatArgument != null) {
@@ -111,18 +124,31 @@ public String toString() {
111124
}
112125

113126
public enum SampleKeyword {
114-
SAMPLE, TABLESAMPLE;
127+
SAMPLE("SAMPLE"), TABLESAMPLE("TABLESAMPLE"), USING_SAMPLE("USING SAMPLE");
128+
129+
String keyword;
130+
131+
SampleKeyword(String keyword) {
132+
this.keyword = keyword;
133+
}
115134

116135
public static SampleKeyword from(String sampleKeyword) {
117-
return Enum.valueOf(SampleKeyword.class, sampleKeyword.toUpperCase());
136+
return Enum.valueOf(SampleKeyword.class,
137+
sampleKeyword.toUpperCase().replaceAll(" ", "_"));
138+
}
139+
140+
@Override
141+
public String toString() {
142+
return keyword;
118143
}
119144
}
120145

121146
public enum SampleMethod {
122147
BERNOULLI, SYSTEM, BLOCK;
123148

124149
public static SampleMethod from(String sampleMethod) {
125-
return Enum.valueOf(SampleMethod.class, sampleMethod.toUpperCase());
150+
return Enum.valueOf(SampleMethod.class,
151+
sampleMethod.toUpperCase().replaceAll(" ", "_"));
126152
}
127153
}
128154
}

‎src/main/java/net/sf/jsqlparser/statement/select/Select.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -379,7 +379,7 @@ public StringBuilder appendTo(StringBuilder builder) {
379379

380380
appendSelectBodyTo(builder);
381381

382-
appendTo(builder, alias, pivot, unPivot);
382+
appendTo(builder, alias, null, pivot, unPivot);
383383

384384
if (forClause != null) {
385385
forClause.appendTo(builder);

‎src/main/java/net/sf/jsqlparser/statement/select/SetOperationList.java

+10
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,16 @@ public <T, S> T accept(FromItemVisitor<T> fromItemVisitor, S context) {
3232
return fromItemVisitor.visit(this, context);
3333
}
3434

35+
@Override
36+
public SampleClause getSampleClause() {
37+
return null;
38+
}
39+
40+
@Override
41+
public FromItem setSampleClause(SampleClause sampleClause) {
42+
return null;
43+
}
44+
3545
public List<OrderByElement> getOrderByElements() {
3646
return orderByElements;
3747
}

‎src/main/java/net/sf/jsqlparser/statement/select/TableFunction.java

+10
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,16 @@ public TableFunction withUnPivot(UnPivot unpivot) {
112112
return (TableFunction) FromItem.super.withUnPivot(unpivot);
113113
}
114114

115+
@Override
116+
public SampleClause getSampleClause() {
117+
return null;
118+
}
119+
120+
@Override
121+
public FromItem setSampleClause(SampleClause sampleClause) {
122+
return null;
123+
}
124+
115125
public StringBuilder appendTo(StringBuilder builder) {
116126
if (prefix != null) {
117127
builder.append(prefix).append(" ");

‎src/main/java/net/sf/jsqlparser/statement/select/TableStatement.java

+10
Original file line numberDiff line numberDiff line change
@@ -60,4 +60,14 @@ public <T, S> T accept(SelectVisitor<T> selectVisitor, S context) {
6060
public <T, S> T accept(FromItemVisitor<T> fromItemVisitor, S context) {
6161
return fromItemVisitor.visit(this, context);
6262
}
63+
64+
@Override
65+
public SampleClause getSampleClause() {
66+
return null;
67+
}
68+
69+
@Override
70+
public FromItem setSampleClause(SampleClause sampleClause) {
71+
return null;
72+
}
6373
}

‎src/main/java/net/sf/jsqlparser/statement/select/Values.java

+10
Original file line numberDiff line numberDiff line change
@@ -108,4 +108,14 @@ public UnPivot getUnPivot() {
108108
public void setUnPivot(UnPivot unpivot) {
109109

110110
}
111+
112+
@Override
113+
public SampleClause getSampleClause() {
114+
return null;
115+
}
116+
117+
@Override
118+
public FromItem setSampleClause(SampleClause sampleClause) {
119+
return null;
120+
}
111121
}

‎src/main/java/net/sf/jsqlparser/util/deparser/SelectDeParser.java

+7
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
import net.sf.jsqlparser.statement.select.PivotVisitor;
5656
import net.sf.jsqlparser.statement.select.PivotXml;
5757
import net.sf.jsqlparser.statement.select.PlainSelect;
58+
import net.sf.jsqlparser.statement.select.SampleClause;
5859
import net.sf.jsqlparser.statement.select.Select;
5960
import net.sf.jsqlparser.statement.select.SelectItem;
6061
import net.sf.jsqlparser.statement.select.SelectItemVisitor;
@@ -139,6 +140,12 @@ public <S> StringBuilder visit(ParenthesedSelect select, S context) {
139140
if (alias != null) {
140141
builder.append(alias);
141142
}
143+
144+
SampleClause sampleClause = select.getSampleClause();
145+
if (sampleClause != null) {
146+
builder.append(sampleClause);
147+
}
148+
142149
Pivot pivot = select.getPivot();
143150
if (pivot != null) {
144151
pivot.accept(this, context);

‎src/main/java/net/sf/jsqlparser/util/deparser/TableStatementDeParser.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -101,8 +101,8 @@ public <S> StringBuilder visit(TableStatement tableStatement, S context) {
101101
// TODO UNION
102102

103103
tableStatement.appendTo(
104-
builder, tableStatement.getAlias(), tableStatement.getPivot(),
105-
tableStatement.getUnPivot());
104+
builder, tableStatement.getAlias(), null,
105+
tableStatement.getPivot(), tableStatement.getUnPivot());
106106

107107
return builder;
108108
}

‎src/main/jjtree/net/sf/jsqlparser/parser/JSqlParserCC.jjt

+18-3
Original file line numberDiff line numberDiff line change
@@ -2246,6 +2246,7 @@ SampleClause SampleClause():
22462246
String keyword;
22472247
String method=null;
22482248
Number percentageArgument;
2249+
String percentageUnit = null;
22492250
Number repeatArgument=null;
22502251
Number seedArgument=null;
22512252
}
@@ -2262,16 +2263,30 @@ SampleClause SampleClause():
22622263
token = <K_TABLESAMPLE> { keyword = token.image; }
22632264
( token = <K_SYSTEM> | token = <K_BERNOULLI> ) { method = token.image; }
22642265
)
2266+
|
2267+
(
2268+
// Duck DB
2269+
<K_USING> <K_SAMPLE> { keyword = "USING SAMPLE"; }
2270+
( token = <K_SYSTEM> | token = <K_BERNOULLI> ) { method = token.image; }
2271+
)
22652272
)
22662273

2267-
"(" percentageArgument = Number() ")"
2274+
"(" percentageArgument = Number()
2275+
[
2276+
"%" { percentageUnit="%"; }
2277+
|
2278+
<K_PERCENT> { percentageUnit="PERCENT"; }
2279+
|
2280+
<K_ROWS> { percentageUnit="ROWS"; }
2281+
]
2282+
")"
22682283

22692284
[ LOOKAHEAD(2) <K_REPEATABLE> "(" repeatArgument = Number() ")" ]
22702285

22712286
[ LOOKAHEAD(2) <K_SEED> "(" seedArgument = Number() ")" ]
22722287

22732288
{
2274-
return new SampleClause(keyword, method, percentageArgument, repeatArgument, seedArgument);
2289+
return new SampleClause(keyword, method, percentageArgument, percentageUnit, repeatArgument, seedArgument);
22752290
}
22762291
}
22772292

@@ -3543,7 +3558,7 @@ FromItem FromItem() #FromItem:
35433558
)
35443559

35453560
[ LOOKAHEAD(2) alias=Alias() { fromItem.setAlias(alias); } ]
3546-
[ LOOKAHEAD(2) sampleClause = SampleClause() { ((Table) fromItem).setSampleClause(sampleClause); } ]
3561+
[ LOOKAHEAD(2) sampleClause = SampleClause() { fromItem.setSampleClause(sampleClause); } ]
35473562
[ LOOKAHEAD(2) unpivot=UnPivot() { fromItem.setUnPivot(unpivot); } ]
35483563
[ LOOKAHEAD(2) ( LOOKAHEAD(2) pivot=PivotXml() | pivot=Pivot() ) { fromItem.setPivot(pivot); } ]
35493564
[

0 commit comments

Comments
 (0)