Adding change for vertica substrait plan

athena-federation-sdk-tools/src/main/java/com/amazonaws/athena/connector/substrait/SubstraitSqlUtils.java

@@ -0,0 +1,117 @@
+/*-
+ * #%L
+ * Amazon Athena Query Federation SDK Tools
+ * %%
+ * Copyright (C) 2019 - 2025 Amazon Web Services
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package com.amazonaws.athena.connector.substrait;
+
+import io.substrait.extension.SimpleExtension;
+import io.substrait.isthmus.SubstraitToCalcite;
+import io.substrait.plan.ProtoPlanConverter;
+import io.substrait.proto.Plan;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.rel2sql.RelToSqlConverter;
+import org.apache.calcite.sql.SqlDialect;
+import org.apache.calcite.sql.SqlNode;
+import io.substrait.isthmus.TypeConverter;
+import org.apache.calcite.sql.type.SqlTypeFactoryImpl;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+/**
+ * Utility class for working with Calcite's abstract syntax tree representation of Substrait plans.
+ */
+public final class SubstraitSqlUtils
+{
+    private SubstraitSqlUtils()
+    {
+    }
+
+    /**
+     * Deserializes a Substrait plan with schema-aware processing.
+     * Uses CustomSubstraitToCalcite to resolve table and column references.
+     */
+    public static SqlNode deserializeSubstraitPlan(String planString, SqlDialect sqlDialect, String schemaName, String tableName, org.apache.arrow.vector.types.pojo.Schema tableSchema)
+    {
+        try {
+            // Create schema-aware converter for table/column resolution
+            CustomSubstraitToCalcite substraitToCalcite = new CustomSubstraitToCalcite(
+                    SimpleExtension.load(loadDefaults()),
+                    new SqlTypeFactoryImpl(sqlDialect.getTypeSystem()),
+                    TypeConverter.DEFAULT,
+                    tableName,
+                    tableSchema
+            );
+
+            return convertSubstraitPlanToSql(planString, sqlDialect, substraitToCalcite);
+        }
+        catch (Exception e) {
+            throw new RuntimeException("Failed to parse Substrait plan with schema: " + e.getMessage(), e);
+        }
+    }
+
+    /**
+     * Deserializes a Substrait plan with standard processing.
+     * Uses standard SubstraitToCalcite converter.
+     */
+    public static SqlNode deserializeSubstraitPlan(String planString, SqlDialect sqlDialect)
+    {
+        try {
+            // Create standard converter
+            SubstraitToCalcite substraitToCalcite = new SubstraitToCalcite(
+                    SimpleExtension.load(loadDefaults()),
+                    new SqlTypeFactoryImpl(sqlDialect.getTypeSystem())
+            );
+
+            return convertSubstraitPlanToSql(planString, sqlDialect, substraitToCalcite);
+        }
+        catch (Exception e) {
+            throw new RuntimeException("Failed to parse Substrait plan: " + e.getMessage(), e);
+        }
+    }
+
+    private static List<String> loadDefaults() {
+        return Stream.of("boolean", "aggregate_generic", "aggregate_approx", "arithmetic_decimal", "arithmetic", "comparison", "datetime", "logarithmic", "rounding", "rounding_decimal", "string")
+                .map((func) -> String.format("/functions_%s.yaml", func)).collect(Collectors.toList());
+    }
+
+    /**
+     * Common logic for converting Substrait plan to SQL using the provided converter.
+     * Handles the core deserialization steps that are identical for both methods.
+     */
+    private static SqlNode convertSubstraitPlanToSql(String planString, SqlDialect sqlDialect, SubstraitToCalcite substraitToCalcite)
+            throws Exception
+    {
+        // Step 1: Convert protobuf plan to Substrait plan object
+        ProtoPlanConverter protoPlanConverter = new ProtoPlanConverter();
+        byte[] planBytes = Base64.getDecoder().decode(planString);
+        Plan substraitPlan = Plan.parseFrom(planBytes);
+
+        // Step 2: Convert Substrait plan to Calcite RelNode
+        io.substrait.plan.Plan root = protoPlanConverter.from(substraitPlan);
+        RelNode node = substraitToCalcite.convert(root.getRoots().get(0).getInput());
+
+        // Step 3: Convert Calcite RelNode to SQL
+        RelToSqlConverter converter = new RelToSqlConverter(sqlDialect);
+        return converter.visitRoot(node).asStatement();
+    }
+}

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/substrait/CustomSubstraitToCalcite.java

