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}