Fix aggregate queries with labels column in prometheus

Remove the map to Block conversion of MetricHeader to resolve
the issue with aggregate queries with labels column in Prometheus

Co-authored-by: Yuya Ebihara <ebyhry@gmail.com>

Author: 2 people

presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusRecordCursor.java

@@ -270,7 +270,7 @@ private Object getFieldValue(int field)
         int columnIndex = fieldToColumnIndex[field];
         switch (columnIndex) {
             case 0:
-                return fields.getLabels();
+                return getBlockFromMap(columnHandles.get(columnIndex).getColumnType(), fields.getLabels());
             case 1:
                 return fields.getTimestamp();
             case 2:
@@ -284,13 +284,12 @@ private void checkFieldType(int field, Type expected)
         Type actual = getType(field);
         checkArgument(actual.equals(expected), "Expected field %s to be type %s but is %s", field, expected, actual);
     }
-
     private List<PrometheusStandardizedRow> prometheusResultsInStandardizedForm(List<PrometheusMetricResult> results)
     {
         return results.stream().map(result ->
                 result.getTimeSeriesValues().getValues().stream().map(prometheusTimeSeriesValue ->
                         new PrometheusStandardizedRow(
-                                getBlockFromMap(columnHandles.get(0).getColumnType(), metricHeaderToMap(result.getMetricHeader())),
+                                result.getMetricHeader(),
                                 prometheusTimeSeriesValue.getTimestamp(),
                                 Double.parseDouble(prometheusTimeSeriesValue.getValue())))
                         .collect(Collectors.toList()))

presto-prometheus/src/main/java/com/facebook/presto/plugin/prometheus/PrometheusStandardizedRow.java

@@ -13,26 +13,25 @@
  */
 package com.facebook.presto.plugin.prometheus;
 
-import com.facebook.presto.common.block.Block;
-
 import java.time.Instant;
+import java.util.Map;
 
 import static java.util.Objects.requireNonNull;
 
 public class PrometheusStandardizedRow
 {
-    private final Block labels;
+    private final Map<String, String> labels;
     private final Instant timestamp;
     private final Double value;
 
-    public PrometheusStandardizedRow(Block labels, Instant timestamp, Double value)
+    public PrometheusStandardizedRow(Map<String, String> labels, Instant timestamp, Double value)
     {
         this.labels = requireNonNull(labels, "labels is null");
         this.timestamp = requireNonNull(timestamp, "timestamp is null");
         this.value = requireNonNull(value, "value is null");
     }
 
-    public Block getLabels()
+    public Map<String, String> getLabels()
     {
         return labels;
     }

presto-prometheus/src/test/java/com/facebook/presto/plugin/prometheus/TestPrometheusMetricsIntegration.java

@@ -118,4 +118,13 @@ public void testPushDown()
         MaterializedResult results = runner.execute(session, "SELECT * FROM prometheus.default.up WHERE timestamp > (NOW() - INTERVAL '15' SECOND)").toTestTypes();
         assertEquals(results.getRowCount(), 1);
     }
+    @Test(priority = 3, dependsOnMethods = "testConfirmMetricAvailableAndCheckUp")
+    public void testCountQuery()
+    {
+        MaterializedResult countResult = runner.execute(session, "SELECT COUNT(*) FROM prometheus.default.up").toTestTypes();
+        assertEquals(countResult.getRowCount(), 1);
+        MaterializedRow countRow = countResult.getMaterializedRows().get(0);
+        long countValue = (Long) countRow.getField(0);
+        assert countValue >= 1 : "Expected COUNT(*) to be >= 1, but got " + countValue;
+    }
 }

presto-prometheus/src/test/java/com/facebook/presto/plugin/prometheus/TestPrometheusRecordSet.java

@@ -36,6 +36,7 @@
 import static com.facebook.presto.plugin.prometheus.PrometheusRecordCursor.getBlockFromMap;
 import static com.facebook.presto.plugin.prometheus.PrometheusRecordCursor.getMapFromBlock;
 import static com.facebook.presto.plugin.prometheus.TestPrometheusTable.TYPE_MANAGER;
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static java.time.Instant.ofEpochMilli;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
@@ -64,27 +65,28 @@ public void testCursorSimple()
         List<PrometheusStandardizedRow> actual = new ArrayList<>();
         while (cursor.advanceNextPosition()) {
             actual.add(new PrometheusStandardizedRow(
-                    (Block) cursor.getObject(0),
+                    getMapFromBlock(varcharMapType, (Block) cursor.getObject(0)).entrySet().stream()
+                            .collect(toImmutableMap(entry -> (String) entry.getKey(), entry -> (String) entry.getValue())),
                     ((Instant) cursor.getObject(1)),
                     cursor.getDouble(2)));
             assertFalse(cursor.isNull(0));
             assertFalse(cursor.isNull(1));
             assertFalse(cursor.isNull(2));
         }
         List<PrometheusStandardizedRow> expected = ImmutableList.<PrometheusStandardizedRow>builder()
-                .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
-                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565962969044L), 1.0))
-                .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
-                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565962984045L), 1.0))
-                .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
-                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565962999044L), 1.0))
-                .add(new PrometheusStandardizedRow(getBlockFromMap(varcharMapType,
-                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus")), ofEpochMilli(1565963014044L), 1.0))
+                .add(new PrometheusStandardizedRow(
+                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus"), ofEpochMilli(1565962969044L), 1.0))
+                .add(new PrometheusStandardizedRow(
+                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus"), ofEpochMilli(1565962984045L), 1.0))
+                .add(new PrometheusStandardizedRow(
+                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus"), ofEpochMilli(1565962999044L), 1.0))
+                .add(new PrometheusStandardizedRow(
+                        ImmutableMap.of("instance", "localhost:9090", "__name__", "up", "job", "prometheus"), ofEpochMilli(1565963014044L), 1.0))
                 .build();
         List<PairLike<PrometheusStandardizedRow, PrometheusStandardizedRow>> pairs = Streams.zip(actual.stream(), expected.stream(), PairLike::new)
                 .collect(Collectors.toList());
         pairs.forEach(pair -> {
-            assertEquals(getMapFromBlock(varcharMapType, pair.getFirst().getLabels()), getMapFromBlock(varcharMapType, pair.getSecond().getLabels()));
+            assertEquals(getMapFromBlock(varcharMapType, getBlockFromMap(varcharMapType, pair.getFirst().getLabels())), getMapFromBlock(varcharMapType, getBlockFromMap(varcharMapType, pair.getSecond().getLabels())));
             assertEquals(pair.getFirst().getTimestamp(), pair.getSecond().getTimestamp());
             assertEquals(pair.getFirst().getValue(), pair.getSecond().getValue());
         });