Search in sources :

Example 6 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class TestExpressionMatching method testMatchCall.

@Test
public void testMatchCall() {
    ConnectorExpression expression = new Call(createDecimalType(21, 2), new FunctionName("add"), List.of(new Variable("first", createDecimalType(10, 2)), new Variable("second", BIGINT)));
    ExpressionPattern pattern = expressionPattern("add(foo: decimal(p, s), bar: bigint)");
    Match match = pattern.getPattern().match(expression).collect(onlyElement());
    MatchContext matchContext = new MatchContext();
    pattern.resolve(match.captures(), matchContext);
    assertThat(matchContext.keys()).containsExactlyInAnyOrder("p", "s", "foo", "bar");
    assertThat(matchContext.get("p")).isEqualTo(10L);
    assertThat(matchContext.get("s")).isEqualTo(2L);
    assertThat(matchContext.get("foo")).isEqualTo(new Variable("first", createDecimalType(10, 2)));
    assertThat(matchContext.get("bar")).isEqualTo(new Variable("second", BIGINT));
}
Also used : Call(io.trino.spi.expression.Call) FunctionName(io.trino.spi.expression.FunctionName) Variable(io.trino.spi.expression.Variable) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Match(io.trino.matching.Match) Test(org.testng.annotations.Test)

Example 7 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class TestExpressionMatching method testExpressionCapture.

@Test
public void testExpressionCapture() {
    ConnectorExpression expression = new Call(createDecimalType(21, 2), new FunctionName("add"), List.of(new Variable("first", createDecimalType(10, 2)), new Variable("second", BIGINT)));
    ExpressionPattern pattern = expressionPattern("foo: decimal(p, s)");
    Match match = pattern.getPattern().match(expression).collect(onlyElement());
    MatchContext matchContext = new MatchContext();
    pattern.resolve(match.captures(), matchContext);
    assertThat(matchContext.keys()).containsExactlyInAnyOrder("p", "s", "foo");
    assertThat(matchContext.get("p")).isEqualTo(21L);
    assertThat(matchContext.get("s")).isEqualTo(2L);
    assertThat(matchContext.get("foo")).isSameAs(expression);
}
Also used : Call(io.trino.spi.expression.Call) FunctionName(io.trino.spi.expression.FunctionName) Variable(io.trino.spi.expression.Variable) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) Match(io.trino.matching.Match) Test(org.testng.annotations.Test)

Example 8 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class TestPruneTableScanColumns method mockApplyProjection.

private Optional<ProjectionApplicationResult<ConnectorTableHandle>> mockApplyProjection(ConnectorSession session, ConnectorTableHandle tableHandle, List<ConnectorExpression> projections, Map<String, ColumnHandle> assignments) {
    MockConnectorTableHandle handle = (MockConnectorTableHandle) tableHandle;
    List<Variable> variables = projections.stream().map(Variable.class::cast).collect(toImmutableList());
    List<ColumnHandle> newColumns = variables.stream().map(variable -> assignments.get(variable.getName())).collect(toImmutableList());
    if (handle.getColumns().isPresent() && newColumns.equals(handle.getColumns().get())) {
        return Optional.empty();
    }
    return Optional.of(new ProjectionApplicationResult<>(new MockConnectorTableHandle(handle.getTableName(), handle.getConstraint(), Optional.of(newColumns)), projections, variables.stream().map(variable -> new Assignment(variable.getName(), assignments.get(variable.getName()), ((MockConnectorColumnHandle) assignments.get(variable.getName())).getType())).collect(toImmutableList()), false));
}
Also used : ColumnMetadata(io.trino.spi.connector.ColumnMetadata) PlanMatchPattern(io.trino.sql.planner.assertions.PlanMatchPattern) Test(org.testng.annotations.Test) Variable(io.trino.spi.expression.Variable) CatalogName(io.trino.connector.CatalogName) ImmutableList(com.google.common.collect.ImmutableList) MockConnectorFactory(io.trino.connector.MockConnectorFactory) TpchTableHandle(io.trino.plugin.tpch.TpchTableHandle) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) Map(java.util.Map) ColumnHandle(io.trino.spi.connector.ColumnHandle) ProjectionApplicationResult(io.trino.spi.connector.ProjectionApplicationResult) Symbol(io.trino.sql.planner.Symbol) TINY_SCHEMA_NAME(io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) RuleTester.defaultRuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester.defaultRuleTester) ImmutableMap(com.google.common.collect.ImmutableMap) BaseRuleTest(io.trino.sql.planner.iterative.rule.test.BaseRuleTest) TINY_SCALE_FACTOR(io.trino.plugin.tpch.TpchMetadata.TINY_SCALE_FACTOR) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) Domain(io.trino.spi.predicate.Domain) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) Assignments(io.trino.sql.planner.plan.Assignments) ConnectorSession(io.trino.spi.connector.ConnectorSession) RuleTester(io.trino.sql.planner.iterative.rule.test.RuleTester) TupleDomain(io.trino.spi.predicate.TupleDomain) SchemaTableName(io.trino.spi.connector.SchemaTableName) TpchTransactionHandle(io.trino.plugin.tpch.TpchTransactionHandle) DOUBLE(io.trino.spi.type.DoubleType.DOUBLE) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle) MockConnectorTransactionHandle(io.trino.connector.MockConnectorTransactionHandle) List(java.util.List) TestingSession.testSessionBuilder(io.trino.testing.TestingSession.testSessionBuilder) TableHandle(io.trino.metadata.TableHandle) ConnectorExpression(io.trino.spi.expression.ConnectorExpression) PlanMatchPattern.strictProject(io.trino.sql.planner.assertions.PlanMatchPattern.strictProject) Assignment(io.trino.spi.connector.Assignment) PlanMatchPattern.strictConstrainedTableScan(io.trino.sql.planner.assertions.PlanMatchPattern.strictConstrainedTableScan) Optional(java.util.Optional) TestingColumnHandle(io.trino.testing.TestingMetadata.TestingColumnHandle) PlanBuilder.expression(io.trino.sql.planner.iterative.rule.test.PlanBuilder.expression) DATE(io.trino.spi.type.DateType.DATE) PlanMatchPattern.tableScan(io.trino.sql.planner.assertions.PlanMatchPattern.tableScan) PlanMatchPattern.strictTableScan(io.trino.sql.planner.assertions.PlanMatchPattern.strictTableScan) Assignment(io.trino.spi.connector.Assignment) ColumnHandle(io.trino.spi.connector.ColumnHandle) TpchColumnHandle(io.trino.plugin.tpch.TpchColumnHandle) MockConnectorColumnHandle(io.trino.connector.MockConnectorColumnHandle) TestingColumnHandle(io.trino.testing.TestingMetadata.TestingColumnHandle) Variable(io.trino.spi.expression.Variable) MockConnectorTableHandle(io.trino.connector.MockConnectorTableHandle)

