[CALCITE-2347] running ElasticSearch in embedded mode for unit tests of ES adapter...
[calcite.git] / core / src / main / java / org / apache / calcite / adapter / elasticsearch / ElasticsearchProject.java
1 /*
2 * Licensed to the Apache Software Foundation (ASF) under one or more
3 * contributor license agreements. See the NOTICE file distributed with
4 * this work for additional information regarding copyright ownership.
5 * The ASF licenses this file to you under the Apache License, Version 2.0
6 * (the "License"); you may not use this file except in compliance with
7 * the License. You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17 package org.apache.calcite.adapter.elasticsearch;
18
19 import org.apache.calcite.adapter.java.JavaTypeFactory;
20 import org.apache.calcite.plan.RelOptCluster;
21 import org.apache.calcite.plan.RelOptCost;
22 import org.apache.calcite.plan.RelOptPlanner;
23 import org.apache.calcite.plan.RelTraitSet;
24 import org.apache.calcite.rel.RelNode;
25 import org.apache.calcite.rel.core.Project;
26 import org.apache.calcite.rel.metadata.RelMetadataQuery;
27 import org.apache.calcite.rel.type.RelDataType;
28 import org.apache.calcite.rex.RexNode;
29 import org.apache.calcite.util.Pair;
30
31 import com.google.common.base.Function;
32 import com.google.common.collect.Lists;
33
34 import java.util.ArrayList;
35 import java.util.List;
36
37 import javax.annotation.Nullable;
38
39 /**
40 * Implementation of {@link org.apache.calcite.rel.core.Project}
41 * relational expression in Elasticsearch.
42 */
43 public class ElasticsearchProject extends Project implements ElasticsearchRel {
44 public ElasticsearchProject(RelOptCluster cluster, RelTraitSet traitSet, RelNode input,
45 List<? extends RexNode> projects, RelDataType rowType) {
46 super(cluster, traitSet, input, projects, rowType);
47 assert getConvention() == ElasticsearchRel.CONVENTION;
48 assert getConvention() == input.getConvention();
49 }
50
51 @Override public Project copy(RelTraitSet relTraitSet, RelNode input, List<RexNode> projects,
52 RelDataType relDataType) {
53 return new ElasticsearchProject(getCluster(), traitSet, input, projects, relDataType);
54 }
55
56 @Override public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {
57 return super.computeSelfCost(planner, mq).multiplyBy(0.1);
58 }
59
60 @Override public void implement(Implementor implementor) {
61 implementor.visitChild(0, getInput());
62
63 final List<String> inFields =
64 ElasticsearchRules.elasticsearchFieldNames(getInput().getRowType());
65 final ElasticsearchRules.RexToElasticsearchTranslator translator =
66 new ElasticsearchRules.RexToElasticsearchTranslator(
67 (JavaTypeFactory) getCluster().getTypeFactory(), inFields);
68
69 final List<String> fields = new ArrayList<>();
70 final List<String> scriptFields = new ArrayList<>();
71 for (Pair<RexNode, String> pair: getNamedProjects()) {
72 final String name = pair.right;
73 final String expr = pair.left.accept(translator);
74
75 if (expr.equals("\"" + name + "\"")) {
76 fields.add(name);
77 } else if (expr.matches("\"literal\":.+")) {
78 scriptFields.add(ElasticsearchRules.quote(name)
79 + ":{\"script\": "
80 + expr.split(":")[1] + "}");
81 } else {
82 scriptFields.add(ElasticsearchRules.quote(name)
83 + ":{\"script\":"
84 // _source (ES2) vs params._source (ES5)
85 + "\"" + implementor.elasticsearchTable.scriptedFieldPrefix() + "."
86 + expr.replaceAll("\"", "") + "\"}");
87 }
88 }
89
90 StringBuilder query = new StringBuilder();
91 if (scriptFields.isEmpty()) {
92 List<String> newList = Lists.transform(fields, new Function<String, String>() {
93 @Nullable
94 @Override public String apply(@Nullable String input) {
95 return ElasticsearchRules.quote(input);
96 }
97 });
98
99 final String findString = String.join(", ", newList);
100 query.append("\"_source\" : [").append(findString).append("]");
101 } else {
102 // if scripted fields are present, ES ignores _source attribute
103 for (String field: fields) {
104 scriptFields.add(ElasticsearchRules.quote(field) + ":{\"script\": "
105 // _source (ES2) vs params._source (ES5)
106 + "\"" + implementor.elasticsearchTable.scriptedFieldPrefix() + "."
107 + field + "\"}");
108 }
109 query.append("\"script_fields\": {" + String.join(", ", scriptFields) + "}");
110 }
111
112 for (String opfield : implementor.list) {
113 if (opfield.startsWith("\"_source\"")) {
114 implementor.list.remove(opfield);
115 }
116 }
117 implementor.add(query.toString());
118 }
119 }
120
121 // End ElasticsearchProject.java