Initial write support

[Jefffrey]

Support for writing the following Arrow array types to an ORC file:

• Int8/Int16/Int32/Int64
• Float32/Float64

Broadly speaking, this writer will delegate to internal types to encode the actual primitive values into a run encoded form (RLEv2 for Int, Byte RLE for bytes, simple copy for floats) into an in-memory buffer until the stripe size exceeds a configured limit in which case the stripe is flushed to the actual writer (e.g. to disk) and the internal buffers cleared.

• Not much focus on optimization, so the internal buffers aren't being reusued, this will be addressed later

Overall architecture

src/arrow_writer.rs contains the public interface for this writer, where an ArrowWriterBuilder is used to specify config options to build an ArrowWriter which allows writing a batch at a time, see example usage:

datafusion-orc/src/arrow_writer.rs

Lines 275 to 280 in 5afcfe2

	let f = File::create("/tmp/new.orc").unwrap();
	let mut writer = ArrowWriterBuilder::new(f, batch.schema())
	.try_build()
	.unwrap();
	writer.write(&batch).unwrap();
	writer.close().unwrap();

This writer will automatically flush a stripe when it is estimated to have exceeded the configured stripe size (default 64mb). Internally, this writer has a StripeWriter to handle encoding the actual stripe in-memory:

datafusion-orc/src/writer/stripe.rs

Lines 40 to 47 in 5afcfe2

	/// Encode a stripe. Will encode columns into an in-memory buffer before flushing
	/// entire stripe to the underlying writer.
	pub struct StripeWriter<W> {
	writer: W,
	/// Flattened columns, in order of their column ID.
	columns: Vec<Box<dyn ColumnStripeEncoder>>,
	pub row_count: usize,
	}

This StripeWriter has an array of ColumnStripeEncoders which are responsible for the encoding of individual columns within the recordbatch/array. Note the flat structure, this will need to be revisited when accounting for nested types (map, struct, etc.).

datafusion-orc/src/writer/column.rs

Lines 27 to 40 in 5afcfe2

	/// Encodes a specific column for a stripe. Will encode to an internal memory
	/// buffer until it is finished, in which case it returns the stream bytes to
	/// be serialized to a writer.
	pub trait ColumnStripeEncoder: EstimateMemory {
	/// Encode entire provided [`ArrayRef`] to internal buffer.
	fn encode_array(&mut self, array: &ArrayRef) -> Result<()>;
	/// Column encoding used for streams.
	fn column_encoding(&self) -> ColumnEncoding;
	/// Emit buffered streams to be written to the writer, and reset state
	/// in preparation for next stripe.
	fn finish(&mut self) -> Vec<Stream>;
	}

Currently only one implementation for ColumnStripeEncoder which is PrimitiveStripeEncoder:

datafusion-orc/src/writer/column.rs

Lines 66 to 76 in 5afcfe2

	/// Encoder for primitive ORC types (e.g. int, float). Uses a specific [`PrimitiveValueEncoder`] to
	/// encode the primitive values into internal memory. When finished, outputs a DATA stream and
	/// optionally a PRESENT stream.
	pub struct PrimitiveStripeEncoder<T: ArrowPrimitiveType, E: PrimitiveValueEncoder<T::Native>> {
	encoder: E,
	column_encoding: ColumnEncoding,
	/// Lazily initialized once we encounter an [`Array`] with a [`NullBuffer`].
	present: Option<PresentStreamEncoder>,
	encoded_count: usize,
	_phantom: PhantomData<T>,
	}

This handles encoding the actual primitive values themselves, which is then handled by the PrimitiveValueEncoder trait, and also optionally handles encoding a present stream.

datafusion-orc/src/writer/column.rs

