From ff6b3f289daf5b09bcd8d44d48595a5f012cb4d5 Mon Sep 17 00:00:00 2001
From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com>
Date: Fri, 17 Jul 2026 11:14:18 +0200
Subject: [PATCH] add decimal and fsb arrow_writer bench

---
 arrow/src/util/data_gen.rs      | 30 ++++++++++++++++++++++++++++++
 parquet/benches/arrow_writer.rs | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/arrow/src/util/data_gen.rs b/arrow/src/util/data_gen.rs
index b5eb9efbc390..1953519df4fe 100644
--- a/arrow/src/util/data_gen.rs
+++ b/arrow/src/util/data_gen.rs
@@ -181,6 +181,8 @@ pub fn create_random_array(
             crate::compute::cast(&v, d)?
         }
         Map(_, _) => create_random_map_array(field, size, null_density, true_density)?,
+        Decimal32(_, _) => create_random_decimal_array(field, size, null_density)?,
+        Decimal64(_, _) => create_random_decimal_array(field, size, null_density)?,
         Decimal128(_, _) => create_random_decimal_array(field, size, null_density)?,
         Decimal256(_, _) => create_random_decimal_array(field, size, null_density)?,
         RunEndEncoded(index, value) => {
@@ -205,6 +207,34 @@ fn create_random_decimal_array(field: &Field, size: usize, null_density: f32) ->
     let mut rng = seedable_rng();
 
     match field.data_type() {
+        DataType::Decimal32(precision, scale) => {
+            let values = (0..size)
+                .map(|_| {
+                    if rng.random::<f32>() < null_density {
+                        None
+                    } else {
+                        Some(rng.random::<i32>())
+                    }
+                })
+                .collect::<Vec<_>>();
+            Ok(Arc::new(
+                Decimal32Array::from(values).with_precision_and_scale(*precision, *scale)?,
+            ))
+        }
+        DataType::Decimal64(precision, scale) => {
+            let values = (0..size)
+                .map(|_| {
+                    if rng.random::<f32>() < null_density {
+                        None
+                    } else {
+                        Some(rng.random::<i64>())
+                    }
+                })
+                .collect::<Vec<_>>();
+            Ok(Arc::new(
+                Decimal64Array::from(values).with_precision_and_scale(*precision, *scale)?,
+            ))
+        }
         DataType::Decimal128(precision, scale) => {
             let values = (0..size)
                 .map(|_| {
diff --git a/parquet/benches/arrow_writer.rs b/parquet/benches/arrow_writer.rs
index 073716e2f3ec..adc4dcc9ea94 100644
--- a/parquet/benches/arrow_writer.rs
+++ b/parquet/benches/arrow_writer.rs
@@ -251,6 +251,22 @@ fn create_float_bench_batch_with_nans(size: usize, nan_density: f32) -> Result<R
     )?)
 }
 
+fn create_decimal_bench_batch(size: usize, null_density: f32) -> Result<RecordBatch> {
+    let fields = vec![
+        Field::new("_1", Decimal32Type::DEFAULT_TYPE, false),
+        Field::new("_2", Decimal64Type::DEFAULT_TYPE, false),
+        Field::new("_3", Decimal128Type::DEFAULT_TYPE, false),
+        Field::new("_4", Decimal256Type::DEFAULT_TYPE, false),
+    ];
+    let schema = Schema::new(fields);
+    Ok(create_random_batch(
+        Arc::new(schema),
+        size,
+        null_density,
+        0.75,
+    )?)
+}
+
 fn create_list_primitive_bench_batch(
     size: usize,
     null_density: f32,
@@ -332,6 +348,17 @@ fn create_struct_bench_batch(size: usize, null_density: f32) -> Result<RecordBat
     )?)
 }
 
+fn create_fsb_bench_batch(size: usize, null_density: f32, len: i32) -> Result<RecordBatch> {
+    let fields = vec![Field::new("_1", DataType::FixedSizeBinary(len), true)];
+    let schema = Schema::new(fields);
+    Ok(create_random_batch(
+        Arc::new(schema),
+        size,
+        null_density,
+        0.75,
+    )?)
+}
+
 fn create_nested_list_bench_batch(size: usize, null_density: f32) -> Result<RecordBatch> {
     // List<List<Int32>> — exercises the nested repetition (non-batched) path
     let fields = vec![Field::new(
@@ -528,6 +555,9 @@ fn create_batches() -> Vec<(&'static str, RecordBatch)> {
     let batch = create_float_bench_batch_with_nans(BATCH_SIZE, 0.5).unwrap();
     batches.push(("float_with_nans", batch));
 
+    let batch = create_decimal_bench_batch(BATCH_SIZE, 0.75).unwrap();
+    batches.push(("decimal", batch));
+
     let batch = create_list_primitive_bench_batch(BATCH_SIZE, 0.25, 0.75).unwrap();
     batches.push(("list_primitive", batch));
 
@@ -552,6 +582,9 @@ fn create_batches() -> Vec<(&'static str, RecordBatch)> {
     let batch = create_struct_bench_batch(BATCH_SIZE, 1.0).unwrap();
     batches.push(("struct_all_null", batch));
 
+    let batch = create_fsb_bench_batch(BATCH_SIZE, 0.9, 16).unwrap();
+    batches.push(("fsb", batch));
+
     let batch = create_nested_list_bench_batch(BATCH_SIZE, 0.25).unwrap();
     batches.push(("list_nested", batch));