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

Example 91 with Constraints

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

the class ConstraintSerializationTest method serializationTest.

@Test
public void serializationTest() throws Exception {
    logger.info("serializationTest - enter");
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col2", SortedRangeSet.copyOf(Types.MinorType.BIGINT.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.BIGINT.getType(), 950L)), false));
    constraintsMap.put("col3", SortedRangeSet.copyOf(Types.MinorType.BIT.getType(), ImmutableList.of(Range.equal(allocator, Types.MinorType.BIT.getType(), false)), false));
    constraintsMap.put("col4", SortedRangeSet.copyOf(Types.MinorType.FLOAT8.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.FLOAT8.getType(), 950.0D)), false));
    constraintsMap.put("col5", SortedRangeSet.copyOf(Types.MinorType.VARCHAR.getType(), ImmutableList.of(Range.equal(allocator, Types.MinorType.VARCHAR.getType(), "8"), Range.equal(allocator, Types.MinorType.VARCHAR.getType(), "9")), false));
    try (GetTableLayoutRequest req = new GetTableLayoutRequest(IdentityUtil.fakeIdentity(), "queryId", "default", new TableName("schema1", "table1"), new Constraints(constraintsMap), SchemaBuilder.newBuilder().build(), new HashSet<>())) {
        ObjectMapperUtil.assertSerialization(req);
    }
    logger.info("serializationTest - exit");
}

Example 92 with Constraints

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

the class GetSplitsRequestSerDeTest method beforeTest.

@Before
public void beforeTest() throws IOException {
    String yearCol = "year";
    String monthCol = "month";
    String dayCol = "day";
    Schema schema = SchemaBuilder.newBuilder().addField(yearCol, new ArrowType.Int(32, true)).addField(monthCol, new ArrowType.Int(32, true)).addField(dayCol, new ArrowType.Int(32, true)).addField("col2", new ArrowType.Utf8()).addField("col3", Types.MinorType.FLOAT8.getType()).addField("col4", Types.MinorType.FLOAT8.getType()).addField("col5", Types.MinorType.FLOAT8.getType()).build();
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col3", SortedRangeSet.copyOf(Types.MinorType.FLOAT8.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.FLOAT8.getType(), -10000D)), false));
    constraintsMap.put("col4", EquatableValueSet.newBuilder(allocator, Types.MinorType.FLOAT8.getType(), false, true).add(1.1D).build());
    constraintsMap.put("col5", new AllOrNoneValueSet(Types.MinorType.FLOAT8.getType(), false, true));
    Constraints constraints = new Constraints(constraintsMap);
    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);
    expected = new GetSplitsRequest(federatedIdentity, "test-query-id", "test-catalog", new TableName("test-schema", "test-table"), partitions, ImmutableList.of(yearCol, monthCol, dayCol), constraints, "test-continuation-token");
    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "GetSplitsRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}

Example 93 with Constraints

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

the class BigQueryRecordHandlerTest method testReadWithConstraint.

