Examples with ValueSet com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet used on opensource projects

Example 21 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class TimestreamRecordHandlerTest method readRecordsTimeSeriesView.

@Test
public void readRecordsTimeSeriesView() throws Exception {
    logger.info("readRecordsTimeSeriesView - enter");
    Schema schemaForReadView = SchemaBuilder.newBuilder().addField("region", Types.MinorType.VARCHAR.getType()).addField("az", Types.MinorType.VARCHAR.getType()).addField("hostname", Types.MinorType.VARCHAR.getType()).addField(FieldBuilder.newBuilder("cpu_utilization", Types.MinorType.LIST.getType()).addField(FieldBuilder.newBuilder("cpu_utilization", Types.MinorType.STRUCT.getType()).addDateMilliField("time").addFloat8Field("measure_value::double").build()).build()).addMetadata(VIEW_METADATA_FIELD, "select az, hostname, region,  CREATE_TIME_SERIES(time, measure_value::double) as cpu_utilization from \"" + DEFAULT_SCHEMA + "\".\"" + TEST_TABLE + "\" WHERE measure_name = 'cpu_utilization' GROUP BY measure_name, az, hostname, region").build();
    String expectedQuery = "WITH t1 AS ( select az, hostname, region,  CREATE_TIME_SERIES(time, measure_value::double) as cpu_utilization from \"my_schema\".\"my_table\" WHERE measure_name = 'cpu_utilization' GROUP BY measure_name, az, hostname, region )  SELECT region, az, hostname, cpu_utilization FROM t1 WHERE (\"az\" IN ('us-east-1a','us-east-1b'))";
    QueryResult mockResult = makeMockQueryResult(schemaForReadView, 1_000);
    when(mockClient.query(any(QueryRequest.class))).thenAnswer((Answer<QueryResult>) invocationOnMock -> {
        QueryRequest request = (QueryRequest) invocationOnMock.getArguments()[0];
        assertEquals("actual: " + request.getQueryString(), expectedQuery, request.getQueryString().replace("\n", ""));
        return mockResult;
    });
    S3SpillLocation splitLoc = S3SpillLocation.newBuilder().withBucket(UUID.randomUUID().toString()).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build();
    Split split = Split.newBuilder(splitLoc, null).build();
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("az", EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, true).add("us-east-1a").add("us-east-1b").build());
    ReadRecordsRequest request = new ReadRecordsRequest(IDENTITY, "default", "queryId-" + System.currentTimeMillis(), new TableName(DEFAULT_SCHEMA, TEST_TABLE), schemaForReadView, split, new Constraints(constraintsMap), // 100GB don't expect this to spill
    100_000_000_000L, 100_000_000_000L);
    RecordResponse rawResponse = handler.doReadRecords(allocator, request);
    ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
    logger.info("readRecordsTimeSeriesView: rows[{}]", response.getRecordCount());
    for (int i = 0; i < response.getRecordCount() && i < 10; i++) {
        logger.info("readRecordsTimeSeriesView: {}", BlockUtils.rowToString(response.getRecords(), i));
    }
    logger.info("readRecordsTimeSeriesView - exit");
}

Example 22 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class SelectQueryBuilderTest method build.

@Test
public void build() {
    logger.info("build: enter");
    String expected = "SELECT col1, col2, col3, col4 FROM \"myDatabase\".\"myTable\" WHERE (\"col4\" IN ('val1','val2')) AND ((\"col2\" < 1)) AND (\"col3\" IN (20000,10000)) AND ((\"col1\" > 1))";
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col1", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.INT.getType(), 1)), false));
    constraintsMap.put("col2", SortedRangeSet.copyOf(Types.MinorType.INT.getType(), ImmutableList.of(Range.lessThan(allocator, Types.MinorType.INT.getType(), 1)), false));
    constraintsMap.put("col3", EquatableValueSet.newBuilder(allocator, Types.MinorType.INT.getType(), true, true).add(20000L).add(10000L).build());
    constraintsMap.put("col4", EquatableValueSet.newBuilder(allocator, Types.MinorType.VARCHAR.getType(), true, true).add("val1").add("val2").build());
    Schema schema = SchemaBuilder.newBuilder().addStringField("col1").addIntField("col2").addBigIntField("col3").addStringField("col4").build();
    String actual = queryFactory.createSelectQueryBuilder(VIEW_METADATA_FIELD).withDatabaseName("myDatabase").withTableName("myTable").withProjection(schema).withConjucts(new Constraints(constraintsMap)).build().replace("\n", "");
    logger.info("build: actual[{}]", actual);
    assertEquals(expected, actual);
    logger.info("build: exit");
}