Lines 42 to 62 in 5afcfe2

	/// Encodes primitive values into an internal buffer, usually with a specialized run length
	/// encoding for better compression.
	pub trait PrimitiveValueEncoder<V>: EstimateMemory
	where
	V: Copy,
	{
	fn new() -> Self;
	fn write_one(&mut self, value: V);
	fn write_slice(&mut self, values: &[V]) {
	for &value in values {
	self.write_one(value);
	}
	}
	/// Take the encoded bytes, replacing it with an empty buffer.
	// TODO: Figure out how to retain the allocation instead of handing
	// it off each time.
	fn take_inner(&mut self) -> Bytes;
	}

Encoding

There are currently three implementations.

Float has the simplest as ORC specifies no special encoding for it, it simply copies. Refer to FloatValueEncoder:

datafusion-orc/src/writer/column.rs

Lines 175 to 184 in 5afcfe2

	/// No special run encoding for floats/doubles, they are stored as their IEEE 754 floating
	/// point bit layout. This encoder simply copies incoming floats/doubles to its internal
	/// byte buffer.
	pub struct FloatValueEncoder<T: ArrowPrimitiveType>
	where
	T::Native: Float,
	{
	data: BytesMut,
	_phantom: PhantomData<T>,
	}

Bytes have some simple logic for run length encoding, see ByteRleWriter:

datafusion-orc/src/reader/decode/byte_rle.rs

Lines 32 to 44 in 5afcfe2

	pub struct ByteRleWriter {
	writer: BytesMut,
	/// Literal values to encode.
	literals: [u8; MAX_LITERAL_LENGTH],
	/// Represents the number of elements currently in `literals` if Literals,
	/// otherwise represents the length of the Run.
	num_literals: usize,
	/// Tracks if current Literal sequence will turn into a Run sequence due to
	/// repeated values at the end of the value sequence.
	tail_run_length: usize,
	/// If in Run sequence or not, and keeps the corresponding value.
	run_value: Option<u8>,
	}

Integers have the most complicated version due to the complex sub0encoding types (as well as the actual logic to determine which sub0encoding to use), see RleWriterV2:

datafusion-orc/src/reader/decode/rle_v2/mod.rs

Lines 208 to 215 in 5afcfe2

	pub struct RleWriterV2<N: NInt, S: EncodingSign> {
	/// Stores the run length encoded sequences.
	data: BytesMut,
	/// Used in state machine for determining which sub-encoding
	/// for a sequence to use.
	state: RleV2EncodingState<N>,
	phantom: PhantomData<S>,
	}

• Write support for V1 encoding seems pointless, so not implemented

For both Byte and Int RLEv2 encodings, the logic has been translated from the Java implementation of ORC in order to capture details the spec doesn't specify as well as use the logic it uses to determine which sub-encodings to use. This logic was refactored to make the logic easier to understand (at least to me).

Other changes

• Bumped rust version to 1.73 (to be able to use div_ceil)
• Introduced proptests, very handy for testing this serde behaviour
• Removed usage of u64 in integer run length encoding code, as this is incorrect since the Java implementation doesn't support unsigned integers (everything is an i64 there)
  • Though the ORC spec states that some types like string use unsigned integer encoding, this is still going to be an i64 so stick with that
• Rename ByteRleIter to ByteRleReader for consistency with other reader interfaces (such as RLEv2)

[Jefffrey]

Writing string/boolean arrays will be done after this since the PR is already quite large

[Jefffrey]

Sorry for such a large PR 😅

I wanted an end to end example, and felt it wouldn't be proper without at least integer support. I didn't want to waste time implementing RLEv1 writing, so the majority of time was spent on trying to understand the RLEv2 write logic in the Java implementation and translate it to code that I think might be easier to understand (I didn't want to just do a one-to-one port from the C++/Java implementation to here).

Pinging for awareness @WenyXu @Xuanwo @waynexia

[Xuanwo]

Thank you very much for this PR; it's really extensive! This review primarily focuses on the API related to IO operations.

[Xuanwo]

What do you think using stripe_size or stripe_length for short?

[Jefffrey]

I like keeping the unit as part of the name, to be unambiguous

[WenyXu]

Rest LGTM