/*
 * 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.facebook.presto.benchmark;

import com.facebook.presto.Session;
import com.facebook.presto.execution.TaskId;
import com.facebook.presto.execution.TaskStateMachine;
import com.facebook.presto.memory.MemoryPool;
import com.facebook.presto.memory.QueryContext;
import com.facebook.presto.operator.Driver;
import com.facebook.presto.operator.TaskContext;
import com.facebook.presto.plugin.memory.MemoryConnectorFactory;
import com.facebook.presto.spi.QueryId;
import com.facebook.presto.spi.memory.MemoryPoolId;
import com.facebook.presto.spiller.SpillSpaceTracker;
import com.facebook.presto.testing.LocalQueryRunner;
import com.facebook.presto.testing.NullOutputOperator;
import com.facebook.presto.tpch.TpchConnectorFactory;
import com.google.common.collect.ImmutableMap;
import io.airlift.units.DataSize;
import org.intellij.lang.annotations.Language;

import java.util.List;
import java.util.concurrent.ExecutorService;

import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
import static io.airlift.units.DataSize.Unit.GIGABYTE;
import static io.airlift.units.DataSize.Unit.MEGABYTE;

public class MemoryLocalQueryRunner
{
    protected LocalQueryRunner localQueryRunner = createMemoryLocalQueryRunner();

    public void execute(@Language("SQL") String query)
    {
        Session session = testSessionBuilder()
                .setSystemProperty("optimizer.optimize-hash-generation", "true")
                .build();
        ExecutorService executor = localQueryRunner.getExecutor();
        MemoryPool memoryPool = new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE));
        MemoryPool systemMemoryPool = new MemoryPool(new MemoryPoolId("testSystem"), new DataSize(1, GIGABYTE));
        SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(new DataSize(1, GIGABYTE));

        TaskContext taskContext = new QueryContext(new QueryId("test"), new DataSize(256, MEGABYTE), memoryPool, systemMemoryPool, executor, new DataSize(1, GIGABYTE), spillSpaceTracker)
                .addTaskContext(new TaskStateMachine(new TaskId("query", 0, 0), executor),
                        session,
                        false,
                        false);

        // Use NullOutputFactory to avoid coping out results to avoid affecting benchmark results
        List<Driver> drivers = localQueryRunner.createDrivers(query, new NullOutputOperator.NullOutputFactory(), taskContext);

        boolean done = false;
        while (!done) {
            boolean processed = false;
            for (Driver driver : drivers) {
                if (!driver.isFinished()) {
                    driver.process();
                    processed = true;
                }
            }
            done = !processed;
        }
    }

    private static LocalQueryRunner createMemoryLocalQueryRunner()
    {
        Session.SessionBuilder sessionBuilder = testSessionBuilder()
                .setCatalog("memory")
                .setSchema("default");

        Session session = sessionBuilder.build();
        LocalQueryRunner localQueryRunner = LocalQueryRunner.queryRunnerWithInitialTransaction(session);

        // add tpch
        localQueryRunner.createCatalog("tpch", new TpchConnectorFactory(1), ImmutableMap.<String, String>of());
        localQueryRunner.createCatalog("memory", new MemoryConnectorFactory(), ImmutableMap.<String, String>of());

        return localQueryRunner;
    }
}

@@ -1624,63 +1624,8 @@
 YTE;
-%0Aimport static io.airlift.units.DataSize.Unit.MEGABYTE;
 %0A%0Apu
@@ -2103,33 +2103,33 @@
 ), new DataSize(
-1
+2
 , GIGABYTE));%0A  
@@ -2216,33 +2216,33 @@
 ), new DataSize(
-1
+2
 , GIGABYTE));%0A  
@@ -2231,32 +2231,33 @@
 (2, GIGABYTE));%0A
+%0A
         SpillSpa
@@ -2328,33 +2328,32 @@
 (1, GIGABYTE));%0A
-%0A
         TaskCont
@@ -2425,15 +2425,13 @@
 ize(
-256, ME
+1, GI
 GABY
@@ -2480,33 +2480,33 @@
 r, new DataSize(
-1
+4
 , GIGABYTE), spi
@@ -3757,32 +3757,16 @@
 ableMap.
-%3CString, String%3E
 of());%0A 
@@ -3807,80 +3807,146 @@
 log(
-%22memory%22, new MemoryConnectorFactory(), ImmutableMap.%3CString, String%3Eof(
+%0A                %22memory%22,%0A                new MemoryConnectorFactory(),%0A                ImmutableMap.of(%22memory.max-data-per-node%22, %224GB%22
 ));%0A

Increase memory limits for memory benchmarks