Skip to content

Commit 2c1b2f9

Browse files
Yury-FridlyandYury-Fridlyand
committed
Further work on pagination.
* Added push down page size from `LogicalPaginate` to `LogicalRelation`. * Improved cursor encoding and decoding. * Added cursor compression. * Fixed issuing `SearchScrollRequest`. * Fixed returning last empty page. * Minor code grooming/commenting. Signed-off-by: Yury-Fridlyand <[email protected]>
1 parent 7b508db commit 2c1b2f9

File tree

16 files changed

+205
-78
lines changed

16 files changed

+205
-78
lines changed

‎core/src/main/java/org/opensearch/sql/executor/PaginatedPlanCache.java‎

Lines changed: 76 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -5,27 +5,25 @@
55

66
package org.opensearch.sql.executor;
77

8-
import java.io.IOException;
9-
import java.io.ObjectInputStream;
8+
import com.google.common.hash.HashCode;
9+
import java.io.ByteArrayInputStream;
10+
import java.io.ByteArrayOutputStream;
1011
import java.util.ArrayList;
1112
import java.util.List;
12-
import java.util.stream.Collectors;
13-
import java.util.stream.Stream;
13+
import java.util.zip.GZIPInputStream;
14+
import java.util.zip.GZIPOutputStream;
1415
import lombok.Data;
1516
import lombok.RequiredArgsConstructor;
17+
import lombok.SneakyThrows;
1618
import org.opensearch.sql.ast.tree.UnresolvedPlan;
1719
import org.opensearch.sql.expression.NamedExpression;
18-
import org.opensearch.sql.expression.ReferenceExpression;
1920
import org.opensearch.sql.expression.serialization.DefaultExpressionSerializer;
2021
import org.opensearch.sql.opensearch.executor.Cursor;
2122
import org.opensearch.sql.planner.PaginateOperator;
2223
import org.opensearch.sql.planner.physical.PhysicalPlan;
23-
import org.opensearch.sql.planner.physical.PhysicalPlanNodeVisitor;
2424
import org.opensearch.sql.planner.physical.ProjectOperator;
2525
import org.opensearch.sql.storage.StorageEngine;
26-
import org.opensearch.sql.storage.Table;
2726
import org.opensearch.sql.storage.TableScanOperator;
28-
import org.opensearch.sql.storage.read.TableScanBuilder;
2927

