You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ya...@apache.org on 2026/04/21 19:19:44 UTC
(pinot) branch master updated: Add LeafOperator tests for aggregation, group-by, and distinct result blocks (#18271)
This is an automated email from the ASF dual-hosted git repository.
yashmayya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 0607efe710d Add LeafOperator tests for aggregation, group-by, and distinct result blocks (#18271)
0607efe710d is described below
commit 0607efe710d40a29852417e55f16c49daf657957
Author: Yash Mayya <ya...@gmail.com>
AuthorDate: Tue Apr 21 12:19:37 2026 -0700
Add LeafOperator tests for aggregation, group-by, and distinct result blocks (#18271)
---
.../query/runtime/operator/LeafOperatorTest.java | 105 ++++++++++++++++++++-
1 file changed, 104 insertions(+), 1 deletion(-)
diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java
index 4f906784c15..e69ffac27cc 100644
--- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java
+++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/LeafOperatorTest.java
@@ -30,10 +30,17 @@ import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.core.data.table.Record;
+import org.apache.pinot.core.data.table.Table;
import org.apache.pinot.core.operator.blocks.InstanceResponseBlock;
+import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock;
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.DistinctResultsBlock;
+import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
import org.apache.pinot.core.operator.blocks.results.MetadataResultsBlock;
import org.apache.pinot.core.operator.blocks.results.SelectionResultsBlock;
+import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
+import org.apache.pinot.core.query.distinct.table.DistinctTable;
import org.apache.pinot.core.query.executor.QueryExecutor;
import org.apache.pinot.core.query.request.ServerQueryRequest;
import org.apache.pinot.core.query.request.context.QueryContext;
@@ -60,7 +67,6 @@ import static org.testng.Assert.assertSame;
import static org.testng.Assert.assertTrue;
-// TODO: add tests for Agg / GroupBy / Distinct result blocks
public class LeafOperatorTest {
private final ExecutorService _executorService = Executors.newCachedThreadPool();
private final AtomicReference<LeafOperator> _operatorRef = new AtomicReference<>();
@@ -401,4 +407,101 @@ public class LeafOperatorTest {
operator.close();
}
+
+ @Test
+ public void shouldReturnAggregationResultBlock() {
+ // Given:
+ QueryContext queryContext =
+ QueryContextConverterUtils.getQueryContext("SELECT COUNT(*), SUM(intCol) FROM tbl");
+ AggregationFunction[] aggFunctions = queryContext.getAggregationFunctions();
+ AggregationResultsBlock aggBlock =
+ new AggregationResultsBlock(aggFunctions, Arrays.asList(2L, 10.0), queryContext);
+ DataSchema schema = aggBlock.getDataSchema();
+ List<BaseResultsBlock> dataBlocks = Collections.singletonList(aggBlock);
+ InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock());
+ QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock);
+ LeafOperator operator =
+ new LeafOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1), schema, queryExecutor,
+ _executorService);
+ _operatorRef.set(operator);
+
+ // When:
+ MseBlock resultBlock = operator.nextBlock();
+
+ // Then:
+ List<Object[]> rows = ((MseBlock.Data) resultBlock).asRowHeap().getRows();
+ assertEquals(rows.size(), 1);
+ assertEquals(rows.get(0), new Object[]{2L, 10.0});
+ assertTrue(operator.nextBlock().isEos(), "Expected EOS after reading the aggregation block");
+
+ operator.close();
+ }
+
+ @Test
+ public void shouldReturnGroupByResultBlock() {
+ // Given:
+ QueryContext queryContext =
+ QueryContextConverterUtils.getQueryContext("SELECT strCol, COUNT(*) FROM tbl GROUP BY strCol");
+ DataSchema schema = new DataSchema(new String[]{"strCol", "count(*)"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.STRING, DataSchema.ColumnDataType.LONG});
+ List<Record> records = Arrays.asList(new Record(new Object[]{"a", 3L}), new Record(new Object[]{"b", 5L}));
+ Table table = mock(Table.class);
+ when(table.getDataSchema()).thenReturn(schema);
+ when(table.size()).thenReturn(records.size());
+ when(table.iterator()).thenAnswer(inv -> records.iterator());
+ GroupByResultsBlock groupByBlock = new GroupByResultsBlock(table, queryContext);
+ List<BaseResultsBlock> dataBlocks = Collections.singletonList(groupByBlock);
+ InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock());
+ QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock);
+ LeafOperator operator =
+ new LeafOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1), schema, queryExecutor,
+ _executorService);
+ _operatorRef.set(operator);
+
+ // When:
+ MseBlock resultBlock = operator.nextBlock();
+
+ // Then:
+ List<Object[]> rows = ((MseBlock.Data) resultBlock).asRowHeap().getRows();
+ assertEquals(rows.size(), 2);
+ assertEquals(rows.get(0), new Object[]{"a", 3L});
+ assertEquals(rows.get(1), new Object[]{"b", 5L});
+ assertTrue(operator.nextBlock().isEos(), "Expected EOS after reading the group-by block");
+
+ operator.close();
+ }
+
+ @Test
+ public void shouldReturnDistinctResultBlock() {
+ // Given:
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext("SELECT DISTINCT intCol FROM tbl");
+ DataSchema schema = new DataSchema(new String[]{"intCol"},
+ new DataSchema.ColumnDataType[]{DataSchema.ColumnDataType.INT});
+ List<Object[]> distinctRows = Arrays.asList(new Object[]{1}, new Object[]{2}, new Object[]{3});
+ DistinctTable distinctTable = mock(DistinctTable.class);
+ when(distinctTable.getDataSchema()).thenReturn(schema);
+ when(distinctTable.size()).thenReturn(distinctRows.size());
+ when(distinctTable.getRows()).thenReturn(distinctRows);
+ DistinctResultsBlock distinctBlock = new DistinctResultsBlock(distinctTable, queryContext);
+ List<BaseResultsBlock> dataBlocks = Collections.singletonList(distinctBlock);
+ InstanceResponseBlock metadataBlock = new InstanceResponseBlock(new MetadataResultsBlock());
+ QueryExecutor queryExecutor = mockQueryExecutor(dataBlocks, metadataBlock);
+ LeafOperator operator =
+ new LeafOperator(OperatorTestUtil.getTracingContext(), mockQueryRequests(1), schema, queryExecutor,
+ _executorService);
+ _operatorRef.set(operator);
+
+ // When:
+ MseBlock resultBlock = operator.nextBlock();
+
+ // Then:
+ List<Object[]> rows = ((MseBlock.Data) resultBlock).asRowHeap().getRows();
+ assertEquals(rows.size(), 3);
+ assertEquals(rows.get(0), new Object[]{1});
+ assertEquals(rows.get(1), new Object[]{2});
+ assertEquals(rows.get(2), new Object[]{3});
+ assertTrue(operator.nextBlock().isEos(), "Expected EOS after reading the distinct block");
+
+ operator.close();
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org