// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership

// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

//! View data source which uses a LogicalPlan as it's input.

use std::{borrow::Cow, sync::Arc};

use crate::Session;

use crate::TableProvider;

use arrow::datatypes::SchemaRef;

use async_trait::async_trait;

use datafusion_common::Column;

use datafusion_common::error::Result;

use datafusion_expr::TableType;

use datafusion_expr::{Expr, LogicalPlan};

use datafusion_expr::{LogicalPlanBuilder, TableProviderFilterPushDown};

use datafusion_physical_plan::ExecutionPlan;

use futures::future::BoxFuture;

/// An implementation of `TableProvider` that uses another logical plan.

#[derive(Debug)]

pub struct ViewTable {

/// LogicalPlan of the view

logical_plan: LogicalPlan,

/// File fields + partition columns

table_schema: SchemaRef,

/// SQL used to create the view, if available

definition: Option<String>,

}

impl ViewTable {

/// Create new view that is executed at query runtime.

///

/// Takes a `LogicalPlan` and optionally the SQL text of the `CREATE`

/// statement.

///

/// Notes: the `LogicalPlan` is not validated or type coerced. If this is

/// needed it should be done after calling this function.

pub fn new(logical_plan: LogicalPlan, definition: Option<String>) -> Self {

let table_schema = Arc::clone(logical_plan.schema().inner());

Self {

logical_plan,

table_schema,

definition,

}

}

/// Get definition ref

pub fn definition(&self) -> Option<&String> {

self.definition.as_ref()

}

/// Get logical_plan ref

pub fn logical_plan(&self) -> &LogicalPlan {

&self.logical_plan

}

}

#[async_trait]

impl TableProvider for ViewTable {

fn get_logical_plan(&'_ self) -> Option<Cow<'_, LogicalPlan>> {

Some(Cow::Borrowed(&self.logical_plan))

}

fn schema(&self) -> SchemaRef {

Arc::clone(&self.table_schema)

}

fn table_type(&self) -> TableType {

TableType::View

}

fn get_table_definition(&self) -> Option<&str> {

self.definition.as_deref()

}

fn supports_filters_pushdown(

&self,

filters: &[&Expr],

) -> Result<Vec<TableProviderFilterPushDown>> {

// A filter is added on the View when given

Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])

}

// Hand-written `#[async_trait]` expansion to reduce compile time. See

// <https://github.com/apache/datafusion/issues/13814#issuecomment-5292709677>

fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(

&'life0 self,

state: &'life1 dyn Session,

projection: Option<&'life2 [usize]>,

filters: &'life3 [Expr],

limit: Option<usize>,

) -> BoxFuture<'async_trait, Result<Arc<dyn ExecutionPlan>>>

where

'life0: 'async_trait,

'life1: 'async_trait,

'life2: 'async_trait,

'life3: 'async_trait,

Self: 'async_trait,

{

self.scan_boxed(state, projection, filters, limit)

}

}

impl ViewTable {

fn scan_boxed<'a>(

&'a self,

state: &'a dyn Session,

projection: Option<&'a [usize]>,

filters: &'a [Expr],

limit: Option<usize>,

) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {

Box::pin(self.scan_inner(state, projection, filters, limit))

}

async fn scan_inner(

&self,

state: &dyn Session,

projection: Option<&[usize]>,

filters: &[Expr],

limit: Option<usize>,

) -> Result<Arc<dyn ExecutionPlan>> {

let filter = filters.iter().cloned().reduce(|acc, new| acc.and(new));

let plan = self.logical_plan().clone();

let mut plan = LogicalPlanBuilder::from(plan);

if let Some(filter) = filter {

plan = plan.filter(filter)?;

}

let mut plan = if let Some(projection) = projection {

// avoiding adding a redundant projection (e.g. SELECT * FROM view)

let current_projection =

(0..plan.schema().fields().len()).collect::<Vec<usize>>();

if projection == current_projection {

plan

} else {

let fields: Vec<Expr> = projection

.iter()

.map(|i| {

Expr::Column(Column::from(

self.logical_plan.schema().qualified_field(*i),

))

})

.collect();

plan.project(fields)?

}

} else {

plan

};

if let Some(limit) = limit {

plan = plan.limit(0, Some(limit))?;

}

state.create_physical_plan(&plan.build()?).await

}

}