Examples with ImmutableList com.google.common.collect.ImmutableList used on opensource projects

Example 31 with ImmutableList

use of com.google.common.collect.ImmutableList in project presto by prestodb.

the class PageProcessorCompiler method generateProjectColumnarMethod.

private static MethodDefinition generateProjectColumnarMethod(ClassDefinition classDefinition, CallSiteBinder callSiteBinder, String methodName, RowExpression projection, MethodDefinition projectionMethod) {
    Parameter session = arg("session", ConnectorSession.class);
    Parameter page = arg("page", Page.class);
    Parameter selectedPositions = arg("selectedPositions", int[].class);
    Parameter pageBuilder = arg("pageBuilder", PageBuilder.class);
    Parameter projectionIndex = arg("projectionIndex", int.class);
    List<Parameter> params = ImmutableList.<Parameter>builder().add(session).add(page).add(selectedPositions).add(pageBuilder).add(projectionIndex).build();
    MethodDefinition method = classDefinition.declareMethod(a(PRIVATE), methodName, type(Block.class), params);
    BytecodeBlock body = method.getBody();
    Scope scope = method.getScope();
    Variable thisVariable = method.getThis();
    ImmutableList.Builder<Variable> builder = ImmutableList.builder();
    for (int channel : getInputChannels(projection)) {
        Variable blockVariable = scope.declareVariable("block_" + channel, body, page.invoke("getBlock", Block.class, constantInt(channel)));
        builder.add(blockVariable);
    }
    List<Variable> inputs = builder.build();
    Variable positionCount = scope.declareVariable("positionCount", body, page.invoke("getPositionCount", int.class));
    Variable position = scope.declareVariable("position", body, constantInt(0));
    Variable cardinality = scope.declareVariable("cardinality", body, selectedPositions.length());
    Variable outputBlock = scope.declareVariable(Block.class, "outputBlock");
    Variable blockBuilder = scope.declareVariable("blockBuilder", body, pageBuilder.invoke("getBlockBuilder", BlockBuilder.class, projectionIndex));
    Variable type = scope.declareVariable("type", body, pageBuilder.invoke("getType", Type.class, projectionIndex));
    BytecodeBlock projectBlock = new BytecodeBlock().append(new ForLoop().initialize(position.set(constantInt(0))).condition(lessThan(position, cardinality)).update(position.increment()).body(invokeProject(thisVariable, session, inputs, selectedPositions.getElement(position), pageBuilder, projectionIndex, projectionMethod))).append(outputBlock.set(blockBuilder.invoke("build", Block.class)));
    if (isIdentityExpression(projection)) {
        // if nothing is filtered out, copy the entire block, else project it
        body.append(new IfStatement().condition(equal(cardinality, positionCount)).ifTrue(new BytecodeBlock().append(inputs.get(0).invoke("assureLoaded", void.class)).append(outputBlock.set(inputs.get(0)))).ifFalse(projectBlock));
    } else if (isConstantExpression(projection)) {
        // if projection is a constant, create RLE block of constant expression with cardinality positions
        ConstantExpression constantExpression = (ConstantExpression) projection;
        verify(getInputChannels(projection).isEmpty());
        BytecodeExpression value = loadConstant(callSiteBinder, constantExpression.getValue(), Object.class);
        body.append(outputBlock.set(invokeStatic(RunLengthEncodedBlock.class, "create", Block.class, type, value, cardinality)));
    } else {
        body.append(projectBlock);
    }
    body.append(outputBlock.ret());
    return method;
}

Example 32 with ImmutableList

use of com.google.common.collect.ImmutableList in project presto by prestodb.

the class TryCodeGenerator method generateExpression.

