Add a runtime error check for `TableProvider::supports_filters_pushdown()` (was for "AND" operators does not work) about arrow-datafusion HOT 10 CLOSED

fn supports_filters_pushdown(
        &self,
        _filters: &[&datafusion::prelude::Expr],
    ) -> datafusion::error::Result<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
        Ok(vec![TableProviderFilterPushDown::Unsupported])
    }

CoalesceBatchesExec {
    input: FilterExec {
        predicate: BinaryExpr {
            left: Column {
                name: "id",
                index: 0,
            },
            op: Gt,
            right: Literal {
                value: Int32(1),
            },
        },
        input: RepartitionExec {
            input: TestExecutionPlan {
                schema: Schema {
                    fields: [
                        Field {
                            name: "id",
                            data_type: Int32,
                            nullable: false,
                            dict_id: 0,
                            dict_is_ordered: false,
                            metadata: {},
                        },
                        Field {
                            name: "title",
                            data_type: Utf8,
                            nullable: false,
                            dict_id: 0,
                            dict_is_ordered: false,
                            metadata: {},
                        },
                    ],
                    metadata: {},
                },
            },
            partitioning: RoundRobinBatch(
                32,
            ),
            state: Mutex {
                data: RepartitionExecState {
                    channels: {},
                    abort_helper: AbortOnDropMany(
                        [],
                    ),
                },
            },
            metrics: ExecutionPlanMetricsSet {
                inner: Mutex {
                    data: MetricsSet {
                        metrics: [],
                    },
                },
            },
            preserve_order: false,
        },
        metrics: ExecutionPlanMetricsSet {
            inner: Mutex {
                data: MetricsSet {
                    metrics: [],
                },
            },
        },
        default_selectivity: 20,
    },
    target_batch_size: 8192,
    metrics: ExecutionPlanMetricsSet {
        inner: Mutex {
            data: MetricsSet {
                metrics: [],
            },
        },
    },
}

use std::{any::Any, fmt, sync::Arc};
use axum::async_trait;
use datafusion::{
    arrow::{
        array::{ArrayRef, Int32Array, RecordBatch, StringArray},
        datatypes::{DataType, Field, Schema, SchemaRef},
        util::pretty::pretty_format_batches,
    },
    catalog::{schema::SchemaProvider, CatalogProvider},
    datasource::{TableProvider, TableType},
    execution::context::{SessionContext, SessionState},
    logical_expr::{Expr, TableProviderFilterPushDown},
    physical_plan::{
        stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionPlan,
    },
};

#[tokio::test]
async fn test_selection_with_and() {
    let ctx = SessionContext::new();
    ctx.register_catalog("test", Arc::new(TestCatalogProvider()));
    let results = ctx
        .sql("SELECT * FROM test.default.test_table WHERE id > 1 AND id < 3")
        .await
        .unwrap()
        .collect()
        .await
        .unwrap();

    assert_eq!(
        &pretty_format_batches(&results).unwrap().to_string(),
        "+----+-------+\n| id | title |\n+----+-------+\n| 2  | b     |\n+----+-------+\n"
    );
}

struct TestCatalogProvider();
impl CatalogProvider for TestCatalogProvider {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
    fn schema_names(&self) -> Vec<String> {
        vec!["default".to_string()]
    }
    fn schema(&self, _name: &str) -> Option<Arc<dyn SchemaProvider>> {
        Some(Arc::new(TestSchemaProvider))
    }
}

struct TestSchemaProvider;

#[async_trait]
impl SchemaProvider for TestSchemaProvider {
    fn as_any(&self) -> &dyn Any {
        self
    }
    fn table_names(&self) -> Vec<String> {
        vec!["test_table".to_string()]
    }
    async fn table(&self, name: &str) -> Option<Arc<dyn TableProvider>> {
        if name == "test_table" {
            Some(Arc::new(TestTableProvider {
                schema: Arc::new(Schema::new(vec![
                    Field::new("id", DataType::Int32, false),
                    Field::new("title", DataType::Utf8, false),
                ])),
            }))
        } else {
            None
        }
    }
    fn table_exist(&self, name: &str) -> bool {
        name == "test_table"
    }
}

struct TestTableProvider {
    schema: SchemaRef,
}

#[async_trait]
impl TableProvider for TestTableProvider {
    fn as_any(&self) -> &dyn Any {
        self
    }
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
    fn table_type(&self) -> TableType {
        TableType::View
    }
    async fn scan(
        &self,
        _state: &SessionState,
        projection: Option<&Vec<usize>>,
        _filters: &[Expr],
        _limit: Option<usize>,
    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
        // schema for execution plan depends on the projection,
        // that's why we need to create a new schema
        let schema = {
            let fields = match projection {
                Some(projection) => projection
                    .iter()
                    .map(|index| self.schema.field(*index).clone())
                    .collect(),
                None => self.schema.fields().clone(),
            };
            Arc::new(Schema::new(fields))
        };

        Ok(Arc::new(TestExecutionPlan { schema }))
    }

    fn supports_filters_pushdown(
        &self,
        _filters: &[&datafusion::prelude::Expr],
    ) -> datafusion::error::Result<Vec<datafusion::logical_expr::TableProviderFilterPushDown>> {
        Ok(vec![TableProviderFilterPushDown::Unsupported])
    }
}

#[derive(Debug)]
struct TestExecutionPlan {
    schema: SchemaRef,
}

impl DisplayAs for TestExecutionPlan {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "TestExecutionPlan")
    }
}

impl ExecutionPlan for TestExecutionPlan {
    fn as_any(&self) -> &dyn Any {
        self
    }
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
    fn output_partitioning(&self) -> datafusion::physical_plan::Partitioning {
        datafusion::physical_plan::Partitioning::UnknownPartitioning(1)
    }
    fn output_ordering(&self) -> Option<&[datafusion::physical_expr::PhysicalSortExpr]> {
        None
    }
    fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
        vec![]
    }
    fn with_new_children(
        self: Arc<Self>,
        _children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
        Ok(self)
    }
    fn execute(
        &self,
        _partition: usize,
        _context: Arc<datafusion::execution::context::TaskContext>,
    ) -> datafusion::error::Result<datafusion::physical_plan::SendableRecordBatchStream> {
        let batch = RecordBatch::try_new(
            self.schema.clone(),
            vec![
                Arc::new(Int32Array::from(vec![1, 2, 3])),
                Arc::new(StringArray::from(vec![
                    "a".to_string(),
                    "b".to_string(),
                    "c".to_string(),
                ])) as ArrayRef,
            ],
        );
        let fut = futures::future::ready(batch.map_err(|e| e.into()));

        let stream = futures::stream::once(fut);

        Ok(Box::pin(RecordBatchStreamAdapter::new(
            self.schema.clone(),
            stream,
        )))
    }
}