Skip to content

Commit 410efcb

Browse files
waralexromwaralexrom
authored
feat: Int96 and Decimal96 support (#21)
1 parent 6bfc9fb commit 410efcb

File tree

30 files changed

+982
-96
lines changed

30 files changed

+982
-96
lines changed

‎.github/workflows/miri.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ jobs:
3030
strategy:
3131
matrix:
3232
arch: [amd64]
33-
rust: [nightly-2021-07-04]
33+
rust: [nightly-2022-06-22]
3434
steps:
3535
- uses: actions/checkout@v2
3636
with:

‎.github/workflows/rust.yml

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ jobs:
3131
strategy:
3232
matrix:
3333
arch: [amd64]
34-
rust: [nightly-2021-07-04]
34+
rust: [nightly-2022-06-22]
3535
container:
3636
image: ${{ matrix.arch }}/rust
3737
env:
@@ -73,7 +73,7 @@ jobs:
7373
strategy:
7474
matrix:
7575
arch: [amd64]
76-
rust: [nightly-2021-07-04]
76+
rust: [nightly-2022-06-22]
7777
container:
7878
image: ${{ matrix.arch }}/rust
7979
env:
@@ -174,7 +174,7 @@ jobs:
174174
strategy:
175175
matrix:
176176
os: [windows-latest, macos-latest]
177-
rust: [nightly-2021-07-04]
177+
rust: [nightly-2022-06-22]
178178
steps:
179179
- uses: actions/checkout@v2
180180
with:
@@ -202,7 +202,7 @@ jobs:
202202
strategy:
203203
matrix:
204204
arch: [amd64]
205-
rust: [nightly-2021-07-04]
205+
rust: [nightly-2022-06-22]
206206
container:
207207
image: ${{ matrix.arch }}/rust
208208
env:
@@ -257,7 +257,7 @@ jobs:
257257
strategy:
258258
matrix:
259259
arch: [amd64]
260-
rust: [nightly-2021-07-04]
260+
rust: [nightly-2022-06-22]
261261
steps:
262262
- uses: actions/checkout@v2
263263
with:
@@ -297,7 +297,7 @@ jobs:
297297
strategy:
298298
matrix:
299299
arch: [amd64]
300-
rust: [nightly-2021-07-04]
300+
rust: [nightly-2022-06-22]
301301
container:
302302
image: ${{ matrix.arch }}/rust
303303
env:
@@ -341,7 +341,7 @@ jobs:
341341
strategy:
342342
matrix:
343343
arch: [amd64]
344-
rust: [nightly-2021-07-04]
344+
rust: [nightly-2022-06-22]
345345
container:
346346
image: ${{ matrix.arch }}/rust
347347
env:

‎arrow/src/array/array.rs

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,7 @@ pub fn make_array(data: ArrayData) -> ArrayRef {
235235
DataType::Int16 => Arc::new(Int16Array::from(data)) as ArrayRef,
236236
DataType::Int32 => Arc::new(Int32Array::from(data)) as ArrayRef,
237237
DataType::Int64 => Arc::new(Int64Array::from(data)) as ArrayRef,
238+
DataType::Int96 => Arc::new(Int96Array::from(data)) as ArrayRef,
238239
DataType::Int64Decimal(0) => Arc::new(Int64Decimal0Array::from(data)) as ArrayRef,
239240
DataType::Int64Decimal(1) => Arc::new(Int64Decimal1Array::from(data)) as ArrayRef,
240241
DataType::Int64Decimal(2) => Arc::new(Int64Decimal2Array::from(data)) as ArrayRef,
@@ -244,6 +245,15 @@ pub fn make_array(data: ArrayData) -> ArrayRef {
244245
DataType::Int64Decimal(10) => {
245246
Arc::new(Int64Decimal10Array::from(data)) as ArrayRef
246247
}
248+
DataType::Int96Decimal(0) => Arc::new(Int96Decimal0Array::from(data)) as ArrayRef,
249+
DataType::Int96Decimal(1) => Arc::new(Int96Decimal1Array::from(data)) as ArrayRef,
250+
DataType::Int96Decimal(2) => Arc::new(Int96Decimal2Array::from(data)) as ArrayRef,
251+
DataType::Int96Decimal(3) => Arc::new(Int96Decimal3Array::from(data)) as ArrayRef,
252+
DataType::Int96Decimal(4) => Arc::new(Int96Decimal4Array::from(data)) as ArrayRef,
253+
DataType::Int96Decimal(5) => Arc::new(Int96Decimal5Array::from(data)) as ArrayRef,
254+
DataType::Int96Decimal(10) => {
255+
Arc::new(Int96Decimal10Array::from(data)) as ArrayRef
256+
}
247257
DataType::UInt8 => Arc::new(UInt8Array::from(data)) as ArrayRef,
248258
DataType::UInt16 => Arc::new(UInt16Array::from(data)) as ArrayRef,
249259
DataType::UInt32 => Arc::new(UInt32Array::from(data)) as ArrayRef,
@@ -401,6 +411,7 @@ pub fn new_null_array(data_type: &DataType, length: usize) -> ArrayRef {
401411
// expanding this into Date23{unit}Type results in needless branching
402412
DataType::Time32(_) => new_null_sized_array::<Int32Type>(data_type, length),
403413
DataType::Int64 => new_null_sized_array::<Int64Type>(data_type, length),
414+
DataType::Int96 => new_null_sized_array::<Int96Type>(data_type, length),
404415
DataType::Int64Decimal(0) => {
405416
new_null_sized_array::<Int64Decimal0Type>(data_type, length)
406417
}
@@ -423,6 +434,28 @@ pub fn new_null_array(data_type: &DataType, length: usize) -> ArrayRef {
423434
new_null_sized_array::<Int64Decimal10Type>(data_type, length)
424435
}
425436
DataType::Int64Decimal(_) => panic!("invalid scale"),
437+
DataType::Int96Decimal(0) => {
438+
new_null_sized_array::<Int96Decimal0Type>(data_type, length)
439+
}
440+
DataType::Int96Decimal(1) => {
441+
new_null_sized_array::<Int96Decimal1Type>(data_type, length)
442+
}
443+
DataType::Int96Decimal(2) => {
444+
new_null_sized_array::<Int96Decimal2Type>(data_type, length)
445+
}
446+
DataType::Int96Decimal(3) => {
447+
new_null_sized_array::<Int96Decimal3Type>(data_type, length)
448+
}
449+
DataType::Int96Decimal(4) => {
450+
new_null_sized_array::<Int96Decimal4Type>(data_type, length)
451+
}
452+
DataType::Int96Decimal(5) => {
453+
new_null_sized_array::<Int96Decimal5Type>(data_type, length)
454+
}
455+
DataType::Int96Decimal(10) => {
456+
new_null_sized_array::<Int96Decimal10Type>(data_type, length)
457+
}
458+
DataType::Int96Decimal(_) => panic!("invalid scale"),
426459
DataType::UInt64 => new_null_sized_array::<UInt64Type>(data_type, length),
427460
DataType::Float64 => new_null_sized_array::<Float64Type>(data_type, length),
428461
DataType::Date64 => new_null_sized_array::<Date64Type>(data_type, length),

‎arrow/src/array/array_primitive.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -393,13 +393,21 @@ def_numeric_from_vec!(Int8Type);
393393
def_numeric_from_vec!(Int16Type);
394394
def_numeric_from_vec!(Int32Type);
395395
def_numeric_from_vec!(Int64Type);
396+
def_numeric_from_vec!(Int96Type);
396397
def_numeric_from_vec!(Int64Decimal0Type);
397398
def_numeric_from_vec!(Int64Decimal1Type);
398399
def_numeric_from_vec!(Int64Decimal2Type);
399400
def_numeric_from_vec!(Int64Decimal3Type);
400401
def_numeric_from_vec!(Int64Decimal4Type);
401402
def_numeric_from_vec!(Int64Decimal5Type);
402403
def_numeric_from_vec!(Int64Decimal10Type);
404+
def_numeric_from_vec!(Int96Decimal0Type);
405+
def_numeric_from_vec!(Int96Decimal1Type);
406+
def_numeric_from_vec!(Int96Decimal2Type);
407+
def_numeric_from_vec!(Int96Decimal3Type);
408+
def_numeric_from_vec!(Int96Decimal4Type);
409+
def_numeric_from_vec!(Int96Decimal5Type);
410+
def_numeric_from_vec!(Int96Decimal10Type);
403411
def_numeric_from_vec!(UInt8Type);
404412
def_numeric_from_vec!(UInt16Type);
405413
def_numeric_from_vec!(UInt32Type);

‎arrow/src/array/builder.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1451,13 +1451,21 @@ pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<ArrayBuilder> {
14511451
DataType::Int16 => Box::new(Int16Builder::new(capacity)),
14521452
DataType::Int32 => Box::new(Int32Builder::new(capacity)),
14531453
DataType::Int64 => Box::new(Int64Builder::new(capacity)),
1454+
DataType::Int96 => Box::new(Int96Builder::new(capacity)),
14541455
DataType::Int64Decimal(0) => Box::new(Int64Decimal0Builder::new(capacity)),
14551456
DataType::Int64Decimal(1) => Box::new(Int64Decimal1Builder::new(capacity)),
14561457
DataType::Int64Decimal(2) => Box::new(Int64Decimal0Builder::new(capacity)),
14571458
DataType::Int64Decimal(3) => Box::new(Int64Decimal0Builder::new(capacity)),
14581459
DataType::Int64Decimal(4) => Box::new(Int64Decimal0Builder::new(capacity)),
14591460
DataType::Int64Decimal(5) => Box::new(Int64Decimal0Builder::new(capacity)),
14601461
DataType::Int64Decimal(10) => Box::new(Int64Decimal0Builder::new(capacity)),
1462+
DataType::Int96Decimal(0) => Box::new(Int96Decimal0Builder::new(capacity)),
1463+
DataType::Int96Decimal(1) => Box::new(Int96Decimal1Builder::new(capacity)),
1464+
DataType::Int96Decimal(2) => Box::new(Int96Decimal0Builder::new(capacity)),
1465+
DataType::Int96Decimal(3) => Box::new(Int96Decimal0Builder::new(capacity)),
1466+
DataType::Int96Decimal(4) => Box::new(Int96Decimal0Builder::new(capacity)),
1467+
DataType::Int96Decimal(5) => Box::new(Int96Decimal0Builder::new(capacity)),
1468+
DataType::Int96Decimal(10) => Box::new(Int96Decimal0Builder::new(capacity)),
14611469
DataType::UInt8 => Box::new(UInt8Builder::new(capacity)),
14621470
DataType::UInt16 => Box::new(UInt16Builder::new(capacity)),
14631471
DataType::UInt32 => Box::new(UInt32Builder::new(capacity)),

‎arrow/src/array/data.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,10 @@ pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuff
7575
MutableBuffer::new(capacity * mem::size_of::<u64>()),
7676
empty_buffer,
7777
],
78+
DataType::Int96Decimal(_) => [
79+
MutableBuffer::new(capacity * mem::size_of::<u128>()),
80+
empty_buffer,
81+
],
7882
DataType::Int8 => [
7983
MutableBuffer::new(capacity * mem::size_of::<i8>()),
8084
empty_buffer,
@@ -91,6 +95,10 @@ pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuff
9195
MutableBuffer::new(capacity * mem::size_of::<i64>()),
9296
empty_buffer,
9397
],
98+
DataType::Int96 => [
99+
MutableBuffer::new(capacity * mem::size_of::<i128>()),
100+
empty_buffer,
101+
],
94102
DataType::Float32 => [
95103
MutableBuffer::new(capacity * mem::size_of::<f32>()),
96104
empty_buffer,
@@ -430,7 +438,9 @@ impl ArrayData {
430438
| DataType::Int16
431439
| DataType::Int32
432440
| DataType::Int64
441+
| DataType::Int96
433442
| DataType::Int64Decimal(_)
443+
| DataType::Int96Decimal(_)
434444
| DataType::Float32
435445
| DataType::Float64
436446
| DataType::Date32

‎arrow/src/array/equal/mod.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,9 +175,15 @@ fn equal_values(
175175
DataType::Int64 => primitive_equal::<i64>(
176176
lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
177177
),
178+
DataType::Int96 => primitive_equal::<i128>(
179+
lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
180+
),
178181
DataType::Int64Decimal(_) => primitive_equal::<i64>(
179182
lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
180183
),
184+
DataType::Int96Decimal(_) => primitive_equal::<i64>(
185+
lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
186+
),
181187
DataType::Float32 => primitive_equal::<f32>(
182188
lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len,
183189
),

‎arrow/src/array/mod.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,13 +137,21 @@ pub type Int8Array = PrimitiveArray<Int8Type>;
137137
pub type Int16Array = PrimitiveArray<Int16Type>;
138138
pub type Int32Array = PrimitiveArray<Int32Type>;
139139
pub type Int64Array = PrimitiveArray<Int64Type>;
140+
pub type Int96Array = PrimitiveArray<Int96Type>;
140141
pub type Int64Decimal0Array = PrimitiveArray<Int64Decimal0Type>;
141142
pub type Int64Decimal1Array = PrimitiveArray<Int64Decimal1Type>;
142143
pub type Int64Decimal2Array = PrimitiveArray<Int64Decimal2Type>;
143144
pub type Int64Decimal3Array = PrimitiveArray<Int64Decimal3Type>;
144145
pub type Int64Decimal4Array = PrimitiveArray<Int64Decimal4Type>;
145146
pub type Int64Decimal5Array = PrimitiveArray<Int64Decimal5Type>;
146147
pub type Int64Decimal10Array = PrimitiveArray<Int64Decimal10Type>;
148+
pub type Int96Decimal0Array = PrimitiveArray<Int96Decimal0Type>;
149+
pub type Int96Decimal1Array = PrimitiveArray<Int96Decimal1Type>;
150+
pub type Int96Decimal2Array = PrimitiveArray<Int96Decimal2Type>;
151+
pub type Int96Decimal3Array = PrimitiveArray<Int96Decimal3Type>;
152+
pub type Int96Decimal4Array = PrimitiveArray<Int96Decimal4Type>;
153+
pub type Int96Decimal5Array = PrimitiveArray<Int96Decimal5Type>;
154+
pub type Int96Decimal10Array = PrimitiveArray<Int96Decimal10Type>;
147155
pub type UInt8Array = PrimitiveArray<UInt8Type>;
148156
pub type UInt16Array = PrimitiveArray<UInt16Type>;
149157
pub type UInt32Array = PrimitiveArray<UInt32Type>;
@@ -200,6 +208,13 @@ pub type Int64Decimal3BufferBuilder = BufferBuilder<Int64Decimal3Type>;
200208
pub type Int64Decimal4BufferBuilder = BufferBuilder<Int64Decimal4Type>;
201209
pub type Int64Decimal5BufferBuilder = BufferBuilder<Int64Decimal5Type>;
202210
pub type Int64Decimal10BufferBuilder = BufferBuilder<Int64Decimal10Type>;
211+
pub type Int96Decimal0BufferBuilder = BufferBuilder<Int96Decimal0Type>;
212+
pub type Int96Decimal1BufferBuilder = BufferBuilder<Int96Decimal1Type>;
213+
pub type Int96Decimal2BufferBuilder = BufferBuilder<Int96Decimal2Type>;
214+
pub type Int96Decimal3BufferBuilder = BufferBuilder<Int96Decimal3Type>;
215+
pub type Int96Decimal4BufferBuilder = BufferBuilder<Int96Decimal4Type>;
216+
pub type Int96Decimal5BufferBuilder = BufferBuilder<Int96Decimal5Type>;
217+
pub type Int96Decimal10BufferBuilder = BufferBuilder<Int96Decimal10Type>;
203218
pub type UInt8BufferBuilder = BufferBuilder<u8>;
204219
pub type UInt16BufferBuilder = BufferBuilder<u16>;
205220
pub type UInt32BufferBuilder = BufferBuilder<u32>;
@@ -246,13 +261,21 @@ pub type Int8Builder = PrimitiveBuilder<Int8Type>;
246261
pub type Int16Builder = PrimitiveBuilder<Int16Type>;
247262
pub type Int32Builder = PrimitiveBuilder<Int32Type>;
248263
pub type Int64Builder = PrimitiveBuilder<Int64Type>;
264+
pub type Int96Builder = PrimitiveBuilder<Int96Type>;
249265
pub type Int64Decimal0Builder = PrimitiveBuilder<Int64Decimal0Type>;
250266
pub type Int64Decimal1Builder = PrimitiveBuilder<Int64Decimal1Type>;
251267
pub type Int64Decimal2Builder = PrimitiveBuilder<Int64Decimal2Type>;
252268
pub type Int64Decimal3Builder = PrimitiveBuilder<Int64Decimal3Type>;
253269
pub type Int64Decimal4Builder = PrimitiveBuilder<Int64Decimal4Type>;
254270
pub type Int64Decimal5Builder = PrimitiveBuilder<Int64Decimal5Type>;
255271
pub type Int64Decimal10Builder = PrimitiveBuilder<Int64Decimal10Type>;
272+
pub type Int96Decimal0Builder = PrimitiveBuilder<Int96Decimal0Type>;
273+
pub type Int96Decimal1Builder = PrimitiveBuilder<Int96Decimal1Type>;
274+
pub type Int96Decimal2Builder = PrimitiveBuilder<Int96Decimal2Type>;
275+
pub type Int96Decimal3Builder = PrimitiveBuilder<Int96Decimal3Type>;
276+
pub type Int96Decimal4Builder = PrimitiveBuilder<Int96Decimal4Type>;
277+
pub type Int96Decimal5Builder = PrimitiveBuilder<Int96Decimal5Type>;
278+
pub type Int96Decimal10Builder = PrimitiveBuilder<Int96Decimal10Type>;
256279
pub type UInt8Builder = PrimitiveBuilder<UInt8Type>;
257280
pub type UInt16Builder = PrimitiveBuilder<UInt16Type>;
258281
pub type UInt32Builder = PrimitiveBuilder<UInt32Type>;

‎arrow/src/array/transform/mod.rs

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,9 @@ fn build_extend(array: &ArrayData) -> Extend {
245245
DataType::Int16 => primitive::build_extend::<i16>(array),
246246
DataType::Int32 => primitive::build_extend::<i32>(array),
247247
DataType::Int64 => primitive::build_extend::<i64>(array),
248+
DataType::Int96 => primitive::build_extend::<i128>(array),
248249
DataType::Int64Decimal(_) => primitive::build_extend::<i64>(array),
250+
DataType::Int96Decimal(_) => primitive::build_extend::<i128>(array),
249251
DataType::Float32 => primitive::build_extend::<f32>(array),
250252
DataType::Float64 => primitive::build_extend::<f64>(array),
251253
DataType::Date32
@@ -291,7 +293,9 @@ fn build_extend_nulls(data_type: &DataType) -> ExtendNulls {
291293
DataType::Int16 => primitive::extend_nulls::<i16>,
292294
DataType::Int32 => primitive::extend_nulls::<i32>,
293295
DataType::Int64 => primitive::extend_nulls::<i64>,
296+
DataType::Int96 => primitive::extend_nulls::<i128>,
294297
DataType::Int64Decimal(_) => primitive::extend_nulls::<i64>,
298+
DataType::Int96Decimal(_) => primitive::extend_nulls::<i128>,
295299
DataType::Float32 => primitive::extend_nulls::<f32>,
296300
DataType::Float64 => primitive::extend_nulls::<f64>,
297301
DataType::Date32
@@ -434,7 +438,9 @@ impl<'a> MutableArrayData<'a> {
434438
| DataType::Int16
435439
| DataType::Int32
436440
| DataType::Int64
441+
| DataType::Int96
437442
| DataType::Int64Decimal(_)
443+
| DataType::Int96Decimal(_)
438444
| DataType::Float32
439445
| DataType::Float64
440446
| DataType::Date32
@@ -616,7 +622,7 @@ impl<'a> MutableArrayData<'a> {
616622
pub fn extend_nulls(&mut self, len: usize) {
617623
self.data.null_count += len;
618624

619-
let null_bytes_count = bit_util::ceil(self.data.len + len, 8);
625+
let null_bytes_count = bit_util::ceil(self.data.len + len, 8);
620626
if null_bytes_count > self.data.null_buffer.len() {
621627
self.data.null_buffer.resize(null_bytes_count, 0x00);
622628
}
@@ -744,7 +750,17 @@ mod tests {
744750
let result = a.freeze();
745751
let array = UInt8Array::from(result);
746752
assert_eq!(array.data().null_buffer().unwrap().len(), 2);
747-
let expected = UInt8Array::from(vec![Some(1), Some(2), Some(3), None, None, None, None, None, None]);
753+
let expected = UInt8Array::from(vec![
754+
Some(1),
755+
Some(2),
756+
Some(3),
757+
None,
758+
None,
759+
None,
760+
None,
761+
None,
762+
None,
763+
]);
748764
assert_eq!(array, expected);
749765

750766
let b = UInt8Array::from(vec![Some(1), Some(2), Some(3)]);
@@ -759,12 +775,23 @@ mod tests {
759775
let result = a.freeze();
760776
let array = UInt8Array::from(result);
761777
assert_eq!(array.data().null_buffer().unwrap().len(), 2);
762-
let expected = UInt8Array::from(vec![Some(1), Some(2), Some(3), None, None, None, None, None, None, Some(1), Some(2), Some(3)]);
778+
let expected = UInt8Array::from(vec![
779+
Some(1),
780+
Some(2),
781+
Some(3),
782+
None,
783+
None,
784+
None,
785+
None,
786+
None,
787+
None,
788+
Some(1),
789+
Some(2),
790+
Some(3),
791+
]);
763792
assert_eq!(array, expected);
764-
765793
}
766794

767-
768795
#[test]
769796
fn test_list_null_offset() -> Result<()> {
770797
let int_builder = Int64Builder::new(24);

0 commit comments

Comments
 (0)