lance_file/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use arrow_array::{ArrayRef, RecordBatch};
7use bytes::Bytes;
8use lance_core::{Result, datatypes::Schema};
9use lance_encoding::decoder::{ColumnInfo, PageEncoding};
10use lance_io::object_store::ObjectStore;
11use object_store::path::Path;
12use prost::Message;
13use prost_types::Any;
14
15use crate::{format::pbfile, versions};
16
17pub(crate) mod structural;
18
19/// Page buffers in current Lance files are aligned to 64 bytes.
20pub(crate) const PAGE_BUFFER_ALIGNMENT: usize = 64;
21pub(crate) const ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES: &str = "LANCE_FILE_WRITER_MAX_PAGE_BYTES";
22
23/// Summary of a completed Lance file write.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct FileWriteSummary {
26    /// The number of rows written to the file.
27    pub num_rows: u64,
28    /// The final size of the file in bytes.
29    pub size_bytes: u64,
30}
31
32/// Runtime options shared by all current-format writers.
33///
34/// These options control buffering and execution only. Select the file grammar
35/// by constructing a writer under [`crate::versions`].
36#[derive(Debug, Clone, Default)]
37pub struct FileWriterOptions {
38    /// How many bytes to use for buffering column data.
39    ///
40    /// The budget is divided evenly across columns. The default is 8 MiB per
41    /// column.
42    pub data_cache_bytes: Option<u64>,
43    /// A best-effort maximum encoded page size.
44    pub max_page_bytes: Option<u64>,
45    /// Keep input arrays instead of copying buffered slices.
46    ///
47    /// Do not enable this for arrays arriving through the Arrow C data
48    /// interface because a small child array can keep an entire batch alive.
49    pub keep_original_array: Option<bool>,
50}
51
52/// A type-erased current-format file writer.
53///
54/// This enum exists for callers that select a concrete file version at
55/// runtime. Each variant owns the complete implementation for exactly one file
56/// grammar; this type only forwards operations without adding format policy.
57pub enum FileWriter {
58    V2_0(Box<versions::v2_0::Writer>),
59    V2_1(Box<versions::v2_1::Writer>),
60    V2_2(Box<versions::v2_2::Writer>),
61    V2_3(Box<versions::v2_3::Writer>),
62}
63
64fn column_info_to_metadata(column: &ColumnInfo) -> Result<pbfile::ColumnMetadata> {
65    let pages = column
66        .page_infos
67        .iter()
68        .map(|page| {
69            let encoding = match &page.encoding {
70                PageEncoding::Legacy(encoding) => Any::from_msg(encoding)?.encode_to_vec(),
71                PageEncoding::Structural(encoding) => Any::from_msg(encoding)?.encode_to_vec(),
72            };
73            let (buffer_offsets, buffer_sizes) =
74                page.buffer_offsets_and_sizes.iter().copied().unzip();
75            Ok(pbfile::column_metadata::Page {
76                buffer_offsets,
77                buffer_sizes,
78                encoding: Some(pbfile::Encoding {
79                    location: Some(pbfile::encoding::Location::Direct(pbfile::DirectEncoding {
80                        encoding,
81                    })),
82                }),
83                length: page.num_rows,
84                priority: page.priority,
85            })
86        })
87        .collect::<Result<Vec<_>>>()?;
88    let (buffer_offsets, buffer_sizes) = column.buffer_offsets_and_sizes.iter().copied().unzip();
89    let encoding = Any::from_msg(&column.encoding)?.encode_to_vec();
90    Ok(pbfile::ColumnMetadata {
91        pages,
92        buffer_offsets,
93        buffer_sizes,
94        encoding: Some(pbfile::Encoding {
95            location: Some(pbfile::encoding::Location::Direct(pbfile::DirectEncoding {
96                encoding,
97            })),
98        }),
99    })
100}
101
102impl From<versions::v2_0::Writer> for FileWriter {
103    fn from(writer: versions::v2_0::Writer) -> Self {
104        Self::V2_0(Box::new(writer))
105    }
106}
107
108impl From<versions::v2_1::Writer> for FileWriter {
109    fn from(writer: versions::v2_1::Writer) -> Self {
110        Self::V2_1(Box::new(writer))
111    }
112}
113
114impl From<versions::v2_2::Writer> for FileWriter {
115    fn from(writer: versions::v2_2::Writer) -> Self {
116        Self::V2_2(Box::new(writer))
117    }
118}
119
120impl From<versions::v2_3::Writer> for FileWriter {
121    fn from(writer: versions::v2_3::Writer) -> Self {
122        Self::V2_3(Box::new(writer))
123    }
124}
125
126impl FileWriter {
127    /// Spill page metadata to a sidecar file.
128    pub fn with_page_metadata_spill(self, object_store: Arc<ObjectStore>, path: Path) -> Self {
129        match self {
130            Self::V2_0(writer) => Self::V2_0(Box::new(
131                (*writer).with_page_metadata_spill(object_store, path),
132            )),
133            Self::V2_1(writer) => Self::V2_1(Box::new(
134                (*writer).with_page_metadata_spill(object_store, path),
135            )),
136            Self::V2_2(writer) => Self::V2_2(Box::new(
137                (*writer).with_page_metadata_spill(object_store, path),
138            )),
139            Self::V2_3(writer) => Self::V2_3(Box::new(
140                (*writer).with_page_metadata_spill(object_store, path),
141            )),
142        }
143    }
144
145    /// Schedule batches of data to be written to the file.
146    pub async fn write_batches(
147        &mut self,
148        batches: impl Iterator<Item = &RecordBatch>,
149    ) -> Result<()> {
150        for batch in batches {
151            self.write_batch(batch).await?;
152        }
153        Ok(())
154    }
155
156    /// Schedule a batch of data to be written to the file.
157    pub async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
158        match self {
159            Self::V2_0(writer) => writer.write_batch(batch).await,
160            Self::V2_1(writer) => writer.write_batch(batch).await,
161            Self::V2_2(writer) => writer.write_batch(batch).await,
162            Self::V2_3(writer) => writer.write_batch(batch).await,
163        }
164    }
165
166    /// Write one top-level column.
167    pub async fn write_column(&mut self, column_index: usize, array: ArrayRef) -> Result<()> {
168        match self {
169            Self::V2_0(writer) => writer.write_column(column_index, array).await,
170            Self::V2_1(writer) => writer.write_column(column_index, array).await,
171            Self::V2_2(writer) => writer.write_column(column_index, array).await,
172            Self::V2_3(writer) => writer.write_column(column_index, array).await,
173        }
174    }
175
176    /// Append a buffer whose page or column metadata is supplied externally.
177    pub async fn write_external_buffer(&mut self, bytes: &[u8]) -> Result<(u64, u64)> {
178        match self {
179            Self::V2_0(writer) => writer.write_external_buffer(bytes).await,
180            Self::V2_1(writer) => writer.write_external_buffer(bytes).await,
181            Self::V2_2(writer) => writer.write_external_buffer(bytes).await,
182            Self::V2_3(writer) => writer.write_external_buffer(bytes).await,
183        }
184    }
185
186    /// Add a metadata entry to the schema.
187    pub fn add_schema_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
188        let key = key.into();
189        let value = value.into();
190        match self {
191            Self::V2_0(writer) => writer.add_schema_metadata(key, value),
192            Self::V2_1(writer) => writer.add_schema_metadata(key, value),
193            Self::V2_2(writer) => writer.add_schema_metadata(key, value),
194            Self::V2_3(writer) => writer.add_schema_metadata(key, value),
195        }
196    }
197
198    /// Prepare a writer from encoded columns whose buffers were produced externally.
199    pub fn initialize_with_external_columns(
200        &mut self,
201        schema: Schema,
202        columns: &[Arc<ColumnInfo>],
203        rows_written: u64,
204    ) -> Result<()> {
205        let column_metadata = columns
206            .iter()
207            .map(|column| column_info_to_metadata(column))
208            .collect::<Result<Vec<_>>>()?;
209        match self {
210            Self::V2_0(writer) => {
211                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
212            }
213            Self::V2_1(writer) => {
214                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
215            }
216            Self::V2_2(writer) => {
217                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
218            }
219            Self::V2_3(writer) => {
220                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
221            }
222        }
223        Ok(())
224    }
225
226    /// Add an arbitrary global buffer and return its one-based index.
227    pub async fn add_global_buffer(&mut self, buffer: Bytes) -> Result<u32> {
228        match self {
229            Self::V2_0(writer) => writer.add_global_buffer(buffer).await,
230            Self::V2_1(writer) => writer.add_global_buffer(buffer).await,
231            Self::V2_2(writer) => writer.add_global_buffer(buffer).await,
232            Self::V2_3(writer) => writer.add_global_buffer(buffer).await,
233        }
234    }
235
236    /// Finish the file and close its object writer.
237    pub async fn finish(&mut self) -> Result<FileWriteSummary> {
238        match self {
239            Self::V2_0(writer) => writer.finish().await,
240            Self::V2_1(writer) => writer.finish().await,
241            Self::V2_2(writer) => writer.finish().await,
242            Self::V2_3(writer) => writer.finish().await,
243        }
244    }
245
246    /// Abandon the file write.
247    pub async fn abort(&mut self) {
248        match self {
249            Self::V2_0(writer) => writer.abort().await,
250            Self::V2_1(writer) => writer.abort().await,
251            Self::V2_2(writer) => writer.abort().await,
252            Self::V2_3(writer) => writer.abort().await,
253        }
254    }
255
256    /// Return the current object-writer position.
257    pub async fn tell(&mut self) -> Result<u64> {
258        match self {
259            Self::V2_0(writer) => writer.tell().await,
260            Self::V2_1(writer) => writer.tell().await,
261            Self::V2_2(writer) => writer.tell().await,
262            Self::V2_3(writer) => writer.tell().await,
263        }
264    }
265
266    /// Return the field-id to physical-column mapping.
267    pub fn field_id_to_column_indices(&self) -> &[(u32, u32)] {
268        match self {
269            Self::V2_0(writer) => writer.field_id_to_column_indices(),
270            Self::V2_1(writer) => writer.field_id_to_column_indices(),
271            Self::V2_2(writer) => writer.field_id_to_column_indices(),
272            Self::V2_3(writer) => writer.field_id_to_column_indices(),
273        }
274    }
275}
276
277#[cfg(test)]
278#[path = "writer_tests.rs"]
279mod writer_tests;
lance_file/writer.rs

lance_file/
writer.rs
