micromegas_analytics/lakehouse/
get_payload_function.rs

1use async_trait::async_trait;
2use datafusion::{
3    arrow::{
4        array::{Array, BinaryBuilder, StringArray},
5        datatypes::DataType,
6    },
7    common::{internal_err, not_impl_err},
8    error::DataFusionError,
9    logical_expr::{
10        ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
11        async_udf::AsyncScalarUDFImpl,
12    },
13};
14use futures::stream::StreamExt;
15use micromegas_ingestion::data_lake_connection::DataLakeConnection;
16use micromegas_tracing::prelude::*;
17use std::sync::Arc;
18
19/// A scalar UDF that retrieves the payload of a block from the data lake.
20#[derive(Debug)]
21pub struct GetPayload {
22    signature: Signature,
23    lake: Arc<DataLakeConnection>,
24}
25
26impl PartialEq for GetPayload {
27    fn eq(&self, other: &Self) -> bool {
28        self.signature == other.signature
29    }
30}
31
32impl Eq for GetPayload {}
33
34impl std::hash::Hash for GetPayload {
35    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
36        self.signature.hash(state);
37    }
38}
39
40impl GetPayload {
41    pub fn new(lake: Arc<DataLakeConnection>) -> Self {
42        Self {
43            signature: Signature::exact(
44                vec![DataType::Utf8, DataType::Utf8, DataType::Utf8],
45                Volatility::Immutable,
46            ),
47            lake,
48        }
49    }
50}
51
52impl ScalarUDFImpl for GetPayload {
53    fn name(&self) -> &str {
54        "get_payload"
55    }
56
57    fn signature(&self) -> &Signature {
58        &self.signature
59    }
60
61    fn return_type(&self, _arg_types: &[DataType]) -> datafusion::error::Result<DataType> {
62        Ok(DataType::Binary)
63    }
64
65    fn invoke_with_args(
66        &self,
67        _args: datafusion::logical_expr::ScalarFunctionArgs,
68    ) -> datafusion::error::Result<ColumnarValue> {
69        not_impl_err!("GetPayload can only be called from async contexts")
70    }
71}
72
73#[async_trait]
74impl AsyncScalarUDFImpl for GetPayload {
75    async fn invoke_async_with_args(
76        &self,
77        args: ScalarFunctionArgs,
78    ) -> datafusion::error::Result<ColumnarValue> {
79        let args = ColumnarValue::values_to_arrays(&args.args)?;
80        if args.len() != 3 {
81            return internal_err!("wrong number of arguments to get_payload()");
82        }
83        let process_ids = args[0]
84            .as_any()
85            .downcast_ref::<StringArray>()
86            .ok_or_else(|| {
87                DataFusionError::Execution("downcasting process_ids in GetPayload".into())
88            })?
89            .clone();
90        let stream_ids = args[1]
91            .as_any()
92            .downcast_ref::<StringArray>()
93            .ok_or_else(|| {
94                DataFusionError::Execution("downcasting stream_ids in GetPayload".into())
95            })?
96            .clone();
97        let block_ids = args[2]
98            .as_any()
99            .downcast_ref::<StringArray>()
100            .ok_or_else(|| {
101                DataFusionError::Execution("downcasting block_ids in GetPayload".into())
102            })?
103            .clone();
104        let lake = self.lake.clone();
105        let mut stream = futures::stream::iter(0..process_ids.len())
106            .map(|i| {
107                let process_id = process_ids.value(i);
108                let stream_id = stream_ids.value(i);
109                let block_id = block_ids.value(i);
110                let obj_path = format!("blobs/{process_id}/{stream_id}/{block_id}");
111                let lake = lake.clone();
112                spawn_with_context(async move { lake.blob_storage.read_blob(&obj_path).await })
113            })
114            .buffered(10);
115        let mut result_builder = BinaryBuilder::with_capacity(block_ids.len(), 1024 * 1024);
116        while let Some(res) = stream.next().await {
117            result_builder.append_value(
118                res.map_err(|e| {
119                    DataFusionError::Execution(format!("error downloading payload: {e:?}"))
120                })?
121                .map_err(|e| {
122                    DataFusionError::Execution(format!("error downloading payload: {e:?}"))
123                })?,
124            );
125        }
126        Ok(ColumnarValue::Array(Arc::new(result_builder.finish())))
127    }
128}
micromegas_analytics/lakehouse/get_payload_function.rs

micromegas_analytics/lakehouse/
get_payload_function.rs
