google_cloud_bigquery/write/format/
arrow.rs1use crate::model::append_rows_request::ArrowData;
16use crate::model::{AppendRowsRequest, ArrowRecordBatch, ArrowSchema};
17
18#[derive(Debug)]
22pub struct Arrow {
23 pub(crate) schema: ArrowSchema,
24}
25
26impl super::DataFormat for Arrow {
27 type Rows = ArrowRecordBatch;
28}
29
30impl super::sealed::DataFormat for Arrow {
31 fn format_name(&self) -> &'static str {
32 "arrow"
33 }
34
35 fn make_request(&self, write_stream: &str, rows: ArrowRecordBatch) -> AppendRowsRequest {
36 AppendRowsRequest::new()
37 .set_write_stream(write_stream)
38 .set_arrow_rows(
39 ArrowData::new()
40 .set_writer_schema(self.schema.clone())
41 .set_rows(rows),
42 )
43 }
44}
45
46#[cfg(test)]
47mod tests {
48 use super::super::sealed::DataFormat;
49 use super::*;
50 use crate::write::test::*;
51
52 #[test]
53 fn format_name() {
54 let f = Arrow { schema: schema() };
55 assert_eq!(f.format_name(), "arrow");
56 }
57
58 #[test]
59 fn request() {
60 let f = Arrow { schema: schema() };
61
62 let req = f.make_request(&write_stream(), rows(1));
63 let data = req.arrow_rows().expect("arrow rows should be set");
64 let s = data.writer_schema.as_ref().expect("schema should be set");
65 assert_eq!(s.serialized_schema, "test");
66 let r = data.rows.as_ref().expect("rows should be set");
67 assert_eq!(r.serialized_record_batch, "1");
68
69 let req = f.make_request(&write_stream(), rows(2));
70 let data = req.arrow_rows().expect("arrow rows should be set");
71 let s = data.writer_schema.as_ref().expect("schema should be set");
72 assert_eq!(s.serialized_schema, "test");
73 let r = data.rows.as_ref().expect("rows should be set");
74 assert_eq!(r.serialized_record_batch, "2");
75 }
76
77 fn rows(id: i64) -> ArrowRecordBatch {
78 ArrowRecordBatch::new().set_serialized_record_batch(id.to_string())
79 }
80}