Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 110 additions & 10 deletions src/catalog/column.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ use std::cmp::{max, min};
use arrow_schema::DataType;
use datafusion::scalar::ScalarValue;
use parquet::file::statistics::Statistics;
use tracing::warn;

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct BoolType {
Expand Down Expand Up @@ -58,33 +59,56 @@ pub enum TypedStatistics {
}

impl TypedStatistics {
pub fn update(self, other: Self) -> Self {
/// Variant name used in logs when the two operands disagree on type.
fn variant_name(&self) -> &'static str {
match self {
TypedStatistics::Bool(_) => "bool",
TypedStatistics::Int(_) => "int",
TypedStatistics::Float(_) => "float",
TypedStatistics::String(_) => "string",
}
}

/// Merge two stat ranges. Returns `None` when the operands disagree on
/// variant (which can happen if the same column was historically written
/// under two different Arrow types — e.g. timestamp vs utf8 — in different
/// parquet files). In that case the caller should drop stats for the
/// column rather than treat them as authoritative; planners then fall
/// back to scanning without min/max pushdown.
pub fn update(self, other: Self) -> Option<Self> {
match (self, other) {
(TypedStatistics::Bool(this), TypedStatistics::Bool(other)) => {
TypedStatistics::Bool(BoolType {
Some(TypedStatistics::Bool(BoolType {
min: min(this.min, other.min),
max: max(this.max, other.max),
})
}))
}
(TypedStatistics::Float(this), TypedStatistics::Float(other)) => {
TypedStatistics::Float(Float64Type {
Some(TypedStatistics::Float(Float64Type {
min: this.min.min(other.min),
max: this.max.max(other.max),
})
}))
}
(TypedStatistics::Int(this), TypedStatistics::Int(other)) => {
TypedStatistics::Int(Int64Type {
Some(TypedStatistics::Int(Int64Type {
min: min(this.min, other.min),
max: max(this.max, other.max),
})
}))
}
(TypedStatistics::String(this), TypedStatistics::String(other)) => {
TypedStatistics::String(Utf8Type {
Some(TypedStatistics::String(Utf8Type {
min: min(this.min, other.min),
max: max(this.max, other.max),
})
}))
}
(this, other) => {
warn!(
"Dropping incompatible column stats: existing={} new={}",
this.variant_name(),
other.variant_name()
);
None
}
_ => panic!("Cannot update wrong types"),
}
}

Expand Down Expand Up @@ -214,3 +238,79 @@ fn int96_to_i64_nanos(int96: &parquet::data_type::Int96) -> i64 {
let days_since_epoch = julian_day - JULIAN_DAY_OF_EPOCH;
days_since_epoch * SECONDS_PER_DAY * NANOS_PER_SECOND + nanos_of_day
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn update_merges_compatible_int_stats() {
let a = TypedStatistics::Int(Int64Type { min: 5, max: 10 });
let b = TypedStatistics::Int(Int64Type { min: 1, max: 7 });
let merged = a.update(b).expect("same variant should merge");
match merged {
TypedStatistics::Int(s) => {
assert_eq!(s.min, 1);
assert_eq!(s.max, 10);
}
_ => panic!("expected Int variant"),
}
}

#[test]
fn update_merges_compatible_string_stats() {
let a = TypedStatistics::String(Utf8Type {
min: "b".into(),
max: "y".into(),
});
let b = TypedStatistics::String(Utf8Type {
min: "a".into(),
max: "z".into(),
});
let merged = a.update(b).expect("same variant should merge");
match merged {
TypedStatistics::String(s) => {
assert_eq!(s.min, "a");
assert_eq!(s.max, "z");
}
_ => panic!("expected String variant"),
}
}

#[test]
fn update_returns_none_on_type_mismatch_instead_of_panicking() {
// Reproduces the "Cannot update wrong types" panic scenario: a column
// that was historically written as Utf8 in some files and Timestamp(ms)
// in others. Timestamps map to Int stats internally.
let str_stats = TypedStatistics::String(Utf8Type {
min: "2025-01-01".into(),
max: "2025-12-31".into(),
});
let int_stats = TypedStatistics::Int(Int64Type {
min: 1_700_000_000_000,
max: 1_800_000_000_000,
});

assert!(str_stats.clone().update(int_stats.clone()).is_none());
assert!(int_stats.update(str_stats).is_none());
}

#[test]
fn update_returns_none_for_each_cross_variant_pair() {
let bool_s = TypedStatistics::Bool(BoolType {
min: false,
max: true,
});
let int_s = TypedStatistics::Int(Int64Type { min: 0, max: 1 });
let float_s = TypedStatistics::Float(Float64Type { min: 0.0, max: 1.0 });
let str_s = TypedStatistics::String(Utf8Type {
min: "a".into(),
max: "b".into(),
});

assert!(bool_s.clone().update(int_s.clone()).is_none());
assert!(int_s.clone().update(float_s.clone()).is_none());
assert!(float_s.clone().update(str_s.clone()).is_none());
assert!(str_s.update(bool_s).is_none());
}
}
2 changes: 1 addition & 1 deletion src/catalog/manifest.rs
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ fn column_statistics(row_groups: &[RowGroupMetaData]) -> HashMap<String, Column>
entry.compressed_size += col.compressed_size() as u64;
entry.uncompressed_size += col.uncompressed_size() as u64;
if let Some(other) = col.statistics().and_then(|stats| stats.try_into().ok()) {
entry.stats = entry.stats.clone().map(|this| this.update(other));
entry.stats = entry.stats.clone().and_then(|this| this.update(other));
}
} else {
columns.insert(
Expand Down
2 changes: 2 additions & 0 deletions src/connectors/kafka/processor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,7 @@ impl ParseableSinkProcessor {
let custom_partition = stream.get_custom_partition();
let static_schema_flag = stream.get_static_schema_flag();
let schema_version = stream.get_schema_version();
let infer_timestamp = stream.get_infer_timestamp();

let mut json_vec = Vec::with_capacity(records.len());
let mut total_payload_size = 0u64;
Expand All @@ -101,6 +102,7 @@ impl ParseableSinkProcessor {
&p_custom_fields,
TelemetryType::Logs,
tenant_id,
infer_timestamp,
)?;

Ok(p_event)
Expand Down
4 changes: 4 additions & 0 deletions src/event/format/json.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ impl EventFormat for Event {
time_partition: Option<&String>,
schema_version: SchemaVersion,
static_schema_flag: bool,
infer_timestamp: bool,
) -> Result<(Self::Data, Vec<Arc<Field>>, bool), anyhow::Error> {
let stream_schema = schema;

Expand Down Expand Up @@ -131,6 +132,7 @@ impl EventFormat for Event {
time_partition,
Some(&value_arr),
schema_version,
infer_timestamp,
);
infer_schema = Schema::new(new_infer_schema.fields().clone());
Schema::try_merge(vec![
Expand Down Expand Up @@ -192,6 +194,7 @@ impl EventFormat for Event {
p_custom_fields: &HashMap<String, String>,
telemetry_type: TelemetryType,
tenant_id: &Option<String>,
infer_timestamp: bool,
) -> Result<super::Event, anyhow::Error> {
let custom_partition_values = match custom_partitions.as_ref() {
Some(custom_partition) => {
Expand All @@ -212,6 +215,7 @@ impl EventFormat for Event {
time_partition,
schema_version,
p_custom_fields,
infer_timestamp,
)?;

Ok(super::Event {
Expand Down
Loading
Loading