3028
@RequiredArgsConstructor
3129
public class PaginatedPlanCache {
@@ -39,7 +37,7 @@ public boolean canConvertToCursor(UnresolvedPlan plan) {
3937

4038
@RequiredArgsConstructor
4139
@Data
42-
static class SeriazationContext {
40+
static class SerializationContext {
4341
private final PaginatedPlanCache cache;
4442
}
4543

@@ -48,74 +46,104 @@ static class SeriazationContext {
4846
*/
4947
public Cursor convertToCursor(PhysicalPlan plan) {
5048
if (plan instanceof PaginateOperator) {
51-
var raw = CURSOR_PREFIX + plan.toCursor();
49+
var cursor = plan.toCursor();
50+
if (cursor == null || cursor.isEmpty()) {
51+
return Cursor.None;
52+
}
53+
var raw = CURSOR_PREFIX + compress(cursor);
5254
return new Cursor(raw.getBytes());
5355
} else {
5456
return Cursor.None;
5557
}
5658
}
5759

60+
@SneakyThrows
61+
public static String compress(String str) {
62+
if (str == null || str.length() == 0) {
63+
return null;
64+
}
65+
66+
ByteArrayOutputStream out = new ByteArrayOutputStream();
67+
GZIPOutputStream gzip = new GZIPOutputStream(out);
68+
gzip.write(str.getBytes());
69+
gzip.close();
70+
return HashCode.fromBytes(out.toByteArray()).toString();
71+
}
72+
73+
@SneakyThrows
74+
public static String decompress(String input) {
75+
if (input == null || input.length() == 0) {
76+
return null;
77+
}
78+
GZIPInputStream gzip = new GZIPInputStream(new ByteArrayInputStream(
79+
HashCode.fromString(input).asBytes()));
80+
return new String(gzip.readAllBytes());
81+
}
82+
83+
/**
84+
* Parse `NamedExpression`s from cursor.
85+
* @param listToFill List to fill with data.
86+
* @param cursor Cursor to parse.
87+
* @return Remaining part of the cursor.
88+
*/
89+
private String parseNamedExpressions(List<NamedExpression> listToFill, String cursor) {
90+
var serializer = new DefaultExpressionSerializer();
91+
while (!cursor.startsWith(")") && !cursor.startsWith("(")) {
92+
listToFill.add((NamedExpression)
93+
serializer.deserialize(cursor.substring(0,
94+
Math.min(cursor.indexOf(','), cursor.indexOf(')')))));
95+
cursor = cursor.substring(cursor.indexOf(',') + 1);
96+
}
97+
return cursor;
98+
}
99+
58100
/**
59101
* Converts a cursor to a physical plan tree.
60102
*/
61103
public PhysicalPlan convertToPlan(String cursor) {
62104
if (cursor.startsWith(CURSOR_PREFIX)) {
63105
try {
64-
String expression = cursor.substring(CURSOR_PREFIX.length());
65-
66-
// TODO Parse expression and initialize variables below.
67-
// storageEngine needs to create the TableScanOperator.
106+
cursor = cursor.substring(CURSOR_PREFIX.length());
107+
cursor = decompress(cursor);
68108

69109
// TODO Parse with ANTLR or serialize as JSON/XML
70-
if (!expression.startsWith("(Paginate,")) {
110+
if (!cursor.startsWith("(Paginate,")) {
71111
throw new UnsupportedOperationException("Unsupported cursor");
72112
}
73-
expression = expression.substring(expression.indexOf(',') + 1);
74-
int currentPageIndex = Integer.parseInt(expression, 0, expression.indexOf(','), 10);
113+
cursor = cursor.substring(cursor.indexOf(',') + 1);
114+
int currentPageIndex = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10);
75115

76-
expression = expression.substring(expression.indexOf(',') + 1);
77-
int pageSize = Integer.parseInt(expression, 0, expression.indexOf(','), 10);
116+
cursor = cursor.substring(cursor.indexOf(',') + 1);
117+
int pageSize = Integer.parseInt(cursor, 0, cursor.indexOf(','), 10);
78118

79-
expression = expression.substring(expression.indexOf(',') + 1);
80-
if (!expression.startsWith("(Project,")) {
119+
cursor = cursor.substring(cursor.indexOf(',') + 1);
120+
if (!cursor.startsWith("(Project,")) {
81121
throw new UnsupportedOperationException("Unsupported cursor");
82122
}
83-
expression = expression.substring(expression.indexOf(',') + 1);
84-
if (!expression.startsWith("(namedParseExpressions,")) {
123+
cursor = cursor.substring(cursor.indexOf(',') + 1);
124+
if (!cursor.startsWith("(namedParseExpressions,")) {
85125
throw new UnsupportedOperationException("Unsupported cursor");
86126
}
87-
expression = expression.substring(expression.indexOf(',') + 1);
88-
var serializer = new DefaultExpressionSerializer();
89-
// TODO parse npe
90-
List<NamedExpression> namedParseExpressions = List.of();
91127

92-
expression = expression.substring(expression.indexOf(',') + 1);
128+
cursor = cursor.substring(cursor.indexOf(',') + 1);
129+
List<NamedExpression> namedParseExpressions = new ArrayList<>();
130+
cursor = parseNamedExpressions(namedParseExpressions, cursor);
131+
132+
cursor = cursor.substring(cursor.indexOf(',') + 1);
93133
List<NamedExpression> projectList = new ArrayList<>();
94-
if (!expression.startsWith("(projectList,")) {
134+
if (!cursor.startsWith("(projectList,")) {
95135
throw new UnsupportedOperationException("Unsupported cursor");
96136
}
97-
expression = expression.substring(expression.indexOf(',') + 1);
98-
while (expression.startsWith("(named,")) {
99-
expression = expression.substring(expression.indexOf(',') + 1);
100-
var name = expression.substring(0, expression.indexOf(','));
101-
expression = expression.substring(expression.indexOf(',') + 1);
102-
var alias = expression.substring(0, expression.indexOf(','));
103-
if (alias.isEmpty()) {
104-
alias = null;
105-
}
106-
expression = expression.substring(expression.indexOf(',') + 1);
107-
projectList.add(new NamedExpression(name,
108-
serializer.deserialize(expression.substring(0, expression.indexOf(')'))), alias));
109-
expression = expression.substring(expression.indexOf(',') + 1);
110-
}
137+
cursor = cursor.substring(cursor.indexOf(',') + 1);
138+
cursor = parseNamedExpressions(projectList, cursor);
111139

112-
if (!expression.startsWith("(OpenSearchPagedIndexScan,")) {
140+
if (!cursor.startsWith("(OpenSearchPagedIndexScan,")) {
113141
throw new UnsupportedOperationException("Unsupported cursor");
114142
}
115-
expression = expression.substring(expression.indexOf(',') + 1);
116-
var indexName = expression.substring(0, expression.indexOf(','));
117-
expression = expression.substring(expression.indexOf(',') + 1);
118-
var scrollId = expression.substring(0, expression.indexOf(')'));
143+
cursor = cursor.substring(cursor.indexOf(',') + 1);
144+
var indexName = cursor.substring(0, cursor.indexOf(','));
145+
cursor = cursor.substring(cursor.indexOf(',') + 1);
146+
var scrollId = cursor.substring(0, cursor.indexOf(')'));
119147
TableScanOperator scan = storageEngine.getTableScan(indexName, scrollId);
120148

121149
return new PaginateOperator(new ProjectOperator(scan, projectList, namedParseExpressions),

‎core/src/main/java/org/opensearch/sql/planner/DefaultImplementor.java‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,6 @@ public PhysicalPlan visitLimit(LogicalLimit node, C context) {
126126
return new LimitOperator(visitChild(node, context), node.getLimit(), node.getOffset());
127127
}
128128

129-
130129
@Override
131130
public PhysicalPlan visitPaginate(LogicalPaginate plan, C context) {
132131
return new PaginateOperator(visitChild(plan, context), plan.getPageSize());
@@ -148,10 +147,8 @@ public PhysicalPlan visitRelation(LogicalRelation node, C context) {
148147
+ "implementing and optimizing logical plan with relation involved");
149148
}
150149

151-
152150
protected PhysicalPlan visitChild(LogicalPlan node, C context) {
153151
// Logical operators visited here must have a single child
154152
return node.getChild().get(0).accept(this, context);
155153
}
156-
157154
}

‎core/src/main/java/org/opensearch/sql/planner/PaginateOperator.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,8 @@ public String toCursor() {
8282
String child = getChild().get(0).toCursor();
8383

8484
var nextPage = getPageIndex() + 1;
85-
return createSection("Paginate", Integer.toString(nextPage),
86-
Integer.toString(getPageSize()), child);
85+
return child == null || child.isEmpty()
86+
? null : createSection("Paginate", Integer.toString(nextPage),
87+
Integer.toString(getPageSize()), child);
8788
}
8889
}

‎core/src/main/java/org/opensearch/sql/planner/logical/LogicalRelation.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import com.google.common.collect.ImmutableList;
1010
import lombok.EqualsAndHashCode;
1111
import lombok.Getter;
12+
import lombok.Setter;
1213
import lombok.ToString;
1314
import org.opensearch.sql.storage.Table;
1415

@@ -25,13 +26,18 @@ public class LogicalRelation extends LogicalPlan {
2526
@Getter
2627
private final Table table;
2728

29+
@Getter
30+
@Setter
31+
private Integer pageSize;
32+
2833
/**
2934
* Constructor of LogicalRelation.
3035
*/
3136
public LogicalRelation(String relationName, Table table) {
3237
super(ImmutableList.of());
3338
this.relationName = relationName;
3439
this.table = table;
40+
this.pageSize = null;
3541
}
3642

3743
@Override

‎core/src/main/java/org/opensearch/sql/planner/optimizer/LogicalPlanOptimizer.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,10 @@
1313
import java.util.List;
1414
import java.util.stream.Collectors;
1515
import org.opensearch.sql.planner.logical.LogicalPlan;
16+
import org.opensearch.sql.planner.optimizer.rule.CreatePagingTableScanBuilder;
1617
import org.opensearch.sql.planner.optimizer.rule.MergeFilterAndFilter;
1718
import org.opensearch.sql.planner.optimizer.rule.PushFilterUnderSort;
19+
import org.opensearch.sql.planner.optimizer.rule.PushPageSize;
1820
import org.opensearch.sql.planner.optimizer.rule.read.CreateTableScanBuilder;
1921
import org.opensearch.sql.planner.optimizer.rule.read.TableScanPushDown;
2022
import org.opensearch.sql.planner.optimizer.rule.write.CreateTableWriteBuilder;
@@ -73,6 +75,7 @@ public static LogicalPlanOptimizer paginationCreate() {
7375
/*
7476
* Phase 2: Transformations that rely on data source push down capability
7577
*/
78+
new PushPageSize(),
7679
new CreatePagingTableScanBuilder(),
7780
TableScanPushDown.PUSH_DOWN_FILTER,
7881
TableScanPushDown.PUSH_DOWN_AGGREGATION,

‎core/src/main/java/org/opensearch/sql/planner/optimizer/pattern/Patterns.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.opensearch.sql.planner.logical.LogicalFilter;
1717
import org.opensearch.sql.planner.logical.LogicalHighlight;
1818
import org.opensearch.sql.planner.logical.LogicalLimit;
19+
import org.opensearch.sql.planner.logical.LogicalPaginate;
1920
import org.opensearch.sql.planner.logical.LogicalPlan;
2021
import org.opensearch.sql.planner.logical.LogicalProject;
2122
import org.opensearch.sql.planner.logical.LogicalRelation;
@@ -112,6 +113,16 @@ public static Property<LogicalPlan, Table> table() {
112113
: Optional.empty());
113114
}
114115

116+
/**
117+
* Logical pagination with page size.
118+
*/
119+
public static Property<LogicalPlan, Integer> pagination() {
120+
return Property.optionalProperty("pagination",
121+
plan -> plan instanceof LogicalPaginate
122+
? Optional.of(((LogicalPaginate) plan).getPageSize())
123+
: Optional.empty());
124+
}
125+
115126
/**
116127
* Logical write with table field.
117128
*/

‎core/src/main/java/org/opensearch/sql/planner/optimizer/CreatePagingTableScanBuilder.java‎ renamed to ‎core/src/main/java/org/opensearch/sql/planner/optimizer/rule/CreatePagingTableScanBuilder.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
package org.opensearch.sql.planner.optimizer;
6+
package org.opensearch.sql.planner.optimizer.rule;
77

88
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.table;
99

@@ -14,6 +14,7 @@
1414
import lombok.experimental.Accessors;
1515
import org.opensearch.sql.planner.logical.LogicalPlan;
1616
import org.opensearch.sql.planner.logical.LogicalRelation;
17+
import org.opensearch.sql.planner.optimizer.Rule;
1718
import org.opensearch.sql.storage.Table;
1819
import org.opensearch.sql.storage.read.TableScanBuilder;
1920

@@ -38,7 +39,8 @@ public CreatePagingTableScanBuilder() {
3839

3940
@Override
4041
public LogicalPlan apply(LogicalRelation plan, Captures captures) {
41-
TableScanBuilder scanBuilder = captures.get(capture).createPagedScanBuilder();
42+
TableScanBuilder scanBuilder = captures.get(capture)
43+
.createPagedScanBuilder(plan.getPageSize());
4244
// TODO: Remove this after Prometheus refactored to new table scan builder too
4345
return (scanBuilder == null) ? plan : scanBuilder;
4446
}

‎core/src/main/java/org/opensearch/sql/planner/optimizer/rule/PushPageSize.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.planner.optimizer.rule;
7+
8+
import static org.opensearch.sql.planner.optimizer.pattern.Patterns.pagination;
9+
10+
import com.facebook.presto.matching.Capture;
11+
import com.facebook.presto.matching.Captures;
12+
import com.facebook.presto.matching.Pattern;
13+
import lombok.Getter;
14+
import lombok.experimental.Accessors;
15+
import org.opensearch.sql.planner.logical.LogicalPaginate;
16+
import org.opensearch.sql.planner.logical.LogicalPlan;
17+
import org.opensearch.sql.planner.logical.LogicalRelation;
18+
import org.opensearch.sql.planner.optimizer.Rule;
19+
20+
import java.util.Objects;
21+
22+
public class PushPageSize
23+
implements Rule<LogicalPaginate> {
24+
/** Capture the table inside matched logical paginate operator. */
25+
private final Capture<Integer> capture;
26+
27+
/** Pattern that matches logical paginate operator. */
28+
@Accessors(fluent = true)
29+
@Getter
30+
private final Pattern<LogicalPaginate> pattern;
31+
32+
/**
33+
* Constructor.
34+
*/
35+
public PushPageSize() {
36+
this.capture = Capture.newCapture();
37+
this.pattern = Pattern.typeOf(LogicalPaginate.class)
38+
.with(pagination().capturedAs(capture));
39+
}
40+
41+
private LogicalRelation findLogicalRelation(LogicalPlan plan) { //TODO TBD multiple relations?
42+
for (var subplan : plan.getChild()) {
43+
if (subplan instanceof LogicalRelation) {
44+
return (LogicalRelation) subplan;
45+
}
46+
var found = findLogicalRelation(subplan);
47+
if (found != null) {
48+
return found;
49+
}
50+
}
51+
return null;
52+
}
53+
54+
@Override
55+
public LogicalPlan apply(LogicalPaginate plan, Captures captures) {
56+
var relation = findLogicalRelation(plan);
57+
if (relation != null) {
58+
relation.setPageSize(captures.get(capture));
59+
}
60+
return plan;
61+
}
62+
}

‎core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlan.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public ExecutionEngine.Schema schema() {
5252
}
5353

5454
public String toCursor() {
55-
throw new IllegalStateException(String.format("%s needs to implement ToCursor",
55+
throw new IllegalStateException(String.format("%s needs to implement toCursor",
5656
this.getClass()));
5757
}
5858

0 commit comments

Comments
 (0)