private static Map<String, DataSource> getDataSourceMap() {
final Map<String, DataSource> dataSourceMap = new HashMap<String, DataSource>();
dataSourceMap.put("column11", _indexSegment.getDataSource("column11"));
dataSourceMap.put("column12", _indexSegment.getDataSource("column12"));
dataSourceMap.put("met_impressionCount", _indexSegment.getDataSource("met_impressionCount"));
return dataSourceMap;
}
private MAggregationGroupByOperator getmAggregationGroupByOperator() {
final List<AggregationFunctionGroupByOperator> aggregationFunctionGroupByOperatorList =
new ArrayList<AggregationFunctionGroupByOperator>();
Operator filterOperator =
new MatchEntireSegmentOperator(_indexSegment.getSegmentMetadata().getTotalRawDocs());
final BReusableFilteredDocIdSetOperator docIdSetOperator =
new BReusableFilteredDocIdSetOperator(
filterOperator, _indexSegment.getSegmentMetadata().getTotalRawDocs(), 5000);
final Map<String, DataSource> dataSourceMap = getDataSourceMap();
final MProjectionOperator projectionOperator =
new MProjectionOperator(dataSourceMap, docIdSetOperator);
for (int i = 0; i < _numAggregations; ++i) {
final MAggregationFunctionGroupByWithDictionaryAndTrieTreeOperator
aggregationFunctionGroupByOperator =
new MAggregationFunctionGroupByWithDictionaryAndTrieTreeOperator(
_aggregationInfos.get(i),
_groupBy,
new UReplicatedProjectionOperator(projectionOperator),
true);
aggregationFunctionGroupByOperatorList.add(aggregationFunctionGroupByOperator);
}
return new MAggregationGroupByOperator(
_indexSegment,
_aggregationInfos,
_groupBy,
projectionOperator,
aggregationFunctionGroupByOperatorList);
}
@Test
public void testMaxQuery() {
BrokerRequest brokerRequest = getMaxQuery();
QuerySource querySource = new QuerySource();
querySource.setTableName("midas");
brokerRequest.setQuerySource(querySource);
InstanceRequest instanceRequest = new InstanceRequest(0, brokerRequest);
instanceRequest.setSearchSegments(new ArrayList<String>());
for (IndexSegment segment : _indexSegmentList) {
instanceRequest.getSearchSegments().add(segment.getSegmentName());
}
try {
DataTable instanceResponse = _queryExecutor.processQuery(instanceRequest);
LOGGER.info("InstanceResponse is " + instanceResponse.getDouble(0, 0));
Assert.assertEquals(instanceResponse.getDouble(0, 0), 200000.0);
LOGGER.info(
"Time used for instanceResponse is " + instanceResponse.getMetadata().get("timeUsedMs"));
} catch (Exception e) {
e.printStackTrace();
// Should never happen
Assert.assertEquals(true, false);
}
}
@AfterClass
public void tearDown() {
if (INDEXES_DIR.exists()) {
FileUtils.deleteQuietly(INDEXES_DIR);
}
for (IndexSegment segment : _indexSegmentList) {
segment.destroy();
}
_indexSegmentList.clear();
}
public static DataSchema extractDataSchema(
List<SelectionSort> sortSequence, List<String> selectionColumns, IndexSegment indexSegment) {
final List<String> columns = new ArrayList<String>();
if (sortSequence != null && !sortSequence.isEmpty()) {
for (final SelectionSort selectionSort : sortSequence) {
columns.add(selectionSort.getColumn());
}
}
String[] selectionColumnArray = selectionColumns.toArray(new String[selectionColumns.size()]);
Arrays.sort(selectionColumnArray);
for (int i = 0; i < selectionColumnArray.length; ++i) {
String selectionColumn = selectionColumnArray[i];
if (!columns.contains(selectionColumn)) {
columns.add(selectionColumn);
}
}
final DataType[] dataTypes = new DataType[columns.size()];
for (int i = 0; i < dataTypes.length; ++i) {
DataSource ds = indexSegment.getDataSource(columns.get(i));
DataSourceMetadata m = ds.getDataSourceMetadata();
dataTypes[i] = m.getDataType();
if (!m.isSingleValue()) {
dataTypes[i] = DataType.valueOf(dataTypes[i] + "_ARRAY");
}
}
return new DataSchema(columns.toArray(new String[0]), dataTypes);
}
@Test
public void testAggregationGroupBys() {
final List<AggregationFunctionGroupByOperator> aggregationFunctionGroupByOperatorList =
new ArrayList<AggregationFunctionGroupByOperator>();
final BReusableFilteredDocIdSetOperator docIdSetOperator =
new BReusableFilteredDocIdSetOperator(null, _indexSegment.getTotalDocs(), 5000);
final Map<String, DataSource> dataSourceMap = getDataSourceMap();
final MProjectionOperator projectionOperator =
new MProjectionOperator(dataSourceMap, docIdSetOperator);
for (int i = 0; i < _numAggregations; ++i) {
final MDefaultAggregationFunctionGroupByOperator aggregationFunctionGroupByOperator =
new MDefaultAggregationFunctionGroupByOperator(
_aggregationInfos.get(i),
_groupBy,
new UReplicatedProjectionOperator(projectionOperator),
true);
aggregationFunctionGroupByOperatorList.add(aggregationFunctionGroupByOperator);
}
final MAggregationGroupByOperator aggregationGroupByOperator =
new MAggregationGroupByOperator(
_indexSegment,
_aggregationInfos,
_groupBy,
projectionOperator,
aggregationFunctionGroupByOperatorList);
final IntermediateResultsBlock block =
(IntermediateResultsBlock) aggregationGroupByOperator.nextBlock();
for (int i = 0; i < _numAggregations; ++i) {
LOGGER.info("Result: {}", block.getAggregationGroupByOperatorResult().get(i));
}
}
public static List<String> getSelectionColumns(
List<String> selectionColumns, IndexSegment indexSegment) {
if ((selectionColumns.size() == 1) && selectionColumns.get(0).equals("*")) {
selectionColumns.clear();
for (final String columnName : indexSegment.getColumnNames()) {
selectionColumns.add(columnName);
}
}
return selectionColumns;
}
public static String[] extractSelectionRelatedColumns(
Selection selection, IndexSegment indexSegment) {
Set<String> selectionColumns = new HashSet<String>();
selectionColumns.addAll(selection.getSelectionColumns());
if ((selectionColumns.size() == 1)
&& ((selectionColumns.toArray(new String[0]))[0].equals("*"))) {
selectionColumns.clear();
selectionColumns.addAll(Arrays.asList(indexSegment.getColumnNames()));
}
if (selection.getSelectionSortSequence() != null) {
for (SelectionSort selectionSort : selection.getSelectionSortSequence()) {
selectionColumns.add(selectionSort.getColumn());
}
}
return selectionColumns.toArray(new String[0]);
}
@Test
public void testAggregationGroupBysWithDataTableEncodeAndDecode() throws Exception {
final List<AggregationFunctionGroupByOperator> aggregationFunctionGroupByOperatorList =
new ArrayList<AggregationFunctionGroupByOperator>();
final BReusableFilteredDocIdSetOperator docIdSetOperator =
new BReusableFilteredDocIdSetOperator(null, _indexSegment.getTotalDocs(), 5000);
final Map<String, DataSource> dataSourceMap = getDataSourceMap();
final MProjectionOperator projectionOperator =
new MProjectionOperator(dataSourceMap, docIdSetOperator);
for (int i = 0; i < _numAggregations; ++i) {
final MDefaultAggregationFunctionGroupByOperator aggregationFunctionGroupByOperator =
new MDefaultAggregationFunctionGroupByOperator(
_aggregationInfos.get(i),
_groupBy,
new UReplicatedProjectionOperator(projectionOperator),
true);
aggregationFunctionGroupByOperatorList.add(aggregationFunctionGroupByOperator);
}
final MAggregationGroupByOperator aggregationGroupByOperator =
new MAggregationGroupByOperator(
_indexSegment,
_aggregationInfos,
_groupBy,
projectionOperator,
aggregationFunctionGroupByOperatorList);
final IntermediateResultsBlock block =
(IntermediateResultsBlock) aggregationGroupByOperator.nextBlock();
for (int i = 0; i < _numAggregations; ++i) {
LOGGER.info("Result 1: {}", block.getAggregationGroupByOperatorResult().get(i));
}
/////////////////////////////////////////////////////////////////////////
final List<AggregationFunctionGroupByOperator> aggregationFunctionGroupByOperatorList1 =
new ArrayList<AggregationFunctionGroupByOperator>();
final BReusableFilteredDocIdSetOperator docIdSetOperator1 =
new BReusableFilteredDocIdSetOperator(null, _indexSegment.getTotalDocs(), 5000);
final Map<String, DataSource> dataSourceMap1 = getDataSourceMap();
final MProjectionOperator projectionOperator1 =
new MProjectionOperator(dataSourceMap1, docIdSetOperator1);
for (int i = 0; i < _numAggregations; ++i) {
final MDefaultAggregationFunctionGroupByOperator aggregationFunctionGroupByOperator1 =
new MDefaultAggregationFunctionGroupByOperator(
_aggregationInfos.get(i),
_groupBy,
new UReplicatedProjectionOperator(projectionOperator1),
true);
aggregationFunctionGroupByOperatorList1.add(aggregationFunctionGroupByOperator1);
}
final MAggregationGroupByOperator aggregationGroupByOperator1 =
new MAggregationGroupByOperator(
_indexSegment,
_aggregationInfos,
_groupBy,
projectionOperator1,
aggregationFunctionGroupByOperatorList1);
final IntermediateResultsBlock block1 =
(IntermediateResultsBlock) aggregationGroupByOperator1.nextBlock();
for (int i = 0; i < _numAggregations; ++i) {
LOGGER.info("Result 2: {}", block1.getAggregationGroupByOperatorResult().get(i));
}
CombineService.mergeTwoBlocks(getAggregationGroupByNoFilterBrokerRequest(), block, block1);
for (int i = 0; i < _numAggregations; ++i) {
LOGGER.info("Combined Result: {}", block.getAggregationGroupByOperatorResult().get(i));
}
final DataTable dataTable = block.getAggregationGroupByResultDataTable();
final List<Map<String, Serializable>> results =
AggregationGroupByOperatorService.transformDataTableToGroupByResult(dataTable);
for (int i = 0; i < _numAggregations; ++i) {
LOGGER.info("Decode AggregationResult from DataTable: {}", results.get(i));
}
}