connectorx/
arrow_batch_iter.rs

1use crate::prelude::*;
2use arrow::record_batch::RecordBatch;
3use itertools::Itertools;
4use log::debug;
5use rayon::prelude::*;
6use std::marker::PhantomData;
7
8pub fn set_global_num_thread(num: usize) {
9    rayon::ThreadPoolBuilder::new()
10        .num_threads(num)
11        .build_global()
12        .unwrap();
13}
14
15/// The iterator that returns arrow in `RecordBatch`
16pub struct ArrowBatchIter<S, TP>
17where
18    S: Source,
19    TP: Transport<
20        TSS = S::TypeSystem,
21        TSD = ArrowStreamTypeSystem,
22        S = S,
23        D = ArrowStreamDestination,
24    >,
25    <S as Source>::Partition: 'static,
26    <S as Source>::TypeSystem: 'static,
27    <TP as Transport>::Error: 'static,
28{
29    dst: ArrowStreamDestination,
30    dst_parts: Option<Vec<ArrowStreamPartitionWriter>>,
31    src_parts: Option<Vec<S::Partition>>,
32    dorder: DataOrder,
33    src_schema: Vec<S::TypeSystem>,
34    dst_schema: Vec<ArrowStreamTypeSystem>,
35    handle: Option<std::thread::JoinHandle<Result<(), TP::Error>>>,
36    _phantom: PhantomData<TP>,
37}
38
39impl<'a, S, TP> ArrowBatchIter<S, TP>
40where
41    S: Source + 'a,
42    TP: Transport<
43        TSS = S::TypeSystem,
44        TSD = ArrowStreamTypeSystem,
45        S = S,
46        D = ArrowStreamDestination,
47    >,
48{
49    pub fn new(
50        src: S,
51        mut dst: ArrowStreamDestination,
52        origin_query: Option<String>,
53        queries: &[CXQuery<String>],
54    ) -> Result<Self, TP::Error> {
55        let dispatcher = Dispatcher::<_, _, TP>::new(src, &mut dst, queries, origin_query);
56        let (dorder, src_parts, dst_parts, src_schema, dst_schema) = dispatcher.prepare()?;
57
58        Ok(Self {
59            dst,
60            dst_parts: Some(dst_parts),
61            src_parts: Some(src_parts),
62            dorder,
63            src_schema,
64            dst_schema,
65            handle: None,
66            _phantom: PhantomData,
67        })
68    }
69
70    fn run(&mut self) {
71        let src_schema = self.src_schema.clone();
72        let dst_schema = self.dst_schema.clone();
73        let src_partitions = self.src_parts.take().unwrap();
74        let dst_partitions = self.dst_parts.take().unwrap();
75        let dorder = self.dorder;
76
77        self.handle = Some(std::thread::spawn(move || -> Result<(), TP::Error> {
78            let schemas: Vec<_> = src_schema
79                .iter()
80                .zip_eq(&dst_schema)
81                .map(|(&src_ty, &dst_ty)| (src_ty, dst_ty))
82                .collect();
83
84            debug!("Start writing");
85            // parse and write
86            dst_partitions
87                .into_par_iter()
88                .zip_eq(src_partitions)
89                .enumerate()
90                .try_for_each(|(i, (mut dst, mut src))| -> Result<(), TP::Error> {
91                    let mut parser = src.parser()?;
92
93                    match dorder {
94                        DataOrder::RowMajor => loop {
95                            let (n, is_last) = parser.fetch_next()?;
96                            dst.aquire_row(n)?;
97                            for _ in 0..n {
98                                #[allow(clippy::needless_range_loop)]
99                                for col in 0..dst.ncols() {
100                                    {
101                                        let (s1, s2) = schemas[col];
102                                        TP::process(s1, s2, &mut parser, &mut dst)?;
103                                    }
104                                }
105                            }
106                            if is_last {
107                                break;
108                            }
109                        },
110                        DataOrder::ColumnMajor => loop {
111                            let (n, is_last) = parser.fetch_next()?;
112                            dst.aquire_row(n)?;
113                            #[allow(clippy::needless_range_loop)]
114                            for col in 0..dst.ncols() {
115                                for _ in 0..n {
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                    }
127
128                    debug!("Finalize partition {}", i);
129                    dst.finalize()?;
130                    debug!("Partition {} finished", i);
131                    Ok(())
132                })?;
133
134            debug!("Writing finished");
135
136            Ok(())
137        }));
138    }
139}
140
141impl<'a, S, TP> Iterator for ArrowBatchIter<S, TP>
142where
143    S: Source + 'a,
144    TP: Transport<
145        TSS = S::TypeSystem,
146        TSD = ArrowStreamTypeSystem,
147        S = S,
148        D = ArrowStreamDestination,
149    >,
150{
151    type Item = RecordBatch;
152    /// NOTE: not thread safe
153    fn next(&mut self) -> Option<Self::Item> {
154        match self.dst.record_batch() {
155            Ok(Some(rb)) => Some(rb),
156            Ok(None) | Err(_) => {
157                // On stream end we must join the producer;
158                // a detached handle would swallow producer panics
159                if let Some(handle) = self.handle.take() {
160                    match handle.join() {
161                        Ok(Ok(())) => {}
162                        Ok(Err(e)) => panic!("cx writer failed: {:?}", e),
163                        Err(payload) => std::panic::resume_unwind(payload),
164                    }
165                }
166                None
167            }
168        }
169    }
170}
171
172pub trait RecordBatchIterator: Send {
173    fn get_schema(&self) -> (RecordBatch, &[String]);
174    fn prepare(&mut self);
175    fn next_batch(&mut self) -> Option<RecordBatch>;
176}
177
178impl<'a, S, TP> RecordBatchIterator for ArrowBatchIter<S, TP>
179where
180    S: Source + 'a,
181    TP: Transport<
182            TSS = S::TypeSystem,
183            TSD = ArrowStreamTypeSystem,
184            S = S,
185            D = ArrowStreamDestination,
186        > + std::marker::Send,
187{
188    fn get_schema(&self) -> (RecordBatch, &[String]) {
189        (self.dst.empty_batch(), self.dst.names())
190    }
191
192    fn prepare(&mut self) {
193        self.run();
194    }
195
196    fn next_batch(&mut self) -> Option<RecordBatch> {
197        self.next()
198    }
199}
connectorx/arrow_batch_iter.rs

connectorx/
arrow_batch_iter.rs
