Examples with SubscriptExpression io.trino.sql.tree.SubscriptExpression used on opensource projects

Example 21 with SubscriptExpression

use of io.trino.sql.tree.SubscriptExpression in project trino by trinodb.

the class PushDownDereferencesThroughTopNRanking method apply.

@Override
public Result apply(ProjectNode projectNode, Captures captures, Context context) {
    TopNRankingNode topNRankingNode = captures.get(CHILD);
    // Extract dereferences from project node assignments for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(projectNode.getAssignments().getExpressions(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    // Exclude dereferences on symbols being used in partitionBy and orderBy
    WindowNode.Specification specification = topNRankingNode.getSpecification();
    dereferences = dereferences.stream().filter(expression -> {
        Symbol symbol = getBase(expression);
        return !specification.getPartitionBy().contains(symbol) && !specification.getOrderingScheme().map(OrderingScheme::getOrderBy).orElse(ImmutableList.of()).contains(symbol);
    }).collect(toImmutableSet());
    if (dereferences.isEmpty()) {
        return Result.empty();
    }
    // Create new symbols for dereference expressions
    Assignments dereferenceAssignments = Assignments.of(dereferences, context.getSession(), context.getSymbolAllocator(), typeAnalyzer);
    // Rewrite project node assignments using new symbols for dereference expressions
    Map<Expression, SymbolReference> mappings = HashBiMap.create(dereferenceAssignments.getMap()).inverse().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
    Assignments newAssignments = projectNode.getAssignments().rewrite(expression -> replaceExpression(expression, mappings));
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), topNRankingNode.replaceChildren(ImmutableList.of(new ProjectNode(context.getIdAllocator().getNextId(), topNRankingNode.getSource(), Assignments.builder().putIdentities(topNRankingNode.getSource().getOutputSymbols()).putAll(dereferenceAssignments).build()))), newAssignments));
}

Example 22 with SubscriptExpression

use of io.trino.sql.tree.SubscriptExpression in project trino by trinodb.

the class TestEqualityInference method testExpressionsThatMayReturnNullOnNonNullInput.

@Test
public void testExpressionsThatMayReturnNullOnNonNullInput() {
    List<Expression> candidates = ImmutableList.of(// try_cast
    new Cast(nameReference("b"), toSqlType(BIGINT), true), functionResolution.functionCallBuilder(QualifiedName.of(TryFunction.NAME)).addArgument(new FunctionType(ImmutableList.of(), VARCHAR), new LambdaExpression(ImmutableList.of(), nameReference("b"))).build(), new NullIfExpression(nameReference("b"), number(1)), new IfExpression(nameReference("b"), number(1), new NullLiteral()), new InPredicate(nameReference("b"), new InListExpression(ImmutableList.of(new NullLiteral()))), new SearchedCaseExpression(ImmutableList.of(new WhenClause(new IsNotNullPredicate(nameReference("b")), new NullLiteral())), Optional.empty()), new SimpleCaseExpression(nameReference("b"), ImmutableList.of(new WhenClause(number(1), new NullLiteral())), Optional.empty()), new SubscriptExpression(new ArrayConstructor(ImmutableList.of(new NullLiteral())), nameReference("b")));
    for (Expression candidate : candidates) {
        EqualityInference inference = EqualityInference.newInstance(metadata, equals(nameReference("b"), nameReference("x")), equals(nameReference("a"), candidate));
        List<Expression> equalities = inference.generateEqualitiesPartitionedBy(symbols("b")).getScopeStraddlingEqualities();
        assertEquals(equalities.size(), 1);
        assertTrue(equalities.get(0).equals(equals(nameReference("x"), nameReference("b"))) || equalities.get(0).equals(equals(nameReference("b"), nameReference("x"))));
    }
}

Example 23 with SubscriptExpression

use of io.trino.sql.tree.SubscriptExpression in project trino by trinodb.

the class TestPartialTranslator method testPartialTranslator.

