Amar3tto commented on code in PR #22584:
URL: https://github.com/apache/beam/pull/22584#discussion_r990165758
##########
sdks/java/io/sparkreceiver/src/main/java/org/apache/beam/sdk/io/sparkreceiver/ReadFromSparkReceiverWithOffsetDoFn.java:
##########
@@ -141,15 +145,42 @@ public boolean hasRecords() {
@Override
public void start(Receiver<V> sparkReceiver) {
this.sparkReceiver = sparkReceiver;
- try {
- new WrappedSupervisor(
- sparkReceiver,
- new SparkConf(),
- objects -> {
- V record = (V) objects[0];
- recordsQueue.offer(record);
+
+ final SerializableFunction<Object[], Void> storeFn =
+ (input) -> {
+ if (input == null) {
return null;
- });
+ }
+
+ if (input[0] instanceof ByteBuffer) {
+ final ByteBuffer byteBuffer = ((ByteBuffer)
input[0]).asReadOnlyBuffer();
+ final byte[] bytes = new byte[byteBuffer.limit()];
+ byteBuffer.get(bytes);
+ final V record = SerializationUtils.deserialize(bytes);
+ recordsQueue.offer(record);
+ } else if (input[0] instanceof Iterator) {
+ final Iterator<V> iterator = (Iterator<V>) input[0];
+ while (iterator.hasNext()) {
+ V record = (V) iterator.next();
+ recordsQueue.offer(record);
+ }
+ } else if (input[0] instanceof ArrayBuffer) {
+ final ArrayBuffer<V> arrayBuffer = (ArrayBuffer<V>) input[0];
+ final Iterator<V> iterator = arrayBuffer.iterator();
+ while (iterator.hasNext()) {
Review Comment:
Added three new tests that cover this logic.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
