use arrow::{array::RecordBatch, datatypes::Schema};

use crate::OnceReader;
use crate::readers::{dataframe_from_stream, log_stream_batch};

pub struct ArrowReader {
    rb: RecordBatch,
}

impl ArrowReader {
    pub fn new(rb: RecordBatch) -> ArrowReader {
        ArrowReader { rb }
    }
}

impl OnceReader for ArrowReader {
    fn read(&mut self) -> crate::Result<autark_dataframe::DataFrame> {
        dataframe_from_stream(self)
    }

    fn schema(&self) -> crate::Result<std::sync::Arc<Schema>> {
        Ok(self.rb.schema())
    }

    fn read_stream(
        &mut self,
        sink: &mut dyn FnMut(RecordBatch) -> crate::Result<()>,
    ) -> crate::Result<()> {
        log_stream_batch("arrow", 0, &self.rb);
        sink(self.rb.clone())
    }
}
