connectorx/
arrow_batch_iter.rs

1use crate::errors::OutResult;
2use crate::prelude::*;
3use arrow::record_batch::RecordBatch;
4use itertools::Itertools;
5use log::debug;
6use rayon::prelude::*;
7use std::marker::PhantomData;
8
9pub fn set_global_num_thread(num: usize) {
10    rayon::ThreadPoolBuilder::new()
11        .num_threads(num)
12        .build_global()
13        .unwrap();
14}
15
16/// The iterator that returns arrow in `RecordBatch`
17pub struct ArrowBatchIter<S, TP>
18where
19    S: Source,
20    TP: Transport<
21        TSS = S::TypeSystem,
22        TSD = ArrowStreamTypeSystem,
23        S = S,
24        D = ArrowStreamDestination,
25    >,
26    <S as Source>::Partition: 'static,
27    <S as Source>::TypeSystem: 'static,
28    <TP as Transport>::Error: 'static,
29{
30    dst: ArrowStreamDestination,
31    dst_parts: Option<Vec<ArrowStreamPartitionWriter>>,
32    src_parts: Option<Vec<S::Partition>>,
33    dorder: DataOrder,
34    src_schema: Vec<S::TypeSystem>,
35    dst_schema: Vec<ArrowStreamTypeSystem>,
36    handle: Option<std::thread::JoinHandle<()>>,
37    _phantom: PhantomData<TP>,
38}
39
40impl<S, TP> ArrowBatchIter<S, TP>
41where
42    S: Source,
43    TP: Transport<
44        TSS = S::TypeSystem,
45        TSD = ArrowStreamTypeSystem,
46        S = S,
47        D = ArrowStreamDestination,
48    >,
49{
50    pub fn new(
51        src: S,
52        mut dst: ArrowStreamDestination,
53        origin_query: Option<String>,
54        queries: &[CXQuery<String>],
55    ) -> Result<Self, TP::Error> {
56        let dispatcher = Dispatcher::<_, _, TP>::new(src, &mut dst, queries, origin_query);
57        let (dorder, src_parts, dst_parts, src_schema, dst_schema) = dispatcher.prepare()?;
58
59        Ok(Self {
60            dst,
61            dst_parts: Some(dst_parts),
62            src_parts: Some(src_parts),
63            dorder,
64            src_schema,
65            dst_schema,
66            handle: None,
67            _phantom: PhantomData,
68        })
69    }
70
71    fn run(&mut self)
72    where
73        ConnectorXOutError: From<TP::Error>,
74    {
75        let src_schema = self.src_schema.clone();
76        let dst_schema = self.dst_schema.clone();
77        let src_partitions = self.src_parts.take().unwrap();
78        let dst_partitions = self.dst_parts.take().unwrap();
79        let dorder = self.dorder;
80
81        self.handle = Some(std::thread::spawn(move || {
82            let schemas: Vec<_> = src_schema
83                .iter()
84                .zip_eq(&dst_schema)
85                .map(|(&src_ty, &dst_ty)| (src_ty, dst_ty))
86                .collect();
87
88            debug!("Start writing");
89            // parse and write
90            //
91            // Each partition reports its own error through the record batch
92            // channel (see `ArrowPartitionWriter::send_error`), so the consumer
93            // learns about it without joining this thread. `Err(())` only tells
94            // rayon to stop scheduling the remaining partitions.
95            let _ = dst_partitions
96                .into_par_iter()
97                .zip_eq(src_partitions)
98                .enumerate()
99                .try_for_each(|(i, (mut dst, mut src))| -> Result<(), ()> {
100                    let mut parser = match src.parser() {
101                        Ok(parser) => parser,
102                        Err(e) => {
103                            dst.send_error(TP::Error::from(e).into());
104                            return Err(());
105                        }
106                    };
107
108                    let result = (|| -> Result<(), TP::Error> {
109                        match dorder {
110                            DataOrder::RowMajor => loop {
111                                let (n, is_last) = parser.fetch_next()?;
112                                dst.aquire_row(n)?;
113                                for _ in 0..n {
114                                    #[allow(clippy::needless_range_loop)]
115                                    for col in 0..dst.ncols() {
116                                        {
117                                            let (s1, s2) = schemas[col];
118                                            TP::process(s1, s2, &mut parser, &mut dst)?;
119                                        }
120                                    }
121                                }
122                                if is_last {
123                                    break;
124                                }
125                            },
126                            DataOrder::ColumnMajor => loop {
127                                let (n, is_last) = parser.fetch_next()?;
128                                dst.aquire_row(n)?;
129                                #[allow(clippy::needless_range_loop)]
130                                for col in 0..dst.ncols() {
131                                    for _ in 0..n {
132                                        {
133                                            let (s1, s2) = schemas[col];
134                                            TP::process(s1, s2, &mut parser, &mut dst)?;
135                                        }
136                                    }
137                                }
138                                if is_last {
139                                    break;
140                                }
141                            },
142                        }
143
144                        debug!("Finalize partition {}", i);
145                        dst.finalize()?;
146                        debug!("Partition {} finished", i);
147                        Ok(())
148                    })();
149
150                    match result {
151                        Ok(()) => Ok(()),
152                        Err(e) => {
153                            debug!("Partition {} failed: {:?}", i, e);
154                            // `parser` is still alive here on purpose; see
155                            // `ArrowPartitionWriter::send_error`.
156                            dst.send_error(e.into());
157                            Err(())
158                        }
159                    }
160                });
161
162            debug!("Writing finished");
163        }));
164    }
165}
166
167impl<S, TP> Iterator for ArrowBatchIter<S, TP>
168where
169    S: Source,
170    TP: Transport<
171        TSS = S::TypeSystem,
172        TSD = ArrowStreamTypeSystem,
173        S = S,
174        D = ArrowStreamDestination,
175    >,
176    ConnectorXOutError: From<TP::Error>,
177{
178    type Item = OutResult<RecordBatch>;
179
180    /// Yields the next batch, or the first error raised by a partition.
181    ///
182    /// The error is returned as soon as the failing partition reports it,
183    /// without waiting for the producer to shut down. Calling `next` again
184    /// after an error joins the producer and returns `None`. That call blocks
185    /// until the other partitions have stopped at their next batch boundary
186    /// and every source has been torn down, which for MySQL includes draining
187    /// the rest of the result set.
188    ///
189    /// NOTE: not thread safe
190    fn next(&mut self) -> Option<Self::Item> {
191        match self.dst.record_batch() {
192            Some(msg) => Some(msg),
193            None => {
194                // On stream end we must join the producer;
195                // a detached handle would swallow producer panics
196                if let Some(handle) = self.handle.take() {
197                    if let Err(payload) = handle.join() {
198                        std::panic::resume_unwind(payload);
199                    }
200                }
201                None
202            }
203        }
204    }
205}
206
207pub trait RecordBatchIterator: Send {
208    fn get_schema(&self) -> (RecordBatch, &[String]);
209    fn prepare(&mut self);
210    /// See [`ArrowBatchIter::next`] for the error semantics.
211    fn next_batch(&mut self) -> Option<OutResult<RecordBatch>>;
212}
213
214impl<S, TP> RecordBatchIterator for ArrowBatchIter<S, TP>
215where
216    S: Source,
217    TP: Transport<
218            TSS = S::TypeSystem,
219            TSD = ArrowStreamTypeSystem,
220            S = S,
221            D = ArrowStreamDestination,
222        > + std::marker::Send,
223    ConnectorXOutError: From<TP::Error>,
224{
225    fn get_schema(&self) -> (RecordBatch, &[String]) {
226        (self.dst.empty_batch(), self.dst.names())
227    }
228
229    fn prepare(&mut self) {
230        self.run();
231    }
232
233    fn next_batch(&mut self) -> Option<OutResult<RecordBatch>> {
234        self.next()
235    }
236}
connectorx/arrow_batch_iter.rs

connectorx/
arrow_batch_iter.rs
