Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix correctness issues in S3 Select pushdown #17563

Merged
merged 3 commits into from
Jul 5, 2023
Merged

Fix correctness issues in S3 Select pushdown #17563

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1f845ab
Update Minio docker image version
alexjo2144 Jun 8, 2023
c8f3db3
Use enhanced switch in IonSqlQueryBuilder
alexjo2144 May 19, 2023
d41a568
Fix correctness issues in S3 Select pushdown
alexjo2144 May 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -332,7 +332,9 @@ jobs:
S3_BUCKET_ENDPOINT: "s3.${{ vars.TRINO_AWS_REGION }}.amazonaws.com"
run: |
if [ "${AWS_ACCESS_KEY_ID}" != "" ]; then
$MAVEN test ${MAVEN_TEST} -pl :trino-hive -P aws-tests
$MAVEN test ${MAVEN_TEST} -pl :trino-hive -P aws-tests \
-Ds3.bucket="${S3_BUCKET}" \
-Ds3.bucket-endpoint="${S3_BUCKET_ENDPOINT}"
fi
- name: Run Hive Azure ABFS Access Key Tests
if: matrix.config != 'config-empty' # Hive 1.x does not support Azure storage
Expand Down
2 changes: 2 additions & 0 deletions plugin/trino-hive/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -532,6 +532,7 @@
<exclude>**/TestHiveGlueMetastore.java</exclude>
<exclude>**/TestHiveS3AndGlueMetastoreTest.java</exclude>
<exclude>**/TestTrinoS3FileSystemAwsS3.java</exclude>
<exclude>**/TestS3SelectQueries.java</exclude>
<exclude>**/TestFullParquetReader.java</exclude>
<exclude>**/Test*FailureRecoveryTest.java</exclude>
</excludes>
Expand Down Expand Up @@ -587,6 +588,7 @@
<include>**/TestHiveGlueMetastore.java</include>
<include>**/TestHiveS3AndGlueMetastoreTest.java</include>
<include>**/TestTrinoS3FileSystemAwsS3.java</include>
<include>**/TestS3SelectQueries.java</include>
</includes>
</configuration>
</plugin>
Expand Down
73 changes: 42 additions & 31 deletions plugin/trino-hive/src/main/java/io/trino/plugin/hive/s3select/IonSqlQueryBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import com.google.common.primitives.SignedBytes;
import io.airlift.slice.Slice;
import io.trino.plugin.hive.HiveColumnHandle;
import io.trino.spi.TrinoException;
import io.trino.spi.predicate.Domain;
import io.trino.spi.predicate.Range;
import io.trino.spi.predicate.TupleDomain;
Expand All @@ -33,11 +32,12 @@

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.Iterables.getOnlyElement;
import static io.trino.plugin.hive.HiveErrorCode.HIVE_UNSUPPORTED_FORMAT;
import static io.trino.plugin.hive.s3select.S3SelectDataType.CSV;
import static io.trino.spi.type.BigintType.BIGINT;
import static io.trino.spi.type.BooleanType.BOOLEAN;
import static io.trino.spi.type.DateType.DATE;
Expand All @@ -62,11 +62,27 @@ public class IonSqlQueryBuilder
private static final String DATA_SOURCE = "S3Object s";
private final TypeManager typeManager;
private final S3SelectDataType s3SelectDataType;
private final String nullPredicate;
private final String notNullPredicate;

public IonSqlQueryBuilder(TypeManager typeManager, S3SelectDataType s3SelectDataType)
public IonSqlQueryBuilder(TypeManager typeManager, S3SelectDataType s3SelectDataType, Optional<String> optionalNullCharacterEncoding)
{
if (optionalNullCharacterEncoding.isPresent()) {
checkArgument(s3SelectDataType == CSV, "Null character encoding should only be provided for CSV data");
}

this.typeManager = requireNonNull(typeManager, "typeManager is null");
this.s3SelectDataType = requireNonNull(s3SelectDataType, "s3SelectDataType is null");

String nullCharacterEncoding = optionalNullCharacterEncoding.orElse("");
this.nullPredicate = switch (s3SelectDataType) {
case JSON -> "IS NULL";
case CSV -> "= '%s'".formatted(nullCharacterEncoding);
};
this.notNullPredicate = switch (s3SelectDataType) {
case JSON -> "IS NOT NULL";
case CSV -> "!= '%s'".formatted(nullCharacterEncoding);
};
}

