Skip to content

Commit bfdf5f5

Browse files
stheppistheppi
authored
Merge pull request #299 from lensesio-dev/test/reproduce-version-detector-bug
fix: ArrayIndexOutOfBoundsException when writing records with different schema versions
2 parents 6ded3f4 + 688871a commit bfdf5f5

10 files changed

Lines changed: 2408 additions & 12 deletions

File tree

‎kafka-connect-aws-s3/src/it/scala/io/lenses/streamreactor/connect/aws/s3/sink/S3SinkTaskParquetSchemaOptimizationTest.scala‎

Lines changed: 353 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,353 @@
1+
/*
2+
* Copyright 2020 Lenses.io
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.lenses.streamreactor.connect.aws.s3.sink
18+
19+
import com.typesafe.scalalogging.LazyLogging
20+
import io.lenses.streamreactor.connect.aws.s3.utils.S3ProxyContainerTest
21+
import io.lenses.streamreactor.connect.cloud.common.config.kcqlprops.PropsKeyEnum.FlushCount
22+
import io.lenses.streamreactor.connect.cloud.common.formats.reader.ParquetFormatReader
23+
import org.apache.avro.generic.GenericRecord
24+
import org.apache.kafka.common.TopicPartition
25+
import org.apache.kafka.common.record.TimestampType
26+
import org.apache.kafka.connect.data.Schema
27+
import org.apache.kafka.connect.data.SchemaAndValue
28+
import org.apache.kafka.connect.data.SchemaBuilder
29+
import org.apache.kafka.connect.data.Struct
30+
import org.apache.kafka.connect.sink.SinkRecord
31+
import org.apache.kafka.connect.sink.SinkTaskContext
32+
import org.mockito.MockitoSugar
33+
import org.scalatest.flatspec.AnyFlatSpec
34+
import org.scalatest.matchers.should.Matchers
35+
36+
import scala.jdk.CollectionConverters.MapHasAsJava
37+
import scala.jdk.CollectionConverters.SeqHasAsJava
38+
39+
/**
40+
* Integration test for Parquet format with schema optimization enabled.
41+
*
42+
* This test verifies that the fix for ArrayIndexOutOfBoundsException works correctly
43+
* when latest.schema.optimization.enabled is true and records with different schema
44+
* versions are interleaved, requiring adaptation to the latest schema.
45+
*/
46+
class S3SinkTaskParquetSchemaOptimizationTest
47+
extends AnyFlatSpec
48+
with Matchers
49+
with S3ProxyContainerTest
50+
with MockitoSugar
51+
with LazyLogging {
52+
53+
private val parquetFormatReader = new ParquetFormatReader()
54+
55+
private val PrefixName = "streamReactorBackups"
56+
private val TopicName = "schemaEvolutionTopic"
57+
58+
private def toSinkRecord(
59+
value: Struct,
60+
topic: String,
61+
partition: Int,
62+
offset: Long,
63+
timestamp: Long,
64+
headers: (String, SchemaAndValue)*,
65+
): SinkRecord = {
66+
val record =
67+
new SinkRecord(topic,
68+
partition,
69+
value.schema(),
70+
value,
71+
value.schema(),
72+
value,
73+
offset,
74+
timestamp,
75+
TimestampType.CREATE_TIME,
76+
)
77+
headers.foreach {
78+
case (name, schemaAndValue) =>
79+
record.headers().add(name, schemaAndValue)
80+
}
81+
record
82+
}
83+
84+
"S3SinkTask" should "write to parquet format using latest schema optimization with interleaved schema versions" in {
85+
val props = (
86+
defaultProps ++
87+
Map(
88+
"connect.s3.kcql" -> s"insert into $BucketName:$PrefixName select * from $TopicName STOREAS PARQUET PROPERTIES('padding.length.partition'='12','padding.length.offset'='12','${FlushCount.entryName}'=6)",
89+
"connect.s3.latest.schema.optimization.enabled" -> "true",
90+
)
91+
).asJava
92+
val expectedFile = "streamReactorBackups/schemaEvolutionTopic/000000000001/000000000006_10001_10006.parquet"
93+
94+
val task = new S3SinkTask()
95+
val ctx = mock[SinkTaskContext]
96+
task.initialize(ctx)
97+
task.start(props)
98+
task.open(Seq(new TopicPartition(TopicName, 1)).asJava)
99+
100+
// Schema V1: Basic user schema
101+
val schemaV1 = SchemaBuilder.struct()
102+
.name("com.example.User")
103+
.field("name", SchemaBuilder.string().required().build())
104+
.field("age", SchemaBuilder.int32().optional().build())
105+
.version(1)
106+
.build()
107+
108+
// Schema V2: Adds optional email field
109+
val schemaV2 = SchemaBuilder.struct()
110+
.name("com.example.User")
111+
.field("name", SchemaBuilder.string().required().build())
112+
.field("age", SchemaBuilder.int32().optional().build())
113+
.field("email", SchemaBuilder.string().optional().build())
114+
.version(2)
115+
.build()
116+
117+
// Schema V3: Adds nested address struct
118+
val addressSchema = SchemaBuilder.struct()
119+
.name("com.example.Address")
120+
.field("street", Schema.OPTIONAL_STRING_SCHEMA)
121+
.field("city", Schema.OPTIONAL_STRING_SCHEMA)
122+
.field("zipCode", Schema.OPTIONAL_STRING_SCHEMA)
123+
.optional()
124+
.build()
125+
126+
val schemaV3 = SchemaBuilder.struct()
127+
.name("com.example.User")
128+
.field("name", SchemaBuilder.string().required().build())
129+
.field("age", SchemaBuilder.int32().optional().build())
130+
.field("email", SchemaBuilder.string().optional().build())
131+
.field("address", addressSchema)
132+
.version(3)
133+
.build()
134+
135+
// Create records with different schema versions interleaved
136+
// Record 1: V1 schema
137+
val struct1 = new Struct(schemaV1).put("name", "Alice").put("age", 30)
138+
139+
// Record 2: V2 schema (introduces email)
140+
val struct2 = new Struct(schemaV2).put("name", "Bob").put("age", 25).put("email", "bob@example.com")
141+
142+
// Record 3: V1 schema again (should be adapted to latest)
143+
val struct3 = new Struct(schemaV1).put("name", "Charlie").put("age", 35)
144+
145+
// Record 4: V3 schema (introduces address)
146+
val address4 = new Struct(addressSchema).put("street", "123 Main St").put("city", "Seattle").put("zipCode", "98101")
147+
val struct4 =
148+
new Struct(schemaV3).put("name", "Diana").put("age", 28).put("email", "diana@example.com").put("address",
149+
address4,
150+
)
151+
152+
// Record 5: V2 schema again (should be adapted to latest V3)
153+
val struct5 = new Struct(schemaV2).put("name", "Eve").put("age", 32).put("email", "eve@example.com")
154+
155+
// Record 6: V1 schema again (should be adapted to latest V3)
156+
val struct6 = new Struct(schemaV1).put("name", "Frank").put("age", 40)
157+
158+
val record1 = toSinkRecord(struct1, TopicName, 1, 1L, 10001L)
159+
val record2 = toSinkRecord(struct2, TopicName, 1, 2L, 10002L)
160+
val record3 = toSinkRecord(struct3, TopicName, 1, 3L, 10003L)
161+
val record4 = toSinkRecord(struct4, TopicName, 1, 4L, 10004L)
162+
val record5 = toSinkRecord(struct5, TopicName, 1, 5L, 10005L)
163+
val record6 = toSinkRecord(struct6, TopicName, 1, 6L, 10006L)
164+
165+
task.put(List(record1, record2, record3, record4, record5, record6).asJava)
166+
task.close(Seq(new TopicPartition(TopicName, 1)).asJava)
167+
task.stop()
168+
169+
// Verify output - all records should be in a single file with the latest schema
170+
val files = listBucketPath(BucketName, "streamReactorBackups/schemaEvolutionTopic/000000000001/")
171+
files.size should be(1)
172+
val bytes = remoteFileAsBytes(BucketName, expectedFile)
173+
174+
val genericRecords = parquetFormatReader.read(bytes)
175+
genericRecords.size should be(6)
176+
177+
// All records should have been adapted to Schema V3
178+
// Record 1: V1 -> V3 (age present, email and address null)
179+
val rec1 = genericRecords.head
180+
rec1.get("name").toString should be("Alice")
181+
rec1.get("age") should be(30)
182+
rec1.get("email") should be(null)
183+
rec1.get("address") should be(null)
184+
185+
// Record 2: V2 -> V3 (email present, address null)
186+
val rec2 = genericRecords(1)
187+
rec2.get("name").toString should be("Bob")
188+
rec2.get("age") should be(25)
189+
rec2.get("email").toString should be("bob@example.com")
190+
rec2.get("address") should be(null)
191+
192+
// Record 3: V1 -> V3
193+
val rec3 = genericRecords(2)
194+
rec3.get("name").toString should be("Charlie")
195+
rec3.get("age") should be(35)
196+
rec3.get("email") should be(null)
197+
rec3.get("address") should be(null)
198+
199+
// Record 4: V3 (full schema)
200+
val rec4 = genericRecords(3)
201+
val address4Rec = rec4.get("address").asInstanceOf[GenericRecord]
202+
rec4.get("name").toString should be("Diana")
203+
rec4.get("age") should be(28)
204+
rec4.get("email").toString should be("diana@example.com")
205+
address4Rec.get("street").toString should be("123 Main St")
206+
address4Rec.get("city").toString should be("Seattle")
207+
address4Rec.get("zipCode").toString should be("98101")
208+
209+
// Record 5: V2 -> V3
210+
val rec5 = genericRecords(4)
211+
rec5.get("name").toString should be("Eve")
212+
rec5.get("age") should be(32)
213+
rec5.get("email").toString should be("eve@example.com")
214+
rec5.get("address") should be(null)
215+
216+
// Record 6: V1 -> V3
217+
val rec6 = genericRecords(5)
218+
rec6.get("name").toString should be("Frank")
219+
rec6.get("age") should be(40)
220+
rec6.get("email") should be(null)
221+
rec6.get("address") should be(null)
222+
}
223+
224+
"S3SinkTask" should "handle nested struct evolution with parquet and schema optimization" in {
225+
val props = (
226+
defaultProps ++
227+
Map(
228+
"connect.s3.kcql" -> s"insert into $BucketName:$PrefixName select * from $TopicName STOREAS PARQUET PROPERTIES('padding.length.partition'='12','padding.length.offset'='12','${FlushCount.entryName}'=4)",
229+
"connect.s3.latest.schema.optimization.enabled" -> "true",
230+
)
231+
).asJava
232+
val expectedFile = "streamReactorBackups/schemaEvolutionTopic/000000000001/000000000004_20001_20004.parquet"
233+
234+
val task = new S3SinkTask()
235+
val ctx = mock[SinkTaskContext]
236+
task.initialize(ctx)
237+
task.start(props)
238+
task.open(Seq(new TopicPartition(TopicName, 1)).asJava)
239+
240+
// Schema V1: metadata with basic fields
241+
val metadataSchemaV1 = SchemaBuilder.struct()
242+
.name("com.example.Metadata")
243+
.field("topic", Schema.STRING_SCHEMA)
244+
.field("partition", Schema.INT32_SCHEMA)
245+
.optional()
246+
.build()
247+
248+
val orderSchemaV1 = SchemaBuilder.struct()
249+
.name("com.example.Order")
250+
.field("orderId", Schema.STRING_SCHEMA)
251+
.field("amount", Schema.FLOAT64_SCHEMA)
252+
.field("metadata", metadataSchemaV1)
253+
.version(1)
254+
.build()
255+
256+
// Schema V2: metadata adds offset and timestamp fields
257+
val metadataSchemaV2 = SchemaBuilder.struct()
258+
.name("com.example.Metadata")
259+
.field("topic", Schema.STRING_SCHEMA)
260+
.field("partition", Schema.INT32_SCHEMA)
261+
.field("offset", Schema.OPTIONAL_INT64_SCHEMA)
262+
.field("timestamp", Schema.OPTIONAL_INT64_SCHEMA)
263+
.optional()
264+
.build()
265+
266+
val orderSchemaV2 = SchemaBuilder.struct()
267+
.name("com.example.Order")
268+
.field("orderId", Schema.STRING_SCHEMA)
269+
.field("amount", Schema.FLOAT64_SCHEMA)
270+
.field("metadata", metadataSchemaV2)
271+
.version(2)
272+
.build()
273+
274+
// Record 1: V1 schema
275+
val meta1 = new Struct(metadataSchemaV1).put("topic", TopicName).put("partition", 1)
276+
val struct1 = new Struct(orderSchemaV1).put("orderId", "ORD-001").put("amount", 100.50).put("metadata", meta1)
277+
278+
// Record 2: V2 schema (introduces offset and timestamp in metadata)
279+
val meta2 =
280+
new Struct(metadataSchemaV2).put("topic", TopicName).put("partition", 1).put("offset", 2L).put("timestamp",
281+
20002L,
282+
)
283+
val struct2 = new Struct(orderSchemaV2).put("orderId", "ORD-002").put("amount", 250.75).put("metadata", meta2)
284+
285+
// Record 3: V1 schema again (should be adapted to V2)
286+
val meta3 = new Struct(metadataSchemaV1).put("topic", TopicName).put("partition", 1)
287+
val struct3 = new Struct(orderSchemaV1).put("orderId", "ORD-003").put("amount", 75.00).put("metadata", meta3)
288+
289+
// Record 4: V2 schema
290+
val meta4 =
291+
new Struct(metadataSchemaV2).put("topic", TopicName).put("partition", 1).put("offset", 4L).put("timestamp",
292+
20004L,
293+
)
294+
val struct4 = new Struct(orderSchemaV2).put("orderId", "ORD-004").put("amount", 500.00).put("metadata", meta4)
295+
296+
val record1 = toSinkRecord(struct1, TopicName, 1, 1L, 20001L)
297+
val record2 = toSinkRecord(struct2, TopicName, 1, 2L, 20002L)
298+
val record3 = toSinkRecord(struct3, TopicName, 1, 3L, 20003L)
299+
val record4 = toSinkRecord(struct4, TopicName, 1, 4L, 20004L)
300+
301+
task.put(List(record1, record2, record3, record4).asJava)
302+
task.close(Seq(new TopicPartition(TopicName, 1)).asJava)
303+
task.stop()
304+
305+
val files = listBucketPath(BucketName, "streamReactorBackups/schemaEvolutionTopic/000000000001/")
306+
files.size should be(1)
307+
val bytes = remoteFileAsBytes(BucketName, expectedFile)
308+
309+
val genericRecords = parquetFormatReader.read(bytes)
310+
genericRecords.size should be(4)
311+
312+
// Record 1: V1 -> V2 (metadata.offset and metadata.timestamp should be null)
313+
val rec1 = genericRecords.head
314+
val metaRec1 = rec1.get("metadata").asInstanceOf[GenericRecord]
315+
rec1.get("orderId").toString should be("ORD-001")
316+
rec1.get("amount") should be(100.50)
317+
metaRec1.get("topic").toString should be(TopicName)
318+
metaRec1.get("partition") should be(1)
319+
metaRec1.get("offset") should be(null)
320+
metaRec1.get("timestamp") should be(null)
321+
322+
// Record 2: V2 (full metadata)
323+
val rec2 = genericRecords(1)
324+
val metaRec2 = rec2.get("metadata").asInstanceOf[GenericRecord]
325+
rec2.get("orderId").toString should be("ORD-002")
326+
rec2.get("amount") should be(250.75)
327+
metaRec2.get("topic").toString should be(TopicName)
328+
metaRec2.get("partition") should be(1)
329+
metaRec2.get("offset") should be(2L)
330+
metaRec2.get("timestamp") should be(20002L)
331+
332+
// Record 3: V1 -> V2
333+
val rec3 = genericRecords(2)
334+
val metaRec3 = rec3.get("metadata").asInstanceOf[GenericRecord]
335+
rec3.get("orderId").toString should be("ORD-003")
336+
rec3.get("amount") should be(75.00)
337+
metaRec3.get("topic").toString should be(TopicName)
338+
metaRec3.get("partition") should be(1)
339+
metaRec3.get("offset") should be(null)
340+
metaRec3.get("timestamp") should be(null)
341+
342+
// Record 4: V2
343+
val rec4 = genericRecords(3)
344+
val metaRec4 = rec4.get("metadata").asInstanceOf[GenericRecord]
345+
rec4.get("orderId").toString should be("ORD-004")
346+
rec4.get("amount") should be(500.00)
347+
metaRec4.get("topic").toString should be(TopicName)
348+
metaRec4.get("partition") should be(1)
349+
metaRec4.get("offset") should be(4L)
350+
metaRec4.get("timestamp") should be(20004L)
351+
}
352+
353+
}

0 commit comments

Comments
 (0)