micromegas-analytics 0.19.0

analytics module of micromegas
Documentation 
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

use async_trait::async_trait;
use datafusion::{
    arrow::{
        array::{Array, BinaryBuilder, StringArray},
        datatypes::DataType,
    },
    common::{internal_err, not_impl_err},
    error::DataFusionError,
    logical_expr::{
        ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
        async_udf::AsyncScalarUDFImpl,
    },
};
use futures::stream::StreamExt;
use micromegas_ingestion::data_lake_connection::DataLakeConnection;
use micromegas_tracing::prelude::*;
use std::sync::Arc;

/// A scalar UDF that retrieves the payload of a block from the data lake.
#[derive(Debug)]
pub struct GetPayload {
    signature: Signature,
    lake: Arc<DataLakeConnection>,
}

impl PartialEq for GetPayload {
    fn eq(&self, other: &Self) -> bool {
        self.signature == other.signature
    }
}

impl Eq for GetPayload {}

impl std::hash::Hash for GetPayload {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.signature.hash(state);
    }
}

impl GetPayload {
    pub fn new(lake: Arc<DataLakeConnection>) -> Self {
        Self {
            signature: Signature::exact(
                vec![DataType::Utf8, DataType::Utf8, DataType::Utf8],
                Volatility::Immutable,
            ),
            lake,
        }
    }
}

impl ScalarUDFImpl for GetPayload {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn name(&self) -> &str {
        "get_payload"
    }

    fn signature(&self) -> &Signature {
        &self.signature
    }

    fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
        Ok(DataType::Binary)
    }

    fn invoke_with_args(
        &self,
        _args: datafusion::logical_expr::ScalarFunctionArgs,
    ) -> datafusion::error::Result<ColumnarValue> {
        not_impl_err!("GetPayload can only be called from async contexts")
    }
}

#[async_trait]
impl AsyncScalarUDFImpl for GetPayload {
    async fn invoke_async_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> datafusion::error::Result<ColumnarValue> {
        let args = ColumnarValue::values_to_arrays(&args.args)?;
        if args.len() != 3 {
            return internal_err!("wrong number of arguments to get_payload()");
        }
        let process_ids = args[0]
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or_else(|| {
                DataFusionError::Execution("downcasting process_ids in GetPayload".into())
            })?
            .clone();
        let stream_ids = args[1]
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or_else(|| {
                DataFusionError::Execution("downcasting stream_ids in GetPayload".into())
            })?
            .clone();
        let block_ids = args[2]
            .as_any()
            .downcast_ref::<StringArray>()
            .ok_or_else(|| {
                DataFusionError::Execution("downcasting block_ids in GetPayload".into())
            })?
            .clone();
        let lake = self.lake.clone();
        let mut stream = futures::stream::iter(0..process_ids.len())
            .map(|i| {
                let process_id = process_ids.value(i);
                let stream_id = stream_ids.value(i);
                let block_id = block_ids.value(i);
                let obj_path = format!("blobs/{process_id}/{stream_id}/{block_id}");
                let lake = lake.clone();
                spawn_with_context(async move { lake.blob_storage.read_blob(&obj_path).await })
            })
            .buffered(10);
        let mut result_builder = BinaryBuilder::with_capacity(block_ids.len(), 1024 * 1024);
        while let Some(res) = stream.next().await {
            result_builder.append_value(
                res.map_err(|e| {
                    DataFusionError::Execution(format!("error downloading payload: {e:?}"))
                })?
                .map_err(|e| {
                    DataFusionError::Execution(format!("error downloading payload: {e:?}"))
                })?,
            );
        }
        Ok(ColumnarValue::Array(Arc::new(result_builder.finish())))
    }
}