Function create_record_batch

datafusion/core/benches/data_utils/mod.rs:89–140

(
    schema: SchemaRef,
    rng: &mut StdRng,
    batch_size: usize,
    batch_index: usize,
)

}

fn create_record_batch(
    schema: SchemaRef,
    rng: &mut StdRng,
    batch_size: usize,
    batch_index: usize,
) -> RecordBatch {
    // Randomly choose from 4 distinct key values; a higher number increases sparseness.
    let key_suffixes = [0, 1, 2, 3];
    let keys = StringArray::from_iter_values(
        (0..batch_size).map(|_| format!("hi{}", key_suffixes.choose(rng).unwrap())),
    );

    let values = create_data(rng, batch_size, 0.5);

    // Integer values between [0, u64::MAX].
    let integer_values_wide = (0..batch_size)
        .map(|_| rng.random::<u64>())
        .collect::<Vec<_>>();

    // Integer values between [0, 1000).
    let integer_values_mid = (0..batch_size)
        .map(|_| rng.random_range(0..1000))
        .collect::<Vec<_>>();

    // Integer values between [0, 10).
    let integer_values_narrow = (0..batch_size)
        .map(|_| rng.random_range(0..10))
        .collect::<Vec<_>>();

    let mut dict_builder = StringDictionaryBuilder::<Int32Type>::new();
    for _ in 0..batch_size {
        if rng.random::<f64>() > 0.9 {
            dict_builder.append_null();
        } else {
            dict_builder.append_value(format!("market_{}", rng.random_range(0..10)));
        }
    }

    RecordBatch::try_new(
        schema,
        vec![
            Arc::new(keys),
            Arc::new(Float32Array::from(vec![batch_index as f32; batch_size])),
            Arc::new(Float64Array::from(values)),
            Arc::new(UInt64Array::from(integer_values_wide)),
            Arc::new(UInt64Array::from(integer_values_mid)),
            Arc::new(UInt64Array::from(integer_values_narrow)),
            Arc::new(dict_builder.finish()),
        ],
    )
    .unwrap()
}

/// Create record batches of `partitions_len` partitions and `batch_size` for each batch,
/// with a total number of `array_len` records