polars-stream 0.53.0

Private crate for the streaming execution engine for the Polars DataFrame library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

use crate::async_primitives::connector;
use crate::async_primitives::morsel_linearizer::{MorselInserter, MorselLinearizer};
use crate::async_primitives::wait_group::WaitGroup;
use crate::morsel::Morsel;

/// This is just a generic enum that can receive from either single or multiple senders.
pub enum FileReaderOutputRecv {
    Connector(connector::Receiver<Morsel>),
    Linearized(MorselLinearizer),
}

/// Wraps either a connector or a linearizer inserter. Also attaches / waits for wait tokens
/// automatically.
pub enum FileReaderOutputSend {
    Connector(connector::Sender<Morsel>, WaitGroup),
    Linearized(MorselInserter, WaitGroup),
}

impl FileReaderOutputRecv {
    pub async fn recv(&mut self) -> Result<Morsel, ()> {
        use FileReaderOutputRecv::*;
        match self {
            Connector(v) => v.recv().await,
            Linearized(v) => v.get().await.ok_or(()),
        }
    }
}

impl FileReaderOutputSend {
    pub fn new_serial() -> (FileReaderOutputSend, FileReaderOutputRecv) {
        let (tx, rx) = connector::connector();
        (
            FileReaderOutputSend::Connector(tx, WaitGroup::default()),
            FileReaderOutputRecv::Connector(rx),
        )
    }

    pub fn new_parallel(num_pipelines: usize) -> (Vec<FileReaderOutputSend>, FileReaderOutputRecv) {
        let (lin, inserters) = MorselLinearizer::new(num_pipelines, 1);
        (
            inserters
                .into_iter()
                .map(|tx| FileReaderOutputSend::Linearized(tx, WaitGroup::default()))
                .collect(),
            FileReaderOutputRecv::Linearized(lin),
        )
    }

    pub async fn send_morsel(&mut self, morsel: Morsel) -> Result<(), Morsel> {
        use FileReaderOutputSend::*;

        // We order to wait first, then send. This is intended to allow the producer create the
        // next morsel while waiting for the current one to be consumed.

        match self {
            Connector(tx, wait_group) => {
                wait_group.wait().await;
                tx.send(morsel).await
            },
            Linearized(tx, wait_group) => {
                wait_group.wait().await;
                tx.insert(morsel).await
            },
        }
    }
}