Example 23 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class VerticaMetadataHandlerTest method getPartitions.

@Test
public void getPartitions() throws Exception {
    Schema tableSchema = SchemaBuilder.newBuilder().addIntField("day").addIntField("month").addIntField("year").addStringField("preparedStmt").addStringField("queryId").addStringField("awsRegionSql").build();
    Set<String> partitionCols = new HashSet<>();
    partitionCols.add("preparedStmt");
    partitionCols.add("queryId");
    partitionCols.add("awsRegionSql");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("day", SortedRangeSet.copyOf(org.apache.arrow.vector.types.Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, org.apache.arrow.vector.types.Types.MinorType.INT.getType(), 0)), false));
    constraintsMap.put("month", SortedRangeSet.copyOf(org.apache.arrow.vector.types.Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, org.apache.arrow.vector.types.Types.MinorType.INT.getType(), 0)), false));
    constraintsMap.put("year", SortedRangeSet.copyOf(org.apache.arrow.vector.types.Types.MinorType.INT.getType(), ImmutableList.of(Range.greaterThan(allocator, org.apache.arrow.vector.types.Types.MinorType.INT.getType(), 2000)), false));
    GetTableLayoutRequest req = null;
    GetTableLayoutResponse res = null;
    String testSql = "Select * from schema1.table1";
    String[] test = new String[] { "Select * from schema1.table1", "Select * from schema1.table1" };
    String[] schema = { "TABLE_SCHEM", "TABLE_NAME", "COLUMN_NAME", "TYPE_NAME" };
    Object[][] values = { { "testSchema", "testTable1", "day", "int" }, { "testSchema", "testTable1", "month", "int" }, { "testSchema", "testTable1", "year", "int" }, { "testSchema", "testTable1", "preparedStmt", "varchar" }, { "testSchema", "testTable1", "queryId", "varchar" }, { "testSchema", "testTable1", "awsRegionSql", "varchar" } };
    int[] types = { Types.INTEGER, Types.INTEGER, Types.INTEGER, Types.VARCHAR, Types.VARCHAR, Types.VARCHAR };
    List<TableName> expectedTables = new ArrayList<>();
    expectedTables.add(new TableName("testSchema", "testTable1"));
    AtomicInteger rowNumber = new AtomicInteger(-1);
    ResultSet resultSet = mockResultSet(schema, types, values, rowNumber);
    Mockito.when(connection.getMetaData().getColumns(null, "schema1", "table1", null)).thenReturn(resultSet);
    Mockito.when(queryFactory.createVerticaExportQueryBuilder()).thenReturn(new VerticaExportQueryBuilder(new ST("templateVerticaExportQuery")));
    Mockito.when(verticaMetadataHandlerMocked.getS3ExportBucket()).thenReturn("testS3Bucket");
    try {
        req = new GetTableLayoutRequest(this.federatedIdentity, "queryId", "default", new TableName("schema1", "table1"), new Constraints(constraintsMap), tableSchema, partitionCols);
        res = verticaMetadataHandlerMocked.doGetTableLayout(allocator, req);
        Block partitions = res.getPartitions();
        String actualQueryID = partitions.getFieldReader("queryId").readText().toString();
        String expectedExportSql = "EXPORT TO PARQUET(directory = 's3://testS3Bucket/" + actualQueryID + "', Compression='snappy', fileSizeMB=16, rowGroupSizeMB=16) " + "AS SELECT day,month,year,preparedStmt,queryId,awsRegionSql " + "FROM \"schema1\".\"table1\" " + "WHERE ((\"day\" > 0 )) AND ((\"month\" > 0 )) AND ((\"year\" > 2000 ))";
        Assert.assertEquals(expectedExportSql, partitions.getFieldReader("preparedStmt").readText().toString());
        for (int row = 0; row < partitions.getRowCount() && row < 1; row++) {
            logger.info("doGetTableLayout:{} {}", row, BlockUtils.rowToString(partitions, row));
        }
        assertTrue(partitions.getRowCount() > 0);
        logger.info("doGetTableLayout: partitions[{}]", partitions.getRowCount());
    } finally {
        try {
            req.close();
            res.close();
        } catch (Exception ex) {
            logger.error("doGetTableLayout: ", ex);
        }
    }
    logger.info("doGetTableLayout - exit");
}