Example 9 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class BaseImplementAvgBigint method rewrite.

@Override
public Optional<JdbcExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<String> context) {
    Variable argument = captures.get(this.argument);
    verify(aggregateFunction.getOutputType() == DOUBLE);
    return Optional.of(new JdbcExpression(format(getRewriteFormatExpression(), context.rewriteExpression(argument).orElseThrow()), new JdbcTypeHandle(Types.DOUBLE, Optional.of("double"), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty())));
}
Also used : JdbcTypeHandle(io.trino.plugin.jdbc.JdbcTypeHandle) Variable(io.trino.spi.expression.Variable) JdbcExpression(io.trino.plugin.jdbc.JdbcExpression)

Example 10 with Variable

use of io.trino.spi.expression.Variable in project trino by trinodb.

the class ImplementAvgDecimal method rewrite.

@Override
public Optional<JdbcExpression> rewrite(AggregateFunction aggregateFunction, Captures captures, RewriteContext<String> context) {
    Variable argument = captures.get(ARGUMENT);
    JdbcColumnHandle columnHandle = (JdbcColumnHandle) context.getAssignment(argument.getName());
    DecimalType type = (DecimalType) columnHandle.getColumnType();
    verify(aggregateFunction.getOutputType().equals(type));
    return Optional.of(new JdbcExpression(format("CAST(avg(%s) AS decimal(%s, %s))", context.rewriteExpression(argument).orElseThrow(), type.getPrecision(), type.getScale()), columnHandle.getJdbcTypeHandle()));
}
Also used : Variable(io.trino.spi.expression.Variable) JdbcColumnHandle(io.trino.plugin.jdbc.JdbcColumnHandle) DecimalType(io.trino.spi.type.DecimalType) JdbcExpression(io.trino.plugin.jdbc.JdbcExpression)

Aggregations

Variable (io.trino.spi.expression.Variable)51 JdbcExpression (io.trino.plugin.jdbc.JdbcExpression)21 ConnectorExpression (io.trino.spi.expression.ConnectorExpression)21 JdbcColumnHandle (io.trino.plugin.jdbc.JdbcColumnHandle)18 Test (org.testng.annotations.Test)12 ImmutableList (com.google.common.collect.ImmutableList)11 ColumnHandle (io.trino.spi.connector.ColumnHandle)11 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)10 Assignment (io.trino.spi.connector.Assignment)10 Map (java.util.Map)9 Optional (java.util.Optional)9 ImmutableMap (com.google.common.collect.ImmutableMap)8 TupleDomain (io.trino.spi.predicate.TupleDomain)8 List (java.util.List)8 ImmutableMap.toImmutableMap (com.google.common.collect.ImmutableMap.toImmutableMap)7 AggregateFunction (io.trino.spi.connector.AggregateFunction)7 ProjectionApplicationResult (io.trino.spi.connector.ProjectionApplicationResult)7 AggregateExpression (io.trino.plugin.pinot.query.AggregateExpression)6 Domain (io.trino.spi.predicate.Domain)6 Verify.verify (com.google.common.base.Verify.verify)5