Skip to content

Commit a3cd42e

Browse files
LantaoJinLantaoJin
authored
Skipping codegen and compile for Scan only plan (opensearch-project#3853)
* Skipping codegen and compile for Scan only plan Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix IT Signed-off-by: Lantao Jin <ltjin@amazon.com> * fix IT Signed-off-by: Lantao Jin <ltjin@amazon.com> * Fix explain IT Signed-off-by: Lantao Jin <ltjin@amazon.com> * Fix explain IT without pushdown Signed-off-by: Lantao Jin <ltjin@amazon.com> * add javadoc Signed-off-by: Lantao Jin <ltjin@amazon.com> --------- Signed-off-by: Lantao Jin <ltjin@amazon.com>
1 parent 877835e commit a3cd42e

5 files changed

Lines changed: 165 additions & 5 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/plan/Scannable.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.calcite.plan;
7+
8+
import org.apache.calcite.linq4j.Enumerable;
9+
import org.checkerframework.checker.nullness.qual.Nullable;
10+
11+
/**
12+
* The customized table scan is implemented in OpenSearch module, to invoke this scan() method in
13+
* core module, we add this interface. Now the only implementation is CalciteEnumerableIndexScan.
14+
* When a RelNode after optimization is a Scannable, we can directly invoke scan() method to get the
15+
* result of the scan instead of codegen and compile via Linq4j expression.
16+
*/
17+
public interface Scannable {
18+
19+
public Enumerable<@Nullable Object> scan();
20+
}

‎core/src/main/java/org/opensearch/sql/calcite/utils/CalciteToolsHelper.java‎

Lines changed: 119 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,26 +27,35 @@
2727

2828
package org.opensearch.sql.calcite.utils;
2929

30+
import static java.util.Objects.requireNonNull;
31+
3032
import com.google.common.collect.ImmutableList;
33+
import java.lang.reflect.Type;
3134
import java.sql.Connection;
3235
import java.sql.PreparedStatement;
3336
import java.sql.SQLException;
3437
import java.time.Instant;
3538
import java.util.Properties;
3639
import java.util.function.Consumer;
40+
import org.apache.calcite.adapter.enumerable.EnumerableConvention;
41+
import org.apache.calcite.adapter.enumerable.EnumerableRel;
3742
import org.apache.calcite.adapter.java.JavaTypeFactory;
3843
import org.apache.calcite.avatica.AvaticaConnection;
3944
import org.apache.calcite.avatica.AvaticaFactory;
45+
import org.apache.calcite.avatica.Meta;
4046
import org.apache.calcite.avatica.UnregisteredDriver;
4147
import org.apache.calcite.config.CalciteConnectionProperty;
48+
import org.apache.calcite.interpreter.BindableConvention;
4249
import org.apache.calcite.interpreter.Bindables;
4350
import org.apache.calcite.jdbc.CalciteFactory;
4451
import org.apache.calcite.jdbc.CalciteJdbc41Factory;
4552
import org.apache.calcite.jdbc.CalcitePrepare;
4653
import org.apache.calcite.jdbc.CalciteSchema;
4754
import org.apache.calcite.jdbc.Driver;
55+
import org.apache.calcite.linq4j.function.Function0;
4856
import org.apache.calcite.plan.Context;
4957
import org.apache.calcite.plan.Contexts;
58+
import org.apache.calcite.plan.Convention;
5059
import org.apache.calcite.plan.RelOptCluster;
5160
import org.apache.calcite.plan.RelOptPlanner;
5261
import org.apache.calcite.plan.RelOptSchema;
@@ -55,30 +64,36 @@
5564
import org.apache.calcite.prepare.CalcitePrepareImpl;
5665
import org.apache.calcite.rel.RelHomogeneousShuttle;
5766
import org.apache.calcite.rel.RelNode;
67+
import org.apache.calcite.rel.RelRoot;
5868
import org.apache.calcite.rel.RelShuttle;
5969
import org.apache.calcite.rel.core.TableScan;
6070
import org.apache.calcite.rel.logical.LogicalTableScan;
71+
import org.apache.calcite.rel.type.RelDataType;
72+
import org.apache.calcite.rel.type.RelDataTypeFactory;
6173
import org.apache.calcite.rel.type.RelDataTypeSystem;
6274
import org.apache.calcite.rex.RexBuilder;
6375
import org.apache.calcite.rex.RexNode;
76+
import org.apache.calcite.runtime.Bindable;
6477
import org.apache.calcite.runtime.Hook;
6578
import org.apache.calcite.schema.SchemaPlus;
6679
import org.apache.calcite.server.CalciteServerStatement;
6780
import org.apache.calcite.sql.SqlAggFunction;
6881
import org.apache.calcite.sql.SqlKind;
6982
import org.apache.calcite.sql.parser.SqlParserPos;
83+
import org.apache.calcite.sql2rel.SqlRexConvertletTable;
7084
import org.apache.calcite.tools.FrameworkConfig;
7185
import org.apache.calcite.tools.Frameworks;
7286
import org.apache.calcite.tools.RelBuilder;
7387
import org.apache.calcite.tools.RelRunner;
7488
import org.apache.calcite.util.Holder;
7589
import org.apache.calcite.util.Util;
7690
import org.opensearch.sql.calcite.CalcitePlanContext;
91+
import org.opensearch.sql.calcite.plan.Scannable;
7792
import org.opensearch.sql.calcite.udf.udaf.NullableSqlAvgAggFunction;
7893

7994
/**
8095
* Calcite Tools Helper. This class is used to create customized: 1. Connection 2. JavaTypeFactory
81-
* 3. RelBuilder 4. RelRunner TODO delete it in future if possible.
96+
* 3. RelBuilder 4. RelRunner 5. CalcitePreparingStmt. TODO delete it in future if possible.
8297
*/
8398
public class CalciteToolsHelper {
8499

@@ -153,6 +168,11 @@ public Connection connect(
153168
this.handler.onConnectionInit(connection);
154169
return connection;
155170
}
171+
172+
@Override
173+
protected Function0<CalcitePrepare> createPrepareFactory() {
174+
return OpenSearchPrepareImpl::new;
175+
}
156176
}
157177

158178
/** do nothing, just extend for a public construct for new */
@@ -214,6 +234,104 @@ public <R> R perform(
214234
final RelOptCluster cluster = createCluster(planner, rexBuilder);
215235
return action.apply(cluster, catalogReader, prepareContext.getRootSchema().plus(), statement);
216236
}
237+
238+
/**
239+
* Customize CalcitePreparingStmt. Override {@link CalcitePrepareImpl#getPreparingStmt} and
240+
* return {@link OpenSearchCalcitePreparingStmt}
241+
*/
242+
@Override
243+
protected CalcitePrepareImpl.CalcitePreparingStmt getPreparingStmt(
244+
CalcitePrepare.Context context,
245+
Type elementType,
246+
CalciteCatalogReader catalogReader,
247+
RelOptPlanner planner) {
248+
final JavaTypeFactory typeFactory = context.getTypeFactory();
249+
final EnumerableRel.Prefer prefer;
250+
if (elementType == Object[].class) {
251+
prefer = EnumerableRel.Prefer.ARRAY;
252+
} else {
253+
prefer = EnumerableRel.Prefer.CUSTOM;
254+
}
255+
final Convention resultConvention =
256+
enableBindable ? BindableConvention.INSTANCE : EnumerableConvention.INSTANCE;
257+
return new OpenSearchCalcitePreparingStmt(
258+
this,
259+
context,
260+
catalogReader,
261+
typeFactory,
262+
context.getRootSchema(),
263+
prefer,
264+
createCluster(planner, new RexBuilder(typeFactory)),
265+
resultConvention,
266+
createConvertletTable());
267+
}
268+
}
269+
270+
/**
271+
* Similar to {@link CalcitePrepareImpl.CalcitePreparingStmt}. Customize the logic to convert an
272+
* EnumerableTableScan to BindableTableScan.
273+
*/
274+
public static class OpenSearchCalcitePreparingStmt
275+
extends CalcitePrepareImpl.CalcitePreparingStmt {
276+
277+
public OpenSearchCalcitePreparingStmt(
278+
CalcitePrepareImpl prepare,
279+
CalcitePrepare.Context context,
280+
CatalogReader catalogReader,
281+
RelDataTypeFactory typeFactory,
282+
CalciteSchema schema,
283+
EnumerableRel.Prefer prefer,
284+
RelOptCluster cluster,
285+
Convention resultConvention,
286+
SqlRexConvertletTable convertletTable) {
287+
super(
288+
prepare,
289+
context,
290+
catalogReader,
291+
typeFactory,
292+
schema,
293+
prefer,
294+
cluster,
295+
resultConvention,
296+
convertletTable);
297+
}
298+
299+
@Override
300+
protected PreparedResult implement(RelRoot root) {
301+
Hook.PLAN_BEFORE_IMPLEMENTATION.run(root);
302+
RelDataType resultType = root.rel.getRowType();
303+
boolean isDml = root.kind.belongsTo(SqlKind.DML);
304+
if (root.rel instanceof Scannable scannable) {
305+
final Bindable bindable = dataContext -> scannable.scan();
306+
307+
return new PreparedResultImpl(
308+
resultType,
309+
requireNonNull(parameterRowType, "parameterRowType"),
310+
requireNonNull(fieldOrigins, "fieldOrigins"),
311+
root.collation.getFieldCollations().isEmpty()
312+
? ImmutableList.of()
313+
: ImmutableList.of(root.collation),
314+
root.rel,
315+
mapTableModOp(isDml, root.kind),
316+
isDml) {
317+
@Override
318+
public String getCode() {
319+
throw new UnsupportedOperationException();
320+
}
321+
322+
@Override
323+
public Bindable getBindable(Meta.CursorFactory cursorFactory) {
324+
return bindable;
325+
}
326+
327+
@Override
328+
public Type getElementType() {
329+
return resultType.getFieldList().size() == 1 ? Object.class : Object[].class;
330+
}
331+
};
332+
}
333+
return super.implement(root);
334+
}
217335
}
218336

