diff --git a/Cargo.lock b/Cargo.lock index 1ecb861658bbe..efb99844a473f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -13116,9 +13116,11 @@ dependencies = [ name = "vector-vrl-functions" version = "0.1.0" dependencies = [ + "chrono", "dnstap-parser", "enrichment", "indoc", + "parquet", "vector-vrl-category", "vector-vrl-metrics", "vrl", diff --git a/lib/vector-vrl/functions/Cargo.toml b/lib/vector-vrl/functions/Cargo.toml index 874ffbd759ad6..1e7e88bf6443e 100644 --- a/lib/vector-vrl/functions/Cargo.toml +++ b/lib/vector-vrl/functions/Cargo.toml @@ -13,6 +13,8 @@ enrichment = { path = "../enrichment" } dnstap-parser = { path = "../dnstap-parser", optional = true } vector-vrl-metrics = { path = "../metrics", optional = true } vector-vrl-category.workspace = true +chrono.workspace = true +parquet = { version = "39.0.0", default-features = false, features = ["snap", "flate2", "zstd"] } [features] default = ["dnstap", "vrl-metrics"] diff --git a/lib/vector-vrl/functions/src/lib.rs b/lib/vector-vrl/functions/src/lib.rs index d3846912aaba8..7ea4872550b5c 100644 --- a/lib/vector-vrl/functions/src/lib.rs +++ b/lib/vector-vrl/functions/src/lib.rs @@ -5,6 +5,7 @@ //! - Standard VRL library functions (`vrl::stdlib::all`) //! - Vector-specific functions (`vector_vrl::secret_functions`) //! - Enrichment table functions (`enrichment::vrl_functions`) +//! - Parquet parsing (`parse_parquet`) //! - DNS tap parsing functions (optional, with `dnstap` feature) #![deny(warnings)] @@ -12,6 +13,7 @@ use vrl::{compiler::Function, path::OwnedTargetPath}; pub mod get_secret; +pub mod parse_parquet; pub mod remove_secret; pub mod set_secret; pub mod set_semantic_meaning; @@ -51,7 +53,10 @@ pub fn all_without_vrl_stdlib() -> Vec> { fn iter_all_without_vrl_stdlib() -> impl Iterator> { let functions = secret_functions() .into_iter() - .chain(enrichment::vrl_functions()); + .chain(enrichment::vrl_functions()) + .chain(std::iter::once( + Box::new(parse_parquet::ParseParquet) as Box + )); #[cfg(feature = "dnstap")] let functions = functions.chain(dnstap_parser::vrl_functions()); diff --git a/lib/vector-vrl/functions/src/parse_parquet.rs b/lib/vector-vrl/functions/src/parse_parquet.rs new file mode 100644 index 0000000000000..96aa78cb41333 --- /dev/null +++ b/lib/vector-vrl/functions/src/parse_parquet.rs @@ -0,0 +1,269 @@ +use std::panic::{self, AssertUnwindSafe}; + +use chrono::{DateTime, NaiveDate, Utc}; +use parquet::{ + file::reader::{FileReader, SerializedFileReader}, + record::{Field as PqField, Row}, +}; +use vrl::prelude::*; + +fn parse_parquet(value: Value) -> Resolved { + let bytes = value.try_bytes()?; + // parquet 39's row reader panics instead of returning errors on corrupt pages. + panic::catch_unwind(AssertUnwindSafe(|| read_rows(bytes))) + .unwrap_or_else(|_| Err("unable to parse parquet: malformed page data".into())) +} + +fn read_rows(bytes: Bytes) -> Resolved { + let reader = SerializedFileReader::new(bytes) + .map_err(|err| format!("unable to parse parquet: {err}"))?; + let rows = reader + .get_row_iter(None) + .map_err(|err| format!("unable to parse parquet: {err}"))?; + Ok(Value::Array(rows.map(|row| row_to_value(&row)).collect())) +} + +fn row_to_value(row: &Row) -> Value { + Value::Object( + row.get_column_iter() + .map(|(name, field)| (name.as_str().into(), field_to_value(field))) + .collect(), + ) +} + +fn field_to_value(field: &PqField) -> Value { + match field { + PqField::Null => Value::Null, + PqField::Bool(v) => Value::Boolean(*v), + PqField::Byte(v) => Value::Integer(i64::from(*v)), + PqField::Short(v) => Value::Integer(i64::from(*v)), + PqField::Int(v) => Value::Integer(i64::from(*v)), + PqField::Long(v) => Value::Integer(*v), + PqField::UByte(v) => Value::Integer(i64::from(*v)), + PqField::UShort(v) => Value::Integer(i64::from(*v)), + PqField::UInt(v) => Value::Integer(i64::from(*v)), + // Values above i64::MAX stay lossless as strings. + PqField::ULong(v) => i64::try_from(*v) + .map(Value::Integer) + .unwrap_or_else(|_| Value::from(v.to_string())), + PqField::Float(v) => float_to_value(f64::from(*v)), + PqField::Double(v) => float_to_value(*v), + // Strings keep full precision; floats would not. + PqField::Decimal(_) => Value::from(field.to_string()), + PqField::Str(v) => Value::from(v.as_str()), + PqField::Bytes(v) => Value::Bytes(Bytes::copy_from_slice(v.data())), + PqField::Date(days) => NaiveDate::from_ymd_opt(1970, 1, 1) + .and_then(|epoch| epoch.checked_add_signed(chrono::Duration::days(i64::from(*days)))) + .and_then(|date| date.and_hms_opt(0, 0, 0)) + .map(|dt| Value::Timestamp(dt.and_utc())) + .unwrap_or(Value::Integer(i64::from(*days))), + PqField::TimestampMillis(v) => DateTime::::from_timestamp_millis(*v) + .map(Value::Timestamp) + .unwrap_or(Value::Integer(*v)), + PqField::TimestampMicros(v) => DateTime::::from_timestamp_micros(*v) + .map(Value::Timestamp) + .unwrap_or(Value::Integer(*v)), + PqField::Group(row) => row_to_value(row), + PqField::ListInternal(list) => { + Value::Array(list.elements().iter().map(field_to_value).collect()) + } + PqField::MapInternal(map) => Value::Object( + map.entries() + .iter() + .map(|(key, value)| (map_key(key).into(), field_to_value(value))) + .collect(), + ), + } +} + +fn float_to_value(v: f64) -> Value { + NotNan::new(v).map(Value::Float).unwrap_or(Value::Null) +} + +fn map_key(key: &PqField) -> String { + match key { + PqField::Str(s) => s.clone(), + other => other.to_string(), + } +} + +#[derive(Clone, Copy, Debug)] +pub struct ParseParquet; + +impl Function for ParseParquet { + fn identifier(&self) -> &'static str { + "parse_parquet" + } + + fn usage(&self) -> &'static str { + indoc! {" + Parses the provided `value` as an Apache Parquet file and returns its rows as an + array of objects. Snappy, gzip and zstd column compression are supported. + Timestamps and dates become timestamps, decimals become strings, and NaN + floats become null. + "} + } + + fn category(&self) -> &'static str { + Category::Parse.as_ref() + } + + fn return_kind(&self) -> u16 { + kind::ARRAY + } + + fn parameters(&self) -> &'static [Parameter] { + const PARAMETERS: &[Parameter] = &[Parameter::required( + "value", + kind::BYTES, + "The Parquet file contents to parse.", + )]; + PARAMETERS + } + + fn examples(&self) -> &'static [Example] { + &[example! { + title: "Invalid Parquet", + source: r#"parse_parquet!("not parquet")"#, + result: Err( + r#"function call error for "parse_parquet" at (0:29): unable to parse parquet: Parquet error: Invalid Parquet file. Corrupt footer"#, + ), + }] + } + + fn compile( + &self, + _state: &TypeState, + _ctx: &mut FunctionCompileContext, + arguments: ArgumentList, + ) -> Compiled { + let value = arguments.required("value"); + Ok(ParseParquetFn { value }.as_expr()) + } +} + +#[derive(Debug, Clone)] +struct ParseParquetFn { + value: Box, +} + +impl FunctionExpression for ParseParquetFn { + fn resolve(&self, ctx: &mut Context) -> Resolved { + let value = self.value.resolve(ctx)?; + parse_parquet(value) + } + + fn type_def(&self, _: &TypeState) -> TypeDef { + TypeDef::array(Collection::from_unknown(Kind::object(Collection::any()))).fallible() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use vrl::value; + + const VPC_FLOW_GZIP: &[u8] = include_bytes!("../tests/data/vpc_flow_gzip.parquet"); + const VPC_FLOW_SNAPPY: &[u8] = include_bytes!("../tests/data/vpc_flow_snappy.parquet"); + const TYPES_ZSTD: &[u8] = include_bytes!("../tests/data/types_zstd.parquet"); + + fn parse(bytes: &[u8]) -> Resolved { + parse_parquet(Value::Bytes(Bytes::copy_from_slice(bytes))) + } + + fn ts(s: &str) -> Value { + Value::Timestamp(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc)) + } + + #[test] + fn flat_rows_across_row_groups_gzip() { + assert_vpc_flow_rows(VPC_FLOW_GZIP); + } + + #[test] + fn flat_rows_across_row_groups_snappy() { + assert_vpc_flow_rows(VPC_FLOW_SNAPPY); + } + + fn assert_vpc_flow_rows(bytes: &[u8]) { + let rows = parse(bytes).unwrap(); + let rows = rows.as_array().unwrap(); + assert_eq!(rows.len(), 3); + assert_eq!( + rows[0], + value!({ + "version": 2, + "account_id": "123456789012", + "interface_id": "eni-0a1b2c3d", + "srcaddr": "10.0.0.1", + "dstaddr": "10.0.1.1", + "srcport": 443, + "dstport": 49152, + "protocol": 6, + "packets": 10, + "bytes": 8400, + "start": 1727136000, + "end": 1727136060, + "action": "ACCEPT", + "log_status": "OK", + }) + ); + assert_eq!(rows[1].get("srcport"), Some(&Value::Null)); + assert_eq!(rows[2].get("action"), Some(&value!("REJECT"))); + } + + #[test] + fn type_mapping() { + let rows = parse(TYPES_ZSTD).unwrap(); + let rows = rows.as_array().unwrap(); + assert_eq!(rows.len(), 1); + let row = &rows[0]; + let expected = [ + ("b", value!(true)), + ("i8", value!(-8)), + ("u64", value!("18446744073709551615")), + ("f32", Value::Null), + ("f64", value!(1.5)), + ("dec", value!("-12345.67")), + ("s", value!("hello")), + ("bin", Value::Bytes(Bytes::from_static(b"\x00\xff"))), + ("d", ts("2024-09-24T00:00:00Z")), + ("ts_ms", ts("2024-09-24T01:02:03.456Z")), + ("ts_us", ts("2024-09-24T01:02:03.456789Z")), + ("st", value!({"x": 1, "y": "a"})), + ("lst", value!([1, null, 3])), + ("m", value!({"k1": 1, "k2": 2})), + ("los", value!([{"n": "a"}, {"n": "b"}])), + ]; + assert_eq!(row.as_object().unwrap().len(), expected.len()); + for (name, want) in expected { + assert_eq!(row.get(name), Some(&want), "column {name}"); + } + } + + #[test] + fn truncated_file_errors() { + let err = parse(&VPC_FLOW_GZIP[..VPC_FLOW_GZIP.len() - 1]).unwrap_err(); + assert!( + err.to_string().starts_with("unable to parse parquet"), + "{err}" + ); + } + + #[test] + fn corrupt_page_header_errors_instead_of_panicking() { + // Flipping the first page-header byte panics parquet 39's row reader. + let mut bytes = VPC_FLOW_GZIP.to_vec(); + bytes[4] ^= 0xff; + let err = parse(&bytes).unwrap_err(); + assert_eq!( + err.to_string(), + "unable to parse parquet: malformed page data" + ); + } + + #[test] + fn non_bytes_input_errors() { + assert!(parse_parquet(value!(1)).is_err()); + } +} diff --git a/lib/vector-vrl/functions/tests/data/types_zstd.parquet b/lib/vector-vrl/functions/tests/data/types_zstd.parquet new file mode 100644 index 0000000000000..5c5b0668d9677 Binary files /dev/null and b/lib/vector-vrl/functions/tests/data/types_zstd.parquet differ diff --git a/lib/vector-vrl/functions/tests/data/vpc_flow_gzip.parquet b/lib/vector-vrl/functions/tests/data/vpc_flow_gzip.parquet new file mode 100644 index 0000000000000..ed81aaee0d3ec Binary files /dev/null and b/lib/vector-vrl/functions/tests/data/vpc_flow_gzip.parquet differ diff --git a/lib/vector-vrl/functions/tests/data/vpc_flow_snappy.parquet b/lib/vector-vrl/functions/tests/data/vpc_flow_snappy.parquet new file mode 100644 index 0000000000000..2b28ab46f8ce8 Binary files /dev/null and b/lib/vector-vrl/functions/tests/data/vpc_flow_snappy.parquet differ