datafusion_ethers/
provider.rs

1use alloy_core::primitives::{Address, B256};
2use alloy_network::AnyNetwork;
3use alloy_provider::DynProvider;
4use alloy_rpc_types_eth::{BlockNumberOrTag, Filter, FilterBlockOption};
5use datafusion::catalog::{CatalogProvider, SchemaProvider, Session};
6use datafusion::common::tree_node::TreeNodeRecursion;
7use datafusion::error::{DataFusionError, Result as DfResult};
8use datafusion::execution::TaskContext;
9use datafusion::logical_expr::BinaryExpr;
10use datafusion::physical_expr::EquivalenceProperties;
11use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
12use datafusion::physical_plan::{DisplayAs, DisplayFormatType, PlanProperties};
13use datafusion::{
14    arrow::{array::RecordBatch, datatypes::SchemaRef},
15    datasource::{TableProvider, TableType},
16    logical_expr::{Operator, TableProviderFilterPushDown},
17    physical_plan::ExecutionPlan,
18    prelude::*,
19    scalar::ScalarValue,
20};
21use futures::{Stream, TryStreamExt};
22use std::sync::Arc;
23
24use crate::config::EthProviderConfig;
25use crate::convert::Transcoder as _;
26use crate::stream::StreamOptions;
27use crate::utils::*;
28
29///////////////////////////////////////////////////////////////////////////////////////////////////
30// Catalog
31///////////////////////////////////////////////////////////////////////////////////////////////////
32
33#[derive(Debug)]
34pub struct EthCatalog {
35    config: EthProviderConfig,
36    rpc_client: DynProvider<AnyNetwork>,
37}
38
39impl EthCatalog {
40    pub fn new(config: EthProviderConfig, rpc_client: DynProvider<AnyNetwork>) -> Self {
41        Self { config, rpc_client }
42    }
43}
44
45impl CatalogProvider for EthCatalog {
46    fn schema_names(&self) -> Vec<String> {
47        vec![self.config.schema_name.clone()]
48    }
49
50    fn schema(&self, name: &str) -> Option<Arc<dyn SchemaProvider>> {
51        if name != self.config.schema_name {
52            return None;
53        }
54
55        Some(Arc::new(EthSchema::new(
56            self.config.clone(),
57            self.rpc_client.clone(),
58        )))
59    }
60}
61
62///////////////////////////////////////////////////////////////////////////////////////////////////
63// Schema
64///////////////////////////////////////////////////////////////////////////////////////////////////
65
66#[derive(Debug)]
67pub struct EthSchema {
68    config: EthProviderConfig,
69    rpc_client: DynProvider<AnyNetwork>,
70}
71
72impl EthSchema {
73    pub fn new(config: EthProviderConfig, rpc_client: DynProvider<AnyNetwork>) -> Self {
74        Self { config, rpc_client }
75    }
76}
77
78#[async_trait::async_trait]
79impl SchemaProvider for EthSchema {
80    fn table_names(&self) -> Vec<String> {
81        vec!["logs".to_string()]
82    }
83
84    async fn table(&self, name: &str) -> DfResult<Option<Arc<dyn TableProvider>>> {
85        match name {
86            "logs" => Ok(Some(Arc::new(EthLogsTable::new(
87                &self.config,
88                self.rpc_client.clone(),
89            )))),
90            _ => Ok(None),
91        }
92    }
93
94    fn table_exist(&self, name: &str) -> bool {
95        name == "logs"
96    }
97}
98
99///////////////////////////////////////////////////////////////////////////////////////////////////
100// Table: eth_logs
101///////////////////////////////////////////////////////////////////////////////////////////////////
102
103#[derive(Debug)]
104pub struct EthLogsTable {
105    schema: SchemaRef,
106    rpc_client: DynProvider<AnyNetwork>,
107}
108
109impl EthLogsTable {
110    pub fn new(config: &EthProviderConfig, rpc_client: DynProvider<AnyNetwork>) -> Self {
111        let encoder = crate::convert::EthRawLogsToArrow::new(&config.stream_options());
112        let schema = encoder.schema();
113
114        Self { schema, rpc_client }
115    }
116
117    fn apply_expr(filter: Filter, expr: &Expr) -> (TableProviderFilterPushDown, Filter) {
118        match expr {
119            Expr::BinaryExpr(e) => match (&*e.left, e.op, &*e.right) {
120                (Expr::Column(left), op, right) => match (left.name.as_str(), op, right) {
121                    (
122                        "address",
123                        Operator::Eq,
124                        Expr::Literal(ScalarValue::Binary(Some(v)), None),
125                    ) => (
126                        TableProviderFilterPushDown::Exact,
127                        filter.address(Address::from_slice(&v[..])),
128                    ),
129                    ("topic0", Operator::Eq, Expr::Literal(ScalarValue::Binary(Some(v)), None)) => {
130                        (
131                            TableProviderFilterPushDown::Exact,
132                            filter.event_signature(B256::from_slice(&v[..])),
133                        )
134                    }
135                    ("topic1", Operator::Eq, Expr::Literal(ScalarValue::Binary(Some(v)), None)) => {
136                        (
137                            TableProviderFilterPushDown::Exact,
138                            filter.topic1(B256::from_slice(&v[..])),
139                        )
140                    }
141                    ("topic2", Operator::Eq, Expr::Literal(ScalarValue::Binary(Some(v)), None)) => {
142                        (
143                            TableProviderFilterPushDown::Exact,
144                            filter.topic2(B256::from_slice(&v[..])),
145                        )
146                    }
147                    ("topic3", Operator::Eq, Expr::Literal(ScalarValue::Binary(Some(v)), None)) => {
148                        (
149                            TableProviderFilterPushDown::Exact,
150                            filter.topic3(B256::from_slice(&v[..])),
151                        )
152                    }
153                    (
154                        "block_number",
155                        Operator::Eq,
156                        Expr::Literal(ScalarValue::UInt64(Some(v)), None),
157                    ) => (
158                        TableProviderFilterPushDown::Exact,
159                        filter.union((*v).into(), (*v).into()),
160                    ),
161                    (
162                        "block_number",
163                        Operator::Gt,
164                        Expr::Literal(ScalarValue::UInt64(Some(v)), None),
165                    ) => (
166                        TableProviderFilterPushDown::Exact,
167                        filter.union((*v + 1).into(), BlockNumberOrTag::Latest),
168                    ),
169                    (
170                        "block_number",
171                        Operator::GtEq,
172                        Expr::Literal(ScalarValue::UInt64(Some(v)), None),
173                    ) => (
174                        TableProviderFilterPushDown::Exact,
175                        filter.union((*v).into(), BlockNumberOrTag::Latest),
176                    ),
177                    (
178                        "block_number",
179                        Operator::Lt,
180                        Expr::Literal(ScalarValue::UInt64(Some(v)), None),
181                    ) => (
182                        TableProviderFilterPushDown::Exact,
183                        filter.union(BlockNumberOrTag::Earliest, (*v - 1).into()),
184                    ),
185                    (
186                        "block_number",
187                        Operator::LtEq,
188                        Expr::Literal(ScalarValue::UInt64(Some(v)), None),
189                    ) => (
190                        TableProviderFilterPushDown::Exact,
191                        filter.union(BlockNumberOrTag::Earliest, (*v).into()),
192                    ),
193                    _ => (TableProviderFilterPushDown::Unsupported, filter),
194                },
195                // expr OR expr OR ...
196                (Expr::BinaryExpr(_), Operator::Or, Expr::BinaryExpr(_)) => {
197                    match Self::collect_or_group(e) {
198                        (TableProviderFilterPushDown::Exact, Some(col), values) => {
199                            match col.name.as_str() {
200                                "address" => (
201                                    TableProviderFilterPushDown::Exact,
202                                    filter.address(
203                                        values
204                                            .into_iter()
205                                            .map(|v| Address::from_slice(&v[..]))
206                                            .collect::<Vec<_>>(),
207                                    ),
208                                ),
209                                "topic0" => (
210                                    TableProviderFilterPushDown::Exact,
211                                    filter.event_signature(
212                                        values
213                                            .into_iter()
214                                            .map(|v| B256::from_slice(&v[..]))
215                                            .collect::<Vec<_>>(),
216                                    ),
217                                ),
218                                "topic1" => (
219                                    TableProviderFilterPushDown::Exact,
220                                    filter.topic1(
221                                        values
222                                            .into_iter()
223                                            .map(|v| B256::from_slice(&v[..]))
224                                            .collect::<Vec<_>>(),
225                                    ),
226                                ),
227                                "topic2" => (
228                                    TableProviderFilterPushDown::Exact,
229                                    filter.topic2(
230                                        values
231                                            .into_iter()
232                                            .map(|v| B256::from_slice(&v[..]))
233                                            .collect::<Vec<_>>(),
234                                    ),
235                                ),
236                                "topic3" => (
237                                    TableProviderFilterPushDown::Exact,
238                                    filter.topic3(
239                                        values
240                                            .into_iter()
241                                            .map(|v| B256::from_slice(&v[..]))
242                                            .collect::<Vec<_>>(),
243                                    ),
244                                ),
245                                _ => (TableProviderFilterPushDown::Unsupported, filter),
246                            }
247                        }
248                        _ => (TableProviderFilterPushDown::Unsupported, filter),
249                    }
250                }
251                _ => (TableProviderFilterPushDown::Unsupported, filter),
252            },
253            _ => (TableProviderFilterPushDown::Unsupported, filter),
254        }
255    }
256
257    /// Converts expressions like:
258    ///   x = v1 or (x = v2 or (x = v3))
259    /// into:
260    ///   [v1, v2, v3]
261    fn collect_or_group(
262        expr: &BinaryExpr,
263    ) -> (TableProviderFilterPushDown, Option<&Column>, Vec<&Vec<u8>>) {
264        match (&*expr.left, expr.op, &*expr.right) {
265            (
266                Expr::Column(col),
267                Operator::Eq,
268                Expr::Literal(ScalarValue::Binary(Some(val)), None),
269            ) => (TableProviderFilterPushDown::Exact, Some(col), vec![val]),
270            (Expr::BinaryExpr(left), Operator::Or, Expr::BinaryExpr(right)) => {
271                // Merge values from two branches, but only if they refer to the same column
272                let mut left = Self::collect_or_group(left);
273                let mut right = Self::collect_or_group(right);
274                if left.0 == right.0
275                    && right.0 == TableProviderFilterPushDown::Exact
276                    && left.1 == right.1
277                {
278                    left.2.append(&mut right.2);
279                    (TableProviderFilterPushDown::Exact, left.1, left.2)
280                } else {
281                    (TableProviderFilterPushDown::Unsupported, None, Vec::new())
282                }
283            }
284            _ => (TableProviderFilterPushDown::Unsupported, None, Vec::new()),
285        }
286    }
287}
288
289#[async_trait::async_trait]
290impl TableProvider for EthLogsTable {
291    fn table_type(&self) -> TableType {
292        TableType::Base
293    }
294
295    fn schema(&self) -> SchemaRef {
296        self.schema.clone()
297    }
298
299    fn supports_filters_pushdown(
300        &self,
301        filters: &[&Expr],
302    ) -> datafusion::error::Result<Vec<TableProviderFilterPushDown>> {
303        tracing::debug!(?filters, "EthLogs: performing filter pushdown");
304
305        let mut filter = EthProviderConfig::default().default_filter();
306        let mut res = Vec::new();
307
308        for expr in filters {
309            let (support, new_filter) = Self::apply_expr(filter, expr);
310            filter = new_filter;
311            res.push(support);
312        }
313
314        tracing::debug!(?filter, "EthLogs: resulting filter");
315        Ok(res)
316    }
317
318    async fn scan(
319        &self,
320        state: &dyn Session,
321        mut projection: Option<&Vec<usize>>,
322        filters: &[Expr],
323        limit: Option<usize>,
324    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
325        let config = state
326            .config_options()
327            .extensions
328            .get::<EthProviderConfig>()
329            .cloned()
330            .unwrap_or_default();
331
332        tracing::debug!(
333            ?config,
334            ?projection,
335            ?filters,
336            ?limit,
337            "EthLogs: creating execution plan"
338        );
339
340        // TODO: Datafusion always passess a projection even when one is not required
341        if let Some(proj) = projection {
342            let is_redundant = proj.len() == self.schema.fields.len()
343                && proj[0] == 0
344                && proj
345                    .iter()
346                    .cloned()
347                    .reduce(|a, b| if a + 1 == b { b } else { a })
348                    == Some(self.schema.fields.len() - 1);
349            if is_redundant {
350                projection = None
351            }
352        }
353
354        let schema = if let Some(proj) = projection {
355            Arc::new(self.schema.project(proj)?)
356        } else {
357            self.schema.clone()
358        };
359
360        // Get filter with range restrictions from the session config level
361        let mut filter = config.default_filter();
362
363        // Push down filter expressions from the query
364        for expr in filters {
365            let (support, new_filter) = Self::apply_expr(filter, expr);
366            assert_eq!(support, TableProviderFilterPushDown::Exact);
367            filter = new_filter;
368        }
369
370        Ok(Arc::new(EthGetLogs::new(
371            self.rpc_client.clone(),
372            schema,
373            projection.cloned(),
374            filter,
375            config.stream_options(),
376            limit,
377        )))
378    }
379}
380
381///////////////////////////////////////////////////////////////////////////////////////////////////
382
383#[derive(Debug, Clone)]
384pub struct EthGetLogs {
385    projected_schema: SchemaRef,
386    projection: Option<Vec<usize>>,
387    rpc_client: DynProvider<AnyNetwork>,
388    filter: Filter,
389    stream_options: StreamOptions,
390    limit: Option<usize>,
391    properties: Arc<PlanProperties>,
392}
393
394impl EthGetLogs {
395    pub fn new(
396        rpc_client: DynProvider<AnyNetwork>,
397        projected_schema: SchemaRef,
398        projection: Option<Vec<usize>>,
399        filter: Filter,
400        stream_options: StreamOptions,
401        limit: Option<usize>,
402    ) -> Self {
403        Self {
404            projected_schema: projected_schema.clone(),
405            projection,
406            rpc_client,
407            filter,
408            stream_options,
409            limit,
410            properties: Arc::new(
411                PlanProperties::new(
412                    EquivalenceProperties::new(projected_schema),
413                    datafusion::physical_expr::Partitioning::UnknownPartitioning(1),
414                    datafusion::physical_plan::execution_plan::EmissionType::Incremental,
415                    // TODO: Change to Unbounded
416                    datafusion::physical_plan::execution_plan::Boundedness::Bounded,
417                )
418                .with_scheduling_type(
419                    // TODO: Validate assumption: Since we pull data from the network
420                    // we will never get a situation where control is never returned
421                    // to the task scheduler.
422                    datafusion::physical_plan::execution_plan::SchedulingType::Cooperative,
423                ),
424            ),
425        }
426    }
427
428    pub fn filter(&self) -> &Filter {
429        &self.filter
430    }
431
432    pub fn with_filter(mut self, filter: Filter) -> Self {
433        self.filter = filter;
434        self
435    }
436
437    fn execute_impl(
438        rpc_client: DynProvider<AnyNetwork>,
439        filter: Filter,
440        options: StreamOptions,
441        projection: Option<Vec<usize>>,
442        limit: Option<usize>,
443    ) -> impl Stream<Item = DfResult<RecordBatch>> {
444        async_stream::try_stream! {
445            let limit = limit.unwrap_or(usize::MAX);
446            let mut coder = crate::convert::EthRawLogsToArrow::new(&options);
447            let mut total = 0;
448
449            // TODO: Streaming/unbounded API
450            let mut log_stream = Box::pin(
451                crate::stream::RawLogsStream::paginate(rpc_client.clone(), filter, options, None)
452            );
453
454            while let Some(batch) = log_stream.try_next().await.map_err(|e| DataFusionError::External(e.into()))? {
455                if total >= limit {
456                    break;
457                }
458                let to_append = usize::min(batch.logs.len(), limit - total);
459                coder.append(&batch.logs[0..to_append]).map_err(|e| DataFusionError::External(e.into()))?;
460                total += to_append;
461
462                // Don't form batches that are too small
463                if coder.len() > 1_000 {
464                    let batch = if let Some(proj) = &projection {
465                        coder.finish().project(proj)?
466                    } else {
467                        coder.finish()
468                    };
469                    yield batch;
470                }
471            }
472
473            // Return at least one batch, even if empty
474            if !coder.is_empty() || total == 0 {
475                let batch = if let Some(proj) = &projection {
476                    coder.finish().project(proj)?
477                } else {
478                    coder.finish()
479                };
480                yield batch;
481            }
482        }
483    }
484}
485
486impl ExecutionPlan for EthGetLogs {
487    fn name(&self) -> &str {
488        Self::static_name()
489    }
490
491    fn properties(&self) -> &Arc<PlanProperties> {
492        &self.properties
493    }
494
495    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
496        // this is a leaf node and has no children
497        vec![]
498    }
499
500    fn apply_expressions(
501        &self,
502        _f: &mut dyn FnMut(
503            &Arc<dyn datafusion::physical_plan::PhysicalExpr>,
504        ) -> DfResult<TreeNodeRecursion>,
505    ) -> DfResult<TreeNodeRecursion> {
506        Ok(TreeNodeRecursion::Continue)
507    }
508
509    fn with_new_children(
510        self: Arc<Self>,
511        _: Vec<Arc<dyn ExecutionPlan>>,
512    ) -> DfResult<Arc<dyn ExecutionPlan>> {
513        Ok(self)
514    }
515
516    fn execute(
517        &self,
518        partition: usize,
519        _context: Arc<TaskContext>,
520    ) -> DfResult<datafusion::execution::SendableRecordBatchStream> {
521        assert_eq!(partition, 0);
522
523        let stream = Self::execute_impl(
524            self.rpc_client.clone(),
525            self.filter.clone(),
526            self.stream_options.clone(),
527            self.projection.clone(),
528            self.limit,
529        );
530
531        let record_batch_stream =
532            RecordBatchStreamAdapter::new(self.projected_schema.clone(), stream);
533
534        Ok(Box::pin(record_batch_stream))
535    }
536}
537
538impl DisplayAs for EthGetLogs {
539    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
540        match t {
541            DisplayFormatType::Default | DisplayFormatType::Verbose => {
542                write!(f, "EthGetLogs: ")?;
543                if self.projection.is_none() {
544                    write!(f, "projection=[*]")?;
545                } else {
546                    let projection_str = self
547                        .projected_schema
548                        .fields
549                        .iter()
550                        .map(|f| f.name().as_str())
551                        .collect::<Vec<_>>()
552                        .join(", ");
553
554                    write!(f, "projection=[{projection_str}]")?;
555                }
556
557                let mut filters = Vec::new();
558
559                match self.filter.block_option {
560                    FilterBlockOption::Range {
561                        from_block,
562                        to_block,
563                    } => filters.push(format!("block_number=[{from_block:?}, {to_block:?}]")),
564                    FilterBlockOption::AtBlockHash(h) => filters.push(format!("block_hash={h}")),
565                }
566
567                if !self.filter.address.is_empty() {
568                    // Provide deterministic order
569                    let mut addrs: Vec<_> =
570                        self.filter.address.iter().map(|h| h.to_string()).collect();
571                    addrs.sort();
572                    filters.push(format!("address=[{}]", addrs.join(", ")));
573                }
574
575                for (i, t) in self
576                    .filter
577                    .topics
578                    .iter()
579                    .enumerate()
580                    .filter(|(_, t)| !t.is_empty())
581                {
582                    // Provide deterministic order
583                    let mut topics: Vec<_> = t.iter().map(|h| h.to_string()).collect();
584                    topics.sort();
585                    filters.push(format!("topic{}=[{}]", i, topics.join(", ")));
586                }
587
588                write!(f, ", filter=[{}]", filters.join(", "))?;
589
590                if let Some(limit) = self.limit {
591                    write!(f, ", limit={limit}")?;
592                }
593
594                Ok(())
595            }
596            DisplayFormatType::TreeRender => {
597                writeln!(f, "EthGetLogs")?;
598
599                if self.projection.is_none() {
600                    writeln!(f, "projection=[*]")?;
601                } else {
602                    let projection_str = self
603                        .projected_schema
604                        .fields
605                        .iter()
606                        .map(|f| f.name().as_str())
607                        .collect::<Vec<_>>()
608                        .join(", ");
609
610                    writeln!(f, "projection=[{projection_str}]")?;
611                }
612
613                let mut filters = Vec::new();
614
615                match self.filter.block_option {
616                    FilterBlockOption::Range {
617                        from_block,
618                        to_block,
619                    } => filters.push(format!("block_number=[{from_block:?}, {to_block:?}]")),
620                    FilterBlockOption::AtBlockHash(h) => filters.push(format!("block_hash={h}")),
621                }
622
623                if !self.filter.address.is_empty() {
624                    // Provide deterministic order
625                    let mut addrs: Vec<_> =
626                        self.filter.address.iter().map(|h| h.to_string()).collect();
627                    addrs.sort();
628                    filters.push(format!("address=[{}]", addrs.join(", ")));
629                }
630
631                for (i, t) in self
632                    .filter
633                    .topics
634                    .iter()
635                    .enumerate()
636                    .filter(|(_, t)| !t.is_empty())
637                {
638                    // Provide deterministic order
639                    let mut topics: Vec<_> = t.iter().map(|h| h.to_string()).collect();
640                    topics.sort();
641                    filters.push(format!("topic{}=[{}]", i, topics.join(", ")));
642                }
643
644                writeln!(f, "filter=[{}]", filters.join(", "))?;
645
646                if let Some(limit) = self.limit {
647                    writeln!(f, "limit={limit}")?;
648                }
649
650                Ok(())
651            }
652        }
653    }
654}
655
656///////////////////////////////////////////////////////////////////////////////////////////////////
datafusion_ethers/provider.rs

datafusion_ethers/
provider.rs
