Function nested_batch

datafusion/core/benches/parquet_struct_projection.rs:313–358

(batch_id: usize)

}

fn nested_batch(batch_id: usize) -> RecordBatch {
    let schema = nested_schema();
    let len = WRITE_RECORD_BATCH_SIZE;

    let base_id = (batch_id * len) as i32;
    let id_values: Vec<i32> = (0..len).map(|i| base_id + i as i32).collect();
    let id_array = Arc::new(Int32Array::from(id_values.clone()));
    let small_int_array = Arc::new(Int32Array::from(id_values));

    let large_string: String = "x".repeat(LARGE_STRING_LEN);

    let mut sb1 = StringBuilder::new();
    let mut sb2 = StringBuilder::new();
    for _ in 0..len {
        sb1.append_value(&large_string);
        sb2.append_value(&large_string);
    }

    let inner_struct = StructArray::from(vec![
        (
            Arc::new(Field::new("large_string", DataType::Utf8, false)),
            Arc::new(sb1.finish()) as ArrayRef,
        ),
        (
            Arc::new(Field::new("small_int", DataType::Int32, false)),
            small_int_array as ArrayRef,
        ),
    ]);

    let inner_fields = Fields::from(vec![
        Field::new("large_string", DataType::Utf8, false),
        Field::new("small_int", DataType::Int32, false),
    ]);
    let outer_struct = StructArray::from(vec![
        (
            Arc::new(Field::new("inner", DataType::Struct(inner_fields), false)),
            Arc::new(inner_struct) as ArrayRef,
        ),
        (
            Arc::new(Field::new("extra_string", DataType::Utf8, false)),
            Arc::new(sb2.finish()) as ArrayRef,
        ),
    ]);

    RecordBatch::try_new(schema, vec![id_array, Arc::new(outer_struct)]).unwrap()
}

fn nested_benchmarks(c: &mut Criterion) {
    let temp_file = generate_file(nested_schema(), nested_batch, "nested_struct");