Problem: In order to be able to use the arrow
crate's infer_file_schema
function with input piped to /dev/stdin
, I need to manually implement the Seek
trait on a Reader, as it doesn't exist for piped /dev/stdin/
by default.
The function that uses that Reader reads n
lines, infers the schema from these n
lines, then seeks back to the start of the "file".
So I don't need to buffer more than these n
lines.
This is what I wrote:
struct SeekableReader<R: std::io::Read> {
inner: R, // underlying reader
buffer: Vec<u8>, // buffer for the first n lines
buffer_size: usize, // size of the buffer in bytes
pos: usize, // current position in the buffer
seekable: bool, // whether seek is still possible
}
const BUFFER_SIZE: usize = 8192;
impl<R: std::io::Read> SeekableReader<R> {
fn new(reader: R, lines_to_buffer: Option<usize>) -> Self {
let mut inner = reader;
let mut buffer = Vec::<u8>::with_capacity(BUFFER_SIZE);
let mut lines = 0;
let mut bytes_read = 0;
let mut bytes_before;
loop {
bytes_before = bytes_read;
buffer.append(&mut vec![0; BUFFER_SIZE - (buffer.len() - bytes_read)]);
bytes_read += inner
.read(&mut buffer[bytes_read..bytes_read + BUFFER_SIZE])
.unwrap();
lines += buffer[bytes_before..bytes_read]
.iter()
.filter(|&&x| x == 10)
.count();
if let Some(lines_to_buffer) = lines_to_buffer {
// +1 because there may be a header
if lines > lines_to_buffer + 1 {
break;
}
}
if bytes_read == 0 {
break;
}
}
SeekableReader {
inner,
buffer,
buffer_size: bytes_read,
pos: 0,
seekable: true,
}
}
}
impl<R: std::io::Read> std::io::Read for SeekableReader<R> {
fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
let buf_len = buf.len();
if self.pos < self.buffer_size {
if self.buffer_size - self.pos < buf_len {
buf[..self.buffer_size - self.pos]
.copy_from_slice(&self.buffer[self.pos..self.buffer_size]);
let len_read = self.buffer_size - self.pos;
self.pos = self.buffer_size;
Ok(len_read)
} else {
buf.copy_from_slice(&self.buffer[self.pos..self.pos + buf_len]);
self.pos += buf_len;
Ok(buf_len)
}
} else {
self.seekable = false;
self.inner.read(buf)
}
}
}
impl<R: std::io::Read> std::io::Seek for SeekableReader<R> {
fn seek(&mut self, pos: SeekFrom) -> Result<u64, std::io::Error> {
let error = Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"Seeking outside of buffer, please report to https://github.com/domoritz/arrow-tools/issues/new".to_string(),
));
if self.seekable {
match pos {
SeekFrom::Start(pos) => {
if pos >= self.buffer_size as u64 {
error
} else {
self.pos = pos as usize;
Ok(pos)
}
}
SeekFrom::Current(pos) => {
let new_pos = self.pos as i64 + pos;
if 0 <= new_pos && new_pos < self.buffer_size as i64 {
self.pos = new_pos as usize;
Ok(new_pos as u64)
} else {
error
}
}
SeekFrom::End(_) => error,
}
} else {
error
}
}
}
You can see this struct in action here: https://github.com/domoritz/arrow-tools/pull/10/files