Example 24 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class SecurityGroupsTableProvider method readWithConstraint.

/**
 * Calls DescribeSecurityGroups on the AWS EC2 Client returning all SecurityGroup rules that match the supplied
 * predicate and attempting to push down certain predicates (namely queries for specific SecurityGroups) to EC2.
 *
 * @See TableProvider
 */
@Override
public void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest recordsRequest, QueryStatusChecker queryStatusChecker) {
    boolean done = false;
    DescribeSecurityGroupsRequest request = new DescribeSecurityGroupsRequest();
    ValueSet idConstraint = recordsRequest.getConstraints().getSummary().get("id");
    if (idConstraint != null && idConstraint.isSingleValue()) {
        request.setGroupIds(Collections.singletonList(idConstraint.getSingleValue().toString()));
    }
    ValueSet nameConstraint = recordsRequest.getConstraints().getSummary().get("name");
    if (nameConstraint != null && nameConstraint.isSingleValue()) {
        request.setGroupNames(Collections.singletonList(nameConstraint.getSingleValue().toString()));
    }
    while (!done) {
        DescribeSecurityGroupsResult response = ec2.describeSecurityGroups(request);
        // Each rule is mapped to a row in the response. SGs have INGRESS and EGRESS rules.
        for (SecurityGroup next : response.getSecurityGroups()) {
            for (IpPermission nextPerm : next.getIpPermissions()) {
                instanceToRow(next, nextPerm, INGRESS, spiller);
            }
            for (IpPermission nextPerm : next.getIpPermissionsEgress()) {
                instanceToRow(next, nextPerm, EGRESS, spiller);
            }
        }
        request.setNextToken(response.getNextToken());
        if (response.getNextToken() == null || !queryStatusChecker.isQueryRunning()) {
            done = true;
        }
    }
}

Example 25 with ValueSet

use of com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet in project aws-athena-query-federation by awslabs.

the class VpcTableProvider method readWithConstraint.

/**
 * Calls DescribeVPCs on the AWS EC2 Client returning all VPCs that match the supplied predicate and attempting
 * to push down certain predicates (namely queries for specific VPCs) to EC2.
 *
 * @See TableProvider
 */
@Override
public void readWithConstraint(BlockSpiller spiller, ReadRecordsRequest recordsRequest, QueryStatusChecker queryStatusChecker) {
    DescribeVpcsRequest request = new DescribeVpcsRequest();
    ValueSet idConstraint = recordsRequest.getConstraints().getSummary().get("id");
    if (idConstraint != null && idConstraint.isSingleValue()) {
        request.setVpcIds(Collections.singletonList(idConstraint.getSingleValue().toString()));
    }
    DescribeVpcsResult response = ec2.describeVpcs(request);
    for (Vpc vpc : response.getVpcs()) {
        instanceToRow(vpc, spiller);
    }
}