Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Stats and timestamp fixes #58

Merged
merged 5 commits into from
Aug 12, 2021
Merged

Stats and timestamp fixes #58

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5131708
Fix stats calculation
xianwill Aug 5, 2021
21fe134
mergefix
xianwill Aug 11, 2021
7f5b706
Cargo fmt and delta-rs bump
mosyp Aug 12, 2021
fe294a3
Fix post merge tests
mosyp Aug 12, 2021
a831800
`Fix post merge tests`
mosyp Aug 12, 2021
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
18 changes: 9 additions & 9 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 5 additions & 3 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,12 +31,14 @@ tokio-stream = { version = "0", features = ["fs"] }
tokio-util = "0.6.3"
uuid = { version = "0.8", features = ["serde", "v4"] }

arrow = { git = "https://github.com/apache/arrow-rs", rev = "6bf1988852f87da21a163203eec4c83a7b692901" }
parquet = { git = "https://github.com/apache/arrow-rs", rev = "6bf1988852f87da21a163203eec4c83a7b692901" }
arrow = { git = "https://github.com/apache/arrow-rs", rev = "fa5acd971c973161f17e69d5c6b50d6e77c7da03" }
parquet = { git = "https://github.com/apache/arrow-rs", rev = "fa5acd971c973161f17e69d5c6b50d6e77c7da03" }

deltalake = { git = "https://github.com/delta-io/delta-rs.git", branch = "writer-map-support", features = ["s3"] }

[dev-dependencies]
utime = "0.3"
serial_test = "*"
tempfile = "3"
parquet = { git = "https://github.com/apache/arrow-rs", rev = "6bf1988852f87da21a163203eec4c83a7b692901", features = ["cli"] }
parquet = { git = "https://github.com/apache/arrow-rs", rev = "fa5acd971c973161f17e69d5c6b50d6e77c7da03", features = ["cli"] }

