Clean up and document the API

Author: the80srobot

Source/santad/Logs/EndpointSecurity/ParquetLogger/column_builder.rs

@@ -18,6 +18,8 @@ pub struct ColumnBuilder {
 }
 
 impl ColumnBuilder {
+    /// Create a new column builder. Providing invalid options will result in
+    /// failures on push, not immediately.
     pub fn new(
         page_size: usize,
         descriptor: Descriptor,
@@ -31,6 +33,12 @@ impl ColumnBuilder {
         }
     }
 
+    /// Compress and return the buffered pages so they can be written to a file.
+    ///
+    /// WARNING: due to parquet2's iterator-centric design, it's necessary to
+    /// drain this iterator before calling push again. Otherwise, it's undefined
+    /// which row group the newly written data will end up, and it could even be
+    /// dropped.
     pub fn drain<'a>(&'a mut self) -> DynStreamingIterator<'a, CompressedPage, Error> {
         let pages: Vec<Result<Page>> = self
             .pages
@@ -45,10 +53,16 @@ impl ColumnBuilder {
         DynStreamingIterator::new(compressor)
     }
 
+    /// Append the value to the most recent page. If the page is full, create a
+    /// new one.
     pub fn push(&mut self, value: Value) -> Result<()> {
         self.page_builder(value.dyn_size())?.push(value)
     }
 
