diff --git a/src/io/parquet/write/file.rs b/src/io/parquet/write/file.rs index 72b95574b18..27c62f6edf5 100644 --- a/src/io/parquet/write/file.rs +++ b/src/io/parquet/write/file.rs @@ -83,8 +83,13 @@ impl FileWriter { } /// Writes the footer of the parquet file. Returns the total size of the file. - pub fn end(self, key_value_metadata: Option>) -> Result<(u64, W)> { + pub fn end(&mut self, key_value_metadata: Option>) -> Result { let key_value_metadata = add_arrow_schema(&self.schema, key_value_metadata); Ok(self.writer.end(key_value_metadata)?) } + + /// Consumes this writer and returns the inner writer + pub fn into_inner(self) -> W { + self.writer.into_inner() + } } diff --git a/tests/it/io/parquet/mod.rs b/tests/it/io/parquet/mod.rs index 4d6a25cfceb..7c8daf23641 100644 --- a/tests/it/io/parquet/mod.rs +++ b/tests/it/io/parquet/mod.rs @@ -731,9 +731,9 @@ fn integration_write(schema: &Schema, batches: &[Chunk>]) -> Resu for group in row_groups { writer.write(group?)?; } - let (_size, writer) = writer.end(None)?; + writer.end(None)?; - Ok(writer.into_inner()) + Ok(writer.into_inner().into_inner()) } type IntegrationRead = (Schema, Vec>>); diff --git a/tests/it/io/parquet/read_indexes.rs b/tests/it/io/parquet/read_indexes.rs index 43bb5faeff3..ade8b268f16 100644 --- a/tests/it/io/parquet/read_indexes.rs +++ b/tests/it/io/parquet/read_indexes.rs @@ -96,7 +96,8 @@ fn read_with_indexes( writer.start()?; writer.write(row_group)?; - let (_, data) = writer.end(None)?; + writer.end(None)?; + let data = writer.into_inner(); let mut reader = Cursor::new(data); diff --git a/tests/it/io/parquet/write.rs b/tests/it/io/parquet/write.rs index a068a8335a1..06f3a706f02 100644 --- a/tests/it/io/parquet/write.rs +++ b/tests/it/io/parquet/write.rs @@ -51,9 +51,9 @@ fn round_trip( for group in row_groups { writer.write(group?)?; } - let (_size, writer) = writer.end(None)?; + writer.end(None)?; - let data = writer.into_inner(); + let data = writer.into_inner().into_inner(); let (result, stats) = read_column(&mut Cursor::new(data), 0, "a1")?; assert_eq!(array.as_ref(), result.as_ref());