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

Example 76 with ValueSet

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

the class SelectQueryBuilderTest method buildWithView.

@Test
public void buildWithView() {
    logger.info("build: buildWithView");
    String expected = "WITH t1 AS ( SELECT col1 from test_table )  SELECT col1, col2, col3, col4 FROM t1 WHERE ((\"col2\" < 1)) AND ((\"col1\" > 1))";
    Schema schema = SchemaBuilder.newBuilder().addStringField("col1").addIntField("col2").addBigIntField("col3").addStringField("col4").addMetadata(VIEW_METADATA_FIELD, "SELECT col1 from test_table").build();
    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));
    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: buildWithView");
}

Example 77 with ValueSet

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

the class TPCDSMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    logger.info("doGetTableLayout - enter");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    Schema schema = SchemaBuilder.newBuilder().build();
    GetTableLayoutRequest req = new GetTableLayoutRequest(identity, "queryId", "default", new TableName("tpcds1", "customer"), new Constraints(constraintsMap), schema, Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout - {}", res.getPartitions());
    assertTrue(res.getPartitions().getRowCount() == 1);
    logger.info("doGetTableLayout - exit");
}

Example 78 with ValueSet

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

the class VerticaMetadataHandlerTest method doGetSplits.

@Test
public void doGetSplits() {
    logger.info("doGetSplits: enter");
    Schema schema = SchemaBuilder.newBuilder().addIntField("day").addIntField("month").addIntField("year").addStringField("preparedStmt").addStringField("queryId").addStringField("awsRegionSql").build();
    List<String> partitionCols = new ArrayList<>();
    partitionCols.add("preparedStmt");
    partitionCols.add("queryId");
    partitionCols.add("awsRegionSql");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    Block partitions = allocator.createBlock(schema);
    int num_partitions = 10;
    for (int i = 0; i < num_partitions; i++) {
        BlockUtils.setValue(partitions.getFieldVector("day"), i, 2016 + i);
        BlockUtils.setValue(partitions.getFieldVector("month"), i, (i % 12) + 1);
        BlockUtils.setValue(partitions.getFieldVector("year"), i, (i % 28) + 1);
        BlockUtils.setValue(partitions.getFieldVector("preparedStmt"), i, "test");
        BlockUtils.setValue(partitions.getFieldVector("queryId"), i, "123");
        BlockUtils.setValue(partitions.getFieldVector("awsRegionSql"), i, "us-west-2");
    }
    List<S3ObjectSummary> s3ObjectSummariesList = new ArrayList<>();
    S3ObjectSummary s3ObjectSummary = new S3ObjectSummary();
    s3ObjectSummary.setBucketName("s3ExportBucket");
    s3ObjectSummary.setKey("testKey");
    s3ObjectSummariesList.add(s3ObjectSummary);
    ListObjectsRequest listObjectsRequestObj = new ListObjectsRequest();
    listObjectsRequestObj.setBucketName("s3ExportBucket");
    listObjectsRequestObj.setPrefix("queryId");
    Mockito.when(verticaMetadataHandlerMocked.getS3ExportBucket()).thenReturn("testS3Bucket");
    Mockito.when(listObjectsRequest.withBucketName(anyString())).thenReturn(listObjectsRequestObj);
    Mockito.when(listObjectsRequest.withPrefix(anyString())).thenReturn(listObjectsRequestObj);
    Mockito.when(amazonS3.listObjects(any(ListObjectsRequest.class))).thenReturn(objectListing);
    Mockito.when(objectListing.getObjectSummaries()).thenReturn(s3ObjectSummariesList);
    GetSplitsRequest originalReq = new GetSplitsRequest(this.federatedIdentity, "queryId", "catalog_name", new TableName("schema", "table_name"), partitions, partitionCols, new Constraints(constraintsMap), null);
    GetSplitsRequest req = new GetSplitsRequest(originalReq, null);
    logger.info("doGetSplits: req[{}]", req);
    MetadataResponse rawResponse = verticaMetadataHandlerMocked.doGetSplits(allocator, req);
    assertEquals(MetadataRequestType.GET_SPLITS, rawResponse.getRequestType());
    GetSplitsResponse response = (GetSplitsResponse) rawResponse;
    String continuationToken = response.getContinuationToken();
    logger.info("doGetSplits: continuationToken[{}] - splits[{}]", continuationToken, response.getSplits());
    for (Split nextSplit : response.getSplits()) {
        assertNotNull(nextSplit.getProperty("query_id"));
        assertNotNull(nextSplit.getProperty("exportBucket"));
        assertNotNull(nextSplit.getProperty("s3ObjectKey"));
    }
    assertTrue(!response.getSplits().isEmpty());
    logger.info("doGetSplits: exit");
}

