algeon-datafusion 0.1.2

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

use std::collections::BTreeMap;

use datafusion::physical_plan::metrics::MetricsSet;

use super::labels::{labels_by_name, optional_usize, optional_usize_label, string_label, usize_label};

const NATIVE_TASK_OPERATOR_ROWS_METRIC_NAME: &str = "native_task_operator_rows";
const NATIVE_TASK_RUNTIME_ROWS_METRIC_NAME: &str = "native_task_runtime_rows";

/// Aggregated task metrics by native operator kind.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct NativeTaskOperatorMetricRow {
    /// Operator-kind label shared by aggregated tasks.
    pub operator_kind: String,
    /// Number of tasks observed for this operator kind.
    pub task_count: usize,
    /// Cumulative elapsed nanoseconds across tasks.
    pub elapsed_ns: usize,
    /// Largest lower-level observed allocator peak for this operator kind.
    pub max_observed_peak_bytes: Option<usize>,
    /// Largest lower-level retained byte count for this operator kind.
    pub max_retained_bytes: Option<usize>,
}

impl NativeTaskOperatorMetricRow {
    /// Serializes this nested row as stable comma-separated key/value pairs.
    pub fn to_stable_kv_row(&self) -> String {
        format!(
            "operator_kind={},task_count={},elapsed_ns={},max_observed_peak_bytes={},max_retained_bytes={}",
            self.operator_kind,
            self.task_count,
            self.elapsed_ns,
            optional_usize(self.max_observed_peak_bytes),
            optional_usize(self.max_retained_bytes)
        )
    }
}

/// Runtime metric row for one compiled pipeline task.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct NativeTaskRuntimeMetricRow {
    /// Stable compiled task id.
    pub compiled_task_id: String,
    /// Operator-kind label for the task.
    pub operator_kind: String,
    /// Cumulative elapsed task time in nanoseconds.
    pub elapsed_ns: usize,
    /// Largest lower-level observed allocator peak for the task.
    pub max_observed_peak_bytes: Option<usize>,
    /// Largest lower-level retained byte count for the task.
    pub max_retained_bytes: Option<usize>,
}

impl NativeTaskRuntimeMetricRow {
    /// Serializes this nested row as stable comma-separated key/value pairs.
    pub fn to_stable_kv_row(&self) -> String {
        format!(
            "compiled_task_id={},operator_kind={},elapsed_ns={},max_observed_peak_bytes={},max_retained_bytes={}",
            self.compiled_task_id,
            self.operator_kind,
            self.elapsed_ns,
            optional_usize(self.max_observed_peak_bytes),
            optional_usize(self.max_retained_bytes)
        )
    }
}

pub(super) fn task_operator_rows(metrics: &MetricsSet) -> Vec<NativeTaskOperatorMetricRow> {
    let mut rows = metrics
        .iter()
        .filter(|metric| metric.value().name() == NATIVE_TASK_OPERATOR_ROWS_METRIC_NAME)
        .filter(|metric| metric.labels().iter().any(|label| label.name() == "operator_kind"))
        .map(|metric| labels_to_task_operator_row(labels_by_name(metric.labels())))
        .collect::<Vec<_>>();
    rows.sort_by(|left, right| left.operator_kind.cmp(&right.operator_kind));
    rows
}

pub(super) fn task_runtime_rows(metrics: &MetricsSet) -> Vec<NativeTaskRuntimeMetricRow> {
    let mut rows = metrics
        .iter()
        .filter(|metric| metric.value().name() == NATIVE_TASK_RUNTIME_ROWS_METRIC_NAME)
        .filter(|metric| metric.labels().iter().any(|label| label.name() == "compiled_task_id"))
        .map(|metric| labels_to_task_runtime_row(labels_by_name(metric.labels())))
        .collect::<Vec<_>>();
    rows.sort_by(|left, right| left.compiled_task_id.cmp(&right.compiled_task_id));
    rows
}

fn labels_to_task_operator_row(labels: BTreeMap<&str, &str>) -> NativeTaskOperatorMetricRow {
    NativeTaskOperatorMetricRow {
        operator_kind: string_label(&labels, "operator_kind"),
        task_count: usize_label(&labels, "task_count"),
        elapsed_ns: usize_label(&labels, "elapsed_ns"),
        max_observed_peak_bytes: optional_usize_label(&labels, "max_observed_peak_bytes"),
        max_retained_bytes: optional_usize_label(&labels, "max_retained_bytes"),
    }
}

fn labels_to_task_runtime_row(labels: BTreeMap<&str, &str>) -> NativeTaskRuntimeMetricRow {
    NativeTaskRuntimeMetricRow {
        compiled_task_id: string_label(&labels, "compiled_task_id"),
        operator_kind: string_label(&labels, "operator_kind"),
        elapsed_ns: usize_label(&labels, "elapsed_ns"),
        max_observed_peak_bytes: optional_usize_label(&labels, "max_observed_peak_bytes"),
        max_retained_bytes: optional_usize_label(&labels, "max_retained_bytes"),
    }
}