public boolean hasAllOutputs(TableScanNode node) { if (!layout.getColumns().isPresent()) { return true; } Set<ColumnHandle> columns = ImmutableSet.copyOf(layout.getColumns().get()); List<ColumnHandle> nodeColumnHandles = node.getOutputSymbols().stream() .map(node.getAssignments()::get) .collect(toImmutableList()); return columns.containsAll(nodeColumnHandles); }
@Test public void testGroupingOperationAllBitsSet() { List<Integer> groupingOrdinals = ImmutableList.of(0, 4, 8); List<Set<Integer>> groupingSetOrdinals = ImmutableList.of(ImmutableSet.of(1), ImmutableSet.of(7, 3, 1), ImmutableSet.of(9, 1)); for (Set<Integer> groupingSet : groupingSetOrdinals) { assertEquals(calculateGrouping(groupingSet, groupingOrdinals), 7L); } }
public Set<String> getAllTables(String schema) throws SchemaNotFoundException { ImmutableSet.Builder<String> builder = ImmutableSet.builder(); builder.addAll(ImmutableList.copyOf(client.getDatabase(schema).listCollectionNames()).stream() .filter(name -> !name.equals(schemaCollection)) .filter(name -> !SYSTEM_TABLES.contains(name)) .collect(toSet())); builder.addAll(getTableMetadataNames(schema)); return builder.build(); }
public MaterializedResult( List<MaterializedRow> rows, List<? extends Type> types, Map<String, String> setSessionProperties, Set<String> resetSessionProperties, Optional<String> updateType, OptionalLong updateCount) { this.rows = ImmutableList.copyOf(requireNonNull(rows, "rows is null")); this.types = ImmutableList.copyOf(requireNonNull(types, "types is null")); this.setSessionProperties = ImmutableMap.copyOf(requireNonNull(setSessionProperties, "setSessionProperties is null")); this.resetSessionProperties = ImmutableSet.copyOf(requireNonNull(resetSessionProperties, "resetSessionProperties is null")); this.updateType = requireNonNull(updateType, "updateType is null"); this.updateCount = requireNonNull(updateCount, "updateCount is null"); }
private static List<Type> toTypes(Map<Symbol, Integer> layout, LocalExecutionPlanContext context) { // verify layout covers all values int channelCount = layout.values().stream().mapToInt(Integer::intValue).max().orElse(-1) + 1; checkArgument( layout.size() == channelCount && ImmutableSet.copyOf(layout.values()).containsAll(ContiguousSet.create(closedOpen(0, channelCount), integers())), "Layout does not have a symbol for every output channel: %s", layout); Map<Integer, Symbol> channelLayout = ImmutableBiMap.copyOf(layout).inverse(); return range(0, channelCount) .mapToObj(channelLayout::get) .map(context.getTypes()::get) .collect(toImmutableList()); }
@Override public PlanNode visitUnnest(UnnestNode node, RewriteContext<Set<Symbol>> context) { List<Symbol> replicateSymbols = node.getReplicateSymbols().stream() .filter(context.get()::contains) .collect(toImmutableList()); Optional<Symbol> ordinalitySymbol = node.getOrdinalitySymbol(); if (ordinalitySymbol.isPresent() && !context.get().contains(ordinalitySymbol.get())) { ordinalitySymbol = Optional.empty(); } Map<Symbol, List<Symbol>> unnestSymbols = node.getUnnestSymbols(); ImmutableSet.Builder<Symbol> expectedInputs = ImmutableSet.<Symbol>builder() .addAll(replicateSymbols) .addAll(unnestSymbols.keySet()); PlanNode source = context.rewrite(node.getSource(), expectedInputs.build()); return new UnnestNode(node.getId(), source, replicateSymbols, unnestSymbols, ordinalitySymbol); }
@Override public List<String> listSchemaNames(Session session, String catalogName) { Optional<CatalogMetadata> catalog = getOptionalCatalogMetadata(session, catalogName); ImmutableSet.Builder<String> schemaNames = ImmutableSet.builder(); if (catalog.isPresent()) { CatalogMetadata catalogMetadata = catalog.get(); ConnectorSession connectorSession = session.toConnectorSession(catalogMetadata.getConnectorId()); for (ConnectorId connectorId : catalogMetadata.listConnectorIds()) { ConnectorMetadata metadata = catalogMetadata.getMetadataFor(connectorId); metadata.listSchemaNames(connectorSession).stream() .map(schema -> schema.toLowerCase(Locale.ENGLISH)) .forEach(schemaNames::add); } } return ImmutableList.copyOf(schemaNames.build()); }
@Test public void testMultipleNodes() { FixedCountScheduler nodeScheduler = new FixedCountScheduler( (node, partition, totalPartitions) -> Optional.of(taskFactory.createTableScanTask( new TaskId("test", 1, 1), node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> {}))), generateRandomNodes(5)); ScheduleResult result = nodeScheduler.schedule(); assertTrue(result.isFinished()); assertTrue(result.getBlocked().isDone()); assertEquals(result.getNewTasks().size(), 5); assertEquals(result.getNewTasks().stream().map(RemoteTask::getNodeId).collect(toImmutableSet()).size(), 5); }
private static SignatureBuilder functionSignature(List<String> arguments, String returnType, List<TypeVariableConstraint> typeVariableConstraints) { ImmutableSet<String> literalParameters = ImmutableSet.of("p", "s", "p1", "s1", "p2", "s2", "p3", "s3"); List<TypeSignature> argumentSignatures = arguments.stream() .map((signature) -> TypeSignature.parseTypeSignature(signature, literalParameters)) .collect(toImmutableList()); return new SignatureBuilder() .returnType(TypeSignature.parseTypeSignature(returnType, literalParameters)) .argumentTypes(argumentSignatures) .typeVariableConstraints(typeVariableConstraints) .kind(SCALAR); }
@Test public void testClientTags() { ResourceGroupId resourceGroupId = new ResourceGroupId(new ResourceGroupId("global"), "foo"); StaticSelector selector = new StaticSelector( Optional.empty(), Optional.empty(), Optional.of(ImmutableList.of("tag1", "tag2")), Optional.empty(), Optional.empty(), new ResourceGroupIdTemplate("global.foo")); assertEquals(selector.match(newSelectionCriteria("userA", null, ImmutableSet.of("tag1", "tag2"), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId)); assertEquals(selector.match(newSelectionCriteria("userB", "source", ImmutableSet.of(), EMPTY_RESOURCE_ESTIMATES)), Optional.empty()); assertEquals(selector.match(newSelectionCriteria("A.user", "a source b", ImmutableSet.of("tag1"), EMPTY_RESOURCE_ESTIMATES)), Optional.empty()); assertEquals(selector.match(newSelectionCriteria("A.user", "a source b", ImmutableSet.of("tag1", "tag2", "tag3"), EMPTY_RESOURCE_ESTIMATES)).map(SelectionContext::getResourceGroupId), Optional.of(resourceGroupId)); }
private void validateMetadata(ExtendedHiveMetastore hiveMetastore) { assertEquals(hiveMetastore.getDatabase("database"), Optional.of(DATABASE)); assertEquals(hiveMetastore.getAllDatabases(), ImmutableList.of("database")); assertEquals(hiveMetastore.getTable("database", "table"), Optional.of(TABLE)); assertEquals(hiveMetastore.getSupportedColumnStatistics(createVarcharType(123)), ImmutableSet.of(MIN_VALUE, MAX_VALUE)); assertEquals(hiveMetastore.getTableStatistics("database", "table"), PARTITION_STATISTICS); assertEquals(hiveMetastore.getPartitionStatistics("database", "table", ImmutableSet.of("value")), ImmutableMap.of("value", PARTITION_STATISTICS)); assertEquals(hiveMetastore.getAllTables("database"), Optional.of(ImmutableList.of("table"))); assertEquals(hiveMetastore.getAllViews("database"), Optional.empty()); assertEquals(hiveMetastore.getPartition("database", "table", ImmutableList.of("value")), Optional.of(PARTITION)); assertEquals(hiveMetastore.getPartitionNames("database", "table"), Optional.of(ImmutableList.of("value"))); assertEquals(hiveMetastore.getPartitionNamesByParts("database", "table", ImmutableList.of("value")), Optional.of(ImmutableList.of("value"))); assertEquals(hiveMetastore.getPartitionsByNames("database", "table", ImmutableList.of("value")), ImmutableMap.of("value", Optional.of(PARTITION))); assertEquals(hiveMetastore.getRoles("user"), ImmutableSet.of("role1", "role2")); assertEquals(hiveMetastore.getDatabasePrivileges("user", "database"), ImmutableSet.of(PRIVILEGE_INFO)); assertEquals(hiveMetastore.getTablePrivileges("user", "database", "table"), ImmutableSet.of(PRIVILEGE_INFO)); }
@Override public PlanNode visitDistinctLimit(DistinctLimitNode node, RewriteContext<Set<Symbol>> context) { Set<Symbol> expectedInputs; if (node.getHashSymbol().isPresent()) { expectedInputs = ImmutableSet.copyOf(concat(node.getDistinctSymbols(), ImmutableList.of(node.getHashSymbol().get()))); } else { expectedInputs = ImmutableSet.copyOf(node.getDistinctSymbols()); } PlanNode source = context.rewrite(node.getSource(), expectedInputs); return new DistinctLimitNode(node.getId(), source, node.getLimit(), node.isPartial(), node.getDistinctSymbols(), node.getHashSymbol()); }
@Override public synchronized void updatePartitionStatistics(String databaseName, String tableName, String partitionName, Function<PartitionStatistics, PartitionStatistics> update) { PartitionStatistics currentStatistics = requireNonNull( getPartitionStatistics(databaseName, tableName, ImmutableSet.of(partitionName)).get(partitionName), "getPartitionStatistics() returned null"); PartitionStatistics updatedStatistics = update.apply(currentStatistics); List<Partition> partitions = getPartitionsByNames(databaseName, tableName, ImmutableList.of(partitionName)); if (partitions.size() != 1) { throw new PrestoException(HIVE_METASTORE_ERROR, "Metastore returned multiple partitions for name: " + partitionName); } Partition originalPartition = getOnlyElement(partitions); Partition modifiedPartition = originalPartition.deepCopy(); HiveBasicStatistics basicStatistics = updatedStatistics.getBasicStatistics(); modifiedPartition.setParameters(updateStatisticsParameters(modifiedPartition.getParameters(), basicStatistics)); alterPartitionWithoutStatistics(databaseName, tableName, modifiedPartition); Map<String, HiveType> columns = modifiedPartition.getSd().getCols().stream() .collect(toImmutableMap(FieldSchema::getName, schema -> HiveType.valueOf(schema.getType()))); setPartitionColumnStatistics(databaseName, tableName, partitionName, columns, updatedStatistics.getColumnStatistics(), basicStatistics.getRowCount()); Set<String> removedStatistics = difference(currentStatistics.getColumnStatistics().keySet(), updatedStatistics.getColumnStatistics().keySet()); removedStatistics.forEach(column -> deletePartitionColumnStatistics(databaseName, tableName, partitionName, column)); }
@Parameterized.Parameters public static Iterable<Object[]> constructorFeeder() { return Iterables.transform( Sets.cartesianProduct( ImmutableList.of( ImmutableSet.of(true, false), ImmutableSet.of("", "MISSING VALUE"), ImmutableSet.of(Optional.of(true), Optional.of(false), Optional.empty()) ) ), List::toArray ); }
public StaticSystemTablesProvider(Set<SystemTable> systemTables) { this.systemTables = ImmutableSet.copyOf(systemTables); this.systemTablesMap = systemTables.stream() .collect(toImmutableMap( table -> table.getTableMetadata().getTable(), identity())); }
static <T> List<T> without(List<T> list, Collection<Integer> indexes) { Set<Integer> indexesSet = ImmutableSet.copyOf(indexes); return IntStream.range(0, list.size()) .filter(index -> !indexesSet.contains(index)) .mapToObj(list::get) .collect(toImmutableList()); }
@JsonCreator public SessionMatchSpec( @JsonProperty("user") Optional<Pattern> userRegex, @JsonProperty("source") Optional<Pattern> sourceRegex, @JsonProperty("clientTags") Optional<List<String>> clientTags, @JsonProperty("queryType") Optional<String> queryType, @JsonProperty("group") Optional<Pattern> resourceGroupRegex, @JsonProperty("sessionProperties") Map<String, String> sessionProperties) { this.userRegex = requireNonNull(userRegex, "userRegex is null"); this.sourceRegex = requireNonNull(sourceRegex, "sourceRegex is null"); requireNonNull(clientTags, "clientTags is null"); this.clientTags = ImmutableSet.copyOf(clientTags.orElse(ImmutableList.of())); this.queryType = requireNonNull(queryType, "queryType is null"); this.resourceGroupRegex = requireNonNull(resourceGroupRegex, "resourceGroupRegex is null"); requireNonNull(sessionProperties, "sessionProperties is null"); this.sessionProperties = ImmutableMap.copyOf(sessionProperties); }
@Test(timeOut = 60 * 1000) public void testTopologyAwareScheduling() throws Exception 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)); 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)))); nonRackLocalSplits = Sets.difference(nonRackLocalSplits, new HashSet<>(assignments.values())); assertEquals(nonRackLocalSplits.size(), 3); 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))); ImmutableSet.Builder<Split> localSplits = ImmutableSet.builder();