219337
public static class OpenSearchRelRunners {

‎integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLExplainIT.java‎

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,15 +44,34 @@ public void testExplainCommand() throws IOException {
4444
}
4545

4646
@Test
47-
public void testExplainCommandExtended() throws IOException {
47+
public void testExplainCommandExtendedWithCodegen() throws IOException {
4848
var result =
49-
executeWithReplace("explain extended source=test | where age = 20 | fields name, age");
49+
executeWithReplace(
50+
"explain extended source=test | where age = 20 | join left=l right=r on l.age=r.age"
51+
+ " test");
5052
assertTrue(
5153
result.contains(
5254
"public org.apache.calcite.linq4j.Enumerable bind(final"
5355
+ " org.apache.calcite.DataContext root)"));
5456
}
5557

58+
@Test
59+
public void testExplainCommandExtendedWithoutCodegen() throws IOException {
60+
var result =
61+
executeWithReplace("explain extended source=test | where age = 20 | fields name, age");
62+
if (isPushdownEnabled()) {
63+
assertFalse(
64+
result.contains(
65+
"public org.apache.calcite.linq4j.Enumerable bind(final"
66+
+ " org.apache.calcite.DataContext root)"));
67+
} else {
68+
assertTrue(
69+
result.contains(
70+
"public org.apache.calcite.linq4j.Enumerable bind(final"
71+
+ " org.apache.calcite.DataContext root)"));
72+
}
73+
}
74+
5675
@Test
5776
public void testExplainCommandCost() throws IOException {
5877
var result = executeWithReplace("explain cost source=test | where age = 20 | fields name, age");

‎opensearch/src/main/java/org/opensearch/sql/opensearch/planner/physical/EnumerableIndexScanRule.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public class EnumerableIndexScanRule extends ConverterRule {
2727
"EnumerableIndexScanRule")
2828
.withRuleFactory(EnumerableIndexScanRule::new);
2929

30-
/** Creates an EnumerableProjectRule. */
30+
/** Creates an EnumerableIndexScanRule. */
3131
protected EnumerableIndexScanRule(Config config) {
3232
super(config);
3333
}

‎opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,13 @@
2828
import org.apache.logging.log4j.Logger;
2929
import org.checkerframework.checker.nullness.qual.Nullable;
3030
import org.opensearch.sql.calcite.plan.OpenSearchRules;
31+
import org.opensearch.sql.calcite.plan.Scannable;
3132
import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
3233
import org.opensearch.sql.opensearch.storage.OpenSearchIndex;
3334

3435
/** The physical relational operator representing a scan of an OpenSearchIndex type. */
35-
public class CalciteEnumerableIndexScan extends AbstractCalciteIndexScan implements EnumerableRel {
36+
public class CalciteEnumerableIndexScan extends AbstractCalciteIndexScan
37+
implements Scannable, EnumerableRel {
3638
private static final Logger LOG = LogManager.getLogger(CalciteEnumerableIndexScan.class);
3739

3840
/**
@@ -85,6 +87,7 @@ public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
8587
* each time to avoid reusing source builder. That's because the source builder has stats like PIT
8688
* or SearchAfter recorded during previous search.
8789
*/
90+
@Override
8891
public Enumerable<@Nullable Object> scan() {
8992
return new AbstractEnumerable<>() {
9093
@Override

0 commit comments

Comments
 (0)