polars_io/utils/
stream_buf_reader.rs

1use std::io::{BufRead, Cursor};
2
3use polars_buffer::Buffer;
4#[cfg(feature = "async")]
5use polars_core::runtime::ASYNC;
6use polars_error::PolarsResult;
7#[cfg(feature = "async")]
8use polars_utils::async_utils::tokio_handle_ext;
9#[cfg(feature = "async")]
10use tokio::sync::OwnedSemaphorePermit;
11
12#[cfg(feature = "async")]
13pub struct OpenReaderState {
14    receiver: tokio::sync::mpsc::Receiver<(
15        tokio_handle_ext::AbortOnDropHandle<PolarsResult<Buffer<u8>>>,
16        OwnedSemaphorePermit,
17    )>,
18    producer_task_handle: tokio_handle_ext::AbortOnDropHandle<std::io::Result<()>>,
19    current: Buffer<u8>,
20}
21
22/// `BufRead` interface for a channel that is receiving `Buffer<u8>` bytes.
23#[cfg(feature = "async")]
24pub enum StreamBufReader {
25    Open(OpenReaderState),
26    Finished,
27}
28
29#[cfg(feature = "async")]
30impl StreamBufReader {
31    pub fn new(
32        receiver: tokio::sync::mpsc::Receiver<(
33            tokio_handle_ext::AbortOnDropHandle<PolarsResult<Buffer<u8>>>,
34            OwnedSemaphorePermit,
35        )>,
36        producer_task_handle: tokio_handle_ext::AbortOnDropHandle<std::io::Result<()>>,
37    ) -> Self {
38        Self::Open(OpenReaderState {
39            receiver,
40            producer_task_handle,
41            current: Buffer::default(),
42        })
43    }
44
45    fn get_open_state(&mut self) -> Option<&mut OpenReaderState> {
46        match self {
47            Self::Open(state) => Some(state),
48            Self::Finished => None,
49        }
50    }
51
52    fn finish(&mut self) -> std::io::Result<()> {
53        let Self::Open(state) = std::mem::replace(self, Self::Finished) else {
54            return Ok(());
55        };
56
57        drop(state.receiver);
58
59        ASYNC.block_in_place_on(state.producer_task_handle)?
60    }
61
62    /// Returns the next fetched buffer without copying, or an empty buffer at EOF.
63    pub fn next_buffer(&mut self) -> std::io::Result<Buffer<u8>> {
64        self.fill_buf()?;
65        Ok(self
66            .get_open_state()
67            .map(|state| std::mem::take(&mut state.current))
68            .unwrap_or_default())
69    }
70}
71
72#[cfg(feature = "async")]
73impl std::io::Read for StreamBufReader {
74    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
75        let remaining = self.fill_buf()?;
76        if remaining.is_empty() {
77            return Ok(0);
78        }
79        let n = buf.len().min(remaining.len());
80        buf[..n].copy_from_slice(&remaining[..n]);
81        self.consume(n);
82        Ok(n)
83    }
84}
85
86#[cfg(feature = "async")]
87impl std::io::BufRead for StreamBufReader {
88    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
89        let Some(state) = self.get_open_state() else {
90            return Ok(&[]);
91        };
92
93        if state.current.is_empty() {
94            match state.receiver.blocking_recv() {
95                Some((handle, _permit)) => {
96                    let fetched_bytes = ASYNC.block_in_place_on(handle).unwrap()?;
97                    state.current = fetched_bytes;
98                },
99                None => {
100                    self.finish()?;
101                    return Ok(&[]);
102                },
103            }
104        }
105
106        let Some(state) = self.get_open_state() else {
107            unreachable!();
108        };
109
110        Ok(state.current.as_ref())
111    }
112
113    fn consume(&mut self, amt: usize) {
114        if let Some(state) = self.get_open_state() {
115            state.current.slice_in_place(amt..);
116        }
117    }
118}
119
120// Supported reader sources for respectively from_memory and streaming.
121pub enum ReaderSource {
122    Memory(Cursor<Buffer<u8>>),
123    #[cfg(feature = "async")]
124    Streaming(StreamBufReader),
125}
126
127impl std::io::Read for ReaderSource {
128    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
129        match self {
130            Self::Memory(r) => r.read(buf),
131            #[cfg(feature = "async")]
132            Self::Streaming(r) => r.read(buf),
133        }
134    }
135}
136
137impl std::io::BufRead for ReaderSource {
138    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
139        match self {
140            Self::Memory(r) => r.fill_buf(),
141            #[cfg(feature = "async")]
142            Self::Streaming(r) => r.fill_buf(),
143        }
144    }
145
146    fn consume(&mut self, amt: usize) {
147        match self {
148            Self::Memory(r) => r.consume(amt),
149            #[cfg(feature = "async")]
150            Self::Streaming(r) => r.consume(amt),
151        }
152    }
153}
polars_io/utils/stream_buf_reader.rs

polars_io/utils/
stream_buf_reader.rs
