alamb commented on code in PR #16986: URL: https://github.com/apache/datafusion/pull/16986#discussion_r2249201460
########## datafusion/proto/src/physical_plan/mod.rs: ########## @@ -2941,12 +2941,126 @@ impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec { } } +/// DataEncoderTuple captures the position of the encoder +/// in the codec list that was used to encode the data and actual encoded data +#[derive(Clone, PartialEq, prost::Message)] +struct DataEncoderTuple { + /// The position of encoder used to encode data + /// (to be used for decoding) + #[prost(uint32, tag = 1)] + pub encoder_position: u32, + + #[prost(bytes, tag = 2)] + pub blob: Vec<u8>, +} + +/// A PhysicalExtensionCodec that tries one of multiple inner codecs +/// until one works +#[derive(Debug)] +pub struct ComposedPhysicalExtensionCodec { + codecs: Vec<Arc<dyn PhysicalExtensionCodec>>, +} + +impl ComposedPhysicalExtensionCodec { + // Position in this codesc list is important as it will be used for decoding. + // If new codec is added it should go to last position. + pub fn new(codecs: Vec<Arc<dyn PhysicalExtensionCodec>>) -> Self { + Self { codecs } + } + + fn decode_protobuf<R>( + &self, + buf: &[u8], + decode: impl FnOnce(&dyn PhysicalExtensionCodec, &[u8]) -> Result<R>, + ) -> Result<R> { + let proto = DataEncoderTuple::decode(buf) + .map_err(|e| DataFusionError::Internal(e.to_string()))?; + + let codec = self.codecs.get(proto.encoder_position as usize).ok_or( + DataFusionError::Internal( + "Can't find required codec in codec list".to_owned(), + ), + )?; + + decode(codec.as_ref(), &proto.blob) + } + + fn encode_protobuf( + &self, + buf: &mut Vec<u8>, + mut encode: impl FnMut(&dyn PhysicalExtensionCodec, &mut Vec<u8>) -> Result<()>, + ) -> Result<()> { + let mut data = vec![]; + let mut last_err = None; Review Comment: Maybe we can add that as a follow on PR to the example -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org For additional commands, e-mail: github-h...@datafusion.apache.org