polars_io/
scan_lines.rs

1use polars_arrow::array::{
2    BINVIEW_ARROW_BUFFER_LEN_LIMIT, BINVIEW_MAX_ROW_BYTE_LEN, Utf8ViewArray, View,
3};
4use polars_arrow::datatypes::ArrowDataType;
5use polars_buffer::Buffer;
6use polars_core::prelude::DataType;
7use polars_core::series::Series;
8use polars_error::{PolarsResult, polars_bail, polars_ensure};
9use polars_utils::pl_str::PlSmallStr;
10
11const CR: u8 = b'\r';
12const LF: u8 = b'\n';
13const BUFFER_SPLIT_THRESHOLD: usize = 4096;
14
15pub fn count_lines(full_bytes: &[u8]) -> usize {
16    let mut n: usize = full_bytes.iter().map(|c| (*c == LF) as usize).sum();
17
18    if let Some(c) = full_bytes.last()
19        && *c != LF
20    {
21        n += 1;
22    }
23
24    n
25}
26
27pub fn split_lines_to_rows(bytes: Buffer<u8>) -> PolarsResult<Series> {
28    split_lines_to_rows_impl(bytes, BINVIEW_MAX_ROW_BYTE_LEN)
29}
30
31fn split_lines_to_rows_impl(bytes: Buffer<u8>, max_row_size: usize) -> PolarsResult<Series> {
32    if bytes.is_empty() {
33        return Ok(Series::new_empty(PlSmallStr::EMPTY, &DataType::String));
34    };
35
36    if simdutf8::basic::from_utf8(&bytes).is_err() {
37        polars_bail!(ComputeError: "invalid utf8")
38    }
39
40    let first_line_len = memchr::memchr(LF, &bytes).unwrap_or(bytes.len());
41    let last_line_len = memchr::memrchr(LF, &bytes).map_or(bytes.len(), |i| bytes.len() - 1 - i);
42
43    let n_lines_estimate = bytes
44        .len()
45        .div_ceil(first_line_len.min(last_line_len).max(1));
46
47    let mut views: Vec<View> = Vec::with_capacity(n_lines_estimate);
48    let mut data_buffers: Vec<Buffer<u8>> = Vec::new();
49    let mut total_bytes_len: usize = 0;
50    let mut total_buffer_len: usize = 0;
51    let mut active_buffer: Option<(usize, usize)> = None;
52
53    let bytes = if bytes.last() == Some(&LF) {
54        let len = bytes.len();
55        bytes.sliced(..len - 1)
56    } else {
57        bytes
58    };
59
60    let slice: &[u8] = &bytes;
61    let mut line_start: usize = 0;
62
63    for line_end in memchr::memchr_iter(LF, slice)
64        .map(|i| i + 1)
65        .chain(std::iter::once(slice.len() + 1))
66    {
67        let start = line_start;
68        let mut end = line_end - 1;
69        line_start = line_end;
70
71        if end > start && slice[end - 1] == CR {
72            end -= 1;
73        }
74
75        let len = end - start;
76
77        polars_ensure!(
78            len <= max_row_size,
79            ComputeError:
80            "line byte length {} exceeds max row byte length {}",
81            len, max_row_size,
82        );
83
84        total_bytes_len += len;
85
86        let line_bytes = unsafe { slice.get_unchecked(start..end) };
87
88        let view = if len <= View::MAX_INLINE_SIZE as usize {
89            unsafe { View::new_inline_unchecked(line_bytes) }
90        } else {
91            // Note: `start > buffer_end`, there is always at least a line terminator in between.
92            if let Some((buffer_start, buffer_end)) = active_buffer
93                && (end - buffer_start > BINVIEW_ARROW_BUFFER_LEN_LIMIT
94                    || start - buffer_end > BUFFER_SPLIT_THRESHOLD)
95            {
96                total_buffer_len += buffer_end - buffer_start;
97                data_buffers.push(bytes.clone().sliced(buffer_start..buffer_end));
98                active_buffer = None;
99            }
100
101            let buffer_start = active_buffer.map_or(start, |(buffer_start, _)| buffer_start);
102            active_buffer = Some((buffer_start, end));
103
104            unsafe {
105                View::new_noninline_unchecked(
106                    line_bytes,
107                    data_buffers.len() as u32,
108                    (start - buffer_start) as u32,
109                )
110            }
111        };
112
113        views.push(view);
114    }
115
116    if let Some((buffer_start, buffer_end)) = active_buffer {
117        total_buffer_len += buffer_end - buffer_start;
118        data_buffers.push(bytes.sliced(buffer_start..buffer_end));
119    }
120
121    let arr = unsafe {
122        Utf8ViewArray::new_unchecked(
123            ArrowDataType::Utf8View,
124            views.into(),
125            data_buffers.into(),
126            None,
127            Some(total_bytes_len),
128            total_buffer_len,
129        )
130    };
131
132    Ok(unsafe {
133        Series::_try_from_arrow_unchecked(
134            PlSmallStr::EMPTY,
135            vec![arr.boxed()],
136            &ArrowDataType::Utf8View,
137        )?
138    })
139}
140
141#[cfg(test)]
142mod tests {
143    use polars_buffer::Buffer;
144    use polars_error::PolarsError;
145
146    use crate::scan_lines::split_lines_to_rows_impl;
147
148    #[test]
149    fn test_split_lines_to_rows_impl() {
150        let data: &'static [u8] = b"
151AAAAABBBBBCCCCCDDDDD
152
153EEEEEFFFFFGGGGGHHHHH
154
155";
156
157        let out = split_lines_to_rows_impl(Buffer::from_static(data), 20).unwrap();
158        let out = out.str().unwrap();
159
160        assert_eq!(
161            out.iter().collect::<Vec<_>>().as_slice(),
162            &[
163                Some(""),
164                Some("AAAAABBBBBCCCCCDDDDD"),
165                Some(""),
166                Some("EEEEEFFFFFGGGGGHHHHH"),
167                Some(""),
168            ]
169        );
170
171        let v: Vec<&[Buffer<u8>]> = out
172            .downcast_iter()
173            .map(|array| array.data_buffers().as_ref())
174            .collect();
175
176        assert_eq!(
177            v.as_slice(),
178            &[&[Buffer::from_static(
179                b"AAAAABBBBBCCCCCDDDDD\n\nEEEEEFFFFFGGGGGHHHHH"
180            )]]
181        );
182
183        let PolarsError::ComputeError(err_str) =
184            split_lines_to_rows_impl(Buffer::from_static(data), 19).unwrap_err()
185        else {
186            unreachable!()
187        };
188
189        assert_eq!(
190            &*err_str,
191            "line byte length 20 exceeds max row byte length 19"
192        );
193    }
194
195    #[test]
196    fn test_split_lines_to_rows_impl_all_inline() {
197        let data: Vec<u8> = [
198            b"AAAABBBBCCCC\n".as_slice(),
199            b"            \n".as_slice(),
200            b"DDDDEEEEFFFF\n".as_slice(),
201            b"            ".as_slice(),
202        ]
203        .concat();
204
205        let out = split_lines_to_rows_impl(Buffer::from_vec(data), 12).unwrap();
206        let out = out.str().unwrap();
207
208        assert_eq!(
209            out.iter().collect::<Vec<_>>().as_slice(),
210            &[
211                Some("AAAABBBBCCCC"),
212                Some("            "),
213                Some("DDDDEEEEFFFF"),
214                Some("            "),
215            ]
216        );
217
218        let v: Vec<&[Buffer<u8>]> = out
219            .downcast_iter()
220            .map(|array| array.data_buffers().as_ref())
221            .collect();
222
223        assert_eq!(v.as_slice(), &[&[][..]]);
224    }
225}
polars_io/scan_lines.rs

polars_io/
scan_lines.rs
