datafusion-tracing 53.0.0

DataFusion tracing of execution plans
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

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.
//
// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2025 Datadog, Inc.

use crate::preview::PreviewFn;
use std::collections::HashMap;
use std::sync::Arc;

/// Configuration options for instrumented execution plans.
#[derive(Clone, Default)]
pub struct InstrumentationOptions {
    /// Whether to record metrics during execution.
    pub record_metrics: bool,

    /// Maximum number of rows to preview per span.
    ///
    /// If set to `0`, batch preview recording will be disabled.
    pub preview_limit: usize,

    /// Optional callback function for formatting previewed record batches.
    ///
    /// The provided function will be invoked for each previewed batch of at most `preview_limit` rows.
    pub preview_fn: Option<Arc<PreviewFn>>,

    /// User-defined custom fields for extensible configuration.
    ///
    /// This can be used to store arbitrary key-value pairs relevant to instrumentation or metadata.
    pub custom_fields: HashMap<String, String>,
}

impl std::fmt::Debug for InstrumentationOptions {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("InstrumentationOptions")
            .field("record_metrics", &self.record_metrics)
            .field("preview_limit", &self.preview_limit)
            .field(
                "preview_fn",
                &if self.preview_fn.is_some() {
                    "Some(PreviewFn)"
                } else {
                    "None"
                },
            )
            .field("custom_fields", &self.custom_fields)
            .finish()
    }
}

impl InstrumentationOptions {
    /// Creates a new builder for `InstrumentationOptions`.
    pub fn builder() -> InstrumentationOptionsBuilder {
        InstrumentationOptionsBuilder::default()
    }
}

/// The builder for `InstrumentationOptions`.
#[derive(Default)]
pub struct InstrumentationOptionsBuilder {
    record_metrics: bool,
    preview_limit: usize,
    preview_fn: Option<Arc<PreviewFn>>,
    custom_fields: HashMap<String, String>,
}

impl InstrumentationOptionsBuilder {
    /// Sets whether to record metrics during execution.
    pub fn record_metrics(mut self, record: bool) -> Self {
        self.record_metrics = record;
        self
    }

    /// Sets the maximum number of rows to preview per span.
    /// Setting this to `0` disables the batch preview.
    pub fn preview_limit(mut self, limit: usize) -> Self {
        self.preview_limit = limit;
        self
    }

    /// Sets the optional callback function for formatting previewed record batches.
    pub fn preview_fn(mut self, func: Arc<PreviewFn>) -> Self {
        self.preview_fn = Some(func);
        self
    }

    /// Adds a single custom field.
    pub fn add_custom_field<K: Into<String>, V: Into<String>>(
        mut self,
        key: K,
        value: V,
    ) -> Self {
        self.custom_fields.insert(key.into(), value.into());
        self
    }

    /// Replaces all custom fields with the provided `HashMap`.
    pub fn custom_fields(mut self, fields: HashMap<String, String>) -> Self {
        self.custom_fields = fields;
        self
    }

    /// Consumes the builder and creates an `InstrumentationOptions` instance.
    pub fn build(self) -> InstrumentationOptions {
        InstrumentationOptions {
            record_metrics: self.record_metrics,
            preview_limit: self.preview_limit,
            preview_fn: self.preview_fn,
            custom_fields: self.custom_fields,
        }
    }
}