@Override
public BytecodeNode generateExpression(Signature signature, BytecodeGeneratorContext context, Type returnType, List<RowExpression> arguments) {
    checkArgument(arguments.size() == 1, "try methods only contain a single expression");
    checkArgument(getOnlyElement(arguments) instanceof CallExpression, "try methods must contain a call expression");
    CallExpression innerCallExpression = (CallExpression) getOnlyElement(arguments);
    checkState(tryMethodsMap.containsKey(innerCallExpression), "try methods map does not contain this try call");
    MethodDefinition definition = tryMethodsMap.get(innerCallExpression);
    ImmutableList<Variable> invokeArguments = definition.getParameters().stream().map(parameter -> context.getScope().getVariable(parameter.getName())).collect(toImmutableList());
    return new BytecodeBlock().append(context.getScope().getThis().invoke(definition, invokeArguments)).append(unboxPrimitiveIfNecessary(context.getScope(), Primitives.wrap(innerCallExpression.getType().getJavaType())));
}

Example 33 with ImmutableList

use of com.google.common.collect.ImmutableList in project presto by prestodb.

the class VarArgsToMapAdapterGenerator method generateVarArgsToMapAdapter.

/**
     * Generate byte code that
     * <p><ul>
     * <li>takes a specified number of variables as arguments (types of the arguments are provided in {@code javaTypes})
     * <li>put the variables in a map (keys of the map are provided in {@code names})
     * <li>invoke the provided {@code function} with the map
     * <li>return with the result of the function call (type must match {@code returnType})
     * </ul></p>
     */
public static MethodHandle generateVarArgsToMapAdapter(Class<?> returnType, List<Class<?>> javaTypes, List<String> names, Function<Map<String, Object>, Object> function) {
    CallSiteBinder callSiteBinder = new CallSiteBinder();
    ClassDefinition classDefinition = new ClassDefinition(a(PUBLIC, FINAL), makeClassName("VarArgsToMapAdapter"), type(Object.class));
    ImmutableList.Builder<Parameter> parameterListBuilder = ImmutableList.builder();
    for (int i = 0; i < javaTypes.size(); i++) {
        Class<?> javaType = javaTypes.get(i);
        parameterListBuilder.add(arg("input_" + i, javaType));
    }
    ImmutableList<Parameter> parameterList = parameterListBuilder.build();
    MethodDefinition methodDefinition = classDefinition.declareMethod(a(PUBLIC, STATIC), "varArgsToMap", ParameterizedType.type(returnType), parameterList);
    BytecodeBlock body = methodDefinition.getBody();
    // ImmutableMap.Builder can not be used here because it doesn't allow nulls.
    Variable map = methodDefinition.getScope().declareVariable(HashMap.class, "map");
    body.append(map.set(invokeStatic(Maps.class, "newHashMapWithExpectedSize", HashMap.class, constantInt(javaTypes.size()))));
    for (int i = 0; i < javaTypes.size(); i++) {
        body.append(map.invoke("put", Object.class, constantString(names.get(i)).cast(Object.class), parameterList.get(i).cast(Object.class)));
    }
    body.append(loadConstant(callSiteBinder, function, Function.class).invoke("apply", Object.class, map.cast(Object.class)).cast(returnType).ret());
    Class<?> generatedClass = defineClass(classDefinition, Object.class, callSiteBinder.getBindings(), new DynamicClassLoader(VarArgsToMapAdapterGenerator.class.getClassLoader()));
    return Reflection.methodHandle(generatedClass, "varArgsToMap", javaTypes.toArray(new Class<?>[javaTypes.size()]));
}

Example 34 with ImmutableList

use of com.google.common.collect.ImmutableList in project presto by prestodb.

the class RelationPlanner method planCrossJoinUnnest.

