Spaces:
Build error
Build error
File size: 4,578 Bytes
84d2a97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
use std::time::Duration;
use common::counter::hardware_accumulator::HwMeasurementAcc;
use futures::Future;
use itertools::Itertools;
use tokio::sync::RwLockReadGuard;
use super::group_by::{group_by, GroupRequest};
use crate::collection::Collection;
use crate::lookup::lookup_ids;
use crate::lookup::types::PseudoId;
use crate::operations::consistency_params::ReadConsistency;
use crate::operations::shard_selector_internal::ShardSelectorInternal;
use crate::operations::types::{CollectionError, CollectionResult, PointGroup};
/// Builds on top of the group_by function to add lookup and possibly other features
pub struct GroupBy<'a, F, Fut>
where
F: Fn(String) -> Fut + Clone,
Fut: Future<Output = Option<RwLockReadGuard<'a, Collection>>>,
{
group_by: GroupRequest,
collection: &'a Collection,
/// `Fn` to get a collection having its name. Obligatory for recommend and lookup
collection_by_name: F,
read_consistency: Option<ReadConsistency>,
shard_selection: ShardSelectorInternal,
timeout: Option<Duration>,
hw_measurement_acc: HwMeasurementAcc,
}
impl<'a, F, Fut> GroupBy<'a, F, Fut>
where
F: Fn(String) -> Fut + Clone,
Fut: Future<Output = Option<RwLockReadGuard<'a, Collection>>>,
{
/// Creates a basic GroupBy builder
pub fn new(
group_by: GroupRequest,
collection: &'a Collection,
collection_by_name: F,
hw_measurement_acc: &HwMeasurementAcc,
) -> Self {
Self {
group_by,
collection,
collection_by_name,
read_consistency: None,
shard_selection: ShardSelectorInternal::All,
timeout: None,
hw_measurement_acc: hw_measurement_acc.new_collector(),
}
}
pub fn set_read_consistency(mut self, read_consistency: Option<ReadConsistency>) -> Self {
self.read_consistency = read_consistency;
self
}
pub fn set_shard_selection(mut self, shard_selection: ShardSelectorInternal) -> Self {
self.shard_selection = shard_selection;
self
}
pub fn set_timeout(mut self, timeout: Option<Duration>) -> Self {
self.timeout = timeout;
self
}
/// Runs the group by operation, optionally with a timeout.
pub async fn execute(self) -> CollectionResult<Vec<PointGroup>> {
if let Some(timeout) = self.timeout {
tokio::time::timeout(timeout, self.run())
.await
.map_err(|_| {
log::debug!("GroupBy timeout reached: {} seconds", timeout.as_secs());
CollectionError::timeout(timeout.as_secs() as usize, "GroupBy")
})?
} else {
self.run().await
}
}
/// Does the actual grouping
async fn run(self) -> CollectionResult<Vec<PointGroup>> {
let start = std::time::Instant::now();
let with_lookup = self.group_by.with_lookup.clone();
let core_group_by = self
.group_by
.into_query_group_request(
self.collection,
self.collection_by_name.clone(),
self.read_consistency,
self.shard_selection.clone(),
self.timeout,
)
.await?;
let mut groups = group_by(
core_group_by,
self.collection,
self.read_consistency,
self.shard_selection.clone(),
self.timeout,
&self.hw_measurement_acc,
)
.await?;
if let Some(lookup) = with_lookup {
// update timeout
let timeout = self
.timeout
.map(|timeout| timeout.saturating_sub(start.elapsed()));
let mut lookups = {
let pseudo_ids = groups
.iter()
.map(|group| group.id.clone())
.map_into()
.collect();
lookup_ids(
lookup,
pseudo_ids,
self.collection_by_name,
self.read_consistency,
&self.shard_selection,
timeout,
)
.await?
};
// Put the lookups in their respective groups
groups.iter_mut().for_each(|group| {
group.lookup = lookups
.remove(&PseudoId::from(group.id.clone()))
.map(api::rest::Record::from);
});
}
Ok(groups)
}
}
|