91 changes: 51 additions & 40 deletions src/deltalake_ext.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,14 @@ use arrow::{
json::reader::Decoder,
record_batch::*,
};
use chrono::prelude::*;
use deltalake::{
action::{Add, ColumnCountStat, ColumnValueStat, Stats},
writer::time_utils::timestamp_to_delta_stats_string,
DeltaDataTypeLong, DeltaDataTypeVersion, DeltaTable, DeltaTableError, DeltaTransactionError,
Schema, StorageBackend, StorageError, UriError,
};
use log::debug;
use parquet::basic::TimestampType;
use parquet::{
arrow::ArrowWriter,
basic::{Compression, LogicalType},
Expand All @@ -35,7 +36,6 @@ use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;

const NANOSECONDS: i64 = 1_000_000_000;
const NULL_PARTITION_VALUE_DATA_PATH: &str = "__HIVE_DEFAULT_PARTITION__";

type MinAndMaxValues = (
Expand Down Expand Up @@ -139,7 +139,12 @@ impl DeltaArrowWriter {
self.arrow_writer.write(&record_batch)?;
self.buffered_record_batch_count += 1;

apply_null_counts(&record_batch.into(), &mut self.null_counts);
apply_null_counts(
partition_columns,
&record_batch.into(),
&mut self.null_counts,
0,
);
Ok(())
}
}
Expand Down Expand Up @@ -416,6 +421,7 @@ pub fn record_batch_from_json(
}

fn min_max_values_from_file_metadata(
partition_values: &HashMap<String, Option<String>>,
file_metadata: &FileMetaData,
) -> Result<MinAndMaxValues, ParquetError> {
let type_ptr = parquet::schema::types::from_thrift(file_metadata.schema.as_slice());
Expand All @@ -442,14 +448,19 @@ fn min_max_values_from_file_metadata(
continue;
}

let column_path = column_descr.path();
let column_path_parts = column_path.parts();

// Do not include partition columns in statistics
if partition_values.contains_key(&column_path_parts[0]) {
continue;
}

let statistics: Vec<&Statistics> = row_group_metadata
.iter()
.filter_map(|g| g.column(i).statistics())
.collect();

let column_path = column_descr.path();
let column_path_parts = column_path.parts();

let _ = apply_min_max_for_column(
statistics.as_slice(),
column_descr.clone(),
Expand All @@ -462,7 +473,12 @@ fn min_max_values_from_file_metadata(
Ok((min_values, max_values))
}

fn apply_null_counts(array: &StructArray, null_counts: &mut HashMap<String, ColumnCountStat>) {
fn apply_null_counts(
partition_columns: &[String],
array: &StructArray,
null_counts: &mut HashMap<String, ColumnCountStat>,
nest_level: i32,
) {
let fields = match array.data_type() {
DataType::Struct(fields) => fields,
_ => unreachable!(),
Expand All @@ -475,6 +491,11 @@ fn apply_null_counts(array: &StructArray, null_counts: &mut HashMap<String, Colu
.for_each(|(column, field)| {
let key = field.name().to_owned();

// Do not include partition columns in statistics
if nest_level == 0 && partition_columns.contains(&key) {
return;
}

match column.data_type() {
// Recursive case
DataType::Struct(_) => {
Expand All @@ -484,7 +505,12 @@ fn apply_null_counts(array: &StructArray, null_counts: &mut HashMap<String, Colu

match col_struct {
ColumnCountStat::Column(map) => {
apply_null_counts(as_struct_array(column), map);
apply_null_counts(
partition_columns,
as_struct_array(column),
map,
nest_level + 1,
);
}
_ => unreachable!(),
}
Expand Down Expand Up @@ -643,26 +669,26 @@ fn min_and_max_from_parquet_statistics(

Ok((min, max))
}
DataType::Int64 if is_timestamp(column_descr.logical_type()) => {
let min_array = as_primitive_array::<arrow::datatypes::Int64Type>(&min_array);
let min = arrow::compute::min(min_array);
let min = min.map(|i| timestamp_ns_to_json_value(i)).flatten();

let max_array = as_primitive_array::<arrow::datatypes::Int64Type>(&max_array);
let max = arrow::compute::max(max_array);
let max = max.map(|i| timestamp_ns_to_json_value(i)).flatten();
Ok((min, max))
}
DataType::Int64 => {
let min_array = as_primitive_array::<arrow::datatypes::Int64Type>(&min_array);
let min = arrow::compute::min(min_array);
let min = min.map(|i| Value::Number(Number::from(i)));

let max_array = as_primitive_array::<arrow::datatypes::Int64Type>(&max_array);
let max = arrow::compute::max(max_array);
let max = max.map(|i| Value::Number(Number::from(i)));

Ok((min, max))
match column_descr.logical_type().as_ref() {
Some(LogicalType::TIMESTAMP(TimestampType { unit, .. })) => {
let min = min.map(|n| Value::String(timestamp_to_delta_stats_string(n, &unit)));
let max = max.map(|n| Value::String(timestamp_to_delta_stats_string(n, &unit)));

Ok((min, max))
}
_ => {
let min = min.map(|i| Value::Number(Number::from(i)));
let max = max.map(|i| Value::Number(Number::from(i)));

Ok((min, max))
}
}
}
DataType::Float32 => {
let min_array = as_primitive_array::<arrow::datatypes::Float32Type>(&min_array);
Expand Down Expand Up @@ -706,14 +732,6 @@ fn is_utf8(opt: Option<LogicalType>) -> bool {
}
}

#[inline]
fn is_timestamp(opt: Option<LogicalType>) -> bool {
match opt.as_ref() {
Some(LogicalType::TIMESTAMP(_)) => true,
_ => false,
}
}

fn min_max_strings_from_stats(
stats_with_min_max: &Vec<&Statistics>,
) -> (Option<Value>, Option<Value>) {
Expand All @@ -736,14 +754,6 @@ fn min_max_strings_from_stats(
return (min_value, max_value);
}

#[inline]
pub fn timestamp_ns_to_json_value(ns: i64) -> Option<Value> {
match Utc.timestamp_opt(ns / NANOSECONDS, (ns % NANOSECONDS) as u32) {
chrono::offset::LocalResult::Single(dt) => Some(Value::String(format!("{:?}", dt))),
_ => None,
}
}

fn arrow_array_from_bytes(
data_type: DataType,
capacity: usize,
Expand Down Expand Up @@ -771,7 +781,8 @@ fn create_add(
size: i64,
file_metadata: &FileMetaData,
) -> Result<Add, DeltaWriterError> {
let (min_values, max_values) = min_max_values_from_file_metadata(file_metadata)?;
let (min_values, max_values) =
min_max_values_from_file_metadata(partition_values, file_metadata)?;

let stats = Stats {
num_records: file_metadata.num_rows,
Expand Down Expand Up @@ -878,7 +889,7 @@ mod tests {
assert_eq!(add.len(), 1);
let stats = add[0].get_stats().unwrap().unwrap();

let min_max_keys = vec!["meta", "some_int", "some_string", "some_bool", "date"];
let min_max_keys = vec!["meta", "some_int", "some_string", "some_bool"];
let mut null_count_keys = vec!["some_list", "some_nested_list"];
null_count_keys.extend_from_slice(min_max_keys.as_slice());

Expand Down
12 changes: 7 additions & 5 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use tokio_util::sync::CancellationToken;

pub mod deltalake_ext;
pub mod instrumentation;
mod transforms;
pub mod transforms;

use crate::transforms::*;
use crate::{
Expand Down Expand Up @@ -592,14 +592,16 @@ impl KafkaJsonToDelta {
return Ok(());
}
Err(e) => match e {
DeltaTransactionError::VersionAlreadyExists { .. }
if attempt_number > DEFAULT_DELTA_MAX_RETRY_COMMIT_ATTEMPTS + 1 =>
{
DeltaTableError::TransactionError {
source: DeltaTransactionError::VersionAlreadyExists { .. },
} if attempt_number > DEFAULT_DELTA_MAX_RETRY_COMMIT_ATTEMPTS + 1 => {
debug!("Transaction attempt failed. Attempts exhausted beyond max_retry_commit_attempts of {} so failing.", DEFAULT_DELTA_MAX_RETRY_COMMIT_ATTEMPTS);
self.log_delta_write_failed().await;
return Err(e.into());
}
DeltaTransactionError::VersionAlreadyExists { .. } => {
DeltaTableError::TransactionError {
source: DeltaTransactionError::VersionAlreadyExists { .. },
} => {
attempt_number += 1;
debug!("Transaction attempt failed. Incrementing attempt number to {} and retrying.", attempt_number);
}
Expand Down
10 changes: 3 additions & 7 deletions src/transforms.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -373,15 +373,11 @@ fn timestamp_value_from_kafka(
match kafka_timestamp {
rdkafka::Timestamp::NotAvailable => serde_json::to_value(None as Option<String>),
// Convert to milliseconds to microseconds for delta format
rdkafka::Timestamp::CreateTime(ms) => serde_json::to_value(millis_to_nanos(ms)),
rdkafka::Timestamp::LogAppendTime(ms) => serde_json::to_value(millis_to_nanos(ms)),
rdkafka::Timestamp::CreateTime(ms) => serde_json::to_value(ms * 1000),
rdkafka::Timestamp::LogAppendTime(ms) => serde_json::to_value(ms * 1000),
}
}

fn millis_to_nanos(ms: i64) -> i64 {
ms * 1_000_000
}

#[cfg(test)]
mod tests {
use rdkafka::message::OwnedMessage;
Expand Down Expand Up @@ -650,7 +646,7 @@ mod tests {
assert_eq!(0i64, kafka_offset);
assert_eq!(0i64, kafka_partition);
assert_eq!("test", kafka_topic);
assert_eq!(1626823098519000000, kafka_timestamp);
assert_eq!(1626823098519000, kafka_timestamp);
assert_eq!(0, kafka_timestamp_type);
}
}
6 changes: 3 additions & 3 deletions tests/integration_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,9 +157,9 @@ async fn e2e_smoke_test() {
assert_eq!(row_num, r.get_int(1).unwrap());
assert_eq!("2021-03-01T14:38:58Z", r.get_string(2).unwrap());
assert_eq!("2021-03-01", r.get_string(3).unwrap());
let kafka_timestamp = r.get_long(4).unwrap();
assert!(start_time.timestamp_nanos() < kafka_timestamp);
assert!(chrono::Utc::now().timestamp_nanos() > kafka_timestamp);
let kafka_timestamp = r.get_timestamp_micros(4).unwrap() as i64;
assert!(start_time.timestamp_millis() * 1000 < kafka_timestamp);
assert!(chrono::Utc::now().timestamp_millis() * 1000 > kafka_timestamp);
assert_eq!(0, r.get_int(5).unwrap());
}

Expand Down
Loading