|
27 | 27 | import com.google.common.collect.ImmutableMap; |
28 | 28 | import com.google.protobuf.ByteString; |
29 | 29 | 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; |
30 | 36 | import org.junit.Test; |
31 | 37 |
|
32 | 38 | public class PslSparkUtilsTest { |
33 | 39 |
|
34 | 40 | @Test |
35 | 41 | public void testToInternalRow() { |
| 42 | + Timestamp publishTimestamp = Timestamp.newBuilder().setSeconds(20000000L).setNanos(20).build(); |
| 43 | + Timestamp eventTimestamp = Timestamp.newBuilder().setSeconds(10000000L).setNanos(10).build(); |
36 | 44 | Message message = |
37 | 45 | Message.builder() |
38 | 46 | .setKey(ByteString.copyFromUtf8("key")) |
39 | 47 | .setData(ByteString.copyFromUtf8("data")) |
40 | | - .setEventTime(Timestamp.newBuilder().setSeconds(10000000L).setNanos(10).build()) |
| 48 | + .setEventTime(eventTimestamp) |
41 | 49 | .setAttributes( |
42 | 50 | ImmutableListMultimap.of( |
43 | 51 | "key1", ByteString.copyFromUtf8("val1"), |
44 | 52 | "key1", ByteString.copyFromUtf8("val2"), |
45 | 53 | "key2", ByteString.copyFromUtf8("val3"))) |
46 | 54 | .build(); |
47 | 55 | 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)); |
57 | 80 | } |
58 | 81 |
|
59 | 82 | @Test |
|
0 commit comments