juarezr / spark-streaming-google-pubsub / 57

Committed 01 Oct 2026 07:02AM UTC coverage: 66.187% (-2.3%) from 68.457%

Build # 57

Build Type

push

github

Committed by

web-flow

Commit Message

fix: update AutoBatchTime logic for receiveTime estimation with write-first approach (#20)

- fix: update AutoBatchTime logic for receiveTime estimation with
write-first approach
- feat: add processingTime configuration to PubSubConfig to hint batch
interval on AutoBatchTime
- fix: remove unneded cap on writeBudget for AutoBatchTime

Coverage Stats

689 of 1169 branches covered (58.94%)

Branch coverage included in aggregate %.

232 of 354 new or added lines in 12 files covered. (65.54%)

15 existing lines in 1 file now uncovered.

1517 of 2164 relevant lines covered (70.1%)

0.7 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

57.58

/src/main/java/io/github/juarezr/spark/pubsub/structured/PubSubPartitionReader.java

package io.github.juarezr.spark.pubsub.structured;

import com.fasterxml.jackson.databind.JsonNode;
import java.io.IOException;
import java.util.Iterator;
import java.util.Map;
import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
import org.apache.spark.sql.catalyst.util.ArrayBasedMapData;
import org.apache.spark.sql.catalyst.util.ArrayData;
import org.apache.spark.sql.catalyst.util.MapData;
import org.apache.spark.sql.connector.read.PartitionReader;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.apache.spark.unsafe.types.UTF8String;

/** Reads {@link PulledMessage} values into Spark {@link InternalRow}s. */
final class PubSubPartitionReader implements PartitionReader<InternalRow> {
  private final Iterator<PulledMessage> iterator;
  private final StructType readSchema;
  private PulledMessage current;
  private JsonNode payloadJson;

  PubSubPartitionReader(PubSubInputPartition partition, StructType readSchema) {
    this.iterator = partition.messages().iterator();
    this.readSchema = readSchema;
  }

  @Override
  public String toString() {
    return "pubsub@" + current == null ? "0" : current.toString();
  }

  @Override
  public boolean next() {
    if (iterator.hasNext()) {
      current = iterator.next();
      payloadJson = null;
      return true;
    }
    return false;
  }

  @Override
  public InternalRow get() {
    GenericInternalRow row = new GenericInternalRow(readSchema.length());
    StructField[] fields = readSchema.fields();
    for (int i = 0; i < fields.length; i++) {
      row.update(i, valueFor(fields[i]));
    }
    return row;
  }

  private Object valueFor(StructField field) {
    String key = PubSubSchema.normalize(field.name());
    switch (key) {
      case PubSubSchema.MESSAGEID:
        return UTF8String.fromString(nullToEmpty(current.messageId()));
      case PubSubSchema.BODY:
      case PubSubSchema.DATA:
        return current.data();
      case PubSubSchema.ATTRIBUTES:
        return toMapData(current.attributes());
      case PubSubSchema.PUBLISHTIME:
        return current.publishTimeMillis() * 1000L;
      case PubSubSchema.ORDERINGKEY:
        return UTF8String.fromString(nullToEmpty(current.orderingKey()));
      case PubSubSchema.ACKID:
        return UTF8String.fromString(nullToEmpty(current.ackId()));
      default:
        return payloadValue(field.name(), field.dataType());
    }
  }

  private Object payloadValue(String name, DataType dataType) {
    if (payloadJson == null) {
      payloadJson = PubSubJsonRowDecoder.parse(current.data());
    }
    return PubSubJsonRowDecoder.value(payloadJson, name, dataType);
  }

  private static String nullToEmpty(String value) {
    return value == null ? "" : value;
  }

  private static MapData toMapData(Map<String, String> attributes) {
    if (attributes == null || attributes.isEmpty()) {
      return new ArrayBasedMapData(
          ArrayData.toArrayData(new Object[0]), ArrayData.toArrayData(new Object[0]));
    }
    Object[] keys = new Object[attributes.size()];
    Object[] values = new Object[attributes.size()];
    int i = 0;
    for (Map.Entry<String, String> e : attributes.entrySet()) {
      keys[i] = UTF8String.fromString(e.getKey());
      values[i] = UTF8String.fromString(e.getValue() == null ? "" : e.getValue());
      i++;
    }
    return new ArrayBasedMapData(ArrayData.toArrayData(keys), ArrayData.toArrayData(values));
  }

  @Override
  public void close() throws IOException {
    // no-op
  }
}

1	package io.github.juarezr.spark.pubsub.structured;
2
3	import com.fasterxml.jackson.databind.JsonNode;
4	import java.io.IOException;
5	import java.util.Iterator;
6	import java.util.Map;
7	import org.apache.spark.sql.catalyst.InternalRow;
8	import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
9	import org.apache.spark.sql.catalyst.util.ArrayBasedMapData;
10	import org.apache.spark.sql.catalyst.util.ArrayData;
11	import org.apache.spark.sql.catalyst.util.MapData;
12	import org.apache.spark.sql.connector.read.PartitionReader;
13	import org.apache.spark.sql.types.DataType;
14	import org.apache.spark.sql.types.StructField;
15	import org.apache.spark.sql.types.StructType;
16	import org.apache.spark.unsafe.types.UTF8String;
17
18	/** Reads {@link PulledMessage} values into Spark {@link InternalRow}s. */
19	final class PubSubPartitionReader implements PartitionReader<InternalRow> {
20	private final Iterator<PulledMessage> iterator;
21	private final StructType readSchema;
22	private PulledMessage current;
23	private JsonNode payloadJson;
24
25	PubSubPartitionReader(PubSubInputPartition partition, StructType readSchema) {	1 ✔
26	this.iterator = partition.messages().iterator();	1 ✔
27	this.readSchema = readSchema;	1 ✔
28	}	1 ✔
29
30	@Override
31	public String toString() {
NEW 32	return "pubsub@" + current == null ? "0" : current.toString();	×
33	}
34
35	@Override
36	public boolean next() {
37	if (iterator.hasNext()) {	1 ✔
38	current = iterator.next();	1 ✔
39	payloadJson = null;	1 ✔
40	return true;	1 ✔
41	}
42	return false;	1 ✔
43	}
44
45	@Override
46	public InternalRow get() {
47	GenericInternalRow row = new GenericInternalRow(readSchema.length());	1 ✔
48	StructField[] fields = readSchema.fields();	1 ✔
49	for (int i = 0; i < fields.length; i++) {	1 ✔
50	row.update(i, valueFor(fields[i]));	1 ✔
51	}
52	return row;	1 ✔
53	}
54
55	private Object valueFor(StructField field) {
56	String key = PubSubSchema.normalize(field.name());	1 ✔
57	switch (key) {	1 !
58	case PubSubSchema.MESSAGEID:
59	return UTF8String.fromString(nullToEmpty(current.messageId()));	1 ✔
60	case PubSubSchema.BODY:
61	case PubSubSchema.DATA:
62	return current.data();	1 ✔
63	case PubSubSchema.ATTRIBUTES:
64	return toMapData(current.attributes());	×
65	case PubSubSchema.PUBLISHTIME:
66	return current.publishTimeMillis() * 1000L;	1 ✔
67	case PubSubSchema.ORDERINGKEY:
68	return UTF8String.fromString(nullToEmpty(current.orderingKey()));	×
69	case PubSubSchema.ACKID:
70	return UTF8String.fromString(nullToEmpty(current.ackId()));	1 ✔
71	default:
72	return payloadValue(field.name(), field.dataType());	1 ✔
73	}
74	}
75
76	private Object payloadValue(String name, DataType dataType) {
77	if (payloadJson == null) {	1 ✔
78	payloadJson = PubSubJsonRowDecoder.parse(current.data());	1 ✔
79	}
80	return PubSubJsonRowDecoder.value(payloadJson, name, dataType);	1 ✔
81	}
82
83	private static String nullToEmpty(String value) {
84	return value == null ? "" : value;	1 !
85	}
86
87	private static MapData toMapData(Map<String, String> attributes) {
88	if (attributes == null \|\| attributes.isEmpty()) {	×
89	return new ArrayBasedMapData(	×
90	ArrayData.toArrayData(new Object[0]), ArrayData.toArrayData(new Object[0]));	×
91	}
92	Object[] keys = new Object[attributes.size()];	×
93	Object[] values = new Object[attributes.size()];	×
94	int i = 0;	×
95	for (Map.Entry<String, String> e : attributes.entrySet()) {	×
96	keys[i] = UTF8String.fromString(e.getKey());	×
97	values[i] = UTF8String.fromString(e.getValue() == null ? "" : e.getValue());	×
98	i++;	×
99	}	×
100	return new ArrayBasedMapData(ArrayData.toArrayData(keys), ArrayData.toArrayData(values));	×
101	}
102
103	@Override
104	public void close() throws IOException {
105	// no-op
106	}	1 ✔
107	}

juarezr / spark-streaming-google-pubsub / 57

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous