This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-mosaic.git
The following commit(s) were added to refs/heads/main by this push:
new 400f2c4 perf(core): compact integer writer buffers and stream
encoding (#98)
400f2c4 is described below
commit 400f2c4e228c5c2d3a89ed5248a0d11930fa29d7
Author: jianguotian <[email protected]>
AuthorDate: Mon Oct 12 16:28:58 2026 +0800
perf(core): compact integer writer buffers and stream encoding (#98)
---
core/INTEGER_WRITER_MEMORY.md | 127 +++++++
core/examples/integer_writer_memory.rs | 140 +++++++
core/src/bucket_writer.rs | 672 +++++++++++++++++++++++----------
core/src/integer_buffer.rs | 287 ++++++++++++++
core/src/lib.rs | 1 +
core/src/writer.rs | 210 +++++++++++
core/src/writer_integer_tests.rs | 426 +++++++++++++++++++++
7 files changed, 1660 insertions(+), 203 deletions(-)
diff --git a/core/INTEGER_WRITER_MEMORY.md b/core/INTEGER_WRITER_MEMORY.md
new file mode 100644
index 0000000..6fa0816
--- /dev/null
+++ b/core/INTEGER_WRITER_MEMORY.md
@@ -0,0 +1,127 @@
+<!--
+ ~ Licensed to the Apache Software Foundation (ASF) under one
+ ~ or more contributor license agreements. See the NOTICE file
+ ~ distributed with this work for additional information
+ ~ regarding copyright ownership. The ASF licenses this file
+ ~ to you under the Apache License, Version 2.0 (the
+ ~ "License"); you may not use this file except in compliance
+ ~ with the License. You may obtain a copy of the License at
+ ~
+ ~ http://www.apache.org/licenses/LICENSE-2.0
+ ~
+ ~ Unless required by applicable law or agreed to in writing,
+ ~ software distributed under the License is distributed on an
+ ~ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ ~ KIND, either express or implied. See the License for the
+ ~ specific language governing permissions and limitations
+ ~ under the License.
+-->
+
+# Integer writer memory
+
+The writer compacts `Int16`, `Int32`, and `Int64` value buffers automatically.
+This includes integer children of lists/maps and their `Int32` length columns.
+`Int8` already occupies one byte. Floating-point, date/time, timestamp,
decimal,
+string and binary caches retain their existing representation. Input Arrow
+schemas, `WriterOptions`, the C/JNI APIs and on-disk column types are
unchanged.
+
+## Representation and range changes
+
+Each integer block contains at most 4,096 non-null values. A block uses 1, 2,
4,
+or 8 bytes per value, capped at the logical type width. A block-sized OR
reduction
+selects the required width before a batch copy. Positive values such as widened
+UInt16 data therefore occupy two bytes even in an Int32 column. Negative values
+retain their logical two's-complement width.
+
+A later outlier widens only the current block, copying backwards in the same
+vector. Completed blocks stay compact. Range checking and promotion never scan
+or expand the whole buffered column. A new block can narrow again. Reset drops
+block payload allocations; the allocator may retain freed pages in process RSS.
+
+Dictionary keys, constant tracking, statistics, dictionary/plain selection,
+row-group accounting and page-layout selection still use logical values and
+logical byte sizes. Saving memory does not cause the same row group to
accumulate
+more rows. Existing large row-group settings remain effective. Constant columns
+keep dictionary tracking pending; the first differing value activates the
+dictionary with the constant prefix as its first key. Entry limits and reset
+retain the same tracking lifecycle as non-integer columns.
+
+## Encoding and compression
+
+Integer values are expanded into a reusable 16 KiB scratch buffer at encoding
+time. Nullable append gathers at most 4,096 i64 values (32 KiB). Dictionary
indices
+are also emitted incrementally. Buckets containing compact integer columns use
+this sink-based encoder rather than assembling the full bucket/page first.
+Uncompressed integer buckets use an additional 16 KiB output buffer to coalesce
+small writes before invoking file or C/JNI output callbacks. Pending bytes are
+drained explicitly at the end of the bucket; errors or unwinding discard them
+without retrying a potentially partial sink write.
+
+Encoded units up to 1 MiB use bounded staging and the original one-shot Zstd
+compressor to avoid streaming overhead on small pages. Larger units stream to
+one Zstd frame per original bucket/page, with the exact uncompressed size
pledged
+to Zstd. Scratch chunks do not create new pages, frames or row groups.
+
+Paged files place slot lengths before their compressed data, so the writer
still
+retains compressed column slots until it can emit the directory. This change
+does not bound all writer memory: row-group caches, null bitmaps, dictionaries,
+Zstd context/window memory, compressed slots and the caller's Arrow input
remain.
+It requires neither seekable output callbacks nor temporary files. The public
+`BucketWriter::finish`/`finish_paged` convenience methods still materialize
their
+returned data; `MosaicWriter` uses the incremental path for integer buckets.
+
+Streaming Zstd output can differ bytewise from one-shot compression for some
+inputs. Schema, decoded values, encoding choices and compression-unit
boundaries
+remain unchanged; file-size validation is part of the test matrix.
+
+## Validation and reproduction
+
+Unit tests cover width changes within/across blocks, signed extremes, nullable
+and sliced Arrow arrays, ALL_NULL/CONST/DICT/PLAIN transitions, mixed columns,
+logical row-group accounting, statistics, reset, bounded expansion, sink
errors,
+and Zstd block boundaries. A wide-buffer reference path exercises the previous
+writer/compression behavior. The comparison matrix checks input-value round
trips,
+row-group/schema/statistics equality and no file-size increase in that corpus.
+
+Run the synthetic example on the base and candidate revisions using independent
+build directories and separate processes:
+
+```sh
+cargo test --locked --workspace
+cargo clippy --all-targets --workspace -- -D warnings
+cargo build --release -p paimon-mosaic-core --example integer_writer_memory
+/usr/bin/time -v target/release/examples/integer_writer_memory 250000 20 paged
unsigned output.mosaic
+```
+
+Use `monolithic` for the other bucket layout, `negative` for full-range signed
+random data, and `outlier` for a late full-range value. Copy the same example
to
+the baseline revision before building. Compare output hashes, kernel peak RSS,
+phase RSS/PSS snapshots, CPU time and append/close timings. Run each pair in
both
+orders; do not compile or run the two workloads concurrently.
+
+ARM64 Linux results against official `main` at `c082354` (2026-10-12, Rust
+1.97.1, Neoverse-N2; 5 repetitions per version/scenario, separate build
directories,
+alternating version order, 20 batches of 250,000 rows, three integer columns,
+Zstd level 1, dictionary disabled). These are fresh measurements of this change
+on that upstream baseline. Medians:
+
+| Layout | Input | Peak RSS before → after (MiB) | Append total before → after
(ms) | Close before → after (ms) |
+| --- | --- | ---: | ---: | ---: |
+| Paged | Widened unsigned | 232.3 → 79.9 | 47.4 → 14.1 | 271.4 → 231.7 |
+| Paged | Full-range signed | 276.7 → 142.6 | 46.7 → 23.4 | 136.8 → 75.4 |
+| Paged | Late outlier | 232.3 → 79.9 | 47.3 → 14.7 | 270.3 → 232.9 |
+| Monolithic | Widened unsigned | 180.7 → 42.3 | 46.8 → 14.2 | 233.1 → 219.4 |
+| Monolithic | Full-range signed | 210.1 → 75.7 | 47.2 → 23.4 | 90.4 → 50.9 |
+| Monolithic | Late outlier | 180.7 → 42.3 | 46.6 → 13.8 | 231.8 → 218.6 |
+
+All six scenarios produced byte-identical files across all baseline/candidate
+runs. All 90 million integer values in the six candidate files were also
decoded
+and checked against the generated input. Native ARM64 core tests passed (183),
+and the release FFI library built successfully. An additional review harness
+performed 1,728 differential comparisons against the original `c082354` bucket
+encoder, including nested List/Map columns, nulls, slices, dictionary limits
and
+reset; this harness is separate from the committed test suite.
+
+Figures are medians of synthetic runs, not measured savings for a particular
+Android application. Source Arrow arrays are still widened by the caller;
+avoiding those arrays requires a separate input API/client integration change.
diff --git a/core/examples/integer_writer_memory.rs
b/core/examples/integer_writer_memory.rs
new file mode 100644
index 0000000..a9c39d6
--- /dev/null
+++ b/core/examples/integer_writer_memory.rs
@@ -0,0 +1,140 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Reproducible write-only workload for comparing integer cache memory.
+//! Run the same example on the base and candidate revisions in separate
processes:
+//! cargo run --release -p paimon-mosaic-core --example integer_writer_memory
-- 250000 20 paged unsigned output.mosaic
+//! Use /usr/bin/time -v on Linux for peak RSS and CPU; this example also
reports
+//! RSS/PSS snapshots at append/close boundaries. Data is synthetic, not a
vehicle trace.
+
+use arrow_array::{ArrayRef, Int16Array, Int32Array, Int64Array, RecordBatch};
+use arrow_schema::{DataType, Field, Schema};
+use paimon_mosaic_core::writer::{FileSink, MosaicWriter, OutputFile,
WriterOptions};
+use std::sync::Arc;
+use std::time::Instant;
+
+fn memory(stage: &str) {
+ let status =
std::fs::read_to_string("/proc/self/status").unwrap_or_default();
+ let smaps =
std::fs::read_to_string("/proc/self/smaps_rollup").unwrap_or_default();
+ let field = |text: &str, name: &str| -> u64 {
+ text.lines()
+ .find_map(|line| line.strip_prefix(name))
+ .and_then(|v| v.split_whitespace().next())
+ .and_then(|v| v.parse().ok())
+ .unwrap_or(0)
+ };
+ println!(
+ "stage={stage} rss_kib={} peak_rss_kib={} pss_kib={} anon_kib={}",
+ field(&status, "VmRSS:"),
+ field(&status, "VmHWM:"),
+ field(&smaps, "Pss:"),
+ field(&status, "RssAnon:")
+ );
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+ let args: Vec<_> = std::env::args().collect();
+ if args.len() != 6 {
+ return Err(
+ "expected: ROWS_PER_BATCH BATCHES paged|monolithic
unsigned|negative|outlier OUTPUT"
+ .into(),
+ );
+ }
+ let rows: usize = args[1].parse()?;
+ let batches: usize = args[2].parse()?;
+ let paged = match args[3].as_str() {
+ "paged" => true,
+ "monolithic" => false,
+ _ => return Err("invalid layout".into()),
+ };
+ let pattern = args[4].as_str();
+ if !matches!(pattern, "unsigned" | "negative" | "outlier") {
+ return Err("invalid pattern".into());
+ }
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("u8_as_i16", DataType::Int16, false),
+ Field::new("u16_as_i32", DataType::Int32, false),
+ Field::new("u32_as_i64", DataType::Int64, false),
+ ]));
+ let options = WriterOptions {
+ num_buckets: 1,
+ row_group_max_size: u64::MAX,
+ zstd_level: 1,
+ max_dict_entries: 0,
+ page_size_threshold: if paged { 1 } else { usize::MAX },
+ ..Default::default()
+ };
+ let sink = FileSink::create(std::path::Path::new(&args[5]))?;
+ let mut writer = MosaicWriter::new(sink, &schema, options)?;
+ let mut append_ns = 0u128;
+ let mut max_append_ns = 0u128;
+ memory("open");
+ for batch in 0..batches {
+ let mut v16 = Vec::with_capacity(rows);
+ let mut v32 = Vec::with_capacity(rows);
+ let mut v64 = Vec::with_capacity(rows);
+ for row in 0..rows {
+ let mut x = (row + batch * rows) as u64;
+ x ^= x >> 30;
+ x = x.wrapping_mul(0xbf58476d1ce4e5b9);
+ x ^= x >> 27;
+ x = x.wrapping_mul(0x94d049bb133111eb);
+ x ^= x >> 31;
+ if pattern == "negative"
+ || (pattern == "outlier" && batch + 1 == batches && row + 1 ==
rows)
+ {
+ v16.push(x as i16);
+ v32.push(x as i32);
+ v64.push(x as i64);
+ } else {
+ v16.push((x as u8) as i16);
+ v32.push((x as u16) as i32);
+ v64.push((x as u32) as i64);
+ }
+ }
+ let arrays: Vec<ArrayRef> = vec![
+ Arc::new(Int16Array::from(v16)),
+ Arc::new(Int32Array::from(v32)),
+ Arc::new(Int64Array::from(v64)),
+ ];
+ let input = RecordBatch::try_new(schema.clone(), arrays)?;
+ let start = Instant::now();
+ writer.write_batch(&input)?;
+ let elapsed = start.elapsed().as_nanos();
+ append_ns += elapsed;
+ max_append_ns = max_append_ns.max(elapsed);
+ if batch + 1 == batches {
+ memory("last_append");
+ }
+ }
+ memory("before_close");
+ let start = Instant::now();
+ writer.close()?;
+ let close_ns = start.elapsed().as_nanos();
+ memory("after_close");
+ println!(
+ "rows={} append_ns={} max_append_ns={} close_ns={} file_bytes={}",
+ rows * batches,
+ append_ns,
+ max_append_ns,
+ close_ns,
+ writer.output().pos()
+ );
+ drop(writer);
+ memory("after_free");
+ Ok(())
+}
diff --git a/core/src/bucket_writer.rs b/core/src/bucket_writer.rs
index 9828c90..f1da9e2 100644
--- a/core/src/bucket_writer.rs
+++ b/core/src/bucket_writer.rs
@@ -16,18 +16,27 @@
// under the License.
use std::collections::HashMap;
-use std::io;
+use std::io::{self, Write};
use std::sync::Arc;
use arrow_array::*;
use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer,
ScalarBuffer};
use arrow_schema::{DataType, Field};
+use crate::integer_buffer::IntegerBuffer;
use crate::spec::*;
use crate::types;
use crate::values;
use crate::varint;
+pub(crate) struct BucketEncodingPlan {
+ pub(crate) encodings: Vec<u8>,
+ pub(crate) has_nulls: Vec<bool>,
+ // The original finish_paged payload sizes, excluding const metadata and
+ // the two-byte page prefix. Keep these for the existing layout heuristic.
+ pub(crate) page_sizes: Vec<Option<usize>>,
+}
+
pub struct PagedBucketOutput {
pub encodings: Vec<u8>,
pub has_nulls: Vec<bool>,
@@ -68,6 +77,7 @@ pub struct BucketWriter {
null_bitmaps: Vec<Vec<u8>>,
value_buffers: Vec<Vec<u8>>,
+ integer_buffers: Vec<Option<Box<IntegerBuffer>>>,
non_null_counts: Vec<usize>,
const_tracking: Vec<bool>,
@@ -94,13 +104,22 @@ impl BucketWriter {
let (physical_types, children) = expand_col_types(col_types);
let total_columns = physical_types.len();
let fixed_widths: Vec<i32> =
physical_types.iter().map(types::fixed_width).collect();
-
+ let integer_buffers: Vec<_> = physical_types
+ .iter()
+ .map(|dt| match dt {
+ DataType::Int16 => Some(Box::new(IntegerBuffer::new(2))),
+ DataType::Int32 => Some(Box::new(IntegerBuffer::new(4))),
+ DataType::Int64 => Some(Box::new(IntegerBuffer::new(8))),
+ _ => None,
+ })
+ .collect();
BucketWriter {
num_primary,
total_columns,
fixed_widths,
null_bitmaps: vec![vec![0u8; 128]; total_columns],
- value_buffers: vec![Vec::with_capacity(1024); total_columns],
+ value_buffers: vec![Vec::new(); total_columns],
+ integer_buffers,
non_null_counts: vec![0; total_columns],
const_tracking: vec![true; total_columns],
first_value_len: vec![0; total_columns],
@@ -325,6 +344,15 @@ impl BucketWriter {
return Ok(0);
}
+ if self.integer_buffers[col].is_some() {
+ return match &typed {
+ TypedArrayRef::Int16(a) => self.append_integer_column(col, a,
start_row),
+ TypedArrayRef::Int32(a) => self.append_integer_column(col, a,
start_row),
+ TypedArrayRef::Int64(a) => self.append_integer_column(col, a,
start_row),
+ _ => unreachable!("integer buffer requires an integer array"),
+ };
+ }
+
if null_count == 0 {
if let Some(col_size) =
self.append_no_null_batch(col, &typed, start_row,
num_new_rows)?
@@ -360,6 +388,100 @@ impl BucketWriter {
Ok(col_size)
}
+ fn append_integer_column<T: arrow_array::types::ArrowPrimitiveType>(
+ &mut self,
+ col: usize,
+ array: &PrimitiveArray<T>,
+ start_row: usize,
+ ) -> io::Result<usize>
+ where
+ T::Native: Into<i64>,
+ {
+ if array.null_count() == 0 {
+ return Ok(self.append_integer_values(col, array.values()));
+ }
+ let nulls = array.nulls().expect("nullable integer batch");
+ append_null_bitmap(&mut self.null_bitmaps[col], start_row, nulls);
+ let mut scratch = [0i64; 4096];
+ let mut used = 0;
+ let mut size = 0;
+ for row in nulls.valid_indices() {
+ scratch[used] = array.value(row).into();
+ used += 1;
+ if used == scratch.len() {
+ size += self.append_integer_values(col, &scratch);
+ used = 0;
+ }
+ }
+ size += self.append_integer_values(col, &scratch[..used]);
+ Ok(size)
+ }
+
+ fn append_integer_values<T: Copy + Into<i64>>(&mut self, col: usize,
values: &[T]) -> usize {
+ if values.is_empty() {
+ return 0;
+ }
+ let width = self.fixed_widths[col] as usize;
+ let mask = u64::MAX >> (64 - width * 8);
+ let mut dict_start = 0;
+ if self.const_tracking[col] {
+ let first = self.integer_buffers[col]
+ .as_ref()
+ .unwrap()
+ .keys()
+ .next()
+ .unwrap_or(values[0].into() as u64 & mask);
+ self.first_value_len[col] = width;
+ if let Some(first_different) = values
+ .iter()
+ .position(|&v| (v.into() as u64 & mask) != first)
+ {
+ // Preserve lazy dictionary activation: seed the constant
prefix
+ // once, then track values starting at the first different key.
+ debug_assert_eq!(self.dict_tracking[col],
DictTracking::Pending);
+ self.const_tracking[col] = false;
+ let dict =
self.long_dict_maps[col].get_or_insert_with(HashMap::new);
+ dict.clear();
+ dict.insert(first, 0);
+ self.dict_tracking[col] = DictTracking::Active;
+ dict_start = first_different;
+ } else {
+ dict_start = values.len();
+ }
+ }
+ if self.dict_tracking[col] == DictTracking::Active {
+ let dict = self.long_dict_maps[col].as_mut().unwrap();
+ if dict.len() <= self.max_dict_entries {
+ for &value in &values[dict_start..] {
+ let key = value.into() as u64 & mask;
+ let next = dict.len();
+ dict.entry(key).or_insert(next);
+ if dict.len() > self.max_dict_entries {
+ break;
+ }
+ }
+ }
+ if dict.len() > self.max_dict_entries {
+ self.long_dict_maps[col] = None;
+ self.dict_tracking[col] = DictTracking::Disabled;
+ }
+ }
+ self.integer_buffers[col]
+ .as_mut()
+ .unwrap()
+ .extend_from_slice(values);
+ self.non_null_counts[col] += values.len();
+ // Row-group boundaries and dictionary decisions use the logical width,
+ // never the compact allocation size.
+ values.len() * width
+ }
+
+ fn logical_value_size(&self, col: usize) -> usize {
+ self.integer_buffers[col]
+ .as_ref()
+ .map_or(self.value_buffers[col].len(), |b| b.logical_len())
+ }
+
fn append_no_null_batch(
&mut self,
col: usize,
@@ -617,7 +739,7 @@ impl BucketWriter {
let dict_size = self.get_dict_size(i);
if dict_size >= 2
&& dict_size <= self.max_dict_entries
- && self.dict_encoded_size(i) < self.value_buffers[i].len()
+ && self.dict_encoded_size(i) < self.logical_value_size(i)
{
encodings[i] = ENCODING_DICT;
} else {
@@ -629,104 +751,142 @@ impl BucketWriter {
(encodings, has_nulls)
}
- #[allow(clippy::needless_range_loop)]
- pub fn finish(&self) -> Vec<u8> {
+ /// Encode into a caller-owned sink without materializing the whole bucket.
+ pub(crate) fn write_encoded(&self, out: &mut impl Write) -> io::Result<()>
{
if self.num_rows == 0 {
- return Vec::new();
+ return Ok(());
}
-
- let (encodings, has_nulls) = self.compute_encodings();
-
- let mut out = Vec::new();
-
- // Header: only written when ARRAY columns exist (backward compatible
with v1)
+ let plan = self.encoding_plan();
if !self.children.is_empty() {
- varint::encode(&mut out, self.num_primary as u32);
- varint::encode(&mut out, self.children.len() as u32);
+ write_varint(out, self.num_primary as u32)?;
+ write_varint(out, self.children.len() as u32)?;
for child in &self.children {
- varint::encode(&mut out, child.num_elements as u32);
+ write_varint(out, child.num_elements as u32)?;
}
}
-
- // Encoding flags: 2 bits per column
- let encoding_flags_bytes = (self.total_columns * 2).div_ceil(8);
- let ef_start = out.len();
- out.resize(ef_start + encoding_flags_bytes, 0);
- for i in 0..self.total_columns {
- let byte_idx = (i * 2) / 8;
- let bit_idx = (i * 2) % 8;
- out[ef_start + byte_idx] |= encodings[i] << bit_idx;
+ let mut flags = vec![0; (self.total_columns * 2).div_ceil(8)];
+ for (i, &encoding) in plan.encodings.iter().enumerate() {
+ flags[i / 4] |= encoding << ((i % 4) * 2);
}
-
- // Has-nulls flags: 1 bit per column
- let has_nulls_bytes = self.total_columns.div_ceil(8);
- let hn_start = out.len();
- out.resize(hn_start + has_nulls_bytes, 0);
- for i in 0..self.total_columns {
- if has_nulls[i] {
- out[hn_start + i / 8] |= 1 << (i % 8);
+ out.write_all(&flags)?;
+ flags.clear();
+ flags.resize(self.total_columns.div_ceil(8), 0);
+ for (i, &has_nulls) in plan.has_nulls.iter().enumerate() {
+ if has_nulls {
+ flags[i / 8] |= 1 << (i % 8);
}
}
-
- // CONST metadata
- for i in 0..self.total_columns {
- if encodings[i] == ENCODING_CONST {
- let len = self.first_value_len[i];
- out.extend_from_slice(&self.value_buffers[i][..len]);
+ out.write_all(&flags)?;
+ for (col, &encoding) in plan.encodings.iter().enumerate() {
+ if encoding == ENCODING_CONST {
+ self.write_const(col, out)?;
}
}
-
- // Dict metadata
- for i in 0..self.total_columns {
- if encodings[i] == ENCODING_DICT {
- if let Some(ref dict) = self.long_dict_maps[i] {
- let num_entries = dict.len();
- varint::encode(&mut out, num_entries as u32);
- let w = self.fixed_widths[i];
- let mut keys = vec![0u64; num_entries];
- for (&key, &idx) in dict {
- keys[idx] = key;
- }
- for key in &keys {
- write_fixed_key_to_vec(&mut out, *key, w);
- }
- } else if let Some(ref dict) = self.byte_dict_maps[i] {
- let num_entries = dict.len();
- varint::encode(&mut out, num_entries as u32);
- let mut keys: Vec<(&Vec<u8>, &usize)> =
dict.iter().collect();
- keys.sort_by_key(|&(_, idx)| *idx);
- for (key, _) in keys {
- out.extend_from_slice(key);
- }
- }
+ for (col, &encoding) in plan.encodings.iter().enumerate() {
+ if encoding == ENCODING_DICT {
+ self.write_dictionary(col, out)?;
}
}
-
- // Null bitmaps (per-column row count)
- for i in 0..self.total_columns {
- if has_nulls[i] && encodings[i] != ENCODING_ALL_NULL {
- let nbytes = self.col_num_rows(i).div_ceil(8);
- out.extend_from_slice(&self.null_bitmaps[i][..nbytes]);
+ for (col, &has_nulls) in plan.has_nulls.iter().enumerate() {
+ if has_nulls && plan.encodings[col] != ENCODING_ALL_NULL {
+ self.write_nulls(col, out)?;
}
}
+ for (col, &encoding) in plan.encodings.iter().enumerate() {
+ self.write_values(col, encoding, out)?;
+ }
+ Ok(())
+ }
- // Column data
- for i in 0..self.total_columns {
- if encodings[i] == ENCODING_PLAIN {
- out.extend_from_slice(&self.value_buffers[i]);
- } else if encodings[i] == ENCODING_DICT {
- let bw = bit_width(self.get_dict_size(i));
- let packed_bytes = (self.non_null_counts[i] * bw).div_ceil(8);
- let data_start = out.len();
- out.resize(data_start + packed_bytes, 0);
- self.write_dict_bit_packed(i, &mut out, data_start);
+ pub fn finish(&self) -> Vec<u8> {
+ let mut out = Vec::new();
+ self.write_encoded(&mut out)
+ .expect("writing into a Vec cannot fail");
+ out
+ }
+
+ pub(crate) fn has_integer_buffers(&self) -> bool {
+ self.integer_buffers.iter().any(Option::is_some)
+ }
+
+ #[cfg(test)]
+ pub(crate) fn use_wide_integer_buffers(&mut self) {
+ assert_eq!(self.num_rows, 0);
+ for buffer in &mut self.integer_buffers {
+ *buffer = None;
+ }
+ }
+
+ pub(crate) fn encoding_plan(&self) -> BucketEncodingPlan {
+ let (encodings, has_nulls) = self.compute_encodings();
+ let page_sizes = (0..self.total_columns)
+ .map(|col| {
+ let nulls = if has_nulls[col] {
+ self.col_num_rows(col).div_ceil(8)
+ } else {
+ 0
+ };
+ match encodings[col] {
+ ENCODING_ALL_NULL => None,
+ ENCODING_CONST => has_nulls[col].then_some(nulls),
+ ENCODING_DICT => Some(nulls + self.dict_encoded_size(col)),
+ ENCODING_PLAIN => Some(nulls +
self.logical_value_size(col)),
+ _ => unreachable!(),
+ }
+ })
+ .collect();
+ BucketEncodingPlan {
+ encodings,
+ has_nulls,
+ page_sizes,
+ }
+ }
+
+ pub(crate) fn page_size(&self, plan: &BucketEncodingPlan, col: usize) ->
usize {
+ if plan.encodings[col] == ENCODING_ALL_NULL {
+ return 0;
+ }
+ 2 + plan.page_sizes[col].unwrap_or(0)
+ + if plan.encodings[col] == ENCODING_CONST {
+ self.first_value_len[col]
+ } else {
+ 0
}
+ }
+
+ pub(crate) fn write_page(
+ &self,
+ plan: &BucketEncodingPlan,
+ col: usize,
+ out: &mut impl Write,
+ ) -> io::Result<()> {
+ let encoding = plan.encodings[col];
+ if encoding == ENCODING_ALL_NULL {
+ return Ok(());
}
+ out.write_all(&[encoding, u8::from(plan.has_nulls[col])])?;
+ if encoding == ENCODING_CONST {
+ self.write_const(col, out)?;
+ }
+ self.write_page_data(plan, col, out)
+ }
- out
+ fn write_page_data(
+ &self,
+ plan: &BucketEncodingPlan,
+ col: usize,
+ out: &mut impl Write,
+ ) -> io::Result<()> {
+ let encoding = plan.encodings[col];
+ if encoding == ENCODING_DICT {
+ self.write_dictionary(col, out)?;
+ }
+ if plan.has_nulls[col] {
+ self.write_nulls(col, out)?;
+ }
+ self.write_values(col, encoding, out)
}
- #[allow(clippy::needless_range_loop)]
pub fn finish_paged(&self) -> PagedBucketOutput {
if self.num_rows == 0 {
return PagedBucketOutput {
@@ -738,73 +898,24 @@ impl BucketWriter {
children: self.children.clone(),
};
}
-
- let (encodings, has_nulls) = self.compute_encodings();
-
+ let plan = self.encoding_plan();
let mut const_data = vec![Vec::new(); self.total_columns];
- let mut column_pages: Vec<Option<Vec<u8>>> = vec![None;
self.total_columns];
-
- for i in 0..self.total_columns {
- let col_rows = self.col_num_rows(i);
- let null_bitmap_bytes = col_rows.div_ceil(8);
-
- match encodings[i] {
- ENCODING_ALL_NULL => {}
- ENCODING_CONST => {
- let len = self.first_value_len[i];
- const_data[i] = self.value_buffers[i][..len].to_vec();
- if has_nulls[i] {
- column_pages[i] =
Some(self.null_bitmaps[i][..null_bitmap_bytes].to_vec());
- }
- }
- ENCODING_DICT => {
- let mut page = Vec::new();
- if let Some(ref dict) = self.long_dict_maps[i] {
- let num_entries = dict.len();
- varint::encode(&mut page, num_entries as u32);
- let w = self.fixed_widths[i];
- let mut keys = vec![0u64; num_entries];
- for (&key, &idx) in dict {
- keys[idx] = key;
- }
- for key in &keys {
- write_fixed_key_to_vec(&mut page, *key, w);
- }
- } else if let Some(ref dict) = self.byte_dict_maps[i] {
- let num_entries = dict.len();
- varint::encode(&mut page, num_entries as u32);
- let mut keys: Vec<(&Vec<u8>, &usize)> =
dict.iter().collect();
- keys.sort_by_key(|&(_, idx)| *idx);
- for (key, _) in keys {
- page.extend_from_slice(key);
- }
- }
- if has_nulls[i] {
-
page.extend_from_slice(&self.null_bitmaps[i][..null_bitmap_bytes]);
- }
- let dict_size = self.get_dict_size(i);
- let bw = bit_width(dict_size);
- let packed_bytes = (self.non_null_counts[i] *
bw).div_ceil(8);
- let data_start = page.len();
- page.resize(data_start + packed_bytes, 0);
- self.write_dict_bit_packed(i, &mut page, data_start);
- column_pages[i] = Some(page);
- }
- ENCODING_PLAIN => {
- let mut page = Vec::new();
- if has_nulls[i] {
-
page.extend_from_slice(&self.null_bitmaps[i][..null_bitmap_bytes]);
- }
- page.extend_from_slice(&self.value_buffers[i]);
- column_pages[i] = Some(page);
- }
- _ => {}
+ let mut column_pages = vec![None; self.total_columns];
+ for col in 0..self.total_columns {
+ if plan.encodings[col] == ENCODING_CONST {
+ self.write_const(col, &mut const_data[col])
+ .expect("Vec write");
+ }
+ if let Some(size) = plan.page_sizes[col] {
+ let mut page = Vec::with_capacity(size);
+ self.write_page_data(&plan, col, &mut page)
+ .expect("Vec write");
+ column_pages[col] = Some(page);
}
}
-
PagedBucketOutput {
- encodings,
- has_nulls,
+ encodings: plan.encodings,
+ has_nulls: plan.has_nulls,
const_data,
column_pages,
num_primary: self.num_primary,
@@ -812,10 +923,110 @@ impl BucketWriter {
}
}
+ fn write_const(&self, col: usize, out: &mut impl Write) -> io::Result<()> {
+ if let Some(buffer) = &self.integer_buffers[col] {
+ let key = buffer.keys().next().expect("non-null constant");
+ out.write_all(&key.to_be_bytes()[8 - self.fixed_widths[col] as
usize..])
+ } else {
+
out.write_all(&self.value_buffers[col][..self.first_value_len[col]])
+ }
+ }
+
+ fn write_dictionary(&self, col: usize, out: &mut impl Write) ->
io::Result<()> {
+ if let Some(dict) = &self.long_dict_maps[col] {
+ write_varint(out, dict.len() as u32)?;
+ let mut keys = vec![0; dict.len()];
+ for (&key, &index) in dict {
+ keys[index] = key;
+ }
+ for key in keys {
+ out.write_all(&key.to_be_bytes()[8 - self.fixed_widths[col] as
usize..])?;
+ }
+ } else if let Some(dict) = &self.byte_dict_maps[col] {
+ write_varint(out, dict.len() as u32)?;
+ let mut keys: Vec<_> = dict.iter().collect();
+ keys.sort_unstable_by_key(|(_, index)| **index);
+ for (key, _) in keys {
+ out.write_all(key)?;
+ }
+ }
+ Ok(())
+ }
+
+ fn write_nulls(&self, col: usize, out: &mut impl Write) -> io::Result<()> {
+
out.write_all(&self.null_bitmaps[col][..self.col_num_rows(col).div_ceil(8)])
+ }
+
+ fn write_values(&self, col: usize, encoding: u8, out: &mut impl Write) ->
io::Result<()> {
+ match encoding {
+ ENCODING_PLAIN => match &self.integer_buffers[col] {
+ Some(buffer) => buffer.write_to(out),
+ None => out.write_all(&self.value_buffers[col]),
+ },
+ ENCODING_DICT => self.write_dict_indices(col, out),
+ _ => Ok(()),
+ }
+ }
+
+ fn write_dict_indices(&self, col: usize, out: &mut impl Write) ->
io::Result<()> {
+ let width = self.fixed_widths[col];
+ let bits = bit_width(self.get_dict_size(col));
+ let mut integer_keys = self.integer_buffers[col].as_ref().map(|b|
b.keys());
+ let mut pos = 0;
+ let mut scratch = [0; 16 * 1024];
+ let mut used = 0;
+ let mut pending = 0u128;
+ let mut pending_bits = 0;
+ // Value buffers contain only non-null values, in row order.
+ for _ in 0..self.non_null_counts[col] {
+ let index = if let Some(dict) = &self.long_dict_maps[col] {
+ let key = if let Some(keys) = &mut integer_keys {
+ keys.next().expect("non-null integer key")
+ } else {
+ let key =
values::extract_fixed_key(&self.value_buffers[col], pos, width);
+ pos += width as usize;
+ key
+ };
+ dict[&key]
+ } else {
+ let len = if width > 0 {
+ width as usize
+ } else {
+ let len = varint::decode(&self.value_buffers[col], &mut
pos.clone())
+ .expect("internal varint in value buffer");
+ varint::encoded_size(len) + len as usize
+ };
+ let key = &self.value_buffers[col][pos..pos + len];
+ pos += len;
+ self.byte_dict_maps[col].as_ref().unwrap()[key]
+ };
+ pending |= (index as u128) << pending_bits;
+ pending_bits += bits;
+ while pending_bits >= 8 {
+ scratch[used] = pending as u8;
+ used += 1;
+ pending >>= 8;
+ pending_bits -= 8;
+ if used == scratch.len() {
+ out.write_all(&scratch)?;
+ used = 0;
+ }
+ }
+ }
+ if pending_bits > 0 {
+ scratch[used] = pending as u8;
+ used += 1;
+ }
+ out.write_all(&scratch[..used])
+ }
+
pub fn reset(&mut self) {
for i in 0..self.total_columns {
self.null_bitmaps[i].fill(0);
self.value_buffers[i].clear();
+ if let Some(buffer) = &mut self.integer_buffers[i] {
+ buffer.clear();
+ }
self.non_null_counts[i] = 0;
self.const_tracking[i] = true;
self.first_value_len[i] = 0;
@@ -848,42 +1059,6 @@ impl BucketWriter {
0
}
- fn write_dict_bit_packed(&self, col: usize, buf: &mut [u8], data_start:
usize) -> usize {
- let dict_size = self.get_dict_size(col);
- let bw = bit_width(dict_size);
- let w = self.fixed_widths[col];
- let mut bit_offset = 0usize;
- let mut val_pos = 0usize;
-
- for r in 0..self.col_num_rows(col) {
- let is_null = (self.null_bitmaps[col][r / 8] & (1 << (r % 8))) !=
0;
- if !is_null {
- let idx = if let Some(ref dict) = self.long_dict_maps[col] {
- let key =
values::extract_fixed_key(&self.value_buffers[col], val_pos, w);
- val_pos += w as usize;
- *dict.get(&key).unwrap()
- } else if let Some(ref dict) = self.byte_dict_maps[col] {
- let value_len = if w > 0 {
- w as usize
- } else {
- let var_len =
- varint::decode(&self.value_buffers[col], &mut
val_pos.clone())
- .expect("internal varint in value buffer");
- varint::encoded_size(var_len) + var_len as usize
- };
- let slice = &self.value_buffers[col][val_pos..val_pos +
value_len];
- val_pos += value_len;
- *dict.get(slice).unwrap()
- } else {
- unreachable!()
- };
- write_bit_packed(buf, data_start, bit_offset, idx, bw);
- bit_offset += bw;
- }
- }
- bit_offset
- }
-
fn dict_encoded_size(&self, col: usize) -> usize {
if self.dict_tracking[col] != DictTracking::Active {
return usize::MAX;
@@ -937,7 +1112,7 @@ impl BucketWriter {
}
}
ENCODING_PLAIN => {
- size += self.value_buffers[i].len();
+ size += self.logical_value_size(i);
}
_ => {}
}
@@ -1662,33 +1837,20 @@ fn bit_width(num_entries: usize) -> usize {
usize::BITS as usize - (num_entries - 1).leading_zeros() as usize
}
-fn write_bit_packed(buf: &mut [u8], byte_base: usize, bit_offset: usize,
value: usize, bw: usize) {
- if bw == 0 {
- return;
- }
- let start_byte = byte_base + bit_offset / 8;
- let bit_shift = bit_offset % 8;
- let mut bits = (value as u64) << bit_shift;
- let total_bits = bit_shift + bw;
- let num_bytes = total_bits.div_ceil(8);
- for i in 0..num_bytes {
- buf[start_byte + i] |= (bits & 0xFF) as u8;
- bits >>= 8;
- }
-}
-
fn equals_first_value(buf: &[u8], offset: usize, len: usize) -> bool {
buf[..len] == buf[offset..offset + len]
}
-fn write_fixed_key_to_vec(buf: &mut Vec<u8>, key: u64, width: i32) {
- match width {
- 1 => buf.push(key as u8),
- 2 => buf.extend_from_slice(&(key as u16).to_be_bytes()),
- 4 => buf.extend_from_slice(&(key as u32).to_be_bytes()),
- 8 => buf.extend_from_slice(&key.to_be_bytes()),
- _ => {}
+fn write_varint(out: &mut impl Write, mut value: u32) -> io::Result<()> {
+ let mut bytes = [0; 5];
+ let mut len = 0;
+ while value >= 128 {
+ bytes[len] = (value as u8) | 0x80;
+ value >>= 7;
+ len += 1;
}
+ bytes[len] = value as u8;
+ out.write_all(&bytes[..len + 1])
}
#[cfg(test)]
@@ -1820,6 +1982,110 @@ mod tests {
assert_eq!(data[h] & 0x03, ENCODING_CONST);
}
+ #[test]
+ fn test_all_null_wide_columns_do_not_preallocate_value_payloads() {
+ let types = vec![&DataType::Float64; 10_000];
+ let mut writer = BucketWriter::new(&types, 32768, 255);
+ let nulls = Float64Array::new_null(65);
+ let arrays = vec![&nulls as &dyn Array; types.len()];
+ for _ in 0..2 {
+ writer.write_columns(&arrays, &types).unwrap();
+ let allocated: usize =
writer.value_buffers.iter().map(Vec::capacity).sum();
+ assert!(allocated <= 1024, "all-null payload capacity:
{allocated}");
+ writer.reset();
+ }
+ }
+
+ fn compact_integer_test_array(dt: &DataType, values: &[Option<i64>]) ->
ArrayRef {
+ match dt {
+ DataType::Int16 => Arc::new(Int16Array::from(
+ values
+ .iter()
+ .map(|v| v.map(|v| v as i16))
+ .collect::<Vec<_>>(),
+ )),
+ DataType::Int32 => Arc::new(Int32Array::from(
+ values
+ .iter()
+ .map(|v| v.map(|v| v as i32))
+ .collect::<Vec<_>>(),
+ )),
+ DataType::Int64 => Arc::new(Int64Array::from(values.to_vec())),
+ _ => unreachable!(),
+ }
+ }
+
+ #[test]
+ fn test_compact_integer_lazy_dictionary_after_nulls_promotion_and_reset() {
+ for dt in [DataType::Int16, DataType::Int32, DataType::Int64] {
+ let mut compact = BucketWriter::new(&[&dt], 32768, 2);
+ let mut wide = BucketWriter::new(&[&dt], 32768, 2);
+ wide.use_wide_integer_buffers();
+ let append =
+ |compact: &mut BucketWriter, wide: &mut BucketWriter, values:
&[Option<i64>]| {
+ // Slice the source and append at non-byte-aligned
destination offsets.
+ let mut padded = vec![Some(99)];
+ padded.extend_from_slice(values);
+ let array = compact_integer_test_array(&dt,
&padded).slice(1, values.len());
+ assert_eq!(
+ compact.write_columns(&[array.as_ref()],
&[&dt]).unwrap(),
+ wide.write_columns(&[array.as_ref()], &[&dt]).unwrap()
+ );
+ assert_eq!(compact.estimated_raw_size(),
wide.estimated_raw_size());
+ assert_eq!(compact.finish(), wide.finish());
+ assert_eq!(
+ compact.finish_paged().column_pages,
+ wide.finish_paged().column_pages
+ );
+ };
+ append(&mut compact, &mut wide, &[None; 65]);
+ append(&mut compact, &mut wide, &vec![Some(255); 4095]);
+ assert_eq!(compact.dict_tracking[0], DictTracking::Pending);
+ assert!(compact.long_dict_maps[0].is_none());
+ append(
+ &mut compact,
+ &mut wide,
+ &[None, Some(255), Some(-1), None, Some(-1)],
+ );
+ assert_eq!(compact.dict_tracking[0], DictTracking::Active);
+ assert_eq!(compact.long_dict_maps[0].as_ref().unwrap().len(), 2);
+ assert_eq!(compact.compute_encodings().0, vec![ENCODING_DICT]);
+
+ compact.reset();
+ wide.reset();
+ append(&mut compact, &mut wide, &vec![Some(7); 100]);
+ assert_eq!(compact.dict_tracking[0], DictTracking::Pending);
+ assert!(compact.long_dict_maps[0].as_ref().unwrap().is_empty());
+ append(&mut compact, &mut wide, &[Some(8), None, Some(7)]);
+ let dict = compact.long_dict_maps[0].as_ref().unwrap();
+ assert_eq!(dict.len(), 2);
+ assert_eq!(dict[&7], 0);
+ assert_eq!(dict[&8], 1);
+ }
+ }
+
+ #[test]
+ fn test_compact_integer_dictionary_limit_stays_disabled_until_reset() {
+ for limit in [0, 1, 2] {
+ let dt = DataType::Int64;
+ let mut compact = BucketWriter::new(&[&dt], 32768, limit);
+ let prefix = Int64Array::from(vec![-1; 100]);
+ compact.write_columns(&[&prefix], &[&dt]).unwrap();
+ assert!(compact.long_dict_maps[0].is_none());
+ let varied = Int64Array::from(vec![Some(255), None, Some(256),
Some(-1)]);
+ compact.write_columns(&[&varied], &[&dt]).unwrap();
+ assert_eq!(compact.dict_tracking[0], DictTracking::Disabled);
+ assert!(compact.long_dict_maps[0].is_none());
+ compact.write_columns(&[&prefix], &[&dt]).unwrap();
+ assert_eq!(compact.dict_tracking[0], DictTracking::Disabled);
+ assert_eq!(compact.compute_encodings().0, vec![ENCODING_PLAIN]);
+ compact.reset();
+ compact.write_columns(&[&prefix], &[&dt]).unwrap();
+ assert_eq!(compact.dict_tracking[0], DictTracking::Pending);
+ assert_eq!(compact.compute_encodings().0, vec![ENCODING_CONST]);
+ }
+ }
+
#[test]
fn test_const_columns_do_not_build_dictionaries() {
let types = [DataType::Int32, DataType::Utf8];
diff --git a/core/src/integer_buffer.rs b/core/src/integer_buffer.rs
new file mode 100644
index 0000000..795d40f
--- /dev/null
+++ b/core/src/integer_buffer.rs
@@ -0,0 +1,287 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::io::{self, Write};
+
+// A late outlier may widen this block, never the entire buffered column.
+const BLOCK_VALUES: usize = 4096;
+const OUTPUT_BYTES: usize = 16 * 1024;
+
+#[derive(Default)]
+struct Block {
+ bytes: Vec<u8>,
+ width: usize,
+}
+
+/// In-memory representation only. Keys retain the original integer's
big-endian
+/// bit pattern; negative values use the logical width, positive values may use
+/// 1/2/4 bytes. The file schema, dictionary keys and encoded bytes never
change.
+pub(crate) struct IntegerBuffer {
+ blocks: Vec<Block>,
+ logical_width: usize,
+ len: usize,
+}
+
+impl IntegerBuffer {
+ pub(crate) fn new(logical_width: usize) -> Self {
+ debug_assert!(matches!(logical_width, 2 | 4 | 8));
+ Self {
+ blocks: Vec::new(),
+ logical_width,
+ len: 0,
+ }
+ }
+
+ #[cfg(test)]
+ fn push(&mut self, value: i64) -> u64 {
+ self.extend_from_slice(&[value]);
+ (value as u64) & (u64::MAX >> (64 - self.logical_width * 8))
+ }
+
+ pub(crate) fn extend_from_slice<T: Copy + Into<i64>>(&mut self, mut
values: &[T]) {
+ let mask = u64::MAX >> (64 - self.logical_width * 8);
+ while !values.is_empty() {
+ let count = values.len().min(BLOCK_VALUES - self.len %
BLOCK_VALUES);
+ let input = &values[..count];
+ // A block-sized range check is vectorizable, unlike a width branch
+ // on every value in the copy loop. Negative keys require full
width.
+ let bits = input.iter().fold(0u64, |bits, &v| bits | v.into() as
u64) & mask;
+ let width = if bits <= u8::MAX as u64 {
+ 1
+ } else if bits <= u16::MAX as u64 {
+ 2
+ } else if bits <= u32::MAX as u64 {
+ 4
+ } else {
+ 8
+ };
+ if self.len.is_multiple_of(BLOCK_VALUES) {
+ self.blocks.push(Block {
+ bytes: Vec::new(),
+ width,
+ });
+ }
+ let block = self.blocks.last_mut().unwrap();
+ if width > block.width {
+ let old_width = block.width;
+ let old_count = block.bytes.len() / old_width;
+ block.bytes.resize(old_count * width, 0);
+ for i in (0..old_count).rev() {
+ let key = read_key(&block.bytes[i * old_width..(i + 1) *
old_width]);
+ block.bytes[i * width..(i + 1) * width]
+ .copy_from_slice(&key.to_be_bytes()[8 - width..]);
+ }
+ block.width = width;
+ }
+ let offset = block.bytes.len();
+ block.bytes.resize(offset + count * block.width, 0);
+ let output = &mut block.bytes[offset..];
+ match block.width {
+ 1 => {
+ for (out, &v) in output.iter_mut().zip(input) {
+ *out = v.into() as u8;
+ }
+ }
+ 2 => {
+ for (out, &v) in output.chunks_exact_mut(2).zip(input) {
+ out.copy_from_slice(&(v.into() as u16).to_be_bytes());
+ }
+ }
+ 4 => {
+ for (out, &v) in output.chunks_exact_mut(4).zip(input) {
+ out.copy_from_slice(&(v.into() as u32).to_be_bytes());
+ }
+ }
+ 8 => {
+ for (out, &v) in output.chunks_exact_mut(8).zip(input) {
+ out.copy_from_slice(&v.into().to_be_bytes());
+ }
+ }
+ _ => unreachable!(),
+ }
+ self.len += count;
+ values = &values[count..];
+ }
+ }
+
+ pub(crate) fn keys(&self) -> impl Iterator<Item = u64> + '_ {
+ self.blocks
+ .iter()
+ .flat_map(|block|
block.bytes.chunks_exact(block.width).map(read_key))
+ }
+
+ pub(crate) fn logical_len(&self) -> usize {
+ self.len * self.logical_width
+ }
+
+ pub(crate) fn write_to(&self, out: &mut impl Write) -> io::Result<()> {
+ let mut scratch = [0; OUTPUT_BYTES];
+ for block in &self.blocks {
+ match (block.width, self.logical_width) {
+ (src, dst) if src == dst => {
+ for chunk in block.bytes.chunks(OUTPUT_BYTES) {
+ out.write_all(chunk)?;
+ }
+ }
+ (1, 2) => widen::<1, 2>(&block.bytes, &mut scratch, out)?,
+ (1, 4) => widen::<1, 4>(&block.bytes, &mut scratch, out)?,
+ (1, 8) => widen::<1, 8>(&block.bytes, &mut scratch, out)?,
+ (2, 4) => widen::<2, 4>(&block.bytes, &mut scratch, out)?,
+ (2, 8) => widen::<2, 8>(&block.bytes, &mut scratch, out)?,
+ (4, 8) => widen::<4, 8>(&block.bytes, &mut scratch, out)?,
+ _ => unreachable!(),
+ }
+ }
+ Ok(())
+ }
+
+ pub(crate) fn clear(&mut self) {
+ // Do not retain a wide allocation after an outlier or a large row
group.
+ self.blocks.clear();
+ self.len = 0;
+ }
+
+ #[cfg(test)]
+ pub(crate) fn allocated_bytes(&self) -> usize {
+ self.blocks.capacity() * std::mem::size_of::<Block>()
+ + self
+ .blocks
+ .iter()
+ .map(|b| b.bytes.capacity())
+ .sum::<usize>()
+ }
+}
+
+fn read_key(bytes: &[u8]) -> u64 {
+ match bytes.len() {
+ 1 => bytes[0] as u64,
+ 2 => u16::from_be_bytes(bytes.try_into().unwrap()) as u64,
+ 4 => u32::from_be_bytes(bytes.try_into().unwrap()) as u64,
+ 8 => u64::from_be_bytes(bytes.try_into().unwrap()),
+ _ => unreachable!(),
+ }
+}
+
+fn widen<const SRC: usize, const DST: usize>(
+ bytes: &[u8],
+ scratch: &mut [u8; OUTPUT_BYTES],
+ out: &mut impl Write,
+) -> io::Result<()> {
+ for input in bytes.chunks(OUTPUT_BYTES / DST * SRC) {
+ let output = &mut scratch[..input.len() / SRC * DST];
+ for (src, dst) in
input.chunks_exact(SRC).zip(output.chunks_exact_mut(DST)) {
+ // Clear the prefix: the previous block may have used a wider
+ // compact representation in the same reusable scratch space.
+ dst[..DST - SRC].fill(0);
+ dst[DST - SRC..].copy_from_slice(src);
+ }
+ out.write_all(output)?;
+ }
+ Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn unsigned_ranges_negative_values_and_late_outliers() {
+ for width in [2, 4, 8] {
+ let max = i64::MAX >> (64 - width * 8);
+ let min = -max - 1;
+ let mut buffer = IntegerBuffer::new(width);
+ let mut expected = Vec::new();
+ for value in (0..BLOCK_VALUES * 3)
+ .map(|i| (i % 256) as i64)
+ .chain([max, min, -1, 0, 127, 128, 255, 256])
+ {
+ buffer.push(value);
+ expected.extend_from_slice(&value.to_be_bytes()[8 - width..]);
+ }
+ assert_eq!(buffer.blocks[0].width, 1);
+ assert_eq!(buffer.blocks[1].width, 1);
+ assert_eq!(buffer.blocks.last().unwrap().width, width);
+ let mut actual = Vec::new();
+ buffer.write_to(&mut actual).unwrap();
+ assert_eq!(actual, expected);
+ assert_eq!(buffer.logical_len(), expected.len());
+ }
+ }
+
+ #[test]
+ fn promotions_within_a_block_preserve_previous_values() {
+ let values = [0, 255, 256, 65535, 65536, u32::MAX as i64, i64::MAX,
-1];
+ let mut buffer = IntegerBuffer::new(8);
+ for value in values {
+ buffer.push(value);
+ }
+ assert_eq!(buffer.keys().collect::<Vec<_>>(), values.map(|v| v as
u64));
+ buffer.clear();
+ buffer.push(255);
+ assert_eq!(buffer.blocks[0].width, 1);
+ assert_eq!(buffer.keys().collect::<Vec<_>>(), vec![255]);
+ }
+
+ #[test]
+ fn uint16_values_use_half_the_int32_payload_plus_bounded_metadata() {
+ let mut buffer = IntegerBuffer::new(4);
+ for i in 0..1_000_000 {
+ buffer.push((i % 65536) as i64);
+ }
+ assert_eq!(buffer.logical_len(), 4_000_000);
+ assert!(buffer.allocated_bytes() < 2_100_000);
+ }
+
+ #[test]
+ fn widening_output_is_bounded_and_propagates_sink_errors() {
+ struct Sink {
+ bytes: usize,
+ fail_after: usize,
+ }
+ impl Write for Sink {
+ fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
+ assert!(bytes.len() <= OUTPUT_BYTES);
+ if self.bytes >= self.fail_after {
+ return Err(io::Error::other("sentinel"));
+ }
+ self.bytes += bytes.len();
+ Ok(bytes.len())
+ }
+ fn flush(&mut self) -> io::Result<()> {
+ Ok(())
+ }
+ }
+ let mut buffer = IntegerBuffer::new(8);
+ for i in 0..100_003 {
+ buffer.push(i);
+ }
+ let mut sink = Sink {
+ bytes: 0,
+ fail_after: usize::MAX,
+ };
+ buffer.write_to(&mut sink).unwrap();
+ assert_eq!(sink.bytes, 800_024);
+ let mut sink = Sink {
+ bytes: 0,
+ fail_after: OUTPUT_BYTES,
+ };
+ assert_eq!(
+ buffer.write_to(&mut sink).unwrap_err().to_string(),
+ "sentinel"
+ );
+ }
+}
diff --git a/core/src/lib.rs b/core/src/lib.rs
index 3f28500..a240073 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -18,6 +18,7 @@
pub mod bpe;
pub mod bucket_reader;
pub mod bucket_writer;
+mod integer_buffer;
pub mod reader;
pub mod schema;
pub mod spec;
diff --git a/core/src/writer.rs b/core/src/writer.rs
index d503a06..6ce79f4 100644
--- a/core/src/writer.rs
+++ b/core/src/writer.rs
@@ -440,6 +440,20 @@ impl<S: OutputFile> MosaicWriter<S> {
if bw.is_empty() {
continue;
}
+ if bw.has_integer_buffers() {
+ let offset = self.out.pos();
+ let (layout, uncompressed_size) = write_integer_bucket(
+ &mut self.out,
+ bw,
+ self.compression,
+ self.zstd_level,
+ self.page_size_threshold,
+ )?;
+ bucket_offsets[b] = offset;
+ bucket_layouts[b] = layout;
+ actual_uncompressed_sizes[b] = uncompressed_size;
+ continue;
+ }
let est_size = bw.estimated_raw_size();
let try_paged =
self.compression == COMPRESSION_ZSTD && est_size >=
self.page_size_threshold;
@@ -723,6 +737,198 @@ impl<S: OutputFile> MosaicWriter<S> {
}
}
+/// Adapt the existing output API without adding seek or temporary-file
+/// requirements to C, JNI or custom output callbacks.
+struct OutputWriter<'a, S>(&'a mut S);
+
+fn terminal_output_error(error: io::Error) -> io::Error {
+ // OutputFile may fail after writing a prefix. Prevent Write::write_all and
+ // Zstd from retrying Interrupted errors and duplicating that prefix.
+ if error.kind() == io::ErrorKind::Interrupted {
+ io::Error::other(error)
+ } else {
+ error
+ }
+}
+
+impl<S: OutputFile> io::Write for OutputWriter<'_, S> {
+ fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
+ self.0.write(bytes).map_err(terminal_output_error)?;
+ Ok(bytes.len())
+ }
+
+ fn flush(&mut self) -> io::Result<()> {
+ self.0.flush()
+ }
+}
+
+const INTEGER_OUTPUT_BUFFER_SIZE: usize = 16 * 1024;
+
+// Coalesce flags, dictionary entries and value chunks before calling the sink.
+// Unlike BufWriter, this has no flushing Drop: a failed callback may already
+// have written a prefix, so neither errors nor unwinding may retry its bytes.
+struct BufferedOutputWriter<'a, S> {
+ out: &'a mut S,
+ buffer: Vec<u8>,
+}
+
+impl<'a, S: OutputFile> BufferedOutputWriter<'a, S> {
+ fn new(out: &'a mut S) -> Self {
+ Self {
+ out,
+ buffer: Vec::with_capacity(INTEGER_OUTPUT_BUFFER_SIZE),
+ }
+ }
+
+ fn drain(&mut self) -> io::Result<()> {
+ if !self.buffer.is_empty() {
+ self.out
+ .write(&self.buffer)
+ .map_err(terminal_output_error)?;
+ self.buffer.clear();
+ }
+ Ok(())
+ }
+
+ fn finish(mut self) -> io::Result<()> {
+ // Drain the encoding buffer only; the sink is flushed when the file
closes.
+ self.drain()
+ }
+}
+
+impl<S: OutputFile> io::Write for BufferedOutputWriter<'_, S> {
+ fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
+ if bytes.len() > INTEGER_OUTPUT_BUFFER_SIZE - self.buffer.len() {
+ self.drain()?;
+ }
+ if bytes.len() >= INTEGER_OUTPUT_BUFFER_SIZE {
+ self.out.write(bytes).map_err(terminal_output_error)?;
+ } else {
+ self.buffer.extend_from_slice(bytes);
+ }
+ Ok(bytes.len())
+ }
+
+ fn flush(&mut self) -> io::Result<()> {
+ self.drain()?;
+ self.out.flush()
+ }
+}
+
+// Small encoded units use the existing one-shot compressor. This bounds the
+// materialized input while avoiding streaming context overhead for tiny pages.
+const INTEGER_ENCODING_BUFFER_LIMIT: usize = 1024 * 1024;
+
+fn compress_integer_unit<W: io::Write>(
+ mut output: W,
+ level: i32,
+ size: usize,
+ encode: impl FnOnce(&mut dyn io::Write) -> io::Result<()>,
+) -> io::Result<W> {
+ check_zstd_block_size(size, "integer encoded unit size")?;
+ if size <= INTEGER_ENCODING_BUFFER_LIMIT {
+ let mut raw = Vec::with_capacity(size);
+ encode(&mut raw)?;
+ if raw.len() != size {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidData,
+ "integer encoded size mismatch",
+ ));
+ }
+ let compressed = zstd::bulk::compress(&raw,
level).map_err(io::Error::other)?;
+ output.write_all(&compressed)?;
+ Ok(output)
+ } else {
+ let mut encoder = zstd::stream::write::Encoder::new(output, level)?;
+ encoder.set_pledged_src_size(Some(size as u64))?;
+ encode(&mut encoder)?;
+ encoder.finish()
+ }
+}
+
+fn write_integer_bucket<S: OutputFile>(
+ out: &mut S,
+ bucket: &BucketWriter,
+ compression: u8,
+ level: i32,
+ page_threshold: usize,
+) -> io::Result<(BucketLayout, usize)> {
+ let raw_size = bucket.estimated_raw_size();
+ let plan = bucket.encoding_plan();
+ let page_count = plan.page_sizes.iter().filter(|p| p.is_some()).count();
+ let page_bytes: usize = plan.page_sizes.iter().flatten().sum();
+ let paged = compression == COMPRESSION_ZSTD
+ && raw_size >= page_threshold
+ && page_bytes
+ .checked_div(page_count)
+ .is_some_and(|avg| avg >= page_threshold);
+ let start = out.pos();
+ if !paged {
+ match compression {
+ COMPRESSION_NONE => {
+ let mut buffered = BufferedOutputWriter::new(out);
+ bucket.write_encoded(&mut buffered)?;
+ buffered.finish()?;
+ }
+ COMPRESSION_ZSTD => {
+ check_zstd_block_size(raw_size, "bucket uncompressed size")?;
+ compress_integer_unit(OutputWriter(out), level, raw_size, |mut
sink| {
+ bucket.write_encoded(&mut sink)
+ })?;
+ }
+ _ => {
+ return Err(io::Error::new(
+ io::ErrorKind::InvalidInput,
+ "Unsupported compression",
+ ))
+ }
+ }
+ return Ok((
+ BucketLayout::Monolithic {
+ compressed_size: (out.pos() - start) as usize,
+ uncompressed_size: raw_size,
+ },
+ raw_size,
+ ));
+ }
+
+ // The format puts compressed slot lengths before the slots. Retain only
+ // compressed slots here, not full-width values or uncompressed pages.
+ let mut slots = Vec::with_capacity(plan.encodings.len());
+ for col in 0..plan.encodings.len() {
+ let size = bucket.page_size(&plan, col);
+ let mut slot = Vec::new();
+ if size != 0 {
+ check_zstd_block_size(size, "page uncompressed size")?;
+ varint::encode(&mut slot, to_u32(size, "page uncompressed size")?);
+ slot = compress_integer_unit(slot, level, size, |mut sink| {
+ bucket.write_page(&plan, col, &mut sink)
+ })?;
+ }
+ slots.push(slot);
+ }
+ if !bucket.children().is_empty() {
+ out.write(&(bucket.children().len() as u16).to_le_bytes())?;
+ for child in bucket.children() {
+ out.write(&(child.num_elements as u32).to_le_bytes())?;
+ }
+ }
+ for slot in &slots {
+ out.write(&to_u32(slot.len(), "paged slot size")?.to_le_bytes())?;
+ }
+ for slot in &slots {
+ if !slot.is_empty() {
+ out.write(slot)?;
+ }
+ }
+ Ok((
+ BucketLayout::Paged {
+ total_size: (out.pos() - start) as usize,
+ },
+ page_bytes,
+ ))
+}
+
impl<S: OutputFile> Drop for MosaicWriter<S> {
fn drop(&mut self) {
if self.state == WriterState::Open {
@@ -1320,3 +1526,7 @@ mod tests {
assert_eq!(stats[1].column_index, 2);
}
}
+
+#[cfg(test)]
+#[path = "writer_integer_tests.rs"]
+mod integer_tests;
diff --git a/core/src/writer_integer_tests.rs b/core/src/writer_integer_tests.rs
new file mode 100644
index 0000000..2983cc5
--- /dev/null
+++ b/core/src/writer_integer_tests.rs
@@ -0,0 +1,426 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use super::*;
+use crate::reader::{InputFile, MosaicReader, ReaderAccess};
+use arrow_schema::{DataType, Field, Schema};
+use std::sync::{Arc, Mutex};
+
+impl InputFile for Bytes {
+ fn read_at(&self, offset: u64, out: &mut [u8]) -> io::Result<()> {
+ let start = offset as usize;
+ out.copy_from_slice(&self.0[start..start + out.len()]);
+ Ok(())
+ }
+}
+
+#[derive(Default)]
+struct Bytes(Vec<u8>);
+impl OutputFile for Bytes {
+ fn write(&mut self, bytes: &[u8]) -> io::Result<()> {
+ self.0.extend_from_slice(bytes);
+ Ok(())
+ }
+ fn flush(&mut self) -> io::Result<()> {
+ Ok(())
+ }
+ fn pos(&self) -> u64 {
+ self.0.len() as u64
+ }
+}
+
+#[derive(Default)]
+struct OutputCalls {
+ bytes: Vec<u8>,
+ writes: Vec<usize>,
+ flushes: usize,
+}
+
+struct RecordingOutput {
+ calls: Arc<Mutex<OutputCalls>>,
+ fail_on_write: Option<(usize, io::ErrorKind)>,
+}
+
+impl OutputFile for RecordingOutput {
+ fn write(&mut self, bytes: &[u8]) -> io::Result<()> {
+ let mut calls = self.calls.lock().unwrap();
+ calls.writes.push(bytes.len());
+ if let Some((_, kind)) = self
+ .fail_on_write
+ .filter(|(write, _)| *write == calls.writes.len())
+ {
+ // OutputFile callbacks can fail after writing a prefix. Retrying
+ // the buffered bytes would duplicate that prefix in the file.
+ calls.bytes.extend_from_slice(&bytes[..bytes.len() / 2]);
+ return Err(io::Error::new(kind, "buffered output failure"));
+ }
+ calls.bytes.extend_from_slice(bytes);
+ Ok(())
+ }
+
+ fn flush(&mut self) -> io::Result<()> {
+ self.calls.lock().unwrap().flushes += 1;
+ Ok(())
+ }
+
+ fn pos(&self) -> u64 {
+ self.calls.lock().unwrap().bytes.len() as u64
+ }
+}
+
+fn dictionary_input(rows: usize) -> RecordBatch {
+ let schema = Arc::new(Schema::new(
+ (0..100)
+ .map(|col| Field::new(format!("c{col:03}"), DataType::Int32,
false))
+ .collect::<Vec<_>>(),
+ ));
+ let values: ArrayRef = Arc::new(Int32Array::from(
+ (0..rows).map(|row| (row % 200) as i32).collect::<Vec<_>>(),
+ ));
+ RecordBatch::try_new(schema, vec![values; 100]).unwrap()
+}
+
+#[test]
+fn uncompressed_integer_output_coalesces_dictionary_writes() {
+ let input = dictionary_input(10_000);
+ let options = || WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 1,
+ row_group_max_size: 1,
+ ..Default::default()
+ };
+ let calls = Arc::new(Mutex::new(OutputCalls::default()));
+ let mut compact = MosaicWriter::new(
+ RecordingOutput {
+ calls: calls.clone(),
+ fail_on_write: None,
+ },
+ &input.schema(),
+ options(),
+ )
+ .unwrap();
+ let mut wide = MosaicWriter::new(Bytes::default(), &input.schema(),
options()).unwrap();
+ for bucket in wide.bucket_writers.iter_mut().flatten() {
+ bucket.use_wide_integer_buffers();
+ }
+ compact.write_batch(&input).unwrap();
+ wide.write_batch(&input).unwrap();
+ {
+ let calls = calls.lock().unwrap();
+ assert_eq!(calls.bytes, wide.output().0);
+ assert!(calls.writes.len() <=
calls.bytes.len().div_ceil(INTEGER_OUTPUT_BUFFER_SIZE / 2));
+ assert!(calls
+ .writes
+ .iter()
+ .all(|&len| len <= INTEGER_OUTPUT_BUFFER_SIZE));
+ assert_eq!(
+ calls.flushes, 0,
+ "bucket completion must not flush the sink"
+ );
+ }
+ compact.close().unwrap();
+ wide.close().unwrap();
+ let calls = calls.lock().unwrap();
+ assert_eq!(calls.bytes, wide.output().0);
+ assert_eq!(calls.flushes, 1);
+}
+
+#[test]
+fn integer_output_failure_does_not_retry_partial_writes() {
+ let plain = RecordBatch::try_from_iter([(
+ "id",
+ Arc::new(Int64Array::from_iter_values(0..200_003)) as ArrayRef,
+ )])
+ .unwrap();
+ // Cover final drain, buffered dictionary writes, direct value chunks,
+ // one-shot compression and streaming compression (> 1 MiB).
+ for (input, compression, fail_on_write) in [
+ (dictionary_input(1), COMPRESSION_NONE, 1),
+ (dictionary_input(10_000), COMPRESSION_NONE, 2),
+ (plain.clone(), COMPRESSION_NONE, 2),
+ (dictionary_input(1_000), COMPRESSION_ZSTD, 1),
+ (plain, COMPRESSION_ZSTD, 1),
+ ] {
+ for kind in [io::ErrorKind::Other, io::ErrorKind::Interrupted] {
+ for row_group_max_size in [1, u64::MAX] {
+ let calls = Arc::new(Mutex::new(OutputCalls::default()));
+ let mut writer = MosaicWriter::new(
+ RecordingOutput {
+ calls: calls.clone(),
+ fail_on_write: Some((fail_on_write, kind)),
+ },
+ &input.schema(),
+ WriterOptions {
+ compression,
+ num_buckets: 1,
+ row_group_max_size,
+ page_size_threshold: usize::MAX,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ let result = writer.write_batch(&input);
+ let error = if row_group_max_size == 1 {
+ result.unwrap_err()
+ } else {
+ result.unwrap();
+ writer.close().unwrap_err()
+ };
+ assert_eq!(error.to_string(), "buffered output failure");
+ assert_eq!(calls.lock().unwrap().writes.len(), fail_on_write);
+ let bytes_after_failure = calls.lock().unwrap().bytes.clone();
+ assert!(writer
+ .write_batch(&input)
+ .unwrap_err()
+ .to_string()
+ .contains("aborted"));
+
assert!(writer.close().unwrap_err().to_string().contains("aborted"));
+ drop(writer);
+ let calls = calls.lock().unwrap();
+ assert_eq!(calls.writes.len(), fail_on_write);
+ assert_eq!(calls.bytes, bytes_after_failure);
+ assert_eq!(calls.flushes, 0);
+ }
+ }
+ }
+}
+
+fn schema() -> Arc<Schema> {
+ Arc::new(Schema::new(vec![
+ Field::new("i8", DataType::Int8, true),
+ Field::new("i16", DataType::Int16, true),
+ Field::new("i32", DataType::Int32, true),
+ Field::new("i64", DataType::Int64, true),
+ Field::new("float", DataType::Float64, false),
+ Field::new("string", DataType::Utf8, true),
+ ]))
+}
+
+fn batch(schema: Arc<Schema>, mode: usize, batch: usize, rows: usize) ->
RecordBatch {
+ let value = |row: usize| -> Option<i64> {
+ let i = row + batch * rows;
+ match mode {
+ 0 => None,
+ 1 => Some(255),
+ 2 => {
+ if i.is_multiple_of(7) {
+ None
+ } else {
+ Some(65535)
+ }
+ }
+ 3 => {
+ if i.is_multiple_of(11) {
+ None
+ } else {
+ Some((i % 103) as i64)
+ }
+ }
+ _ => {
+ if i.is_multiple_of(13) {
+ None
+ } else {
+ Some(i as i64)
+ }
+ }
+ }
+ };
+ let ints: Vec<_> = (0..rows + 2).map(value).collect();
+ let i8 = Int8Array::from(ints.iter().map(|v| v.map(|v| v as
i8)).collect::<Vec<_>>());
+ let i16 = Int16Array::from(
+ ints.iter()
+ .map(|v| v.map(|v| (v % 256) as i16))
+ .collect::<Vec<_>>(),
+ );
+ let mut i32 = ints
+ .iter()
+ .map(|v| v.map(|v| (v % 65536) as i32))
+ .collect::<Vec<_>>();
+ let mut i64 = ints
+ .iter()
+ .map(|v| v.map(|v| (v * 2654435761) & 0xffff_ffff))
+ .collect::<Vec<_>>();
+ if mode == 5 && batch == 3 {
+ i32[rows - 1] = Some(i32::MIN);
+ i32[rows] = Some(i32::MAX);
+ i64[rows - 1] = Some(i64::MIN);
+ i64[rows] = Some(i64::MAX);
+ }
+ let arrays: Vec<ArrayRef> = vec![
+ Arc::new(i8),
+ Arc::new(i16),
+ Arc::new(Int32Array::from(i32)),
+ Arc::new(Int64Array::from(i64)),
+ Arc::new(Float64Array::from(
+ (0..rows + 2).map(|i| i as f64 / 7.0).collect::<Vec<_>>(),
+ )),
+ Arc::new(StringArray::from(
+ (0..rows + 2)
+ .map(|i| {
+ if i.is_multiple_of(17) {
+ None
+ } else {
+ Some("mixed-column")
+ }
+ })
+ .collect::<Vec<_>>(),
+ )),
+ ];
+ // A nonzero offset exercises Arrow slicing and nullable arrays together.
+ RecordBatch::try_new(
+ schema,
+ arrays.into_iter().map(|a| a.slice(1, rows)).collect(),
+ )
+ .unwrap()
+}
+
+#[test]
+fn integer_stream_preserves_values_layout_stats_and_compression() {
+ let schema = schema();
+ for compression in [COMPRESSION_NONE, COMPRESSION_ZSTD] {
+ for threshold in [1, 4096, usize::MAX] {
+ for row_group in [60_000, u64::MAX] {
+ for mode in 0..6 {
+ let options = || WriterOptions {
+ compression,
+ num_buckets: 1,
+ page_size_threshold: threshold,
+ row_group_max_size: row_group,
+ zstd_level: 3,
+ stats_columns: vec!["i16".into(), "i32".into(),
"i64".into()],
+ ..Default::default()
+ };
+ let mut compact =
+ MosaicWriter::new(Bytes::default(), &schema,
options()).unwrap();
+ let mut wide = MosaicWriter::new(Bytes::default(),
&schema, options()).unwrap();
+ for bucket in wide.bucket_writers.iter_mut().flatten() {
+ bucket.use_wide_integer_buffers();
+ }
+ for i in 0..5 {
+ let input = batch(schema.clone(), mode, i, 10_003);
+ compact.write_batch(&input).unwrap();
+ wide.write_batch(&input).unwrap();
+ assert_eq!(compact.current_buffered_size,
wide.current_buffered_size);
+ assert_eq!(compact.row_group_metas.len(),
wide.row_group_metas.len());
+ assert_eq!(compact.estimated_file_size(),
wide.estimated_file_size());
+ }
+ compact.close().unwrap();
+ wide.close().unwrap();
+ let actual = &compact.output().0;
+ let expected = &wide.output().0;
+ assert!(actual.len() <= expected.len(),
+ "larger file: compression={compression},
threshold={threshold}, row_group={row_group}, mode={mode}; {} vs {}",
actual.len(), expected.len());
+ if compression == COMPRESSION_NONE {
+ assert!(actual == expected, "uncompressed encoding
changed");
+ }
+ let a = MosaicReader::new(Bytes(actual.clone()),
actual.len() as u64).unwrap();
+ let b =
+ MosaicReader::new(Bytes(expected.clone()),
expected.len() as u64).unwrap();
+ assert_eq!(a.num_row_groups(), b.num_row_groups());
+ assert_eq!(a.schema().serialize(), b.schema().serialize());
+ let mut consumed = 0;
+ for rg in 0..a.num_row_groups() {
+ assert_eq!(
+ a.row_group_num_rows(rg).unwrap(),
+ b.row_group_num_rows(rg).unwrap()
+ );
+ assert_eq!(
+ stats::serialize_stats(
+ a.row_group_stats(rg).unwrap(),
+ &a.schema().columns
+ ),
+ stats::serialize_stats(
+ b.row_group_stats(rg).unwrap(),
+ &b.schema().columns
+ )
+ );
+ let a_batch =
a.row_group_reader(rg).unwrap().read_columns().unwrap();
+ let b_batch =
b.row_group_reader(rg).unwrap().read_columns().unwrap();
+ assert!(a_batch == b_batch, "decoded values changed");
+ let mut row = 0;
+ while row < a_batch.num_rows() {
+ let input = batch(schema.clone(), mode, consumed /
10_003, 10_003);
+ let offset = consumed % 10_003;
+ let count = (10_003 -
offset).min(a_batch.num_rows() - row);
+ for field in schema.fields() {
+ let actual = a_batch
+ .column_by_name(field.name())
+ .unwrap()
+ .slice(row, count);
+ let expected = input
+ .column_by_name(field.name())
+ .unwrap()
+ .slice(offset, count);
+ assert!(
+ actual.to_data() == expected.to_data(),
+ "input values changed: {}",
+ field.name()
+ );
+ }
+ row += count;
+ consumed += count;
+ }
+ }
+ assert_eq!(consumed, 50_015);
+ }
+ }
+ }
+ }
+}
+
+#[test]
+fn integer_only_stream_matches_one_shot_at_zstd_block_boundaries() {
+ let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]);
+ for rows in [
+ 0, 1, 4095, 4096, 4097, 16383, 16384, 16385, 100_003, 200_003,
+ ] {
+ for level in [-1, 1, 9] {
+ for threshold in [1, usize::MAX] {
+ let options = || WriterOptions {
+ num_buckets: 1,
+ max_dict_entries: 0,
+ page_size_threshold: threshold,
+ zstd_level: level,
+ row_group_max_size: u64::MAX,
+ ..Default::default()
+ };
+ let mut compact = MosaicWriter::new(Bytes::default(), &schema,
options()).unwrap();
+ let mut wide = MosaicWriter::new(Bytes::default(), &schema,
options()).unwrap();
+ for bucket in wide.bucket_writers.iter_mut().flatten() {
+ bucket.use_wide_integer_buffers();
+ }
+ let input = RecordBatch::try_new(
+ Arc::new(schema.clone()),
+ vec![Arc::new(Int64Array::from(
+ (0..rows)
+ .map(|i| ((i as u64 * 2654435761) & 0xffff_ffff)
as i64)
+ .collect::<Vec<_>>(),
+ ))],
+ )
+ .unwrap();
+ compact.write_batch(&input).unwrap();
+ wide.write_batch(&input).unwrap();
+ compact.close().unwrap();
+ wide.close().unwrap();
+ assert!(
+ compact.output().0 == wide.output().0,
+ "rows={rows}, level={level}, threshold={threshold}"
+ );
+ }
+ }
+ }
+}
