Method extract_message

arrow-flight/src/decode.rs:263–338

Extracts flight data from the next message, updating decoding state as necessary.

(&mut self, data: FlightData)

    /// Extracts flight data from the next message, updating decoding
    /// state as necessary.
    fn extract_message(&mut self, data: FlightData) -> Result<Option<DecodedFlightData>> {
        use arrow_ipc::MessageHeader;
        let message = arrow_ipc::root_as_message(&data.data_header[..])
            .map_err(|e| FlightError::DecodeError(format!("Error decoding root message: {e}")))?;

        match message.header_type() {
            MessageHeader::NONE => Ok(Some(DecodedFlightData::new_none(data))),
            MessageHeader::Schema => {
                let schema = Schema::try_from(&data)
                    .map_err(|e| FlightError::DecodeError(format!("Error decoding schema: {e}")))?;

                let schema = Arc::new(schema);
                let dictionaries_by_field = HashMap::new();

                self.state = Some(FlightStreamState {
                    schema: Arc::clone(&schema),
                    dictionaries_by_field,
                });
                Ok(Some(DecodedFlightData::new_schema(data, schema)))
            }
            MessageHeader::DictionaryBatch => {
                let state = if let Some(state) = self.state.as_mut() {
                    state
                } else {
                    return Err(FlightError::protocol(
                        "Received DictionaryBatch prior to Schema",
                    ));
                };

                let buffer = Buffer::from(data.data_body);
                let dictionary_batch = message.header_as_dictionary_batch().ok_or_else(|| {
                    FlightError::protocol(
                        "Could not get dictionary batch from DictionaryBatch message",
                    )
                })?;

                arrow_ipc::reader::read_dictionary(
                    &buffer,
                    dictionary_batch,
                    &state.schema,
                    &mut state.dictionaries_by_field,
                    &message.version(),
                )
                .map_err(|e| {
                    FlightError::DecodeError(format!("Error decoding ipc dictionary: {e}"))
                })?;

                // Updated internal state, but no decoded message
                Ok(None)
            }
            MessageHeader::RecordBatch => {
                let state = if let Some(state) = self.state.as_ref() {
                    state
                } else {
                    return Err(FlightError::protocol(
                        "Received RecordBatch prior to Schema",
                    ));
                };