connectorx/
dispatcher.rs

1//! This module provides [`dispatcher::Dispatcher`], the core struct in ConnectorX
2//! that drives the data loading from a source to a destination.
3use crate::{
4    data_order::{coordinate, DataOrder},
5    destinations::{Destination, DestinationPartition},
6    errors::{ConnectorXError, Result as CXResult},
7    sources::{PartitionParser, Source, SourcePartition},
8    sql::CXQuery,
9    typesystem::Transport,
10};
11use itertools::Itertools;
12use log::debug;
13use rayon::prelude::*;
14use std::marker::PhantomData;
15
16/// What [`Dispatcher::prepare`] returns: the negotiated data order, the source and
17/// destination partitions, and the source and destination schemas.
18pub type Prepared<'w, S, D> = (
19    DataOrder,
20    Vec<<S as Source>::Partition>,
21    Vec<<D as Destination>::Partition<'w>>,
22    Vec<<S as Source>::TypeSystem>,
23    Vec<<D as Destination>::TypeSystem>,
24);
25
26/// A dispatcher takes a `S: Source`, a `D: Destination`, a `TP: Transport` and a vector of `queries` as input to
27/// load data from `S` to `D` using the queries.
28pub struct Dispatcher<'a, S, D, TP> {
29    src: S,
30    dst: &'a mut D,
31    queries: Vec<CXQuery<String>>,
32    origin_query: Option<String>,
33    _phantom: PhantomData<TP>,
34}
35
36impl<'w, S, D, TP> Dispatcher<'w, S, D, TP>
37where
38    S: Source,
39    D: Destination,
40    TP: Transport<TSS = S::TypeSystem, TSD = D::TypeSystem, S = S, D = D>,
41{
42    /// Create a new dispatcher by providing a source, a destination and the queries.
43    pub fn new<Q>(src: S, dst: &'w mut D, queries: &[Q], origin_query: Option<String>) -> Self
44    where
45        for<'a> &'a Q: Into<CXQuery>,
46    {
47        Self {
48            src,
49            dst,
50            queries: queries.iter().map(Into::into).collect(),
51            origin_query,
52            _phantom: PhantomData,
53        }
54    }
55
56    pub fn set_pre_execution_queries(&mut self, pre_execution_queries: Option<&[String]>) {
57        self.src.set_pre_execution_queries(pre_execution_queries);
58    }
59
60    pub fn prepare(mut self) -> Result<Prepared<'w, S, D>, TP::Error> {
61        debug!("Prepare");
62        let dorder = coordinate(S::DATA_ORDERS, D::DATA_ORDERS)?;
63        self.src.set_data_order(dorder)?;
64        self.src.set_queries(self.queries.as_slice());
65        self.src.set_origin_query(self.origin_query);
66
67        debug!("Fetching metadata");
68        self.src.fetch_metadata()?;
69        let src_schema = self.src.schema();
70        let dst_schema = src_schema
71            .iter()
72            .map(|&s| TP::convert_typesystem(s))
73            .collect::<CXResult<Vec<_>>>()?;
74        let names = self.src.names();
75
76        let mut total_rows = if self.dst.needs_count() {
77            // return None if cannot derive total count
78            debug!("Try get row rounts for entire result");
79            self.src.result_rows()?
80        } else {
81            debug!("Do not need counts in advance");
82            Some(0)
83        };
84        let mut src_partitions: Vec<S::Partition> = self.src.partition()?;
85        if self.dst.needs_count() && total_rows.is_none() {
86            debug!("Manually count rows of each partitioned query and sum up");
87            // run queries
88            src_partitions
89                .par_iter_mut()
90                .try_for_each(|partition| -> Result<(), S::Error> { partition.result_rows() })?;
91
92            // get number of row of each partition from the source
93            let part_rows: Vec<usize> = src_partitions
94                .iter()
95                .map(|partition| partition.nrows())
96                .collect();
97            total_rows = Some(part_rows.iter().sum());
98        }
99        let total_rows = total_rows.ok_or_else(ConnectorXError::CountError)?;
100
101        debug!(
102            "Allocate destination memory: {}x{}",
103            total_rows,
104            src_schema.len()
105        );
106        self.dst.allocate(total_rows, &names, &dst_schema, dorder)?;
107
108        debug!("Create destination partition");
109        let dst_partitions = self.dst.partition(self.queries.len())?;
110
111        Ok((
112            dorder,
113            src_partitions,
114            dst_partitions,
115            src_schema,
116            dst_schema,
117        ))
118    }
119
120    /// Start the data loading process.
121    pub fn run(self) -> Result<(), TP::Error> {
122        debug!("Run dispatcher");
123        let (dorder, src_partitions, dst_partitions, src_schema, dst_schema) = self.prepare()?;
124
125        #[cfg(all(not(feature = "branch"), not(feature = "fptr")))]
126        compile_error!("branch or fptr, pick one");
127
128        #[cfg(feature = "branch")]
129        let schemas: Vec<_> = src_schema
130            .iter()
131            .zip_eq(&dst_schema)
132            .map(|(&src_ty, &dst_ty)| (src_ty, dst_ty))
133            .collect();
134
135        debug!("Start writing");
136        // parse and write
137        dst_partitions
138            .into_par_iter()
139            .zip_eq(src_partitions)
140            .enumerate()
141            .try_for_each(|(i, (mut dst, mut src))| -> Result<(), TP::Error> {
142                #[cfg(feature = "fptr")]
143                let f: Vec<_> = src_schema
144                    .iter()
145                    .zip_eq(&dst_schema)
146                    .map(|(&src_ty, &dst_ty)| TP::processor(src_ty, dst_ty))
147                    .collect::<CXResult<Vec<_>>>()?;
148
149                let mut parser = src.parser()?;
150
151                match dorder {
152                    DataOrder::RowMajor => loop {
153                        let (n, is_last) = parser.fetch_next()?;
154                        dst.aquire_row(n)?;
155                        for _ in 0..n {
156                            #[allow(clippy::needless_range_loop)]
157                            for col in 0..dst.ncols() {
158                                #[cfg(feature = "fptr")]
159                                f[col](&mut parser, &mut dst)?;
160
161                                #[cfg(feature = "branch")]
162                                {
163                                    let (s1, s2) = schemas[col];
164                                    TP::process(s1, s2, &mut parser, &mut dst)?;
165                                }
166                            }
167                        }
168                        if is_last {
169                            break;
170                        }
171                    },
172                    DataOrder::ColumnMajor => loop {
173                        let (n, is_last) = parser.fetch_next()?;
174                        dst.aquire_row(n)?;
175                        #[allow(clippy::needless_range_loop)]
176                        for col in 0..dst.ncols() {
177                            for _ in 0..n {
178                                #[cfg(feature = "fptr")]
179                                f[col](&mut parser, &mut dst)?;
180                                #[cfg(feature = "branch")]
181                                {
182                                    let (s1, s2) = schemas[col];
183                                    TP::process(s1, s2, &mut parser, &mut dst)?;
184                                }
185                            }
186                        }
187                        if is_last {
188                            break;
189                        }
190                    },
191                }
192
193                debug!("Finalize partition {}", i);
194                dst.finalize()?;
195                debug!("Partition {} finished", i);
196                Ok(())
197            })?;
198
199        debug!("Writing finished");
200
201        Ok(())
202    }
203
204    /// Only fetch the metadata (header) of the destination.
205    pub fn get_meta(&mut self) -> Result<(), TP::Error> {
206        let dorder = coordinate(S::DATA_ORDERS, D::DATA_ORDERS)?;
207        self.src.set_data_order(dorder)?;
208        self.src.set_queries(self.queries.as_slice());
209        self.src.set_origin_query(self.origin_query.clone());
210        self.src.fetch_metadata()?;
211        let src_schema = self.src.schema();
212        let dst_schema = src_schema
213            .iter()
214            .map(|&s| TP::convert_typesystem(s))
215            .collect::<CXResult<Vec<_>>>()?;
216        let names = self.src.names();
217        self.dst.allocate(0, &names, &dst_schema, dorder)?;
218        Ok(())
219    }
220}
connectorx/dispatcher.rs

connectorx/
dispatcher.rs