@Test
public void testPartialTranslator() {
    Expression rowSymbolReference = new SymbolReference("row_symbol_1");
    Expression dereferenceExpression1 = new SubscriptExpression(rowSymbolReference, new LongLiteral("1"));
    Expression dereferenceExpression2 = new SubscriptExpression(rowSymbolReference, new LongLiteral("2"));
    Expression stringLiteral = new StringLiteral("abcd");
    Expression symbolReference1 = new SymbolReference("double_symbol_1");
    assertFullTranslation(symbolReference1);
    assertFullTranslation(dereferenceExpression1);
    assertFullTranslation(stringLiteral);
    Expression binaryExpression = new ArithmeticBinaryExpression(ADD, symbolReference1, dereferenceExpression1);
    assertPartialTranslation(binaryExpression, ImmutableList.of(symbolReference1, dereferenceExpression1));
    List<Expression> functionArguments = ImmutableList.of(stringLiteral, dereferenceExpression2);
    Expression functionCallExpression = new FunctionCall(QualifiedName.of("concat"), functionArguments);
    assertFullTranslation(functionCallExpression);
}

Example 24 with SubscriptExpression

use of io.trino.sql.tree.SubscriptExpression in project trino by trinodb.

the class TestConnectorPushdownRulesWithHive method testPushdownWithDuplicateExpressions.

@Test
public void testPushdownWithDuplicateExpressions() {
    String tableName = "duplicate_expressions";
    tester().getQueryRunner().execute(format("CREATE TABLE  %s (struct_of_bigint, just_bigint) AS SELECT cast(row(5, 6) AS row(a bigint, b bigint)) AS struct_of_int, 5 AS just_bigint WHERE false", tableName));
    PushProjectionIntoTableScan pushProjectionIntoTableScan = new PushProjectionIntoTableScan(tester().getPlannerContext(), tester().getTypeAnalyzer(), new ScalarStatsCalculator(tester().getPlannerContext(), tester().getTypeAnalyzer()));
    HiveTableHandle hiveTable = new HiveTableHandle(SCHEMA_NAME, tableName, ImmutableMap.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty());
    TableHandle table = new TableHandle(new CatalogName(HIVE_CATALOG_NAME), hiveTable, new HiveTransactionHandle(false));
    HiveColumnHandle bigintColumn = createBaseColumn("just_bigint", 1, toHiveType(BIGINT), BIGINT, REGULAR, Optional.empty());
    HiveColumnHandle partialColumn = new HiveColumnHandle("struct_of_bigint", 0, toHiveType(ROW_TYPE), ROW_TYPE, Optional.of(new HiveColumnProjectionInfo(ImmutableList.of(0), ImmutableList.of("a"), toHiveType(BIGINT), BIGINT)), REGULAR, Optional.empty());
    // Test projection pushdown with duplicate column references
    tester().assertThat(pushProjectionIntoTableScan).on(p -> {
        SymbolReference column = p.symbol("just_bigint", BIGINT).toSymbolReference();
        Expression negation = new ArithmeticUnaryExpression(MINUS, column);
        return p.project(Assignments.of(// The column reference is part of both the assignments
        p.symbol("column_ref", BIGINT), column, p.symbol("negated_column_ref", BIGINT), negation), p.tableScan(table, ImmutableList.of(p.symbol("just_bigint", BIGINT)), ImmutableMap.of(p.symbol("just_bigint", BIGINT), bigintColumn)));
    }).matches(project(ImmutableMap.of("column_ref", expression("just_bigint_0"), "negated_column_ref", expression("- just_bigint_0")), tableScan(hiveTable.withProjectedColumns(ImmutableSet.of(bigintColumn))::equals, TupleDomain.all(), ImmutableMap.of("just_bigint_0", bigintColumn::equals))));
    // Test Dereference pushdown
    tester().assertThat(pushProjectionIntoTableScan).on(p -> {
        SubscriptExpression subscript = new SubscriptExpression(p.symbol("struct_of_bigint", ROW_TYPE).toSymbolReference(), new LongLiteral("1"));
        Expression sum = new ArithmeticBinaryExpression(ADD, subscript, new LongLiteral("2"));
        return p.project(Assignments.of(// The subscript expression instance is part of both the assignments
        p.symbol("expr_deref", BIGINT), subscript, p.symbol("expr_deref_2", BIGINT), sum), p.tableScan(table, ImmutableList.of(p.symbol("struct_of_bigint", ROW_TYPE)), ImmutableMap.of(p.symbol("struct_of_bigint", ROW_TYPE), partialColumn.getBaseColumn())));
    }).matches(project(ImmutableMap.of("expr_deref", expression(new SymbolReference("struct_of_bigint#a")), "expr_deref_2", expression(new ArithmeticBinaryExpression(ADD, new SymbolReference("struct_of_bigint#a"), new LongLiteral("2")))), tableScan(hiveTable.withProjectedColumns(ImmutableSet.of(partialColumn))::equals, TupleDomain.all(), ImmutableMap.of("struct_of_bigint#a", partialColumn::equals))));
    metastore.dropTable(SCHEMA_NAME, tableName, true);
}

