Configure spill space limits in LocalQueryRunner
@@ -97,6 +97,7 @@ import com.facebook.presto.spi.connector.ConnectorFactory;
import com.facebook.presto.spi.type.Type;
import com.facebook.presto.spiller.FileSingleStreamSpillerFactory;
import com.facebook.presto.spiller.GenericSpillerFactory;
+import com.facebook.presto.spiller.NodeSpillConfig;
import com.facebook.presto.spiller.SpillerFactory;
import com.facebook.presto.spiller.SpillerStats;
import com.facebook.presto.split.PageSinkManager;
@@ -185,7 +186,6 @@ import static com.facebook.presto.spi.type.BigintType.BIGINT;
import static com.facebook.presto.sql.testing.TreeAssertions.assertFormattedSql;
import static com.facebook.presto.testing.TestingSession.TESTING_CATALOG;
import static com.facebook.presto.testing.TestingSession.createBogusTestingCatalog;
-import static com.facebook.presto.testing.TestingTaskContext.createTaskContext;
import static com.facebook.presto.transaction.TransactionBuilder.transaction;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
@@ -231,6 +231,7 @@ public class LocalQueryRunner
private final ImmutableMap<Class<? extends Statement>, DataDefinitionTask<?>> dataDefinitionTask;
private final boolean alwaysRevokeMemory;
+ private final NodeSpillConfig nodeSpillConfig;
private boolean printPlan;
private final ReadWriteLock lock = new ReentrantReadWriteLock();
@@ -262,10 +263,16 @@ public class LocalQueryRunner
}
public LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig, boolean withInitialTransaction, boolean alwaysRevokeMemory)
+ {
+ this(defaultSession, featuresConfig, new NodeSpillConfig(), withInitialTransaction, alwaysRevokeMemory);
+ }
+
+ public LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig, NodeSpillConfig nodeSpillConfig, boolean withInitialTransaction, boolean alwaysRevokeMemory)
{
requireNonNull(defaultSession, "defaultSession is null");
checkArgument(!defaultSession.getTransactionId().isPresent() || !withInitialTransaction, "Already in transaction");
+ this.nodeSpillConfig = requireNonNull(nodeSpillConfig, "nodeSpillConfig is null");
this.alwaysRevokeMemory = alwaysRevokeMemory;
this.executor = newCachedThreadPool(daemonThreadsNamed("local-query-runner-%s"));
this.transactionCheckExecutor = newSingleThreadScheduledExecutor(daemonThreadsNamed("transaction-idle-check"));
@@ -386,7 +393,7 @@ public class LocalQueryRunner
public static LocalQueryRunner queryRunnerWithInitialTransaction(Session defaultSession)
{
checkArgument(!defaultSession.getTransactionId().isPresent(), "Already in transaction!");
- return new LocalQueryRunner(defaultSession, new FeaturesConfig(), true, false);
+ return new LocalQueryRunner(defaultSession, new FeaturesConfig(), new NodeSpillConfig(), true, false);
}
@Override
@@ -541,7 +548,10 @@ public class LocalQueryRunner
return builder.get()::page;
});
- TaskContext taskContext = createTaskContext(executor, session);
+ TaskContext taskContext = TestingTaskContext.builder(executor, session)
+ .setMaxSpillSize(nodeSpillConfig.getMaxSpillPerNode())
+ .setQueryMaxSpillSize(nodeSpillConfig.getQueryMaxSpillPerNode())
+ .build();
List<Driver> drivers = createDrivers(session, sql, outputFactory, taskContext);
drivers.forEach(closer::register);
