Examples with Session io.trino.Session used on opensource projects

Example 11 with Session

use of io.trino.Session in project trino by trinodb.

the class MetadataManager method resolve.

@VisibleForTesting
public ResolvedFunction resolve(Session session, FunctionBinding functionBinding, FunctionMetadata functionMetadata, FunctionDependencyDeclaration declaration) {
    Map<TypeSignature, Type> dependentTypes = declaration.getTypeDependencies().stream().map(typeSignature -> applyBoundVariables(typeSignature, functionBinding)).collect(toImmutableMap(Function.identity(), typeManager::getType, (left, right) -> left));
    ImmutableSet.Builder<ResolvedFunction> functions = ImmutableSet.builder();
    declaration.getFunctionDependencies().stream().map(functionDependency -> {
        try {
            List<TypeSignature> argumentTypes = applyBoundVariables(functionDependency.getArgumentTypes(), functionBinding);
            return resolvedFunctionInternal(session, functionDependency.getName(), fromTypeSignatures(argumentTypes));
        } catch (TrinoException e) {
            if (functionDependency.isOptional()) {
                return null;
            }
            throw e;
        }
    }).filter(Objects::nonNull).forEach(functions::add);
    declaration.getOperatorDependencies().stream().map(operatorDependency -> {
        try {
            List<TypeSignature> argumentTypes = applyBoundVariables(operatorDependency.getArgumentTypes(), functionBinding);
            return resolvedFunctionInternal(session, QualifiedName.of(mangleOperatorName(operatorDependency.getOperatorType())), fromTypeSignatures(argumentTypes));
        } catch (TrinoException e) {
            if (operatorDependency.isOptional()) {
                return null;
            }
            throw e;
        }
    }).filter(Objects::nonNull).forEach(functions::add);
    declaration.getCastDependencies().stream().map(castDependency -> {
        try {
            Type fromType = typeManager.getType(applyBoundVariables(castDependency.getFromType(), functionBinding));
            Type toType = typeManager.getType(applyBoundVariables(castDependency.getToType(), functionBinding));
            return getCoercion(session, fromType, toType);
        } catch (TrinoException e) {
            if (castDependency.isOptional()) {
                return null;
            }
            throw e;
        }
    }).filter(Objects::nonNull).forEach(functions::add);
    return new ResolvedFunction(functionBinding.getBoundSignature(), functionBinding.getFunctionId(), functionMetadata.getKind(), functionMetadata.isDeterministic(), functionMetadata.getFunctionNullability(), dependentTypes, functions.build());
}

Example 12 with Session

use of io.trino.Session in project trino by trinodb.

the class MetadataManager method applyFilter.

@Override
public Optional<ConstraintApplicationResult<TableHandle>> applyFilter(Session session, TableHandle table, Constraint constraint) {
    CatalogName catalogName = table.getCatalogName();
    ConnectorMetadata metadata = getMetadata(session, catalogName);
    ConnectorSession connectorSession = session.toConnectorSession(catalogName);
    return metadata.applyFilter(connectorSession, table.getConnectorHandle(), constraint).map(result -> result.transform(handle -> new TableHandle(catalogName, handle, table.getTransaction())));
}

Example 13 with Session

use of io.trino.Session in project trino by trinodb.

the class ExtractSpatialJoins method loadKdbTree.

private static KdbTree loadKdbTree(String tableName, Session session, Metadata metadata, SplitManager splitManager, PageSourceManager pageSourceManager) {
    QualifiedObjectName name = toQualifiedObjectName(tableName, session.getCatalog().get(), session.getSchema().get());
    TableHandle tableHandle = metadata.getTableHandle(session, name).orElseThrow(() -> new TrinoException(INVALID_SPATIAL_PARTITIONING, format("Table not found: %s", name)));
    Map<String, ColumnHandle> columnHandles = metadata.getColumnHandles(session, tableHandle);
    List<ColumnHandle> visibleColumnHandles = columnHandles.values().stream().filter(handle -> !metadata.getColumnMetadata(session, tableHandle, handle).isHidden()).collect(toImmutableList());
    checkSpatialPartitioningTable(visibleColumnHandles.size() == 1, "Expected single column for table %s, but found %s columns", name, columnHandles.size());
    ColumnHandle kdbTreeColumn = Iterables.getOnlyElement(visibleColumnHandles);
    Optional<KdbTree> kdbTree = Optional.empty();
    try (SplitSource splitSource = splitManager.getSplits(session, tableHandle, UNGROUPED_SCHEDULING, EMPTY, alwaysTrue())) {
        while (!Thread.currentThread().isInterrupted()) {
            SplitBatch splitBatch = getFutureValue(splitSource.getNextBatch(NOT_PARTITIONED, Lifespan.taskWide(), 1000));
            List<Split> splits = splitBatch.getSplits();
            for (Split split : splits) {
                try (ConnectorPageSource pageSource = pageSourceManager.createPageSource(session, split, tableHandle, ImmutableList.of(kdbTreeColumn), DynamicFilter.EMPTY)) {
                    do {
                        getFutureValue(pageSource.isBlocked());
                        Page page = pageSource.getNextPage();
                        if (page != null && page.getPositionCount() > 0) {
                            checkSpatialPartitioningTable(kdbTree.isEmpty(), "Expected exactly one row for table %s, but found more", name);
                            checkSpatialPartitioningTable(page.getPositionCount() == 1, "Expected exactly one row for table %s, but found %s rows", name, page.getPositionCount());
                            String kdbTreeJson = VARCHAR.getSlice(page.getBlock(0), 0).toStringUtf8();
                            try {
                                kdbTree = Optional.of(KdbTreeUtils.fromJson(kdbTreeJson));
                            } catch (IllegalArgumentException e) {
                                checkSpatialPartitioningTable(false, "Invalid JSON string for KDB tree: %s", e.getMessage());
                            }
                        }
                    } while (!pageSource.isFinished());
                } catch (IOException e) {
                    throw new UncheckedIOException(e);
                }
            }
            if (splitBatch.isLastBatch()) {
                break;
            }
        }
    }
    checkSpatialPartitioningTable(kdbTree.isPresent(), "Expected exactly one row for table %s, but got none", name);
    return kdbTree.get();
}

