partitioning without using airbyteValues

Author: frifriSF59

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/message/DestinationMessage.kt

@@ -183,8 +183,8 @@ data class EnrichedDestinationRecordAirbyteValue(
 
 data class DestinationRecordRaw(
     val stream: DestinationStream.Descriptor,
-    private val rawData: AirbyteMessage,
-    private val serialized: String
+    val rawData: AirbyteMessage,
+    val serialized: String
 ) {
     fun asDestinationRecordAirbyteValue(): DestinationRecordAirbyteValue {
         return DestinationRecordAirbyteValue(

airbyte-integrations/connectors/destination-s3-data-lake/src/main/kotlin/io/airbyte/integrations/destination/s3_data_lake/S3DataLakePartitioner.kt

@@ -4,9 +4,11 @@
 
 package io.airbyte.integrations.destination.s3_data_lake
 
+import com.fasterxml.jackson.databind.JsonNode
 import io.airbyte.cdk.load.command.Dedupe
 import io.airbyte.cdk.load.command.DestinationCatalog
 import io.airbyte.cdk.load.data.ObjectValue
+import io.airbyte.cdk.load.data.json.toAirbyteValue
 import io.airbyte.cdk.load.message.DestinationRecordAirbyteValue
 import io.airbyte.cdk.load.message.DestinationRecordRaw
 import io.airbyte.cdk.load.pipeline.InputPartitioner
@@ -27,15 +29,18 @@ class S3DataLakePartitioner(catalog: DestinationCatalog) : InputPartitioner {
     private val random = Random(System.currentTimeMillis())
 
     override fun getPartition(record: DestinationRecordRaw, numParts: Int): Int {
-        val recordAirbyteValue = record.asDestinationRecordAirbyteValue()
-
         if (numParts == 1) {
             return 0
         }
 
-        streamToPrimaryKeyFieldNames[recordAirbyteValue.stream]?.let { primaryKey ->
-            val primaryKeyValues =
-                primaryKey.map { it.map { key -> (recordAirbyteValue.data as ObjectValue).values[key] } }
+        streamToPrimaryKeyFieldNames[record.stream]?.let { primaryKey ->
+            val jsonData = record.rawData.record.data as JsonNode
+
+            val primaryKeyValues = primaryKey.map { keys ->
+                keys.map { key ->
+                    if (jsonData.has(key)) jsonData.get(key) else null
+                }
+            }
             val hash = primaryKeyValues.hashCode()
             /** abs(MIN_VALUE) == MIN_VALUE, so we need to handle this case separately */
             if (hash == Int.MIN_VALUE) {