|
| 1 | +package io.substrait.isthmus; |
| 2 | + |
| 3 | +import java.util.Map; |
| 4 | +import org.apache.calcite.adapter.tpcds.TpcdsSchema; |
| 5 | +import org.apache.calcite.rel.RelRoot; |
| 6 | +import org.apache.calcite.rex.RexFieldAccess; |
| 7 | +import org.apache.calcite.sql.parser.SqlParseException; |
| 8 | +import org.apache.calcite.sql.parser.SqlParser; |
| 9 | +import org.apache.calcite.sql.validate.SqlConformanceEnum; |
| 10 | +import org.junit.jupiter.api.Assertions; |
| 11 | +import org.junit.jupiter.api.Test; |
| 12 | + |
| 13 | +public class ApplyJoinPlanTest { |
| 14 | + |
| 15 | + private static RelRoot getCalcitePlan(SqlToSubstrait s, TpcdsSchema schema, String sql) |
| 16 | + throws SqlParseException { |
| 17 | + var pair = s.registerSchema("tpcds", schema); |
| 18 | + var converter = s.createSqlToRelConverter(pair.left, pair.right); |
| 19 | + SqlParser parser = SqlParser.create(sql, s.parserConfig); |
| 20 | + var root = s.getBestExpRelRoot(converter, parser.parseQuery()); |
| 21 | + return root; |
| 22 | + } |
| 23 | + |
| 24 | + private static void validateOuterRef( |
| 25 | + Map<RexFieldAccess, Integer> fieldAccessDepthMap, String refName, String colName, int depth) { |
| 26 | + var entry = |
| 27 | + fieldAccessDepthMap.entrySet().stream() |
| 28 | + .filter(f -> f.getKey().getReferenceExpr().toString().equals(refName)) |
| 29 | + .filter(f -> f.getKey().getField().getName().equals(colName)) |
| 30 | + .filter(f -> f.getValue() == depth) |
| 31 | + .findFirst(); |
| 32 | + Assertions.assertTrue(entry.isPresent()); |
| 33 | + } |
| 34 | + |
| 35 | + private static Map<RexFieldAccess, Integer> buildOuterFieldRefMap(RelRoot root) { |
| 36 | + final OuterReferenceResolver resolver = new OuterReferenceResolver(); |
| 37 | + var fieldAccessDepthMap = resolver.getFieldAccessDepthMap(); |
| 38 | + Assertions.assertEquals(0, fieldAccessDepthMap.size()); |
| 39 | + resolver.apply(root.rel); |
| 40 | + return fieldAccessDepthMap; |
| 41 | + } |
| 42 | + |
| 43 | + @Test |
| 44 | + public void lateralJoinQuery() throws SqlParseException { |
| 45 | + TpcdsSchema schema = new TpcdsSchema(1.0); |
| 46 | + String sql; |
| 47 | + sql = |
| 48 | + """ |
| 49 | + SELECT ss_sold_date_sk, ss_item_sk, ss_customer_sk |
| 50 | + FROM store_sales CROSS JOIN LATERAL |
| 51 | + (select i_item_sk from item where item.i_item_sk = store_sales.ss_item_sk)"""; |
| 52 | + |
| 53 | + /* the calcite plan for the above query is: |
| 54 | + LogicalProject(SS_SOLD_DATE_SK=[$0], SS_ITEM_SK=[$2], SS_CUSTOMER_SK=[$3]) |
| 55 | + LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{2}]) |
| 56 | + LogicalTableScan(table=[[tpcds, STORE_SALES]]) |
| 57 | + LogicalProject(I_ITEM_SK=[$0]) |
| 58 | + LogicalFilter(condition=[=($0, $cor0.SS_ITEM_SK)]) |
| 59 | + LogicalTableScan(table=[[tpcds, ITEM]]) |
| 60 | + */ |
| 61 | + |
| 62 | + // validate outer reference map |
| 63 | + RelRoot root = getCalcitePlan(new SqlToSubstrait(), schema, sql); |
| 64 | + Map<RexFieldAccess, Integer> fieldAccessDepthMap = buildOuterFieldRefMap(root); |
| 65 | + Assertions.assertEquals(1, fieldAccessDepthMap.size()); |
| 66 | + validateOuterRef(fieldAccessDepthMap, "$cor0", "SS_ITEM_SK", 1); |
| 67 | + |
| 68 | + // TODO validate end to end conversion |
| 69 | + var sE2E = new SqlToSubstrait(); |
| 70 | + Assertions.assertThrows( |
| 71 | + UnsupportedOperationException.class, |
| 72 | + () -> sE2E.execute(sql, "tpcds", schema), |
| 73 | + "Lateral join is not supported"); |
| 74 | + } |
| 75 | + |
| 76 | + @Test |
| 77 | + public void outerApplyQuery() throws SqlParseException { |
| 78 | + TpcdsSchema schema = new TpcdsSchema(1.0); |
| 79 | + String sql; |
| 80 | + sql = |
| 81 | + """ |
| 82 | + SELECT ss_sold_date_sk, ss_item_sk, ss_customer_sk |
| 83 | + FROM store_sales OUTER APPLY |
| 84 | + (select i_item_sk from item where item.i_item_sk = store_sales.ss_item_sk)"""; |
| 85 | + |
| 86 | + FeatureBoard featureBoard = |
| 87 | + ImmutableFeatureBoard.builder() |
| 88 | + .sqlConformanceMode(SqlConformanceEnum.SQL_SERVER_2008) |
| 89 | + .build(); |
| 90 | + SqlToSubstrait s = new SqlToSubstrait(featureBoard); |
| 91 | + RelRoot root = getCalcitePlan(s, schema, sql); |
| 92 | + |
| 93 | + Map<RexFieldAccess, Integer> fieldAccessDepthMap = buildOuterFieldRefMap(root); |
| 94 | + Assertions.assertEquals(1, fieldAccessDepthMap.size()); |
| 95 | + validateOuterRef(fieldAccessDepthMap, "$cor0", "SS_ITEM_SK", 1); |
| 96 | + |
| 97 | + // TODO validate end to end conversion |
| 98 | + Assertions.assertThrows( |
| 99 | + UnsupportedOperationException.class, |
| 100 | + () -> s.execute(sql, "tpcds", schema), |
| 101 | + "APPLY is not supported"); |
| 102 | + } |
| 103 | + |
| 104 | + @Test |
| 105 | + public void nestedApplyJoinQuery() throws SqlParseException { |
| 106 | + TpcdsSchema schema = new TpcdsSchema(1.0); |
| 107 | + String sql; |
| 108 | + sql = |
| 109 | + """ |
| 110 | + SELECT ss_sold_date_sk, ss_item_sk, ss_customer_sk |
| 111 | + FROM store_sales CROSS APPLY |
| 112 | + ( SELECT i_item_sk |
| 113 | + FROM item CROSS APPLY |
| 114 | + ( SELECT p_promo_sk |
| 115 | + FROM promotion |
| 116 | + WHERE p_item_sk = i_item_sk AND p_item_sk = ss_item_sk ) |
| 117 | + WHERE item.i_item_sk = store_sales.ss_item_sk )"""; |
| 118 | + |
| 119 | + /* the calcite plan for the above query is: |
| 120 | + LogicalProject(SS_SOLD_DATE_SK=[$0], SS_ITEM_SK=[$2], SS_CUSTOMER_SK=[$3]) |
| 121 | + LogicalCorrelate(correlation=[$cor2], joinType=[inner], requiredColumns=[{2}]) |
| 122 | + LogicalTableScan(table=[[tpcds, STORE_SALES]]) |
| 123 | + LogicalProject(I_ITEM_SK=[$0]) |
| 124 | + LogicalFilter(condition=[=($0, $cor2.SS_ITEM_SK)]) |
| 125 | + LogicalCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0}]) |
| 126 | + LogicalTableScan(table=[[tpcds, ITEM]]) |
| 127 | + LogicalProject(P_PROMO_SK=[$0]) |
| 128 | + LogicalFilter(condition=[AND(=($4, $cor0.I_ITEM_SK), =($4, $cor2.SS_ITEM_SK))]) |
| 129 | + LogicalTableScan(table=[[tpcds, PROMOTION]]) |
| 130 | + */ |
| 131 | + FeatureBoard featureBoard = |
| 132 | + ImmutableFeatureBoard.builder() |
| 133 | + .sqlConformanceMode(SqlConformanceEnum.SQL_SERVER_2008) |
| 134 | + .build(); |
| 135 | + SqlToSubstrait s = new SqlToSubstrait(featureBoard); |
| 136 | + RelRoot root = getCalcitePlan(s, schema, sql); |
| 137 | + |
| 138 | + Map<RexFieldAccess, Integer> fieldAccessDepthMap = buildOuterFieldRefMap(root); |
| 139 | + Assertions.assertEquals(3, fieldAccessDepthMap.size()); |
| 140 | + validateOuterRef(fieldAccessDepthMap, "$cor2", "SS_ITEM_SK", 1); |
| 141 | + validateOuterRef(fieldAccessDepthMap, "$cor2", "SS_ITEM_SK", 2); |
| 142 | + validateOuterRef(fieldAccessDepthMap, "$cor0", "I_ITEM_SK", 1); |
| 143 | + |
| 144 | + // TODO validate end to end conversion |
| 145 | + Assertions.assertThrows( |
| 146 | + UnsupportedOperationException.class, |
| 147 | + () -> s.execute(sql, "tpcds", schema), |
| 148 | + "APPLY is not supported"); |
| 149 | + } |
| 150 | + |
| 151 | + @Test |
| 152 | + public void crossApplyQuery() throws SqlParseException { |
| 153 | + TpcdsSchema schema = new TpcdsSchema(1.0); |
| 154 | + String sql; |
| 155 | + sql = |
| 156 | + """ |
| 157 | + SELECT ss_sold_date_sk, ss_item_sk, ss_customer_sk |
| 158 | + FROM store_sales CROSS APPLY |
| 159 | + (select i_item_sk from item where item.i_item_sk = store_sales.ss_item_sk)"""; |
| 160 | + |
| 161 | + FeatureBoard featureBoard = |
| 162 | + ImmutableFeatureBoard.builder() |
| 163 | + .sqlConformanceMode(SqlConformanceEnum.SQL_SERVER_2008) |
| 164 | + .build(); |
| 165 | + SqlToSubstrait s = new SqlToSubstrait(featureBoard); |
| 166 | + |
| 167 | + // TODO validate end to end conversion |
| 168 | + Assertions.assertThrows( |
| 169 | + UnsupportedOperationException.class, |
| 170 | + () -> s.execute(sql, "tpcds", schema), |
| 171 | + "APPLY is not supported"); |
| 172 | + } |
| 173 | +} |
0 commit comments