Skip to main content

datafusion_ethers/convert/
hybrid.rs

1use alloy_core::dyn_abi::Error as DynAbiError;
2use alloy_core::dyn_abi::{DynSolEvent, Specifier};
3use alloy_core::json_abi::Event;
4use alloy_rpc_types_eth::Log;
5use datafusion::arrow::array::RecordBatch;
6use datafusion::arrow::datatypes::{DataType, Field, SchemaBuilder, SchemaRef};
7use std::sync::Arc;
8
9use crate::stream::StreamOptions;
10
11use super::{AppendError, Transcoder};
12
13///////////////////////////////////////////////////////////////////////////////////////////////////
14
15/// Transcodes raw Ethereum log events into Arrow record batches while also creating a nested
16/// struct with fields decoded according to the provided log signature
17pub struct EthRawAndDecodedLogsToArrow {
18    schema: SchemaRef,
19    raw: super::EthRawLogsToArrow,
20    decoded: super::EthDecodedLogsToArrow,
21}
22
23impl EthRawAndDecodedLogsToArrow {
24    pub fn new(options: &StreamOptions, event_type: Event, resolved_type: DynSolEvent) -> Self {
25        let raw = super::EthRawLogsToArrow::new(options);
26        let decoded = super::EthDecodedLogsToArrow::new(event_type, resolved_type);
27
28        let mut builder = SchemaBuilder::from(&raw.schema().fields);
29        builder.push(Field::new(
30            "event",
31            DataType::Struct(decoded.schema().fields.clone()),
32            false,
33        ));
34        Self {
35            schema: Arc::new(builder.finish()),
36            raw,
37            decoded,
38        }
39    }
40
41    pub fn new_from_signature(
42        options: &StreamOptions,
43        signature: &str,
44    ) -> Result<Self, DynAbiError> {
45        let event_type = Event::parse(signature)?;
46        let resolved_type = event_type.resolve()?;
47        Ok(Self::new(options, event_type, resolved_type))
48    }
49}
50
51impl Transcoder for EthRawAndDecodedLogsToArrow {
52    fn schema(&self) -> SchemaRef {
53        self.schema.clone()
54    }
55
56    #[allow(clippy::get_first)]
57    fn append(&mut self, logs: &[Log]) -> Result<(), AppendError> {
58        self.raw.append(logs)?;
59        self.decoded.append(logs)?;
60        Ok(())
61    }
62
63    fn len(&self) -> usize {
64        self.raw.len()
65    }
66
67    fn finish(&mut self) -> RecordBatch {
68        let raw = self.raw.finish();
69        let decoded = self.decoded.finish();
70
71        let event_col = datafusion::arrow::array::StructArray::new(
72            decoded.schema().fields().clone(),
73            decoded.columns().to_vec(),
74            None,
75        );
76
77        let mut columns = raw.columns().to_vec();
78        columns.push(Arc::new(event_col));
79
80        RecordBatch::try_new(self.schema.clone(), columns).unwrap()
81    }
82}
83
84///////////////////////////////////////////////////////////////////////////////////////////////////