ella-engine 0.1.5

Core engine implementation for the ella datastore.
Documentation 
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

use std::sync::Arc;

use arrow_schema::{DataType, Field, Schema};
use datafusion::{
    arrow::{
        array::{ArrayData, ArrayRef, FixedSizeListArray},
        compute,
        record_batch::RecordBatch,
    },
    error::Result as DfResult,
    physical_plan::{
        expressions::{CastExpr, Column},
        projection::ProjectionExec,
        ExecutionPlan, PhysicalExpr,
    },
};

pub fn parquet_compat_schema(schema: Arc<Schema>) -> Option<Arc<Schema>> {
    let fields = schema
        .fields()
        .iter()
        .map(|f| {
            let dt = datatype_to_parquet_compat(f.data_type());
            Arc::new((**f).clone().with_data_type(dt))
        })
        .collect::<Vec<_>>();
    let out = Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()));
    if out != schema {
        Some(out)
    } else {
        None
    }
}

fn datatype_to_parquet_compat(dt: &DataType) -> DataType {
    use DataType::*;
    match dt {
        Duration(_) => Int64,
        &FixedSizeList(ref inner, size) => {
            let inner = Arc::new(
                (**inner)
                    .clone()
                    .with_data_type(datatype_to_parquet_compat(inner.data_type())),
            );
            FixedSizeList(inner, size)
        }
        List(inner) => {
            let inner = Arc::new(
                (**inner)
                    .clone()
                    .with_data_type(datatype_to_parquet_compat(inner.data_type())),
            );
            List(inner)
        }
        LargeList(inner) => {
            let inner = Arc::new(
                (**inner)
                    .clone()
                    .with_data_type(datatype_to_parquet_compat(inner.data_type())),
            );
            LargeList(inner)
        }
        _ => dt.clone(),
    }
}

pub fn cast_batch(batch: &RecordBatch, schema: Arc<Schema>) -> crate::Result<RecordBatch> {
    let columns = batch
        .columns()
        .iter()
        .zip(schema.fields().iter())
        .map(|(col, f)| cast_array_inner(col, f))
        .collect::<Result<Vec<_>, _>>()?;

    Ok(RecordBatch::try_new(schema, columns)?)
}

fn cast_array_inner(col: &ArrayRef, f: &Arc<Field>) -> crate::Result<ArrayRef> {
    if col.data_type() != f.data_type() {
        // cast kernel doesn't support fixed-size lists for some reason, so we have to do this weird hack.
        if let DataType::FixedSizeList(inner, _size) = f.data_type() {
            let values = col
                .as_any()
                .downcast_ref::<FixedSizeListArray>()
                .expect("expected fixed-size list array")
                .values();
            let values = compute::cast(values, inner.data_type())?;
            let data = unsafe {
                ArrayData::builder(f.data_type().clone())
                    .len(col.len())
                    .nulls(col.nulls().cloned())
                    .add_child_data(values.to_data())
                    .build_unchecked()
            };
            Ok(Arc::new(FixedSizeListArray::from(data)))
        } else {
            Ok(compute::cast(col, f.data_type())?)
        }
    } else {
        Ok(col.clone())
    }
}

pub fn cast_batch_plan(
    plan: Arc<dyn ExecutionPlan>,
    src_schema: &Arc<Schema>,
    dst_schema: &Arc<Schema>,
) -> DfResult<Arc<dyn ExecutionPlan>> {
    let columns = src_schema
        .fields()
        .iter()
        .zip(dst_schema.fields().iter())
        .enumerate()
        .map(|(i, (src, dst))| {
            let col = Arc::new(Column::new(dst.name(), i));
            if src == dst {
                (col as Arc<dyn PhysicalExpr>, dst.name().clone())
            } else {
                (
                    Arc::new(CastExpr::new(col, dst.data_type().clone(), None)) as Arc<_>,
                    dst.name().clone(),
                )
            }
        })
        .collect::<Vec<_>>();
    Ok(Arc::new(ProjectionExec::try_new(columns, plan)?))
}