()
| 1832 | |
| 1833 | #[test] |
| 1834 | fn test_split_batch_for_grpc_response() { |
| 1835 | let max_flight_data_size = 1024; |
| 1836 | |
| 1837 | // no split |
| 1838 | let c = UInt32Array::from(vec![1, 2, 3, 4, 5, 6]); |
| 1839 | let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(c) as ArrayRef)]) |
| 1840 | .expect("cannot create record batch"); |
| 1841 | let split = split_batch_for_grpc_response(batch.clone(), max_flight_data_size); |
| 1842 | assert_eq!(split.len(), 1); |
| 1843 | assert_eq!(batch, split[0]); |
| 1844 | |
| 1845 | // split once |
| 1846 | let n_rows = max_flight_data_size + 1; |
| 1847 | assert!(n_rows % 2 == 1, "should be an odd number"); |
| 1848 | let c = UInt8Array::from((0..n_rows).map(|i| (i % 256) as u8).collect::<Vec<_>>()); |
| 1849 | let batch = RecordBatch::try_from_iter(vec![("a", Arc::new(c) as ArrayRef)]) |
| 1850 | .expect("cannot create record batch"); |
| 1851 | let split = split_batch_for_grpc_response(batch.clone(), max_flight_data_size); |
| 1852 | assert_eq!(split.len(), 3); |
| 1853 | assert_eq!( |
| 1854 | split.iter().map(|batch| batch.num_rows()).sum::<usize>(), |
| 1855 | n_rows |
| 1856 | ); |
| 1857 | let a = pretty_format_batches(&split).unwrap().to_string(); |
| 1858 | let b = pretty_format_batches(&[batch]).unwrap().to_string(); |
| 1859 | assert_eq!(a, b); |
| 1860 | } |
| 1861 | |
| 1862 | #[test] |
| 1863 | fn test_split_batch_for_grpc_response_sizes() { |
nothing calls this directly
no test coverage detected