| 438 | } |
| 439 | |
| 440 | async fn roundtrip_with_encoder( |
| 441 | encoder: FlightDataEncoderBuilder, |
| 442 | input_batches: Vec<RecordBatch>, |
| 443 | expected_batches: Vec<RecordBatch>, |
| 444 | ) { |
| 445 | println!("Round tripping with encoder:\n{encoder:#?}"); |
| 446 | |
| 447 | let input_batch_stream = futures::stream::iter(input_batches.clone()).map(Ok); |
| 448 | |
| 449 | let encode_stream = encoder.build(input_batch_stream); |
| 450 | |
| 451 | let decode_stream = FlightRecordBatchStream::new_from_flight_data(encode_stream); |
| 452 | let output_batches: Vec<_> = decode_stream.try_collect().await.expect("encode / decode"); |
| 453 | |
| 454 | // remove any empty batches from input as they are not transmitted |
| 455 | let expected_batches: Vec<_> = expected_batches |
| 456 | .into_iter() |
| 457 | .filter(|b| b.num_rows() > 0) |
| 458 | .collect(); |
| 459 | |
| 460 | assert_eq!(expected_batches, output_batches); |
| 461 | } |
| 462 | |
| 463 | /// Workaround for https://github.com/apache/arrow-rs/issues/1206 |
| 464 | fn prepare_schema_for_flight(schema: &Schema) -> Schema { |