Skip to main content

datafusion_ethers/
udf.rs

1use std::any::Any;
2use std::sync::Arc;
3
4use alloy_core::dyn_abi::{DecodedEvent, DynSolValue, EventExt as _};
5use alloy_core::hex::ToHexExt;
6use alloy_core::json_abi::Event;
7use alloy_core::primitives::B256;
8use datafusion::arrow::array::{self, Array as _};
9use datafusion::error::DataFusionError;
10use datafusion::execution::FunctionRegistry;
11use datafusion::logical_expr::ScalarFunctionArgs;
12use datafusion::{
13    arrow::datatypes::DataType,
14    common::plan_err,
15    logical_expr::{ColumnarValue, ScalarUDF, ScalarUDFImpl, Signature, Volatility},
16    scalar::ScalarValue,
17};
18
19///////////////////////////////////////////////////////////////////////////////////////////////////
20
21pub fn event_to_json(schema: &Event, event: &DecodedEvent) -> serde_json::Value {
22    let mut obj = serde_json::Map::new();
23    obj.insert("name".to_string(), schema.name.clone().into());
24
25    let mut i_indexed = 0;
26    let mut i_body = 0;
27
28    for input in &schema.inputs {
29        let value = if input.indexed {
30            let v = &event.indexed[i_indexed];
31            i_indexed += 1;
32            v
33        } else {
34            let v = &event.body[i_body];
35            i_body += 1;
36            v
37        };
38
39        obj.insert(input.name.clone(), solidity_value_to_json(value));
40    }
41
42    obj.into()
43}
44
45pub fn solidity_value_to_json(value: &DynSolValue) -> serde_json::Value {
46    use serde_json::Value;
47    match value {
48        DynSolValue::Bool(v) => Value::Bool(*v),
49        DynSolValue::Int(v, bits) if *bits <= 64 => Value::Number(v.as_i64().into()),
50        DynSolValue::Int(v, _) => Value::String(v.to_string()),
51        DynSolValue::Uint(v, bits) if *bits <= 64 => Value::Number(v.to::<u64>().into()),
52        DynSolValue::Uint(v, _) => Value::String(v.to_string()),
53        DynSolValue::FixedBytes(v, _) => Value::String(v.encode_hex()),
54        DynSolValue::Address(v) => Value::String(v.encode_hex()),
55        DynSolValue::Function(_) => todo!(),
56        DynSolValue::Bytes(v) => Value::String(v.encode_hex()),
57        DynSolValue::String(v) => Value::String(v.clone()),
58        DynSolValue::Array(_) => todo!(),
59        DynSolValue::FixedArray(_) => todo!(),
60        DynSolValue::Tuple(_) => todo!(),
61        #[cfg(feature = "eip712")]
62        DynSolValue::CustomStruct { .. } => todo!(),
63    }
64}
65
66///////////////////////////////////////////////////////////////////////////////////////////////////
67
68#[derive(Debug, PartialEq, Eq, Hash)]
69struct UdfEthDecodeEvent {
70    signature: Signature,
71}
72
73impl UdfEthDecodeEvent {
74    fn new() -> Self {
75        Self {
76            signature: Signature::variadic_any(Volatility::Immutable),
77        }
78    }
79}
80
81impl ScalarUDFImpl for UdfEthDecodeEvent {
82    fn as_any(&self) -> &dyn Any {
83        self
84    }
85
86    fn name(&self) -> &str {
87        "eth_decode_event"
88    }
89
90    fn signature(&self) -> &Signature {
91        &self.signature
92    }
93
94    #[allow(clippy::get_first)]
95    fn return_type(&self, args: &[DataType]) -> datafusion::error::Result<DataType> {
96        if args.len() != 6 {
97            return plan_err!(
98                "eth_decode_event accepts 6 arguments: event_signature, topic0, topic1, topic2, topic3, data"
99            );
100        }
101        if !matches!(args.get(0), Some(&DataType::Utf8)) {
102            return plan_err!("event_signature must be a Utf8 scalar");
103        }
104        if !matches!(args.get(1), Some(&DataType::Binary)) {
105            return plan_err!("topic0 must be a Binary");
106        }
107        if !matches!(args.get(2), Some(&DataType::Binary)) {
108            return plan_err!("topic1 must be a Binary");
109        }
110        if !matches!(args.get(3), Some(&DataType::Binary)) {
111            return plan_err!("topic2 must be a Binary");
112        }
113        if !matches!(args.get(4), Some(&DataType::Binary)) {
114            return plan_err!("topic3 must be a Binary");
115        }
116        if !matches!(args.get(5), Some(&DataType::Binary)) {
117            return plan_err!("data must be a Binary");
118        }
119        Ok(DataType::Utf8)
120    }
121
122    fn invoke_with_args(
123        &self,
124        args: ScalarFunctionArgs,
125    ) -> datafusion::error::Result<ColumnarValue> {
126        let signature = match &args.args[0] {
127            ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) => v,
128            _ => return plan_err!("Event signature must be a Utf8 scalar"),
129        };
130
131        let event = Event::parse(signature)
132            .map_err(|e| datafusion::error::DataFusionError::External(e.into()))?;
133
134        let args = ColumnarValue::values_to_arrays(&args.args[1..])?;
135        let c_topic0 = datafusion::common::cast::as_binary_array(&args[0])?;
136        let c_topic1 = datafusion::common::cast::as_binary_array(&args[1])?;
137        let c_topic2 = datafusion::common::cast::as_binary_array(&args[2])?;
138        let c_topic3 = datafusion::common::cast::as_binary_array(&args[3])?;
139        let c_data = datafusion::common::cast::as_binary_array(&args[4])?;
140
141        let mut builder = array::StringBuilder::new();
142
143        for i in 0..c_topic0.len() {
144            let mut topics = Vec::with_capacity(4);
145
146            for c_topic in [c_topic0, c_topic1, c_topic2, c_topic3] {
147                if c_topic.is_null(i) {
148                    break;
149                }
150                topics.push(B256::try_from(c_topic.value(i)).expect("??"));
151            }
152            let data = c_data.value(i);
153
154            let decoded = event
155                .decode_log_parts(topics, data)
156                .map_err(|e| DataFusionError::External(e.into()))?;
157
158            builder.append_value(event_to_json(&event, &decoded).to_string());
159        }
160
161        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
162    }
163}
164
165///////////////////////////////////////////////////////////////////////////////////////////////////
166
167#[derive(Debug, PartialEq, Eq, Hash)]
168struct UdfEthTryDecodeEvent {
169    signature: Signature,
170}
171
172impl UdfEthTryDecodeEvent {
173    fn new() -> Self {
174        Self {
175            signature: Signature::variadic_any(Volatility::Immutable),
176        }
177    }
178}
179
180impl ScalarUDFImpl for UdfEthTryDecodeEvent {
181    fn as_any(&self) -> &dyn Any {
182        self
183    }
184
185    fn name(&self) -> &str {
186        "eth_try_decode_event"
187    }
188
189    fn signature(&self) -> &Signature {
190        &self.signature
191    }
192
193    #[allow(clippy::get_first)]
194    fn return_type(&self, args: &[DataType]) -> datafusion::error::Result<DataType> {
195        if args.len() != 6 {
196            return plan_err!(
197                "eth_decode_event accepts 6 arguments: event_signature, topic0, topic1, topic2, topic3, data"
198            );
199        }
200        if !matches!(args.get(0), Some(&DataType::Utf8)) {
201            return plan_err!("event_signature must be a Utf8 scalar");
202        }
203        if !matches!(args.get(1), Some(&DataType::Binary)) {
204            return plan_err!("topic0 must be a Binary");
205        }
206        if !matches!(args.get(2), Some(&DataType::Binary)) {
207            return plan_err!("topic1 must be a Binary");
208        }
209        if !matches!(args.get(3), Some(&DataType::Binary)) {
210            return plan_err!("topic2 must be a Binary");
211        }
212        if !matches!(args.get(4), Some(&DataType::Binary)) {
213            return plan_err!("topic3 must be a Binary");
214        }
215        if !matches!(args.get(5), Some(&DataType::Binary)) {
216            return plan_err!("data must be a Binary");
217        }
218        Ok(DataType::Utf8)
219    }
220
221    fn invoke_with_args(
222        &self,
223        args: ScalarFunctionArgs,
224    ) -> datafusion::error::Result<ColumnarValue> {
225        let signature = match &args.args[0] {
226            ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) => v,
227            _ => return plan_err!("Event signature must be a Utf8 scalar"),
228        };
229
230        let event = Event::parse(signature)
231            .map_err(|e| datafusion::error::DataFusionError::External(e.into()))?;
232
233        let args = ColumnarValue::values_to_arrays(&args.args[1..])?;
234        let c_topic0 = datafusion::common::cast::as_binary_array(&args[0])?;
235        let c_topic1 = datafusion::common::cast::as_binary_array(&args[1])?;
236        let c_topic2 = datafusion::common::cast::as_binary_array(&args[2])?;
237        let c_topic3 = datafusion::common::cast::as_binary_array(&args[3])?;
238        let c_data = datafusion::common::cast::as_binary_array(&args[4])?;
239
240        let mut builder = array::StringBuilder::new();
241
242        for i in 0..c_topic0.len() {
243            let mut topics = Vec::with_capacity(4);
244
245            for c_topic in [c_topic0, c_topic1, c_topic2, c_topic3] {
246                if c_topic.is_null(i) {
247                    break;
248                }
249                topics.push(B256::try_from(c_topic.value(i)).expect("??"));
250            }
251            let data = c_data.value(i);
252
253            match event.decode_log_parts(topics, data) {
254                Ok(decoded) => builder.append_value(event_to_json(&event, &decoded).to_string()),
255                Err(_) => builder.append_null(),
256            }
257        }
258
259        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
260    }
261}
262
263///////////////////////////////////////////////////////////////////////////////////////////////////
264
265#[derive(Debug, PartialEq, Eq, Hash)]
266struct UdfEthEventSelector {
267    signature: Signature,
268}
269
270impl UdfEthEventSelector {
271    fn new() -> Self {
272        Self {
273            signature: Signature::exact(vec![DataType::Utf8], Volatility::Immutable),
274        }
275    }
276}
277
278impl ScalarUDFImpl for UdfEthEventSelector {
279    fn as_any(&self) -> &dyn Any {
280        self
281    }
282
283    fn name(&self) -> &str {
284        "eth_event_selector"
285    }
286
287    fn signature(&self) -> &Signature {
288        &self.signature
289    }
290
291    fn return_type(&self, _args: &[DataType]) -> datafusion::error::Result<DataType> {
292        Ok(DataType::Binary)
293    }
294
295    fn invoke_with_args(
296        &self,
297        args: ScalarFunctionArgs,
298    ) -> datafusion::error::Result<ColumnarValue> {
299        let signature = match &args.args[0] {
300            ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) => v,
301            _ => return plan_err!("Event signature must be a Utf8 scalar"),
302        };
303
304        let event = Event::parse(signature)
305            .map_err(|e| datafusion::error::DataFusionError::External(e.into()))?;
306
307        let mut builder = array::BinaryBuilder::new();
308        builder.append_value(event.selector().as_slice());
309        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
310    }
311}
312
313///////////////////////////////////////////////////////////////////////////////////////////////////
314
315pub fn register_all(registry: &mut dyn FunctionRegistry) -> datafusion::error::Result<()> {
316    registry.register_udf(Arc::new(ScalarUDF::from(UdfEthDecodeEvent::new())))?;
317    registry.register_udf(Arc::new(ScalarUDF::from(UdfEthTryDecodeEvent::new())))?;
318    registry.register_udf(Arc::new(ScalarUDF::from(UdfEthEventSelector::new())))?;
319    Ok(())
320}
321
322///////////////////////////////////////////////////////////////////////////////////////////////////