vineyard_datafusion/ds/
dataframe_test.rs

1// Copyright 2020-2023 Alibaba Group Holding Limited.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#[cfg(test)]
16mod tests {
17    use futures::executor::block_on;
18    use std::sync::Arc;
19
20    use arrow_array as array;
21    use arrow_schema as schema;
22    use datafusion::prelude::*;
23    use spectral::prelude::*;
24
25    use crate::ds::dataframe::{error, ArrowDataFrameBuilder, DataFrame};
26    use vineyard::client::*;
27    use vineyard::{get, put};
28
29    #[test]
30    fn test_polars_dataframe() -> Result<()> {
31        const N: usize = 1024;
32        let mut client = IPCClient::default()?;
33
34        // prepare data
35        let vec0 = (0..N).map(|i| i as i32).collect::<Vec<_>>();
36        let vec1 = (0..N).map(|i| i as f64).collect::<Vec<_>>();
37        let vec2 = (0..N).map(|i| format!("{:?}", i)).collect::<Vec<_>>();
38        let array0 = array::Int32Array::from(vec0);
39        let array1 = array::Float64Array::from(vec1);
40        let array2 = array::LargeStringArray::from(vec2);
41
42        let schema = schema::Schema::new(vec![
43            schema::Field::new("f0", schema::DataType::Int32, false),
44            schema::Field::new("f1", schema::DataType::Float64, false),
45            schema::Field::new("f2", schema::DataType::LargeUtf8, false),
46        ]);
47        let batch = array::RecordBatch::try_new(
48            Arc::new(schema),
49            vec![Arc::new(array0), Arc::new(array1), Arc::new(array2)],
50        )?;
51        let batches = vec![batch];
52
53        let object = put!(&mut client, ArrowDataFrameBuilder, &batches)?;
54        let object_id = object.id();
55        let dataframe = get!(client, DataFrame, object_id)?;
56
57        // test sql on dataframe
58        let ctx = SessionContext::new();
59        let table = ctx.read_table(dataframe.table_provider()).map_err(error)?;
60
61        assert_that!(block_on(table.count()).map_err(error)?).is_equal_to(N);
62        return Ok(());
63    }
64}
vineyard_datafusion/ds/dataframe_test.rs

vineyard_datafusion/ds/
dataframe_test.rs