@Test
public void testReadWithConstraint() throws Exception {
    try (ReadRecordsRequest request = new ReadRecordsRequest(federatedIdentity, BigQueryTestUtils.PROJECT_1_NAME, "queryId", new TableName("dataset1", "table1"), BigQueryTestUtils.getBlockTestSchema(), Split.newBuilder(S3SpillLocation.newBuilder().withBucket(bucket).withPrefix(prefix).withSplitId(UUID.randomUUID().toString()).withQueryId(UUID.randomUUID().toString()).withIsDirectory(true).build(), keyFactory.create()).build(), new Constraints(Collections.EMPTY_MAP), // This is ignored when directly calling readWithConstraints.
    0, 0)) {
        // This is ignored when directly calling readWithConstraints.
        // Always return try for the evaluator to keep all rows.
        ConstraintEvaluator evaluator = mock(ConstraintEvaluator.class);
        when(evaluator.apply(any(String.class), any(Object.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
            return true;
        });
        // Populate the schema and data that the mocked Google BigQuery client will return.
        com.google.cloud.bigquery.Schema tableSchema = BigQueryTestUtils.getTestSchema();
        List<FieldValueList> tableRows = Arrays.asList(BigQueryTestUtils.getBigQueryFieldValueList(false, 1000, "test1", 123123.12312), BigQueryTestUtils.getBigQueryFieldValueList(true, 500, "test2", 5345234.22111), BigQueryTestUtils.getBigQueryFieldValueList(false, 700, "test3", 324324.23423), BigQueryTestUtils.getBigQueryFieldValueList(true, 900, null, null), BigQueryTestUtils.getBigQueryFieldValueList(null, null, "test5", 2342.234234), BigQueryTestUtils.getBigQueryFieldValueList(true, 1200, "test6", 1123.12312), BigQueryTestUtils.getBigQueryFieldValueList(false, 100, "test7", 1313.12312), BigQueryTestUtils.getBigQueryFieldValueList(true, 120, "test8", 12313.1312), BigQueryTestUtils.getBigQueryFieldValueList(false, 300, "test9", 12323.1312));
        Page<FieldValueList> fieldValueList = new BigQueryPage<>(tableRows);
        TableResult result = new TableResult(tableSchema, tableRows.size(), fieldValueList);
        // Mock out the Google BigQuery Job.
        Job mockBigQueryJob = mock(Job.class);
        when(mockBigQueryJob.isDone()).thenReturn(false).thenReturn(true);
        when(mockBigQueryJob.getQueryResults()).thenReturn(result);
        when(bigQuery.create(any(JobInfo.class))).thenReturn(mockBigQueryJob);
        QueryStatusChecker queryStatusChecker = mock(QueryStatusChecker.class);
        when(queryStatusChecker.isQueryRunning()).thenReturn(true);
        // Execute the test
        bigQueryRecordHandler.readWithConstraint(spillWriter, request, queryStatusChecker);
        PowerMockito.mockStatic(System.class);
        PowerMockito.when(System.getenv(anyString())).thenReturn("test");
        logger.info("Project Name: " + BigQueryUtils.getProjectName(request.getCatalogName()));
        // Ensure that there was a spill so that we can read the spilled block.
        assertTrue(spillWriter.spilled());
    }
}

Example 94 with Constraints

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

the class BigQuerySqlUtilsTest method testSqlWithConstraintsRanges.

@Test
public void testSqlWithConstraintsRanges() throws Exception {
    Map<String, ValueSet> constraintMap = new LinkedHashMap<>();
    ValueSet rangeSet = SortedRangeSet.newBuilder(INT_TYPE, true).add(new Range(Marker.above(new BlockAllocatorImpl(), INT_TYPE, 10), Marker.exactly(new BlockAllocatorImpl(), INT_TYPE, 20))).build();
    ValueSet isNullRangeSet = SortedRangeSet.newBuilder(INT_TYPE, true).build();
    ValueSet isNonNullRangeSet = SortedRangeSet.newBuilder(INT_TYPE, false).add(new Range(Marker.lowerUnbounded(new BlockAllocatorImpl(), INT_TYPE), Marker.upperUnbounded(new BlockAllocatorImpl(), INT_TYPE))).build();
    ValueSet stringRangeSet = SortedRangeSet.newBuilder(STRING_TYPE, false).add(new Range(Marker.exactly(new BlockAllocatorImpl(), STRING_TYPE, "a_low"), Marker.below(new BlockAllocatorImpl(), STRING_TYPE, "z_high"))).build();
    ValueSet booleanRangeSet = SortedRangeSet.newBuilder(BOOLEAN_TYPE, false).add(new Range(Marker.exactly(new BlockAllocatorImpl(), BOOLEAN_TYPE, true), Marker.exactly(new BlockAllocatorImpl(), BOOLEAN_TYPE, true))).build();
    ValueSet integerInRangeSet = SortedRangeSet.newBuilder(INT_TYPE, false).add(new Range(Marker.exactly(new BlockAllocatorImpl(), INT_TYPE, 10), Marker.exactly(new BlockAllocatorImpl(), INT_TYPE, 10))).add(new Range(Marker.exactly(new BlockAllocatorImpl(), INT_TYPE, 1000_000), Marker.exactly(new BlockAllocatorImpl(), INT_TYPE, 1000_000))).build();
    constraintMap.put("integerRange", rangeSet);
    constraintMap.put("isNullRange", isNullRangeSet);
    constraintMap.put("isNotNullRange", isNonNullRangeSet);
    constraintMap.put("stringRange", stringRangeSet);
    constraintMap.put("booleanRange", booleanRangeSet);
    constraintMap.put("integerInRange", integerInRangeSet);
    Mockito.when(split.getProperties()).thenReturn(Collections.emptyMap());
    final List<QueryParameterValue> expectedParameterValues = ImmutableList.of(QueryParameterValue.int64(10), QueryParameterValue.int64(20), QueryParameterValue.string("a_low"), QueryParameterValue.string("z_high"), QueryParameterValue.bool(true), QueryParameterValue.int64(10), QueryParameterValue.int64(1000000));
    try (Constraints constraints = new Constraints(constraintMap)) {
        List<QueryParameterValue> parameterValues = new ArrayList<>();
        String sql = BigQuerySqlUtils.buildSqlFromSplit(tableName, makeSchema(constraintMap), constraints, split, parameterValues);
        assertEquals(expectedParameterValues, parameterValues);
        assertEquals("SELECT `integerRange`,`isNullRange`,`isNotNullRange`,`stringRange`,`booleanRange`,`integerInRange` from `schema`.`table` " + "WHERE ((integerRange IS NULL) OR (`integerRange` > ? AND `integerRange` <= ?)) " + "AND (isNullRange IS NULL) AND (isNotNullRange IS NOT NULL) " + "AND ((`stringRange` >= ? AND `stringRange` < ?)) " + "AND (`booleanRange` = ?) " + "AND (`integerInRange` IN (?,?))", sql);
    }
}

Example 95 with Constraints

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

the class HbaseMetadataHandlerTest method doGetTableLayout.

@Test
public void doGetTableLayout() throws Exception {
    GetTableLayoutRequest req = new GetTableLayoutRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME, new Constraints(new HashMap<>()), SchemaBuilder.newBuilder().build(), Collections.EMPTY_SET);
    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);
    logger.info("doGetTableLayout - {}", res);
    Block partitions = res.getPartitions();
    for (int row = 0; row < partitions.getRowCount() && row < 10; row++) {
        logger.info("doGetTableLayout:{} {}", row, BlockUtils.rowToString(partitions, row));
    }
    assertTrue(partitions.getRowCount() > 0);
}