Skip to main content

connectorx/destinations/
mod.rs

1//! This module defines three traits [`Destination`], [`DestinationPartition`], and [`Consume`] to define a destination.
2//! This module also contains destination implementations for various dataframes.
3
4#[cfg(feature = "dst_arrow")]
5pub mod arrow;
6#[cfg(feature = "dst_arrow")]
7pub mod arrowstream;
8
9use crate::data_order::DataOrder;
10use crate::errors::ConnectorXError;
11use crate::typesystem::{TypeAssoc, TypeSystem};
12
13/// A `Destination` is associated with a `TypeSystem` and a `PartitionDestination`.
14/// `PartitionDestination` allows multiple threads write data into the buffer owned by `Destination`.
15///
16/// [`Dispatcher`](crate::dispatcher::Dispatcher) calls [`needs_count`](Destination::needs_count),
17/// then [`allocate`](Destination::allocate), then [`partition`](Destination::partition), and
18/// hands each partition to a worker thread.
19pub trait Destination: Sized {
20    /// Supported data orders, ordering by preference.
21    const DATA_ORDERS: &'static [DataOrder];
22    /// The type system this `Destination` is associated with.
23    type TypeSystem: TypeSystem;
24    /// The partition type created by [`partition`](Destination::partition). It borrows the
25    /// destination's buffer.
26    type Partition<'a>: DestinationPartition<'a, TypeSystem = Self::TypeSystem, Error = Self::Error>
27    where
28        Self: 'a;
29    /// The error type returned by this destination and its partitions.
30    type Error: From<ConnectorXError> + Send;
31
32    /// Specify whether the destination needs total rows in advance
33    /// in order to pre-allocate the buffer.
34    fn needs_count(&self) -> bool;
35
36    /// Construct the `Destination`.
37    /// This allocates the memory based on the types of each columns
38    /// and the number of rows.
39    fn allocate<S: AsRef<str>>(
40        &mut self,
41        nrow: usize,
42        names: &[S],
43        schema: &[Self::TypeSystem],
44        data_order: DataOrder,
45    ) -> Result<(), Self::Error>;
46
47    /// Create `counts` partition destinations, one per source partition, in the same order.
48    fn partition(&mut self, counts: usize) -> Result<Vec<Self::Partition<'_>>, Self::Error>;
49    /// Return the schema of the destination.
50    fn schema(&self) -> &[Self::TypeSystem];
51}
52
53/// `PartitionDestination` writes values to its own region. `PartitionDestination` is parameterized
54/// on lifetime `'a`, which is the lifetime of the parent `Destination`. Usually,
55/// a `PartitionDestination` can never live longer than the parent.
56pub trait DestinationPartition<'a>: Send {
57    /// The type system of the parent [`Destination`].
58    type TypeSystem: TypeSystem;
59    /// The error type of the parent [`Destination`].
60    type Error: From<ConnectorXError> + Send;
61
62    /// Write a value of type T to the location (row, col). If T mismatch with the
63    /// schema, `ConnectorXError::TypeCheckFailed` will return.
64    fn write<T>(&mut self, value: T) -> Result<(), <Self as DestinationPartition<'a>>::Error>
65    where
66        T: TypeAssoc<Self::TypeSystem>,
67        Self: Consume<T, Error = <Self as DestinationPartition<'a>>::Error>,
68    {
69        self.consume(value)
70    }
71
72    /// Number of columns this `PartitionDestination` writes.
73    fn ncols(&self) -> usize;
74
75    /// Final clean ups, called once after the last value of the partition is written.
76    fn finalize(&mut self) -> Result<(), Self::Error>;
77
78    /// Aquire n rows in final destination.
79    ///
80    /// Called before each batch of `n` rows is written. Destinations that write into a shared,
81    /// pre-allocated buffer (see [`needs_count`](Destination::needs_count)) use this to claim the
82    /// next `n` rows for this partition. Returns the starting row index of the claimed range; the
83    /// dispatcher currently ignores it.
84    fn aquire_row(&mut self, n: usize) -> Result<usize, Self::Error>;
85}
86
87/// A type implemented `Consume<T>` means that it can consume a value `T` by adding it to it's own buffer.
88pub trait Consume<T> {
89    /// The error returned when `T` cannot be written.
90    type Error: From<ConnectorXError> + Send;
91    /// Write `value` to the next position of the buffer.
92    fn consume(&mut self, value: T) -> Result<(), Self::Error>;
93}