diff --git a/Cargo.lock b/Cargo.lock index f656897525541..da30a38df43f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -374,19 +374,39 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" dependencies = [ - "arrow-arith", - "arrow-array", - "arrow-buffer", - "arrow-cast", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", "arrow-csv", - "arrow-data", - "arrow-ipc", - "arrow-json", - "arrow-ord", - "arrow-row", - "arrow-schema", - "arrow-select", - "arrow-string", + "arrow-data 58.3.0", + "arrow-ipc 58.3.0", + "arrow-json 58.3.0", + "arrow-ord 58.3.0", + "arrow-row 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", + "arrow-string 58.3.0", +] + +[[package]] +name = "arrow" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b952ca5a8046ad741b60f142d6eca4aeebcad615694202bc64c5341f23e32c5b" +dependencies = [ + "arrow-arith 59.1.0", + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-cast 59.1.0", + "arrow-data 59.1.0", + "arrow-ipc 59.1.0", + "arrow-json 59.1.0", + "arrow-ord 59.1.0", + "arrow-row 59.1.0", + "arrow-schema 59.1.0", + "arrow-select 59.1.0", + "arrow-string 59.1.0", ] [[package]] @@ -395,10 +415,24 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "num-traits", +] + +[[package]] +name = "arrow-arith" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64a13b8d3008c4e9063c597a08f46446fe3fd5789277127672d6c0bdbb43b1ff" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-schema 59.1.0", "chrono", "num-traits", ] @@ -410,9 +444,28 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" dependencies = [ "ahash 0.8.11", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "chrono", + "chrono-tz", + "half", + "hashbrown 0.17.1", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-array" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9486151b2f0785bafc6fa04fc5c99fcb4495455662e58787ea32eaaed33c4192" +dependencies = [ + "ahash 0.8.11", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-schema 59.1.0", "chrono", "chrono-tz", "half", @@ -434,18 +487,30 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrow-buffer" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4776577a87794bfdf0b4e90e2ea12454fa7738ea2823c4be5b9d1851da7b434" +dependencies = [ + "bytes", + "half", + "num-bigint", + "num-traits", +] + [[package]] name = "arrow-cast" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-ord", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "atoi", "base64 0.22.1", "chrono", @@ -456,15 +521,36 @@ dependencies = [ "ryu", ] +[[package]] +name = "arrow-cast" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9ad451ce4f98710828a455b96991b8f031deb2e67f5fcad6773f017e4a69c3a" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-ord 59.1.0", + "arrow-schema 59.1.0", + "arrow-select 59.1.0", + "atoi", + "base64 0.22.1", + "chrono", + "half", + "lexical-core", + "num-traits", + "ryu", +] + [[package]] name = "arrow-csv" version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e94e8cf7e517657a52b91ea1263acf38c4ca62a84655d72458a3359b12ab97de" dependencies = [ - "arrow-array", - "arrow-cast", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-cast 58.3.0", + "arrow-schema 58.3.0", "chrono", "csv", "csv-core", @@ -477,8 +563,21 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" dependencies = [ - "arrow-buffer", - "arrow-schema", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-data" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38fe43e2e8704360f1464e6e8cc4fc381ef02cc4fb0192afa8df1aaa0115c66" +dependencies = [ + "arrow-buffer 59.1.0", + "arrow-schema 59.1.0", "half", "num-integer", "num-traits", @@ -490,11 +589,32 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28abfe8bf9f124e5fc83b334af4fa58f8d0323ad25312ccb2d1da50178415704" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-ipc", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-ipc 58.3.0", + "arrow-schema 58.3.0", + "base64 0.22.1", + "bytes", + "futures", + "prost 0.14.3", + "prost-types 0.14.3", + "tonic 0.14.5", + "tonic-prost", +] + +[[package]] +name = "arrow-flight" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42115e09dbb694b5955da998912121451c6910b338228cb80a5701370dba43ff" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-cast 59.1.0", + "arrow-data 59.1.0", + "arrow-ipc 59.1.0", + "arrow-schema 59.1.0", "base64 0.22.1", "bytes", "futures", @@ -510,11 +630,27 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", + "flatbuffers 25.9.23", + "lz4_flex", + "zstd 0.13.2", +] + +[[package]] +name = "arrow-ipc" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29dac499fcbc6ba74ee0324057821d381929a48526a3966bd9dffb44aa06d98c" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-schema 59.1.0", + "arrow-select 59.1.0", "flatbuffers 25.9.23", "lz4_flex", "zstd 0.13.2", @@ -526,12 +662,37 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-ord", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", + "chrono", + "half", + "indexmap 2.12.0", + "itoa", + "lexical-core", + "memchr", + "num-traits", + "ryu", + "serde_core", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-json" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe05e916ddc50f4c7a363cd69c0ef5894fcee063517e9a0b8582f0c56746af6" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-cast 59.1.0", + "arrow-ord 59.1.0", + "arrow-schema 59.1.0", + "arrow-select 59.1.0", "chrono", "half", "indexmap 2.12.0", @@ -551,11 +712,24 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", +] + +[[package]] +name = "arrow-ord" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e13dbdc2a9c053c10c7baa6e30faee04a180aa7ce88e471835850ce37abd20b" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-schema 59.1.0", + "arrow-select 59.1.0", ] [[package]] @@ -564,10 +738,23 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "half", +] + +[[package]] +name = "arrow-row" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d5a1f8c733d15260b305683472ee8ad89c62cbd706703ca873b90d051b41592" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-schema 59.1.0", "half", ] @@ -581,6 +768,16 @@ dependencies = [ "serde_core", ] +[[package]] +name = "arrow-schema" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9e4969dc350d571766247143ab36a5187d095d3d3690970408bc630d47c69e5" +dependencies = [ + "serde_core", + "serde_json", +] + [[package]] name = "arrow-select" version = "58.3.0" @@ -588,10 +785,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" dependencies = [ "ahash 0.8.11", - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "num-traits", +] + +[[package]] +name = "arrow-select" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "402770dba90865359d98d1ef92ef16e23d75c0cca9c2c880c8a05468b7743bf9" +dependencies = [ + "ahash 0.8.11", + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-schema 59.1.0", "num-traits", ] @@ -601,11 +812,28 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", + "memchr", + "num-traits", + "regex", + "regex-syntax", +] + +[[package]] +name = "arrow-string" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b0afbb8b9016700938291123df30838b89decc3213dba00852021988b170d3" +dependencies = [ + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-schema 59.1.0", + "arrow-select 59.1.0", "memchr", "num-traits", "regex", @@ -2578,7 +2806,7 @@ name = "codecs" version = "0.1.0" dependencies = [ "apache-avro", - "arrow", + "arrow 59.1.0", "async-trait", "bytes", "chrono", @@ -2589,6 +2817,7 @@ dependencies = [ "dyn-clone", "flate2", "futures", + "hex", "indoc", "influxdb-line-protocol", "memchr", @@ -2596,6 +2825,8 @@ dependencies = [ "opentelemetry-proto", "ordered-float 4.6.0", "parquet", + "parquet-variant", + "parquet-variant-json", "pin-project", "prost 0.12.6", "prost-reflect", @@ -3409,14 +3640,14 @@ dependencies = [ [[package]] name = "databricks-zerobus-ingest-sdk" -version = "2.3.1" +version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6b258735a4a1c5bb6008eaafd069505d7c2f540e2a4c8c6e944d66de8bb4bca" +checksum = "3d70a18f4f33f665aba0861e4d90d70af24f3c6a91e885c9faaf1af3539186b5" dependencies = [ - "arrow-array", - "arrow-flight", - "arrow-ipc", - "arrow-schema", + "arrow-array 59.1.0", + "arrow-flight 59.1.0", + "arrow-ipc 59.1.0", + "arrow-schema 59.1.0", "async-trait", "bytes", "futures", @@ -4701,11 +4932,11 @@ version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a5ee8f8f967c2b8183136920bdb0a36d06618cc9402f901495c0da66cb2cc74b" dependencies = [ - "arrow", - "arrow-array", - "arrow-flight", - "arrow-ipc", - "arrow-schema", + "arrow 58.3.0", + "arrow-array 58.3.0", + "arrow-flight 58.3.0", + "arrow-ipc 58.3.0", + "arrow-schema 58.3.0", "async-stream", "async-trait", "base64 0.22.1", @@ -5863,12 +6094,6 @@ dependencies = [ "syn 2.0.117", ] -[[package]] -name = "integer-encoding" -version = "3.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" - [[package]] name = "inventory" version = "0.3.24" @@ -7978,17 +8203,17 @@ dependencies = [ [[package]] name = "parquet" -version = "58.3.0" +version = "59.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" +checksum = "5302d4da74d6596a1f11f9928767995b53bca657cbeea1e4e8c5074f8a1157dd" dependencies = [ "ahash 0.8.11", - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-ipc", - "arrow-schema", - "arrow-select", + "arrow-array 59.1.0", + "arrow-buffer 59.1.0", + "arrow-data 59.1.0", + "arrow-ipc 59.1.0", + "arrow-schema 59.1.0", + "arrow-select 59.1.0", "base64 0.22.1", "bytes", "chrono", @@ -8002,11 +8227,40 @@ dependencies = [ "paste", "seq-macro", "snap", - "thrift", "twox-hash", "zstd 0.13.2", ] +[[package]] +name = "parquet-variant" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57db04c7f8598c9e12c935a7e953a441009189da3f1d504b63c5dc3346a4d59b" +dependencies = [ + "arrow 59.1.0", + "arrow-schema 59.1.0", + "chrono", + "half", + "indexmap 2.12.0", + "num-traits", + "simdutf8", + "uuid", +] + +[[package]] +name = "parquet-variant-json" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5fc56a025dc3d5537140fc95d80879750fc1aa4324cf2452f3d1471854c2f8d" +dependencies = [ + "arrow-schema 59.1.0", + "base64 0.22.1", + "chrono", + "parquet-variant", + "serde_json", + "uuid", +] + [[package]] name = "parse-size" version = "1.1.0" @@ -11684,17 +11938,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "thrift" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" -dependencies = [ - "byteorder", - "integer-encoding", - "ordered-float 2.10.1", -] - [[package]] name = "tikv-jemalloc-sys" version = "0.7.1+5.3.1-0-g81034ce1f1373e37dc865038e1bc8eeecf559ce8" @@ -12957,8 +13200,8 @@ dependencies = [ "approx", "arc-swap", "arr_macro", - "arrow", - "arrow-schema", + "arrow 59.1.0", + "arrow-schema 59.1.0", "assert_cmd", "async-compression", "async-nats", diff --git a/Cargo.toml b/Cargo.toml index 285aaa9456b0b..6abec980ce5a6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -352,7 +352,7 @@ prost-reflect = { workspace = true, optional = true } prost-types = { workspace = true, optional = true } # Databricks Zerobus -databricks-zerobus-ingest-sdk = { version = "2.3.1", optional = true, features = ["arrow-flight"] } +databricks-zerobus-ingest-sdk = { version = "2.4.0", optional = true, features = ["arrow-flight"] } # GCP goauth = { version = "0.16.0", optional = true } @@ -373,9 +373,9 @@ greptimedb-ingester = { version = "0.17.0", default-features = false, optional = # External libs arc-swap = { workspace = true, default-features = false, optional = true } async-compression = { workspace = true, features = ["zstd"], optional = true } -arrow = { version = "58.2.0", default-features = false, features = ["ipc"], optional = true } -arrow-schema = { version = "58.2.0", default-features = false, optional = true } -parquet = { version = "58.2.0", default-features = false, features = [ +arrow = { version = "59.1.0", default-features = false, features = ["ipc"], optional = true } +arrow-schema = { version = "59.1.0", default-features = false, optional = true } +parquet = { version = "59.1.0", default-features = false, features = [ "arrow", "snap", "zstd", diff --git a/LICENSE-3rdparty.csv b/LICENSE-3rdparty.csv index 0ccbde72aafbe..33401fd7d3e6b 100644 --- a/LICENSE-3rdparty.csv +++ b/LICENSE-3rdparty.csv @@ -390,7 +390,6 @@ inotify,https://github.com/hannobraun/inotify,ISC,"Hanno Braun inout,https://github.com/RustCrypto/utils,MIT OR Apache-2.0,RustCrypto Developers instability,https://github.com/ratatui-org/instability,MIT,"Stephen M. Coakley , Joshka" -integer-encoding,https://github.com/dermesser/integer-encoding-rs,MIT,Lewin Bormann inventory,https://github.com/dtolnay/inventory,MIT OR Apache-2.0,David Tolnay ipconfig,https://github.com/liranringel/ipconfig,MIT OR Apache-2.0,Liran Ringel ipcrypt-rs,https://github.com/jedisct1/rust-ipcrypt2,ISC,Frank Denis @@ -549,6 +548,8 @@ parking,https://github.com/smol-rs/parking,Apache-2.0 OR MIT,"Stjepan Glavina parking_lot_core,https://github.com/Amanieu/parking_lot,MIT OR Apache-2.0,Amanieu d'Antras parquet,https://github.com/apache/arrow-rs,Apache-2.0,Apache Arrow +parquet-variant,https://github.com/apache/arrow-rs,Apache-2.0,Apache Arrow +parquet-variant-json,https://github.com/apache/arrow-rs,Apache-2.0,Apache Arrow parse-size,https://github.com/kennytm/parse-size,MIT,kennytm paste,https://github.com/dtolnay/paste,MIT OR Apache-2.0,David Tolnay pastey,https://github.com/as1100k/pastey,MIT OR Apache-2.0,"Aditya Kumar , David Tolnay " @@ -777,7 +778,6 @@ terminal_size,https://github.com/eminence/terminal-size,MIT OR Apache-2.0,Andrew thiserror,https://github.com/dtolnay/thiserror,MIT OR Apache-2.0,David Tolnay thiserror-impl,https://github.com/dtolnay/thiserror,MIT OR Apache-2.0,David Tolnay thread_local,https://github.com/Amanieu/thread_local-rs,MIT OR Apache-2.0,Amanieu d'Antras -thrift,https://github.com/apache/thrift/tree/master/lib/rs,Apache-2.0,Apache Thrift Developers tikv-jemalloc-sys,https://github.com/tikv/jemallocator,MIT OR Apache-2.0,"Alex Crichton , Gonzalo Brito Gadeschi , The TiKV Project Developers" tikv-jemallocator,https://github.com/tikv/jemallocator,MIT OR Apache-2.0,"Alex Crichton , Gonzalo Brito Gadeschi , Simon Sapin , Steven Fackler , The TiKV Project Developers" time,https://github.com/time-rs/time,MIT OR Apache-2.0,"Jacob Pratt , Time contributors" diff --git a/changelog.d/zerobus_ingest_sdk_2_4_0.enhancement.md b/changelog.d/zerobus_ingest_sdk_2_4_0.enhancement.md new file mode 100644 index 0000000000000..69c52280cecb7 --- /dev/null +++ b/changelog.d/zerobus_ingest_sdk_2_4_0.enhancement.md @@ -0,0 +1,3 @@ +The `databricks_zerobus` sink now supports Unity Catalog `VARIANT` columns. + +authors: yorickvanzweeden diff --git a/lib/codecs/Cargo.toml b/lib/codecs/Cargo.toml index cb9bad72b67d5..27a0e50a640f9 100644 --- a/lib/codecs/Cargo.toml +++ b/lib/codecs/Cargo.toml @@ -14,14 +14,16 @@ path = "tests/bin/generate-avro-fixtures.rs" [dependencies] apache-avro = { version = "0.20.0", default-features = false } -arrow = { version = "58.2.0", default-features = false, features = ["ipc", "json"], optional = true } -parquet = { version = "58.2.0", default-features = false, features = [ +arrow = { version = "59.1.0", default-features = false, features = ["ipc", "json"], optional = true } +parquet = { version = "59.1.0", default-features = false, features = [ "arrow", "snap", "zstd", "lz4", "flate2-rust_backened", ], optional = true } +parquet-variant = { version = "59.1.0", default-features = false, optional = true } +parquet-variant-json = { version = "59.1.0", default-features = false, optional = true } async-trait.workspace = true bytes.workspace = true chrono.workspace = true @@ -33,6 +35,7 @@ derive_more = { version = "2.1.1", optional = true, features = ["from", "display dyn-clone = { version = "1", default-features = false } flate2.workspace = true futures.workspace = true +hex = { version = "0.4.3", default-features = false, optional = true } influxdb-line-protocol = { version = "2", default-features = false } lookup = { package = "vector-lookup", path = "../vector-lookup", default-features = false, features = ["test"] } memchr = { version = "2", default-features = false } @@ -77,7 +80,13 @@ uuid.workspace = true vrl.workspace = true [features] -arrow = ["dep:arrow", "arrow/chrono-tz"] +arrow = [ + "dep:arrow", + "dep:hex", + "dep:parquet-variant", + "dep:parquet-variant-json", + "arrow/chrono-tz", +] parquet = ["dep:parquet", "arrow"] opentelemetry = ["dep:opentelemetry-proto"] syslog = ["dep:syslog_loose", "dep:strum", "dep:derive_more", "dep:serde-aux", "dep:toml"] diff --git a/lib/codecs/src/encoding/format/arrow.rs b/lib/codecs/src/encoding/format/arrow.rs index 1476f7600c7f7..87a745fe34b85 100644 --- a/lib/codecs/src/encoding/format/arrow.rs +++ b/lib/codecs/src/encoding/format/arrow.rs @@ -4,6 +4,8 @@ //! This implements the streaming variant of the Arrow IPC protocol, which writes //! a continuous stream of record batches without a file footer. +use std::{collections::HashSet, sync::Arc}; + use arrow::{ datatypes::{DataType, Field, Fields, Schema, SchemaRef}, error::ArrowError, @@ -13,6 +15,8 @@ use arrow::{ }; use async_trait::async_trait; use bytes::{BufMut, Bytes, BytesMut}; +use parquet_variant::VariantBuilder; +use parquet_variant_json::JsonToVariant; use snafu::{ResultExt, Snafu, ensure}; use vector_config::configurable_component; use vector_core::event::Event; @@ -49,6 +53,15 @@ pub struct ArrowStreamSerializerConfig { #[serde(default)] #[configurable(derived)] pub allow_nullable_fields: bool, + + /// Arrow field paths whose physical struct representation contains Parquet VARIANT data. + /// + /// This is populated programmatically by schema providers that retain logical type + /// information. It is not user-configurable because the physical Arrow shape alone is not + /// sufficient to distinguish VARIANT from an ordinary struct. + #[serde(skip)] + #[configurable(derived)] + pub variant_fields: Vec>, } impl std::fmt::Debug for ArrowStreamSerializerConfig { @@ -62,6 +75,7 @@ impl std::fmt::Debug for ArrowStreamSerializerConfig { .map(|s| format!("{} fields", s.fields().len())), ) .field("allow_nullable_fields", &self.allow_nullable_fields) + .field("variant_fields", &self.variant_fields) .finish() } } @@ -72,9 +86,16 @@ impl ArrowStreamSerializerConfig { Self { schema: Some(schema), allow_nullable_fields: false, + variant_fields: Vec::new(), } } + /// Mark logical Parquet VARIANT fields by their Arrow schema paths. + pub fn with_variant_fields(mut self, variant_fields: Vec>) -> Self { + self.variant_fields = variant_fields; + self + } + /// The data type of events that are accepted by `ArrowStreamEncoder`. pub fn input_type(&self) -> vector_core::config::DataType { vector_core::config::DataType::Log @@ -90,6 +111,7 @@ impl ArrowStreamSerializerConfig { #[derive(Clone, Debug)] pub struct ArrowStreamSerializer { schema: SchemaRef, + variant_fields: Arc>>, } impl ArrowStreamSerializer { @@ -103,12 +125,15 @@ impl ArrowStreamSerializer { source: arrow::error::ArrowError::JsonError(e.to_string()), } })?; - build_record_batch(self.schema.clone(), &values) + build_record_batch_with_variant_fields(self.schema.clone(), &values, &self.variant_fields) } /// Create a new ArrowStreamSerializer with the given configuration pub fn new(config: ArrowStreamSerializerConfig) -> Result { let schema = config.schema.ok_or(ArrowEncodingError::MissingSchema)?; + let variant_fields: HashSet<_> = config.variant_fields.into_iter().collect(); + + validate_variant_fields(&schema, &variant_fields)?; // If allow_nullable_fields is enabled, transform the schema once here // instead of on every batch encoding @@ -116,7 +141,9 @@ impl ArrowStreamSerializer { let nullable_fields: Fields = schema .fields() .iter() - .map(|f| make_field_nullable(f)) + .map(|field| { + make_field_nullable(field, &variant_fields, &[field.name().to_owned()]) + }) .collect::, _>>()? .into(); Schema::new_with_metadata(nullable_fields, schema.metadata().clone()) @@ -126,6 +153,7 @@ impl ArrowStreamSerializer { Ok(Self { schema: SchemaRef::new(schema), + variant_fields: Arc::new(variant_fields), }) } } @@ -138,7 +166,11 @@ impl tokio_util::codec::Encoder> for ArrowStreamSerializer { return Err(ArrowEncodingError::NoEvents); } - let bytes = encode_events_to_arrow_ipc_stream(&events, self.schema.clone())?; + let bytes = encode_events_to_arrow_ipc_stream_with_variant_fields( + &events, + self.schema.clone(), + &self.variant_fields, + )?; buffer.extend_from_slice(&bytes); Ok(()) @@ -198,6 +230,24 @@ pub enum ArrowEncodingError { source: arrow::error::ArrowError, }, + /// Failed to encode a value using the Parquet Variant binary format + #[snafu(display("Failed to encode VARIANT field '{field_name}': {source}"))] + VariantEncoding { + /// The field containing the invalid Variant JSON + field_name: String, + /// The underlying Variant encoding error + source: arrow::error::ArrowError, + }, + + /// A configured VARIANT field path is missing or has an incompatible physical type. + #[snafu(display("Invalid VARIANT field '{}': {reason}", field_path.join(".")))] + InvalidVariantField { + /// Path to the field in the Arrow schema. + field_path: Vec, + /// Description of the invalid path or physical type. + reason: String, + }, + /// Invalid Map schema structure #[snafu(display("Invalid Map schema for field '{field_name}': {reason}"))] InvalidMapSchema { @@ -208,10 +258,10 @@ pub enum ArrowEncodingError { }, } -/// Encodes a batch of events into Arrow IPC streaming format -pub fn encode_events_to_arrow_ipc_stream( +fn encode_events_to_arrow_ipc_stream_with_variant_fields( events: &[Event], schema: SchemaRef, + variant_fields: &HashSet>, ) -> Result { if events.is_empty() { return Err(ArrowEncodingError::NoEvents); @@ -223,7 +273,8 @@ pub fn encode_events_to_arrow_ipc_stream( } })?; - let record_batch = build_record_batch(schema, &json_values)?; + let record_batch = + build_record_batch_with_variant_fields(schema, &json_values, variant_fields)?; let mut buffer = BytesMut::new().writer(); let mut writer = @@ -234,14 +285,53 @@ pub fn encode_events_to_arrow_ipc_stream( Ok(buffer.into_inner().freeze()) } +#[cfg(test)] +fn encode_events_to_arrow_ipc_stream( + events: &[Event], + schema: SchemaRef, +) -> Result { + encode_events_to_arrow_ipc_stream_with_variant_fields(events, schema, &HashSet::new()) +} + /// Recursively makes a Field and all its nested fields nullable -fn make_field_nullable(field: &Field) -> Result { +fn make_field_nullable( + field: &Field, + variant_fields: &HashSet>, + path: &[String], +) -> Result { + if variant_fields.contains(path) { + return Ok(field.clone().with_nullable(true)); + } + let new_data_type = match field.data_type() { - DataType::List(inner_field) => DataType::List(make_field_nullable(inner_field)?.into()), + DataType::List(inner_field) => { + let mut inner_path = path.to_vec(); + inner_path.push(inner_field.name().to_owned()); + DataType::List(make_field_nullable(inner_field, variant_fields, &inner_path)?.into()) + } + DataType::LargeList(inner_field) => { + let mut inner_path = path.to_vec(); + inner_path.push(inner_field.name().to_owned()); + DataType::LargeList( + make_field_nullable(inner_field, variant_fields, &inner_path)?.into(), + ) + } + DataType::FixedSizeList(inner_field, length) => { + let mut inner_path = path.to_vec(); + inner_path.push(inner_field.name().to_owned()); + DataType::FixedSizeList( + make_field_nullable(inner_field, variant_fields, &inner_path)?.into(), + *length, + ) + } DataType::Struct(fields) => DataType::Struct( fields .iter() - .map(|f| make_field_nullable(f)) + .map(|child| { + let mut child_path = path.to_vec(); + child_path.push(child.name().to_owned()); + make_field_nullable(child, variant_fields, &child_path) + }) .collect::, _>>()? .into(), ), @@ -264,9 +354,14 @@ fn make_field_nullable(field: &Field) -> Result { ); let key_field = &fields[0]; let value_field = &fields[1]; + let mut value_path = path.to_vec(); + value_path.push(value_field.name().to_owned()); - let new_struct_fields: Fields = - [key_field.clone(), make_field_nullable(value_field)?.into()].into(); + let new_struct_fields: Fields = [ + key_field.clone(), + make_field_nullable(value_field, variant_fields, &value_path)?.into(), + ] + .into(); // Reconstruct the inner "entries" field // The inner field itself must be non-nullable (only the Map wrapper is nullable) @@ -320,14 +415,31 @@ pub(crate) fn vector_log_events_to_json_values( } /// Build an Arrow RecordBatch from a slice of events using the provided schema. +#[cfg(feature = "parquet")] pub(crate) fn build_record_batch( schema: SchemaRef, values: &[serde_json::Value], +) -> Result { + build_record_batch_with_variant_fields(schema, values, &HashSet::new()) +} + +fn build_record_batch_with_variant_fields( + schema: SchemaRef, + values: &[serde_json::Value], + variant_fields: &HashSet>, ) -> Result { if values.is_empty() { return Err(ArrowEncodingError::NoEvents); } + let mut encoded_values = None; + if !variant_fields.is_empty() { + let mut values = values.to_vec(); + encode_variant_fields(&schema, &mut values, variant_fields)?; + encoded_values = Some(values); + } + let values = encoded_values.as_deref().unwrap_or(values); + let missing = find_null_non_nullable_fields(&schema, values); if !missing.is_empty() { let error: vector_common::Error = Box::new(ArrowEncodingError::NullConstraint { @@ -373,6 +485,191 @@ pub(crate) fn build_record_batch( .ok_or(ArrowEncodingError::NoEvents) } +fn is_variant_data_type(data_type: &DataType) -> bool { + let DataType::Struct(fields) = data_type else { + return false; + }; + + fields.len() == 2 + && fields[0].name() == "metadata" + && fields[0].data_type() == &DataType::LargeBinary + && !fields[0].is_nullable() + && fields[1].name() == "value" + && fields[1].data_type() == &DataType::LargeBinary + && !fields[1].is_nullable() +} + +fn validate_variant_fields( + schema: &Schema, + variant_fields: &HashSet>, +) -> Result<(), ArrowEncodingError> { + for path in variant_fields { + if path.is_empty() { + return Err(ArrowEncodingError::InvalidVariantField { + field_path: path.clone(), + reason: "field path must not be empty".to_owned(), + }); + } + + let Some(field) = field_at_path(schema, path) else { + return Err(ArrowEncodingError::InvalidVariantField { + field_path: path.clone(), + reason: "field does not exist in the Arrow schema".to_owned(), + }); + }; + + if !is_variant_data_type(field.data_type()) { + return Err(ArrowEncodingError::InvalidVariantField { + field_path: path.clone(), + reason: format!( + "expected Struct, found {:?}", + field.data_type() + ), + }); + } + } + + Ok(()) +} + +fn field_at_path<'a>(schema: &'a Schema, path: &[String]) -> Option<&'a Field> { + let (first, remaining) = path.split_first()?; + let field = schema.fields().iter().find(|field| field.name() == first)?; + nested_field_at_path(field, remaining) +} + +fn nested_field_at_path<'a>(field: &'a Field, path: &[String]) -> Option<&'a Field> { + let Some((next, remaining)) = path.split_first() else { + return Some(field); + }; + + let child = match field.data_type() { + DataType::Struct(fields) => fields.iter().find(|field| field.name() == next)?, + DataType::List(child) | DataType::LargeList(child) | DataType::FixedSizeList(child, _) => { + (child.name() == next).then_some(child.as_ref())? + } + DataType::Map(entries, _) => { + let DataType::Struct(fields) = entries.data_type() else { + return None; + }; + fields.get(1).filter(|field| field.name() == next)? + } + _ => return None, + }; + + nested_field_at_path(child, remaining) +} + +fn encode_variant_fields( + schema: &Schema, + values: &mut [serde_json::Value], + variant_fields: &HashSet>, +) -> Result<(), ArrowEncodingError> { + for value in values { + let Some(object) = value.as_object_mut() else { + continue; + }; + + for field in schema.fields() { + if let Some(value) = object.get_mut(field.name()) { + encode_variant_value( + field, + value, + &mut vec![field.name().to_owned()], + variant_fields, + )?; + } + } + } + + Ok(()) +} + +fn encode_variant_value( + field: &Field, + value: &mut serde_json::Value, + path: &mut Vec, + variant_fields: &HashSet>, +) -> Result<(), ArrowEncodingError> { + if value.is_null() { + return Ok(()); + } + + if variant_fields.contains(path.as_slice()) { + let owned_json; + let json = if let Some(json) = value.as_str() { + json + } else { + owned_json = value.to_string(); + &owned_json + }; + + let mut builder = VariantBuilder::new(); + builder + .append_json(json) + .map_err(|source| ArrowEncodingError::VariantEncoding { + field_name: path.join("."), + source, + })?; + let (metadata, encoded_value) = builder.finish(); + *value = serde_json::json!({ + "metadata": hex::encode(metadata), + "value": hex::encode(encoded_value), + }); + return Ok(()); + } + + match field.data_type() { + DataType::Struct(fields) => { + if let Some(object) = value.as_object_mut() { + for child in fields { + if let Some(value) = object.get_mut(child.name()) { + path.push(child.name().to_owned()); + encode_variant_value(child, value, path, variant_fields)?; + path.pop(); + } + } + } + } + DataType::List(child) | DataType::LargeList(child) => { + if let Some(values) = value.as_array_mut() { + path.push(child.name().to_owned()); + for value in values { + encode_variant_value(child, value, path, variant_fields)?; + } + path.pop(); + } + } + DataType::FixedSizeList(child, _) => { + if let Some(values) = value.as_array_mut() { + path.push(child.name().to_owned()); + for value in values { + encode_variant_value(child, value, path, variant_fields)?; + } + path.pop(); + } + } + DataType::Map(entries, _) => { + let DataType::Struct(fields) = entries.data_type() else { + return Ok(()); + }; + let Some(value_field) = fields.get(1) else { + return Ok(()); + }; + if let Some(object) = value.as_object_mut() { + path.push(value_field.name().to_owned()); + for value in object.values_mut() { + encode_variant_value(value_field, value, path, variant_fields)?; + } + path.pop(); + } + } + _ => {} + } + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -391,7 +688,27 @@ mod tests { events: Vec, schema: SchemaRef, ) -> Result> { - let bytes = encode_events_to_arrow_ipc_stream(&events, schema.clone())?; + let bytes = encode_events_to_arrow_ipc_stream_with_variant_fields( + &events, + schema, + &HashSet::new(), + )?; + let cursor = Cursor::new(bytes); + let mut reader = StreamReader::try_new(cursor, None)?; + Ok(reader.next().unwrap()?) + } + + fn encode_and_decode_with_variant_fields( + events: Vec, + schema: SchemaRef, + variant_fields: Vec>, + ) -> Result> { + let variant_fields = variant_fields.into_iter().collect(); + let bytes = encode_events_to_arrow_ipc_stream_with_variant_fields( + &events, + schema, + &variant_fields, + )?; let cursor = Cursor::new(bytes); let mut reader = StreamReader::try_new(cursor, None)?; Ok(reader.next().unwrap()?) @@ -409,6 +726,135 @@ mod tests { Event::Log(log) } + fn variant_data_type() -> DataType { + DataType::Struct(Fields::from(vec![ + Field::new("metadata", DataType::LargeBinary, false), + Field::new("value", DataType::LargeBinary, false), + ])) + } + + mod variant { + use super::*; + use parquet_variant::Variant; + use vrl::value::ObjectMap; + + #[test] + fn encodes_json_text_as_variant_binary() { + let schema = + Schema::new(vec![Field::new("attributes", variant_data_type(), false)]).into(); + let event = create_event(vec![("attributes", r#"{"name":"Alice","age":30}"#)]); + + let batch = encode_and_decode_with_variant_fields( + vec![event], + schema, + vec![vec!["attributes".to_owned()]], + ) + .unwrap(); + let array = batch.column(0).as_struct(); + let metadata = array.column(0).as_binary::().value(0); + let value = array.column(1).as_binary::().value(0); + let variant = Variant::try_new(metadata, value).unwrap(); + + assert!(matches!(variant, Variant::Object(_))); + } + + #[test] + fn encodes_native_values_in_nested_variant_fields() { + let variant = variant_data_type(); + let map_entries = Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::LargeUtf8, false), + Field::new("values", variant.clone(), true), + ])), + false, + ); + let payload_type = DataType::Struct(Fields::from(vec![ + Field::new("nested", variant.clone(), true), + Field::new( + "items", + DataType::List(Field::new("item", variant, true).into()), + true, + ), + Field::new("properties", DataType::Map(map_entries.into(), false), true), + ])); + let schema = Schema::new(vec![Field::new("payload", payload_type, true)]).into(); + + let mut nested = ObjectMap::new(); + nested.insert("enabled".into(), Value::Boolean(true)); + let mut properties = ObjectMap::new(); + properties.insert("first".into(), Value::Integer(42)); + let mut payload = ObjectMap::new(); + payload.insert("nested".into(), Value::Object(nested)); + payload.insert( + "items".into(), + Value::Array(vec![Value::Integer(1), Value::Boolean(false)]), + ); + payload.insert("properties".into(), Value::Object(properties)); + let event = create_event(vec![("payload", Value::Object(payload))]); + + let batch = encode_and_decode_with_variant_fields( + vec![event], + schema, + vec![ + vec!["payload".to_owned(), "nested".to_owned()], + vec!["payload".to_owned(), "items".to_owned(), "item".to_owned()], + vec![ + "payload".to_owned(), + "properties".to_owned(), + "values".to_owned(), + ], + ], + ) + .unwrap(); + assert_eq!(batch.num_rows(), 1); + } + + #[test] + fn rejects_invalid_variant_json_text() { + let schema = + Schema::new(vec![Field::new("attributes", variant_data_type(), true)]).into(); + let event = create_event(vec![("attributes", "not valid JSON")]); + + let error = encode_and_decode_with_variant_fields( + vec![event], + schema, + vec![vec!["attributes".to_owned()]], + ) + .unwrap_err(); + assert!(error.to_string().contains("attributes")); + } + + #[test] + fn preserves_unmarked_struct_with_variant_physical_shape() { + let schema = + Schema::new(vec![Field::new("payload", variant_data_type(), false)]).into(); + let mut payload = ObjectMap::new(); + payload.insert("metadata".into(), Value::Bytes("0102".into())); + payload.insert("value".into(), Value::Bytes("0304".into())); + let event = create_event(vec![("payload", Value::Object(payload))]); + + let batch = encode_and_decode(vec![event], schema).unwrap(); + let array = batch.column(0).as_struct(); + + assert_eq!(array.column(0).as_binary::().value(0), &[1, 2]); + assert_eq!(array.column(1).as_binary::().value(0), &[3, 4]); + } + + #[test] + fn rejects_variant_marker_for_non_variant_physical_type() { + let schema = Schema::new(vec![Field::new("payload", DataType::LargeUtf8, false)]); + let config = ArrowStreamSerializerConfig::new(schema) + .with_variant_fields(vec![vec!["payload".to_owned()]]); + + let error = ArrowStreamSerializer::new(config).unwrap_err(); + assert!(matches!( + error, + ArrowEncodingError::InvalidVariantField { .. } + )); + } + } + mod comprehensive { use super::*; @@ -748,6 +1194,10 @@ mod tests { use super::*; use tokio_util::codec::Encoder; + fn make_nullable_without_variants(field: &Field) -> Result { + make_field_nullable(field, &HashSet::new(), &[field.name().to_owned()]) + } + #[test] fn test_config_allow_nullable_fields_overrides_schema() { let mut log1 = LogEvent::default(); @@ -804,7 +1254,7 @@ mod tests { let outer_struct = DataType::Struct(arrow::datatypes::Fields::from(vec![outer_field])); let original_field = Field::new("root", outer_struct, false); - let nullable_field = make_field_nullable(&original_field).unwrap(); + let nullable_field = make_nullable_without_variants(&original_field).unwrap(); assert!( nullable_field.is_nullable(), @@ -832,6 +1282,22 @@ mod tests { } } + #[test] + fn test_make_field_nullable_preserves_variant_children() { + let original_field = Field::new("variant", variant_data_type(), false); + let variant_fields = HashSet::from([vec!["variant".to_owned()]]); + let nullable_field = + make_field_nullable(&original_field, &variant_fields, &["variant".to_owned()]) + .unwrap(); + + assert!(nullable_field.is_nullable()); + let DataType::Struct(fields) = nullable_field.data_type() else { + panic!("Expected Struct type for Variant field"); + }; + assert!(!fields[0].is_nullable()); + assert!(!fields[1].is_nullable()); + } + #[test] fn test_make_field_nullable_with_map_type() { let key_field = Field::new("key", DataType::Utf8, false); @@ -842,7 +1308,7 @@ mod tests { let map_type = DataType::Map(entries_field.into(), false); let original_field = Field::new("my_map", map_type, false); - let nullable_field = make_field_nullable(&original_field).unwrap(); + let nullable_field = make_nullable_without_variants(&original_field).unwrap(); assert!( nullable_field.is_nullable(), diff --git a/src/sinks/databricks_zerobus/service.rs b/src/sinks/databricks_zerobus/service.rs index 29845ed6bfac2..cb7126869ef5d 100644 --- a/src/sinks/databricks_zerobus/service.rs +++ b/src/sinks/databricks_zerobus/service.rs @@ -348,7 +348,7 @@ impl ZerobusService { async fn resolve_arrow_schema( config: &ZerobusSinkConfig, http_client: &HttpClient, - ) -> Result { + ) -> Result { let (client_id, client_secret) = config.auth.credentials(); let table_schema = unity_catalog_schema::fetch_table_schema( @@ -367,11 +367,11 @@ impl ZerobusService { pub(super) async fn ensure_schema(&self) -> Result<&ResolvedSchema, ZerobusSinkError> { self.schema .get_or_try_init(|| async { - let arrow_schema = - Self::resolve_arrow_schema(&self.config, &self.http_client).await?; + let generated = Self::resolve_arrow_schema(&self.config, &self.http_client).await?; let batch_serializer = BatchSerializerConfig::ArrowStream( - ArrowStreamSerializerConfig::new(arrow_schema.clone()), + ArrowStreamSerializerConfig::new(generated.schema.clone()) + .with_variant_fields(generated.variant_fields), ) .build_batch_serializer() .map_err(|e| ZerobusSinkError::ConfigError { @@ -380,7 +380,7 @@ impl ZerobusService { Ok(ResolvedSchema { encoder: BatchEncoder::new(batch_serializer), - arrow_schema: Arc::new(arrow_schema), + arrow_schema: Arc::new(generated.schema), }) }) .await diff --git a/src/sinks/databricks_zerobus/unity_catalog_schema.rs b/src/sinks/databricks_zerobus/unity_catalog_schema.rs index c4637cb58e9c1..7ab5ba3136be3 100644 --- a/src/sinks/databricks_zerobus/unity_catalog_schema.rs +++ b/src/sinks/databricks_zerobus/unity_catalog_schema.rs @@ -35,6 +35,17 @@ fn status_is_retryable(status: StatusCode) -> bool { use databricks_zerobus_ingest_sdk::schema::UcColumn as UnityCatalogColumn; pub use databricks_zerobus_ingest_sdk::schema::UcTableSchema as UnityCatalogTableSchema; +/// Arrow schema plus the logical Unity Catalog VARIANT fields it contains. +/// +/// VARIANT uses the same physical Arrow shape as a possible user-defined struct, so logical type +/// identity must be retained separately from the wire schema. +pub struct GeneratedArrowSchema { + /// Canonical schema sent to the Databricks Arrow Flight server. + pub schema: arrow::datatypes::Schema, + /// Logical VARIANT paths used only by Vector's Arrow encoder. + pub variant_fields: Vec>, +} + /// OAuth token response from Databricks #[derive(Debug, Deserialize)] struct OAuthTokenResponse { @@ -217,11 +228,12 @@ async fn get_oauth_token( /// `RecordBatch` schema in lock-step with the stream's declared schema. pub fn generate_arrow_schema_from_schema( schema: &UnityCatalogTableSchema, -) -> Result { +) -> Result { let arrow_schema = arrow_schema_from_uc_schema(schema).map_err(|e| ZerobusSinkError::ConfigError { message: format!("Failed to convert Unity Catalog schema to Arrow: {}", e), })?; + let variant_fields = collect_variant_fields(schema)?; if tracing::enabled!(tracing::Level::INFO) { info!( @@ -230,7 +242,101 @@ pub fn generate_arrow_schema_from_schema( ); } - Ok(arrow_schema) + Ok(GeneratedArrowSchema { + schema: arrow_schema, + variant_fields, + }) +} + +fn collect_variant_fields( + schema: &UnityCatalogTableSchema, +) -> Result>, ZerobusSinkError> { + let mut columns: Vec<_> = schema + .columns + .iter() + .filter(|column| column.position >= 0) + .collect(); + columns.sort_by_key(|column| column.position); + + let mut variant_fields = Vec::new(); + for column in columns { + let path = vec![column.name.clone()]; + if column.type_name == "VARIANT" { + variant_fields.push(path); + } else if matches!(column.type_name.as_str(), "STRUCT" | "ARRAY" | "MAP") { + let data_type: serde_json::Value = + serde_json::from_str(&column.type_json).map_err(|error| { + ZerobusSinkError::ConfigError { + message: format!( + "Failed to inspect Unity Catalog type for column '{}': {}", + column.name, error + ), + } + })?; + collect_variant_fields_from_type(&data_type, &path, &mut variant_fields); + } + } + + Ok(variant_fields) +} + +fn collect_variant_fields_from_type( + data_type: &serde_json::Value, + path: &[String], + variant_fields: &mut Vec>, +) { + if data_type + .as_str() + .is_some_and(|data_type| data_type == "variant") + { + variant_fields.push(path.to_vec()); + return; + } + + let Some(data_type) = data_type.as_object() else { + return; + }; + + match data_type.get("type").and_then(serde_json::Value::as_str) { + Some("variant") => variant_fields.push(path.to_vec()), + Some("struct") => { + let Some(fields) = data_type + .get("fields") + .and_then(serde_json::Value::as_array) + else { + return; + }; + for field in fields { + let Some(field) = field.as_object() else { + continue; + }; + let (Some(name), Some(data_type)) = ( + field.get("name").and_then(serde_json::Value::as_str), + field.get("type"), + ) else { + continue; + }; + let mut child_path = path.to_vec(); + child_path.push(name.to_owned()); + collect_variant_fields_from_type(data_type, &child_path, variant_fields); + } + } + Some("array") => { + if let Some(element_type) = data_type.get("elementType") { + let mut element_path = path.to_vec(); + element_path.push("item".to_owned()); + collect_variant_fields_from_type(element_type, &element_path, variant_fields); + } + } + Some("map") => { + if let Some(value_type) = data_type.get("valueType") { + let mut value_path = path.to_vec(); + value_path.push("values".to_owned()); + collect_variant_fields_from_type(value_type, &value_path, variant_fields); + } + } + _ => {} + } } #[cfg(test)] @@ -286,8 +392,9 @@ mod tests { ], }; - let arrow_schema = + let generated = generate_arrow_schema_from_schema(&schema).expect("arrow schema should be generated"); + let arrow_schema = generated.schema; assert_eq!(arrow_schema.fields().len(), 2); let id = arrow_schema.field_with_name("id").expect("id field"); @@ -310,8 +417,9 @@ mod tests { let schema: UnityCatalogTableSchema = serde_json::from_str(json).expect("Failed to parse nested_structs_complete schema"); - let arrow_schema = + let generated = generate_arrow_schema_from_schema(&schema).expect("Failed to generate arrow schema"); + let arrow_schema = generated.schema; // Primitive fields map to their canonical Arrow types. assert_eq!( @@ -352,4 +460,135 @@ mod tests { ); } } + + #[test] + fn tracks_variant_fields_without_marking_same_shaped_structs() { + let schema = UnityCatalogTableSchema { + name: "test_table".to_owned(), + catalog_name: "test_catalog".to_owned(), + schema_name: "test_schema".to_owned(), + columns: vec![ + UnityCatalogColumn { + name: "top_level".to_owned(), + type_text: "variant".to_owned(), + type_name: "VARIANT".to_owned(), + position: 0, + nullable: true, + type_json: "\"variant\"".to_owned(), + }, + UnityCatalogColumn { + name: "payload".to_owned(), + type_text: "struct".to_owned(), + type_name: "STRUCT".to_owned(), + position: 1, + nullable: true, + type_json: serde_json::json!({ + "type": "struct", + "fields": [ + { + "name": "metadata", + "type": "binary", + "nullable": false, + "metadata": {} + }, + { + "name": "value", + "type": "binary", + "nullable": false, + "metadata": {} + }, + { + "name": "nested", + "type": "variant", + "nullable": true, + "metadata": {} + } + ] + }) + .to_string(), + }, + UnityCatalogColumn { + name: "items".to_owned(), + type_text: "array".to_owned(), + type_name: "ARRAY".to_owned(), + position: 2, + nullable: true, + type_json: serde_json::json!({ + "type": "array", + "elementType": "variant", + "containsNull": true + }) + .to_string(), + }, + UnityCatalogColumn { + name: "properties".to_owned(), + type_text: "map".to_owned(), + type_name: "MAP".to_owned(), + position: 3, + nullable: true, + type_json: serde_json::json!({ + "type": "map", + "keyType": "string", + "valueType": "variant", + "valueContainsNull": true + }) + .to_string(), + }, + UnityCatalogColumn { + name: "binary_pair".to_owned(), + type_text: "struct".to_owned(), + type_name: "STRUCT".to_owned(), + position: 4, + nullable: true, + type_json: serde_json::json!({ + "type": "struct", + "fields": [ + { + "name": "metadata", + "type": "binary", + "nullable": false, + "metadata": {} + }, + { + "name": "value", + "type": "binary", + "nullable": false, + "metadata": {} + } + ] + }) + .to_string(), + }, + ], + }; + + let generated = + generate_arrow_schema_from_schema(&schema).expect("arrow schema should be generated"); + + assert_eq!( + generated.variant_fields, + vec![ + vec!["top_level".to_owned()], + vec!["payload".to_owned(), "nested".to_owned()], + vec!["items".to_owned(), "item".to_owned()], + vec!["properties".to_owned(), "values".to_owned()], + ] + ); + assert!( + !generated + .variant_fields + .contains(&vec!["payload".to_owned()]) + ); + assert!( + !generated + .variant_fields + .contains(&vec!["binary_pair".to_owned()]) + ); + + let config = + vector_lib::codecs::encoding::ArrowStreamSerializerConfig::new(generated.schema) + .with_variant_fields(generated.variant_fields); + vector_lib::codecs::encoding::ArrowStreamSerializer::new(config) + .expect("generated VARIANT paths should match the Arrow schema"); + } }