1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use crate::connection::Connection;
use arrow::compute::concat_batches;
use arrow::datatypes::{Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use datafusion_common::{DFSchema, ScalarValue};
use datafusion_expr::{
    expr, window_function, BuiltInWindowFunction, Expr, WindowFrame, WindowFrameBound,
    WindowFrameUnits,
};
use std::any::Any;
use std::sync::Arc;
use vegafusion_common::data::table::VegaFusionTable;
use vegafusion_common::error::{Result, ResultWithContext, VegaFusionError};

#[async_trait]
pub trait DataFrame: Send + Sync + 'static {
    fn as_any(&self) -> &dyn Any;

    fn schema(&self) -> Schema;

    fn schema_df(&self) -> Result<DFSchema> {
        Ok(DFSchema::try_from(self.schema())?)
    }

    fn connection(&self) -> Arc<dyn Connection>;

    fn fingerprint(&self) -> u64;

    async fn collect(&self) -> Result<VegaFusionTable>;

    async fn collect_flat(&self) -> Result<RecordBatch> {
        let mut arrow_schema = Arc::new(self.schema()) as SchemaRef;
        let table = self.collect().await?;
        if let Some(batch) = table.batches.get(0) {
            arrow_schema = batch.schema()
        }
        concat_batches(&arrow_schema, table.batches.as_slice())
            .with_context(|| String::from("Failed to concatenate RecordBatches"))
    }

    async fn sort(&self, _expr: Vec<Expr>, _limit: Option<i32>) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported("sort not supported"))
    }

    async fn select(&self, _expr: Vec<Expr>) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported("select not supported"))
    }

    async fn aggregate(
        &self,
        _group_expr: Vec<Expr>,
        _aggr_expr: Vec<Expr>,
    ) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported(
            "aggregate not supported",
        ))
    }

    async fn joinaggregate(
        &self,
        _group_expr: Vec<Expr>,
        _aggr_expr: Vec<Expr>,
    ) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported(
            "joinaggregate not supported",
        ))
    }

    async fn filter(&self, _predicate: Expr) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported("filter not supported"))
    }

    async fn limit(&self, _limit: i32) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported("limit not supported"))
    }

    async fn fold(
        &self,
        _fields: &[String],
        _value_col: &str,
        _key_col: &str,
        _order_field: Option<&str>,
    ) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported("fold not supported"))
    }

    async fn stack(
        &self,
        _field: &str,
        _orderby: Vec<Expr>,
        _groupby: &[String],
        _start_field: &str,
        _stop_field: &str,
        _mode: StackMode,
    ) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported("stack not supported"))
    }

    async fn impute(
        &self,
        _field: &str,
        _value: ScalarValue,
        _key: &str,
        _groupby: &[String],
        _order_field: Option<&str>,
    ) -> Result<Arc<dyn DataFrame>> {
        Err(VegaFusionError::sql_not_supported("impute not supported"))
    }

    async fn with_index(&self, index_name: &str) -> Result<Arc<dyn DataFrame>> {
        if self.schema().column_with_name(index_name).is_some() {
            // Column is already present, don't overwrite
            self.select(vec![Expr::Wildcard]).await
        } else {
            let selections = vec![
                Expr::WindowFunction(expr::WindowFunction {
                    fun: window_function::WindowFunction::BuiltInWindowFunction(
                        BuiltInWindowFunction::RowNumber,
                    ),
                    args: vec![],
                    partition_by: vec![],
                    order_by: vec![],
                    window_frame: WindowFrame {
                        units: WindowFrameUnits::Rows,
                        start_bound: WindowFrameBound::Preceding(ScalarValue::Null),
                        end_bound: WindowFrameBound::CurrentRow,
                    },
                })
                .alias(index_name),
                Expr::Wildcard,
            ];
            self.select(selections).await
        }
    }
}

#[derive(Debug, Clone)]
pub enum StackMode {
    Zero,
    Center,
    Normalize,
}