datui_lib/formats/
ipc_stream.rs

1//! Arrow IPC streams, the format Hugging Face `datasets` caches in: the IPC file format
2//! without the `ARROW1` magic and footer, so Polars cannot scan it. An open converts the
3//! stream (or a directory's stream shards) once into a temp IPC file and scans that; IPC
4//! files among the shards are scanned in place ([`Part`]). One record batch is held at a
5//! time: 0.2 GiB peak for a 3.0 GiB stream that read eagerly took 3.6 GiB.
6
7use std::fs::File;
8use std::io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11use std::sync::atomic::{AtomicU64, Ordering};
12
13use color_eyre::{Result, eyre::eyre};
14use polars_arrow::io::ipc::format::ipc::planus::ReadAsRoot;
15use polars_arrow::io::ipc::format::ipc::{MessageHeaderRef, MessageRef};
16use polars_arrow::io::ipc::read::{StreamReader, StreamState, read_stream_metadata};
17use polars_arrow::io::ipc::write::{FileWriter, WriteOptions};
18
19use crate::cloud::download::TempDownload;
20use crate::error_display::{FileError, user_message_from_io};
21use crate::loading::unfinished::Writer;
22
23/// What a stream's messages start with since Arrow 0.15. Older streams start with the
24/// schema message's length.
25const CONTINUATION: [u8; 4] = [0xff; 4];
26
27/// The longest schema message read to tell a stream from other bytes. A schema with
28/// thousands of columns and its metadata fits well inside.
29const MAX_SCHEMA: usize = 16 << 20;
30
31/// Whether `head` begins an Arrow IPC stream: a schema message, with or (older streams)
32/// without the continuation marker. A whole message must be a schema message; a cut one
33/// counts if marked. Fields are not parsed: Polars panics on unimplemented column types,
34/// and this runs on any opened file.
35pub fn is_stream_head(head: &[u8]) -> bool {
36    let marked = head.starts_with(&CONTINUATION);
37    let rest = if marked { &head[4..] } else { head };
38    let Some(length) = rest.get(..4) else {
39        return false;
40    };
41    let length = i32::from_le_bytes([length[0], length[1], length[2], length[3]]);
42    let Ok(length) = usize::try_from(length) else {
43        return false;
44    };
45    if length == 0 || length > MAX_SCHEMA {
46        return false;
47    }
48    let Some(message) = rest.get(4..4 + length) else {
49        return marked;
50    };
51    begins_schema(message)
52}
53
54/// How much of a long schema message is read to see that it begins like one before the
55/// rest is: the schema table is near the front, its fields and metadata after it.
56const SCHEMA_PREFIX: usize = 64 << 10;
57
58/// Whether the file at `path` is an Arrow IPC stream, by its contents.
59pub fn is_stream_file(path: &Path) -> bool {
60    File::open(path).is_ok_and(is_stream)
61}
62
63/// Whether `source` holds an Arrow IPC stream. A file that only happens to start with
64/// a small number, as many binary files do, is read no further than [`SCHEMA_PREFIX`].
65fn is_stream(mut source: impl Read) -> bool {
66    let mut head = Vec::new();
67    // The marker and the length, then the message the length names.
68    if (&mut source).take(8).read_to_end(&mut head).is_err() {
69        return false;
70    }
71    let at = if head.starts_with(&CONTINUATION) {
72        4
73    } else {
74        0
75    };
76    let Some(length) = head
77        .get(at..at + 4)
78        .map(|b| i32::from_le_bytes([b[0], b[1], b[2], b[3]]))
79        .and_then(|l| usize::try_from(l).ok())
80        .filter(|l| (1..=MAX_SCHEMA).contains(l))
81    else {
82        return false;
83    };
84    let mut read_to = |end: usize, head: &mut Vec<u8>| {
85        let more = end.saturating_sub(head.len()) as u64;
86        (&mut source).take(more).read_to_end(head).is_ok()
87    };
88    let start = at + 4;
89    if length > SCHEMA_PREFIX
90        && !(read_to(start + SCHEMA_PREFIX, &mut head) && begins_schema(&head[start..]))
91    {
92        return false;
93    }
94    read_to(start + length, &mut head) && is_stream_head(&head)
95}
96
97/// Whether `message`, all or the front of one, is a schema message as far as it goes.
98fn begins_schema(message: &[u8]) -> bool {
99    matches!(
100        MessageRef::read_as_root(message).and_then(|m| m.header()),
101        Ok(Some(MessageHeaderRef::Schema(_)))
102    )
103}
104
105/// Whether Arrow `paths` are converted: when the first is a stream (only it is opened).
106/// A stream behind an IPC file is found by [`any_stream`] once the scan fails on it.
107pub fn starts_with_stream(paths: &[PathBuf]) -> bool {
108    paths.first().is_some_and(|p| is_stream_file(p))
109}
110
111/// Whether any of `paths` is a stream: asked once a scan of them as IPC files failed.
112pub fn any_stream(paths: &[PathBuf]) -> bool {
113    paths.iter().any(|p| is_stream_file(p))
114}
115
116/// Where one input's rows are read from once its streams are converted.
117#[derive(Debug, Clone, PartialEq, Eq)]
118pub enum Part {
119    /// An IPC file, scanned where it is: a local path or an object's URL.
120    InPlace(PathBuf),
121    /// A stream, `source`, converted: `rows` rows of the converted file from `offset`.
122    Converted {
123        source: PathBuf,
124        offset: u64,
125        rows: u64,
126    },
127}
128
129/// What a conversion wrote: the IPC file of the streams, and each input's place.
130#[derive(Debug)]
131pub(crate) struct Converted {
132    pub file: TempDownload,
133    pub parts: Vec<Part>,
134}
135
136/// Whether an Arrow file starts as an IPC file does, with its footer at the end.
137pub(crate) fn is_ipc_file_head(head: &[u8]) -> bool {
138    head.starts_with(b"ARROW1")
139}
140
141/// Convert the Arrow streams at `paths`, in order, into one IPC file in `temp_dir`
142/// (system temp if `None`), claimed via `writer` and removed on stop or failure. IPC
143/// files among them are scanned in place. `read` counts bytes looked at. One batch in
144/// memory at a time; LZ4 and ZSTD buffers are written uncompressed so the file maps.
145pub(crate) fn convert(
146    paths: &[PathBuf],
147    temp_dir: Option<&Path>,
148    writer: &Writer,
149    read: &AtomicU64,
150) -> Result<Converted> {
151    let mut merge = Merge::create(temp_dir, writer)?;
152    let mut parts = Vec::with_capacity(paths.len());
153    let mut before = 0;
154    for path in paths {
155        let mut source = File::open(path)?;
156        let size = source.metadata()?.len();
157        let mut head = Vec::with_capacity(6);
158        (&mut source).take(6).read_to_end(&mut head)?;
159        if is_ipc_file_head(&head) {
160            parts.push(Part::InPlace(path.clone()));
161        } else {
162            source.seek(SeekFrom::Start(0))?;
163            has_room(size, temp_dir)?;
164            parts.push(merge.append(source, path, before, read)?);
165        }
166        before += size;
167        read.store(before, Ordering::Relaxed);
168    }
169    Ok(Converted {
170        file: merge.finish()?,
171        parts,
172    })
173}
174
175/// One IPC file being written from the batches of Arrow streams, appended one at a
176/// time: what a conversion writes, and what a bucket's streams are read into as they
177/// download. Dropped unfinished, the file goes.
178pub(crate) struct Merge<'a> {
179    /// The writer, the first input's name and its columns, once one is appended. Its
180    /// handle on the file is let go before the file is removed.
181    out: Option<(FileWriter<BufWriter<File>>, PathBuf, Columns)>,
182    /// The file, then its claim: dropped in that order.
183    file: tempfile::NamedTempFile,
184    claim: crate::loading::unfinished::Claim,
185    dir: PathBuf,
186    writer: &'a Writer,
187    /// The rows written so far.
188    rows: u64,
189}
190
191type Columns = Vec<(
192    polars::prelude::PlSmallStr,
193    polars_arrow::datatypes::ArrowDataType,
194)>;
195
196fn columns(schema: &polars_arrow::datatypes::ArrowSchema) -> Columns {
197    schema
198        .iter_values()
199        .map(|f| (f.name.clone(), f.dtype.clone()))
200        .collect()
201}
202
203impl<'a> Merge<'a> {
204    /// The empty file, in `temp_dir`, claimed through `writer`.
205    pub(crate) fn create(temp_dir: Option<&Path>, writer: &'a Writer) -> Result<Self> {
206        let Some((file, claim)) =
207            writer.create(|| TempDownload::create(temp_dir, Some("arrow")))?
208        else {
209            return Err(stopped());
210        };
211        Ok(Self {
212            file,
213            claim,
214            dir: temp_dir
215                .map(Path::to_path_buf)
216                .unwrap_or_else(std::env::temp_dir),
217            writer,
218            out: None,
219            rows: 0,
220        })
221    }
222
223    /// Append the batches of the stream `source` reads, which errors call `name`,
224    /// counting its bytes into `read` after the `before` read ahead of it. Its place
225    /// in the file is returned.
226    pub(crate) fn append(
227        &mut self,
228        source: impl Read,
229        name: &Path,
230        before: u64,
231        read: &AtomicU64,
232    ) -> Result<Part> {
233        let offset = self.rows;
234        match self.batches(source, name, before, read) {
235            Ok(true) => Ok(Part::Converted {
236                source: name.to_path_buf(),
237                offset,
238                rows: self.rows - offset,
239            }),
240            Ok(false) => Err(stopped()),
241            Err(e) => Err(out_of_room(e, &self.dir)),
242        }
243    }
244
245    /// The finished file, held with its claim.
246    pub(crate) fn finish(self) -> Result<TempDownload> {
247        let Some((mut out, _, _)) = self.out else {
248            return Err(eyre!("No Arrow IPC stream to read."));
249        };
250        let finished = out
251            .finish()
252            .map_err(color_eyre::Report::from)
253            .and_then(|()| Ok(out.into_inner().flush()?));
254        if let Err(e) = finished {
255            return Err(out_of_room(e, &self.dir));
256        }
257        Ok(TempDownload::held(self.file, Some(self.claim)))
258    }
259
260    /// `false` when the open was stopped first.
261    fn batches(
262        &mut self,
263        source: impl Read,
264        name: &Path,
265        before: u64,
266        read: &AtomicU64,
267    ) -> Result<bool> {
268        let mut reader = Forward {
269            inner: BufReader::with_capacity(
270                1 << 20,
271                Counting {
272                    inner: source,
273                    at: 0,
274                    before,
275                    read,
276                },
277            ),
278            at: 0,
279        };
280        let unreadable = move |e: &dyn std::fmt::Display| -> color_eyre::Report {
281            FileError::new(name, format!("not a readable Arrow IPC stream: {e}")).into()
282        };
283        // A read that failed under the stream (a download cut off) says so, not that
284        // the stream is damaged.
285        let failed = move |e: polars::prelude::PolarsError| match e {
286            polars::prelude::PolarsError::IO { error, .. } => {
287                FileError::new(name, user_message_from_io(&error, None)).into()
288            }
289            e => unreadable(&e),
290        };
291        // Polars panics on a column type it has not implemented, such as run-end
292        // encoding: that is a stream it cannot read, not a crash.
293        let metadata = crate::logging::catch_panic(|| read_stream_metadata(&mut reader))
294            .map_err(|_| unreadable(&"it has a column type Polars cannot read"))?
295            .map_err(failed)?;
296        self.start(
297            name,
298            &metadata.schema,
299            &metadata.ipc_schema.fields,
300            metadata.custom_schema_metadata.as_ref(),
301        )?;
302        let mut batches = StreamReader::new(reader, metadata, None);
303        let (out, _, _) = self.out.as_mut().expect("started just above");
304        loop {
305            if self.writer.stopped() {
306                return Ok(false);
307            }
308            // Polars also panics on some malformed record batches, rather than erring.
309            let next = crate::logging::catch_panic(|| batches.next())
310                .map_err(|_| unreadable(&"a record batch in it is damaged"))?;
311            let batch = match next {
312                Some(Ok(StreamState::Some(batch))) => batch,
313                // The end of a stream written without its end-of-stream marker.
314                Some(Ok(StreamState::Waiting)) | None => break,
315                Some(Err(e)) => return Err(failed(e)),
316            };
317            self.rows += batch.len() as u64;
318            out.write(&batch, None)?;
319        }
320        Ok(true)
321    }
322
323    /// Start the file with the first input's schema, or check a later one has the same
324    /// columns.
325    fn start(
326        &mut self,
327        name: &Path,
328        schema: &polars_arrow::datatypes::ArrowSchema,
329        fields: &[polars_arrow::io::ipc::IpcField],
330        custom: Option<&polars_arrow::datatypes::Metadata>,
331    ) -> Result<()> {
332        match &self.out {
333            None => {
334                let mut out = FileWriter::try_new(
335                    BufWriter::with_capacity(1 << 20, self.file.as_file().try_clone()?),
336                    Arc::new(schema.clone()),
337                    Some(fields.to_vec()),
338                    WriteOptions { compression: None },
339                )?;
340                if let Some(custom) = custom {
341                    out.set_custom_schema_metadata(Arc::new(custom.clone()));
342                }
343                self.out = Some((out, name.to_path_buf(), columns(schema)));
344            }
345            Some((_, first, first_columns)) => {
346                if columns(schema) != *first_columns {
347                    return Err(eyre!(
348                        "{} has different columns from {}, so they cannot be read as one table.",
349                        name.display(),
350                        first.display()
351                    ));
352                }
353            }
354        }
355        Ok(())
356    }
357}
358
359fn stopped() -> color_eyre::Report {
360    eyre!("Converting the Arrow stream was stopped.")
361}
362
363/// What to do about a temp directory too small for the copy.
364const ELSEWHERE: &str = "Choose another place with --temp-dir or the temp_dir setting.";
365
366/// Whether `temp_dir` (the system temp directory when `None`) has room for a copy of
367/// `needs` bytes of Arrow; see [`room`].
368pub(crate) fn has_room(needs: u64, temp_dir: Option<&Path>) -> Result<()> {
369    let dir = temp_dir
370        .map(Path::to_path_buf)
371        .unwrap_or_else(std::env::temp_dir);
372    room(needs, crate::cloud::local_copy::free_space(&dir), &dir)
373}
374
375/// The copy is about the size of the streams, larger where their buffers are
376/// compressed: refused before any of it is written where `dir` has less free.
377fn room(needs: u64, free: Option<u64>, dir: &Path) -> Result<()> {
378    match free {
379        Some(free) if free < needs => Err(eyre!(
380            "Converting the Arrow stream needs {} free in {}, which has {}. {ELSEWHERE}",
381            crate::numfmt::bytes(needs),
382            dir.display(),
383            crate::numfmt::bytes(free),
384        )),
385        _ => Ok(()),
386    }
387}
388
389/// A write that ran out of space, said so with where and what to do.
390fn out_of_room(error: color_eyre::Report, dir: &Path) -> color_eyre::Report {
391    let full = error.chain().any(|cause| {
392        cause
393            .downcast_ref::<std::io::Error>()
394            .is_some_and(|e| e.kind() == std::io::ErrorKind::StorageFull)
395    });
396    if full {
397        eyre!(
398            "{} ran out of space for the converted Arrow stream, which is written uncompressed. {ELSEWHERE}",
399            dir.display()
400        )
401    } else {
402        error
403    }
404}
405
406/// A stream being read, counting its bytes into the open's progress.
407struct Counting<'a, R> {
408    inner: R,
409    at: u64,
410    /// The bytes of the inputs before this one.
411    before: u64,
412    read: &'a AtomicU64,
413}
414
415impl<R: Read> Read for Counting<'_, R> {
416    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
417        let n = self.inner.read(buf)?;
418        self.at += n as u64;
419        self.read.store(self.before + self.at, Ordering::Relaxed);
420        Ok(n)
421    }
422}
423
424/// A stream read front to back, by the stream reader that seeks: only forward, over
425/// the padding after a message, so a download can be read as it arrives.
426struct Forward<R> {
427    inner: R,
428    at: u64,
429}
430
431impl<R: BufRead> Read for Forward<R> {
432    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
433        let n = self.inner.read(buf)?;
434        self.at += n as u64;
435        Ok(n)
436    }
437}
438
439impl<R: BufRead> Seek for Forward<R> {
440    fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
441        let to = match pos {
442            SeekFrom::Start(to) => Some(to),
443            SeekFrom::Current(by) => self.at.checked_add_signed(by),
444            SeekFrom::End(_) => None,
445        };
446        let skip = to
447            .and_then(|to| to.checked_sub(self.at))
448            .ok_or_else(|| std::io::Error::other("an Arrow stream is read front to back"))?;
449        let skipped = std::io::copy(&mut (&mut self.inner).take(skip), &mut std::io::sink())?;
450        self.at += skipped;
451        Ok(self.at)
452    }
453}
454
455#[cfg(test)]
456pub(crate) mod tests;
datui_lib/formats/ipc_stream.rs

datui_lib/formats/
ipc_stream.rs