private RelationPlan planCrossJoinUnnest(RelationPlan leftPlan, Join joinNode, Unnest node) {
    RelationType unnestOutputDescriptor = analysis.getOutputDescriptor(node);
    // Create symbols for the result of unnesting
    ImmutableList.Builder<Symbol> unnestedSymbolsBuilder = ImmutableList.builder();
    for (Field field : unnestOutputDescriptor.getVisibleFields()) {
        Symbol symbol = symbolAllocator.newSymbol(field);
        unnestedSymbolsBuilder.add(symbol);
    }
    ImmutableList<Symbol> unnestedSymbols = unnestedSymbolsBuilder.build();
    // Add a projection for all the unnest arguments
    PlanBuilder planBuilder = initializePlanBuilder(leftPlan);
    planBuilder = planBuilder.appendProjections(node.getExpressions(), symbolAllocator, idAllocator);
    TranslationMap translations = planBuilder.getTranslations();
    ProjectNode projectNode = (ProjectNode) planBuilder.getRoot();
    ImmutableMap.Builder<Symbol, List<Symbol>> unnestSymbols = ImmutableMap.builder();
    UnmodifiableIterator<Symbol> unnestedSymbolsIterator = unnestedSymbols.iterator();
    for (Expression expression : node.getExpressions()) {
        Type type = analysis.getType(expression);
        Symbol inputSymbol = translations.get(expression);
        if (type instanceof ArrayType) {
            unnestSymbols.put(inputSymbol, ImmutableList.of(unnestedSymbolsIterator.next()));
        } else if (type instanceof MapType) {
            unnestSymbols.put(inputSymbol, ImmutableList.of(unnestedSymbolsIterator.next(), unnestedSymbolsIterator.next()));
        } else {
            throw new IllegalArgumentException("Unsupported type for UNNEST: " + type);
        }
    }
    Optional<Symbol> ordinalitySymbol = node.isWithOrdinality() ? Optional.of(unnestedSymbolsIterator.next()) : Optional.empty();
    checkState(!unnestedSymbolsIterator.hasNext(), "Not all output symbols were matched with input symbols");
    UnnestNode unnestNode = new UnnestNode(idAllocator.getNextId(), projectNode, leftPlan.getFieldMappings(), unnestSymbols.build(), ordinalitySymbol);
    return new RelationPlan(unnestNode, analysis.getScope(joinNode), unnestNode.getOutputSymbols());
}

Example 35 with ImmutableList

use of com.google.common.collect.ImmutableList in project presto by prestodb.

the class TestNodeScheduler method testTopologyAwareScheduling.

