Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Various refactoring #39

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Various refactoring #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
6c94b3d
Refactor retrieval of int decoders
Jefffrey Jan 1, 2025
6df9d31
Remove read_stream methods from Column
Jefffrey Jan 1, 2025
279efea
Minor refactoring
Jefffrey Jan 2, 2025
8c0bd77
Minor refactoring
Jefffrey Jan 2, 2025
2059515
Minor refactoring around stripe/column
Jefffrey Jan 2, 2025
4989285
Allow RootDataType to lookup all transitive column indices
Jefffrey Jan 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions src/array_decoder/decimal.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use arrow::datatypes::Decimal128Type;
use snafu::ResultExt;

use crate::encoding::decimal::UnboundedVarintStreamDecoder;
use crate::encoding::integer::get_rle_reader;
use crate::encoding::integer::get_signed_int_decoder;
use crate::encoding::PrimitiveValueDecoder;
use crate::error::ArrowSnafu;
use crate::proto::stream::Kind;
Expand All @@ -38,13 +38,13 @@ pub fn new_decimal_decoder(
stripe: &Stripe,
precision: u32,
fixed_scale: u32,
) -> Result<Box<dyn ArrayBatchDecoder>> {
) -> Box<dyn ArrayBatchDecoder> {
let varint_iter = stripe.stream_map().get(column, Kind::Data);
let varint_iter = Box::new(UnboundedVarintStreamDecoder::new(varint_iter));

// Scale is specified on a per varint basis (in addition to being encoded in the type)
let scale_iter = stripe.stream_map().get(column, Kind::Secondary);
let scale_iter = get_rle_reader::<i32, _>(column, scale_iter)?;
let scale_iter = get_signed_int_decoder::<i32>(scale_iter, column.rle_version());

let present = PresentDecoder::from_stripe(stripe, column);

Expand All @@ -55,12 +55,12 @@ pub fn new_decimal_decoder(
};
let iter = Box::new(iter);

Ok(Box::new(DecimalArrayDecoder::new(
Box::new(DecimalArrayDecoder::new(
precision as u8,
fixed_scale as i8,
iter,
present,
)))
))
}

/// Wrapper around PrimitiveArrayDecoder to allow specifying the precision and scale
Expand Down
6 changes: 3 additions & 3 deletions src/array_decoder/list.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use snafu::ResultExt;

use crate::array_decoder::derive_present_vec;
use crate::column::Column;
use crate::encoding::integer::get_unsigned_rle_reader;
use crate::encoding::integer::get_unsigned_int_decoder;
use crate::encoding::PrimitiveValueDecoder;
use crate::proto::stream::Kind;

