@VisibleForTesting
    public static Optional<QueryId> getMaxMemoryConsumingQuery(ListMultimap<QueryId, SqlTask> queryIDToSqlTaskMap)
    {
        if (queryIDToSqlTaskMap.isEmpty()) {
            return Optional.empty();
        }

        Comparator<Map.Entry<QueryId, Long>> comparator = Comparator.comparingLong(Map.Entry::getValue);

        Optional<QueryId> maxMemoryConsumpingQueryId = queryIDToSqlTaskMap.asMap().entrySet().stream()
                .map(entry ->
                        new AbstractMap.SimpleEntry<>(entry.getKey(), entry.getValue().stream()
                                .map(SqlTask::getTaskInfo)
                                .map(TaskInfo::getStats)
                                .mapToLong(stats -> stats.getUserMemoryReservationInBytes() + stats.getSystemMemoryReservationInBytes() + stats.getRevocableMemoryReservationInBytes())
                                .sum())
                ).max(comparator).map(Map.Entry::getKey);

        return maxMemoryConsumpingQueryId;
    }
@Test
    public void testMaxMemoryConsumingQuery()
            throws Exception
    {
        QueryId highMemoryQueryId = new QueryId("query1");
        SqlTask highMemoryTask = createInitialTask(highMemoryQueryId);
        updateTaskMemory(highMemoryTask, 200);

        QueryId lowMemoryQueryId = new QueryId("query2");
        SqlTask lowMemoryTask = createInitialTask(lowMemoryQueryId);
        updateTaskMemory(lowMemoryTask, 100);

        List<SqlTask> activeTasks = ImmutableList.of(highMemoryTask, lowMemoryTask);

        ListMultimap<QueryId, SqlTask> activeQueriesToTasksMap = activeTasks.stream()
                .collect(toImmutableListMultimap(task -> task.getQueryContext().getQueryId(), Function.identity()));

        Optional<QueryId> optionalQueryId = HighMemoryTaskKiller.getMaxMemoryConsumingQuery(activeQueriesToTasksMap);

        assertTrue(optionalQueryId.isPresent());
        assertEquals(optionalQueryId.get(), highMemoryQueryId);
    }