memory_arrow.rs

memory_arrow.rs

use std::sync::{Mutex, Arc, Barrier};
use std::env;
use std::time::Instant;
use std::sync::mpsc::channel;
use std::ptr::NonNull;
use std::ptr;
use std::alloc::Layout;
use std::mem;

use threadpool::ThreadPool;

use arrow::array::UInt64Array;
use arrow::datatypes::*;

// test memory bandwidth
fn main() {
    let args: Vec<String> = env::args().collect();
    let threads: usize = args[1].parse().unwrap();
    let block_size: usize = args[2].parse().unwrap();

    let counter = Arc::new(Mutex::new(0));
    let jobs = (10000000000u64 / block_size as u64) as usize;
    let pool = ThreadPool::new(threads);

    let start = Instant::now();

    let (tx, rx) = channel();
    for _ in 0..jobs {
        let counter = counter.clone();
        let tx = tx.clone();

        pool.execute(move || {

            unsafe {
                let length = (block_size - 0) as usize;
                let layout = Layout::from_size_align_unchecked(length * mem::size_of::<u64>(), mem::size_of::<u64>());
                let p = std::alloc::alloc(layout) as *mut u64;

                for i in 0 ..block_size as u64 {
                    *p.offset((i-0) as isize) = i;
                }

                let buffer = arrow::buffer::Buffer::from_raw_parts(
                    NonNull::new(p as *mut u8).unwrap(), block_size, block_size);

                let arr_data = arrow::array::ArrayData::builder(DataType::UInt64)
                .len(block_size)
                .offset(0)
                .add_buffer(buffer)
                .build();

                let arr = UInt64Array::from(arr_data);

                 // Create a new builder with a capacity of 100
                // Build the array
                // let builder = arrow::array::ArrayDataBuilder::new(DataType::UInt64);

                let mut num = counter.lock().unwrap();
                *num += arr.len();
            }

            tx.send(1).expect("channel will be there waiting for the pool");
        });
    }
    assert_eq!(rx.iter().take(jobs).fold(0, |a, b| a + b), jobs);

    let cost = start.elapsed();
    println!("Memory Bandwidth: {:?} GB/s, cost: {:?}", (8.0*1e10/1e6) as u128 / cost.as_millis(), cost);
}