fix: Assign correct partition number to span

Author: jan-auer

src/sentry/spans/consumers/process/factory.py

@@ -64,6 +64,7 @@ def create_with_partitions(
         committer = CommitOffsets(commit)
 
         buffer = SpansBuffer(assigned_shards=[p.index for p in partitions])
+        first_partition = next((p.index for p in partitions), 0)
 
         # patch onto self just for testing
         flusher: ProcessingStrategy[FilteredPayload | int]
@@ -75,7 +76,7 @@ def create_with_partitions(
 
         if self.num_processes != 1:
             run_task = run_task_with_multiprocessing(
-                function=partial(process_batch, buffer),
+                function=partial(process_batch, buffer, first_partition),
                 next_step=flusher,
                 max_batch_size=self.max_batch_size,
                 max_batch_time=self.max_batch_time,
@@ -85,7 +86,7 @@ def create_with_partitions(
             )
         else:
             run_task = RunTask(
-                function=partial(process_batch, buffer),
+                function=partial(process_batch, buffer, first_partition),
                 next_step=flusher,
             )
 
@@ -119,7 +120,9 @@ def shutdown(self) -> None:
 
 
 def process_batch(
-    buffer: SpansBuffer, values: Message[ValuesBatch[tuple[int, KafkaPayload]]]
+    buffer: SpansBuffer,
+    first_partition: int,
+    values: Message[ValuesBatch[tuple[int, KafkaPayload]]],
 ) -> int:
     min_timestamp = None
     spans = []
@@ -130,10 +133,9 @@ def process_batch(
 
         val = rapidjson.loads(payload.value)
 
-        partition_id: int | None = None
-
+        partition_id: int = first_partition
         if len(value.committable) == 1:
-            partition_id = value.committable[next(iter(value.committable))]
+            partition_id = next(iter(value.committable)).index
 
         if killswitches.killswitch_matches_context(
             "spans.drop-in-buffer",
@@ -147,7 +149,7 @@ def process_batch(
             continue
 
         span = Span(
-            partition=partition_id if partition_id is not None else -1,  # TODO: Fallback OK?
+            partition=partition_id,
             trace_id=val["trace_id"],
             span_id=val["span_id"],
             parent_span_id=val.get("parent_span_id"),