Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion .github/workflows/flink.yml
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ jobs:
export fmt_SOURCE=BUNDLED
export folly_SOURCE=BUNDLED
git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git
cd velox4j && git reset --hard 97fc1edafebd0f505e613d260f77f92f5252d048
cd velox4j && git reset --hard f112d178576e2010954db01ccae3761aeb200a47
git apply $GITHUB_WORKSPACE/gluten-flink/patches/fix-velox4j.patch
$GITHUB_WORKSPACE/build/mvn clean install -DskipTests -Dgpg.skip -Dspotless.skip=true
cd ..
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
import org.apache.flink.contrib.streaming.state.RocksDBKeyedStateBackend;
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
import org.apache.flink.streaming.api.operators.BoundedOneInput;
import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
Expand All @@ -50,7 +51,10 @@

/** Calculate operator in gluten, which will call Velox to run. */
public class GlutenOneInputOperator<IN, OUT> extends TableStreamOperator<OUT>
implements OneInputStreamOperator<IN, OUT>, GlutenOperator, NativeCallbackTarget {
implements OneInputStreamOperator<IN, OUT>,
BoundedOneInput,
GlutenOperator,
NativeCallbackTarget {

private final StatefulPlanNode glutenPlan;
private final String id;
Expand Down Expand Up @@ -218,24 +222,46 @@ private void drainTaskOutput() {
}
UpIterator.State state = task.advance();
if (state == UpIterator.State.AVAILABLE) {
final StatefulElement statefulElement = task.statefulGet();
try {
if (statefulElement.isWatermark()) {
StatefulWatermark watermark = statefulElement.asWatermark();
output.emitWatermark(new Watermark(watermark.getTimestamp()));
} else {
outputBridge.collect(
output, statefulElement.asRecord(), sessionResource.getAllocator(), outputType);
}
} finally {
statefulElement.close();
}
processAvailableElement();
} else {
break;
}
}
}

private void processAvailableElement() {
final StatefulElement statefulElement = task.statefulGet();
try {
if (statefulElement.isWatermark()) {
StatefulWatermark watermark = statefulElement.asWatermark();
output.emitWatermark(new Watermark(watermark.getTimestamp()));
} else {
outputBridge.collect(
output, statefulElement.asRecord(), sessionResource.getAllocator(), outputType);
}
} finally {
statefulElement.close();
}
}

private void finishTask() {
while (true) {
UpIterator.State state = task.advance();
switch (state) {
case AVAILABLE:
processAvailableElement();
break;
case BLOCKED:
task.waitFor();
break;
case FINISHED:
return;
default:
throw new IllegalStateException("Unknown Velox task state: " + state);
}
}
}

