diff --git a/src/datatypes/src/series/samples.rs b/src/datatypes/src/series/samples.rs index a19329c62e0..aeccfdf47e3 100644 --- a/src/datatypes/src/series/samples.rs +++ b/src/datatypes/src/series/samples.rs @@ -353,9 +353,16 @@ impl Samples { let shared_batch = histogram_batches .next() .filter(|first| histogram_batches.all(|batch| Arc::ptr_eq(first, batch))); + let needs_float = chunks.iter().any(|chunk| { + matches!( + chunk.storage, + SampleStorage::Float(_) | SampleStorage::Mixed(_) + ) + }); let mut tables: Vec<&dyn Array> = Vec::new(); let mut indices = UInt32Builder::with_capacity(total_len); - let mut floats = arrow::array::Float64Builder::with_capacity(total_len); + let mut floats = + needs_float.then(|| arrow::array::Float64Builder::with_capacity(total_len)); let mut base = 0u32; for chunk in chunks.iter() { let (float_values, histogram) = match &chunk.storage { @@ -372,11 +379,13 @@ impl Samples { tables.push(histogram.batch.as_arrow()); } for position in 0..chunk.len() { - match float_values { - Some(values) if values.is_valid(position) => { - floats.append_value(values.value(position)) + if let Some(floats) = &mut floats { + match float_values { + Some(values) if values.is_valid(position) => { + floats.append_value(values.value(position)); + } + _ => floats.append_null(), } - _ => floats.append_null(), } match histogram { Some(histogram) if histogram.indices.is_valid(position) => { @@ -422,14 +431,8 @@ impl Samples { Arc::new(HistogramBatch::try_from_arrow(struct_array)?) }; let histogram = HistogramSamples::try_new(indices.finish(), batch)?; - let floats = floats.finish(); - if chunks.iter().any(|chunk| { - matches!( - chunk.storage, - SampleStorage::Float(_) | SampleStorage::Mixed(_) - ) - }) { - Self::try_from_mixed(floats, histogram) + if let Some(mut floats) = floats { + Self::try_from_mixed(floats.finish(), histogram) } else { Ok(Self { storage: SampleStorage::Histogram(histogram), diff --git a/src/datatypes/src/series/samples/tests.rs b/src/datatypes/src/series/samples/tests.rs index 75f4bca75cb..3f672ed418a 100644 --- a/src/datatypes/src/series/samples/tests.rs +++ b/src/datatypes/src/series/samples/tests.rs @@ -318,6 +318,34 @@ fn densifying_histogram_and_mixed_chunks_keeps_shared_table() { } assert_eq!(values.histograms.indices().value(3), 0); assert_eq!(values.histograms.indices().value(6), 0); + + let all_null_mixed = Samples::try_from_mixed( + Float64Array::new_null(2), + HistogramSamples::try_new(UInt32Array::from(vec![None, Some(0)]), table.clone()).unwrap(), + ) + .unwrap(); + let dense = Samples::try_from_chunks(vec![ + Samples::try_from_histogram(UInt32Array::from(vec![None]), table.clone()).unwrap(), + all_null_mixed, + ]) + .unwrap() + .into_dense() + .unwrap(); + let DenseSamplesRef::Mixed(values) = dense.dense().unwrap() else { + panic!() + }; + assert_eq!(dense.len(), 3); + assert!(Arc::ptr_eq(values.histograms.batch(), &table)); + assert_eq!(values.floats.null_count(), 3); + assert!(values.histograms.indices().is_null(0)); + assert!(values.histograms.indices().is_null(1)); + assert_eq!(values.histograms.indices().value(2), 0); + assert!(matches!(dense.sample(0), SampleRef::Missing)); + assert!(matches!(dense.sample(1), SampleRef::Missing)); + assert!(same_sample( + dense.sample(2), + SampleRef::Histogram(table.view(0).unwrap()) + )); } #[test]