use of com.amazonaws.athena.connector.lambda.domain.predicate.Range in project aws-athena-query-federation by awslabs.
the class BigQuerySqlUtils method toPredicate.
private static String toPredicate(String columnName, ValueSet valueSet, ArrowType type, List<QueryParameterValue> parameterValues) {
List<String> disjuncts = new ArrayList<>();
List<Object> singleValues = new ArrayList<>();
if (valueSet instanceof SortedRangeSet) {
if (valueSet.isNone() && valueSet.isNullAllowed()) {
return String.format("(%s IS NULL)", columnName);
}
if (valueSet.isNullAllowed()) {
disjuncts.add(String.format("(%s IS NULL)", columnName));
}
Range rangeSpan = ((SortedRangeSet) valueSet).getSpan();
if (!valueSet.isNullAllowed() && rangeSpan.getLow().isLowerUnbounded() && rangeSpan.getHigh().isUpperUnbounded()) {
return String.format("(%s IS NOT NULL)", columnName);
}
for (Range range : valueSet.getRanges().getOrderedRanges()) {
if (range.isSingleValue()) {
singleValues.add(range.getLow().getValue());
} else {
List<String> rangeConjuncts = new ArrayList<>();
if (!range.getLow().isLowerUnbounded()) {
switch(range.getLow().getBound()) {
case ABOVE:
rangeConjuncts.add(toPredicate(columnName, ">", range.getLow().getValue(), type, parameterValues));
break;
case EXACTLY:
rangeConjuncts.add(toPredicate(columnName, ">=", range.getLow().getValue(), type, parameterValues));
break;
case BELOW:
throw new IllegalArgumentException("Low marker should never use BELOW bound");
default:
throw new AssertionError("Unhandled bound: " + range.getLow().getBound());
}
}
if (!range.getHigh().isUpperUnbounded()) {
switch(range.getHigh().getBound()) {
case ABOVE:
throw new IllegalArgumentException("High marker should never use ABOVE bound");
case EXACTLY:
rangeConjuncts.add(toPredicate(columnName, "<=", range.getHigh().getValue(), type, parameterValues));
break;
case BELOW:
rangeConjuncts.add(toPredicate(columnName, "<", range.getHigh().getValue(), type, parameterValues));
break;
default:
throw new AssertionError("Unhandled bound: " + range.getHigh().getBound());
}
}
// If rangeConjuncts is null, then the range was ALL, which should already have been checked for
Preconditions.checkState(!rangeConjuncts.isEmpty());
disjuncts.add("(" + Joiner.on(" AND ").join(rangeConjuncts) + ")");
}
}
// Add back all of the possible single values either as an equality or an IN predicate
if (singleValues.size() == 1) {
disjuncts.add(toPredicate(columnName, "=", Iterables.getOnlyElement(singleValues), type, parameterValues));
} else if (singleValues.size() > 1) {
for (Object value : singleValues) {
parameterValues.add(getValueForWhereClause(columnName, value, type));
}
String values = Joiner.on(",").join(Collections.nCopies(singleValues.size(), "?"));
disjuncts.add(quote(columnName) + " IN (" + values + ")");
}
}
return "(" + Joiner.on(" OR ").join(disjuncts) + ")";
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Range 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);
}
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Range in project aws-athena-query-federation by awslabs.
the class PostGreSqlRecordHandlerTest method getRangeSet.
private ValueSet getRangeSet(Marker.Bound lowerBound, Object lowerValue, Marker.Bound upperBound, Object upperValue) {
Range range = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(range.isSingleValue()).thenReturn(false);
Mockito.when(range.getLow().getBound()).thenReturn(lowerBound);
Mockito.when(range.getLow().getValue()).thenReturn(lowerValue);
Mockito.when(range.getHigh().getBound()).thenReturn(upperBound);
Mockito.when(range.getHigh().getValue()).thenReturn(upperValue);
ValueSet valueSet = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(valueSet.getRanges().getOrderedRanges()).thenReturn(Collections.singletonList(range));
return valueSet;
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Range in project aws-athena-query-federation by awslabs.
the class PostGreSqlRecordHandlerTest method buildSplitSqlTest.
@Test
public void buildSplitSqlTest() throws SQLException {
logger.info("buildSplitSqlTest - enter");
TableName tableName = new TableName("testSchema", "testTable");
SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
schemaBuilder.addField(FieldBuilder.newBuilder("testCol1", Types.MinorType.INT.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol2", Types.MinorType.VARCHAR.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol3", Types.MinorType.BIGINT.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol4", Types.MinorType.FLOAT4.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol5", Types.MinorType.SMALLINT.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol6", Types.MinorType.TINYINT.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol7", Types.MinorType.FLOAT8.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol8", Types.MinorType.BIT.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("testCol9", new ArrowType.Decimal(8, 2)).build());
schemaBuilder.addField(FieldBuilder.newBuilder("partition_schema_name", Types.MinorType.VARCHAR.getType()).build());
schemaBuilder.addField(FieldBuilder.newBuilder("partition_name", Types.MinorType.VARCHAR.getType()).build());
Schema schema = schemaBuilder.build();
Split split = Mockito.mock(Split.class);
Mockito.when(split.getProperties()).thenReturn(ImmutableMap.of("partition_schema_name", "s0", "partition_name", "p0"));
Mockito.when(split.getProperty(Mockito.eq(com.amazonaws.athena.connectors.postgresql.PostGreSqlMetadataHandler.BLOCK_PARTITION_SCHEMA_COLUMN_NAME))).thenReturn("s0");
Mockito.when(split.getProperty(Mockito.eq(com.amazonaws.athena.connectors.postgresql.PostGreSqlMetadataHandler.BLOCK_PARTITION_COLUMN_NAME))).thenReturn("p0");
Range range1a = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(range1a.isSingleValue()).thenReturn(true);
Mockito.when(range1a.getLow().getValue()).thenReturn(1);
Range range1b = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(range1b.isSingleValue()).thenReturn(true);
Mockito.when(range1b.getLow().getValue()).thenReturn(2);
ValueSet valueSet1 = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(valueSet1.getRanges().getOrderedRanges()).thenReturn(ImmutableList.of(range1a, range1b));
ValueSet valueSet2 = getRangeSet(Marker.Bound.EXACTLY, "1", Marker.Bound.BELOW, "10");
ValueSet valueSet3 = getRangeSet(Marker.Bound.ABOVE, 2L, Marker.Bound.EXACTLY, 20L);
ValueSet valueSet4 = getSingleValueSet(1.1F);
ValueSet valueSet5 = getSingleValueSet(1);
ValueSet valueSet6 = getSingleValueSet(0);
ValueSet valueSet7 = getSingleValueSet(1.2d);
ValueSet valueSet8 = getSingleValueSet(true);
ValueSet valueSet9 = getSingleValueSet(BigDecimal.valueOf(12.34));
Constraints constraints = Mockito.mock(Constraints.class);
Mockito.when(constraints.getSummary()).thenReturn(new ImmutableMap.Builder<String, ValueSet>().put("testCol1", valueSet1).put("testCol2", valueSet2).put("testCol3", valueSet3).put("testCol4", valueSet4).put("testCol5", valueSet5).put("testCol6", valueSet6).put("testCol7", valueSet7).put("testCol8", valueSet8).put("testCol9", valueSet9).build());
String expectedSql = "SELECT \"testCol1\", \"testCol2\", \"testCol3\", \"testCol4\", \"testCol5\", \"testCol6\", \"testCol7\", \"testCol8\", \"testCol9\" FROM \"s0\".\"p0\" WHERE (\"testCol1\" IN (?,?)) AND ((\"testCol2\" >= ? AND \"testCol2\" < ?)) AND ((\"testCol3\" > ? AND \"testCol3\" <= ?)) AND (\"testCol4\" = ?) AND (\"testCol5\" = ?) AND (\"testCol6\" = ?) AND (\"testCol7\" = ?) AND (\"testCol8\" = ?) AND (\"testCol9\" = ?)";
PreparedStatement expectedPreparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(Mockito.eq(expectedSql))).thenReturn(expectedPreparedStatement);
PreparedStatement preparedStatement = this.postGreSqlRecordHandler.buildSplitSql(this.connection, "testCatalogName", tableName, schema, constraints, split);
Assert.assertEquals(expectedPreparedStatement, preparedStatement);
Mockito.verify(preparedStatement, Mockito.times(1)).setInt(1, 1);
Mockito.verify(preparedStatement, Mockito.times(1)).setInt(2, 2);
Mockito.verify(preparedStatement, Mockito.times(1)).setString(3, "1");
Mockito.verify(preparedStatement, Mockito.times(1)).setString(4, "10");
Mockito.verify(preparedStatement, Mockito.times(1)).setLong(5, 2L);
Mockito.verify(preparedStatement, Mockito.times(1)).setLong(6, 20L);
Mockito.verify(preparedStatement, Mockito.times(1)).setFloat(7, 1.1F);
Mockito.verify(preparedStatement, Mockito.times(1)).setShort(8, (short) 1);
Mockito.verify(preparedStatement, Mockito.times(1)).setByte(9, (byte) 0);
Mockito.verify(preparedStatement, Mockito.times(1)).setDouble(10, 1.2d);
Mockito.verify(preparedStatement, Mockito.times(1)).setBoolean(11, true);
Mockito.verify(preparedStatement, Mockito.times(1)).setBigDecimal(12, BigDecimal.valueOf(12.34));
logger.info("buildSplitSqlTest - exit");
}
use of com.amazonaws.athena.connector.lambda.domain.predicate.Range in project aws-athena-query-federation by awslabs.
the class PostGreSqlRecordHandlerTest method getSingleValueSet.
private ValueSet getSingleValueSet(Object value) {
Range range = Mockito.mock(Range.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(range.isSingleValue()).thenReturn(true);
Mockito.when(range.getLow().getValue()).thenReturn(value);
ValueSet valueSet = Mockito.mock(SortedRangeSet.class, Mockito.RETURNS_DEEP_STUBS);
Mockito.when(valueSet.getRanges().getOrderedRanges()).thenReturn(Collections.singletonList(range));
return valueSet;
}
Aggregations