Skip to content

Instantly share code, notes, and snippets.

@sundy-li
Last active January 16, 2021 03:39
Show Gist options
  • Save sundy-li/73ac9378c00878bf91603cd90deb5ca5 to your computer and use it in GitHub Desktop.
Save sundy-li/73ac9378c00878bf91603cd90deb5ca5 to your computer and use it in GitHub Desktop.
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);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment