arrow_string/
length.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Defines kernel for length of string arrays and binary arrays
19
20use arrow_array::*;
21use arrow_array::{cast::AsArray, types::*};
22use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer};
23use arrow_schema::{ArrowError, DataType};
24use std::sync::Arc;
25
26fn length_impl<P: ArrowPrimitiveType>(
27    offsets: &OffsetBuffer<P::Native>,
28    nulls: Option<&NullBuffer>,
29) -> ArrayRef {
30    let v: Vec<_> = offsets
31        .windows(2)
32        .map(|w| w[1].sub_wrapping(w[0]))
33        .collect();
34    Arc::new(PrimitiveArray::<P>::new(v.into(), nulls.cloned()))
35}
36
37fn bit_length_impl<P: ArrowPrimitiveType>(
38    offsets: &OffsetBuffer<P::Native>,
39    nulls: Option<&NullBuffer>,
40) -> ArrayRef {
41    let bits = P::Native::usize_as(8);
42    let c = |w: &[P::Native]| w[1].sub_wrapping(w[0]).mul_wrapping(bits);
43    let v: Vec<_> = offsets.windows(2).map(c).collect();
44    Arc::new(PrimitiveArray::<P>::new(v.into(), nulls.cloned()))
45}
46
47/// Returns an array of Int32/Int64 denoting the length of each value in the array.
48///
49/// For list array, length is the number of elements in each list.
50/// For string array and binary array, length is the number of bytes of each value.
51///
52/// * this only accepts ListArray/LargeListArray, StringArray/LargeStringArray/StringViewArray, BinaryArray/LargeBinaryArray, and FixedSizeListArray,
53///   or DictionaryArray with above Arrays as values
54/// * length of null is null.
55pub fn length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
56    if let Some(d) = array.as_any_dictionary_opt() {
57        let lengths = length(d.values().as_ref())?;
58        return Ok(d.with_values(lengths));
59    }
60
61    match array.data_type() {
62        DataType::List(_) => {
63            let list = array.as_list::<i32>();
64            Ok(length_impl::<Int32Type>(list.offsets(), list.nulls()))
65        }
66        DataType::LargeList(_) => {
67            let list = array.as_list::<i64>();
68            Ok(length_impl::<Int64Type>(list.offsets(), list.nulls()))
69        }
70        DataType::Utf8 => {
71            let list = array.as_string::<i32>();
72            Ok(length_impl::<Int32Type>(list.offsets(), list.nulls()))
73        }
74        DataType::LargeUtf8 => {
75            let list = array.as_string::<i64>();
76            Ok(length_impl::<Int64Type>(list.offsets(), list.nulls()))
77        }
78        DataType::Utf8View => {
79            let list = array.as_string_view();
80            let v = list.views().iter().map(|v| *v as i32).collect::<Vec<_>>();
81            Ok(Arc::new(PrimitiveArray::<Int32Type>::try_new(
82                v.into(),
83                list.nulls().cloned(),
84            )?))
85        }
86        DataType::Binary => {
87            let list = array.as_binary::<i32>();
88            Ok(length_impl::<Int32Type>(list.offsets(), list.nulls()))
89        }
90        DataType::LargeBinary => {
91            let list = array.as_binary::<i64>();
92            Ok(length_impl::<Int64Type>(list.offsets(), list.nulls()))
93        }
94        DataType::FixedSizeBinary(len) | DataType::FixedSizeList(_, len) => Ok(Arc::new(
95            Int32Array::try_new(vec![*len; array.len()].into(), array.nulls().cloned())?,
96        )),
97        DataType::BinaryView => {
98            let list = array.as_binary_view();
99            let v = list.views().iter().map(|v| *v as i32).collect::<Vec<_>>();
100            Ok(Arc::new(PrimitiveArray::<Int32Type>::try_new(
101                v.into(),
102                list.nulls().cloned(),
103            )?))
104        }
105        other => Err(ArrowError::ComputeError(format!(
106            "length not supported for {other:?}"
107        ))),
108    }
109}
110
111/// Returns an array of Int32/Int64 denoting the number of bits in each value in the array.
112///
113/// * this only accepts StringArray/Utf8, LargeString/LargeUtf8, BinaryArray and LargeBinaryArray,
114///   or DictionaryArray with above Arrays as values
115/// * bit_length of null is null.
116/// * bit_length is in number of bits
117pub fn bit_length(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
118    if let Some(d) = array.as_any_dictionary_opt() {
119        let lengths = bit_length(d.values().as_ref())?;
120        return Ok(d.with_values(lengths));
121    }
122
123    match array.data_type() {
124        DataType::List(_) => {
125            let list = array.as_list::<i32>();
126            Ok(bit_length_impl::<Int32Type>(list.offsets(), list.nulls()))
127        }
128        DataType::LargeList(_) => {
129            let list = array.as_list::<i64>();
130            Ok(bit_length_impl::<Int64Type>(list.offsets(), list.nulls()))
131        }
132        DataType::Utf8 => {
133            let list = array.as_string::<i32>();
134            Ok(bit_length_impl::<Int32Type>(list.offsets(), list.nulls()))
135        }
136        DataType::LargeUtf8 => {
137            let list = array.as_string::<i64>();
138            Ok(bit_length_impl::<Int64Type>(list.offsets(), list.nulls()))
139        }
140        DataType::Utf8View => {
141            let list = array.as_string_view();
142            let values = list
143                .views()
144                .iter()
145                .map(|view| (*view as i32).wrapping_mul(8))
146                .collect();
147            Ok(Arc::new(Int32Array::try_new(
148                values,
149                array.nulls().cloned(),
150            )?))
151        }
152        DataType::Binary => {
153            let list = array.as_binary::<i32>();
154            Ok(bit_length_impl::<Int32Type>(list.offsets(), list.nulls()))
155        }
156        DataType::LargeBinary => {
157            let list = array.as_binary::<i64>();
158            Ok(bit_length_impl::<Int64Type>(list.offsets(), list.nulls()))
159        }
160        DataType::FixedSizeBinary(len) => Ok(Arc::new(Int32Array::try_new(
161            vec![*len * 8; array.len()].into(),
162            array.nulls().cloned(),
163        )?)),
164        other => Err(ArrowError::ComputeError(format!(
165            "bit_length not supported for {other:?}"
166        ))),
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use arrow_buffer::Buffer;
174    use arrow_data::ArrayData;
175    use arrow_schema::Field;
176
177    fn length_cases_string() -> Vec<(Vec<&'static str>, usize, Vec<i32>)> {
178        // a large array
179        let values = [
180            "one",
181            "on",
182            "o",
183            "",
184            "this is a longer string to test string array with",
185        ];
186        let values = values.into_iter().cycle().take(4096).collect();
187        let expected = [3, 2, 1, 0, 49].into_iter().cycle().take(4096).collect();
188
189        vec![
190            (vec!["hello", " ", "world"], 3, vec![5, 1, 5]),
191            (vec!["hello", " ", "world", "!"], 4, vec![5, 1, 5, 1]),
192            (vec!["💖"], 1, vec![4]),
193            (values, 4096, expected),
194        ]
195    }
196
197    macro_rules! length_binary_helper {
198        ($offset_ty: ty, $result_ty: ty, $kernel: ident, $value: expr, $expected: expr) => {{
199            let array = GenericBinaryArray::<$offset_ty>::from($value);
200            let result = $kernel(&array).unwrap();
201            let result = result.as_any().downcast_ref::<$result_ty>().unwrap();
202            let expected: $result_ty = $expected.into();
203            assert_eq!(&expected, result);
204        }};
205    }
206
207    macro_rules! length_list_helper {
208        ($offset_ty: ty, $result_ty: ty, $element_ty: ty, $value: expr, $expected: expr) => {{
209            let array =
210                GenericListArray::<$offset_ty>::from_iter_primitive::<$element_ty, _, _>($value);
211            let result = length(&array).unwrap();
212            let result = result.as_any().downcast_ref::<$result_ty>().unwrap();
213            let expected: $result_ty = $expected.into();
214            assert_eq!(&expected, result);
215        }};
216    }
217
218    #[test]
219    fn length_test_string() {
220        length_cases_string()
221            .into_iter()
222            .for_each(|(input, len, expected)| {
223                let array = StringArray::from(input);
224                let result = length(&array).unwrap();
225                assert_eq!(len, result.len());
226                let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
227                expected.iter().enumerate().for_each(|(i, value)| {
228                    assert_eq!(*value, result.value(i));
229                });
230            })
231    }
232
233    #[test]
234    fn length_test_large_string() {
235        length_cases_string()
236            .into_iter()
237            .for_each(|(input, len, expected)| {
238                let array = LargeStringArray::from(input);
239                let result = length(&array).unwrap();
240                assert_eq!(len, result.len());
241                let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
242                expected.iter().enumerate().for_each(|(i, value)| {
243                    assert_eq!(*value as i64, result.value(i));
244                });
245            })
246    }
247
248    #[test]
249    fn length_test_string_view() {
250        length_cases_string()
251            .into_iter()
252            .for_each(|(input, len, expected)| {
253                let array = StringViewArray::from(input);
254                let result = length(&array).unwrap();
255                assert_eq!(len, result.len());
256                let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
257                expected.iter().enumerate().for_each(|(i, value)| {
258                    assert_eq!(*value, result.value(i));
259                });
260            })
261    }
262
263    #[test]
264    fn length_test_binary() {
265        let value: Vec<&[u8]> = vec![b"zero", b"one", &[0xff, 0xf8]];
266        let result: Vec<i32> = vec![4, 3, 2];
267        length_binary_helper!(i32, Int32Array, length, value, result)
268    }
269
270    #[test]
271    fn length_test_large_binary() {
272        let value: Vec<&[u8]> = vec![b"zero", &[0xff, 0xf8], b"two"];
273        let result: Vec<i64> = vec![4, 2, 3];
274        length_binary_helper!(i64, Int64Array, length, value, result)
275    }
276
277    #[test]
278    fn length_test_binary_view() {
279        let value: Vec<&[u8]> = vec![
280            b"zero",
281            &[0xff, 0xf8],
282            b"two",
283            b"this is a longer string to test binary array with",
284        ];
285        let expected: Vec<i32> = vec![4, 2, 3, 49];
286
287        let array = BinaryViewArray::from(value);
288        let result = length(&array).unwrap();
289        let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
290        let expected: Int32Array = expected.into();
291        assert_eq!(&expected, result);
292    }
293
294    #[test]
295    fn length_test_list() {
296        let value = vec![
297            Some(vec![]),
298            Some(vec![Some(1), Some(2), Some(4)]),
299            Some(vec![Some(0)]),
300        ];
301        let result: Vec<i32> = vec![0, 3, 1];
302        length_list_helper!(i32, Int32Array, Int32Type, value, result)
303    }
304
305    #[test]
306    fn length_test_large_list() {
307        let value = vec![
308            Some(vec![]),
309            Some(vec![Some(1.1), Some(2.2), Some(3.3)]),
310            Some(vec![None]),
311        ];
312        let result: Vec<i64> = vec![0, 3, 1];
313        length_list_helper!(i64, Int64Array, Float32Type, value, result)
314    }
315
316    type OptionStr = Option<&'static str>;
317
318    fn length_null_cases_string() -> Vec<(Vec<OptionStr>, usize, Vec<Option<i32>>)> {
319        vec![(
320            vec![Some("one"), None, Some("three"), Some("four")],
321            4,
322            vec![Some(3), None, Some(5), Some(4)],
323        )]
324    }
325
326    #[test]
327    fn length_null_string() {
328        length_null_cases_string()
329            .into_iter()
330            .for_each(|(input, len, expected)| {
331                let array = StringArray::from(input);
332                let result = length(&array).unwrap();
333                assert_eq!(len, result.len());
334                let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
335
336                let expected: Int32Array = expected.into();
337                assert_eq!(&expected, result);
338            })
339    }
340
341    #[test]
342    fn length_null_large_string() {
343        length_null_cases_string()
344            .into_iter()
345            .for_each(|(input, len, expected)| {
346                let array = LargeStringArray::from(input);
347                let result = length(&array).unwrap();
348                assert_eq!(len, result.len());
349                let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
350
351                // convert to i64
352                let expected: Int64Array = expected
353                    .iter()
354                    .map(|e| e.map(|e| e as i64))
355                    .collect::<Vec<_>>()
356                    .into();
357                assert_eq!(&expected, result);
358            })
359    }
360
361    #[test]
362    fn length_null_binary() {
363        let value: Vec<Option<&[u8]>> =
364            vec![Some(b"zero"), None, Some(&[0xff, 0xf8]), Some(b"three")];
365        let result: Vec<Option<i32>> = vec![Some(4), None, Some(2), Some(5)];
366        length_binary_helper!(i32, Int32Array, length, value, result)
367    }
368
369    #[test]
370    fn length_null_large_binary() {
371        let value: Vec<Option<&[u8]>> =
372            vec![Some(&[0xff, 0xf8]), None, Some(b"two"), Some(b"three")];
373        let result: Vec<Option<i64>> = vec![Some(2), None, Some(3), Some(5)];
374        length_binary_helper!(i64, Int64Array, length, value, result)
375    }
376
377    #[test]
378    fn length_null_list() {
379        let value = vec![
380            Some(vec![]),
381            None,
382            Some(vec![Some(1), None, Some(2), Some(4)]),
383            Some(vec![Some(0)]),
384        ];
385        let result: Vec<Option<i32>> = vec![Some(0), None, Some(4), Some(1)];
386        length_list_helper!(i32, Int32Array, Int8Type, value, result)
387    }
388
389    #[test]
390    fn length_null_large_list() {
391        let value = vec![
392            Some(vec![]),
393            None,
394            Some(vec![Some(1.1), None, Some(4.0)]),
395            Some(vec![Some(0.1)]),
396        ];
397        let result: Vec<Option<i64>> = vec![Some(0), None, Some(3), Some(1)];
398        length_list_helper!(i64, Int64Array, Float32Type, value, result)
399    }
400
401    /// Tests that length is not valid for u64.
402    #[test]
403    fn length_wrong_type() {
404        let array: UInt64Array = vec![1u64].into();
405
406        assert!(length(&array).is_err());
407    }
408
409    /// Tests with an offset
410    #[test]
411    fn length_offsets_string() {
412        let a = StringArray::from(vec![Some("hello"), Some(" "), Some("world"), None]);
413        let b = a.slice(1, 3);
414        let result = length(&b).unwrap();
415        let result: &Int32Array = result.as_primitive();
416
417        let expected = Int32Array::from(vec![Some(1), Some(5), None]);
418        assert_eq!(&expected, result);
419    }
420
421    #[test]
422    fn length_offsets_binary() {
423        let value: Vec<Option<&[u8]>> = vec![Some(b"hello"), Some(b" "), Some(&[0xff, 0xf8]), None];
424        let a = BinaryArray::from(value);
425        let b = a.slice(1, 3);
426        let result = length(&b).unwrap();
427        let result: &Int32Array = result.as_primitive();
428
429        let expected = Int32Array::from(vec![Some(1), Some(2), None]);
430        assert_eq!(&expected, result);
431    }
432
433    fn bit_length_cases() -> Vec<(Vec<&'static str>, usize, Vec<i32>)> {
434        // a large array
435        let values = ["one", "on", "o", ""];
436        let values = values.into_iter().cycle().take(4096).collect();
437        let expected = [24, 16, 8, 0].into_iter().cycle().take(4096).collect();
438
439        vec![
440            (vec!["hello", " ", "world", "!"], 4, vec![40, 8, 40, 8]),
441            (vec!["💖"], 1, vec![32]),
442            (vec!["josé"], 1, vec![40]),
443            (values, 4096, expected),
444        ]
445    }
446
447    #[test]
448    fn bit_length_test_string() {
449        bit_length_cases()
450            .into_iter()
451            .for_each(|(input, len, expected)| {
452                let array = StringArray::from(input);
453                let result = bit_length(&array).unwrap();
454                assert_eq!(len, result.len());
455                let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
456                expected.iter().enumerate().for_each(|(i, value)| {
457                    assert_eq!(*value, result.value(i));
458                });
459            })
460    }
461
462    #[test]
463    fn bit_length_test_large_string() {
464        bit_length_cases()
465            .into_iter()
466            .for_each(|(input, len, expected)| {
467                let array = LargeStringArray::from(input);
468                let result = bit_length(&array).unwrap();
469                assert_eq!(len, result.len());
470                let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
471                expected.iter().enumerate().for_each(|(i, value)| {
472                    assert_eq!(*value as i64, result.value(i));
473                });
474            })
475    }
476
477    #[test]
478    fn bit_length_test_utf8view() {
479        bit_length_cases()
480            .into_iter()
481            .for_each(|(input, len, expected)| {
482                let string_array = StringViewArray::from(input);
483                let result = bit_length(&string_array).unwrap();
484                assert_eq!(len, result.len());
485                let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
486                expected.iter().enumerate().for_each(|(i, value)| {
487                    assert_eq!(*value, result.value(i));
488                });
489            })
490    }
491
492    #[test]
493    fn bit_length_null_utf8view() {
494        bit_length_null_cases()
495            .into_iter()
496            .for_each(|(input, len, expected)| {
497                let array = StringArray::from(input);
498                let result = bit_length(&array).unwrap();
499                assert_eq!(len, result.len());
500                let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
501
502                let expected: Int32Array = expected.into();
503                assert_eq!(&expected, result);
504            })
505    }
506    #[test]
507    fn bit_length_binary() {
508        let value: Vec<&[u8]> = vec![b"one", &[0xff, 0xf8], b"three"];
509        let expected: Vec<i32> = vec![24, 16, 40];
510        length_binary_helper!(i32, Int32Array, bit_length, value, expected)
511    }
512
513    #[test]
514    fn bit_length_large_binary() {
515        let value: Vec<&[u8]> = vec![b"zero", b" ", &[0xff, 0xf8]];
516        let expected: Vec<i64> = vec![32, 8, 16];
517        length_binary_helper!(i64, Int64Array, bit_length, value, expected)
518    }
519
520    fn bit_length_null_cases() -> Vec<(Vec<OptionStr>, usize, Vec<Option<i32>>)> {
521        vec![(
522            vec![Some("one"), None, Some("three"), Some("four")],
523            4,
524            vec![Some(24), None, Some(40), Some(32)],
525        )]
526    }
527
528    #[test]
529    fn bit_length_null_string() {
530        bit_length_null_cases()
531            .into_iter()
532            .for_each(|(input, len, expected)| {
533                let array = StringArray::from(input);
534                let result = bit_length(&array).unwrap();
535                assert_eq!(len, result.len());
536                let result = result.as_any().downcast_ref::<Int32Array>().unwrap();
537
538                let expected: Int32Array = expected.into();
539                assert_eq!(&expected, result);
540            })
541    }
542
543    #[test]
544    fn bit_length_null_large_string() {
545        bit_length_null_cases()
546            .into_iter()
547            .for_each(|(input, len, expected)| {
548                let array = LargeStringArray::from(input);
549                let result = bit_length(&array).unwrap();
550                assert_eq!(len, result.len());
551                let result = result.as_any().downcast_ref::<Int64Array>().unwrap();
552
553                // convert to i64
554                let expected: Int64Array = expected
555                    .iter()
556                    .map(|e| e.map(|e| e as i64))
557                    .collect::<Vec<_>>()
558                    .into();
559                assert_eq!(&expected, result);
560            })
561    }
562
563    #[test]
564    fn bit_length_null_binary() {
565        let value: Vec<Option<&[u8]>> =
566            vec![Some(b"one"), None, Some(b"three"), Some(&[0xff, 0xf8])];
567        let expected: Vec<Option<i32>> = vec![Some(24), None, Some(40), Some(16)];
568        length_binary_helper!(i32, Int32Array, bit_length, value, expected)
569    }
570
571    #[test]
572    fn bit_length_null_large_binary() {
573        let value: Vec<Option<&[u8]>> =
574            vec![Some(b"one"), None, Some(&[0xff, 0xf8]), Some(b"four")];
575        let expected: Vec<Option<i64>> = vec![Some(24), None, Some(16), Some(32)];
576        length_binary_helper!(i64, Int64Array, bit_length, value, expected)
577    }
578
579    /// Tests that bit_length is not valid for u64.
580    #[test]
581    fn bit_length_wrong_type() {
582        let array: UInt64Array = vec![1u64].into();
583
584        assert!(bit_length(&array).is_err());
585    }
586
587    /// Tests with an offset
588    #[test]
589    fn bit_length_offsets_string() {
590        let a = StringArray::from(vec![Some("hello"), Some(" "), Some("world"), None]);
591        let b = a.slice(1, 3);
592        let result = bit_length(&b).unwrap();
593        let result: &Int32Array = result.as_primitive();
594
595        let expected = Int32Array::from(vec![Some(8), Some(40), None]);
596        assert_eq!(&expected, result);
597    }
598
599    #[test]
600    fn bit_length_offsets_binary() {
601        let value: Vec<Option<&[u8]>> = vec![Some(b"hello"), Some(&[]), Some(b"world"), None];
602        let a = BinaryArray::from(value);
603        let b = a.slice(1, 3);
604        let result = bit_length(&b).unwrap();
605        let result: &Int32Array = result.as_primitive();
606
607        let expected = Int32Array::from(vec![Some(0), Some(40), None]);
608        assert_eq!(&expected, result);
609    }
610
611    #[test]
612    fn length_dictionary() {
613        _length_dictionary::<Int8Type>();
614        _length_dictionary::<Int16Type>();
615        _length_dictionary::<Int32Type>();
616        _length_dictionary::<Int64Type>();
617        _length_dictionary::<UInt8Type>();
618        _length_dictionary::<UInt16Type>();
619        _length_dictionary::<UInt32Type>();
620        _length_dictionary::<UInt64Type>();
621    }
622
623    fn _length_dictionary<K: ArrowDictionaryKeyType>() {
624        const TOTAL: i32 = 100;
625
626        let v = ["aaaa", "bb", "ccccc", "ddd", "eeeeee"];
627        let data: Vec<Option<&str>> = (0..TOTAL)
628            .map(|n| {
629                let i = n % 5;
630                if i == 3 { None } else { Some(v[i as usize]) }
631            })
632            .collect();
633
634        let dict_array: DictionaryArray<K> = data.clone().into_iter().collect();
635
636        let expected: Vec<Option<i32>> =
637            data.iter().map(|opt| opt.map(|s| s.len() as i32)).collect();
638
639        let res = length(&dict_array).unwrap();
640        let actual = res.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
641        let actual: Vec<Option<i32>> = actual
642            .values()
643            .as_any()
644            .downcast_ref::<Int32Array>()
645            .unwrap()
646            .take_iter(dict_array.keys_iter())
647            .collect();
648
649        for i in 0..TOTAL as usize {
650            assert_eq!(expected[i], actual[i],);
651        }
652    }
653
654    #[test]
655    fn bit_length_dictionary() {
656        _bit_length_dictionary::<Int8Type>();
657        _bit_length_dictionary::<Int16Type>();
658        _bit_length_dictionary::<Int32Type>();
659        _bit_length_dictionary::<Int64Type>();
660        _bit_length_dictionary::<UInt8Type>();
661        _bit_length_dictionary::<UInt16Type>();
662        _bit_length_dictionary::<UInt32Type>();
663        _bit_length_dictionary::<UInt64Type>();
664    }
665
666    fn _bit_length_dictionary<K: ArrowDictionaryKeyType>() {
667        const TOTAL: i32 = 100;
668
669        let v = ["aaaa", "bb", "ccccc", "ddd", "eeeeee"];
670        let data: Vec<Option<&str>> = (0..TOTAL)
671            .map(|n| {
672                let i = n % 5;
673                if i == 3 { None } else { Some(v[i as usize]) }
674            })
675            .collect();
676
677        let dict_array: DictionaryArray<K> = data.clone().into_iter().collect();
678
679        let expected: Vec<Option<i32>> = data
680            .iter()
681            .map(|opt| opt.map(|s| (s.chars().count() * 8) as i32))
682            .collect();
683
684        let res = bit_length(&dict_array).unwrap();
685        let actual = res.as_any().downcast_ref::<DictionaryArray<K>>().unwrap();
686        let actual: Vec<Option<i32>> = actual
687            .values()
688            .as_any()
689            .downcast_ref::<Int32Array>()
690            .unwrap()
691            .take_iter(dict_array.keys_iter())
692            .collect();
693
694        for i in 0..TOTAL as usize {
695            assert_eq!(expected[i], actual[i],);
696        }
697    }
698
699    #[test]
700    fn test_fixed_size_list_length() {
701        // Construct a value array
702        let value_data = ArrayData::builder(DataType::Int32)
703            .len(9)
704            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8]))
705            .build()
706            .unwrap();
707        let list_data_type =
708            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
709        let nulls = NullBuffer::from(vec![true, false, true]);
710        let list_data = ArrayData::builder(list_data_type)
711            .len(3)
712            .add_child_data(value_data)
713            .nulls(Some(nulls))
714            .build()
715            .unwrap();
716        let list_array = FixedSizeListArray::from(list_data);
717
718        let lengths = length(&list_array).unwrap();
719        let lengths = lengths.as_primitive::<Int32Type>();
720
721        assert_eq!(lengths.len(), 3);
722        assert_eq!(lengths.value(0), 3);
723        assert!(lengths.is_null(1));
724        assert_eq!(lengths.value(2), 3);
725    }
726
727    #[test]
728    fn test_fixed_size_binary() {
729        let array = FixedSizeBinaryArray::new(4, [0; 16].into(), None);
730        let result = length(&array).unwrap();
731        assert_eq!(result.as_ref(), &Int32Array::from(vec![4; 4]));
732
733        let result = bit_length(&array).unwrap();
734        assert_eq!(result.as_ref(), &Int32Array::from(vec![32; 4]));
735    }
736}
arrow_string/length.rs

arrow_string/
length.rs
