algeon-datafusion 0.1.0

DataFusion physical-plan adapter for native GPU execution.
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

//! Hand-built engine query plans for adapter exec-node tests.
//!
//! These plans exercise `AlgeonNativeExec` / `AlgeonNativeComposedExec` /
//! `AlgeonNativeGpuProducerExec` directly, without going through SQL or the
//! optimizer rule. Only tests that need access to the hidden API
//! repair use these builders.

use std::sync::Arc;

use arrow::datatypes::{DataType, Field, Schema};
use datafusion::parquet::file::reader::{FileReader, SerializedFileReader};
use query_planning::{
    plan::{QueryPlan, QueryPlanBuilder},
    source::{GpuResidentSource, LocalParquetFile, LocalParquetSource, SourceNode},
};

pub(super) fn host_sink_plan() -> QueryPlan {
    let mut builder = QueryPlanBuilder::new();
    let source = builder.add_source(SourceNode::LocalParquet(local_parquet_source()));
    builder.finish_with_host_sink(source).unwrap()
}

pub(super) fn gpu_resident_plan() -> QueryPlan {
    let mut builder = QueryPlanBuilder::new();
    let source = builder.add_source(SourceNode::LocalParquet(local_parquet_source()));
    builder.finish_gpu_resident(source).unwrap()
}

pub(super) fn local_parquet_gpu_resident_plan(path: &std::path::Path, schema: Arc<Schema>) -> QueryPlan {
    let source = local_parquet_source_with_exact_metadata(
        path,
        schema,
        vec!["key".to_string(), "payload".to_string(), "m00".to_string(), "m01".to_string()],
    );
    let mut builder = QueryPlanBuilder::new();
    let source = builder.add_source(SourceNode::LocalParquet(source));
    builder.finish_gpu_resident(source).unwrap()
}

pub(super) fn local_parquet_host_sink_plan(path: &std::path::Path, schema: Arc<Schema>) -> QueryPlan {
    let source = local_parquet_source_with_exact_metadata(
        path,
        schema,
        vec!["key".to_string(), "payload".to_string(), "m00".to_string(), "m01".to_string()],
    );
    let mut builder = QueryPlanBuilder::new();
    let source = builder.add_source(SourceNode::LocalParquet(source));
    builder.finish_with_host_sink(source).unwrap()
}

pub(super) fn local_parquet_source_with_exact_metadata(
    path: &std::path::Path,
    schema: Arc<Schema>,
    read_columns: Vec<String>,
) -> LocalParquetSource {
    LocalParquetSource::from_files(vec![exact_local_parquet_file(path)], schema, read_columns)
}

fn exact_local_parquet_file(path: &std::path::Path) -> LocalParquetFile {
    let reader = SerializedFileReader::try_from(path).unwrap_or_else(|error| {
        panic!("written native fixture must be readable as Parquet: {}: {error}", path.display())
    });
    let metadata = reader.metadata();
    let file_size = std::fs::metadata(path)
        .unwrap_or_else(|error| panic!("written native fixture must be stat-able: {}: {error}", path.display()))
        .len();
    let row_group_rows = metadata
        .row_groups()
        .iter()
        .map(|row_group| u64::try_from(row_group.num_rows()).expect("fixture row count fits u64"))
        .collect::<Vec<_>>();
    let row_group_bytes = metadata
        .row_groups()
        .iter()
        .map(|row_group| u64::try_from(row_group.compressed_size()).expect("fixture compressed bytes fit u64"))
        .collect::<Vec<_>>();
    let projected_uncompressed_bytes = metadata
        .row_groups()
        .iter()
        .map(|row_group| u64::try_from(row_group.total_byte_size()).expect("fixture uncompressed bytes fit u64"))
        .collect::<Vec<_>>();
    LocalParquetFile::new(path)
        .with_file_size_bytes(file_size)
        .with_row_groups((0..metadata.num_row_groups()).collect())
        .with_row_group_rows(row_group_rows)
        .with_row_group_bytes(row_group_bytes.clone())
        .with_projected_row_group_bytes(row_group_bytes, projected_uncompressed_bytes)
}

pub(super) fn gpu_resident_host_sink_plan(input_name: &str, schema: Arc<Schema>) -> QueryPlan {
    let mut builder = QueryPlanBuilder::new();
    let source = builder.add_source(SourceNode::GpuResident(GpuResidentSource::new(input_name, schema)));
    builder.finish_with_host_sink(source).unwrap()
}

fn local_parquet_source() -> LocalParquetSource {
    LocalParquetSource::new("/tmp/algeon-datafusion-native-exec-test.parquet", schema(), vec!["key".to_string()])
}

fn schema() -> Arc<Schema> {
    Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)]))
}