Skip to content

Commit ab6a124

Browse files
feat: Fixes timestamps and more thorough testing (#38)
1 parent 0faf688 commit ab6a124

2 files changed

Lines changed: 35 additions & 12 deletions

File tree

src/main/java/com/google/cloud/pubsublite/spark/PslSparkUtils.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,9 +75,9 @@ public static InternalRow toInternalRow(
7575
msg.offset().value(),
7676
ByteArray.concat(msg.message().key().toByteArray()),
7777
ByteArray.concat(msg.message().data().toByteArray()),
78-
Timestamps.toMillis(msg.publishTime()),
78+
Timestamps.toMicros(msg.publishTime()),
7979
msg.message().eventTime().isPresent()
80-
? Timestamps.toMillis(msg.message().eventTime().get())
80+
? Timestamps.toMicros(msg.message().eventTime().get())
8181
: null,
8282
convertAttributesToSparkMap(msg.message().attributes())));
8383
return InternalRow.apply(asScalaBufferConverter(list).asScala());

src/test/java/com/google/cloud/pubsublite/spark/PslSparkUtilsTest.java

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,33 +27,56 @@
2727
import com.google.common.collect.ImmutableMap;
2828
import com.google.protobuf.ByteString;
2929
import com.google.protobuf.Timestamp;
30+
import com.google.protobuf.util.Timestamps;
31+
import java.nio.charset.StandardCharsets;
32+
import org.apache.spark.sql.catalyst.InternalRow;
33+
import org.apache.spark.sql.catalyst.util.ArrayData;
34+
import org.apache.spark.sql.catalyst.util.GenericArrayData;
35+
import org.apache.spark.sql.types.DataTypes;
3036
import org.junit.Test;
3137

3238
public class PslSparkUtilsTest {
3339

3440
@Test
3541
public void testToInternalRow() {
42+
Timestamp publishTimestamp = Timestamp.newBuilder().setSeconds(20000000L).setNanos(20).build();
43+
Timestamp eventTimestamp = Timestamp.newBuilder().setSeconds(10000000L).setNanos(10).build();
3644
Message message =
3745
Message.builder()
3846
.setKey(ByteString.copyFromUtf8("key"))
3947
.setData(ByteString.copyFromUtf8("data"))
40-
.setEventTime(Timestamp.newBuilder().setSeconds(10000000L).setNanos(10).build())
48+
.setEventTime(eventTimestamp)
4149
.setAttributes(
4250
ImmutableListMultimap.of(
4351
"key1", ByteString.copyFromUtf8("val1"),
4452
"key1", ByteString.copyFromUtf8("val2"),
4553
"key2", ByteString.copyFromUtf8("val3")))
4654
.build();
4755
SequencedMessage sequencedMessage =
48-
SequencedMessage.of(
49-
message,
50-
Timestamp.newBuilder().setSeconds(10000000L).setNanos(10).build(),
51-
Offset.of(10L),
52-
10L);
53-
PslSparkUtils.toInternalRow(
54-
sequencedMessage,
55-
UnitTestExamples.exampleSubscriptionPath(),
56-
UnitTestExamples.examplePartition());
56+
SequencedMessage.of(message, publishTimestamp, UnitTestExamples.exampleOffset(), 10L);
57+
InternalRow row =
58+
PslSparkUtils.toInternalRow(
59+
sequencedMessage,
60+
UnitTestExamples.exampleSubscriptionPath(),
61+
UnitTestExamples.examplePartition());
62+
assertThat(row.getString(0)).isEqualTo(UnitTestExamples.exampleSubscriptionPath().toString());
63+
assertThat(row.getLong(1)).isEqualTo(UnitTestExamples.examplePartition().value());
64+
assertThat(row.getLong(2)).isEqualTo(UnitTestExamples.exampleOffset().value());
65+
assertThat(row.getBinary(3)).isEqualTo("key".getBytes(StandardCharsets.UTF_8));
66+
assertThat(row.getBinary(4)).isEqualTo("data".getBytes(StandardCharsets.UTF_8));
67+
assertThat(row.getLong(5)).isEqualTo(Timestamps.toMicros(publishTimestamp));
68+
assertThat(row.getLong(6)).isEqualTo(Timestamps.toMicros(eventTimestamp));
69+
ArrayData keys = row.getMap(7).keyArray();
70+
ArrayData values = row.getMap(7).valueArray();
71+
assertThat(keys.get(0, DataTypes.StringType).toString()).isEqualTo("key1");
72+
assertThat(keys.get(1, DataTypes.StringType).toString()).isEqualTo("key2");
73+
GenericArrayData valueOfKey1 =
74+
(GenericArrayData) values.get(0, DataTypes.createArrayType(DataTypes.BinaryType));
75+
GenericArrayData valueOfKey2 =
76+
(GenericArrayData) values.get(1, DataTypes.createArrayType(DataTypes.BinaryType));
77+
assertThat(valueOfKey1.getBinary(0)).isEqualTo("val1".getBytes(StandardCharsets.UTF_8));
78+
assertThat(valueOfKey1.getBinary(1)).isEqualTo("val2".getBytes(StandardCharsets.UTF_8));
79+
assertThat(valueOfKey2.getBinary(0)).isEqualTo("val3".getBytes(StandardCharsets.UTF_8));
5780
}
5881

5982
@Test

0 commit comments

Comments
 (0)