Extract `Schema` or `RecordBatch`es from the `FlightData` wire representation
(
flight_data: FlightData,
arrow_schema_ref: &SchemaRef,
)
| 629 | |
| 630 | /// Extract `Schema` or `RecordBatch`es from the `FlightData` wire representation |
| 631 | pub fn arrow_data_from_flight_data( |
| 632 | flight_data: FlightData, |
| 633 | arrow_schema_ref: &SchemaRef, |
| 634 | ) -> std::result::Result<ArrowFlightData, ArrowError> { |
| 635 | let ipc_message = root_as_message(&flight_data.data_header[..]) |
| 636 | .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))?; |
| 637 | |
| 638 | match ipc_message.header_type() { |
| 639 | MessageHeader::RecordBatch => { |
| 640 | let ipc_record_batch = ipc_message.header_as_record_batch().ok_or_else(|| { |
| 641 | ArrowError::ComputeError( |
| 642 | "Unable to convert flight data header to a record batch".to_string(), |
| 643 | ) |
| 644 | })?; |
| 645 | |
| 646 | let dictionaries_by_field = HashMap::new(); |
| 647 | let record_batch = read_record_batch( |
| 648 | &Buffer::from(flight_data.data_body), |
| 649 | ipc_record_batch, |
| 650 | arrow_schema_ref.clone(), |
| 651 | &dictionaries_by_field, |
| 652 | None, |
| 653 | &ipc_message.version(), |
| 654 | )?; |
| 655 | Ok(ArrowFlightData::RecordBatch(record_batch)) |
| 656 | } |
| 657 | MessageHeader::Schema => { |
| 658 | let ipc_schema = ipc_message.header_as_schema().ok_or_else(|| { |
| 659 | ArrowError::ComputeError( |
| 660 | "Unable to convert flight data header to a schema".to_string(), |
| 661 | ) |
| 662 | })?; |
| 663 | |
| 664 | let arrow_schema = fb_to_schema(ipc_schema); |
| 665 | Ok(ArrowFlightData::Schema(arrow_schema)) |
| 666 | } |
| 667 | MessageHeader::DictionaryBatch => { |
| 668 | let _ = ipc_message.header_as_dictionary_batch().ok_or_else(|| { |
| 669 | ArrowError::ComputeError( |
| 670 | "Unable to convert flight data header to a dictionary batch".to_string(), |
| 671 | ) |
| 672 | })?; |
| 673 | Err(ArrowError::NotYetImplemented( |
| 674 | "no idea on how to convert an ipc dictionary batch to an arrow type".to_string(), |
| 675 | )) |
| 676 | } |
| 677 | MessageHeader::Tensor => { |
| 678 | let _ = ipc_message.header_as_tensor().ok_or_else(|| { |
| 679 | ArrowError::ComputeError( |
| 680 | "Unable to convert flight data header to a tensor".to_string(), |
| 681 | ) |
| 682 | })?; |
| 683 | Err(ArrowError::NotYetImplemented( |
| 684 | "no idea on how to convert an ipc tensor to an arrow type".to_string(), |
| 685 | )) |
| 686 | } |
| 687 | MessageHeader::SparseTensor => { |
| 688 | let _ = ipc_message.header_as_sparse_tensor().ok_or_else(|| { |
nothing calls this directly
no test coverage detected