perf: Optimizing sorting by id/parent id (open-telemetry#2121)

# Change Summary

This PR optimizes some id sorting logic that we have for split using
techniques for two column sorts that we also use elsewhere in the
transport optimize code.

| Scenario | Baseline (RowConverter) | Optimized | Speedup |

|----------------------------------------|-------------------------|-----------|---------|
| u16 id only (random) | 7.23 us | 7.27 us | ~same |
| u16 id only (sorted) | 1.20 us | 71 ns | 17x |
| u16 pid + u16 id (native) | 25.4 us | 7.33 us | 3.5x |
| u32 id + Dict<u8,u32> pid | 28.7 us | 7.89 us | 3.6x |
| u32 id + Dict<u16,u32> pid | 28.3 us | 7.93 us | 3.6x |
| u32 id (25% null) + Dict<u8,u32> pid | 30.4 us | 10.6 us | 2.9x |

## What issue does this PR close?

* Closes open-telemetry#2122

## How are these changes tested?

New unit testing suite comparing results with the old and also
indirectly covered by all the split tests.

## Are there any user-facing changes?

No

---------

Co-authored-by: Tom Tan <lilotom@gmail.com>
Co-authored-by: Drew Relmas <drewrelmas@gmail.com>

Author: 3 people

rust/otap-dataflow/crates/pdata/Cargo.toml

@@ -50,3 +50,7 @@ bench = []
 [[bench]]
 name = "concatenate"
 harness = false
+
+[[bench]]
+name = "sort_by_parent_then_id"
+harness = false

rust/otap-dataflow/crates/pdata/benches/sort_by_parent_then_id.rs

@@ -0,0 +1,186 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Benchmarks for `sort_by_parent_then_id`.
+
+use std::sync::Arc;
+
+use arrow::array::{
+    ArrayRef, DictionaryArray, PrimitiveArray, RecordBatch, UInt16Array, UInt32Array,
+};
+use arrow::datatypes::{DataType, Field, Schema, UInt8Type, UInt16Type};
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+
+use otap_df_pdata::otap::transform::util::sort_by_parent_then_id;
+
+const NUM_ROWS: usize = 1_000;
+const SEED: u64 = 42;
+
+criterion_group!(benches, bench_sort);
+criterion_main!(benches);
+
+fn bench_sort(c: &mut Criterion) {
+    let mut group = c.benchmark_group("sort_by_parent_then_id");
+    let mut rng = StdRng::seed_from_u64(SEED);
+
+    // 1) u16 id only - random order
+    {
+        let id_col = Arc::new(random_u16_array(&mut rng, NUM_ROWS)) as ArrayRef;
+        let batch = make_batch(vec![("id", DataType::UInt16, id_col)]);
+        let _ = group.bench_with_input(
+            BenchmarkId::new("u16_id_only", "random"),
+            &batch,
+            |b, batch| {
+                b.iter(|| sort_by_parent_then_id(batch.clone()).expect("sort failed"));
+            },
+        );
+    }
+
+    // 2) u16 id only - already sorted
+    {
+        let id_col = Arc::new(sorted_u16_array(NUM_ROWS)) as ArrayRef;
+        let batch = make_batch(vec![("id", DataType::UInt16, id_col)]);
+        let _ = group.bench_with_input(
+            BenchmarkId::new("u16_id_only", "sorted"),
+            &batch,
+            |b, batch| {
+                b.iter(|| sort_by_parent_then_id(batch.clone()).expect("sort failed"));
+            },
+        );
+    }
+
+    // 3) u16 parent_id + u16 id - random (native)
+    {
+        let pid = Arc::new(random_u16_array(&mut rng, NUM_ROWS)) as ArrayRef;
+        let id_col = Arc::new(random_u16_array(&mut rng, NUM_ROWS)) as ArrayRef;
+        let batch = make_batch(vec![
+            ("parent_id", DataType::UInt16, pid),
+            ("id", DataType::UInt16, id_col),
+        ]);
+        let _ = group.bench_with_input(
+            BenchmarkId::new("u16_pid_u16_id_native", "random"),
+            &batch,
+            |b, batch| {
+                b.iter(|| sort_by_parent_then_id(batch.clone()).expect("sort failed"));
+            },
+        );
+    }
+
+    // 4) u32 id + Dict<UInt8, UInt32> parent_id
+    {
+        let pid = random_dict_u8_u32(&mut rng, NUM_ROWS, 50);
+        let id_col = Arc::new(random_u32_array(&mut rng, NUM_ROWS)) as ArrayRef;
+        let dict_dt = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32));
+        let batch = make_batch(vec![
+            ("parent_id", dict_dt, pid),
+            ("id", DataType::UInt32, id_col),
+        ]);
+        let _ = group.bench_with_input(
+            BenchmarkId::new("u32_id_dict_u8_u32_pid", "random"),
+            &batch,
+            |b, batch| {
+                b.iter(|| sort_by_parent_then_id(batch.clone()).expect("sort failed"));
+            },
+        );
+    }
+
+    // 5) u32 id + Dict<UInt16, UInt32> parent_id
+    {
+        let pid = random_dict_u16_u32(&mut rng, NUM_ROWS, 200);
+        let id_col = Arc::new(random_u32_array(&mut rng, NUM_ROWS)) as ArrayRef;
+        let dict_dt = DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::UInt32));
+        let batch = make_batch(vec![
+            ("parent_id", dict_dt, pid),
+            ("id", DataType::UInt32, id_col),
+        ]);
+        let _ = group.bench_with_input(
+            BenchmarkId::new("u32_id_dict_u16_u32_pid", "random"),
+            &batch,
+            |b, batch| {
+                b.iter(|| sort_by_parent_then_id(batch.clone()).expect("sort failed"));
+            },
+        );
+    }
+
+    // 6) u32 id (25% null) + Dict<UInt8, UInt32> parent_id
+    {
+        let pid = random_dict_u8_u32(&mut rng, NUM_ROWS, 50);
+        let id_col = Arc::new(random_nullable_u32_array(&mut rng, NUM_ROWS)) as ArrayRef;
+        let dict_dt = DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32));
+        let batch = make_batch(vec![
+            ("parent_id", dict_dt, pid),
+            ("id", DataType::UInt32, id_col),
+        ]);
+        let _ = group.bench_with_input(
+            BenchmarkId::new("u32_id_null25_dict_u8_u32_pid", "random"),
+            &batch,
+            |b, batch| {
+                b.iter(|| sort_by_parent_then_id(batch.clone()).expect("sort failed"));
+            },
+        );
+    }
+
+    group.finish();
+}
+
+fn make_batch(fields: Vec<(&str, DataType, ArrayRef)>) -> RecordBatch {
+    let schema = Arc::new(Schema::new(
+        fields
+            .iter()
+            .map(|(name, dt, _)| Field::new(*name, dt.clone(), true))
+            .collect::<Vec<_>>(),
+    ));
+    let columns: Vec<ArrayRef> = fields.into_iter().map(|(_, _, arr)| arr).collect();
+    RecordBatch::try_new(schema, columns).expect("valid record batch")
+}
+
+fn random_u16_array(rng: &mut StdRng, n: usize) -> UInt16Array {
+    PrimitiveArray::from_iter_values((0..n).map(|_| rng.random::<u16>()))
+}
+
+fn sorted_u16_array(n: usize) -> UInt16Array {
+    PrimitiveArray::from_iter_values((0..n).map(|i| i as u16))
+}
+
+fn random_u32_array(rng: &mut StdRng, n: usize) -> UInt32Array {
+    PrimitiveArray::from_iter_values((0..n).map(|_| rng.random::<u32>()))
+}
+
+/// Build a Dict<UInt8, UInt32> array with `n` rows, `n_values` distinct values.
+fn random_dict_u8_u32(rng: &mut StdRng, n: usize, n_values: usize) -> ArrayRef {
+    let values: Vec<u32> = (0..n_values).map(|_| rng.random::<u32>()).collect();
+    let keys: Vec<u8> = (0..n)
+        .map(|_| rng.random_range(0..n_values as u8))
+        .collect();
+    let keys_arr = PrimitiveArray::<UInt8Type>::from_iter_values(keys);
+    let values_arr = Arc::new(UInt32Array::from(values));
+    Arc::new(DictionaryArray::new(keys_arr, values_arr))
+}
+
+/// Build a Dict<UInt16, UInt32> array with `n` rows, `n_values` distinct values.
+fn random_dict_u16_u32(rng: &mut StdRng, n: usize, n_values: usize) -> ArrayRef {
+    let values: Vec<u32> = (0..n_values).map(|_| rng.random::<u32>()).collect();
+    let keys: Vec<u16> = (0..n)
+        .map(|_| rng.random_range(0..n_values as u16))
+        .collect();
+    let keys_arr = PrimitiveArray::<UInt16Type>::from_iter_values(keys);
+    let values_arr = Arc::new(UInt32Array::from(values));
+    Arc::new(DictionaryArray::new(keys_arr, values_arr))
+}
+
+/// Build a UInt32 array with `n` rows where ~25% are null.
+fn random_nullable_u32_array(rng: &mut StdRng, n: usize) -> UInt32Array {
+    UInt32Array::from(
+        (0..n)
+            .map(|_| {
+                if rng.random_range(0u8..4) == 0 {
+                    None
+                } else {
+                    Some(rng.random::<u32>())
+                }
+            })
+            .collect::<Vec<_>>(),
+    )
+}

rust/otap-dataflow/crates/pdata/src/error.rs

@@ -112,6 +112,9 @@ pub enum Error {
     )]
     UnsupportedParentIdType { actual: DataType },
 
+    #[error("parent_id column must not contain nulls")]
+    NullParentId,
+
     #[error("Unsupported payload type, got: {}", actual)]
     UnsupportedPayloadType { actual: i32 },