micromegas_datafusion_extensions/histogram/
expand.rs

1use super::histogram_udaf::HistogramArray;
2use async_trait::async_trait;
3use datafusion::arrow::array::{ArrayRef, Float64Array, StructArray, UInt64Array};
4use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
5use datafusion::arrow::record_batch::RecordBatch;
6use datafusion::catalog::Session;
7use datafusion::catalog::TableFunctionArgs;
8use datafusion::catalog::TableFunctionImpl;
9use datafusion::catalog::TableProvider;
10use datafusion::datasource::TableType;
11use datafusion::datasource::memory::{DataSourceExec, MemorySourceConfig};
12use datafusion::error::DataFusionError;
13use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder};
14use datafusion::physical_plan::ExecutionPlan;
15use datafusion::prelude::Expr;
16use datafusion::scalar::ScalarValue;
17use std::sync::Arc;
18
19/// A DataFusion `TableFunctionImpl` that expands a histogram struct into rows of (bin_center, count).
20///
21/// Usage:
22/// ```sql
23/// SELECT bin_center, count
24/// FROM expand_histogram(
25///   (SELECT make_histogram(0.0, 100.0, 100, value)
26///    FROM measures WHERE name = 'cpu_usage')
27/// )
28/// ```
29#[derive(Debug)]
30pub struct ExpandHistogramTableFunction {}
31
32impl ExpandHistogramTableFunction {
33    pub fn new() -> Self {
34        Self {}
35    }
36}
37
38impl Default for ExpandHistogramTableFunction {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44/// The source of histogram data - either a literal value or a subquery to evaluate.
45#[derive(Debug, Clone)]
46enum HistogramSource {
47    Literal(ScalarValue),
48    Subquery(Arc<LogicalPlan>),
49}
50
51impl TableFunctionImpl for ExpandHistogramTableFunction {
52    fn call_with_args(
53        &self,
54        args: TableFunctionArgs,
55    ) -> datafusion::error::Result<Arc<dyn TableProvider>> {
56        let args = args.exprs();
57        if args.len() != 1 {
58            return Err(DataFusionError::Plan(
59                "expand_histogram requires exactly one argument (a histogram)".into(),
60            ));
61        }
62
63        // Extract the histogram from the expression
64        let source = match &args[0] {
65            Expr::Literal(scalar, _metadata) => HistogramSource::Literal(scalar.clone()),
66            Expr::ScalarSubquery(subquery) => HistogramSource::Subquery(subquery.subquery.clone()),
67            other => {
68                let plan = LogicalPlanBuilder::empty(true)
69                    .project(vec![other.clone()])?
70                    .build()?;
71                HistogramSource::Subquery(Arc::new(plan))
72            }
73        };
74
75        Ok(Arc::new(ExpandHistogramTableProvider { source }))
76    }
77}
78
79fn output_schema() -> SchemaRef {
80    Arc::new(Schema::new(vec![
81        Field::new("bin_center", DataType::Float64, false),
82        Field::new("count", DataType::UInt64, false),
83    ]))
84}
85
86fn expand_histogram_to_batch(
87    histo_array: &HistogramArray,
88    index: usize,
89) -> Result<RecordBatch, DataFusionError> {
90    if histo_array.is_null_at(index) {
91        return Ok(RecordBatch::new_empty(output_schema()));
92    }
93    let start = histo_array.get_start(index)?;
94    let end = histo_array.get_end(index)?;
95    let bins = histo_array.get_bins(index)?;
96
97    let num_bins = bins.len();
98    if num_bins == 0 {
99        return Ok(RecordBatch::new_empty(output_schema()));
100    }
101
102    // Handle edge case where start == end (all values in a single point)
103    let bin_width = if (end - start).abs() < f64::EPSILON {
104        1.0 // Use unit width when range is zero
105    } else {
106        (end - start) / (num_bins as f64)
107    };
108
109    let mut bin_centers = Vec::with_capacity(num_bins);
110    let mut counts = Vec::with_capacity(num_bins);
111
112    for i in 0..num_bins {
113        let bin_center = start + (i as f64 + 0.5) * bin_width;
114        bin_centers.push(bin_center);
115        counts.push(bins.value(i));
116    }
117
118    let bin_center_array: ArrayRef = Arc::new(Float64Array::from(bin_centers));
119    let count_array: ArrayRef = Arc::new(UInt64Array::from(counts));
120
121    RecordBatch::try_new(output_schema(), vec![bin_center_array, count_array])
122        .map_err(|e| DataFusionError::External(e.into()))
123}
124
125fn extract_histogram_from_struct(
126    struct_array: &Arc<StructArray>,
127) -> Result<RecordBatch, DataFusionError> {
128    let histo_array = HistogramArray::new(struct_array.clone());
129    if histo_array.is_empty() {
130        return Ok(RecordBatch::new_empty(output_schema()));
131    }
132    expand_histogram_to_batch(&histo_array, 0)
133}
134
135fn scalar_to_batch(scalar: &ScalarValue) -> Result<RecordBatch, DataFusionError> {
136    match scalar {
137        ScalarValue::Struct(struct_array) => extract_histogram_from_struct(struct_array),
138        ScalarValue::Dictionary(_, inner) => scalar_to_batch(inner.as_ref()),
139        _ => Err(DataFusionError::Plan(format!(
140            "expand_histogram argument must be a struct (histogram), got: {:?}",
141            scalar.data_type()
142        ))),
143    }
144}
145
146/// Table provider for expanding histogram data.
147#[derive(Debug)]
148pub struct ExpandHistogramTableProvider {
149    source: HistogramSource,
150}
151
152impl ExpandHistogramTableProvider {
153    /// Creates a new provider from a histogram scalar value.
154    pub fn from_scalar(scalar: ScalarValue) -> Result<Self, DataFusionError> {
155        if !matches!(scalar, ScalarValue::Struct(_)) {
156            return Err(DataFusionError::Plan(format!(
157                "expand_histogram argument must be a struct (histogram), got: {:?}",
158                scalar.data_type()
159            )));
160        }
161        Ok(Self {
162            source: HistogramSource::Literal(scalar),
163        })
164    }
165}
166
167#[async_trait]
168impl TableProvider for ExpandHistogramTableProvider {
169    fn schema(&self) -> SchemaRef {
170        output_schema()
171    }
172
173    fn table_type(&self) -> TableType {
174        TableType::Temporary
175    }
176
177    async fn scan(
178        &self,
179        state: &dyn Session,
180        projection: Option<&Vec<usize>>,
181        _filters: &[Expr],
182        limit: Option<usize>,
183    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
184        let mut record_batch = match &self.source {
185            HistogramSource::Literal(scalar) => scalar_to_batch(scalar)?,
186            HistogramSource::Subquery(plan) => {
187                // Execute the subquery to get the histogram scalar
188                let physical_plan = state.create_physical_plan(plan).await?;
189                let task_ctx = state.task_ctx();
190                let batches = datafusion::physical_plan::collect(physical_plan, task_ctx).await?;
191
192                if batches.is_empty() || batches[0].num_rows() == 0 {
193                    return Err(DataFusionError::Execution(
194                        "expand_histogram subquery returned no rows".into(),
195                    ));
196                }
197
198                let batch = &batches[0];
199                if batch.num_columns() != 1 {
200                    return Err(DataFusionError::Execution(format!(
201                        "expand_histogram subquery must return exactly one column, got {}",
202                        batch.num_columns()
203                    )));
204                }
205
206                // Extract the struct from the first row
207                let column = batch.column(0);
208                let struct_array = column.as_any().downcast_ref::<StructArray>().ok_or_else(
209                    || {
210                        DataFusionError::Execution(format!(
211                            "expand_histogram subquery must return a struct (histogram), got {:?}",
212                            column.data_type()
213                        ))
214                    },
215                )?;
216
217                let histo_array = HistogramArray::new(Arc::new(struct_array.clone()));
218                if histo_array.is_empty() {
219                    RecordBatch::new_empty(output_schema())
220                } else {
221                    expand_histogram_to_batch(&histo_array, 0)?
222                }
223            }
224        };
225
226        // Apply limit if specified
227        if let Some(n) = limit
228            && n < record_batch.num_rows()
229        {
230            record_batch = record_batch.slice(0, n);
231        }
232
233        let source = MemorySourceConfig::try_new(
234            &[vec![record_batch]],
235            self.schema(),
236            projection.map(|v| v.to_owned()),
237        )?;
238        Ok(DataSourceExec::from_data_source(source))
239    }
240}
micromegas_datafusion_extensions/histogram/expand.rs

micromegas_datafusion_extensions/histogram/
expand.rs