Example 25 with SubscriptExpression

use of io.trino.sql.tree.SubscriptExpression in project trino by trinodb.

the class TestConnectorPushdownRulesWithHive method testProjectionPushdown.

@Test
public void testProjectionPushdown() {
    String tableName = "projection_test";
    PushProjectionIntoTableScan pushProjectionIntoTableScan = new PushProjectionIntoTableScan(tester().getPlannerContext(), tester().getTypeAnalyzer(), new ScalarStatsCalculator(tester().getPlannerContext(), tester().getTypeAnalyzer()));
    tester().getQueryRunner().execute(format("CREATE TABLE  %s (struct_of_int) AS " + "SELECT cast(row(5, 6) as row(a bigint, b bigint)) as struct_of_int where false", tableName));
    Type baseType = ROW_TYPE;
    HiveColumnHandle partialColumn = new HiveColumnHandle("struct_of_int", 0, toHiveType(baseType), baseType, Optional.of(new HiveColumnProjectionInfo(ImmutableList.of(0), ImmutableList.of("a"), toHiveType(BIGINT), BIGINT)), REGULAR, Optional.empty());
    HiveTableHandle hiveTable = new HiveTableHandle(SCHEMA_NAME, tableName, ImmutableMap.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty());
    TableHandle table = new TableHandle(new CatalogName(HIVE_CATALOG_NAME), hiveTable, new HiveTransactionHandle(false));
    HiveColumnHandle fullColumn = partialColumn.getBaseColumn();
    // Test projected columns pushdown to HiveTableHandle in case of full column references
    tester().assertThat(pushProjectionIntoTableScan).on(p -> p.project(Assignments.of(p.symbol("struct_of_int", baseType), p.symbol("struct_of_int", baseType).toSymbolReference()), p.tableScan(table, ImmutableList.of(p.symbol("struct_of_int", baseType)), ImmutableMap.of(p.symbol("struct_of_int", baseType), fullColumn)))).matches(project(ImmutableMap.of("expr", expression("col")), tableScan(hiveTable.withProjectedColumns(ImmutableSet.of(fullColumn))::equals, TupleDomain.all(), ImmutableMap.of("col", fullColumn::equals))));
    // Rule should return Optional.empty after projected ColumnHandles have been added to HiveTableHandle
    tester().assertThat(pushProjectionIntoTableScan).on(p -> p.project(Assignments.of(p.symbol("struct_of_int", baseType), p.symbol("struct_of_int", baseType).toSymbolReference()), p.tableScan(new TableHandle(new CatalogName(HIVE_CATALOG_NAME), hiveTable.withProjectedColumns(ImmutableSet.of(fullColumn)), new HiveTransactionHandle(false)), ImmutableList.of(p.symbol("struct_of_int", baseType)), ImmutableMap.of(p.symbol("struct_of_int", baseType), fullColumn)))).doesNotFire();
    // Test Dereference pushdown
    tester().assertThat(pushProjectionIntoTableScan).on(p -> p.project(Assignments.of(p.symbol("expr_deref", BIGINT), new SubscriptExpression(p.symbol("struct_of_int", baseType).toSymbolReference(), new LongLiteral("1"))), p.tableScan(table, ImmutableList.of(p.symbol("struct_of_int", baseType)), ImmutableMap.of(p.symbol("struct_of_int", baseType), fullColumn)))).matches(project(ImmutableMap.of("expr_deref", expression(new SymbolReference("struct_of_int#a"))), tableScan(hiveTable.withProjectedColumns(ImmutableSet.of(partialColumn))::equals, TupleDomain.all(), ImmutableMap.of("struct_of_int#a", partialColumn::equals))));
    metastore.dropTable(SCHEMA_NAME, tableName, true);
}