Skip to content

Commit

Permalink
[Starrocks On ES] support explain for ES table (StarRocks#8202)
Browse files Browse the repository at this point in the history
  • Loading branch information
@RowenWoo @melt-code
RowenWoo authored and melt-code committed Jul 17, 2022
1 parent 7f71aea commit 13ceb4e
Show file tree
Hide file tree
Showing 5 changed files with 502 additions and 33 deletions.
9 changes: 9 additions & 0 deletions fe/fe-core/src/main/java/com/starrocks/common/io/FastByteArrayOutputStream.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;

/**
* ByteArrayOutputStream implementation that doesn't synchronize methods
Expand Down Expand Up @@ -95,4 +96,12 @@ public InputStream getInputStream() {
return new FastByteArrayInputStream(buf, size);
}

@Override
public String toString() {
if (buf == null || size == 0) {
return "";
} else {
return new String(buf, 0, size, StandardCharsets.UTF_8);
}
}
}
90 changes: 68 additions & 22 deletions fe/fe-core/src/main/java/com/starrocks/external/elasticsearch/QueryBuilders.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,19 @@

package com.starrocks.external.elasticsearch;

import com.fasterxml.jackson.core.JsonEncoding;
import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.starrocks.common.io.FastByteArrayOutputStream;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
Expand All @@ -30,6 +38,32 @@
*/
public final class QueryBuilders {

/**
* Base class to build various ES queries
*/
public abstract static class QueryBuilder {

/**
* Convert query to JSON format
*
* @param out used to generate JSON elements
* @throws IOException if IO error occurred
*/
public abstract void toJson(JsonGenerator out) throws IOException;

@Override
public String toString() throws StarRocksESException {
FastByteArrayOutputStream out = new FastByteArrayOutputStream(256);
JsonFactory factory = new JsonFactory();
try (JsonGenerator jsonGenerator = factory.createGenerator(out, JsonEncoding.UTF8)) {
toJson(jsonGenerator);
} catch (IOException e) {
throw new StarRocksESException(e.getMessage());
}
return out.toString();
}
}

/**
* A query that matches on all documents.
*/
Expand Down Expand Up @@ -157,20 +191,6 @@ public static RangeQueryBuilder rangeQuery(String name) {
return new RangeQueryBuilder(name);
}

/**
* Base class to build various ES queries
*/
abstract static class QueryBuilder {

/**
* Convert query to JSON format
*
* @param out used to generate JSON elements
* @throws IOException if IO error occurred
*/
abstract void toJson(JsonGenerator out) throws IOException;
}

/**
* A Query that matches documents matching boolean combinations of other queries.
*/
Expand Down Expand Up @@ -206,7 +226,7 @@ BoolQueryBuilder should(QueryBuilder queryBuilder) {
}

@Override
protected void toJson(JsonGenerator out) throws IOException {
public void toJson(JsonGenerator out) throws IOException {
out.writeStartObject();
out.writeFieldName("bool");
out.writeStartObject();
Expand Down Expand Up @@ -250,7 +270,7 @@ private TermQueryBuilder(final String fieldName, final Object value) {
}

@Override
void toJson(final JsonGenerator out) throws IOException {
public void toJson(final JsonGenerator out) throws IOException {
out.writeStartObject();
out.writeFieldName("term");
out.writeStartObject();
Expand All @@ -274,7 +294,7 @@ private TermsQueryBuilder(final String fieldName, final Iterable<?> values) {
}

@Override
void toJson(final JsonGenerator out) throws IOException {
public void toJson(final JsonGenerator out) throws IOException {
out.writeStartObject();
out.writeFieldName("terms");
out.writeStartObject();
Expand Down Expand Up @@ -341,7 +361,7 @@ RangeQueryBuilder format(String format) {
}

@Override
void toJson(final JsonGenerator out) throws IOException {
public void toJson(final JsonGenerator out) throws IOException {
if (lt == null && gt == null) {
throw new IllegalStateException("Either lower or upper bound should be provided");
}
Expand Down Expand Up @@ -390,7 +410,7 @@ public WildcardQueryBuilder(String fieldName, String value) {
}

@Override
void toJson(JsonGenerator out) throws IOException {
public void toJson(JsonGenerator out) throws IOException {
out.writeStartObject();
out.writeFieldName("wildcard");
out.writeStartObject();
Expand All @@ -413,7 +433,7 @@ static class ExistsQueryBuilder extends QueryBuilder {
}

@Override
void toJson(JsonGenerator out) throws IOException {
public void toJson(JsonGenerator out) throws IOException {
out.writeStartObject();
out.writeFieldName("exists");
out.writeStartObject();
Expand All @@ -427,13 +447,13 @@ void toJson(JsonGenerator out) throws IOException {
/**
* A query that matches on all documents
*/
static class MatchAllQueryBuilder extends QueryBuilder {
public static class MatchAllQueryBuilder extends QueryBuilder {

private MatchAllQueryBuilder() {
}

@Override
void toJson(final JsonGenerator out) throws IOException {
public void toJson(final JsonGenerator out) throws IOException {
out.writeStartObject();
out.writeFieldName("match_all");
out.writeStartObject();
Expand All @@ -452,4 +472,30 @@ void toJson(final JsonGenerator out) throws IOException {
private static void writeObject(JsonGenerator out, Object value) throws IOException {
out.writeObject(value);
}

public static class RawQueryBuilder extends QueryBuilder {
private static final ObjectMapper MAPPER =
new ObjectMapper()
.configure(JsonParser.Feature.ALLOW_COMMENTS, true);
private final String value;

public RawQueryBuilder(String value) {
this.value = value;
Map<String, Object> map;
try {
map = MAPPER.readValue(value, HashMap.class);
} catch (JsonProcessingException e) {
throw new StarRocksESException(e.getMessage());
}
if (map == null || map.size() != 1) {
throw new StarRocksESException("raw query dsl in `esquery` must have one (only one) root json element");
}

}

@Override
public void toJson(JsonGenerator out) throws IOException {
out.writeRaw(value);
}
}
}

0 comments on commit 13ceb4e

Please sign in to comment.