Expand All @@ -45,10 +45,10 @@ impl ListArrayDecoder {
let present = PresentDecoder::from_stripe(stripe, column);

let child = &column.children()[0];
let inner = array_decoder_factory(child, field.clone(), stripe)?;
let inner = array_decoder_factory(child, field.data_type(), stripe)?;

let reader = stripe.stream_map().get(column, Kind::Length);
let lengths = get_unsigned_rle_reader(column, reader);
let lengths = get_unsigned_int_decoder(reader, column.rle_version());

Ok(Self {
inner,
Expand Down
8 changes: 4 additions & 4 deletions src/array_decoder/map.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use snafu::ResultExt;

use crate::array_decoder::derive_present_vec;
use crate::column::Column;
use crate::encoding::integer::get_unsigned_rle_reader;
use crate::encoding::integer::get_unsigned_int_decoder;
use crate::encoding::PrimitiveValueDecoder;
use crate::error::{ArrowSnafu, Result};
use crate::proto::stream::Kind;
Expand All @@ -50,13 +50,13 @@ impl MapArrayDecoder {
let present = PresentDecoder::from_stripe(stripe, column);

let keys_column = &column.children()[0];
let keys = array_decoder_factory(keys_column, keys_field.clone(), stripe)?;
let keys = array_decoder_factory(keys_column, keys_field.data_type(), stripe)?;

let values_column = &column.children()[1];
let values = array_decoder_factory(values_column, values_field.clone(), stripe)?;
let values = array_decoder_factory(values_column, values_field.data_type(), stripe)?;

let reader = stripe.stream_map().get(column, Kind::Length);
let lengths = get_unsigned_rle_reader(column, reader);
let lengths = get_unsigned_int_decoder(reader, column.rle_version());

let fields = Fields::from(vec![keys_field, values_field]);

Expand Down
30 changes: 13 additions & 17 deletions src/array_decoder/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use arrow::array::{ArrayRef, BooleanArray, BooleanBufferBuilder, PrimitiveArray}
use arrow::buffer::NullBuffer;
use arrow::datatypes::ArrowNativeTypeOp;
use arrow::datatypes::ArrowPrimitiveType;
use arrow::datatypes::{DataType as ArrowDataType, Field};
use arrow::datatypes::DataType as ArrowDataType;
use arrow::datatypes::{
Date32Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, SchemaRef,
};
Expand All @@ -32,7 +32,7 @@ use crate::column::Column;
use crate::encoding::boolean::BooleanDecoder;
use crate::encoding::byte::ByteRleDecoder;
use crate::encoding::float::FloatDecoder;
use crate::encoding::integer::get_rle_reader;
use crate::encoding::integer::get_signed_int_decoder;
use crate::encoding::PrimitiveValueDecoder;
use crate::error::{
self, MismatchedSchemaSnafu, Result, UnexpectedSnafu, UnsupportedTypeVariantSnafu,
Expand Down Expand Up @@ -258,10 +258,10 @@ impl Iterator for NaiveStripeDecoder {

pub fn array_decoder_factory(
column: &Column,
field: Arc<Field>,
hinted_arrow_type: &ArrowDataType,
stripe: &Stripe,
) -> Result<Box<dyn ArrayBatchDecoder>> {
let decoder: Box<dyn ArrayBatchDecoder> = match (column.data_type(), field.data_type()) {
let decoder: Box<dyn ArrayBatchDecoder> = match (column.data_type(), hinted_arrow_type) {
// TODO: try make branches more generic, reduce duplication
(DataType::Boolean { .. }, ArrowDataType::Boolean) => {
let iter = stripe.stream_map().get(column, Kind::Data);
Expand All @@ -277,19 +277,19 @@ pub fn array_decoder_factory(
}
(DataType::Short { .. }, ArrowDataType::Int16) => {
let iter = stripe.stream_map().get(column, Kind::Data);
let iter = get_rle_reader(column, iter)?;
let iter = get_signed_int_decoder(iter, column.rle_version());
let present = PresentDecoder::from_stripe(stripe, column);
Box::new(Int16ArrayDecoder::new(iter, present))
}
(DataType::Int { .. }, ArrowDataType::Int32) => {
let iter = stripe.stream_map().get(column, Kind::Data);
let iter = get_rle_reader(column, iter)?;
let iter = get_signed_int_decoder(iter, column.rle_version());
let present = PresentDecoder::from_stripe(stripe, column);
Box::new(Int32ArrayDecoder::new(iter, present))
}
(DataType::Long { .. }, ArrowDataType::Int64) => {
let iter = stripe.stream_map().get(column, Kind::Data);
let iter = get_rle_reader(column, iter)?;
let iter = get_signed_int_decoder(iter, column.rle_version());
let present = PresentDecoder::from_stripe(stripe, column);
Box::new(Int64ArrayDecoder::new(iter, present))
}
Expand All @@ -315,7 +315,7 @@ pub fn array_decoder_factory(
},
ArrowDataType::Decimal128(a_precision, a_scale),
) if *precision as u8 == *a_precision && *scale as i8 == *a_scale => {
new_decimal_decoder(column, stripe, *precision, *scale)?
new_decimal_decoder(column, stripe, *precision, *scale)
}
(DataType::Timestamp { .. }, field_type) => {
new_timestamp_decoder(column, field_type.clone(), stripe)?
Expand All @@ -326,7 +326,7 @@ pub fn array_decoder_factory(
(DataType::Date { .. }, ArrowDataType::Date32) => {
// TODO: allow Date64
let iter = stripe.stream_map().get(column, Kind::Data);
let iter = get_rle_reader(column, iter)?;
let iter = get_signed_int_decoder(iter, column.rle_version());
let present = PresentDecoder::from_stripe(stripe, column);
Box::new(DateArrayDecoder::new(iter, present))
}
Expand Down Expand Up @@ -433,17 +433,13 @@ impl NaiveStripeDecoder {
}

pub fn new(stripe: Stripe, schema_ref: SchemaRef, batch_size: usize) -> Result<Self> {
let mut decoders = Vec::with_capacity(stripe.columns().len());
let number_of_rows = stripe.number_of_rows();

for (col, field) in stripe
let decoders = stripe
.columns()
.iter()
.zip(schema_ref.fields.iter().cloned())
{
let decoder = array_decoder_factory(col, field, &stripe)?;
decoders.push(decoder);
}
.zip(schema_ref.fields.iter())
.map(|(col, field)| array_decoder_factory(col, field.data_type(), &stripe))
.collect::<Result<Vec<_>>>()?;

Ok(Self {
stripe,
Expand Down
8 changes: 4 additions & 4 deletions src/array_decoder/string.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ use snafu::ResultExt;
use crate::array_decoder::derive_present_vec;
use crate::column::Column;
use crate::compression::Decompressor;
use crate::encoding::integer::get_unsigned_rle_reader;
use crate::encoding::integer::get_unsigned_int_decoder;
use crate::encoding::PrimitiveValueDecoder;
use crate::error::{ArrowSnafu, IoSnafu, Result};
use crate::proto::column_encoding::Kind as ColumnEncodingKind;
Expand All @@ -42,7 +42,7 @@ pub fn new_binary_decoder(column: &Column, stripe: &Stripe) -> Result<Box<dyn Ar
let present = PresentDecoder::from_stripe(stripe, column);

let lengths = stripe.stream_map().get(column, Kind::Length);
let lengths = get_unsigned_rle_reader(column, lengths);
let lengths = get_unsigned_int_decoder(lengths, column.rle_version());

let bytes = Box::new(stripe.stream_map().get(column, Kind::Data));
Ok(Box::new(BinaryArrayDecoder::new(bytes, lengths, present)))
Expand All @@ -53,7 +53,7 @@ pub fn new_string_decoder(column: &Column, stripe: &Stripe) -> Result<Box<dyn Ar
let present = PresentDecoder::from_stripe(stripe, column);

let lengths = stripe.stream_map().get(column, Kind::Length);
let lengths = get_unsigned_rle_reader(column, lengths);
let lengths = get_unsigned_int_decoder(lengths, column.rle_version());

match kind {
ColumnEncodingKind::Direct | ColumnEncodingKind::DirectV2 => {
Expand All @@ -72,7 +72,7 @@ pub fn new_string_decoder(column: &Column, stripe: &Stripe) -> Result<Box<dyn Ar
let dictionary_strings = Arc::new(dictionary_strings);

let indexes = stripe.stream_map().get(column, Kind::Data);
let indexes = get_unsigned_rle_reader(column, indexes);
let indexes = get_unsigned_int_decoder(indexes, column.rle_version());
let indexes = Int64ArrayDecoder::new(indexes, present);

Ok(Box::new(DictionaryStringArrayDecoder::new(
Expand Down
4 changes: 2 additions & 2 deletions src/array_decoder/struct_decoder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@ impl StructArrayDecoder {
let decoders = column
.children()
.iter()
.zip(fields.iter().cloned())
.map(|(child, field)| array_decoder_factory(child, field, stripe))
.zip(fields.iter())
.map(|(child, field)| array_decoder_factory(child, field.data_type(), stripe))
.collect::<Result<Vec<_>>>()?;

Ok(Self {
Expand Down
Loading
Loading