From 885bcbc1edf6b6642e3fbc3287b41fcbd471b987 Mon Sep 17 00:00:00 2001 From: Navin Date: Mon, 26 Jul 2021 22:00:56 +1000 Subject: [PATCH] Doctests for UnionBuilder and UnionArray. (#603) * Doctests for UnionBuilder and UnionArray. * Update arrow/src/array/array_union.rs Co-authored-by: Andrew Lamb Co-authored-by: Andrew Lamb --- arrow/src/array/array_union.rs | 69 ++++++---------------------------- arrow/src/array/builder.rs | 41 ++++++++++++++++++++ 2 files changed, 52 insertions(+), 58 deletions(-) diff --git a/arrow/src/array/array_union.rs b/arrow/src/array/array_union.rs index e369037244ac..2cc13c201ed3 100644 --- a/arrow/src/array/array_union.rs +++ b/arrow/src/array/array_union.rs @@ -15,64 +15,8 @@ // specific language governing permissions and limitations // under the License. -//! Contains the `UnionArray` type. -//! -//! Each slot in a `UnionArray` can have a value chosen from a number of types. Each of the -//! possible types are named like the fields of a [`StructArray`](crate::array::StructArray). -//! A `UnionArray` can have two possible memory layouts, "dense" or "sparse". For more information -//! on please see the [specification](https://arrow.apache.org/docs/format/Columnar.html#union-layout). -//! -//! Builders are provided for `UnionArray`'s involving primitive types. `UnionArray`'s of nested -//! types are also supported but not via `UnionBuilder`, see the tests for examples. -//! -//! # Example: Dense Memory Layout -//! -//! ``` -//! use arrow::array::UnionBuilder; -//! use arrow::datatypes::{Float64Type, Int32Type}; -//! -//! # fn main() -> arrow::error::Result<()> { -//! let mut builder = UnionBuilder::new_dense(3); -//! builder.append::("a", 1).unwrap(); -//! builder.append::("b", 3.0).unwrap(); -//! builder.append::("a", 4).unwrap(); -//! let union = builder.build().unwrap(); -//! -//! assert_eq!(union.type_id(0), 0_i8); -//! assert_eq!(union.type_id(1), 1_i8); -//! assert_eq!(union.type_id(2), 0_i8); -//! -//! assert_eq!(union.value_offset(0), 0_i32); -//! assert_eq!(union.value_offset(1), 0_i32); -//! assert_eq!(union.value_offset(2), 1_i32); -//! -//! # Ok(()) -//! # } -//! ``` -//! -//! # Example: Sparse Memory Layout -//! ``` -//! use arrow::array::UnionBuilder; -//! use arrow::datatypes::{Float64Type, Int32Type}; -//! -//! # fn main() -> arrow::error::Result<()> { -//! let mut builder = UnionBuilder::new_sparse(3); -//! builder.append::("a", 1).unwrap(); -//! builder.append::("b", 3.0).unwrap(); -//! builder.append::("a", 4).unwrap(); -//! let union = builder.build().unwrap(); -//! -//! assert_eq!(union.type_id(0), 0_i8); -//! assert_eq!(union.type_id(1), 1_i8); -//! assert_eq!(union.type_id(2), 0_i8); -//! -//! assert_eq!(union.value_offset(0), 0_i32); -//! assert_eq!(union.value_offset(1), 1_i32); -//! assert_eq!(union.value_offset(2), 2_i32); -//! -//! # Ok(()) -//! # } -//! ``` +/// Contains the `UnionArray` type. +/// use crate::array::{data::count_nulls, make_array, Array, ArrayData, ArrayRef}; use crate::buffer::Buffer; use crate::datatypes::*; @@ -83,6 +27,15 @@ use std::any::Any; use std::mem::size_of; /// An Array that can represent slots of varying types. +/// +/// Each slot in a `UnionArray` can have a value chosen from a number of types. Each of the +/// possible types are named like the fields of a [`StructArray`](crate::array::StructArray). +/// A `UnionArray` can have two possible memory layouts, "dense" or "sparse". For more information +/// on please see the [specification](https://arrow.apache.org/docs/format/Columnar.html#union-layout). +/// +/// [`UnionBuilder`]can be used to create `UnionArray`'s of primitive types. `UnionArray`'s of nested +/// types are also supported but not via `UnionBuilder`, see the tests for examples. +/// pub struct UnionArray { data: ArrayData, boxed_fields: Vec, diff --git a/arrow/src/array/builder.rs b/arrow/src/array/builder.rs index ac548e206085..55c4e598cd8e 100644 --- a/arrow/src/array/builder.rs +++ b/arrow/src/array/builder.rs @@ -1706,6 +1706,47 @@ impl FieldData { } /// Builder type for creating a new `UnionArray`. +/// +/// Example: **Dense Memory Layout** +/// +/// ``` +/// use arrow::array::UnionBuilder; +/// use arrow::datatypes::{Float64Type, Int32Type}; +/// +/// let mut builder = UnionBuilder::new_dense(3); +/// builder.append::("a", 1).unwrap(); +/// builder.append::("b", 3.0).unwrap(); +/// builder.append::("a", 4).unwrap(); +/// let union = builder.build().unwrap(); +/// +/// assert_eq!(union.type_id(0), 0_i8); +/// assert_eq!(union.type_id(1), 1_i8); +/// assert_eq!(union.type_id(2), 0_i8); +/// +/// assert_eq!(union.value_offset(0), 0_i32); +/// assert_eq!(union.value_offset(1), 0_i32); +/// assert_eq!(union.value_offset(2), 1_i32); +/// ``` +/// +/// Example: **Sparse Memory Layout** +/// ``` +/// use arrow::array::UnionBuilder; +/// use arrow::datatypes::{Float64Type, Int32Type}; +/// +/// let mut builder = UnionBuilder::new_sparse(3); +/// builder.append::("a", 1).unwrap(); +/// builder.append::("b", 3.0).unwrap(); +/// builder.append::("a", 4).unwrap(); +/// let union = builder.build().unwrap(); +/// +/// assert_eq!(union.type_id(0), 0_i8); +/// assert_eq!(union.type_id(1), 1_i8); +/// assert_eq!(union.type_id(2), 0_i8); +/// +/// assert_eq!(union.value_offset(0), 0_i32); +/// assert_eq!(union.value_offset(1), 1_i32); +/// assert_eq!(union.value_offset(2), 2_i32); +/// ``` #[derive(Debug)] pub struct UnionBuilder { /// The current number of slots in the array