+    /// Return the most recent, partially built page. If the page can't fit the
+    /// size_hint without going over page_size, a new page is created.
+    ///
+    /// This call can only fail if the schema is invalid.
     pub fn page_builder(&mut self, size_hint: usize) -> Result<&mut PageBuilder> {
         let last_page = match self.pages.last_mut() {
             Some(page) => page,
@@ -71,10 +85,13 @@ impl ColumnBuilder {
         Ok(self.pages.last_mut().unwrap())
     }
 
+    /// Returns the current size of the column in bytes, as a sum of the sizes
+    /// of all buffered pages.
     pub fn size(&self) -> usize {
         self.pages.iter().map(|page| page.size()).sum()
     }
 
+    /// Returns the current number of buffered values.
     pub fn count(&self) -> usize {
         self.pages.iter().map(|page| page.count()).sum()
     }

Source/santad/Logs/EndpointSecurity/ParquetLogger/cpp_api.rs

@@ -135,10 +135,11 @@ fn table_push_string(table: &mut Table, column_no: usize, value: &CxxString) ->
     table.push(column_no, Value::Bytes(value.as_bytes()))
 }
 
-fn table_flush(table: &mut Table) -> Result<(usize), Error> {
+fn table_flush(table: &mut Table) -> Result<usize, Error> {
     table.flush()
 }
 
-fn table_end(mut table: Box<Table>) -> Result<u64, Error> {
-    table.end()
+fn table_end(table: Box<Table>) -> Result<u64, Error> {
+    let (n, _writer) = table.end()?;
+    Ok(n)
 }

Source/santad/Logs/EndpointSecurity/ParquetLogger/page_builder.rs

@@ -11,9 +11,9 @@ use parquet2::{
 };
 use std::cmp::PartialOrd;
 
-// A page builder serializes primitive values into bytes and appends them to a
-// page (buffer). Implementations are provided for NativeType and &[u8] (byte
-// array).
+/// A page builder serializes primitive values into bytes and appends them to a
+/// page (buffer). Implementations are provided for NativeType and &[u8] (byte
+/// array).
 pub struct PageBuilder {
     page_builder: InnerBuilder,
 }
@@ -93,11 +93,8 @@ enum InnerBuilder {
     F64(NativePage<f64>),
 }
 
-// A page builder serializes primitive values into bytes and appends them to a
-// page (buffer). Implementations are provided for NativeType and &[u8] (byte
-// array).
-
-// Builds a page of variable length by arrays. Used for strings and other blobs.
+/// Builds a page of variable length by arrays. Used for strings and other
+/// blobs.
 pub struct ByteArrayPage {
     buffer: Vec<u8>,
     count: usize,
@@ -151,8 +148,8 @@ impl ByteArrayPage {
     }
 }
 
-// A page of numbers using plain encoding. This is implemented (and fast) for
-// most native numeric types. (Int96 isn't used at the moment.)
+/// A page of numbers using plain encoding. This is implemented (and fast) for
+/// most native numeric types. (Int96 isn't used at the moment.)
 pub struct NativePage<T: NativeType + PartialOrd> {
     buffer: Vec<u8>,
     count: usize,

Source/santad/Logs/EndpointSecurity/ParquetLogger/parquet_logger.rs

@@ -1,48 +1,48 @@
 //! This package provides an opinionated API for producing a Parquet file
 //! containing a simple table. It's intended to be easy to use from both Rust
 //! and C++ code, and uses Cxx to expose a C++ API. (See cpp_api.rs.)
-//! 
+//!
 //! We take the following simplifying assumptions:
-//! 
+//!
 //! * All fields are always required (no NULLs).
 //! * All fields are simple types: integers, floats and strings.
 //! * All files are brotli-compressed.
-//! 
+//!
 //! The API provides reasonable defaults for many of the knobs Parquet exposes,
 //! and doesn't allow overriding most of them. This is intentional - the goal is
 //! to be as simple to use as possible.
-//! 
+//!
 //! To get started from C++, look at cpp_api.rs. To get started from Rust, look
 //! at the Table type in table.rs.
-//! 
+//!
 //! # Implementation Notes
-//! 
+//!
 //! The API is implemented on top of parquet2, a minimal reimplamentation of the
 //! official arrow crate. We chose parquet2 for its simplicity, compilation
 //! speed and lack of unsafe code. (The official arrow project is extremely
 //! large and depends on Boost in C++.)
-//! 
+//!
 //! The code structure roughly mirrors that of a parquet file:
-//! 
+//!
 //! * Table: represents a parquet file, which consists of one or more row
 //!   groups.
 //! * ColumnBuilder: represents a column chunk in a row group.
 //! * PageBuilder: represents a data page in a column chunk.
 //! * Value: represents a single scalar (number or byte blob) in a data page.
-//! 
+//!
 //! Correctness, including of types, is enforced at runtime. Value, rather than
 //! being a generic type, is an enumeration (discriminated union) that can hold
 //! any of the supported types. This is done for two reasons:
-//! 
+//!
 //! 1. It makes the code eaiser to understand - multiple layers of generic
 //!    traits are required for static type checking of column chunks and pages.
 //! 2. The Table type must expose a runtime-generic way of setting a cell in a
 //!    column, and this is the most common way of using the API, so any savings
 //!    gained from static type checking would be bypassed by the most common
 //!    code path anyway.
-//! 
+//!
 //! # Future Work
-//! 
+//!
 //! * Support fixed-length byte arrays.
 //! * Reimplement FileWriter to use an arena-style buffer instead of nested
 //!   iterators.
@@ -67,6 +67,9 @@ pub extern "C" fn parquet2_1337_bloom_filter_contains(x: i64) -> bool {
     bloom_filter::is_in_set(&bits, bloom_filter::hash_native(x))
 }
 
+// This is the main public API.
+pub use crate::table::Table;
+
 #[cfg(test)]
 mod test {
     use crate::{
@@ -89,7 +92,6 @@ mod test {
                 version: Version::V1,
             },
             compression_options: CompressionOptions::Brotli(Some(BrotliLevel::try_new(5).unwrap())),
-            // compression_options: CompressionOptions::Uncompressed,
             page_size: 1024,
         };
 
@@ -117,9 +119,8 @@ mod test {
                 .expect("push failed");
         }
         table.flush().expect("flush failed");
-        table.end().expect("end failed");
+        let (_, writer) = table.end().expect("end failed");
 
-        let (_schema, writer, _options) = table.into_inner();
         let writer = if let Writer::Memory(w) = writer {
             w
         } else {

Source/santad/Logs/EndpointSecurity/ParquetLogger/table.rs

@@ -18,8 +18,6 @@ pub struct Options {
 /// easy to construct.)
 pub struct Table {
     columns: Vec<ColumnBuilder>,
-    schema: SchemaDescriptor,
-    options: Options,
     writer: Writer,
 }
 
@@ -36,18 +34,22 @@ impl Table {
                 )
             })
             .collect::<Vec<_>>();
-        Self {
-            columns,
-            schema,
-            options,
-            writer,
-        }
+        Self { columns, writer }
     }
 
+    /// Pushes a value to the column. The column number is the index the column
+    /// had in the schema. The type of Value must match the type of the column
+    /// as specified in new.
+    ///
+    /// It is not necessary to push values in order of column number. It is also
+    /// not required to push one row at a time (it's fine to write column by
+    /// column). However, the same number of values must be pushed to each
+    /// column by the time flush is called.
     pub fn push(&mut self, column_no: usize, value: Value) -> Result<()> {
         self.columns[column_no].push(value)
     }
 
+    /// Convenience method for pushing a row of values at a time. See push.
     pub fn push_row<'a, I>(&mut self, values: I) -> Result<()>
     where
         I: Iterator<Item = Value<'a>>,
@@ -58,6 +60,7 @@ impl Table {
         Ok(())
     }
 
+    /// Convenience method for pushing a column of values at a time. See push
     pub fn push_column<'a, I>(&mut self, column_no: usize, values: I) -> Result<()>
     where
         I: Iterator<Item = Value<'a>>,
@@ -68,15 +71,15 @@ impl Table {
         Ok(())
     }
 
-    // Checks that the table is well-formed and returns the number of rows in
-    // the buffer.
-    pub fn validate(&self) -> Result<(usize)> {
+    /// Checks that the table is well-formed and returns the number of rows in
+    /// the buffer.
+    pub fn validate(&self) -> Result<usize> {
         match self.columns.len() {
             0 => Err(parquet2::error::Error::OutOfSpec("No columns".to_string())),
             _ => {
                 let n = self.columns[0].count();
                 if self.columns.iter().all(|column| column.count() == n) {
-                    Ok((n))
+                    Ok(n)
                 } else {
                     Err(parquet2::error::Error::OutOfSpec(
                         "Column counts don't match".to_string(),
@@ -86,25 +89,30 @@ impl Table {
         }
     }
 
-    // Flushes a rowg group to the writer and returns the number of rows
-    // flushed. Does nothing if no rows are buffered.
-    pub fn flush(&mut self) -> Result<(usize)> {
+    /// Flushes a rowg group to the writer and returns the number of rows
+    /// flushed. Does nothing if no rows are buffered.
+    pub fn flush(&mut self) -> Result<usize> {
         match self.validate() {
-            Ok(0) => Ok((0)),
+            Ok(0) => Ok(0),
             Ok(n) => {
                 write_row_group(&mut self.writer, &mut self.columns)?;
-                Ok((n))
+                Ok(n)
             }
             Err(e) => Err(e),
         }
     }
 
-    pub fn end(&mut self) -> Result<u64> {
+    /// Flushes all buffered data and ends the file, writing the footer.
+    pub fn end(mut self) -> Result<(u64, Writer)> {
         self.flush()?;
-        self.writer.end()
+        let n = self.writer.end()?;
+        Ok((n, self.writer))
     }
 
-    pub fn into_inner(self) -> (SchemaDescriptor, Writer, Options) {
-        (self.schema, self.writer, self.options)
+    /// Return the total buffered size of the table in bytes. This does not
+    /// count bytes already written to disk, or any metadata and header and
+    /// footer.
+    pub fn size(&self) -> usize {
+        self.columns.iter().map(|column| column.size()).sum()
     }
 }

Source/santad/Logs/EndpointSecurity/ParquetLogger/value.rs

@@ -1,11 +1,11 @@
 use parquet2::error::Result;
 
-// A value can be written to a page in a column chunk in a row group in a
-// parquet file in the house that Jack built.
-//
-// Parquet only supports a handful of physical types. In addition to what's
-// listed in this enum, parquet supports int96 and fixed length arrays, which
-// are not yet implemented here.
+/// A value can be written to a page in a column chunk in a row group in a
+/// parquet file in the house that Jack built.
+///
+/// Parquet only supports a handful of physical types. In addition to what's
+/// listed in this enum, parquet supports int96 and fixed length arrays, which
+/// are not yet implemented here.
 pub enum Value<'a> {
     I32(i32),
     I64(i64),

Source/santad/Logs/EndpointSecurity/ParquetLogger/writer.rs

@@ -9,8 +9,8 @@ use parquet2::{
 
 use crate::column_builder::ColumnBuilder;
 
-// Wraps the FileWriter for Table to allow constructing the latter from C++.
-// (FileWriter is generic, but Table cannot be.)
+/// Wraps the FileWriter for Table to allow constructing the latter from C++.
+/// (FileWriter is generic, but Table cannot be.)
 pub enum Writer {
     Memory(FileWriter<Vec<u8>>),
     File(FileWriter<File>),