@Test(timeOut = 60 * 1000)
public void testTopologyAwareScheduling() throws Exception {
    TestingTransactionHandle transactionHandle = TestingTransactionHandle.create();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    InMemoryNodeManager nodeManager = new InMemoryNodeManager();
    ImmutableList.Builder<Node> nodeBuilder = ImmutableList.builder();
    nodeBuilder.add(new PrestoNode("node1", URI.create("http://host1.rack1:11"), NodeVersion.UNKNOWN, false));
    nodeBuilder.add(new PrestoNode("node2", URI.create("http://host2.rack1:12"), NodeVersion.UNKNOWN, false));
    nodeBuilder.add(new PrestoNode("node3", URI.create("http://host3.rack2:13"), NodeVersion.UNKNOWN, false));
    ImmutableList<Node> nodes = nodeBuilder.build();
    nodeManager.addNode(CONNECTOR_ID, nodes);
    // contents of taskMap indicate the node-task map for the current stage
    Map<Node, RemoteTask> taskMap = new HashMap<>();
    NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig().setMaxSplitsPerNode(25).setIncludeCoordinator(false).setNetworkTopology("test").setMaxPendingSplitsPerTask(20);
    TestNetworkTopology topology = new TestNetworkTopology();
    NetworkLocationCache locationCache = new NetworkLocationCache(topology) {

        @Override
        public NetworkLocation get(HostAddress host) {
            // Bypass the cache for workers, since we only look them up once and they would all be unresolved otherwise
            if (host.getHostText().startsWith("host")) {
                return topology.locate(host);
            } else {
                return super.get(host);
            }
        }
    };
    NodeScheduler nodeScheduler = new NodeScheduler(locationCache, topology, nodeManager, nodeSchedulerConfig, nodeTaskMap);
    NodeSelector nodeSelector = nodeScheduler.createNodeSelector(CONNECTOR_ID);
    // Fill up the nodes with non-local data
    ImmutableSet.Builder<Split> nonRackLocalBuilder = ImmutableSet.builder();
    for (int i = 0; i < (25 + 11) * 3; i++) {
        nonRackLocalBuilder.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("data.other_rack", 1))));
    }
    Set<Split> nonRackLocalSplits = nonRackLocalBuilder.build();
    Multimap<Node, Split> assignments = nodeSelector.computeAssignments(nonRackLocalSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor);
    int task = 0;
    for (Node node : assignments.keySet()) {
        TaskId taskId = new TaskId("test", 1, task);
        task++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        remoteTask.startSplits(25);
        nodeTaskMap.addTask(node, remoteTask);
        taskMap.put(node, remoteTask);
    }
    // Continue assigning to fill up part of the queue
    nonRackLocalSplits = Sets.difference(nonRackLocalSplits, new HashSet<>(assignments.values()));
    assignments = nodeSelector.computeAssignments(nonRackLocalSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (Node node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    nonRackLocalSplits = Sets.difference(nonRackLocalSplits, new HashSet<>(assignments.values()));
    // Check that 3 of the splits were rejected, since they're non-local
    assertEquals(nonRackLocalSplits.size(), 3);
    // Assign rack-local splits
    ImmutableSet.Builder<Split> rackLocalSplits = ImmutableSet.builder();
    HostAddress dataHost1 = HostAddress.fromParts("data.rack1", 1);
    HostAddress dataHost2 = HostAddress.fromParts("data.rack2", 1);
    for (int i = 0; i < 6 * 2; i++) {
        rackLocalSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(dataHost1)));
    }
    for (int i = 0; i < 6; i++) {
        rackLocalSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(dataHost2)));
    }
    assignments = nodeSelector.computeAssignments(rackLocalSplits.build(), ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (Node node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    Set<Split> unassigned = Sets.difference(rackLocalSplits.build(), new HashSet<>(assignments.values()));
    // Compute the assignments a second time to account for the fact that some splits may not have been assigned due to asynchronous
    // loading of the NetworkLocationCache
    boolean cacheRefreshed = false;
    while (!cacheRefreshed) {
        cacheRefreshed = true;
        if (locationCache.get(dataHost1).equals(ROOT_LOCATION)) {
            cacheRefreshed = false;
        }
        if (locationCache.get(dataHost2).equals(ROOT_LOCATION)) {
            cacheRefreshed = false;
        }
        MILLISECONDS.sleep(10);
    }
    assignments = nodeSelector.computeAssignments(unassigned, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (Node node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    unassigned = Sets.difference(unassigned, new HashSet<>(assignments.values()));
    assertEquals(unassigned.size(), 3);
    int rack1 = 0;
    int rack2 = 0;
    for (Split split : unassigned) {
        String rack = topology.locate(split.getAddresses().get(0)).getSegments().get(0);
        switch(rack) {
            case "rack1":
                rack1++;
                break;
            case "rack2":
                rack2++;
                break;
            default:
                fail();
        }
    }
    assertEquals(rack1, 2);
    assertEquals(rack2, 1);
    // Assign local splits
    ImmutableSet.Builder<Split> localSplits = ImmutableSet.builder();
    localSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("host1.rack1", 1))));
    localSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("host2.rack1", 1))));
    localSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote(HostAddress.fromParts("host3.rack2", 1))));
    assignments = nodeSelector.computeAssignments(localSplits.build(), ImmutableList.copyOf(taskMap.values())).getAssignments();
    assertEquals(assignments.size(), 3);
    assertEquals(assignments.keySet().size(), 3);
}