/**
* Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com)
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.linkedin.pinot.tools.scan.query;
import com.linkedin.pinot.core.query.utils.Pair;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.codehaus.jackson.annotate.JsonProperty;
import org.codehaus.jackson.map.ObjectMapper;
import org.codehaus.jackson.map.annotate.JsonSerialize;
public class QueryResponse {
private int _numDocsScanned = 0;
private int _totalDocs = 0;
private long _timeUsedMs = 0;
SelectionResults _selectionResults;
List<AggregationResult> _aggregationResults;
QueryResponse(ResultTable resultTable) {
if (resultTable == null) {
return;
}
_numDocsScanned = resultTable.getNumDocsScanned();
_totalDocs = resultTable.getTotalDocs();
_timeUsedMs = resultTable.getProcessingTime();
switch (resultTable.getResultType()) {
case Selection:
buildSelectionResult(resultTable);
break;
case Aggregation:
buildAggregationResult(resultTable);
break;
case AggregationGroupBy:
buildAggregationGroupByResult(resultTable);
break;
default:
// Nothing to do here.
}
}
private void buildSelectionResult(ResultTable resultTable) {
List<List<String>> results = new ArrayList<>();
List<String> columns = new ArrayList<>();
for (Pair<String, String> pair : resultTable.getColumnList()) {
columns.add(pair.getFirst());
}
for (ResultTable.Row row : resultTable) {
List<String> columnValues = new ArrayList<>();
for (Object value : row) {
columnValues.add(value.toString());
}
results.add(columnValues);
}
_selectionResults = new SelectionResults(columns, results);
}
private void buildAggregationResult(ResultTable resultTable) {
_aggregationResults = new ArrayList<>();
for (ResultTable.Row row : resultTable) {
int columnId = 0;
for (Object value : row) {
AggregationResult aggregationResult =
new AggregationResult(resultTable.getFunction(columnId), value.toString());
_aggregationResults.add(aggregationResult);
++columnId;
}
}
}
private void buildAggregationGroupByResult(ResultTable resultTable) {
List<Pair> columnList = resultTable.getColumnList();
List<String> groupByColumns = new ArrayList<>();
_aggregationResults = new ArrayList<>();
for (Pair pair : columnList) {
if (pair.getSecond() == null) {
groupByColumns.add(pair.getFirst().toString());
}
}
int numGroupByColumns = groupByColumns.size();
List<List<String>> groupValueStrings = new ArrayList<>();
for (ResultTable.Row row : resultTable) {
int colId = 0;
List<String> group = new ArrayList<>();
for (Object value : row) {
// All group by columns come first.
if (columnList.get(colId).getSecond() != null) {
break;
}
if (value instanceof Object []) {
Object[] array = (Object[]) value;
for (Object obj : array) {
group.add(obj.toString());
}
} else {
group.add(value.toString());
}
++colId;
}
groupValueStrings.add(group);
}
for (int colId = numGroupByColumns; colId < columnList.size(); ++colId) {
String function = resultTable.getFunction(colId);
if (function.equalsIgnoreCase("count_*")) {
function = "count_star";
}
int rowId = 0;
List<GroupValue> groupValues = new ArrayList<>();
for (ResultTable.Row row : resultTable) {
String value = row.get(colId).toString();
GroupValue groupValue = new GroupValue(value, groupValueStrings.get(rowId));
groupValues.add(groupValue);
++rowId;
}
AggregationResult aggregationResult = new AggregationResult(groupValues, groupByColumns, function);
_aggregationResults.add(aggregationResult);
}
}
public int getNumDocsScanned() {
return _numDocsScanned;
}
public int getTotalDocs() {
return _totalDocs;
}
public long getTimeUsedMs() {
return _timeUsedMs;
}
@JsonProperty("selectionResults")
@JsonSerialize(include= JsonSerialize.Inclusion.NON_NULL)
SelectionResults getSelectionResults() {
return _selectionResults;
}
@JsonProperty("aggregationResults")
@JsonSerialize(include= JsonSerialize.Inclusion.NON_NULL)
List<AggregationResult> getAggregationResults() {
return _aggregationResults;
}
class SelectionResults {
private final List<String> _columnList;
private final List<List<String>> _results;
SelectionResults(List<String> columnList, List<List<String>> results) {
_columnList = columnList;
_results = results;
}
@JsonProperty("columns")
public List<String> getColumnList() {
return _columnList;
}
@JsonProperty("results")
public List<List<String>> getResults() {
return _results;
}
}
class AggregationResult {
private String _value;
private String _function;
List<GroupValue> _groupValues;
List<String> _groupByColumns;
AggregationResult(String function, String value) {
_function = function;
if (_function.equalsIgnoreCase("count_*")) {
_function = "count_star";
}
_value = value;
_groupValues = null;
}
public AggregationResult(List<GroupValue> group, List<String> groupByColumns, String function) {
_groupValues = group;
_groupByColumns = groupByColumns;
_function = function;
_value = null;
}
@JsonProperty("function")
@JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
public String getFunction() {
return _function;
}
@JsonProperty("value")
@JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
public String getValue() {
return _value;
}
@JsonProperty("groupByResult")
@JsonSerialize(include = JsonSerialize.Inclusion.NON_NULL)
public List<GroupValue> getGroupValues() {
return _groupValues;
}
@JsonProperty("groupByColumns")
public List<String> getGroupByColumns() {
return _groupByColumns;
}
}
class GroupValue {
private final String _value;
private final List<String> _group;
GroupValue(String value, List<String> group) {
_value = value;
_group = group;
}
@JsonProperty("value")
public String getValue() {
return _value;
}
@JsonProperty("group")
public List<String> getGroup() {
return _group;
}
}
@Override
public String toString() {
try {
return new ObjectMapper().writeValueAsString(this);
} catch (IOException e) {
return null;
}
}
}