Example 79 with ValueSet

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

the class ExampleMetadataHandlerTest method doGetSplits.

@Test
public void doGetSplits() {
    if (!enableTests) {
        // We do this because until you complete the tutorial these tests will fail. When you attempt to publis
        // using ../toos/publish.sh ...  it will set the publishing flag and force these tests. This is how we
        // avoid breaking the build but still have a useful tutorial. We are also duplicateing this block
        // on purpose since this is a somewhat odd pattern.
        logger.info("doGetSplits: Tests are disabled, to enable them set the 'publishing' environment variable " + "using maven clean install -Dpublishing=true");
        return;
    }
    logger.info("doGetSplits: enter");
    String yearCol = "year";
    String monthCol = "month";
    String dayCol = "day";
    // This is the schema that ExampleMetadataHandler has layed out for a 'Partition' so we need to populate this
    // minimal set of info here.
    Schema schema = SchemaBuilder.newBuilder().addIntField(yearCol).addIntField(monthCol).addIntField(dayCol).build();
    List<String> partitionCols = new ArrayList<>();
    partitionCols.add(yearCol);
    partitionCols.add(monthCol);
    partitionCols.add(dayCol);
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    Block partitions = allocator.createBlock(schema);
    int num_partitions = 10;
    for (int i = 0; i < num_partitions; i++) {
        BlockUtils.setValue(partitions.getFieldVector(yearCol), i, 2016 + i);
        BlockUtils.setValue(partitions.getFieldVector(monthCol), i, (i % 12) + 1);
        BlockUtils.setValue(partitions.getFieldVector(dayCol), i, (i % 28) + 1);
    }
    partitions.setRowCount(num_partitions);
    String continuationToken = null;
    GetSplitsRequest originalReq = new GetSplitsRequest(fakeIdentity(), "queryId", "catalog_name", new TableName("schema", "table_name"), partitions, partitionCols, new Constraints(constraintsMap), continuationToken);
    int numContinuations = 0;
    do {
        GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken);
        logger.info("doGetSplits: req[{}]", req);
        MetadataResponse rawResponse = handler.doGetSplits(allocator, req);
        assertEquals(MetadataRequestType.GET_SPLITS, rawResponse.getRequestType());
        GetSplitsResponse response = (GetSplitsResponse) rawResponse;
        continuationToken = response.getContinuationToken();
        logger.info("doGetSplits: continuationToken[{}] - splits[{}]", continuationToken, response.getSplits());
        for (Split nextSplit : response.getSplits()) {
            assertNotNull(nextSplit.getProperty("year"));
            assertNotNull(nextSplit.getProperty("month"));
            assertNotNull(nextSplit.getProperty("day"));
        }
        assertTrue(!response.getSplits().isEmpty());
        if (continuationToken != null) {
            numContinuations++;
        }
    } while (continuationToken != null);
    assertTrue(numContinuations == 0);
    logger.info("doGetSplits: exit");
}

Example 80 with ValueSet

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

the class ExampleRecordHandlerTest method doReadRecordsNoSpill.

@Test
public void doReadRecordsNoSpill() throws Exception {
    if (!enableTests) {
        // We do this because until you complete the tutorial these tests will fail. When you attempt to publis
        // using ../toos/publish.sh ...  it will set the publishing flag and force these tests. This is how we
        // avoid breaking the build but still have a useful tutorial. We are also duplicateing this block
        // on purpose since this is a somewhat odd pattern.
        logger.info("doReadRecordsNoSpill: Tests are disabled, to enable them set the 'publishing' environment variable " + "using maven clean install -Dpublishing=true");
        return;
    }
    for (int i = 0; i < 2; i++) {
        Map<String, ValueSet> constraintsMap = new HashMap<>();
        ReadRecordsRequest request = new ReadRecordsRequest(fakeIdentity(), "catalog", "queryId-" + System.currentTimeMillis(), new TableName("schema", "table"), schemaForRead, Split.newBuilder(makeSpillLocation(), null).add("year", "2017").add("month", "11").add("day", "1").build(), new Constraints(constraintsMap), // 100GB don't expect this to spill
        100_000_000_000L, 100_000_000_000L);
        RecordResponse rawResponse = handler.doReadRecords(allocator, request);
        assertTrue(rawResponse instanceof ReadRecordsResponse);
        ReadRecordsResponse response = (ReadRecordsResponse) rawResponse;
        logger.info("doReadRecordsNoSpill: rows[{}]", response.getRecordCount());
        assertTrue(response.getRecords().getRowCount() > 0);
        logger.info("doReadRecordsNoSpill: {}", BlockUtils.rowToString(response.getRecords(), 0));
    }
}