Skip to content

Commit 229643f

Browse files
committed
feat: Add SQL support for multi-condition joins
- Enhanced WayangMultiConditionJoinVisitor with SQL implementation - Updated JdbcJoinOperator to generate AND clauses for multi-column joins - Added test coverage for multi-condition join SQL generation - Maintains backward compatibility with single-condition joins
1 parent 8d8930a commit 229643f

3 files changed

Lines changed: 146 additions & 45 deletions

File tree

wayang-api/wayang-api-sql/src/main/java/org/apache/wayang/api/sql/calcite/converter/WayangMultiConditionJoinVisitor.java

Lines changed: 39 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,15 @@
2323
import java.util.stream.Collectors;
2424

2525
import org.apache.calcite.rel.core.Join;
26+
import org.apache.calcite.rel.type.RelDataTypeField;
2627
import org.apache.calcite.rex.RexCall;
2728
import org.apache.calcite.rex.RexInputRef;
2829
import org.apache.calcite.rex.RexNode;
2930

3031
import org.apache.wayang.api.sql.calcite.converter.functions.JoinFlattenResult;
3132
import org.apache.wayang.api.sql.calcite.converter.functions.MultiConditionJoinKeyExtractor;
3233
import org.apache.wayang.api.sql.calcite.rel.WayangJoin;
34+
import org.apache.wayang.api.sql.calcite.rel.WayangTableScan;
3335
import org.apache.wayang.basic.data.Record;
3436
import org.apache.wayang.basic.data.Tuple2;
3537
import org.apache.wayang.basic.operators.JoinOperator;
@@ -48,26 +50,21 @@ public class WayangMultiConditionJoinVisitor extends WayangRelNodeVisitor<Wayang
4850
*
4951
* @param wayangRelConverter
5052
*/
51-
WayangMultiConditionJoinVisitor(final WayangRelConverter wayangRelConverter) {
53+
public WayangMultiConditionJoinVisitor(final WayangRelConverter wayangRelConverter) {
5254
super(wayangRelConverter);
5355
}
5456

5557
@Override
56-
Operator visit(WayangJoin wayangRelNode) {
58+
public Operator visit(WayangJoin wayangRelNode) {
5759
final Operator childOpLeft = wayangRelConverter.convert(wayangRelNode.getInput(0));
5860
final Operator childOpRight = wayangRelConverter.convert(wayangRelNode.getInput(1));
5961
final RexNode condition = ((Join) wayangRelNode).getCondition();
6062
final RexCall call = (RexCall) condition;
6163

62-
//
6364
final List<RexCall> subConditions = call.operands.stream()
6465
.map(RexCall.class::cast)
6566
.collect(Collectors.toList());
6667

67-
// calcite generates the RexInputRef indexes via looking at the union
68-
// field list of the left and right input of a join.
69-
// since the left input is always the first in this joined field list
70-
// we can eagerly get the fields in the left input
7168
final List<RexInputRef> leftTableInputRefs = subConditions.stream()
7269
.map(sub -> sub.getOperands().stream()
7370
.map(RexInputRef.class::cast)
@@ -79,9 +76,6 @@ Operator visit(WayangJoin wayangRelNode) {
7976
.map(RexInputRef::getIndex)
8077
.toArray(Integer[]::new);
8178

82-
// for the right table input refs, the indexes are offset by the amount of rows
83-
// in the left
84-
// input to the join
8579
final List<RexInputRef> rightTableInputRefs = subConditions.stream()
8680
.map(sub -> sub.getOperands().stream()
8781
.map(RexInputRef.class::cast)
@@ -91,36 +85,40 @@ Operator visit(WayangJoin wayangRelNode) {
9185

9286
final Integer[] rightTableKeyIndexes = rightTableInputRefs.stream()
9387
.map(RexInputRef::getIndex)
94-
.map(key -> key - wayangRelNode.getLeft().getRowType().getFieldCount()) // apply offset
88+
.map(key -> key - wayangRelNode.getLeft().getRowType().getFieldCount())
9589
.toArray(Integer[]::new);
9690

97-
/*
98-
final List<RelDataTypeField> leftFields = Arrays.stream(leftTableKeyIndexes)
99-
.map(key -> wayangRelNode.getLeft().getRowType().getFieldList().get(key))
91+
final List<RelDataTypeField> leftFields = leftTableInputRefs.stream()
92+
.map(ref -> wayangRelNode.getLeft().getRowType().getFieldList().get(ref.getIndex()))
10093
.collect(Collectors.toList());
10194

102-
final List<RelDataTypeField> rightFields = Arrays.stream(rightTableKeyIndexes)
103-
.map(key -> wayangRelNode.getRight().getRowType().getFieldList().get(key))
95+
final List<RelDataTypeField> rightFields = rightTableInputRefs.stream()
96+
.map(ref -> wayangRelNode.getRight().getRowType().getFieldList().get(ref.getIndex() - wayangRelNode.getLeft().getRowType().getFieldCount()))
10497
.collect(Collectors.toList());
10598

106-
final String joiningTableName = childOpLeft instanceof WayangTableScan ? childOpLeft.getName() : childOpRight.getName();
107-
*/
108-
109-
// if join is joining the LHS of a join condition "JOIN left ON left = right"
110-
// then we pick the first case, otherwise the 2nd "JOIN right ON left = right"
111-
final JoinOperator<Record, Record, Record> join = this.getJoinOperator(
99+
final String leftTableName = extractTableName(wayangRelNode.getLeft());
100+
final String rightTableName = extractTableName(wayangRelNode.getRight());
101+
102+
final String leftFieldNames = leftFields.stream()
103+
.map(RelDataTypeField::getName)
104+
.collect(Collectors.joining(","));
105+
106+
final String rightFieldNames = rightFields.stream()
107+
.map(RelDataTypeField::getName)
108+
.collect(Collectors.joining(","));
109+
110+
final JoinOperator<Record, Record, Record> join = getJoinOperator(
112111
leftTableKeyIndexes,
113112
rightTableKeyIndexes,
114113
wayangRelNode,
115-
"",
116-
"",
117-
"",
118-
"");
114+
leftTableName,
115+
leftFieldNames,
116+
rightTableName,
117+
rightFieldNames);
119118

120119
childOpLeft.connectTo(0, join, 0);
121120
childOpRight.connectTo(0, join, 1);
122121

123-
// Join returns Tuple2 - map to a Record
124122
final SerializableFunction<Tuple2<Record, Record>, Record> mp = new JoinFlattenResult();
125123

126124
final MapOperator<Tuple2<Record, Record>, Record> mapOperator = new MapOperator<Tuple2<Record, Record>, Record>(
@@ -133,33 +131,34 @@ Operator visit(WayangJoin wayangRelNode) {
133131
return mapOperator;
134132
}
135133

136-
/**
137-
* This method handles the {@link JoinOperator} creation
138-
*
139-
* @param wayangRelNode
140-
* @param leftKeyIndex
141-
* @param rightKeyIndex
142-
* @return
143-
*/
134+
private String extractTableName(org.apache.calcite.rel.RelNode relNode) {
135+
if (relNode instanceof WayangTableScan) {
136+
return ((WayangTableScan) relNode).getTableName();
137+
}
138+
if (relNode.getInputs() != null && !relNode.getInputs().isEmpty()) {
139+
return extractTableName(relNode.getInput(0));
140+
}
141+
return "UNKNOWN";
142+
}
143+
144144
protected JoinOperator<Record, Record, Record> getJoinOperator(final Integer[] leftKeyIndexes,
145145
final Integer[] rightKeyIndexes,
146146
final WayangJoin wayangRelNode, final String leftTableName, final String leftFieldNames,
147147
final String rightTableName, final String rightFieldNames) {
148-
// TODO: needs withSqlImplementation() for sql support
149148

150149
if (wayangRelNode.getInputs().size() != 2)
151150
throw new UnsupportedOperationException("Join had an unexpected amount of inputs, found: "
152151
+ wayangRelNode.getInputs().size() + ", expected: 2");
153152

154153
final TransformationDescriptor<Record, Record> leftProjectionDescriptor = new TransformationDescriptor<Record, Record>(
155154
new MultiConditionJoinKeyExtractor(leftKeyIndexes),
156-
Record.class, Record.class);
157-
// .withSqlImplementation(""," ")
155+
Record.class, Record.class)
156+
.withSqlImplementation(leftTableName, leftFieldNames);
158157

159158
final TransformationDescriptor<Record, Record> rightProjectionDescriptor = new TransformationDescriptor<Record, Record>(
160159
new MultiConditionJoinKeyExtractor(rightKeyIndexes),
161-
Record.class, Record.class);
162-
// .withSqlImplementation(""," ")
160+
Record.class, Record.class)
161+
.withSqlImplementation(rightTableName, rightFieldNames);
163162

164163
final JoinOperator<Record, Record, Record> join = new JoinOperator<>(
165164
leftProjectionDescriptor,

wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/operators/JdbcJoinOperator.java

Lines changed: 29 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,13 +69,37 @@ public String createSqlClause(Connection connection, FunctionCompiler compiler)
6969
final Tuple<String, String> left = this.keyDescriptor0.getSqlImplementation();
7070
final Tuple<String, String> right = this.keyDescriptor1.getSqlImplementation();
7171
final String leftTableName = left.field0;
72-
final String leftKey = left.field1;
72+
final String leftKeys = left.field1;
7373
final String rightTableName = right.field0;
74-
final String rightKey = right.field1;
74+
final String rightKeys = right.field1;
7575

76-
return "JOIN " + rightTableName + " ON " +
77-
rightTableName + "." + rightKey
78-
+ "=" + leftTableName + "." + leftKey;
76+
if (leftKeys.contains(",") && rightKeys.contains(",")) {
77+
final String[] leftColumns = leftKeys.split(",");
78+
final String[] rightColumns = rightKeys.split(",");
79+
80+
if (leftColumns.length != rightColumns.length) {
81+
throw new IllegalStateException(
82+
"Mismatch in join key counts: left has " + leftColumns.length +
83+
" keys, right has " + rightColumns.length + " keys");
84+
}
85+
86+
final StringBuilder joinCondition = new StringBuilder();
87+
for (int i = 0; i < leftColumns.length; i++) {
88+
if (i > 0) {
89+
joinCondition.append(" AND ");
90+
}
91+
joinCondition.append(leftTableName).append(".").append(leftColumns[i].trim())
92+
.append("=")
93+
.append(rightTableName).append(".").append(rightColumns[i].trim());
94+
}
95+
96+
return "JOIN " + rightTableName + " ON " + joinCondition.toString();
97+
} else {
98+
// Backward compatibility
99+
return "JOIN " + rightTableName + " ON " +
100+
rightTableName + "." + rightKeys
101+
+ "=" + leftTableName + "." + leftKeys;
102+
}
79103
}
80104

81105
@Override

wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/operators/JdbcJoinOperatorTest.java

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,4 +123,82 @@ void testWithHsqldb() throws SQLException {
123123
sqlQueryChannelInstance.getSqlQuery()
124124
);
125125
}
126+
127+
@Test
128+
void testMultiConditionJoinWithHsqldb() throws SQLException {
129+
Configuration configuration = new Configuration();
130+
131+
Job job = mock(Job.class);
132+
when(job.getConfiguration()).thenReturn(configuration);
133+
when(job.getCrossPlatformExecutor()).thenReturn(new CrossPlatformExecutor(job, new NoInstrumentationStrategy()));
134+
SqlQueryChannel.Descriptor sqlChannelDescriptor = HsqldbPlatform.getInstance().getSqlQueryChannelDescriptor();
135+
136+
HsqldbPlatform hsqldbPlatform = new HsqldbPlatform();
137+
138+
ExecutionStage sqlStage = mock(ExecutionStage.class);
139+
140+
// Create test data with multiple join keys
141+
try (Connection jdbcConnection = hsqldbPlatform.createDatabaseDescriptor(configuration).createJdbcConnection()) {
142+
final Statement statement = jdbcConnection.createStatement();
143+
statement.execute("CREATE TABLE orders (order_id INT, customer_id INT, product_id INT, quantity INT);");
144+
statement.execute("INSERT INTO orders VALUES (1, 100, 1001, 5);");
145+
statement.execute("INSERT INTO orders VALUES (2, 101, 1002, 3);");
146+
statement.execute("CREATE TABLE shipments (order_id INT, customer_id INT, ship_date VARCHAR(10));");
147+
statement.execute("INSERT INTO shipments VALUES (1, 100, '2024-01-15');");
148+
statement.execute("INSERT INTO shipments VALUES (2, 101, '2024-01-16');");
149+
}
150+
151+
JdbcTableSource tableSourceOrders = new HsqldbTableSource("orders");
152+
JdbcTableSource tableSourceShipments = new HsqldbTableSource("shipments");
153+
154+
ExecutionTask tableSourceOrdersTask = new ExecutionTask(tableSourceOrders);
155+
tableSourceOrdersTask.setOutputChannel(0, new SqlQueryChannel(sqlChannelDescriptor, tableSourceOrders.getOutput(0)));
156+
tableSourceOrdersTask.setStage(sqlStage);
157+
158+
ExecutionTask tableSourceShipmentsTask = new ExecutionTask(tableSourceShipments);
159+
tableSourceShipmentsTask.setOutputChannel(0, new SqlQueryChannel(sqlChannelDescriptor, tableSourceShipments.getOutput(0)));
160+
tableSourceShipmentsTask.setStage(sqlStage);
161+
162+
// Create multi-condition join: JOIN ON orders.order_id = shipments.order_id AND orders.customer_id = shipments.customer_id
163+
final ExecutionOperator joinOperator = new HsqldbJoinOperator<Record>(
164+
new TransformationDescriptor<Record, Record>(
165+
(record) -> new Record(record.getField(0), record.getField(1)),
166+
Record.class,
167+
Record.class
168+
).withSqlImplementation("orders", "order_id,customer_id"),
169+
new TransformationDescriptor<Record, Record>(
170+
(record) -> new Record(record.getField(0), record.getField(1)),
171+
Record.class,
172+
Record.class
173+
).withSqlImplementation("shipments", "order_id,customer_id")
174+
);
175+
176+
ExecutionTask joinTask = new ExecutionTask(joinOperator);
177+
tableSourceOrdersTask.getOutputChannel(0).addConsumer(joinTask, 0);
178+
tableSourceShipmentsTask.getOutputChannel(0).addConsumer(joinTask, 1);
179+
joinTask.setOutputChannel(0, new SqlQueryChannel(sqlChannelDescriptor, joinOperator.getOutput(0)));
180+
joinTask.setStage(sqlStage);
181+
182+
when(sqlStage.getStartTasks()).thenReturn(Collections.singleton(tableSourceOrdersTask));
183+
when(sqlStage.getTerminalTasks()).thenReturn(Collections.singleton(joinTask));
184+
185+
ExecutionStage nextStage = mock(ExecutionStage.class);
186+
187+
SqlToStreamOperator sqlToStreamOperator = new SqlToStreamOperator(HsqldbPlatform.getInstance());
188+
ExecutionTask sqlToStreamTask = new ExecutionTask(sqlToStreamOperator);
189+
joinTask.getOutputChannel(0).addConsumer(sqlToStreamTask, 0);
190+
sqlToStreamTask.setStage(nextStage);
191+
192+
JdbcExecutor executor = new JdbcExecutor(HsqldbPlatform.getInstance(), job);
193+
executor.execute(sqlStage, new DefaultOptimizationContext(job), job.getCrossPlatformExecutor());
194+
195+
SqlQueryChannel.Instance sqlQueryChannelInstance =
196+
(SqlQueryChannel.Instance) job.getCrossPlatformExecutor().getChannelInstance(sqlToStreamTask.getInputChannel(0));
197+
198+
// Verify that multi-condition join generates proper SQL with AND
199+
assertEquals(
200+
"SELECT * FROM orders JOIN shipments ON orders.order_id=shipments.order_id AND orders.customer_id=shipments.customer_id;",
201+
sqlQueryChannelInstance.getSqlQuery()
202+
);
203+
}
126204
}

0 commit comments

Comments
 (0)