public String buildSql(List<HiveColumnHandle> columns, TupleDomain<HiveColumnHandle> tupleDomain)
Expand Down Expand Up @@ -105,14 +121,10 @@ public String buildSql(List<HiveColumnHandle> columns, TupleDomain<HiveColumnHan

private String getFullyQualifiedColumnName(HiveColumnHandle column)
{
switch (s3SelectDataType) {
case JSON:
return format("s.%s", column.getBaseColumnName());
case CSV:
return format("s._%d", column.getBaseHiveColumnIndex() + 1);
default:
throw new TrinoException(HIVE_UNSUPPORTED_FORMAT, "Attempted to build SQL for unknown S3SelectDataType");
}
return switch (s3SelectDataType) {
case JSON -> "s.%s".formatted(column.getBaseColumnName());
case CSV -> "s._%d".formatted(column.getBaseHiveColumnIndex() + 1);
};
}

private List<String> toConjuncts(List<HiveColumnHandle> columns, TupleDomain<HiveColumnHandle> tupleDomain)
Expand Down Expand Up @@ -149,7 +161,7 @@ private String toPredicate(Domain domain, Type type, HiveColumnHandle column)

if (domain.getValues().isNone()) {
if (domain.isNullAllowed()) {
return getFullyQualifiedColumnName(column) + " = '' ";
return getFullyQualifiedColumnName(column) + " " + nullPredicate;
}
return "FALSE";
}
Expand All @@ -158,7 +170,7 @@ private String toPredicate(Domain domain, Type type, HiveColumnHandle column)
if (domain.isNullAllowed()) {
return "TRUE";
}
return getFullyQualifiedColumnName(column) + " <> '' ";
return getFullyQualifiedColumnName(column) + " " + notNullPredicate;
}

List<String> disjuncts = new ArrayList<>();
Expand All @@ -178,26 +190,35 @@ private String toPredicate(Domain domain, Type type, HiveColumnHandle column)
}
// If rangeConjuncts is null, then the range was ALL, which should already have been checked for
checkState(!rangeConjuncts.isEmpty());
disjuncts.add("(" + Joiner.on(" AND ").join(rangeConjuncts) + ")");
if (rangeConjuncts.size() == 1) {
disjuncts.add("%s %s AND %s".formatted(getFullyQualifiedColumnName(column), notNullPredicate, getOnlyElement(rangeConjuncts)));
}
else {
disjuncts.add("(%s %s AND %s)".formatted(getFullyQualifiedColumnName(column), notNullPredicate, 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("=", getOnlyElement(singleValues), type, column));
disjuncts.add("%s %s AND %s".formatted(getFullyQualifiedColumnName(column), notNullPredicate, toPredicate("=", getOnlyElement(singleValues), type, column)));
}
else if (singleValues.size() > 1) {
List<String> values = new ArrayList<>();
for (Object value : singleValues) {
checkType(type);
values.add(valueToQuery(type, value));
}
disjuncts.add(createColumn(type, column) + " IN (" + Joiner.on(",").join(values) + ")");
disjuncts.add("%s %s AND %s IN (%s)".formatted(
getFullyQualifiedColumnName(column),
notNullPredicate,
createColumn(type, column),
Joiner.on(",").join(values)));
}

// Add nullability disjuncts
checkState(!disjuncts.isEmpty());
if (domain.isNullAllowed()) {
disjuncts.add(getFullyQualifiedColumnName(column) + " = '' ");
disjuncts.add(getFullyQualifiedColumnName(column) + " " + nullPredicate);
}

return "(" + Joiner.on(" OR ").join(disjuncts) + ")";
Expand Down Expand Up @@ -233,7 +254,8 @@ private static String valueToQuery(Type type, Object value)
return String.valueOf((boolean) value);
}
if (type.equals(DATE)) {
return "`" + FORMATTER.print(DAYS.toMillis((long) value)) + "`";
// CAST('2007-04-05T14:30Z' AS TIMESTAMP)
return "'" + FORMATTER.print(DAYS.toMillis((long) value)) + "'";
}
if (type.equals(VarcharType.VARCHAR)) {
return "'" + ((Slice) value).toStringUtf8() + "'";
Expand All @@ -252,22 +274,11 @@ private String createColumn(Type type, HiveColumnHandle columnHandle)
String column = getFullyQualifiedColumnName(columnHandle);

if (type.equals(BIGINT) || type.equals(INTEGER) || type.equals(SMALLINT) || type.equals(TINYINT)) {
return formatPredicate(column, "INT");
return "CAST(" + column + " AS INT)";
}
if (type.equals(BOOLEAN)) {
return formatPredicate(column, "BOOL");
}
if (type.equals(DATE)) {
return formatPredicate(column, "TIMESTAMP");
}
if (type instanceof DecimalType decimalType) {
return formatPredicate(column, format("DECIMAL(%s,%s)", decimalType.getPrecision(), decimalType.getScale()));
return "CAST(" + column + " AS BOOL)";
}
return column;
}

