Skip to content

Commit 0571f75

Browse files
feat: fix offsets type check
1 parent 54dc6ae commit 0571f75

3 files changed

Lines changed: 9 additions & 15 deletions

File tree

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,7 @@ public PslContinuousInputPartition(
5353
@Override
5454
public InputPartitionReader<InternalRow> createContinuousReader(PartitionOffset offset) {
5555
checkArgument(
56-
SparkPartitionOffset.class.isAssignableFrom(offset.getClass()),
57-
"offset is not assignable to SparkPartitionOffset");
56+
offset instanceof SparkPartitionOffset, "offset is not instance of SparkPartitionOffset");
5857

5958
SparkPartitionOffset sparkPartitionOffset = (SparkPartitionOffset) offset;
6059
PslPartitionOffset pslPartitionOffset =

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

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,6 @@ public PslContinuousReader(
6262

6363
@Override
6464
public Offset mergeOffsets(PartitionOffset[] offsets) {
65-
checkArgument(
66-
SparkPartitionOffset.class.isAssignableFrom(offsets.getClass().getComponentType()),
67-
"PartitionOffset object is not assignable to SparkPartitionOffset.");
6865
return SparkSourceOffset.merge(
6966
Arrays.copyOf(offsets, offsets.length, SparkPartitionOffset[].class));
7067
}
@@ -83,8 +80,8 @@ public Offset getStartOffset() {
8380
public void setStartOffset(Optional<Offset> start) {
8481
if (start.isPresent()) {
8582
checkArgument(
86-
SparkSourceOffset.class.isAssignableFrom(start.get().getClass()),
87-
"start offset is not assignable to PslSourceOffset.");
83+
start.get() instanceof SparkSourceOffset,
84+
"start offset is not instance of SparkSourceOffset.");
8885
startOffset = (SparkSourceOffset) start.get();
8986
return;
9087
}
@@ -95,8 +92,7 @@ public void setStartOffset(Optional<Offset> start) {
9592
@Override
9693
public void commit(Offset end) {
9794
checkArgument(
98-
SparkSourceOffset.class.isAssignableFrom(end.getClass()),
99-
"end offset is not assignable to SparkSourceOffset.");
95+
end instanceof SparkSourceOffset, "end offset is not assignable to SparkSourceOffset.");
10096
committer.commit(PslSparkUtils.toPslSourceOffset((SparkSourceOffset) end));
10197
}
10298

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -66,17 +66,17 @@ public PslMicroBatchReader(
6666
public void setOffsetRange(Optional<Offset> start, Optional<Offset> end) {
6767
if (start.isPresent()) {
6868
checkArgument(
69-
SparkSourceOffset.class.isAssignableFrom(start.get().getClass()),
70-
"start offset is not assignable to PslSourceOffset.");
69+
start.get() instanceof SparkSourceOffset,
70+
"start offset is not instance of SparkSourceOffset.");
7171
startOffset = (SparkSourceOffset) start.get();
7272
} else {
7373
startOffset =
7474
PslSparkUtils.getSparkStartOffset(cursorClient, subscriptionPath, topicPartitionCount);
7575
}
7676
if (end.isPresent()) {
7777
checkArgument(
78-
SparkSourceOffset.class.isAssignableFrom(end.get().getClass()),
79-
"start offset is not assignable to PslSourceOffset.");
78+
end.get() instanceof SparkSourceOffset,
79+
"end offset is not instance of SparkSourceOffset.");
8080
endOffset = (SparkSourceOffset) end.get();
8181
} else {
8282
endOffset = PslSparkUtils.toSparkSourceOffset(headOffsetReader.getHeadOffset());
@@ -101,8 +101,7 @@ public Offset deserializeOffset(String json) {
101101
@Override
102102
public void commit(Offset end) {
103103
checkArgument(
104-
SparkSourceOffset.class.isAssignableFrom(end.getClass()),
105-
"end offset is not assignable to SparkSourceOffset.");
104+
end instanceof SparkSourceOffset, "end offset is not instance of SparkSourceOffset.");
106105
committer.commit(PslSparkUtils.toPslSourceOffset((SparkSourceOffset) end));
107106
}
108107

0 commit comments

Comments
 (0)