diff --git a/Cargo.lock b/Cargo.lock index 7495e00..7c2bc32 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -884,7 +884,6 @@ dependencies = [ "ouroboros", "prost", "prost-build", - "prost-reflect", "protobuf", "strum 0.26.3", ] @@ -982,9 +981,9 @@ dependencies = [ [[package]] name = "fixedbitset" -version = "0.4.2" +version = "0.5.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" [[package]] name = "flate2" @@ -1090,6 +1089,12 @@ dependencies = [ "rayon", ] +[[package]] +name = "hashbrown" +version = "0.15.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" + [[package]] name = "hashlink" version = "0.8.4" @@ -1161,12 +1166,12 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.2.6" +version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "168fb715dda47215e360912c096649d23d58bf392ac62f73919e831745e40f26" +checksum = "fe4cd85333e22411419a0bcae1297d25e58c9443848b11dc6a86fefe8c78a661" dependencies = [ "equivalent", - "hashbrown 0.14.5", + "hashbrown 0.15.4", ] [[package]] @@ -1405,9 +1410,9 @@ dependencies = [ [[package]] name = "multimap" -version = "0.10.0" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "defc4c55412d89136f966bbb339008b474350e5e6e78d2714439c386b3137a03" +checksum = "1d87ecb2933e8aeadb3e3a02b828fed80a7528047e68b4f424523a0981a3a084" [[package]] name = "multiversion" @@ -1595,9 +1600,9 @@ checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" [[package]] name = "petgraph" -version = "0.6.5" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db" +checksum = "3672b37090dbd86368a4145bc067582552b29c27377cad4e0a306c97f9bd7772" dependencies = [ "fixedbitset", "indexmap", @@ -1981,9 +1986,9 @@ dependencies = [ [[package]] name = "prost" -version = "0.13.2" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b2ecbe40f08db5c006b5764a2645f7f3f141ce756412ac9e1dd6087e6d32995" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" dependencies = [ "bytes", "prost-derive", @@ -1991,11 +1996,10 @@ dependencies = [ [[package]] name = "prost-build" -version = "0.13.1" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5bb182580f71dd070f88d01ce3de9f4da5021db7115d2e1c3605a754153b77c1" +checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ - "bytes", "heck 0.5.0", "itertools", "log", @@ -2012,9 +2016,9 @@ dependencies = [ [[package]] name = "prost-derive" -version = "0.13.2" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "acf0c195eebb4af52c752bec4f52f645da98b6e92077a04110c7f349477ae5ac" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", "itertools", @@ -2023,22 +2027,11 @@ dependencies = [ "syn 2.0.77", ] -[[package]] -name = "prost-reflect" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55a6a9143ae25c25fa7b6a48d6cc08b10785372060009c25140a4e7c340e95af" -dependencies = [ - "once_cell", - "prost", - "prost-types", -] - [[package]] name = "prost-types" -version = "0.13.1" +version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cee5168b05f49d4b0ca581206eb14a7b22fafd963efe729ac48eb03266e25cc2" +checksum = "52c2c1bf36ddb1a1c396b3601a3cec27c2462e45f07c386894ec3ccf5332bd16" dependencies = [ "prost", ] @@ -2046,9 +2039,8 @@ dependencies = [ [[package]] name = "protobuf" version = "3.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df67496db1a89596beaced1579212e9b7c53c22dca1d9745de00ead76573d514" dependencies = [ + "bytes", "once_cell", "protobuf-support", "thiserror", diff --git a/Cargo.toml b/Cargo.toml index 207c44d..465ee49 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ members = [ "packages/vendor/duckdb", "packages/vendor/duckdb-loadable-macros", "packages/vendor/libduckdb-sys", + "packages/vendor/protobuf", "packages/duckdb_metadata", "packages/duckdb_metadata_bin", "packages/duckdb_protobuf" @@ -14,3 +15,4 @@ members = [ duckdb = { path = 'packages/vendor/duckdb' } duckdb-loadable-macros = { path = 'packages/vendor/duckdb-loadable-macros' } libduckdb-sys = { path = 'packages/vendor/libduckdb-sys' } +protobuf = { path = 'packages/vendor/protobuf' } diff --git a/flake.nix b/flake.nix index 8d50eea..c479404 100644 --- a/flake.nix +++ b/flake.nix @@ -71,15 +71,28 @@ patches = [ patches/libduckdb-sys+1.0.0.patch ]; }; + protobufCrate = applyPatch { + src = pkgs.fetchCrate { + pname = "protobuf"; + version = "3.5.0"; + sha256 = "sha256-xU9CAl/LQT9Gk0OdnbZuViQUeeD1tbvJcmxc9ohuZM8="; + }; + patches = [ patches/protobuf+3.5.0-range.patch patches/protobuf+3.5.0-reflection-range.patch ]; + }; + vendorScript = pkgs.writeShellScriptBin "vendor-deps" '' set -euo pipefail mkdir -p packages/vendor/duckdb mkdir -p packages/vendor/duckdb-loadable-macros mkdir -p packages/vendor/libduckdb-sys + mkdir -p packages/vendor/protobuf cp -r ${duckdbCrate}/* packages/vendor/duckdb/ cp -r ${duckdbLoadableMacrosCrate}/* packages/vendor/duckdb-loadable-macros/ cp -r ${libduckdbSysCrate}/* packages/vendor/libduckdb-sys/ + cp -r ${protobufCrate}/* packages/vendor/protobuf/ + + chmod -R +w packages/vendor ''; vendoredSrc = pkgs.stdenvNoCC.mkDerivation { diff --git a/packages/duckdb_protobuf/Cargo.toml b/packages/duckdb_protobuf/Cargo.toml index a650600..6c67ced 100644 --- a/packages/duckdb_protobuf/Cargo.toml +++ b/packages/duckdb_protobuf/Cargo.toml @@ -8,10 +8,11 @@ crate-type = ["cdylib"] [dependencies] anyhow = "1.0.86" -duckdb = { version = "1.0.0", features = ["vtab-loadable", "loadable_extension"] } +duckdb = { version = "1.0.0", features = [ + "vtab-loadable", + "loadable_extension", +] } duckdb-loadable-macros = "0.1.2" -prost = "0.13.0" -prost-reflect = "0.14.0" protobuf = "3.5.0" glob = "0.3.1" byteorder = "1.5.0" @@ -28,4 +29,3 @@ prost-build = "0.13.1" [[test]] name = "it" path = "tests/it/main.rs" - diff --git a/packages/duckdb_protobuf/src/filtered_dynamic_message.rs b/packages/duckdb_protobuf/src/filtered_dynamic_message.rs deleted file mode 100644 index 106d1f6..0000000 --- a/packages/duckdb_protobuf/src/filtered_dynamic_message.rs +++ /dev/null @@ -1,59 +0,0 @@ -use prost::bytes::{Buf, BufMut}; -use prost::encoding::{DecodeContext, WireType}; -use prost::{DecodeError, Message}; -use prost_reflect::{DynamicMessage, UnknownField}; -use std::collections::HashSet; - -#[derive(Debug, Clone)] -pub struct FilteredDynamicMessage { - message: DynamicMessage, - accepted_fields: HashSet, -} - -impl FilteredDynamicMessage { - pub fn new(message: DynamicMessage, accepted_fields: HashSet) -> FilteredDynamicMessage { - FilteredDynamicMessage { - message, - accepted_fields, - } - } - - pub fn into(self) -> DynamicMessage { - self.message - } -} - -impl Message for FilteredDynamicMessage { - fn encode_raw(&self, buf: &mut impl BufMut) - where - Self: Sized, - { - self.message.encode_raw(buf) - } - - fn merge_field( - &mut self, - number: u32, - wire_type: WireType, - buf: &mut impl Buf, - ctx: DecodeContext, - ) -> Result<(), DecodeError> - where - Self: Sized, - { - if !self.accepted_fields.contains(&number) { - let _field = UnknownField::decode_value(number, wire_type, buf, ctx)?; - return Ok(()); - } - - self.message.merge_field(number, wire_type, buf, ctx) - } - - fn encoded_len(&self) -> usize { - self.message.encoded_len() - } - - fn clear(&mut self) { - self.message.clear() - } -} diff --git a/packages/duckdb_protobuf/src/lib.rs b/packages/duckdb_protobuf/src/lib.rs index 1dd6105..6923c71 100644 --- a/packages/duckdb_protobuf/src/lib.rs +++ b/packages/duckdb_protobuf/src/lib.rs @@ -1,4 +1,3 @@ -mod filtered_dynamic_message; mod io; mod read; mod types; diff --git a/packages/duckdb_protobuf/src/read.rs b/packages/duckdb_protobuf/src/read.rs index 2e73499..f0bc3d0 100644 --- a/packages/duckdb_protobuf/src/read.rs +++ b/packages/duckdb_protobuf/src/read.rs @@ -1,22 +1,29 @@ use std::collections::HashMap; -use std::ffi::CString; use std::marker::PhantomData; +use std::ops::Deref; use std::slice; use anyhow::{bail, format_err}; -use duckdb::vtab::{DataChunk, LogicalType, LogicalTypeId}; -use prost_reflect::{Cardinality, DynamicMessage, FieldDescriptor, Kind, ReflectMessage, Value}; +use duckdb::vtab::{DataChunk, LogicalTypeId}; +use protobuf::reflect::{ + FieldDescriptor, MessageDescriptor, MessageRef, ReflectFieldRef, ReflectOptionalRef, + ReflectValueRef, RuntimeFieldType, RuntimeType, +}; +use protobuf::MessageDyn; pub fn write_to_output( mappings: &[u64], columns_state: &mut HashMap, - value: &DynamicMessage, + value: &dyn MessageDyn, output: &DataChunk, max_rows: usize, row_idx: usize, + include_position: bool, + starting_offset: u64, ) -> Result<(), anyhow::Error> { let column_key = &ColumnKey::empty(); - let fields = value.descriptor().fields().collect::>(); + let message_descriptor = value.descriptor_dyn(); + let fields: Vec = message_descriptor.fields().collect(); for (output_field_idx, field_idx) in mappings.iter().enumerate() { let field_idx = *field_idx as usize; if field_idx >= fields.len() { @@ -25,18 +32,20 @@ pub fn write_to_output( let field_descriptor = &fields[field_idx]; let column_vector = output.get_vector(output_field_idx); - let value = value.get_field(&field_descriptor); + let field_ref = field_descriptor.get_reflect(value); let column_key = column_key.field(&field_descriptor); write_column( columns_state, &column_key, - &value, + Some(field_ref), &field_descriptor, column_vector, max_rows, row_idx, + include_position, + starting_offset, )?; } @@ -46,28 +55,58 @@ pub fn write_to_output( pub fn write_message( columns_state: &mut HashMap, column_key: &ColumnKey, - value: &DynamicMessage, + message_descriptor: &MessageDescriptor, + value: Option, output: &impl VectorAccessor, max_rows: usize, row_idx: usize, + include_position: bool, + starting_offset: u64, ) -> Result<(), anyhow::Error> { - for (field_idx, field_descriptor) in value.descriptor().fields().enumerate() { + let mut last_field_idx = 0; + for (field_idx, field_descriptor) in message_descriptor.fields().enumerate() { + last_field_idx = field_idx; + let column_vector = output.get_vector(field_idx); - let value = value.get_field(&field_descriptor); + let field_ref = value + .as_ref() + .map(|it| field_descriptor.get_reflect(it.deref())); let column_key = column_key.field(&field_descriptor); write_column( columns_state, &column_key, - &value, + field_ref, &field_descriptor, column_vector, max_rows, row_idx, + include_position, + starting_offset, )?; } + // if include_position { + // if let Some(value) = value { + // let (start_pos, end_pos) = value.deref().special_fields_dyn().range(); + + // let mut position_vector = { + // let column = output.get_vector(last_field_idx + 1); + // unsafe { MyFlatVector::::with_capacity(column, max_rows) } + // }; + + // let mut length_vector = { + // let column = output.get_vector(last_field_idx + 2); + // unsafe { MyFlatVector::::with_capacity(column, max_rows) } + // }; + + // position_vector.as_mut_slice()[row_idx] = start_pos + starting_offset; + // length_vector.as_mut_slice()[row_idx] = end_pos - start_pos; + // } + // } + + Ok(()) } @@ -98,14 +137,16 @@ impl MyFlatVector { pub fn write_column( columns_state: &mut HashMap, column_key: &ColumnKey, - value: &Value, + field_ref: Option, field_descriptor: &FieldDescriptor, column: duckdb::ffi::duckdb_vector, max_rows: usize, row_idx: usize, + include_position: bool, + starting_offset: u64, ) -> Result<(), anyhow::Error> { - match field_descriptor.cardinality() { - Cardinality::Repeated => { + match field_descriptor.runtime_field_type() { + RuntimeFieldType::Repeated(element_type) => { let column_key = column_key.extending(ColumnKeyElement::List); let mut list_entries_vector = unsafe { @@ -113,10 +154,6 @@ pub fn write_column( }; let list_entry = &mut list_entries_vector.as_mut_slice()[row_idx]; - let values = value - .as_list() - .ok_or_else(|| format_err!("expected list"))?; - let next_offset_ref = columns_state.get_mut(&column_key); let next_offset = if let Some(it) = &next_offset_ref { **it @@ -124,7 +161,11 @@ pub fn write_column( 0 }; - let len_u64 = u64::try_from(values.len())?; + let len_u64 = if let Some(ReflectFieldRef::Repeated(values)) = &field_ref { + u64::try_from(values.len())? + } else { + 0 + }; list_entry.offset = next_offset; list_entry.length = len_u64; @@ -142,33 +183,46 @@ pub fn write_column( unsafe { duckdb::ffi::duckdb_list_vector_reserve(column, new_length) }; unsafe { duckdb::ffi::duckdb_list_vector_set_size(column, new_length) }; - let child_vector = unsafe { duckdb::ffi::duckdb_list_vector_get_child(column) }; - - for (idx, value) in values.iter().enumerate() { - let row_idx = next_offset as usize + idx; - - write_single_column( - columns_state, - &column_key, - value, - field_descriptor, - child_vector, - new_length as usize, - row_idx, - )?; + if let Some(ReflectFieldRef::Repeated(values)) = field_ref { + let child_vector = unsafe { duckdb::ffi::duckdb_list_vector_get_child(column) }; + + let element_type = values.element_type(); + + for (idx, value) in values.into_iter().enumerate() { + let row_idx = next_offset as usize + idx; + + write_single_column( + columns_state, + &column_key, + ReflectOptionalRef::some(value), + &element_type, + child_vector, + new_length as usize, + row_idx, + include_position, + starting_offset + )?; + } } } - Cardinality::Optional | Cardinality::Required => { + RuntimeFieldType::Singular(runtime_type) => { write_single_column( columns_state, column_key, - value, - field_descriptor, + if let Some(ReflectFieldRef::Optional(values)) = field_ref { + values + } else { + ReflectOptionalRef::none(runtime_type.clone()) + }, + &runtime_type, column, max_rows, row_idx, + include_position, + starting_offset )?; } + _ => return Err(format_err!("unknown type")), } Ok(()) @@ -177,73 +231,73 @@ pub fn write_column( pub fn write_single_column( columns_state: &mut HashMap, column_key: &ColumnKey, - value: &Value, - field_descriptor: &FieldDescriptor, + value: ReflectOptionalRef, + runtime_type: &RuntimeType, column: duckdb::ffi::duckdb_vector, max_rows: usize, row_idx: usize, + include_position: bool, + starting_offset: u64, ) -> Result<(), anyhow::Error> { - match field_descriptor.kind() { - Kind::Message(message_descriptor) + match runtime_type { + RuntimeType::Message(message_descriptor) if message_descriptor.full_name() == "google.protobuf.Timestamp" => { - let message = value - .as_message() - .ok_or_else(|| format_err!("expected message"))?; - let seconds = - message - .get_field(&message_descriptor.get_field(1).ok_or_else(|| { - format_err!("expected field 1 for google.protobuf.Timestamp") - })?) - .as_i64() + let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; + vector.as_mut_slice()[row_idx] = if let Some(ReflectValueRef::Message(message_value)) = + value.value() + { + let seconds = message_descriptor + .field_by_number(1) + .ok_or_else(|| format_err!("expected field 1 for google.protobuf.Timestamp"))? + .get_singular_field_or_default(message_value.deref()) + .to_i64() .ok_or_else(|| format_err!("expected i64"))?; - let nanos = - message - .get_field(&message_descriptor.get_field(2).ok_or_else(|| { - format_err!("expected field 2 for google.protobuf.Timestamp") - })?) - .as_i32() - .ok_or_else(|| format_err!("expected i32"))?; + let nanos = message_descriptor + .field_by_number(2) + .ok_or_else(|| format_err!("expected field 2 for google.protobuf.Timestamp"))? + .get_singular_field_or_default(message_value.deref()) + .to_i32() + .ok_or_else(|| format_err!("expected i64"))?; - let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = seconds * 1000000 + (nanos as i64 / 1000); + seconds * 1000000 + (nanos as i64 / 1000) + } else { + 0 + } } - Kind::Message(..) => { - let message = value - .as_message() - .ok_or_else(|| format_err!("expected message"))?; - + RuntimeType::Message(message_descriptor) => { let source = unsafe { StructVector::new(column) }; write_message( columns_state, column_key, - message, + &message_descriptor, + if let Some(ReflectValueRef::Message(message)) = value.value() { + Some(message) + } else { + None + }, &source, max_rows, row_idx, + include_position, + starting_offset, )?; } - Kind::Enum(enum_descriptor) => { - let enum_value = value - .as_enum_number() - .ok_or_else(|| format_err!("expected enum value"))?; + RuntimeType::Enum(enum_descriptor) => { + let enum_value_descriptor = + if let Some(ReflectValueRef::Enum(.., value)) = value.value() { + enum_descriptor.value_by_number_or_default(value) + } else { + enum_descriptor.default_value() + }; + let idx = enum_value_descriptor.value(); - let enum_value_descriptor = enum_descriptor - .get_value(enum_value) - .unwrap_or_else(|| enum_descriptor.default_value()); + let column_type = unsafe { duckdb::ffi::duckdb_vector_get_column_type(column) }; - let (idx, _) = enum_descriptor - .values() - .enumerate() - .find(|(_, it)| it.number() == enum_value_descriptor.number()) - .unwrap(); - - let column_type = - unsafe { duckdb::ffi::duckdb_vector_get_column_type(column) }; - - let logical_type = LogicalTypeId::from(unsafe { duckdb::ffi::duckdb_enum_internal_type(column_type) }); + let logical_type = + LogicalTypeId::from(unsafe { duckdb::ffi::duckdb_enum_internal_type(column_type) }); match logical_type { LogicalTypeId::UTinyint => { @@ -263,73 +317,103 @@ pub fn write_single_column( _ => bail!("unknown enum column type {:?}", logical_type), } } - Kind::String => { - let value = value - .as_str() - .ok_or_else(|| format_err!("expected string"))?; - let value = CString::new(value)?; - - unsafe { - duckdb::ffi::duckdb_vector_assign_string_element( - column, - row_idx as u64, - value.as_ptr(), - ) - }; - } - Kind::Double => { - let value = value - .as_f64() - .ok_or_else(|| format_err!("expected double"))?; - let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = value; + RuntimeType::U32 => { + let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; + vector.as_mut_slice()[row_idx] = + if let Some(ReflectValueRef::U32(value)) = value.value() { + value + } else { + u32::default() + }; } - Kind::Float => { - let value = value - .as_f32() - .ok_or_else(|| format_err!("expected float"))?; - let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = value; + RuntimeType::U64 => { + let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; + vector.as_mut_slice()[row_idx] = + if let Some(ReflectValueRef::U64(value)) = value.value() { + value + } else { + u64::default() + }; } - Kind::Int32 => { - let value = value - .as_i32() - .ok_or_else(|| format_err!("expected int32"))?; + RuntimeType::I32 => { let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = value; + vector.as_mut_slice()[row_idx] = + if let Some(ReflectValueRef::I32(value)) = value.value() { + value + } else { + i32::default() + }; } - Kind::Int64 => { - let value = value - .as_i64() - .ok_or_else(|| format_err!("expected int64"))?; + RuntimeType::I64 => { let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = value; + vector.as_mut_slice()[row_idx] = + if let Some(ReflectValueRef::I64(value)) = value.value() { + value + } else { + i64::default() + }; } - Kind::Uint32 => { - let value = value - .as_u32() - .ok_or_else(|| format_err!("expected uint32"))?; - let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = value; + RuntimeType::F32 => { + let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; + vector.as_mut_slice()[row_idx] = + if let Some(ReflectValueRef::F32(value)) = value.value() { + value + } else { + f32::default() + }; } - Kind::Uint64 => { - let value = value - .as_u64() - .ok_or_else(|| format_err!("expected uint64"))?; - let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = value; + RuntimeType::F64 => { + let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; + vector.as_mut_slice()[row_idx] = + if let Some(ReflectValueRef::F64(value)) = value.value() { + value + } else { + f64::default() + }; } - Kind::Bool => { - let value = value - .as_bool() - .ok_or_else(|| format_err!("expected bool"))?; + RuntimeType::Bool => { let mut vector = unsafe { MyFlatVector::::with_capacity(column, max_rows) }; - vector.as_mut_slice()[row_idx] = value; + vector.as_mut_slice()[row_idx] = + if let Some(ReflectValueRef::Bool(value)) = value.value() { + value + } else { + bool::default() + }; + } + RuntimeType::String => { + let value = if let Some(ReflectValueRef::String(value)) = value.value() { + value + } else { + "" + }; + let value = value.as_bytes(); + + unsafe { + duckdb::ffi::duckdb_vector_assign_string_element_len( + column, + row_idx as u64, + value.as_ptr() as _, + value.len() as _, + ) + }; } - _ => { - bail!("unhandled field type"); + RuntimeType::VecU8 => { + let value = if let Some(ReflectValueRef::Bytes(value)) = value.value() { + value + } else { + &[] + }; + + unsafe { + duckdb::ffi::duckdb_vector_assign_string_element_len( + column, + row_idx as u64, + value.as_ptr() as _, + value.len() as _, + ) + }; } - }; + } Ok(()) } @@ -348,7 +432,7 @@ pub struct ColumnKey { impl ColumnKey { pub fn field(&self, field: &FieldDescriptor) -> ColumnKey { self.extending(ColumnKeyElement::Field { - field_tag: field.number(), + field_tag: field.proto().number() as u32, }) } diff --git a/packages/duckdb_protobuf/src/types.rs b/packages/duckdb_protobuf/src/types.rs index bbb6e1a..8b010de 100644 --- a/packages/duckdb_protobuf/src/types.rs +++ b/packages/duckdb_protobuf/src/types.rs @@ -1,54 +1,60 @@ use anyhow::format_err; use duckdb::vtab::{LogicalType, LogicalTypeId}; -use prost_reflect::{Cardinality, FieldDescriptor, Kind}; +use protobuf::reflect::{FieldDescriptor, RuntimeFieldType, RuntimeType}; -pub fn into_logical_type(field: &FieldDescriptor) -> Result { - Ok(match field.cardinality() { - Cardinality::Optional | Cardinality::Required => into_logical_type_single(field)?, - Cardinality::Repeated => LogicalType::list(&into_logical_type_single(field)?), +pub fn into_logical_type( + field: &FieldDescriptor, + with_location: bool, +) -> Result { + Ok(match field.runtime_field_type() { + RuntimeFieldType::Singular(field) => into_logical_type_single(field, with_location)?, + RuntimeFieldType::Repeated(field) => { + LogicalType::list(&into_logical_type_single(field, with_location)?) + } + RuntimeFieldType::Map(_, _) => return Err(format_err!("map unimplemented")), }) } -fn into_logical_type_single(field: &FieldDescriptor) -> Result { - let value = match field.kind() { - Kind::Message(message_descriptor) - if message_descriptor.full_name() == "google.protobuf.Timestamp" => - { - LogicalType::new(LogicalTypeId::Timestamp) - } - Kind::Message(message_descriptor) => { - let fields = message_descriptor - .fields() - .collect::>(); +fn into_logical_type_single( + field: RuntimeType, + with_location: bool, +) -> Result { + let value = match field { + RuntimeType::Message(message_descriptor) => { + if message_descriptor.full_name() == "google.protobuf.Timestamp" { + LogicalType::new(LogicalTypeId::Timestamp) + } else { + let fields = message_descriptor + .fields() + .collect::>(); - let fields = fields - .iter() - .map(|field| Ok((field.name(), into_logical_type(&field)?))) - .collect::, anyhow::Error>>()?; + let mut fields = fields + .iter() + .map(|field| Ok((field.name(), into_logical_type(&field, with_location)?))) + .collect::, anyhow::Error>>()?; - LogicalType::struct_type(fields.as_slice()) + if with_location { + fields.push(("position", LogicalType::new(LogicalTypeId::UBigint))); + fields.push(("length", LogicalType::new(LogicalTypeId::UBigint))) + } + + LogicalType::struct_type(fields.as_slice()) + } } - Kind::Enum(descriptor) => { - let names = descriptor.values().collect::>(); + RuntimeType::Enum(enum_descriptor) => { + let names = enum_descriptor.values().collect::>(); let names = names.iter().map(|it| it.name()).collect::>(); LogicalType::enumeration(names.as_slice()) } - Kind::Double => LogicalType::new(LogicalTypeId::Double), - Kind::Float => LogicalType::new(LogicalTypeId::Float), - Kind::Int32 => LogicalType::new(LogicalTypeId::Integer), - Kind::Int64 => LogicalType::new(LogicalTypeId::Bigint), - Kind::Uint32 => LogicalType::new(LogicalTypeId::UInteger), - Kind::Uint64 => LogicalType::new(LogicalTypeId::UBigint), - Kind::Bool => LogicalType::new(LogicalTypeId::Boolean), - Kind::String => LogicalType::new(LogicalTypeId::Varchar), - logical_type => { - return Err(format_err!( - "unhandled field: {}, type: {:?}", - field.name(), - logical_type, - ) - .into()) - } + RuntimeType::F64 => LogicalType::new(LogicalTypeId::Double), + RuntimeType::F32 => LogicalType::new(LogicalTypeId::Float), + RuntimeType::I32 => LogicalType::new(LogicalTypeId::Integer), + RuntimeType::I64 => LogicalType::new(LogicalTypeId::Bigint), + RuntimeType::U32 => LogicalType::new(LogicalTypeId::UInteger), + RuntimeType::U64 => LogicalType::new(LogicalTypeId::UBigint), + RuntimeType::Bool => LogicalType::new(LogicalTypeId::Boolean), + RuntimeType::String => LogicalType::new(LogicalTypeId::Varchar), + RuntimeType::VecU8 => LogicalType::new(LogicalTypeId::Blob), }; Ok(value) diff --git a/packages/duckdb_protobuf/src/vtab.rs b/packages/duckdb_protobuf/src/vtab.rs index 2c45e67..6daf165 100644 --- a/packages/duckdb_protobuf/src/vtab.rs +++ b/packages/duckdb_protobuf/src/vtab.rs @@ -1,4 +1,3 @@ -use crate::filtered_dynamic_message::FilteredDynamicMessage; use crate::io::{parse, DelimitedLengthKind, LengthDelimitedRecordsReader, LengthKind, Record}; use crate::read::{write_to_output, MyFlatVector, VectorAccessor}; use crate::types::into_logical_type; @@ -8,8 +7,9 @@ use duckdb::vtab::{ BindInfo, DataChunk, Free, FunctionInfo, InitInfo, LogicalType, LogicalTypeId, VTab, VTabLocalData, }; -use prost::Message; -use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor, ReflectMessage}; +use protobuf::descriptor::FileDescriptorSet; +use protobuf::reflect::{FileDescriptor, MessageDescriptor}; +use protobuf::CodedInputStream; use std::error::Error; use std::ffi::CString; use std::fs::File; @@ -22,11 +22,11 @@ pub struct Parameters { pub files: String, pub descriptor_bytes: Vec, pub message_name: String, - pub shared_message_descriptor: MessageDescriptor, + pub file_descriptors: Vec, + pub message_descriptor: MessageDescriptor, pub length_kind: LengthKind, pub include_filename: bool, pub include_position: bool, - pub include_size: bool, } impl Parameters { @@ -50,16 +50,34 @@ impl Parameters { })() .with_context(|| format_err!("field `descriptors`"))?; - let shared_descriptor_pool = DescriptorPool::decode(descriptor_bytes.as_slice())?; + let descriptor_set: FileDescriptorSet = + protobuf::Message::parse_from_bytes(&descriptor_bytes)?; + let file_descriptors = { + let mut file_descriptors = Vec::with_capacity(descriptor_set.file.len()); + for fd_proto in descriptor_set.file { + let fd = FileDescriptor::new_dynamic(fd_proto, &file_descriptors)?; + file_descriptors.push(fd); + } + + file_descriptors + }; let message_name = bind .get_named_parameter("message_type") .ok_or_else(|| format_err!("missing parameter `message_type`"))? .to_string(); - let message_descriptor = shared_descriptor_pool - .get_message_by_name(&message_name.as_str()) - .ok_or_else(|| format_err!("message type not found in `descriptor`"))?; + let message_descriptor = (|| { + for file_descriptor in &file_descriptors { + if let Some(message_descriptor) = + file_descriptor.message_by_full_name(&format!(".{}", message_name)) + { + return Ok(message_descriptor); + } + } + + Err(anyhow::anyhow!("message type '{}' not found", message_name)) + })()?; let length_kind = bind .get_named_parameter("delimiter") @@ -78,31 +96,20 @@ impl Parameters { .map(|value| value.to_int64() != 0) .unwrap_or(false); - let include_size = bind - .get_named_parameter("size") - .map(|value| value.to_int64() != 0) - .unwrap_or(false); - Ok(Self { files, descriptor_bytes, message_name, - shared_message_descriptor: message_descriptor, + file_descriptors, + message_descriptor, length_kind, include_filename, include_position, - include_size, }) } pub fn message_descriptor(&self) -> Result { - let descriptor_pool = DescriptorPool::decode(self.descriptor_bytes.as_slice())?; - - let message_descriptor = descriptor_pool - .get_message_by_name(&self.message_name) - .unwrap(); - - Ok(message_descriptor) + Ok(self.message_descriptor.clone()) } pub fn values() -> Vec<(String, LogicalType)> { @@ -131,7 +138,6 @@ impl Parameters { "position".to_string(), LogicalType::new(LogicalTypeId::Boolean), ), - ("size".to_string(), LogicalType::new(LogicalTypeId::Boolean)), ] } } @@ -220,10 +226,10 @@ impl ProtobufVTab { let params = Parameters::from_bind_info(bind)?; - for field_descriptor in params.shared_message_descriptor.fields() { + for field_descriptor in params.message_descriptor.fields() { bind.add_result_column( field_descriptor.name().as_ref(), - into_logical_type(&field_descriptor)?, + into_logical_type(&field_descriptor, params.include_position)?, ); } @@ -233,9 +239,6 @@ impl ProtobufVTab { if params.include_position { bind.add_result_column("position", LogicalType::new(LogicalTypeId::UBigint)); - } - - if params.include_size { bind.add_result_column("size", LogicalType::new(LogicalTypeId::UBigint)); } @@ -282,28 +285,7 @@ impl ProtobufVTab { let mut column_information = Default::default(); - let message = { - let message = DynamicMessage::new(local_descriptor.clone()); - let fields: Vec<_> = local_descriptor.fields().collect(); - - let message = FilteredDynamicMessage::new( - message, - init_data - .column_indices - .iter() - .filter_map(|it| { - let it = *it as usize; - if it >= fields.len() { - return None; - } - - Some(fields[it].number()) - }) - .collect(), - ); - - message - }; + let mut message_template = local_descriptor.new_instance(); for output_row_idx in 0..available_chunk_size { let StateContainerValue { @@ -316,21 +298,22 @@ impl ProtobufVTab { Some(message_info) => message_info, }; - let mut message = message.clone(); - message.merge(bytes.as_slice())?; - let message = message.into(); + let mut cis = CodedInputStream::from_bytes(&bytes); + message_template.merge_from_dyn(&mut cis)?; + + let mut field_offset = message_template.descriptor_dyn().fields().count(); write_to_output( &init_data.column_indices, &mut column_information, - &message, + message_template.as_ref(), output, available_chunk_size, output_row_idx, + parameters.include_position, + position, )?; - let mut field_offset = message.descriptor().fields().len(); - if parameters.include_filename { if let Some((field_offset, _)) = init_data .column_indices @@ -380,9 +363,7 @@ impl ProtobufVTab { } field_offset += 1; - } - if parameters.include_size { if let Some((field_offset, _)) = init_data .column_indices .iter() diff --git a/patches/protobuf+3.5.0-range.patch b/patches/protobuf+3.5.0-range.patch new file mode 100644 index 0000000..b3f84db --- /dev/null +++ b/patches/protobuf+3.5.0-range.patch @@ -0,0 +1,35 @@ +diff -ruN a/src/special.rs b/src/special.rs +--- a/src/special.rs 2006-07-23 20:21:28 ++++ b/src/special.rs 2024-11-30 18:16:16 +@@ -6,6 +6,7 @@ + pub struct SpecialFields { + unknown_fields: UnknownFields, + cached_size: CachedSize, ++ range: Option<(u64, u64)>, + } + + impl SpecialFields { +@@ -14,6 +15,7 @@ + SpecialFields { + unknown_fields: UnknownFields::new(), + cached_size: CachedSize::new(), ++ range: None, + } + } + +@@ -36,5 +38,15 @@ + /// Setter. + pub fn mut_unknown_fields(&mut self) -> &mut UnknownFields { + &mut self.unknown_fields ++ } ++ ++ /// Set range ++ pub fn set_range(&mut self, range: (u64, u64)) { ++ self.range = Some(range); ++ } ++ ++ /// Get range ++ pub fn range(&self) -> (u64, u64) { ++ self.range.unwrap() + } + } diff --git a/patches/protobuf+3.5.0-reflection-range.patch b/patches/protobuf+3.5.0-reflection-range.patch new file mode 100644 index 0000000..681162d --- /dev/null +++ b/patches/protobuf+3.5.0-reflection-range.patch @@ -0,0 +1,23 @@ +diff -ruN a/src/reflect/dynamic/mod.rs b/src/reflect/dynamic/mod.rs +--- a/src/reflect/dynamic/mod.rs 2025-07-21 16:11:49 ++++ b/src/reflect/dynamic/mod.rs 2025-07-21 16:11:08 +@@ -341,6 +341,8 @@ + } + + fn merge_from_dyn(&mut self, is: &mut CodedInputStream) -> crate::Result<()> { ++ let start_pos = is.pos(); ++ + while !is.eof()? { + let (field, wire_type) = is.read_tag_unpack()?; + let field_desc = match self.descriptor.field_by_number(field) { +@@ -387,6 +389,10 @@ + } + } + } ++ ++ let end_pos = is.pos(); ++ self.special_fields.set_range((start_pos, end_pos)); ++ + Ok(()) + } +