Skip to main content

google_cloud_bigquery/write/format/
arrow.rs

1// Copyright 2026 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::model::append_rows_request::ArrowData;
16use crate::model::{AppendRowsRequest, ArrowRecordBatch, ArrowSchema};
17
18/// Represents the [Arrow] data format for a writer.
19///
20/// [Arrow]: https://arrow.apache.org/
21#[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}