public <NIN, NOUT> GlutenOneInputOperator<NIN, NOUT> cloneWithInputOutputClasses(
StatefulPlanNode plan, Class<NIN> newInClass, Class<NOUT> newOutClass) {
return new GlutenOneInputOperator<>(
Expand All @@ -259,6 +285,16 @@ public void processWatermark2(Watermark mark) throws Exception {
throw new UnsupportedOperationException("Not implemented for GlutenOneInputOperator");
}

@Override
public void endInput() throws Exception {
if (inputQueue != null) {
inputQueue.noMoreInput();
}
if (task != null) {
finishTask();
}
}

@Override
public void close() throws Exception {
closing = true;
Expand All @@ -268,6 +304,7 @@ public void close() throws Exception {
task.unbindNativeCallbackTarget();
}
},
() -> {},
() -> {
if (task != null) {
task.close();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import org.apache.flink.runtime.state.StateInitializationContext;
import org.apache.flink.runtime.state.StateSnapshotContext;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.streaming.api.operators.BoundedMultiInput;
import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
Expand All @@ -51,7 +52,10 @@
* instead of flink RowData.
*/
public class GlutenTwoInputOperator<IN, OUT> extends AbstractStreamOperator<OUT>
implements TwoInputStreamOperator<IN, IN, OUT>, GlutenOperator, NativeCallbackTarget {
implements TwoInputStreamOperator<IN, IN, OUT>,
BoundedMultiInput,
GlutenOperator,
NativeCallbackTarget {

private static final Logger LOG = LoggerFactory.getLogger(GlutenTwoInputOperator.class);

Expand All @@ -75,6 +79,8 @@ public class GlutenTwoInputOperator<IN, OUT> extends AbstractStreamOperator<OUT>
private VectorOutputBridge<OUT> outputBridge;
private String description;
private final GlutenMailboxHolder mailboxHolder = new GlutenMailboxHolder();
private boolean leftInputEnded;
private boolean rightInputEnded;

public GlutenTwoInputOperator(
StatefulPlanNode plan,
Expand Down Expand Up @@ -198,24 +204,46 @@ private void drainTaskOutput() {
}
UpIterator.State state = task.advance();
if (state == UpIterator.State.AVAILABLE) {
final StatefulElement element = task.statefulGet();
try {
if (element.isWatermark()) {
StatefulWatermark watermark = element.asWatermark();
output.emitWatermark(new Watermark(watermark.getTimestamp()));
} else {
outputBridge.collect(
output, element.asRecord(), sessionResource.getAllocator(), outputType);
}
} finally {
element.close();
}
processAvailableElement();
} else {
break;
}
}
}

private void processAvailableElement() {
final StatefulElement element = task.statefulGet();
try {
if (element.isWatermark()) {
StatefulWatermark watermark = element.asWatermark();
output.emitWatermark(new Watermark(watermark.getTimestamp()));
} else {
outputBridge.collect(
output, element.asRecord(), sessionResource.getAllocator(), outputType);
}
} finally {
element.close();
}
}

private void finishTask() {
while (true) {
UpIterator.State state = task.advance();
switch (state) {
case AVAILABLE:
processAvailableElement();
break;
case BLOCKED:
task.waitFor();
break;
case FINISHED:
return;
default:
throw new IllegalStateException("Unknown Velox task state: " + state);
}
}
}

@Override
public void processWatermark(Watermark mark) throws Exception {
task.notifyWatermark(mark.getTimestamp());
Expand All @@ -234,30 +262,64 @@ public void processWatermark2(Watermark mark) throws Exception {
processElementInternal();
}

@Override
public void endInput(int inputId) throws Exception {
switch (inputId) {
case 1:
leftInputEnded = true;
if (leftInputQueue != null) {
leftInputQueue.noMoreInput();
}
break;
case 2:
rightInputEnded = true;
if (rightInputQueue != null) {
rightInputQueue.noMoreInput();
}
break;
default:
throw new IllegalArgumentException("Unknown input id: " + inputId);
}
if (leftInputEnded && rightInputEnded && task != null) {
finishTask();
}
}

@Override
public void close() throws Exception {
closing = true;
GlutenCloseables.runWithCleanup(
() -> {
if (leftInputQueue != null) {
leftInputQueue.close();
leftInputQueue.noMoreInput();
}
},
() -> {
if (rightInputQueue != null) {
rightInputQueue.close();
rightInputQueue.noMoreInput();
}
},
() -> {
if (task != null) {
task.unbindNativeCallbackTarget();
}
},
() -> {},
() -> {
if (task != null) {
task.close();
}
},
() -> {
if (leftInputQueue != null) {
leftInputQueue.close();
}
},
() -> {
if (rightInputQueue != null) {
rightInputQueue.close();
}
},
() -> {
if (sessionResource != null) {
sessionResource.close();
Expand Down
30 changes: 30 additions & 0 deletions gluten-flink/ut/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,36 @@
<version>${flink.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-orc</artifactId>
<version>${flink.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.orc</groupId>
<artifactId>orc-core</artifactId>
<version>1.9.2</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.hive</groupId>
<artifactId>hive-storage-api</artifactId>
<version>2.8.1</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java</artifactId>
<version>3.25.5</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.apache.hadoop</groupId>
<artifactId>hadoop-client</artifactId>
<version>2.7.4</version>
<scope>test</scope>
</dependency>
</dependencies>

<build>
Expand Down
Loading
Loading