private String formatPredicate(String column, String type)
{
return format("case %s when '' then null else CAST(%s AS %s) end", column, column, type);
alexjo2144 marked this conversation as resolved.
Show resolved Hide resolved
}
}
7 changes: 6 additions & 1 deletion .../trino-hive/src/main/java/io/trino/plugin/hive/s3select/S3SelectRecordCursorProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import io.trino.plugin.hive.HiveColumnHandle;
import io.trino.plugin.hive.HiveRecordCursorProvider;
import io.trino.plugin.hive.ReaderColumns;
import io.trino.plugin.hive.s3select.csv.S3SelectCsvRecordReader;
import io.trino.plugin.hive.type.TypeInfo;
import io.trino.spi.TrinoException;
import io.trino.spi.connector.ConnectorSession;
Expand Down Expand Up @@ -106,7 +107,11 @@ public Optional<ReaderRecordCursorWithProjections> createRecordCursor(
if (s3SelectDataTypeOptional.isPresent()) {
S3SelectDataType s3SelectDataType = s3SelectDataTypeOptional.get();

IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(typeManager, s3SelectDataType);
Optional<String> nullCharacterEncoding = Optional.empty();
if (s3SelectDataType == S3SelectDataType.CSV) {
nullCharacterEncoding = S3SelectCsvRecordReader.nullCharacterEncoding(schema);
}
IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(typeManager, s3SelectDataType, nullCharacterEncoding);
String ionSqlQuery = queryBuilder.buildSql(readerColumns, effectivePredicate);
Optional<S3SelectLineRecordReader> recordReader = S3SelectLineRecordReaderProvider.get(configuration, path, start, length, schema,
ionSqlQuery, s3ClientFactory, s3SelectDataType);
Expand Down
7 changes: 7 additions & 0 deletions ...n/trino-hive/src/main/java/io/trino/plugin/hive/s3select/csv/S3SelectCsvRecordReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,11 @@
import com.amazonaws.services.s3.model.OutputSerialization;
import io.trino.plugin.hive.s3select.S3SelectLineRecordReader;
import io.trino.plugin.hive.s3select.TrinoS3ClientFactory;
import io.trino.plugin.hive.util.SerdeConstants;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;

import java.util.Optional;
import java.util.Properties;

import static io.trino.plugin.hive.util.SerdeConstants.ESCAPE_CHAR;
Expand Down Expand Up @@ -104,4 +106,9 @@ public boolean shouldEnableScanRange()
buildInputSerialization().getCsv().getAllowQuotedRecordDelimiter());
return CompressionType.NONE.equals(getCompressionType()) && !isQuotedRecordDelimiterAllowed;
}

public static Optional<String> nullCharacterEncoding(Properties schema)
{
return Optional.ofNullable(schema.getProperty(SerdeConstants.SERIALIZATION_NULL_FORMAT));
}
}
Loading