@@ -0,0 +1,88 @@
+/*-
+ * #%L
+ * Amazon Athena Query Federation SDK
+ * %%
+ * Copyright (C) 2019 - 2025 Amazon Web Services
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package com.amazonaws.athena.connector.substrait;
+
+import io.substrait.extension.SimpleExtension;
+import io.substrait.isthmus.SubstraitToCalcite;
+import io.substrait.isthmus.TypeConverter;
+import io.substrait.relation.Rel;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.apache.calcite.jdbc.CalciteSchema;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.schema.impl.AbstractTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+
+public class CustomSubstraitToCalcite extends SubstraitToCalcite {
+
+    private final String tableName;
+    private final Schema tableSchema;
+
+    public CustomSubstraitToCalcite(SimpleExtension.ExtensionCollection extensions,
+                                    RelDataTypeFactory typeFactory,
+                                    TypeConverter typeConverter,
+                                    String tableName,
+                                    Schema tableSchema) {
+        super(extensions, typeFactory, typeConverter);
+        this.tableName = tableName;
+        this.tableSchema = tableSchema;
+    }
+
+    @Override
+    protected CalciteSchema toSchema(Rel rel) {
+        CalciteSchema rootSchema = CalciteSchema.createRootSchema(false);
+
+        rootSchema.add(tableName, new AbstractTable() {
+            @Override
+            public RelDataType getRowType(RelDataTypeFactory factory) {
+                RelDataTypeFactory.Builder builder = factory.builder();
+
+                for (Field field : tableSchema.getFields()) {
+                    SqlTypeName sqlType = mapArrowTypeToSqlType(field.getType());
+                    builder.add(field.getName(), factory.createSqlType(sqlType));
+                }
+
+                return builder.build();
+            }
+        });
+
+        return rootSchema;
+    }
+
+    private SqlTypeName mapArrowTypeToSqlType(org.apache.arrow.vector.types.pojo.ArrowType arrowType) {
+        if (arrowType instanceof org.apache.arrow.vector.types.pojo.ArrowType.Int) {
+            return SqlTypeName.INTEGER;
+        } else if (arrowType instanceof org.apache.arrow.vector.types.pojo.ArrowType.FloatingPoint) {
+            return SqlTypeName.DOUBLE;
+        } else if (arrowType instanceof org.apache.arrow.vector.types.pojo.ArrowType.Bool) {
+            return SqlTypeName.BOOLEAN;
+        } else if (arrowType instanceof org.apache.arrow.vector.types.pojo.ArrowType.Date) {
+            return SqlTypeName.DATE;
+        } else if (arrowType instanceof org.apache.arrow.vector.types.pojo.ArrowType.Timestamp) {
+            return SqlTypeName.TIMESTAMP;
+        } else if (arrowType instanceof org.apache.arrow.vector.types.pojo.ArrowType.Decimal) {
+            return SqlTypeName.DECIMAL;
+        } else {
+            // Default to VARCHAR for string types and unknown types
+            return SqlTypeName.VARCHAR;
+        }
+    }
+}

athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/SubstraitTypeAndValue.java

@@ -0,0 +1,61 @@
+/*-
+ * #%L
+ * athena-jdbc
+ * %%
+ * Copyright (C) 2019 - 2023 Amazon Web Services
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package com.amazonaws.athena.connectors.jdbc.manager;
+
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.commons.lang3.Validate;
+
+public class SubstraitTypeAndValue
+{
+    private final SqlTypeName type;
+    private final Object value;
+    private final String columnName;
+
+    public SubstraitTypeAndValue(SqlTypeName type, Object value, String columnName)
+    {
+        this.type = Validate.notNull(type, "type is null");
+        this.value = Validate.notNull(value, "value is null");
+        this.columnName = Validate.notNull(columnName, "value is null");
+    }
+
+    public SqlTypeName getType()
+    {
+        return type;
+    }
+
+    public Object getValue()
+    {
+        return value;
+    }
+
+    public String getColumnName()
+    {
+        return columnName;
+    }
+
+    @Override
+    public String toString()
+    {
+        return "TypeAndValue{" +
+                "type=" + type +
+                ", value=" + value +
+                '}';
+    }
+}

athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/visitor/FilterRemovalVisitor.java

@@ -0,0 +1,81 @@
+/*-
+ * #%L
+ * athena-jdbc
+ * %%
+ * Copyright (C) 2019 - 2025 Amazon Web Services
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package com.amazonaws.athena.connectors.jdbc.visitor;
+
+import org.apache.calcite.sql.*;
+import org.apache.calcite.sql.parser.SqlParserPos;
+import org.apache.calcite.sql.util.SqlShuttle;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+
+public class FilterRemovalVisitor extends SqlShuttle
+{
+    private final Set<String> targetColumns;
+    private final boolean insideOr;
+
+    public FilterRemovalVisitor(Set<String> targetColumns)
+    {
+        this(targetColumns, false);
+    }
+
+    private FilterRemovalVisitor(Set<String> targetColumns, boolean insideOr)
+    {
+        this.targetColumns = targetColumns;
+        this.insideOr = insideOr;
+    }
+
+    @Override
+    public SqlNode visit(SqlCall call)
+    {
+        SqlKind kind = call.getKind();
+
+        if (isDirectConditionOnTarget(call)) {
+            return SqlLiteral.createBoolean(!insideOr, SqlParserPos.ZERO);
+        }
+
+        if (kind == SqlKind.AND || kind == SqlKind.OR) {
+            List<SqlNode> newOperands = new ArrayList<>();
+            for (SqlNode operand : call.getOperandList()) {
+                SqlNode newOperand = operand.accept(new FilterRemovalVisitor(
+                        targetColumns, kind == SqlKind.OR
+                ));
+
+                newOperands.add(newOperand);
+            }
+            return call.getOperator().createCall(SqlParserPos.ZERO, newOperands);
+        }
+
+        return super.visit(call);
+    }
+
+    private boolean isDirectConditionOnTarget(SqlCall call)
+    {
+        if (call.operandCount() < 1) {
+            return false;
+        }
+        SqlNode first = call.operand(0);
+        return first instanceof SqlIdentifier &&
+                targetColumns.stream()
+                        .anyMatch(targetColumn ->
+                                targetColumn.equalsIgnoreCase(((SqlIdentifier) first).getSimple()));
+    }
+}