algeon-datafusion 0.1.1

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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

//! DataFusion statistics translated at the native output occurrence boundary.

use std::sync::Arc;

use arrow::datatypes::SchemaRef;
use datafusion::{
    common::stats::Precision,
    physical_plan::{ExecutionPlan, StatisticsArgs, StatisticsContext},
};
use query_planning::{
    error::QueryError,
    plan::{
        AdvisoryOutputBinding, AdvisoryOutputStatistics, AdvisoryStatistic, AdvisoryUnknownReason, NodeId,
        QueryPlanBuilder,
    },
};

use super::LoweredNativeInput;
use crate::native::normalize::{NormalizedNativeInput, NormalizedNativeLiteralSource};

pub(super) fn output_statistics(plan: &dyn ExecutionPlan) -> Option<AdvisoryOutputStatistics> {
    let statistics = StatisticsContext::new().compute(plan, &StatisticsArgs::new()).ok()?;
    Some(AdvisoryOutputStatistics {
        rows: lower_statistic(statistics.num_rows),
        bytes: lower_statistic(statistics.total_byte_size),
    })
}

fn lower_statistic(precision: Precision<usize>) -> AdvisoryStatistic {
    match precision {
        Precision::Exact(value) => AdvisoryStatistic::Exact(value as u64),
        Precision::Inexact(value) => AdvisoryStatistic::Inexact(value as u64),
        Precision::Absent => AdvisoryStatistic::Absent,
    }
}

pub(super) fn attach(
    builder: &mut QueryPlanBuilder,
    node: NodeId,
    schema: &SchemaRef,
    frontend: &dyn ExecutionPlan,
) -> Result<(), QueryError> {
    let (mut statistics, mut unknown_reason) = match output_statistics(frontend) {
        Some(statistics) => {
            let reason = (statistics.rows == AdvisoryStatistic::Absent
                || statistics.bytes == AdvisoryStatistic::Absent)
                .then_some(AdvisoryUnknownReason::MissingStatistics);
            (statistics, reason)
        },
        None => (absent(), Some(AdvisoryUnknownReason::StatisticsComputationFailed)),
    };
    // Lowering may prune columns or rename them. Renaming preserves byte layout;
    // projection to a different layout does not preserve frontend output bytes.
    let frontend_schema = frontend.schema();
    if schema.fields().len() != frontend_schema.fields().len()
        || schema.fields().iter().zip(frontend_schema.fields()).any(|(native, frontend)| {
            native.data_type() != frontend.data_type() || native.is_nullable() != frontend.is_nullable()
        })
    {
        statistics.bytes = AdvisoryStatistic::Absent;
        unknown_reason = Some(AdvisoryUnknownReason::OutputSchemaChanged);
    }
    builder.set_advisory_output(AdvisoryOutputBinding {
        node,
        schema: Arc::clone(schema),
        frontend: "DataFusion".to_owned(),
        frontend_version: datafusion::DATAFUSION_VERSION.to_owned(),
        statistics,
        unknown_reason,
    })?;
    Ok(())
}

fn absent() -> AdvisoryOutputStatistics {
    AdvisoryOutputStatistics { rows: AdvisoryStatistic::Absent, bytes: AdvisoryStatistic::Absent }
}

pub(super) fn attach_input(
    builder: &mut QueryPlanBuilder,
    input: &NormalizedNativeInput<'_>,
    output: &LoweredNativeInput,
) -> Result<(), QueryError> {
    // Row-count weights carry one row per read unit. Frontend statistics do not
    // describe this input, and no advisory is attached.
    if output.row_count_weights {
        return Ok(());
    }
    let frontend: &dyn ExecutionPlan = match input {
        NormalizedNativeInput::Filter(filter) => filter.filter,
        NormalizedNativeInput::Projection(projection) => projection.projection,
        NormalizedNativeInput::Aggregate(aggregate) => aggregate.statistics_output,
        NormalizedNativeInput::HashJoin(join) => join.join,
        NormalizedNativeInput::NestedLoopJoin(join) => join.join,
        NormalizedNativeInput::CrossJoin(_) => {
            builder.set_advisory_output(AdvisoryOutputBinding {
                node: output.node,
                schema: Arc::clone(&output.schema),
                frontend: "DataFusion".to_owned(),
                frontend_version: datafusion::DATAFUSION_VERSION.to_owned(),
                statistics: absent(),
                unknown_reason: Some(AdvisoryUnknownReason::CartesianFallback),
            })?;
            return Ok(());
        },
        NormalizedNativeInput::Source(source) => {
            if let Some(filter) = source.filter {
                filter
            } else {
                source.scan
            }
        },
        NormalizedNativeInput::LiteralSource(source) => match source {
            NormalizedNativeLiteralSource::PlaceholderRow(plan) => *plan,
            NormalizedNativeLiteralSource::Empty(plan) => *plan,
            NormalizedNativeLiteralSource::Memory(plan) => *plan,
        },
        NormalizedNativeInput::Limit(_)
        | NormalizedNativeInput::GpuResidentSource(_)
        | NormalizedNativeInput::Sort(_)
        | NormalizedNativeInput::Window(_)
        | NormalizedNativeInput::Union(_) => return Ok(()),
    };
    attach(builder, output.node, &output.schema, frontend)
}