Example 14 with Session

use of io.trino.Session in project trino by trinodb.

the class InlineProjections method extractInliningTargets.

private static Set<Symbol> extractInliningTargets(PlannerContext plannerContext, ProjectNode parent, ProjectNode child, Session session, TypeAnalyzer typeAnalyzer, TypeProvider types) {
    // candidates for inlining are
    // 1. references to simple constants or symbol references
    // 2. references to complex expressions that
    // a. are not inputs to try() expressions
    // b. appear only once across all expressions
    // c. are not identity projections
    // which come from the child, as opposed to an enclosing scope.
    Set<Symbol> childOutputSet = ImmutableSet.copyOf(child.getOutputSymbols());
    Map<Symbol, Long> dependencies = parent.getAssignments().getExpressions().stream().flatMap(expression -> SymbolsExtractor.extractAll(expression).stream()).filter(childOutputSet::contains).collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
    // find references to simple constants or symbol references
    Set<Symbol> basicReferences = dependencies.keySet().stream().filter(input -> isEffectivelyLiteral(plannerContext, session, child.getAssignments().get(input)) || child.getAssignments().get(input) instanceof SymbolReference).filter(// skip identities, otherwise, this rule will keep firing forever
    input -> !child.getAssignments().isIdentity(input)).collect(toSet());
    // exclude any complex inputs to TRY expressions. Inlining them would potentially
    // change the semantics of those expressions
    Set<Symbol> tryArguments = parent.getAssignments().getExpressions().stream().flatMap(expression -> extractTryArguments(expression).stream()).collect(toSet());
    Set<Symbol> singletons = dependencies.entrySet().stream().filter(// reference appears just once across all expressions in parent project node
    entry -> entry.getValue() == 1).filter(// they are not inputs to TRY. Otherwise, inlining might change semantics
    entry -> !tryArguments.contains(entry.getKey())).filter(// skip identities, otherwise, this rule will keep firing forever
    entry -> !child.getAssignments().isIdentity(entry.getKey())).filter(entry -> {
        // skip dereferences, otherwise, inlining can cause conflicts with PushdownDereferences
        Expression assignment = child.getAssignments().get(entry.getKey());
        if (assignment instanceof SubscriptExpression) {
            if (typeAnalyzer.getType(session, types, ((SubscriptExpression) assignment).getBase()) instanceof RowType) {
                return false;
            }
        }
        return true;
    }).map(Map.Entry::getKey).collect(toSet());
    return Sets.union(singletons, basicReferences);
}

Example 15 with Session

use of io.trino.Session in project trino by trinodb.

the class AddIntermediateAggregations method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    Lookup lookup = context.getLookup();
    PlanNodeIdAllocator idAllocator = context.getIdAllocator();
    Session session = context.getSession();
    Optional<PlanNode> rewrittenSource = recurseToPartial(lookup.resolve(aggregation.getSource()), lookup, idAllocator);
    if (rewrittenSource.isEmpty()) {
        return Result.empty();
    }
    PlanNode source = rewrittenSource.get();
    if (getTaskConcurrency(session) > 1) {
        source = ExchangeNode.partitionedExchange(idAllocator.getNextId(), ExchangeNode.Scope.LOCAL, source, new PartitioningScheme(Partitioning.create(FIXED_ARBITRARY_DISTRIBUTION, ImmutableList.of()), source.getOutputSymbols()));
        source = new AggregationNode(idAllocator.getNextId(), source, inputsAsOutputs(aggregation.getAggregations()), aggregation.getGroupingSets(), aggregation.getPreGroupedSymbols(), AggregationNode.Step.INTERMEDIATE, aggregation.getHashSymbol(), aggregation.getGroupIdSymbol());
        source = ExchangeNode.gatheringExchange(idAllocator.getNextId(), ExchangeNode.Scope.LOCAL, source);
    }
    return Result.ofPlanNode(aggregation.replaceChildren(ImmutableList.of(source)));
}