1#![doc(
164 html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg",
165 html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg"
166)]
167#![cfg_attr(docsrs, feature(doc_cfg))]
168#![warn(missing_docs)]
169use std::cmp::Ordering;
170use std::hash::{Hash, Hasher};
171use std::iter::Map;
172use std::slice::Windows;
173use std::sync::Arc;
174
175use arrow_array::cast::*;
176use arrow_array::types::{ArrowDictionaryKeyType, ByteArrayType, ByteViewType};
177use arrow_array::*;
178use arrow_buffer::{ArrowNativeType, Buffer, OffsetBuffer, ScalarBuffer};
179use arrow_schema::*;
180use variable::{decode_binary_view, decode_string_view};
181
182use crate::fixed::{decode_bool, decode_fixed_size_binary, decode_primitive};
183use crate::list::{compute_lengths_fixed_size_list, encode_fixed_size_list};
184use crate::variable::{decode_binary, decode_string};
185use arrow_array::types::{Int16Type, Int32Type, Int64Type};
186
187mod fixed;
188mod list;
189mod run;
190mod variable;
191
192#[derive(Debug)]
575pub struct RowConverter {
576 fields: Arc<[SortField]>,
577 codecs: Vec<Codec>,
579}
580
581#[derive(Debug)]
582enum Codec {
583 Stateless,
585 Dictionary(RowConverter, OwnedRow),
588 Struct(RowConverter, OwnedRow),
591 List(RowConverter),
593 Map(RowConverter),
595 RunEndEncoded(RowConverter),
597 Union(Vec<RowConverter>, Vec<i8>, Vec<OwnedRow>),
600}
601
602fn compute_list_view_bounds<O: OffsetSizeTrait>(array: &GenericListViewArray<O>) -> (usize, usize) {
605 if array.is_empty() {
606 return (0, 0);
607 }
608
609 let offsets = array.value_offsets();
610 let sizes = array.value_sizes();
611 let values_len = array.values().len();
612
613 let mut min_offset = usize::MAX;
614 let mut max_end = 0usize;
615
616 for i in 0..array.len() {
617 let offset = offsets[i].as_usize();
618 let size = sizes[i].as_usize();
619 let end = offset + size;
620
621 if size > 0 {
622 min_offset = min_offset.min(offset);
623 max_end = max_end.max(end);
624 }
625
626 if min_offset == 0 && max_end == values_len {
630 break;
631 }
632 }
633
634 if min_offset == usize::MAX {
635 (0, 0)
637 } else {
638 (min_offset, max_end)
639 }
640}
641
642impl Codec {
643 fn new(sort_field: &SortField) -> Result<Self, ArrowError> {
644 match &sort_field.data_type {
645 DataType::Dictionary(_, values) => {
646 let sort_field =
647 SortField::new_with_options(values.as_ref().clone(), sort_field.options);
648
649 let converter = RowConverter::new(vec![sort_field])?;
650 let null_array = new_null_array(values.as_ref(), 1);
651 let nulls = converter.convert_columns(&[null_array])?;
652
653 let owned = OwnedRow {
654 data: nulls.buffer.into(),
655 config: nulls.config,
656 };
657 Ok(Self::Dictionary(converter, owned))
658 }
659 DataType::RunEndEncoded(_, values) => {
660 let options = SortOptions {
662 descending: false,
663 nulls_first: sort_field.options.nulls_first != sort_field.options.descending,
664 };
665
666 let field = SortField::new_with_options(values.data_type().clone(), options);
667 let converter = RowConverter::new(vec![field])?;
668 Ok(Self::RunEndEncoded(converter))
669 }
670 d if !d.is_nested() => Ok(Self::Stateless),
671 DataType::List(f)
672 | DataType::LargeList(f)
673 | DataType::ListView(f)
674 | DataType::LargeListView(f) => {
675 let options = SortOptions {
679 descending: false,
680 nulls_first: sort_field.options.nulls_first != sort_field.options.descending,
681 };
682
683 let field = SortField::new_with_options(f.data_type().clone(), options);
684 let converter = RowConverter::new(vec![field])?;
685 Ok(Self::List(converter))
686 }
687 DataType::Map(f, _) => {
688 let options = SortOptions {
692 descending: false,
693 nulls_first: sort_field.options.nulls_first != sort_field.options.descending,
694 };
695
696 let DataType::Struct(fields) = f.data_type() else {
697 return Err(ArrowError::InvalidArgumentError(format!(
698 "expected struct field in map, got {:?}",
699 f.data_type()
700 )));
701 };
702
703 let fields = fields
705 .iter()
706 .map(|struct_field| {
707 SortField::new_with_options(struct_field.data_type().clone(), options)
708 })
709 .collect::<Vec<_>>();
710 assert_eq!(fields.len(), 2);
711 let converter = RowConverter::new(fields)?;
712 Ok(Self::Map(converter))
713 }
714 DataType::FixedSizeList(f, _) => {
715 let field = SortField::new_with_options(f.data_type().clone(), sort_field.options);
716 let converter = RowConverter::new(vec![field])?;
717 Ok(Self::List(converter))
718 }
719 DataType::Struct(f) => {
720 let sort_fields = f
721 .iter()
722 .map(|x| SortField::new_with_options(x.data_type().clone(), sort_field.options))
723 .collect();
724
725 let converter = RowConverter::new(sort_fields)?;
726 let nulls: Vec<_> = f.iter().map(|x| new_null_array(x.data_type(), 1)).collect();
727
728 let nulls = converter.convert_columns(&nulls)?;
729 let owned = OwnedRow {
730 data: nulls.buffer.into(),
731 config: nulls.config,
732 };
733
734 Ok(Self::Struct(converter, owned))
735 }
736 DataType::Union(fields, _mode) => {
737 let options = SortOptions {
740 descending: false,
741 nulls_first: sort_field.options.nulls_first != sort_field.options.descending,
742 };
743
744 let mut converters = Vec::with_capacity(fields.len());
745 let mut type_ids = Vec::with_capacity(fields.len());
746 let mut null_rows = Vec::with_capacity(fields.len());
747
748 for (type_id, field) in fields.iter() {
749 let sort_field =
750 SortField::new_with_options(field.data_type().clone(), options);
751 let converter = RowConverter::new(vec![sort_field])?;
752
753 let null_array = new_null_array(field.data_type(), 1);
754 let nulls = converter.convert_columns(&[null_array])?;
755 let owned = OwnedRow {
756 data: nulls.buffer.into(),
757 config: nulls.config,
758 };
759
760 converters.push(converter);
761 type_ids.push(type_id);
762 null_rows.push(owned);
763 }
764
765 Ok(Self::Union(converters, type_ids, null_rows))
766 }
767 _ => Err(ArrowError::NotYetImplemented(format!(
768 "not yet implemented: {:?}",
769 sort_field.data_type
770 ))),
771 }
772 }
773
774 fn encoder(&self, array: &dyn Array) -> Result<Encoder<'_>, ArrowError> {
775 match self {
776 Codec::Stateless => Ok(Encoder::Stateless),
777 Codec::Dictionary(converter, nulls) => {
778 let values = array.as_any_dictionary().values().clone();
779 let rows = converter.convert_columns(&[values])?;
780 Ok(Encoder::Dictionary(rows, nulls.row()))
781 }
782 Codec::Struct(converter, null) => {
783 let v = as_struct_array(array);
784 let rows = converter.convert_columns(v.columns())?;
785 Ok(Encoder::Struct(rows, null.row()))
786 }
787 Codec::List(converter) => {
788 let values = match array.data_type() {
789 DataType::List(_) => {
790 let list_array = as_list_array(array);
791 let first_offset = list_array.offsets()[0] as usize;
792 let last_offset =
793 list_array.offsets()[list_array.offsets().len() - 1] as usize;
794
795 list_array
798 .values()
799 .slice(first_offset, last_offset - first_offset)
800 }
801 DataType::LargeList(_) => {
802 let list_array = as_large_list_array(array);
803
804 let first_offset = list_array.offsets()[0] as usize;
805 let last_offset =
806 list_array.offsets()[list_array.offsets().len() - 1] as usize;
807
808 list_array
811 .values()
812 .slice(first_offset, last_offset - first_offset)
813 }
814 DataType::ListView(_) => {
815 let list_view_array = array.as_list_view::<i32>();
816 let (min_offset, max_end) = compute_list_view_bounds(list_view_array);
817 list_view_array
818 .values()
819 .slice(min_offset, max_end - min_offset)
820 }
821 DataType::LargeListView(_) => {
822 let list_view_array = array.as_list_view::<i64>();
823 let (min_offset, max_end) = compute_list_view_bounds(list_view_array);
824 list_view_array
825 .values()
826 .slice(min_offset, max_end - min_offset)
827 }
828 DataType::FixedSizeList(_, _) => {
829 as_fixed_size_list_array(array).values().clone()
830 }
831 _ => unreachable!(),
832 };
833 let rows = converter.convert_columns(&[values])?;
834 Ok(Encoder::List(rows))
835 }
836 Codec::Map(converter) => {
837 let map_array = as_map_array(array);
838
839 let first_offset = map_array.offsets()[0] as usize;
840 let last_offset = map_array.offsets()[map_array.offsets().len() - 1] as usize;
841
842 let sliced_entries = map_array
845 .entries()
846 .slice(first_offset, last_offset - first_offset);
847
848 let rows = converter.convert_columns(sliced_entries.columns())?;
850 Ok(Encoder::Map(rows))
851 }
852 Codec::RunEndEncoded(converter) => {
853 let values = match array.data_type() {
854 DataType::RunEndEncoded(r, _) => match r.data_type() {
855 DataType::Int16 => array.as_run::<Int16Type>().values_slice(),
856 DataType::Int32 => array.as_run::<Int32Type>().values_slice(),
857 DataType::Int64 => array.as_run::<Int64Type>().values_slice(),
858 _ => unreachable!("Unsupported run end index type: {r:?}"),
859 },
860 _ => unreachable!(),
861 };
862 let rows = converter.convert_columns(std::slice::from_ref(&values))?;
863 Ok(Encoder::RunEndEncoded(rows))
864 }
865 Codec::Union(converters, field_to_type_ids, _) => {
866 let union_array = array
867 .as_any()
868 .downcast_ref::<UnionArray>()
869 .expect("expected Union array");
870
871 let type_ids = union_array.type_ids().clone();
872 let offsets = union_array.offsets().cloned();
873
874 let mut child_rows = Vec::with_capacity(converters.len());
875 for (field_idx, converter) in converters.iter().enumerate() {
876 let type_id = field_to_type_ids[field_idx];
877 let child_array = union_array.child(type_id);
878 let rows = converter.convert_columns(std::slice::from_ref(child_array))?;
879 child_rows.push(rows);
880 }
881
882 Ok(Encoder::Union {
883 child_rows,
884 field_to_type_ids: field_to_type_ids.clone(),
885 type_ids,
886 offsets,
887 })
888 }
889 }
890 }
891
892 fn size(&self) -> usize {
893 match self {
894 Codec::Stateless => 0,
895 Codec::Dictionary(converter, nulls) => converter.size() + nulls.data.len(),
896 Codec::Struct(converter, nulls) => converter.size() + nulls.data.len(),
897 Codec::List(converter) => converter.size(),
898 Codec::Map(converter) => converter.size(),
899 Codec::RunEndEncoded(converter) => converter.size(),
900 Codec::Union(converters, _, null_rows) => {
901 converters.iter().map(|c| c.size()).sum::<usize>()
902 + null_rows.iter().map(|n| n.data.len()).sum::<usize>()
903 }
904 }
905 }
906}
907
908#[derive(Debug)]
909enum Encoder<'a> {
910 Stateless,
912 Dictionary(Rows, Row<'a>),
914 Struct(Rows, Row<'a>),
920 List(Rows),
922 Map(Rows),
924 RunEndEncoded(Rows),
926 Union {
928 child_rows: Vec<Rows>,
929 field_to_type_ids: Vec<i8>,
930 type_ids: ScalarBuffer<i8>,
931 offsets: Option<ScalarBuffer<i32>>,
932 },
933}
934
935#[derive(Debug, Clone, PartialEq, Eq)]
937pub struct SortField {
938 options: SortOptions,
940 data_type: DataType,
942}
943
944impl SortField {
945 pub fn new(data_type: DataType) -> Self {
947 Self::new_with_options(data_type, Default::default())
948 }
949
950 pub fn new_with_options(data_type: DataType, options: SortOptions) -> Self {
952 Self { options, data_type }
953 }
954
955 pub fn size(&self) -> usize {
959 self.data_type.size() + std::mem::size_of::<Self>() - std::mem::size_of::<DataType>()
960 }
961}
962
963impl RowConverter {
964 pub fn new(fields: Vec<SortField>) -> Result<Self, ArrowError> {
966 if !Self::supports_fields(&fields) {
967 return Err(ArrowError::NotYetImplemented(format!(
968 "Row format support not yet implemented for: {fields:?}"
969 )));
970 }
971
972 let codecs = fields.iter().map(Codec::new).collect::<Result<_, _>>()?;
973 Ok(Self {
974 fields: fields.into(),
975 codecs,
976 })
977 }
978
979 pub fn supports_fields(fields: &[SortField]) -> bool {
981 fields.iter().all(|x| Self::supports_datatype(&x.data_type))
982 }
983
984 fn supports_datatype(d: &DataType) -> bool {
985 match d {
986 _ if !d.is_nested() => true,
987 DataType::List(f)
988 | DataType::LargeList(f)
989 | DataType::ListView(f)
990 | DataType::LargeListView(f)
991 | DataType::FixedSizeList(f, _)
992 | DataType::Map(f, _) => Self::supports_datatype(f.data_type()),
993 DataType::Struct(f) => f.iter().all(|x| Self::supports_datatype(x.data_type())),
994 DataType::RunEndEncoded(_, values) => Self::supports_datatype(values.data_type()),
995 DataType::Union(fs, _mode) => fs
996 .iter()
997 .all(|(_, f)| Self::supports_datatype(f.data_type())),
998 _ => false,
999 }
1000 }
1001
1002 pub fn convert_columns(&self, columns: &[ArrayRef]) -> Result<Rows, ArrowError> {
1012 let num_rows = columns.first().map(|x| x.len()).unwrap_or(0);
1013 let mut rows = self.empty_rows(num_rows, 0);
1014 self.append(&mut rows, columns)?;
1015 Ok(rows)
1016 }
1017
1018 pub fn append(&self, rows: &mut Rows, columns: &[ArrayRef]) -> Result<(), ArrowError> {
1049 assert!(
1050 Arc::ptr_eq(&rows.config.fields, &self.fields),
1051 "rows were not produced by this RowConverter"
1052 );
1053
1054 if columns.len() != self.fields.len() {
1055 return Err(ArrowError::InvalidArgumentError(format!(
1056 "Incorrect number of arrays provided to RowConverter, expected {} got {}",
1057 self.fields.len(),
1058 columns.len()
1059 )));
1060 }
1061 for column in columns.iter().skip(1) {
1062 if column.len() != columns[0].len() {
1063 return Err(ArrowError::InvalidArgumentError(format!(
1064 "RowConverter columns must all have the same length, expected {} got {}",
1065 columns[0].len(),
1066 column.len()
1067 )));
1068 }
1069 }
1070
1071 let encoders = columns
1072 .iter()
1073 .zip(&self.codecs)
1074 .zip(self.fields.iter())
1075 .map(|((column, codec), field)| {
1076 if !column.data_type().equals_datatype(&field.data_type) {
1077 return Err(ArrowError::InvalidArgumentError(format!(
1078 "RowConverter column schema mismatch, expected {} got {}",
1079 field.data_type,
1080 column.data_type()
1081 )));
1082 }
1083 codec.encoder(column.as_ref())
1084 })
1085 .collect::<Result<Vec<_>, _>>()?;
1086
1087 let write_offset = rows.num_rows();
1088 let lengths = row_lengths(columns, &encoders);
1089 let total = lengths.extend_offsets(rows.offsets[write_offset], &mut rows.offsets);
1090 rows.buffer.resize(total, 0);
1091
1092 for ((column, field), encoder) in columns.iter().zip(self.fields.iter()).zip(encoders) {
1093 encode_column(
1095 &mut rows.buffer,
1096 &mut rows.offsets[write_offset..],
1097 column.as_ref(),
1098 field.options,
1099 &encoder,
1100 )
1101 }
1102
1103 if cfg!(debug_assertions) {
1104 assert_eq!(*rows.offsets.last().unwrap(), rows.buffer.len());
1105 rows.offsets
1106 .windows(2)
1107 .for_each(|w| assert!(w[0] <= w[1], "offsets should be monotonic"));
1108 }
1109
1110 Ok(())
1111 }
1112
1113 pub fn convert_rows<'a, I>(&self, rows: I) -> Result<Vec<ArrayRef>, ArrowError>
1121 where
1122 I: IntoIterator<Item = Row<'a>>,
1123 {
1124 let mut validate_utf8 = false;
1125 let mut rows: Vec<_> = rows
1126 .into_iter()
1127 .map(|row| {
1128 assert!(
1129 Arc::ptr_eq(&row.config.fields, &self.fields),
1130 "rows were not produced by this RowConverter"
1131 );
1132 validate_utf8 |= row.config.validate_utf8;
1133 row.data
1134 })
1135 .collect();
1136
1137 let result = unsafe { self.convert_raw(&mut rows, validate_utf8) }?;
1141
1142 if cfg!(debug_assertions) {
1143 for (i, row) in rows.iter().enumerate() {
1144 if !row.is_empty() {
1145 return Err(ArrowError::InvalidArgumentError(format!(
1146 "Codecs {codecs:?} did not consume all bytes for row {i}, remaining bytes: {row:?}",
1147 codecs = self.codecs
1148 )));
1149 }
1150 }
1151 }
1152
1153 Ok(result)
1154 }
1155
1156 pub fn empty_rows(&self, row_capacity: usize, data_capacity: usize) -> Rows {
1185 let mut offsets = Vec::with_capacity(row_capacity.saturating_add(1));
1186 offsets.push(0);
1187
1188 Rows {
1189 offsets,
1190 buffer: Vec::with_capacity(data_capacity),
1191 config: RowConfig {
1192 fields: self.fields.clone(),
1193 validate_utf8: false,
1194 },
1195 }
1196 }
1197
1198 pub fn from_binary(&self, array: BinaryArray) -> Rows {
1225 assert_eq!(
1226 array.null_count(),
1227 0,
1228 "can't construct Rows instance from array with nulls"
1229 );
1230 let (offsets, values, _) = array.into_parts();
1231 let offsets = offsets.iter().map(|&i| i.as_usize()).collect();
1232 let buffer = values.into_vec().unwrap_or_else(|values| values.to_vec());
1234 Rows {
1235 buffer,
1236 offsets,
1237 config: RowConfig {
1238 fields: Arc::clone(&self.fields),
1239 validate_utf8: true,
1240 },
1241 }
1242 }
1243
1244 unsafe fn convert_raw(
1250 &self,
1251 rows: &mut [&[u8]],
1252 validate_utf8: bool,
1253 ) -> Result<Vec<ArrayRef>, ArrowError> {
1254 self.fields
1255 .iter()
1256 .zip(&self.codecs)
1257 .map(|(field, codec)| unsafe { decode_column(field, rows, codec, validate_utf8) })
1258 .collect()
1259 }
1260
1261 pub fn parser(&self) -> RowParser {
1263 RowParser::new(Arc::clone(&self.fields))
1264 }
1265
1266 pub unsafe fn parser_skip_utf8_validation(&self) -> RowParser {
1271 unsafe { RowParser::with_skip_utf8_validate(Arc::clone(&self.fields)) }
1272 }
1273
1274 pub fn size(&self) -> usize {
1278 std::mem::size_of::<Self>()
1279 + self.fields.iter().map(|x| x.size()).sum::<usize>()
1280 + self.codecs.capacity() * std::mem::size_of::<Codec>()
1281 + self.codecs.iter().map(Codec::size).sum::<usize>()
1282 }
1283}
1284
1285#[derive(Debug)]
1287pub struct RowParser {
1288 config: RowConfig,
1289}
1290
1291impl RowParser {
1292 fn new(fields: Arc<[SortField]>) -> Self {
1293 Self {
1294 config: RowConfig {
1295 fields,
1296 validate_utf8: true,
1297 },
1298 }
1299 }
1300 unsafe fn with_skip_utf8_validate(fields: Arc<[SortField]>) -> Self {
1305 Self {
1306 config: RowConfig {
1307 fields,
1308 validate_utf8: false,
1309 },
1310 }
1311 }
1312
1313 pub fn parse<'a>(&'a self, bytes: &'a [u8]) -> Row<'a> {
1318 Row {
1319 data: bytes,
1320 config: &self.config,
1321 }
1322 }
1323}
1324
1325#[derive(Debug, Clone)]
1327struct RowConfig {
1328 fields: Arc<[SortField]>,
1330 validate_utf8: bool,
1332}
1333
1334#[derive(Debug, Clone)]
1338pub struct Rows {
1339 buffer: Vec<u8>,
1341 offsets: Vec<usize>,
1343 config: RowConfig,
1345}
1346
1347pub type RowLengthIter<'a> = Map<Windows<'a, usize>, fn(&'a [usize]) -> usize>;
1349
1350impl Rows {
1351 pub fn push(&mut self, row: Row<'_>) {
1357 assert!(
1358 Arc::ptr_eq(&row.config.fields, &self.config.fields),
1359 "row was not produced by this RowConverter"
1360 );
1361 self.config.validate_utf8 |= row.config.validate_utf8;
1362 self.buffer.extend_from_slice(row.data);
1363 self.offsets.push(self.buffer.len())
1364 }
1365
1366 pub fn reserve(&mut self, row_capacity: usize, data_capacity: usize) {
1368 self.buffer.reserve(data_capacity);
1369 self.offsets.reserve(row_capacity);
1370 }
1371
1372 pub fn row(&self, row: usize) -> Row<'_> {
1378 self.checked_row_end(row);
1379 unsafe { self.row_unchecked(row) }
1380 }
1381
1382 fn checked_row_end(&self, row: usize) -> usize {
1383 row.checked_add(1)
1384 .filter(|end| *end < self.offsets.len())
1385 .expect("row index out of bounds")
1386 }
1387
1388 pub unsafe fn row_unchecked(&self, index: usize) -> Row<'_> {
1393 let end = unsafe { self.offsets.get_unchecked(index + 1) };
1394 let start = unsafe { self.offsets.get_unchecked(index) };
1395 let data = unsafe { self.buffer.get_unchecked(*start..*end) };
1396 Row {
1397 data,
1398 config: &self.config,
1399 }
1400 }
1401
1402 pub fn row_len(&self, row: usize) -> usize {
1405 let end = self.checked_row_end(row);
1406
1407 self.offsets[end] - self.offsets[row]
1408 }
1409
1410 pub fn lengths(&self) -> RowLengthIter<'_> {
1412 self.offsets.windows(2).map(|w| w[1] - w[0])
1413 }
1414
1415 pub fn clear(&mut self) {
1417 self.offsets.truncate(1);
1418 self.buffer.clear();
1419 }
1420
1421 pub fn num_rows(&self) -> usize {
1423 self.offsets.len() - 1
1424 }
1425
1426 pub fn iter(&self) -> RowsIter<'_> {
1428 self.into_iter()
1429 }
1430
1431 pub fn size(&self) -> usize {
1435 std::mem::size_of::<Self>()
1437 + self.buffer.capacity()
1438 + self.offsets.capacity() * std::mem::size_of::<usize>()
1439 }
1440
1441 pub fn try_into_binary(self) -> Result<BinaryArray, ArrowError> {
1471 if self.buffer.len() > i32::MAX as usize {
1472 return Err(ArrowError::InvalidArgumentError(format!(
1473 "{}-byte rows buffer too long to convert into a i32-indexed BinaryArray",
1474 self.buffer.len()
1475 )));
1476 }
1477 let offsets_scalar = ScalarBuffer::from_iter(self.offsets.into_iter().map(i32::usize_as));
1479 let array = unsafe {
1481 BinaryArray::new_unchecked(
1482 OffsetBuffer::new_unchecked(offsets_scalar),
1483 Buffer::from_vec(self.buffer),
1484 None,
1485 )
1486 };
1487 Ok(array)
1488 }
1489}
1490
1491impl<'a> IntoIterator for &'a Rows {
1492 type Item = Row<'a>;
1493 type IntoIter = RowsIter<'a>;
1494
1495 fn into_iter(self) -> Self::IntoIter {
1496 RowsIter {
1497 rows: self,
1498 start: 0,
1499 end: self.num_rows(),
1500 }
1501 }
1502}
1503
1504#[derive(Debug)]
1506pub struct RowsIter<'a> {
1507 rows: &'a Rows,
1508 start: usize,
1509 end: usize,
1510}
1511
1512impl<'a> Iterator for RowsIter<'a> {
1513 type Item = Row<'a>;
1514
1515 fn next(&mut self) -> Option<Self::Item> {
1516 if self.end == self.start {
1517 return None;
1518 }
1519
1520 let row = unsafe { self.rows.row_unchecked(self.start) };
1522 self.start += 1;
1523 Some(row)
1524 }
1525
1526 fn size_hint(&self) -> (usize, Option<usize>) {
1527 let len = self.len();
1528 (len, Some(len))
1529 }
1530}
1531
1532impl ExactSizeIterator for RowsIter<'_> {
1533 fn len(&self) -> usize {
1534 self.end - self.start
1535 }
1536}
1537
1538impl DoubleEndedIterator for RowsIter<'_> {
1539 fn next_back(&mut self) -> Option<Self::Item> {
1540 if self.end == self.start {
1541 return None;
1542 }
1543
1544 self.end -= 1;
1545
1546 let row = unsafe { self.rows.row_unchecked(self.end) };
1549 Some(row)
1550 }
1551}
1552
1553#[derive(Debug, Copy, Clone)]
1562pub struct Row<'a> {
1563 data: &'a [u8],
1564 config: &'a RowConfig,
1565}
1566
1567impl<'a> Row<'a> {
1568 pub fn owned(&self) -> OwnedRow {
1570 OwnedRow {
1571 data: self.data.into(),
1572 config: self.config.clone(),
1573 }
1574 }
1575
1576 pub fn data(&self) -> &'a [u8] {
1578 self.data
1579 }
1580}
1581
1582impl PartialEq for Row<'_> {
1585 #[inline]
1586 fn eq(&self, other: &Self) -> bool {
1587 self.data.eq(other.data)
1588 }
1589}
1590
1591impl Eq for Row<'_> {}
1592
1593impl PartialOrd for Row<'_> {
1594 #[inline]
1595 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1596 Some(self.cmp(other))
1597 }
1598}
1599
1600impl Ord for Row<'_> {
1601 #[inline]
1602 fn cmp(&self, other: &Self) -> Ordering {
1603 self.data.cmp(other.data)
1604 }
1605}
1606
1607impl Hash for Row<'_> {
1608 #[inline]
1609 fn hash<H: Hasher>(&self, state: &mut H) {
1610 self.data.hash(state)
1611 }
1612}
1613
1614impl AsRef<[u8]> for Row<'_> {
1615 #[inline]
1616 fn as_ref(&self) -> &[u8] {
1617 self.data
1618 }
1619}
1620
1621#[derive(Debug, Clone)]
1625pub struct OwnedRow {
1626 data: Box<[u8]>,
1627 config: RowConfig,
1628}
1629
1630impl OwnedRow {
1631 pub fn row(&self) -> Row<'_> {
1635 Row {
1636 data: &self.data,
1637 config: &self.config,
1638 }
1639 }
1640}
1641
1642impl PartialEq for OwnedRow {
1645 #[inline]
1646 fn eq(&self, other: &Self) -> bool {
1647 self.row().eq(&other.row())
1648 }
1649}
1650
1651impl Eq for OwnedRow {}
1652
1653impl PartialOrd for OwnedRow {
1654 #[inline]
1655 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1656 Some(self.cmp(other))
1657 }
1658}
1659
1660impl Ord for OwnedRow {
1661 #[inline]
1662 fn cmp(&self, other: &Self) -> Ordering {
1663 self.row().cmp(&other.row())
1664 }
1665}
1666
1667impl Hash for OwnedRow {
1668 #[inline]
1669 fn hash<H: Hasher>(&self, state: &mut H) {
1670 self.row().hash(state)
1671 }
1672}
1673
1674impl AsRef<[u8]> for OwnedRow {
1675 #[inline]
1676 fn as_ref(&self) -> &[u8] {
1677 &self.data
1678 }
1679}
1680
1681#[inline]
1683fn null_sentinel(options: SortOptions) -> u8 {
1684 match options.nulls_first {
1685 true => 0,
1686 false => 0xFF,
1687 }
1688}
1689
1690enum LengthTracker {
1692 Fixed { length: usize, num_rows: usize },
1694 Variable {
1696 fixed_length: usize,
1697 lengths: Vec<usize>,
1698 },
1699}
1700
1701impl LengthTracker {
1702 fn new(num_rows: usize) -> Self {
1703 Self::Fixed {
1704 length: 0,
1705 num_rows,
1706 }
1707 }
1708
1709 fn push_fixed(&mut self, new_length: usize) {
1711 match self {
1712 LengthTracker::Fixed { length, .. } => *length += new_length,
1713 LengthTracker::Variable { fixed_length, .. } => *fixed_length += new_length,
1714 }
1715 }
1716
1717 fn push_variable(&mut self, new_lengths: impl ExactSizeIterator<Item = usize>) {
1719 match self {
1720 LengthTracker::Fixed { length, .. } => {
1721 *self = LengthTracker::Variable {
1722 fixed_length: *length,
1723 lengths: new_lengths.collect(),
1724 }
1725 }
1726 LengthTracker::Variable { lengths, .. } => {
1727 assert_eq!(lengths.len(), new_lengths.len());
1728 lengths
1729 .iter_mut()
1730 .zip(new_lengths)
1731 .for_each(|(length, new_length)| *length += new_length);
1732 }
1733 }
1734 }
1735
1736 fn materialized(&mut self) -> &mut [usize] {
1738 if let LengthTracker::Fixed { length, num_rows } = *self {
1739 *self = LengthTracker::Variable {
1740 fixed_length: length,
1741 lengths: vec![0; num_rows],
1742 };
1743 }
1744
1745 match self {
1746 LengthTracker::Variable { lengths, .. } => lengths,
1747 LengthTracker::Fixed { .. } => unreachable!(),
1748 }
1749 }
1750
1751 fn extend_offsets(&self, initial_offset: usize, offsets: &mut Vec<usize>) -> usize {
1769 match self {
1770 LengthTracker::Fixed { length, num_rows } => {
1771 offsets.extend((0..*num_rows).map(|i| initial_offset + i * length));
1772
1773 initial_offset + num_rows * length
1774 }
1775 LengthTracker::Variable {
1776 fixed_length,
1777 lengths,
1778 } => {
1779 let mut acc = initial_offset;
1780
1781 offsets.extend(lengths.iter().map(|length| {
1782 let current = acc;
1783 acc += length + fixed_length;
1784 current
1785 }));
1786
1787 acc
1788 }
1789 }
1790 }
1791}
1792
1793fn row_lengths(cols: &[ArrayRef], encoders: &[Encoder]) -> LengthTracker {
1795 use fixed::FixedLengthEncoding;
1796
1797 let num_rows = cols.first().map(|x| x.len()).unwrap_or(0);
1798 let mut tracker = LengthTracker::new(num_rows);
1799
1800 for (array, encoder) in cols.iter().zip(encoders) {
1801 match encoder {
1802 Encoder::Stateless => {
1803 downcast_primitive_array! {
1804 array => tracker.push_fixed(fixed::encoded_len(array)),
1805 DataType::Null => tracker.push_fixed(2)
1806 DataType::Boolean => tracker.push_fixed(bool::ENCODED_LEN),
1807 DataType::Binary => push_generic_byte_array_lengths(&mut tracker, as_generic_binary_array::<i32>(array)),
1808 DataType::LargeBinary => push_generic_byte_array_lengths(&mut tracker, as_generic_binary_array::<i64>(array)),
1809 DataType::BinaryView => push_byte_view_array_lengths(&mut tracker, array.as_binary_view()),
1810 DataType::Utf8 => push_generic_byte_array_lengths(&mut tracker, array.as_string::<i32>()),
1811 DataType::LargeUtf8 => push_generic_byte_array_lengths(&mut tracker, array.as_string::<i64>()),
1812 DataType::Utf8View => push_byte_view_array_lengths(&mut tracker, array.as_string_view()),
1813 DataType::FixedSizeBinary(len) => {
1814 let len = len.to_usize().unwrap();
1815 tracker.push_fixed(1 + len)
1816 }
1817 _ => unimplemented!("unsupported data type: {}", array.data_type()),
1818 }
1819 }
1820 Encoder::Dictionary(values, null) => {
1821 downcast_dictionary_array! {
1822 array => {
1823 tracker.push_variable(
1824 array.keys().iter().map(|v| match v {
1825 Some(k) => values.row_len(k.as_usize()),
1826 None => null.data.len(),
1827 })
1828 )
1829 }
1830 _ => unreachable!(),
1831 }
1832 }
1833 Encoder::Struct(rows, null) => {
1834 let array = as_struct_array(array);
1835 if rows.num_rows() > 0 {
1836 tracker.push_variable((0..array.len()).map(|idx| match array.is_valid(idx) {
1838 true => 1 + rows.row_len(idx),
1839 false => 1 + null.data.len(),
1840 }));
1841 } else {
1842 tracker.push_variable((0..array.len()).map(|idx| match array.is_valid(idx) {
1844 true => 1,
1845 false => 1 + null.data.len(),
1846 }));
1847 }
1848 }
1849 Encoder::List(rows) => match array.data_type() {
1850 DataType::List(_) => {
1851 list::compute_lengths(tracker.materialized(), rows, as_list_array(array))
1852 }
1853 DataType::LargeList(_) => {
1854 list::compute_lengths(tracker.materialized(), rows, as_large_list_array(array))
1855 }
1856 DataType::ListView(_) => {
1857 let list_view = array.as_list_view::<i32>();
1858 let (min_offset, _) = compute_list_view_bounds(list_view);
1859 list::compute_lengths_list_view(
1860 tracker.materialized(),
1861 rows,
1862 list_view,
1863 min_offset,
1864 )
1865 }
1866 DataType::LargeListView(_) => {
1867 let list_view = array.as_list_view::<i64>();
1868 let (min_offset, _) = compute_list_view_bounds(list_view);
1869 list::compute_lengths_list_view(
1870 tracker.materialized(),
1871 rows,
1872 list_view,
1873 min_offset,
1874 )
1875 }
1876 DataType::FixedSizeList(_, _) => compute_lengths_fixed_size_list(
1877 &mut tracker,
1878 rows,
1879 as_fixed_size_list_array(array),
1880 ),
1881 _ => unreachable!(),
1882 },
1883 Encoder::Map(rows) => {
1884 list::compute_lengths(tracker.materialized(), rows, as_map_array(array))
1885 }
1886 Encoder::RunEndEncoded(rows) => match array.data_type() {
1887 DataType::RunEndEncoded(r, _) => match r.data_type() {
1888 DataType::Int16 => run::compute_lengths(
1889 tracker.materialized(),
1890 rows,
1891 array.as_run::<Int16Type>(),
1892 ),
1893 DataType::Int32 => run::compute_lengths(
1894 tracker.materialized(),
1895 rows,
1896 array.as_run::<Int32Type>(),
1897 ),
1898 DataType::Int64 => run::compute_lengths(
1899 tracker.materialized(),
1900 rows,
1901 array.as_run::<Int64Type>(),
1902 ),
1903 _ => unreachable!("Unsupported run end index type: {r:?}"),
1904 },
1905 _ => unreachable!(),
1906 },
1907 Encoder::Union {
1908 child_rows,
1909 field_to_type_ids,
1910 type_ids,
1911 offsets,
1912 } => {
1913 let union_array = array
1914 .as_any()
1915 .downcast_ref::<UnionArray>()
1916 .expect("expected UnionArray");
1917
1918 let mut type_id_to_field_idx = [0usize; 128];
1919 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
1920 type_id_to_field_idx[type_id as usize] = field_idx;
1921 }
1922
1923 let lengths = (0..union_array.len()).map(|i| {
1924 let type_id = type_ids[i];
1925 let field_idx = type_id_to_field_idx[type_id as usize];
1926 let child_row_i = offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i);
1927 let child_row_len = child_rows[field_idx].row_len(child_row_i);
1928
1929 1 + child_row_len
1931 });
1932
1933 tracker.push_variable(lengths);
1934 }
1935 }
1936 }
1937
1938 tracker
1939}
1940
1941fn push_generic_byte_array_lengths<T: ByteArrayType>(
1943 tracker: &mut LengthTracker,
1944 array: &GenericByteArray<T>,
1945) {
1946 if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
1947 tracker.push_variable(
1948 array
1949 .offsets()
1950 .lengths()
1951 .zip(nulls.iter())
1952 .map(|(length, is_valid)| if is_valid { Some(length) } else { None })
1953 .map(variable::padded_length),
1954 )
1955 } else {
1956 tracker.push_variable(
1957 array
1958 .offsets()
1959 .lengths()
1960 .map(variable::non_null_padded_length),
1961 )
1962 }
1963}
1964
1965fn push_byte_view_array_lengths<T: ByteViewType>(
1967 tracker: &mut LengthTracker,
1968 array: &GenericByteViewArray<T>,
1969) {
1970 if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
1971 tracker.push_variable(
1972 array
1973 .lengths()
1974 .zip(nulls.iter())
1975 .map(|(length, is_valid)| {
1976 if is_valid {
1977 Some(length as usize)
1978 } else {
1979 None
1980 }
1981 })
1982 .map(variable::padded_length),
1983 )
1984 } else {
1985 tracker.push_variable(
1986 array
1987 .lengths()
1988 .map(|len| variable::padded_length(Some(len as usize))),
1989 )
1990 }
1991}
1992
1993fn encode_column(
1995 data: &mut [u8],
1996 offsets: &mut [usize],
1997 column: &dyn Array,
1998 opts: SortOptions,
1999 encoder: &Encoder<'_>,
2000) {
2001 match encoder {
2002 Encoder::Stateless => {
2003 downcast_primitive_array! {
2004 column => {
2005 if let Some(nulls) = column.nulls().filter(|n| n.null_count() > 0){
2006 fixed::encode(data, offsets, column.values(), nulls, opts)
2007 } else {
2008 fixed::encode_not_null(data, offsets, column.values(), opts)
2009 }
2010 }
2011 DataType::Null => {
2012 for offset in offsets.iter_mut().skip(1) {
2013 variable::encode_null_value(&mut data[*offset..], opts);
2014 *offset += 2;
2015 }
2016 }
2017 DataType::Boolean => {
2018 if let Some(nulls) = column.nulls().filter(|n| n.null_count() > 0){
2019 fixed::encode_boolean(data, offsets, column.as_boolean().values(), nulls, opts)
2020 } else {
2021 fixed::encode_boolean_not_null(data, offsets, column.as_boolean().values(), opts)
2022 }
2023 }
2024 DataType::Binary => {
2025 variable::encode_generic_byte_array(data, offsets, as_generic_binary_array::<i32>(column), opts)
2026 }
2027 DataType::BinaryView => {
2028 variable::encode(data, offsets, column.as_binary_view().iter(), opts)
2029 }
2030 DataType::LargeBinary => {
2031 variable::encode_generic_byte_array(data, offsets, as_generic_binary_array::<i64>(column), opts)
2032 }
2033 DataType::Utf8 => variable::encode_generic_byte_array(
2034 data, offsets,
2035 column.as_string::<i32>(),
2036 opts,
2037 ),
2038 DataType::LargeUtf8 => variable::encode_generic_byte_array(
2039 data, offsets,
2040 column.as_string::<i64>(),
2041 opts,
2042 ),
2043 DataType::Utf8View => variable::encode(
2044 data, offsets,
2045 column.as_string_view().iter().map(|x| x.map(|x| x.as_bytes())),
2046 opts,
2047 ),
2048 DataType::FixedSizeBinary(_) => {
2049 let array = column.as_any().downcast_ref().unwrap();
2050 fixed::encode_fixed_size_binary(data, offsets, array, opts)
2051 }
2052 _ => unimplemented!("unsupported data type: {}", column.data_type()),
2053 }
2054 }
2055 Encoder::Dictionary(values, nulls) => {
2056 downcast_dictionary_array! {
2057 column => encode_dictionary_values(data, offsets, column, values, nulls),
2058 _ => unreachable!()
2059 }
2060 }
2061 Encoder::Struct(rows, null) => {
2062 fn struct_encode_helper<const NO_CHILD_FIELDS: bool>(
2063 array: &StructArray,
2064 offsets: &mut [usize],
2065 null_sentinel: u8,
2066 rows: &Rows,
2067 null: &Row<'_>,
2068 data: &mut [u8],
2069 ) {
2070 let empty_row = Row {
2071 data: &[],
2072 config: &rows.config,
2073 };
2074
2075 offsets
2076 .iter_mut()
2077 .skip(1)
2078 .enumerate()
2079 .for_each(|(idx, offset)| {
2080 let (row, sentinel) = match array.is_valid(idx) {
2081 true => (
2082 if NO_CHILD_FIELDS {
2083 empty_row
2084 } else {
2085 rows.row(idx)
2086 },
2087 0x01,
2088 ),
2089 false => (*null, null_sentinel),
2090 };
2091 let end_offset = *offset + 1 + row.as_ref().len();
2092 data[*offset] = sentinel;
2093 data[*offset + 1..end_offset].copy_from_slice(row.as_ref());
2094 *offset = end_offset;
2095 })
2096 }
2097
2098 let array = as_struct_array(column);
2099 let null_sentinel = null_sentinel(opts);
2100 if rows.num_rows() == 0 {
2101 struct_encode_helper::<true>(array, offsets, null_sentinel, rows, null, data);
2103 } else {
2104 struct_encode_helper::<false>(array, offsets, null_sentinel, rows, null, data);
2105 }
2106 }
2107 Encoder::List(rows) => match column.data_type() {
2108 DataType::List(_) => list::encode(data, offsets, rows, opts, as_list_array(column)),
2109 DataType::LargeList(_) => {
2110 list::encode(data, offsets, rows, opts, as_large_list_array(column))
2111 }
2112 DataType::ListView(_) => {
2113 let list_view = column.as_list_view::<i32>();
2114 let (min_offset, _) = compute_list_view_bounds(list_view);
2115 list::encode_list_view(data, offsets, rows, opts, list_view, min_offset)
2116 }
2117 DataType::LargeListView(_) => {
2118 let list_view = column.as_list_view::<i64>();
2119 let (min_offset, _) = compute_list_view_bounds(list_view);
2120 list::encode_list_view(data, offsets, rows, opts, list_view, min_offset)
2121 }
2122 DataType::FixedSizeList(_, _) => {
2123 encode_fixed_size_list(data, offsets, rows, opts, as_fixed_size_list_array(column))
2124 }
2125 _ => unreachable!(),
2126 },
2127 Encoder::Map(rows) => list::encode(data, offsets, rows, opts, as_map_array(column)),
2128 Encoder::RunEndEncoded(rows) => match column.data_type() {
2129 DataType::RunEndEncoded(r, _) => match r.data_type() {
2130 DataType::Int16 => {
2131 run::encode(data, offsets, rows, opts, column.as_run::<Int16Type>())
2132 }
2133 DataType::Int32 => {
2134 run::encode(data, offsets, rows, opts, column.as_run::<Int32Type>())
2135 }
2136 DataType::Int64 => {
2137 run::encode(data, offsets, rows, opts, column.as_run::<Int64Type>())
2138 }
2139 _ => unreachable!("Unsupported run end index type: {r:?}"),
2140 },
2141 _ => unreachable!(),
2142 },
2143 Encoder::Union {
2144 child_rows,
2145 field_to_type_ids,
2146 type_ids,
2147 offsets: offsets_buf,
2148 } => {
2149 let mut type_id_to_field_idx = [0usize; 128];
2150 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
2151 type_id_to_field_idx[type_id as usize] = field_idx;
2152 }
2153
2154 offsets
2155 .iter_mut()
2156 .skip(1)
2157 .enumerate()
2158 .for_each(|(i, offset)| {
2159 let type_id = type_ids[i];
2160 let field_idx = type_id_to_field_idx[type_id as usize];
2161
2162 let child_row_idx = offsets_buf.as_ref().map(|o| o[i] as usize).unwrap_or(i);
2163 let child_row = child_rows[field_idx].row(child_row_idx);
2164 let child_bytes = child_row.as_ref();
2165
2166 let type_id_byte = if opts.descending {
2167 !(type_id as u8)
2168 } else {
2169 type_id as u8
2170 };
2171 data[*offset] = type_id_byte;
2172
2173 let child_start = *offset + 1;
2174 let child_end = child_start + child_bytes.len();
2175 data[child_start..child_end].copy_from_slice(child_bytes);
2176
2177 *offset = child_end;
2178 });
2179 }
2180 }
2181}
2182
2183pub fn encode_dictionary_values<K: ArrowDictionaryKeyType>(
2185 data: &mut [u8],
2186 offsets: &mut [usize],
2187 column: &DictionaryArray<K>,
2188 values: &Rows,
2189 null: &Row<'_>,
2190) {
2191 for (offset, k) in offsets.iter_mut().skip(1).zip(column.keys()) {
2192 let row = match k {
2193 Some(k) => values.row(k.as_usize()).data,
2194 None => null.data,
2195 };
2196 let end_offset = *offset + row.len();
2197 data[*offset..end_offset].copy_from_slice(row);
2198 *offset = end_offset;
2199 }
2200}
2201
2202macro_rules! decode_primitive_helper {
2203 ($t:ty, $rows:ident, $data_type:ident, $options:ident) => {
2204 Arc::new(decode_primitive::<$t>($rows, $data_type, $options))
2205 };
2206}
2207
2208unsafe fn decode_column(
2214 field: &SortField,
2215 rows: &mut [&[u8]],
2216 codec: &Codec,
2217 validate_utf8: bool,
2218) -> Result<ArrayRef, ArrowError> {
2219 let options = field.options;
2220
2221 let array: ArrayRef = match codec {
2222 Codec::Stateless => {
2223 let data_type = field.data_type.clone();
2224 downcast_primitive! {
2225 data_type => (decode_primitive_helper, rows, data_type, options),
2226 DataType::Null => {
2227 variable::decode_null_value(rows, options);
2228 Arc::new(NullArray::new(rows.len()))
2229 }
2230 DataType::Boolean => Arc::new(decode_bool(rows, options)),
2231 DataType::Binary => Arc::new(decode_binary::<i32>(rows, options)),
2232 DataType::LargeBinary => Arc::new(decode_binary::<i64>(rows, options)),
2233 DataType::BinaryView => Arc::new(decode_binary_view(rows, options)),
2234 DataType::FixedSizeBinary(size) => Arc::new(decode_fixed_size_binary(rows, size, options)),
2235 DataType::Utf8 => Arc::new(unsafe{ decode_string::<i32>(rows, options, validate_utf8) }),
2236 DataType::LargeUtf8 => Arc::new(unsafe { decode_string::<i64>(rows, options, validate_utf8) }),
2237 DataType::Utf8View => Arc::new(unsafe { decode_string_view(rows, options, validate_utf8) }),
2238 _ => return Err(ArrowError::NotYetImplemented(format!("unsupported data type: {data_type}" )))
2239 }
2240 }
2241 Codec::Dictionary(converter, _) => {
2242 let cols = unsafe { converter.convert_raw(rows, validate_utf8) }?;
2243 cols.into_iter().next().unwrap()
2244 }
2245 Codec::Struct(converter, _) => {
2246 let nulls = fixed::decode_nulls(rows);
2247 rows.iter_mut().for_each(|row| *row = &row[1..]);
2248 let children = unsafe { converter.convert_raw(rows, validate_utf8) }?;
2249
2250 let corrected_fields: Vec<Field> = match &field.data_type {
2253 DataType::Struct(struct_fields) => struct_fields
2254 .iter()
2255 .zip(children.iter())
2256 .map(|(orig_field, child_array)| {
2257 orig_field
2258 .as_ref()
2259 .clone()
2260 .with_data_type(child_array.data_type().clone())
2261 })
2262 .collect(),
2263 _ => unreachable!("Only Struct types should be corrected here"),
2264 };
2265
2266 Arc::new(unsafe {
2267 StructArray::new_unchecked_with_length(
2268 corrected_fields.into(),
2269 children,
2270 nulls,
2271 rows.len(),
2272 )
2273 })
2274 }
2275 Codec::List(converter) => match &field.data_type {
2276 DataType::List(_) => Arc::new(unsafe {
2277 list::decode::<GenericListArray<i32>>(converter, rows, field, validate_utf8)
2278 }?),
2279 DataType::LargeList(_) => Arc::new(unsafe {
2280 list::decode::<GenericListArray<i64>>(converter, rows, field, validate_utf8)
2281 }?),
2282 DataType::ListView(_) => Arc::new(unsafe {
2283 list::decode_list_view::<i32>(converter, rows, field, validate_utf8)
2284 }?),
2285 DataType::LargeListView(_) => Arc::new(unsafe {
2286 list::decode_list_view::<i64>(converter, rows, field, validate_utf8)
2287 }?),
2288 DataType::FixedSizeList(_, value_length) => Arc::new(unsafe {
2289 list::decode_fixed_size_list(
2290 converter,
2291 rows,
2292 field,
2293 validate_utf8,
2294 value_length.as_usize(),
2295 )
2296 }?),
2297 _ => unreachable!(),
2298 },
2299 Codec::Map(converter) => {
2300 Arc::new(unsafe { list::decode::<MapArray>(converter, rows, field, validate_utf8) }?)
2301 }
2302 Codec::RunEndEncoded(converter) => match &field.data_type {
2303 DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() {
2304 DataType::Int16 => Arc::new(unsafe {
2305 run::decode::<Int16Type>(converter, rows, field, validate_utf8)
2306 }?),
2307 DataType::Int32 => Arc::new(unsafe {
2308 run::decode::<Int32Type>(converter, rows, field, validate_utf8)
2309 }?),
2310 DataType::Int64 => Arc::new(unsafe {
2311 run::decode::<Int64Type>(converter, rows, field, validate_utf8)
2312 }?),
2313 _ => unreachable!(),
2314 },
2315 _ => unreachable!(),
2316 },
2317 Codec::Union(converters, field_to_type_ids, null_rows) => {
2318 let len = rows.len();
2319
2320 let DataType::Union(union_fields, mode) = &field.data_type else {
2321 unreachable!()
2322 };
2323
2324 let mut type_id_to_field_idx = [0usize; 128];
2325 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
2326 type_id_to_field_idx[type_id as usize] = field_idx;
2327 }
2328
2329 let mut type_ids = Vec::with_capacity(len);
2330 let mut rows_by_field: Vec<Vec<(usize, &[u8])>> = vec![Vec::new(); converters.len()];
2331
2332 for (idx, row) in rows.iter_mut().enumerate() {
2333 let type_id_byte = {
2334 let id = row[0];
2335 if options.descending { !id } else { id }
2336 };
2337
2338 let type_id = type_id_byte as i8;
2339 type_ids.push(type_id);
2340
2341 let field_idx = type_id_to_field_idx[type_id as usize];
2342
2343 let child_row = &row[1..];
2344 rows_by_field[field_idx].push((idx, child_row));
2345 }
2346
2347 let mut child_arrays: Vec<ArrayRef> = Vec::with_capacity(converters.len());
2348 let mut offsets = (*mode == UnionMode::Dense).then(|| Vec::with_capacity(len));
2349
2350 for (field_idx, converter) in converters.iter().enumerate() {
2351 let field_rows = &rows_by_field[field_idx];
2352
2353 match &mode {
2354 UnionMode::Dense => {
2355 if field_rows.is_empty() {
2356 let (_, field) = union_fields.iter().nth(field_idx).unwrap();
2357 child_arrays.push(arrow_array::new_empty_array(field.data_type()));
2358 continue;
2359 }
2360
2361 let mut child_data = field_rows
2362 .iter()
2363 .map(|(_, bytes)| *bytes)
2364 .collect::<Vec<_>>();
2365
2366 let child_array =
2367 unsafe { converter.convert_raw(&mut child_data, validate_utf8) }?;
2368
2369 for ((row_idx, original_bytes), remaining_bytes) in
2371 field_rows.iter().zip(child_data)
2372 {
2373 let consumed_length = 1 + original_bytes.len() - remaining_bytes.len();
2374 rows[*row_idx] = &rows[*row_idx][consumed_length..];
2375 }
2376
2377 child_arrays.push(child_array.into_iter().next().unwrap());
2378 }
2379 UnionMode::Sparse => {
2380 let mut sparse_data: Vec<&[u8]> = Vec::with_capacity(len);
2381 let mut field_row_iter = field_rows.iter().peekable();
2382 let null_row_bytes: &[u8] = &null_rows[field_idx].data;
2383
2384 for idx in 0..len {
2385 if let Some((next_idx, bytes)) = field_row_iter.peek()
2386 && *next_idx == idx
2387 {
2388 sparse_data.push(*bytes);
2389
2390 field_row_iter.next();
2391 continue;
2392 }
2393 sparse_data.push(null_row_bytes);
2394 }
2395
2396 let child_array =
2397 unsafe { converter.convert_raw(&mut sparse_data, validate_utf8) }?;
2398
2399 for (row_idx, child_row) in field_rows {
2401 let remaining_len = sparse_data[*row_idx].len();
2402 let consumed_length = 1 + child_row.len() - remaining_len;
2403 rows[*row_idx] = &rows[*row_idx][consumed_length..];
2404 }
2405
2406 child_arrays.push(child_array.into_iter().next().unwrap());
2407 }
2408 }
2409 }
2410
2411 if let Some(ref mut offsets_vec) = offsets {
2413 let mut count = vec![0i32; converters.len()];
2414 for type_id in &type_ids {
2415 let field_idx = *type_id as usize;
2416 offsets_vec.push(count[field_idx]);
2417
2418 count[field_idx] += 1;
2419 }
2420 }
2421
2422 let type_ids_buffer = ScalarBuffer::from(type_ids);
2423 let offsets_buffer = offsets.map(ScalarBuffer::from);
2424
2425 let union_array = UnionArray::try_new(
2426 union_fields.clone(),
2427 type_ids_buffer,
2428 offsets_buffer,
2429 child_arrays,
2430 )?;
2431
2432 Arc::new(union_array)
2435 }
2436 };
2437 Ok(array)
2438}
2439
2440#[cfg(test)]
2441mod tests {
2442 use arrow_array::builder::*;
2443 use arrow_array::types::*;
2444 use arrow_buffer::{Buffer, OffsetBuffer};
2445 use arrow_buffer::{NullBuffer, i256};
2446 use arrow_cast::display::{ArrayFormatter, FormatOptions};
2447 use arrow_ord::sort::{LexicographicalComparator, SortColumn};
2448 use rand::distr::uniform::SampleUniform;
2449 use rand::distr::{Distribution, StandardUniform};
2450 use rand::prelude::StdRng;
2451 use rand::{RngExt, SeedableRng};
2452
2453 use super::*;
2454
2455 fn all_sort_options() -> [SortOptions; 4] {
2456 [
2457 SortOptions {
2458 descending: false,
2459 nulls_first: false,
2460 },
2461 SortOptions {
2462 descending: false,
2463 nulls_first: true,
2464 },
2465 SortOptions {
2466 descending: true,
2467 nulls_first: false,
2468 },
2469 SortOptions {
2470 descending: true,
2471 nulls_first: true,
2472 },
2473 ]
2474 }
2475
2476 #[test]
2477 fn test_fixed_width() {
2478 let cols = [
2479 Arc::new(Int16Array::from_iter([
2480 Some(1),
2481 Some(2),
2482 None,
2483 Some(-5),
2484 Some(2),
2485 Some(2),
2486 Some(0),
2487 ])) as ArrayRef,
2488 Arc::new(Float32Array::from_iter([
2489 Some(1.3),
2490 Some(2.5),
2491 None,
2492 Some(4.),
2493 Some(0.1),
2494 Some(-4.),
2495 Some(-0.),
2496 ])) as ArrayRef,
2497 ];
2498
2499 let converter = RowConverter::new(vec![
2500 SortField::new(DataType::Int16),
2501 SortField::new(DataType::Float32),
2502 ])
2503 .unwrap();
2504 let rows = converter.convert_columns(&cols).unwrap();
2505
2506 assert_eq!(rows.offsets, &[0, 8, 16, 24, 32, 40, 48, 56]);
2507 assert_eq!(
2508 rows.buffer,
2509 &[
2510 1, 128, 1, 1, 191, 166, 102, 102, 1, 128, 2, 1, 192, 32, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 127, 251, 1, 192, 128, 0, 0, 1, 128, 2, 1, 189, 204, 204, 205, 1, 128, 2, 1, 63, 127, 255, 255, 1, 128, 0, 1, 127, 255, 255, 255 ]
2525 );
2526
2527 assert!(rows.row(3) < rows.row(6));
2528 assert!(rows.row(0) < rows.row(1));
2529 assert!(rows.row(3) < rows.row(0));
2530 assert!(rows.row(4) < rows.row(1));
2531 assert!(rows.row(5) < rows.row(4));
2532
2533 let back = converter.convert_rows(&rows).unwrap();
2534 for (expected, actual) in cols.iter().zip(&back) {
2535 assert_eq!(expected, actual);
2536 }
2537 }
2538
2539 fn test_roundtrip(sort_option: SortOptions, col: ArrayRef) {
2540 let converter = RowConverter::new(vec![SortField::new_with_options(
2541 col.data_type().clone(),
2542 sort_option,
2543 )])
2544 .unwrap();
2545 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2546 let back = converter.convert_rows(&rows).unwrap();
2547 assert_eq!(back.len(), 1);
2548 assert_eq!(&back[0], &col);
2549 back[0].to_data().validate_full().unwrap();
2550 }
2551
2552 #[test]
2553 fn test_zero_width_fixed_size_binary_roundtrip() {
2554 for sort_option in all_sort_options() {
2555 for with_null in [true, false] {
2558 let nulls = if with_null {
2559 Some(NullBuffer::from(vec![true, false, true, false, true]))
2560 } else {
2561 None
2562 };
2563 let col: ArrayRef = Arc::new(
2564 FixedSizeBinaryArray::try_new_with_len(0, Buffer::default(), nulls, 5).unwrap(),
2565 );
2566
2567 test_roundtrip(sort_option, col);
2568 }
2569 }
2570 }
2571
2572 #[test]
2573 fn test_zero_width_fixed_size_list_roundtrip() {
2574 for sort_option in all_sort_options() {
2575 for with_null in [true, false] {
2578 let nulls = if with_null {
2579 Some(NullBuffer::from(vec![true, false, true, false, true]))
2580 } else {
2581 None
2582 };
2583 let col: ArrayRef = Arc::new(
2584 FixedSizeListArray::try_new_with_length(
2585 Arc::new(Field::new("item", DataType::Boolean, false)),
2586 0,
2587 new_empty_array(&DataType::Boolean),
2588 nulls,
2589 5,
2590 )
2591 .unwrap(),
2592 );
2593
2594 test_roundtrip(sort_option, col);
2595 }
2596 }
2597 }
2598
2599 #[test]
2600 fn test_decimal32() {
2601 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal32(
2602 DECIMAL32_MAX_PRECISION,
2603 7,
2604 ))])
2605 .unwrap();
2606 let col = Arc::new(
2607 Decimal32Array::from_iter([
2608 None,
2609 Some(i32::MIN),
2610 Some(-13),
2611 Some(46_i32),
2612 Some(5456_i32),
2613 Some(i32::MAX),
2614 ])
2615 .with_precision_and_scale(9, 7)
2616 .unwrap(),
2617 ) as ArrayRef;
2618
2619 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2620 for i in 0..rows.num_rows() - 1 {
2621 assert!(rows.row(i) < rows.row(i + 1));
2622 }
2623
2624 let back = converter.convert_rows(&rows).unwrap();
2625 assert_eq!(back.len(), 1);
2626 assert_eq!(col.as_ref(), back[0].as_ref())
2627 }
2628
2629 #[test]
2630 fn test_decimal64() {
2631 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal64(
2632 DECIMAL64_MAX_PRECISION,
2633 7,
2634 ))])
2635 .unwrap();
2636 let col = Arc::new(
2637 Decimal64Array::from_iter([
2638 None,
2639 Some(i64::MIN),
2640 Some(-13),
2641 Some(46_i64),
2642 Some(5456_i64),
2643 Some(i64::MAX),
2644 ])
2645 .with_precision_and_scale(18, 7)
2646 .unwrap(),
2647 ) as ArrayRef;
2648
2649 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2650 for i in 0..rows.num_rows() - 1 {
2651 assert!(rows.row(i) < rows.row(i + 1));
2652 }
2653
2654 let back = converter.convert_rows(&rows).unwrap();
2655 assert_eq!(back.len(), 1);
2656 assert_eq!(col.as_ref(), back[0].as_ref())
2657 }
2658
2659 #[test]
2660 fn test_decimal128() {
2661 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal128(
2662 DECIMAL128_MAX_PRECISION,
2663 7,
2664 ))])
2665 .unwrap();
2666 let col = Arc::new(
2667 Decimal128Array::from_iter([
2668 None,
2669 Some(i128::MIN),
2670 Some(-13),
2671 Some(46_i128),
2672 Some(5456_i128),
2673 Some(i128::MAX),
2674 ])
2675 .with_precision_and_scale(38, 7)
2676 .unwrap(),
2677 ) as ArrayRef;
2678
2679 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2680 for i in 0..rows.num_rows() - 1 {
2681 assert!(rows.row(i) < rows.row(i + 1));
2682 }
2683
2684 let back = converter.convert_rows(&rows).unwrap();
2685 assert_eq!(back.len(), 1);
2686 assert_eq!(col.as_ref(), back[0].as_ref())
2687 }
2688
2689 #[test]
2690 fn test_decimal256() {
2691 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal256(
2692 DECIMAL256_MAX_PRECISION,
2693 7,
2694 ))])
2695 .unwrap();
2696 let col = Arc::new(
2697 Decimal256Array::from_iter([
2698 None,
2699 Some(i256::MIN),
2700 Some(i256::from_parts(0, -1)),
2701 Some(i256::from_parts(u128::MAX, -1)),
2702 Some(i256::from_parts(u128::MAX, 0)),
2703 Some(i256::from_parts(0, 46_i128)),
2704 Some(i256::from_parts(5, 46_i128)),
2705 Some(i256::MAX),
2706 ])
2707 .with_precision_and_scale(DECIMAL256_MAX_PRECISION, 7)
2708 .unwrap(),
2709 ) as ArrayRef;
2710
2711 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2712 for i in 0..rows.num_rows() - 1 {
2713 assert!(rows.row(i) < rows.row(i + 1));
2714 }
2715
2716 let back = converter.convert_rows(&rows).unwrap();
2717 assert_eq!(back.len(), 1);
2718 assert_eq!(col.as_ref(), back[0].as_ref())
2719 }
2720
2721 #[test]
2722 fn test_bool() {
2723 let converter = RowConverter::new(vec![SortField::new(DataType::Boolean)]).unwrap();
2724
2725 let col = Arc::new(BooleanArray::from_iter([None, Some(false), Some(true)])) as ArrayRef;
2726
2727 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2728 assert!(rows.row(2) > rows.row(1));
2729 assert!(rows.row(2) > rows.row(0));
2730 assert!(rows.row(1) > rows.row(0));
2731
2732 let cols = converter.convert_rows(&rows).unwrap();
2733 assert_eq!(&cols[0], &col);
2734
2735 let converter = RowConverter::new(vec![SortField::new_with_options(
2736 DataType::Boolean,
2737 SortOptions::default().desc().with_nulls_first(false),
2738 )])
2739 .unwrap();
2740
2741 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2742 assert!(rows.row(2) < rows.row(1));
2743 assert!(rows.row(2) < rows.row(0));
2744 assert!(rows.row(1) < rows.row(0));
2745 let cols = converter.convert_rows(&rows).unwrap();
2746 assert_eq!(&cols[0], &col);
2747 }
2748
2749 #[test]
2750 fn test_timezone() {
2751 let a =
2752 TimestampNanosecondArray::from(vec![1, 2, 3, 4, 5]).with_timezone("+01:00".to_string());
2753 let d = a.data_type().clone();
2754
2755 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
2756 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
2757 let back = converter.convert_rows(&rows).unwrap();
2758 assert_eq!(back.len(), 1);
2759 assert_eq!(back[0].data_type(), &d);
2760
2761 let mut a = PrimitiveDictionaryBuilder::<Int32Type, TimestampNanosecondType>::new();
2763 a.append(34).unwrap();
2764 a.append_null();
2765 a.append(345).unwrap();
2766
2767 let dict = a.finish();
2769 let values = TimestampNanosecondArray::from(dict.values().to_data());
2770 let dict_with_tz = dict.with_values(Arc::new(values.with_timezone("+02:00")));
2771 let v = DataType::Timestamp(TimeUnit::Nanosecond, Some("+02:00".into()));
2772 let d = DataType::Dictionary(Box::new(DataType::Int32), Box::new(v.clone()));
2773
2774 assert_eq!(dict_with_tz.data_type(), &d);
2775 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
2776 let rows = converter
2777 .convert_columns(&[Arc::new(dict_with_tz) as _])
2778 .unwrap();
2779 let back = converter.convert_rows(&rows).unwrap();
2780 assert_eq!(back.len(), 1);
2781 assert_eq!(back[0].data_type(), &v);
2782 }
2783
2784 #[test]
2785 fn test_null_encoding() {
2786 let col = Arc::new(NullArray::new(10));
2787 let converter = RowConverter::new(vec![SortField::new(DataType::Null)]).unwrap();
2788 let rows = converter.convert_columns(&[col]).unwrap();
2789 assert_eq!(rows.num_rows(), 10);
2790 assert_eq!(rows.row(1).data.len(), 2);
2792 }
2793
2794 #[test]
2795 fn test_variable_width() {
2796 let col = Arc::new(StringArray::from_iter([
2797 Some("hello"),
2798 Some("he"),
2799 None,
2800 Some("foo"),
2801 Some(""),
2802 ])) as ArrayRef;
2803
2804 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
2805 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2806
2807 assert!(rows.row(1) < rows.row(0));
2808 assert!(rows.row(2) < rows.row(4));
2809 assert!(rows.row(3) < rows.row(0));
2810 assert!(rows.row(3) < rows.row(1));
2811
2812 let cols = converter.convert_rows(&rows).unwrap();
2813 assert_eq!(&cols[0], &col);
2814
2815 let col = Arc::new(BinaryArray::from_iter([
2816 None,
2817 Some(vec![0_u8; 0]),
2818 Some(vec![0_u8; 6]),
2819 Some(vec![0_u8; variable::MINI_BLOCK_SIZE]),
2820 Some(vec![0_u8; variable::MINI_BLOCK_SIZE + 1]),
2821 Some(vec![0_u8; variable::BLOCK_SIZE]),
2822 Some(vec![0_u8; variable::BLOCK_SIZE + 1]),
2823 Some(vec![1_u8; 6]),
2824 Some(vec![1_u8; variable::MINI_BLOCK_SIZE]),
2825 Some(vec![1_u8; variable::MINI_BLOCK_SIZE + 1]),
2826 Some(vec![1_u8; variable::BLOCK_SIZE]),
2827 Some(vec![1_u8; variable::BLOCK_SIZE + 1]),
2828 Some(vec![0xFF_u8; 6]),
2829 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE]),
2830 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE + 1]),
2831 Some(vec![0xFF_u8; variable::BLOCK_SIZE]),
2832 Some(vec![0xFF_u8; variable::BLOCK_SIZE + 1]),
2833 ])) as ArrayRef;
2834
2835 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
2836 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2837
2838 for i in 0..rows.num_rows() {
2839 for j in i + 1..rows.num_rows() {
2840 assert!(
2841 rows.row(i) < rows.row(j),
2842 "{} < {} - {:?} < {:?}",
2843 i,
2844 j,
2845 rows.row(i),
2846 rows.row(j)
2847 );
2848 }
2849 }
2850
2851 let cols = converter.convert_rows(&rows).unwrap();
2852 assert_eq!(&cols[0], &col);
2853
2854 let converter = RowConverter::new(vec![SortField::new_with_options(
2855 DataType::Binary,
2856 SortOptions::default().desc().with_nulls_first(false),
2857 )])
2858 .unwrap();
2859 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2860
2861 for i in 0..rows.num_rows() {
2862 for j in i + 1..rows.num_rows() {
2863 assert!(
2864 rows.row(i) > rows.row(j),
2865 "{} > {} - {:?} > {:?}",
2866 i,
2867 j,
2868 rows.row(i),
2869 rows.row(j)
2870 );
2871 }
2872 }
2873
2874 let cols = converter.convert_rows(&rows).unwrap();
2875 assert_eq!(&cols[0], &col);
2876 }
2877
2878 fn dictionary_eq(a: &dyn Array, b: &dyn Array) {
2880 match b.data_type() {
2881 DataType::Dictionary(_, v) => {
2882 assert_eq!(a.data_type(), v.as_ref());
2883 let b = arrow_cast::cast(b, v).unwrap();
2884 assert_eq!(a, b.as_ref())
2885 }
2886 _ => assert_eq!(a, b),
2887 }
2888 }
2889
2890 #[test]
2891 fn test_string_dictionary() {
2892 let a = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2893 Some("foo"),
2894 Some("hello"),
2895 Some("he"),
2896 None,
2897 Some("hello"),
2898 Some(""),
2899 Some("hello"),
2900 Some("hello"),
2901 ])) as ArrayRef;
2902
2903 let field = SortField::new(a.data_type().clone());
2904 let converter = RowConverter::new(vec![field]).unwrap();
2905 let rows_a = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2906
2907 assert!(rows_a.row(3) < rows_a.row(5));
2908 assert!(rows_a.row(2) < rows_a.row(1));
2909 assert!(rows_a.row(0) < rows_a.row(1));
2910 assert!(rows_a.row(3) < rows_a.row(0));
2911
2912 assert_eq!(rows_a.row(1), rows_a.row(4));
2913 assert_eq!(rows_a.row(1), rows_a.row(6));
2914 assert_eq!(rows_a.row(1), rows_a.row(7));
2915
2916 let cols = converter.convert_rows(&rows_a).unwrap();
2917 dictionary_eq(&cols[0], &a);
2918
2919 let b = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2920 Some("hello"),
2921 None,
2922 Some("cupcakes"),
2923 ])) as ArrayRef;
2924
2925 let rows_b = converter.convert_columns(&[Arc::clone(&b)]).unwrap();
2926 assert_eq!(rows_a.row(1), rows_b.row(0));
2927 assert_eq!(rows_a.row(3), rows_b.row(1));
2928 assert!(rows_b.row(2) < rows_a.row(0));
2929
2930 let cols = converter.convert_rows(&rows_b).unwrap();
2931 dictionary_eq(&cols[0], &b);
2932
2933 let converter = RowConverter::new(vec![SortField::new_with_options(
2934 a.data_type().clone(),
2935 SortOptions::default().desc().with_nulls_first(false),
2936 )])
2937 .unwrap();
2938
2939 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2940 assert!(rows_c.row(3) > rows_c.row(5));
2941 assert!(rows_c.row(2) > rows_c.row(1));
2942 assert!(rows_c.row(0) > rows_c.row(1));
2943 assert!(rows_c.row(3) > rows_c.row(0));
2944
2945 let cols = converter.convert_rows(&rows_c).unwrap();
2946 dictionary_eq(&cols[0], &a);
2947
2948 let converter = RowConverter::new(vec![SortField::new_with_options(
2949 a.data_type().clone(),
2950 SortOptions::default().desc().with_nulls_first(true),
2951 )])
2952 .unwrap();
2953
2954 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2955 assert!(rows_c.row(3) < rows_c.row(5));
2956 assert!(rows_c.row(2) > rows_c.row(1));
2957 assert!(rows_c.row(0) > rows_c.row(1));
2958 assert!(rows_c.row(3) < rows_c.row(0));
2959
2960 let cols = converter.convert_rows(&rows_c).unwrap();
2961 dictionary_eq(&cols[0], &a);
2962 }
2963
2964 #[test]
2965 fn test_struct() {
2966 let a = Arc::new(Int32Array::from(vec![1, 1, 2, 2])) as ArrayRef;
2968 let a_f = Arc::new(Field::new("int", DataType::Int32, false));
2969 let u = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as ArrayRef;
2970 let u_f = Arc::new(Field::new("s", DataType::Utf8, false));
2971 let s1 = Arc::new(StructArray::from(vec![(a_f, a), (u_f, u)])) as ArrayRef;
2972
2973 let sort_fields = vec![SortField::new(s1.data_type().clone())];
2974 let converter = RowConverter::new(sort_fields).unwrap();
2975 let r1 = converter.convert_columns(&[Arc::clone(&s1)]).unwrap();
2976
2977 for (a, b) in r1.iter().zip(r1.iter().skip(1)) {
2978 assert!(a < b);
2979 }
2980
2981 let back = converter.convert_rows(&r1).unwrap();
2982 assert_eq!(back.len(), 1);
2983 assert_eq!(&back[0], &s1);
2984
2985 let data = s1
2987 .to_data()
2988 .into_builder()
2989 .null_bit_buffer(Some(Buffer::from_slice_ref([0b00001010])))
2990 .null_count(2)
2991 .build()
2992 .unwrap();
2993
2994 let s2 = Arc::new(StructArray::from(data)) as ArrayRef;
2995 let r2 = converter.convert_columns(&[Arc::clone(&s2)]).unwrap();
2996 assert_eq!(r2.row(0), r2.row(2)); assert!(r2.row(0) < r2.row(1)); assert_ne!(r1.row(0), r2.row(0)); assert_eq!(r1.row(1), r2.row(1)); let back = converter.convert_rows(&r2).unwrap();
3002 assert_eq!(back.len(), 1);
3003 assert_eq!(&back[0], &s2);
3004
3005 back[0].to_data().validate_full().unwrap();
3006 }
3007
3008 #[test]
3009 fn test_dictionary_in_struct() {
3010 let builder = StringDictionaryBuilder::<Int32Type>::new();
3011 let mut struct_builder = StructBuilder::new(
3012 vec![Field::new_dictionary(
3013 "foo",
3014 DataType::Int32,
3015 DataType::Utf8,
3016 true,
3017 )],
3018 vec![Box::new(builder)],
3019 );
3020
3021 let dict_builder = struct_builder
3022 .field_builder::<StringDictionaryBuilder<Int32Type>>(0)
3023 .unwrap();
3024
3025 dict_builder.append_value("a");
3027 dict_builder.append_null();
3028 dict_builder.append_value("a");
3029 dict_builder.append_value("b");
3030
3031 for _ in 0..4 {
3032 struct_builder.append(true);
3033 }
3034
3035 let s = Arc::new(struct_builder.finish()) as ArrayRef;
3036 let sort_fields = vec![SortField::new(s.data_type().clone())];
3037 let converter = RowConverter::new(sort_fields).unwrap();
3038 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3039
3040 let back = converter.convert_rows(&r).unwrap();
3041 let [s2] = back.try_into().unwrap();
3042
3043 assert_ne!(&s.data_type(), &s2.data_type());
3046 s2.to_data().validate_full().unwrap();
3047
3048 let s1_struct = s.as_struct();
3052 let s1_0 = s1_struct.column(0);
3053 let s1_idx_0 = s1_0.as_dictionary::<Int32Type>();
3054 let keys = s1_idx_0.keys();
3055 let values = s1_idx_0.values().as_string::<i32>();
3056 let s2_struct = s2.as_struct();
3058 let s2_0 = s2_struct.column(0);
3059 let s2_idx_0 = s2_0.as_string::<i32>();
3060
3061 for i in 0..keys.len() {
3062 if keys.is_null(i) {
3063 assert!(s2_idx_0.is_null(i));
3064 } else {
3065 let dict_index = keys.value(i) as usize;
3066 assert_eq!(values.value(dict_index), s2_idx_0.value(i));
3067 }
3068 }
3069 }
3070
3071 #[test]
3072 fn test_dictionary_in_struct_empty() {
3073 let ty = DataType::Struct(
3074 vec![Field::new_dictionary(
3075 "foo",
3076 DataType::Int32,
3077 DataType::Int32,
3078 false,
3079 )]
3080 .into(),
3081 );
3082 let s = arrow_array::new_empty_array(&ty);
3083
3084 let sort_fields = vec![SortField::new(s.data_type().clone())];
3085 let converter = RowConverter::new(sort_fields).unwrap();
3086 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3087
3088 let back = converter.convert_rows(&r).unwrap();
3089 let [s2] = back.try_into().unwrap();
3090
3091 assert_ne!(&s.data_type(), &s2.data_type());
3094 s2.to_data().validate_full().unwrap();
3095 assert_eq!(s.len(), 0);
3096 assert_eq!(s2.len(), 0);
3097 }
3098
3099 #[test]
3100 fn test_list_of_string_dictionary() {
3101 let mut builder = ListBuilder::<StringDictionaryBuilder<Int32Type>>::default();
3102 builder.values().append("a").unwrap();
3104 builder.values().append("b").unwrap();
3105 builder.values().append("zero").unwrap();
3106 builder.values().append_null();
3107 builder.values().append("c").unwrap();
3108 builder.values().append("b").unwrap();
3109 builder.values().append("d").unwrap();
3110 builder.append(true);
3111 builder.append(false);
3113 builder.values().append("e").unwrap();
3115 builder.values().append("zero").unwrap();
3116 builder.values().append("a").unwrap();
3117 builder.append(true);
3118
3119 let a = Arc::new(builder.finish()) as ArrayRef;
3120 let data_type = a.data_type().clone();
3121
3122 let field = SortField::new(data_type.clone());
3123 let converter = RowConverter::new(vec![field]).unwrap();
3124 let rows = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
3125
3126 let back = converter.convert_rows(&rows).unwrap();
3127 assert_eq!(back.len(), 1);
3128 let [a2] = back.try_into().unwrap();
3129
3130 assert_ne!(&a.data_type(), &a2.data_type());
3133
3134 a2.to_data().validate_full().unwrap();
3135
3136 let a2_list = a2.as_list::<i32>();
3137 let a1_list = a.as_list::<i32>();
3138
3139 let a1_0 = a1_list.value(0);
3142 let a1_idx_0 = a1_0.as_dictionary::<Int32Type>();
3143 let keys = a1_idx_0.keys();
3144 let values = a1_idx_0.values().as_string::<i32>();
3145 let a2_0 = a2_list.value(0);
3146 let a2_idx_0 = a2_0.as_string::<i32>();
3147
3148 for i in 0..keys.len() {
3149 if keys.is_null(i) {
3150 assert!(a2_idx_0.is_null(i));
3151 } else {
3152 let dict_index = keys.value(i) as usize;
3153 assert_eq!(values.value(dict_index), a2_idx_0.value(i));
3154 }
3155 }
3156
3157 assert!(a1_list.is_null(1));
3159 assert!(a2_list.is_null(1));
3160
3161 let a1_2 = a1_list.value(2);
3163 let a1_idx_2 = a1_2.as_dictionary::<Int32Type>();
3164 let keys = a1_idx_2.keys();
3165 let values = a1_idx_2.values().as_string::<i32>();
3166 let a2_2 = a2_list.value(2);
3167 let a2_idx_2 = a2_2.as_string::<i32>();
3168
3169 for i in 0..keys.len() {
3170 if keys.is_null(i) {
3171 assert!(a2_idx_2.is_null(i));
3172 } else {
3173 let dict_index = keys.value(i) as usize;
3174 assert_eq!(values.value(dict_index), a2_idx_2.value(i));
3175 }
3176 }
3177 }
3178
3179 #[test]
3180 fn test_primitive_dictionary() {
3181 let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
3182 builder.append(2).unwrap();
3183 builder.append(3).unwrap();
3184 builder.append(0).unwrap();
3185 builder.append_null();
3186 builder.append(5).unwrap();
3187 builder.append(3).unwrap();
3188 builder.append(-1).unwrap();
3189
3190 let a = builder.finish();
3191 let data_type = a.data_type().clone();
3192 let columns = [Arc::new(a) as ArrayRef];
3193
3194 let field = SortField::new(data_type.clone());
3195 let converter = RowConverter::new(vec![field]).unwrap();
3196 let rows = converter.convert_columns(&columns).unwrap();
3197 assert!(rows.row(0) < rows.row(1));
3198 assert!(rows.row(2) < rows.row(0));
3199 assert!(rows.row(3) < rows.row(2));
3200 assert!(rows.row(6) < rows.row(2));
3201 assert!(rows.row(3) < rows.row(6));
3202
3203 let back = converter.convert_rows(&rows).unwrap();
3204 assert_eq!(back.len(), 1);
3205 back[0].to_data().validate_full().unwrap();
3206 }
3207
3208 #[test]
3209 fn test_dictionary_nulls() {
3210 let values = Int32Array::from_iter([Some(1), Some(-1), None, Some(4), None]).into_data();
3211 let keys =
3212 Int32Array::from_iter([Some(0), Some(0), Some(1), Some(2), Some(4), None]).into_data();
3213
3214 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3215 let data = keys
3216 .into_builder()
3217 .data_type(data_type.clone())
3218 .child_data(vec![values])
3219 .build()
3220 .unwrap();
3221
3222 let columns = [Arc::new(DictionaryArray::<Int32Type>::from(data)) as ArrayRef];
3223 let field = SortField::new(data_type.clone());
3224 let converter = RowConverter::new(vec![field]).unwrap();
3225 let rows = converter.convert_columns(&columns).unwrap();
3226
3227 assert_eq!(rows.row(0), rows.row(1));
3228 assert_eq!(rows.row(3), rows.row(4));
3229 assert_eq!(rows.row(4), rows.row(5));
3230 assert!(rows.row(3) < rows.row(0));
3231 }
3232
3233 #[test]
3234 fn test_from_binary_shared_buffer() {
3235 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3236 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3237 let rows = converter.convert_columns(&[array]).unwrap();
3238 let binary_rows = rows.try_into_binary().expect("known-small rows");
3239 let _binary_rows_shared_buffer = binary_rows.clone();
3240
3241 let parsed = converter.from_binary(binary_rows);
3242
3243 converter.convert_rows(parsed.iter()).unwrap();
3244 }
3245
3246 #[test]
3247 #[should_panic(expected = "Encountered non UTF-8 data")]
3248 fn test_invalid_utf8() {
3249 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3250 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3251 let rows = converter.convert_columns(&[array]).unwrap();
3252 let binary_row = rows.row(0);
3253
3254 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3255 let parser = converter.parser();
3256 let utf8_row = parser.parse(binary_row.as_ref());
3257
3258 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3259 }
3260
3261 #[test]
3262 #[should_panic(expected = "Encountered non UTF-8 data")]
3263 fn test_invalid_utf8_array() {
3264 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3265 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3266 let rows = converter.convert_columns(&[array]).unwrap();
3267 let binary_rows = rows.try_into_binary().expect("known-small rows");
3268
3269 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3270 let parsed = converter.from_binary(binary_rows);
3271
3272 converter.convert_rows(parsed.iter()).unwrap();
3273 }
3274
3275 #[test]
3276 #[should_panic(expected = "index out of bounds")]
3277 fn test_invalid_empty() {
3278 let binary_row: &[u8] = &[];
3279
3280 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3281 let parser = converter.parser();
3282 let utf8_row = parser.parse(binary_row.as_ref());
3283
3284 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3285 }
3286
3287 #[test]
3288 #[should_panic(expected = "index out of bounds")]
3289 fn test_invalid_empty_array() {
3290 let row: &[u8] = &[];
3291 let binary_rows = BinaryArray::from(vec![row]);
3292
3293 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3294 let parsed = converter.from_binary(binary_rows);
3295
3296 converter.convert_rows(parsed.iter()).unwrap();
3297 }
3298
3299 #[test]
3300 #[should_panic(expected = "index out of bounds")]
3301 fn test_invalid_truncated() {
3302 let binary_row: &[u8] = &[0x02];
3303
3304 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3305 let parser = converter.parser();
3306 let utf8_row = parser.parse(binary_row.as_ref());
3307
3308 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3309 }
3310
3311 #[test]
3312 #[should_panic(expected = "index out of bounds")]
3313 fn test_invalid_truncated_array() {
3314 let row: &[u8] = &[0x02];
3315 let binary_rows = BinaryArray::from(vec![row]);
3316
3317 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3318 let parsed = converter.from_binary(binary_rows);
3319
3320 converter.convert_rows(parsed.iter()).unwrap();
3321 }
3322
3323 #[test]
3324 #[should_panic(expected = "rows were not produced by this RowConverter")]
3325 fn test_different_converter() {
3326 let values = Arc::new(Int32Array::from_iter([Some(1), Some(-1)]));
3327 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3328 let rows = converter.convert_columns(&[values]).unwrap();
3329
3330 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3331 let _ = converter.convert_rows(&rows);
3332 }
3333
3334 fn test_single_list<O: OffsetSizeTrait>() {
3335 let mut builder = GenericListBuilder::<O, _>::new(Int32Builder::new());
3336 builder.values().append_value(32);
3337 builder.values().append_value(52);
3338 builder.values().append_value(32);
3339 builder.append(true);
3340 builder.values().append_value(32);
3341 builder.values().append_value(52);
3342 builder.values().append_value(12);
3343 builder.append(true);
3344 builder.values().append_value(32);
3345 builder.values().append_value(52);
3346 builder.append(true);
3347 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3350 builder.values().append_value(32);
3351 builder.values().append_null();
3352 builder.append(true);
3353 builder.append(true);
3354 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3357
3358 let list = Arc::new(builder.finish()) as ArrayRef;
3359 let d = list.data_type().clone();
3360
3361 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3362
3363 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3364 assert!(rows.row(0) > rows.row(1)); assert!(rows.row(2) < rows.row(1)); assert!(rows.row(3) < rows.row(2)); assert!(rows.row(4) < rows.row(2)); assert!(rows.row(5) < rows.row(2)); assert!(rows.row(3) < rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3373 assert_eq!(back.len(), 1);
3374 back[0].to_data().validate_full().unwrap();
3375 assert_eq!(&back[0], &list);
3376
3377 let options = SortOptions::default().asc().with_nulls_first(false);
3378 let field = SortField::new_with_options(d.clone(), options);
3379 let converter = RowConverter::new(vec![field]).unwrap();
3380 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3381
3382 assert!(rows.row(0) > rows.row(1)); assert!(rows.row(2) < rows.row(1)); assert!(rows.row(3) > rows.row(2)); assert!(rows.row(4) > rows.row(2)); assert!(rows.row(5) < rows.row(2)); assert!(rows.row(3) > rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3391 assert_eq!(back.len(), 1);
3392 back[0].to_data().validate_full().unwrap();
3393 assert_eq!(&back[0], &list);
3394
3395 let options = SortOptions::default().desc().with_nulls_first(false);
3396 let field = SortField::new_with_options(d.clone(), options);
3397 let converter = RowConverter::new(vec![field]).unwrap();
3398 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3399
3400 assert!(rows.row(0) < rows.row(1)); assert!(rows.row(2) > rows.row(1)); assert!(rows.row(3) > rows.row(2)); assert!(rows.row(4) > rows.row(2)); assert!(rows.row(5) > rows.row(2)); assert!(rows.row(3) > rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3409 assert_eq!(back.len(), 1);
3410 back[0].to_data().validate_full().unwrap();
3411 assert_eq!(&back[0], &list);
3412
3413 let options = SortOptions::default().desc().with_nulls_first(true);
3414 let field = SortField::new_with_options(d, options);
3415 let converter = RowConverter::new(vec![field]).unwrap();
3416 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3417
3418 assert!(rows.row(0) < rows.row(1)); assert!(rows.row(2) > rows.row(1)); assert!(rows.row(3) < rows.row(2)); assert!(rows.row(4) < rows.row(2)); assert!(rows.row(5) > rows.row(2)); assert!(rows.row(3) < rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3427 assert_eq!(back.len(), 1);
3428 back[0].to_data().validate_full().unwrap();
3429 assert_eq!(&back[0], &list);
3430
3431 let sliced_list = list.slice(1, 5);
3432 let rows_on_sliced_list = converter
3433 .convert_columns(&[Arc::clone(&sliced_list)])
3434 .unwrap();
3435
3436 assert!(rows_on_sliced_list.row(1) > rows_on_sliced_list.row(0)); assert!(rows_on_sliced_list.row(2) < rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(3) < rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(4) > rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(2) < rows_on_sliced_list.row(4)); let back = converter.convert_rows(&rows_on_sliced_list).unwrap();
3443 assert_eq!(back.len(), 1);
3444 back[0].to_data().validate_full().unwrap();
3445 assert_eq!(&back[0], &sliced_list);
3446 }
3447
3448 fn test_nested_list<O: OffsetSizeTrait>() {
3449 let mut builder =
3450 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(Int32Builder::new()));
3451
3452 builder.values().values().append_value(1);
3453 builder.values().values().append_value(2);
3454 builder.values().append(true);
3455 builder.values().values().append_value(1);
3456 builder.values().values().append_null();
3457 builder.values().append(true);
3458 builder.append(true);
3459
3460 builder.values().values().append_value(1);
3461 builder.values().values().append_null();
3462 builder.values().append(true);
3463 builder.values().values().append_value(1);
3464 builder.values().values().append_null();
3465 builder.values().append(true);
3466 builder.append(true);
3467
3468 builder.values().values().append_value(1);
3469 builder.values().values().append_null();
3470 builder.values().append(true);
3471 builder.values().append(false);
3472 builder.append(true);
3473 builder.append(false);
3474
3475 builder.values().values().append_value(1);
3476 builder.values().values().append_value(2);
3477 builder.values().append(true);
3478 builder.append(true);
3479
3480 let list = Arc::new(builder.finish()) as ArrayRef;
3481 let d = list.data_type().clone();
3482
3483 let options = SortOptions::default().asc().with_nulls_first(true);
3491 let field = SortField::new_with_options(d.clone(), options);
3492 let converter = RowConverter::new(vec![field]).unwrap();
3493 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3494
3495 assert!(rows.row(0) > rows.row(1));
3496 assert!(rows.row(1) > rows.row(2));
3497 assert!(rows.row(2) > rows.row(3));
3498 assert!(rows.row(4) < rows.row(0));
3499 assert!(rows.row(4) > rows.row(1));
3500
3501 let back = converter.convert_rows(&rows).unwrap();
3502 assert_eq!(back.len(), 1);
3503 back[0].to_data().validate_full().unwrap();
3504 assert_eq!(&back[0], &list);
3505
3506 let options = SortOptions::default().desc().with_nulls_first(true);
3507 let field = SortField::new_with_options(d.clone(), options);
3508 let converter = RowConverter::new(vec![field]).unwrap();
3509 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3510
3511 assert!(rows.row(0) > rows.row(1));
3512 assert!(rows.row(1) > rows.row(2));
3513 assert!(rows.row(2) > rows.row(3));
3514 assert!(rows.row(4) > rows.row(0));
3515 assert!(rows.row(4) > rows.row(1));
3516
3517 let back = converter.convert_rows(&rows).unwrap();
3518 assert_eq!(back.len(), 1);
3519 back[0].to_data().validate_full().unwrap();
3520 assert_eq!(&back[0], &list);
3521
3522 let options = SortOptions::default().desc().with_nulls_first(false);
3523 let field = SortField::new_with_options(d, options);
3524 let converter = RowConverter::new(vec![field]).unwrap();
3525 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3526
3527 assert!(rows.row(0) < rows.row(1));
3528 assert!(rows.row(1) < rows.row(2));
3529 assert!(rows.row(2) < rows.row(3));
3530 assert!(rows.row(4) > rows.row(0));
3531 assert!(rows.row(4) < rows.row(1));
3532
3533 let back = converter.convert_rows(&rows).unwrap();
3534 assert_eq!(back.len(), 1);
3535 back[0].to_data().validate_full().unwrap();
3536 assert_eq!(&back[0], &list);
3537
3538 let sliced_list = list.slice(1, 3);
3539 let rows = converter
3540 .convert_columns(&[Arc::clone(&sliced_list)])
3541 .unwrap();
3542
3543 assert!(rows.row(0) < rows.row(1));
3544 assert!(rows.row(1) < rows.row(2));
3545
3546 let back = converter.convert_rows(&rows).unwrap();
3547 assert_eq!(back.len(), 1);
3548 back[0].to_data().validate_full().unwrap();
3549 assert_eq!(&back[0], &sliced_list);
3550 }
3551
3552 #[test]
3553 fn test_list() {
3554 test_single_list::<i32>();
3555 test_nested_list::<i32>();
3556 }
3557
3558 #[test]
3559 fn test_large_list() {
3560 test_single_list::<i64>();
3561 test_nested_list::<i64>();
3562 }
3563
3564 fn test_single_list_view<O: OffsetSizeTrait>() {
3565 let mut builder = GenericListViewBuilder::<O, _>::new(Int32Builder::new());
3566 builder.values().append_value(32);
3567 builder.values().append_value(52);
3568 builder.values().append_value(32);
3569 builder.append(true);
3570 builder.values().append_value(32);
3571 builder.values().append_value(52);
3572 builder.values().append_value(12);
3573 builder.append(true);
3574 builder.values().append_value(32);
3575 builder.values().append_value(52);
3576 builder.append(true);
3577 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3580 builder.values().append_value(32);
3581 builder.values().append_null();
3582 builder.append(true);
3583 builder.append(true);
3584 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3587
3588 let list = Arc::new(builder.finish()) as ArrayRef;
3589 let d = list.data_type().clone();
3590
3591 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3592
3593 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3594 assert!(rows.row(0) > rows.row(1)); assert!(rows.row(2) < rows.row(1)); assert!(rows.row(3) < rows.row(2)); assert!(rows.row(4) < rows.row(2)); assert!(rows.row(5) < rows.row(2)); assert!(rows.row(3) < rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3603 assert_eq!(back.len(), 1);
3604 back[0].to_data().validate_full().unwrap();
3605
3606 let back_list_view = back[0]
3608 .as_any()
3609 .downcast_ref::<GenericListViewArray<O>>()
3610 .unwrap();
3611 let orig_list_view = list
3612 .as_any()
3613 .downcast_ref::<GenericListViewArray<O>>()
3614 .unwrap();
3615
3616 assert_eq!(back_list_view.len(), orig_list_view.len());
3617 for i in 0..back_list_view.len() {
3618 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3619 if back_list_view.is_valid(i) {
3620 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3621 }
3622 }
3623
3624 let options = SortOptions::default().asc().with_nulls_first(false);
3625 let field = SortField::new_with_options(d.clone(), options);
3626 let converter = RowConverter::new(vec![field]).unwrap();
3627 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3628
3629 assert!(rows.row(0) > rows.row(1)); assert!(rows.row(2) < rows.row(1)); assert!(rows.row(3) > rows.row(2)); assert!(rows.row(4) > rows.row(2)); assert!(rows.row(5) < rows.row(2)); assert!(rows.row(3) > rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3638 assert_eq!(back.len(), 1);
3639 back[0].to_data().validate_full().unwrap();
3640
3641 let options = SortOptions::default().desc().with_nulls_first(false);
3642 let field = SortField::new_with_options(d.clone(), options);
3643 let converter = RowConverter::new(vec![field]).unwrap();
3644 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3645
3646 assert!(rows.row(0) < rows.row(1)); assert!(rows.row(2) > rows.row(1)); assert!(rows.row(3) > rows.row(2)); assert!(rows.row(4) > rows.row(2)); assert!(rows.row(5) > rows.row(2)); assert!(rows.row(3) > rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3655 assert_eq!(back.len(), 1);
3656 back[0].to_data().validate_full().unwrap();
3657
3658 let options = SortOptions::default().desc().with_nulls_first(true);
3659 let field = SortField::new_with_options(d, options);
3660 let converter = RowConverter::new(vec![field]).unwrap();
3661 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3662
3663 assert!(rows.row(0) < rows.row(1)); assert!(rows.row(2) > rows.row(1)); assert!(rows.row(3) < rows.row(2)); assert!(rows.row(4) < rows.row(2)); assert!(rows.row(5) > rows.row(2)); assert!(rows.row(3) < rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
3672 assert_eq!(back.len(), 1);
3673 back[0].to_data().validate_full().unwrap();
3674
3675 let sliced_list = list.slice(1, 5);
3676 let rows_on_sliced_list = converter
3677 .convert_columns(&[Arc::clone(&sliced_list)])
3678 .unwrap();
3679
3680 assert!(rows_on_sliced_list.row(1) > rows_on_sliced_list.row(0)); assert!(rows_on_sliced_list.row(2) < rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(3) < rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(4) > rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(2) < rows_on_sliced_list.row(4)); let back = converter.convert_rows(&rows_on_sliced_list).unwrap();
3687 assert_eq!(back.len(), 1);
3688 back[0].to_data().validate_full().unwrap();
3689 }
3690
3691 fn test_nested_list_view<O: OffsetSizeTrait>() {
3692 let mut builder = GenericListViewBuilder::<O, _>::new(GenericListViewBuilder::<O, _>::new(
3693 Int32Builder::new(),
3694 ));
3695
3696 builder.values().values().append_value(1);
3698 builder.values().values().append_value(2);
3699 builder.values().append(true);
3700 builder.values().values().append_value(1);
3701 builder.values().values().append_null();
3702 builder.values().append(true);
3703 builder.append(true);
3704
3705 builder.values().values().append_value(1);
3707 builder.values().values().append_null();
3708 builder.values().append(true);
3709 builder.values().values().append_value(1);
3710 builder.values().values().append_null();
3711 builder.values().append(true);
3712 builder.append(true);
3713
3714 builder.values().values().append_value(1);
3716 builder.values().values().append_null();
3717 builder.values().append(true);
3718 builder.values().append(false);
3719 builder.append(true);
3720
3721 builder.append(false);
3723
3724 builder.values().values().append_value(1);
3726 builder.values().values().append_value(2);
3727 builder.values().append(true);
3728 builder.append(true);
3729
3730 let list = Arc::new(builder.finish()) as ArrayRef;
3731 let d = list.data_type().clone();
3732
3733 let options = SortOptions::default().asc().with_nulls_first(true);
3741 let field = SortField::new_with_options(d.clone(), options);
3742 let converter = RowConverter::new(vec![field]).unwrap();
3743 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3744
3745 assert!(rows.row(0) > rows.row(1));
3746 assert!(rows.row(1) > rows.row(2));
3747 assert!(rows.row(2) > rows.row(3));
3748 assert!(rows.row(4) < rows.row(0));
3749 assert!(rows.row(4) > rows.row(1));
3750
3751 let back = converter.convert_rows(&rows).unwrap();
3752 assert_eq!(back.len(), 1);
3753 back[0].to_data().validate_full().unwrap();
3754
3755 let back_list_view = back[0]
3757 .as_any()
3758 .downcast_ref::<GenericListViewArray<O>>()
3759 .unwrap();
3760 let orig_list_view = list
3761 .as_any()
3762 .downcast_ref::<GenericListViewArray<O>>()
3763 .unwrap();
3764
3765 assert_eq!(back_list_view.len(), orig_list_view.len());
3766 for i in 0..back_list_view.len() {
3767 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3768 if back_list_view.is_valid(i) {
3769 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3770 }
3771 }
3772
3773 let options = SortOptions::default().desc().with_nulls_first(true);
3774 let field = SortField::new_with_options(d.clone(), options);
3775 let converter = RowConverter::new(vec![field]).unwrap();
3776 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3777
3778 assert!(rows.row(0) > rows.row(1));
3779 assert!(rows.row(1) > rows.row(2));
3780 assert!(rows.row(2) > rows.row(3));
3781 assert!(rows.row(4) > rows.row(0));
3782 assert!(rows.row(4) > rows.row(1));
3783
3784 let back = converter.convert_rows(&rows).unwrap();
3785 assert_eq!(back.len(), 1);
3786 back[0].to_data().validate_full().unwrap();
3787
3788 let back_list_view = back[0]
3790 .as_any()
3791 .downcast_ref::<GenericListViewArray<O>>()
3792 .unwrap();
3793
3794 assert_eq!(back_list_view.len(), orig_list_view.len());
3795 for i in 0..back_list_view.len() {
3796 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3797 if back_list_view.is_valid(i) {
3798 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3799 }
3800 }
3801
3802 let options = SortOptions::default().desc().with_nulls_first(false);
3803 let field = SortField::new_with_options(d.clone(), options);
3804 let converter = RowConverter::new(vec![field]).unwrap();
3805 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3806
3807 assert!(rows.row(0) < rows.row(1));
3808 assert!(rows.row(1) < rows.row(2));
3809 assert!(rows.row(2) < rows.row(3));
3810 assert!(rows.row(4) > rows.row(0));
3811 assert!(rows.row(4) < rows.row(1));
3812
3813 let back = converter.convert_rows(&rows).unwrap();
3814 assert_eq!(back.len(), 1);
3815 back[0].to_data().validate_full().unwrap();
3816
3817 let back_list_view = back[0]
3819 .as_any()
3820 .downcast_ref::<GenericListViewArray<O>>()
3821 .unwrap();
3822
3823 assert_eq!(back_list_view.len(), orig_list_view.len());
3824 for i in 0..back_list_view.len() {
3825 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3826 if back_list_view.is_valid(i) {
3827 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3828 }
3829 }
3830
3831 let sliced_list = list.slice(1, 3);
3832 let rows = converter
3833 .convert_columns(&[Arc::clone(&sliced_list)])
3834 .unwrap();
3835
3836 assert!(rows.row(0) < rows.row(1));
3837 assert!(rows.row(1) < rows.row(2));
3838
3839 let back = converter.convert_rows(&rows).unwrap();
3840 assert_eq!(back.len(), 1);
3841 back[0].to_data().validate_full().unwrap();
3842 }
3843
3844 #[test]
3845 fn test_list_view() {
3846 test_single_list_view::<i32>();
3847 test_nested_list_view::<i32>();
3848 }
3849
3850 #[test]
3851 fn test_large_list_view() {
3852 test_single_list_view::<i64>();
3853 test_nested_list_view::<i64>();
3854 }
3855
3856 fn test_list_view_with_shared_values<O: OffsetSizeTrait>() {
3857 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
3859 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
3860
3861 let offsets = ScalarBuffer::<O>::from(vec![
3869 O::from_usize(0).unwrap(),
3870 O::from_usize(0).unwrap(),
3871 O::from_usize(5).unwrap(),
3872 O::from_usize(2).unwrap(),
3873 O::from_usize(1).unwrap(),
3874 O::from_usize(2).unwrap(),
3875 ]);
3876 let sizes = ScalarBuffer::<O>::from(vec![
3877 O::from_usize(3).unwrap(),
3878 O::from_usize(3).unwrap(),
3879 O::from_usize(2).unwrap(),
3880 O::from_usize(2).unwrap(),
3881 O::from_usize(4).unwrap(),
3882 O::from_usize(1).unwrap(),
3883 ]);
3884
3885 let list_view: GenericListViewArray<O> =
3886 GenericListViewArray::try_new(field, offsets, sizes, Arc::new(values), None).unwrap();
3887
3888 let d = list_view.data_type().clone();
3889 let list = Arc::new(list_view) as ArrayRef;
3890
3891 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3892 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3893
3894 assert_eq!(rows.row(0), rows.row(1));
3896
3897 assert!(rows.row(0) < rows.row(2));
3899
3900 assert!(rows.row(3) > rows.row(0));
3902
3903 assert!(rows.row(4) > rows.row(0));
3905
3906 assert!(rows.row(5) < rows.row(3));
3908
3909 assert!(rows.row(5) > rows.row(4));
3911
3912 let back = converter.convert_rows(&rows).unwrap();
3914 assert_eq!(back.len(), 1);
3915 back[0].to_data().validate_full().unwrap();
3916
3917 let back_list_view = back[0]
3919 .as_any()
3920 .downcast_ref::<GenericListViewArray<O>>()
3921 .unwrap();
3922 let orig_list_view = list
3923 .as_any()
3924 .downcast_ref::<GenericListViewArray<O>>()
3925 .unwrap();
3926
3927 assert_eq!(back_list_view.len(), orig_list_view.len());
3928 for i in 0..back_list_view.len() {
3929 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3930 if back_list_view.is_valid(i) {
3931 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3932 }
3933 }
3934
3935 let options = SortOptions::default().desc();
3937 let field = SortField::new_with_options(d, options);
3938 let converter = RowConverter::new(vec![field]).unwrap();
3939 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3940
3941 assert_eq!(rows.row(0), rows.row(1)); assert!(rows.row(0) > rows.row(2)); assert!(rows.row(3) < rows.row(0)); let back = converter.convert_rows(&rows).unwrap();
3947 assert_eq!(back.len(), 1);
3948 back[0].to_data().validate_full().unwrap();
3949 }
3950
3951 #[test]
3952 fn test_list_view_shared_values() {
3953 test_list_view_with_shared_values::<i32>();
3954 }
3955
3956 #[test]
3957 fn test_large_list_view_shared_values() {
3958 test_list_view_with_shared_values::<i64>();
3959 }
3960
3961 #[test]
3962 fn test_fixed_size_list() {
3963 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
3964 builder.values().append_value(32);
3965 builder.values().append_value(52);
3966 builder.values().append_value(32);
3967 builder.append(true);
3968 builder.values().append_value(32);
3969 builder.values().append_value(52);
3970 builder.values().append_value(12);
3971 builder.append(true);
3972 builder.values().append_value(32);
3973 builder.values().append_value(52);
3974 builder.values().append_null();
3975 builder.append(true);
3976 builder.values().append_value(32); builder.values().append_value(52); builder.values().append_value(13); builder.append(false);
3980 builder.values().append_value(32);
3981 builder.values().append_null();
3982 builder.values().append_null();
3983 builder.append(true);
3984 builder.values().append_null();
3985 builder.values().append_null();
3986 builder.values().append_null();
3987 builder.append(true);
3988 builder.values().append_value(17); builder.values().append_null(); builder.values().append_value(77); builder.append(false);
3992
3993 let list = Arc::new(builder.finish()) as ArrayRef;
3994 let d = list.data_type().clone();
3995
3996 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3998
3999 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4000 assert!(rows.row(0) > rows.row(1)); assert!(rows.row(2) < rows.row(1)); assert!(rows.row(3) < rows.row(2)); assert!(rows.row(4) < rows.row(2)); assert!(rows.row(5) < rows.row(2)); assert!(rows.row(3) < rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
4009 assert_eq!(back.len(), 1);
4010 back[0].to_data().validate_full().unwrap();
4011 assert_eq!(&back[0], &list);
4012
4013 let options = SortOptions::default().asc().with_nulls_first(false);
4015 let field = SortField::new_with_options(d.clone(), options);
4016 let converter = RowConverter::new(vec![field]).unwrap();
4017 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4018 assert!(rows.row(0) > rows.row(1)); assert!(rows.row(2) > rows.row(1)); assert!(rows.row(3) > rows.row(2)); assert!(rows.row(4) > rows.row(2)); assert!(rows.row(5) > rows.row(2)); assert!(rows.row(3) > rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
4027 assert_eq!(back.len(), 1);
4028 back[0].to_data().validate_full().unwrap();
4029 assert_eq!(&back[0], &list);
4030
4031 let options = SortOptions::default().desc().with_nulls_first(false);
4033 let field = SortField::new_with_options(d.clone(), options);
4034 let converter = RowConverter::new(vec![field]).unwrap();
4035 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4036 assert!(rows.row(0) < rows.row(1)); assert!(rows.row(2) > rows.row(1)); assert!(rows.row(3) > rows.row(2)); assert!(rows.row(4) > rows.row(2)); assert!(rows.row(5) > rows.row(2)); assert!(rows.row(3) > rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
4045 assert_eq!(back.len(), 1);
4046 back[0].to_data().validate_full().unwrap();
4047 assert_eq!(&back[0], &list);
4048
4049 let options = SortOptions::default().desc().with_nulls_first(true);
4051 let field = SortField::new_with_options(d, options);
4052 let converter = RowConverter::new(vec![field]).unwrap();
4053 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4054
4055 assert!(rows.row(0) < rows.row(1)); assert!(rows.row(2) < rows.row(1)); assert!(rows.row(3) < rows.row(2)); assert!(rows.row(4) < rows.row(2)); assert!(rows.row(5) < rows.row(2)); assert!(rows.row(3) < rows.row(5)); assert_eq!(rows.row(3), rows.row(6)); let back = converter.convert_rows(&rows).unwrap();
4064 assert_eq!(back.len(), 1);
4065 back[0].to_data().validate_full().unwrap();
4066 assert_eq!(&back[0], &list);
4067
4068 let sliced_list = list.slice(1, 5);
4069 let rows_on_sliced_list = converter
4070 .convert_columns(&[Arc::clone(&sliced_list)])
4071 .unwrap();
4072
4073 assert!(rows_on_sliced_list.row(2) < rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(3) < rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(4) < rows_on_sliced_list.row(1)); assert!(rows_on_sliced_list.row(2) < rows_on_sliced_list.row(4)); let back = converter.convert_rows(&rows_on_sliced_list).unwrap();
4079 assert_eq!(back.len(), 1);
4080 back[0].to_data().validate_full().unwrap();
4081 assert_eq!(&back[0], &sliced_list);
4082 }
4083
4084 #[test]
4085 fn test_two_fixed_size_lists() {
4086 let mut first = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4087 first.values().append_value(100);
4089 first.append(true);
4090 first.values().append_value(101);
4092 first.append(true);
4093 first.values().append_value(102);
4095 first.append(true);
4096 first.values().append_null();
4098 first.append(true);
4099 first.values().append_null(); first.append(false);
4102 let first = Arc::new(first.finish()) as ArrayRef;
4103 let first_type = first.data_type().clone();
4104
4105 let mut second = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4106 second.values().append_value(200);
4108 second.append(true);
4109 second.values().append_value(201);
4111 second.append(true);
4112 second.values().append_value(202);
4114 second.append(true);
4115 second.values().append_null();
4117 second.append(true);
4118 second.values().append_null(); second.append(false);
4121 let second = Arc::new(second.finish()) as ArrayRef;
4122 let second_type = second.data_type().clone();
4123
4124 let converter = RowConverter::new(vec![
4125 SortField::new(first_type.clone()),
4126 SortField::new(second_type.clone()),
4127 ])
4128 .unwrap();
4129
4130 let rows = converter
4131 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4132 .unwrap();
4133
4134 let back = converter.convert_rows(&rows).unwrap();
4135 assert_eq!(back.len(), 2);
4136 back[0].to_data().validate_full().unwrap();
4137 assert_eq!(&back[0], &first);
4138 back[1].to_data().validate_full().unwrap();
4139 assert_eq!(&back[1], &second);
4140 }
4141
4142 #[test]
4143 fn test_fixed_size_list_with_variable_width_content() {
4144 let mut first = FixedSizeListBuilder::new(
4145 StructBuilder::from_fields(
4146 vec![
4147 Field::new(
4148 "timestamp",
4149 DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))),
4150 false,
4151 ),
4152 Field::new("offset_minutes", DataType::Int16, false),
4153 Field::new("time_zone", DataType::Utf8, false),
4154 ],
4155 1,
4156 ),
4157 1,
4158 );
4159 first
4161 .values()
4162 .field_builder::<TimestampMicrosecondBuilder>(0)
4163 .unwrap()
4164 .append_null();
4165 first
4166 .values()
4167 .field_builder::<Int16Builder>(1)
4168 .unwrap()
4169 .append_null();
4170 first
4171 .values()
4172 .field_builder::<StringBuilder>(2)
4173 .unwrap()
4174 .append_null();
4175 first.values().append(false);
4176 first.append(false);
4177 first
4179 .values()
4180 .field_builder::<TimestampMicrosecondBuilder>(0)
4181 .unwrap()
4182 .append_null();
4183 first
4184 .values()
4185 .field_builder::<Int16Builder>(1)
4186 .unwrap()
4187 .append_null();
4188 first
4189 .values()
4190 .field_builder::<StringBuilder>(2)
4191 .unwrap()
4192 .append_null();
4193 first.values().append(false);
4194 first.append(true);
4195 first
4197 .values()
4198 .field_builder::<TimestampMicrosecondBuilder>(0)
4199 .unwrap()
4200 .append_value(0);
4201 first
4202 .values()
4203 .field_builder::<Int16Builder>(1)
4204 .unwrap()
4205 .append_value(0);
4206 first
4207 .values()
4208 .field_builder::<StringBuilder>(2)
4209 .unwrap()
4210 .append_value("UTC");
4211 first.values().append(true);
4212 first.append(true);
4213 first
4215 .values()
4216 .field_builder::<TimestampMicrosecondBuilder>(0)
4217 .unwrap()
4218 .append_value(1126351800123456);
4219 first
4220 .values()
4221 .field_builder::<Int16Builder>(1)
4222 .unwrap()
4223 .append_value(120);
4224 first
4225 .values()
4226 .field_builder::<StringBuilder>(2)
4227 .unwrap()
4228 .append_value("Europe/Warsaw");
4229 first.values().append(true);
4230 first.append(true);
4231 let first = Arc::new(first.finish()) as ArrayRef;
4232 let first_type = first.data_type().clone();
4233
4234 let mut second = StringBuilder::new();
4235 second.append_value("somewhere near");
4236 second.append_null();
4237 second.append_value("Greenwich");
4238 second.append_value("Warsaw");
4239 let second = Arc::new(second.finish()) as ArrayRef;
4240 let second_type = second.data_type().clone();
4241
4242 let converter = RowConverter::new(vec![
4243 SortField::new(first_type.clone()),
4244 SortField::new(second_type.clone()),
4245 ])
4246 .unwrap();
4247
4248 let rows = converter
4249 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4250 .unwrap();
4251
4252 let back = converter.convert_rows(&rows).unwrap();
4253 assert_eq!(back.len(), 2);
4254 back[0].to_data().validate_full().unwrap();
4255 assert_eq!(&back[0], &first);
4256 back[1].to_data().validate_full().unwrap();
4257 assert_eq!(&back[1], &second);
4258 }
4259
4260 #[test]
4261 fn test_single_map() {
4262 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4263 builder.keys().append_value("hello");
4265 builder.values().append_value(1);
4266 builder.keys().append_value("world");
4267 builder.values().append_value(2);
4268 builder.append(true).unwrap();
4269
4270 builder.keys().append_value("foo");
4272 builder.values().append_value(3);
4273 builder.append(true).unwrap();
4274
4275 builder.append(true).unwrap();
4277
4278 builder.keys().append_value("masked_key");
4280 builder.values().append_value(999);
4281 builder.append(false).unwrap();
4282
4283 builder.append(false).unwrap();
4285
4286 builder.keys().append_value("bar");
4288 builder.values().append_null();
4289 builder.append(true).unwrap();
4290
4291 builder.keys().append_value("other_masked");
4293 builder.values().append_value(0);
4294 builder.append(false).unwrap();
4295
4296 builder.keys().append_value("a");
4298 builder.values().append_value(10);
4299 builder.keys().append_value("b");
4300 builder.values().append_value(20);
4301 builder.keys().append_value("c");
4302 builder.values().append_value(30);
4303 builder.append(true).unwrap();
4304
4305 let map = Arc::new(builder.finish()) as ArrayRef;
4306 let d = map.data_type().clone();
4307
4308 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4309
4310 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4311
4312 assert_eq!(rows.row(3), rows.row(4));
4314 assert_eq!(rows.row(4), rows.row(6));
4315
4316 let back = converter.convert_rows(&rows).unwrap();
4317 assert_eq!(back.len(), 1);
4318 back[0].to_data().validate_full().unwrap();
4319 assert_eq!(&back[0], &map);
4320
4321 let sliced_map = map.slice(1, map.len() - 2);
4322 let rows_on_sliced = converter
4323 .convert_columns(&[Arc::clone(&sliced_map)])
4324 .unwrap();
4325
4326 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4327 assert_eq!(back.len(), 1);
4328 back[0].to_data().validate_full().unwrap();
4329 assert_eq!(&back[0], &sliced_map);
4330 }
4331
4332 #[test]
4333 fn two_maps_with_different_keys_order_should_sort_by_entry_order() {
4334 let map_1: ArrayRef =
4335 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4336 vec![Some(vec![("hello", Some(1)), ("world", Some(2))])],
4337 false,
4338 ));
4339 let map_2: ArrayRef =
4341 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4342 vec![Some(vec![("world", Some(2)), ("hello", Some(1))])],
4343 false,
4344 ));
4345
4346 let converter = RowConverter::new(vec![SortField::new(map_1.data_type().clone())]).unwrap();
4347
4348 let map_1_rows = converter.convert_columns(&[Arc::clone(&map_1)]).unwrap();
4349 let map_2_rows = converter.convert_columns(&[Arc::clone(&map_2)]).unwrap();
4350
4351 assert_ne!(map_1_rows.row(0), map_2_rows.row(0));
4352 assert!(map_1_rows.row(0) < map_2_rows.row(0));
4353
4354 let back_1 = converter.convert_rows(&map_1_rows).unwrap();
4355 let back_2 = converter.convert_rows(&map_2_rows).unwrap();
4356 assert_eq!(&back_1[0], &map_1);
4357 assert_eq!(&back_2[0], &map_2);
4358 }
4359
4360 #[test]
4361 fn test_nested_map() {
4362 let mut builder = MapBuilder::new(
4364 None,
4365 StringBuilder::new(),
4366 MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()),
4367 );
4368
4369 builder.keys().append_value("outer1");
4371 builder.values().keys().append_value("inner_a");
4372 builder.values().values().append_value(1);
4373 builder.values().keys().append_value("inner_b");
4374 builder.values().values().append_value(2);
4375 builder.values().append(true).unwrap();
4376 builder.keys().append_value("outer2");
4377 builder.values().keys().append_value("inner_c");
4378 builder.values().values().append_value(3);
4379 builder.values().append(true).unwrap();
4380 builder.append(true).unwrap();
4381
4382 builder.keys().append_value("x");
4384 builder.values().append(true).unwrap();
4385 builder.append(true).unwrap();
4386
4387 builder.keys().append_value("y");
4389 builder.values().keys().append_value("masked"); builder.values().values().append_value(0); builder.values().append(false).unwrap();
4392 builder.append(true).unwrap();
4393
4394 builder.keys().append_value("y");
4396 builder.values().append(false).unwrap(); builder.append(true).unwrap();
4398
4399 builder.keys().append_value("masked_outer"); builder.values().keys().append_value("masked_inner"); builder.values().values().append_value(0); builder.values().append(true).unwrap(); builder.append(false).unwrap();
4405
4406 builder.keys().append_value("masked_outer"); builder.values().append(false).unwrap(); builder.append(false).unwrap();
4410
4411 builder.append(false).unwrap(); builder.append(true).unwrap();
4416
4417 let map = Arc::new(builder.finish()) as ArrayRef;
4418 let d = map.data_type().clone();
4419
4420 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4421
4422 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4423
4424 let back = converter.convert_rows(&rows).unwrap();
4425 assert_eq!(back.len(), 1);
4426 back[0].to_data().validate_full().unwrap();
4427 assert_eq!(&back[0], &map);
4428
4429 let sliced_map = map.slice(1, 3);
4430 let rows_on_sliced = converter
4431 .convert_columns(&[Arc::clone(&sliced_map)])
4432 .unwrap();
4433
4434 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4435 assert_eq!(back.len(), 1);
4436 back[0].to_data().validate_full().unwrap();
4437 assert_eq!(&back[0], &sliced_map);
4438 }
4439
4440 #[test]
4441 fn test_single_map_with_non_nullable_values() {
4442 let value_field = Arc::new(Field::new(
4444 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4445 DataType::Int32,
4446 false,
4447 ));
4448 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4449 .with_values_field(value_field);
4450 builder.keys().append_value("a");
4452 builder.values().append_value(1);
4453 builder.keys().append_value("b");
4454 builder.values().append_value(2);
4455 builder.append(true).unwrap();
4456 builder.append(false).unwrap();
4458 builder.keys().append_value("c");
4460 builder.values().append_value(3);
4461 builder.append(true).unwrap();
4462 builder.append(true).unwrap();
4464 builder.keys().append_value("masked"); builder.values().append_value(0); builder.append(false).unwrap();
4468
4469 let map = Arc::new(builder.finish()) as ArrayRef;
4470 let d = map.data_type().clone();
4471
4472 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4473
4474 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4475
4476 let back = converter.convert_rows(&rows).unwrap();
4477 assert_eq!(back.len(), 1);
4478 back[0].to_data().validate_full().unwrap();
4479 assert_eq!(&back[0], &map);
4480 }
4481
4482 #[test]
4483 fn test_single_map_with_non_nullable_map_but_with_nullable_values() {
4484 let value_field = Arc::new(Field::new(
4486 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4487 DataType::Int32,
4488 true,
4489 ));
4490 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4491 .with_values_field(value_field);
4492
4493 builder.keys().append_value("a");
4495 builder.values().append_value(1);
4496 builder.keys().append_value("b");
4497 builder.values().append_null();
4498 builder.append(true).unwrap();
4499 builder.keys().append_value("c");
4501 builder.values().append_null();
4502 builder.keys().append_value("d");
4503 builder.values().append_null();
4504 builder.append(true).unwrap();
4505 builder.append(true).unwrap();
4507 builder.keys().append_value("e");
4509 builder.values().append_value(5);
4510 builder.append(true).unwrap();
4511
4512 let map = Arc::new(builder.finish()) as ArrayRef;
4513 let d = map.data_type().clone();
4514
4515 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4516
4517 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4518
4519 let back = converter.convert_rows(&rows).unwrap();
4520 assert_eq!(back.len(), 1);
4521 back[0].to_data().validate_full().unwrap();
4522 assert_eq!(&back[0], &map);
4523 }
4524
4525 #[test]
4526 fn test_map_all_nulls() {
4527 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4528 builder.keys().append_value("m1"); builder.values().append_value(1); builder.append(false).unwrap();
4532 builder.keys().append_value("m2"); builder.values().append_value(2); builder.append(false).unwrap();
4535
4536 builder.append(false).unwrap(); builder.keys().append_value("m3"); builder.values().append_value(3); builder.append(false).unwrap();
4541
4542 let map = Arc::new(builder.finish()) as ArrayRef;
4543 let d = map.data_type().clone();
4544
4545 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4546
4547 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4548
4549 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4551
4552 let back = converter.convert_rows(&rows).unwrap();
4553 assert_eq!(back.len(), 1);
4554 back[0].to_data().validate_full().unwrap();
4555 assert_eq!(&back[0], &map);
4556 }
4557
4558 #[test]
4559 fn test_map_all_empty() {
4560 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4561 builder.append(true).unwrap();
4563 builder.append(true).unwrap();
4564 builder.append(true).unwrap();
4565
4566 let map = Arc::new(builder.finish()) as ArrayRef;
4567 let d = map.data_type().clone();
4568
4569 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4570
4571 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4572
4573 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4575
4576 let back = converter.convert_rows(&rows).unwrap();
4577 assert_eq!(back.len(), 1);
4578 back[0].to_data().validate_full().unwrap();
4579 assert_eq!(&back[0], &map);
4580 }
4581
4582 #[test]
4583 fn test_map_empty_array() {
4584 let builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4586 let map = Arc::new(builder.finish_cloned()) as ArrayRef;
4587 let d = map.data_type().clone();
4588
4589 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4590
4591 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4592
4593 let back = converter.convert_rows(&rows).unwrap();
4594 assert_eq!(back.len(), 1);
4595 back[0].to_data().validate_full().unwrap();
4596 assert_eq!(&back[0], &map);
4597 }
4598
4599 fn generate_primitive_array<K>(
4600 rng: &mut StdRng,
4601 len: usize,
4602 valid_percent: f64,
4603 ) -> PrimitiveArray<K>
4604 where
4605 K: ArrowPrimitiveType,
4606 StandardUniform: Distribution<K::Native>,
4607 {
4608 (0..len)
4609 .map(|_| rng.random_bool(valid_percent).then(|| rng.random()))
4610 .collect()
4611 }
4612
4613 fn generate_all_unique_primitive_array<K>(rng: &mut StdRng, len: usize) -> PrimitiveArray<K>
4614 where
4615 K: ArrowPrimitiveType,
4616 K::Native: Hash + Eq,
4617 StandardUniform: Distribution<K::Native>,
4618 {
4619 let possible_number_of_values = 2i32.saturating_pow(size_of::<K::Native>() as u32 * 8);
4620 assert!(
4621 len <= possible_number_of_values as usize,
4622 "len {len} is larger than the number of possible values {possible_number_of_values}"
4623 );
4624
4625 let mut seen = std::collections::HashSet::new();
4626 (0..len)
4627 .map(|_| {
4628 let mut value;
4629 loop {
4630 value = rng.random();
4631
4632 if seen.insert(value) {
4633 break;
4634 }
4635 }
4636
4637 Some(value)
4638 })
4639 .collect()
4640 }
4641
4642 fn generate_boolean_array(rng: &mut StdRng, len: usize, valid_percent: f64) -> BooleanArray {
4643 (0..len)
4644 .map(|_| rng.random_bool(valid_percent).then(|| rng.random_bool(0.5)))
4645 .collect()
4646 }
4647
4648 fn generate_strings<O: OffsetSizeTrait>(
4649 rng: &mut StdRng,
4650 len: usize,
4651 valid_percent: f64,
4652 ) -> GenericStringArray<O> {
4653 (0..len)
4654 .map(|_| {
4655 rng.random_bool(valid_percent).then(|| {
4656 let len = rng.random_range(0..100);
4657 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4658 String::from_utf8(bytes).unwrap()
4659 })
4660 })
4661 .collect()
4662 }
4663
4664 fn generate_string_view(rng: &mut StdRng, len: usize, valid_percent: f64) -> StringViewArray {
4665 (0..len)
4666 .map(|_| {
4667 rng.random_bool(valid_percent).then(|| {
4668 let len = rng.random_range(0..100);
4669 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4670 String::from_utf8(bytes).unwrap()
4671 })
4672 })
4673 .collect()
4674 }
4675
4676 fn generate_byte_view(rng: &mut StdRng, len: usize, valid_percent: f64) -> BinaryViewArray {
4677 (0..len)
4678 .map(|_| {
4679 rng.random_bool(valid_percent).then(|| {
4680 let len = rng.random_range(0..100);
4681 let bytes: Vec<_> = (0..len).map(|_| rng.random_range(0..128)).collect();
4682 bytes
4683 })
4684 })
4685 .collect()
4686 }
4687
4688 fn generate_fixed_stringview_column(len: usize) -> StringViewArray {
4689 let edge_cases = vec![
4690 Some("bar".to_string()),
4691 Some("bar\0".to_string()),
4692 Some("LongerThan12Bytes".to_string()),
4693 Some("LongerThan12Bytez".to_string()),
4694 Some("LongerThan12Bytes\0".to_string()),
4695 Some("LongerThan12Byt".to_string()),
4696 Some("backend one".to_string()),
4697 Some("backend two".to_string()),
4698 Some("a".repeat(257)),
4699 Some("a".repeat(300)),
4700 ];
4701
4702 let mut values = Vec::with_capacity(len);
4704 for i in 0..len {
4705 values.push(
4706 edge_cases
4707 .get(i % edge_cases.len())
4708 .cloned()
4709 .unwrap_or(None),
4710 );
4711 }
4712
4713 StringViewArray::from(values)
4714 }
4715
4716 fn generate_dictionary<K>(
4717 rng: &mut StdRng,
4718 values: ArrayRef,
4719 len: usize,
4720 valid_percent: f64,
4721 ) -> DictionaryArray<K>
4722 where
4723 K: ArrowDictionaryKeyType,
4724 K::Native: SampleUniform,
4725 {
4726 let min_key = K::Native::from_usize(0).unwrap();
4727 let max_key = K::Native::from_usize(values.len()).unwrap();
4728 let keys: PrimitiveArray<K> = (0..len)
4729 .map(|_| {
4730 rng.random_bool(valid_percent)
4731 .then(|| rng.random_range(min_key..max_key))
4732 })
4733 .collect();
4734
4735 let data_type =
4736 DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
4737
4738 let data = keys
4739 .into_data()
4740 .into_builder()
4741 .data_type(data_type)
4742 .add_child_data(values.to_data())
4743 .build()
4744 .unwrap();
4745
4746 DictionaryArray::from(data)
4747 }
4748
4749 fn generate_fixed_size_binary(
4750 rng: &mut StdRng,
4751 len: usize,
4752 valid_percent: f64,
4753 ) -> FixedSizeBinaryArray {
4754 let width = rng.random_range(0..20);
4755 let mut builder = FixedSizeBinaryBuilder::new(width);
4756
4757 let mut b = vec![0; width as usize];
4758 for _ in 0..len {
4759 match rng.random_bool(valid_percent) {
4760 true => {
4761 b.iter_mut().for_each(|x| *x = rng.random());
4762 builder.append_value(&b).unwrap();
4763 }
4764 false => builder.append_null(),
4765 }
4766 }
4767
4768 builder.finish()
4769 }
4770
4771 fn generate_struct(rng: &mut StdRng, len: usize, valid_percent: f64) -> StructArray {
4772 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4773 let a = generate_primitive_array::<Int32Type>(rng, len, valid_percent);
4774 let b = generate_strings::<i32>(rng, len, valid_percent);
4775 let fields = Fields::from(vec![
4776 Field::new("a", DataType::Int32, true),
4777 Field::new("b", DataType::Utf8, true),
4778 ]);
4779 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
4780 StructArray::new(fields, values, Some(nulls))
4781 }
4782
4783 fn generate_list<F>(rng: &mut StdRng, len: usize, valid_percent: f64, values: F) -> ListArray
4784 where
4785 F: FnOnce(&mut StdRng, usize) -> ArrayRef,
4786 {
4787 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4788 let values_len = offsets.last().as_usize();
4789 let values = values(rng, values_len);
4790 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4791 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4792 ListArray::new(field, offsets, values, Some(nulls))
4793 }
4794
4795 fn generate_list_view<F>(
4796 rng: &mut StdRng,
4797 len: usize,
4798 valid_percent: f64,
4799 values: F,
4800 ) -> ListViewArray
4801 where
4802 F: FnOnce(&mut StdRng, usize) -> ArrayRef,
4803 {
4804 let sizes: Vec<i32> = (0..len).map(|_| rng.random_range(0..10)).collect();
4806 let values_len: usize = sizes.iter().map(|s| *s as usize).sum::<usize>().max(1);
4807 let values = values(rng, values_len);
4808
4809 let offsets: Vec<i32> = sizes
4811 .iter()
4812 .map(|&size| {
4813 if size == 0 {
4814 0
4815 } else {
4816 rng.random_range(0..=(values_len as i32 - size))
4817 }
4818 })
4819 .collect();
4820
4821 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4822 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4823 ListViewArray::new(
4824 field,
4825 ScalarBuffer::from(offsets),
4826 ScalarBuffer::from(sizes),
4827 values,
4828 Some(nulls),
4829 )
4830 }
4831
4832 fn generate_map<KeysFn, ValuesFn>(
4833 rng: &mut StdRng,
4834 len: usize,
4835 valid_percent: f64,
4836 gen_keys: KeysFn,
4837 gen_values: ValuesFn,
4838 ) -> MapArray
4839 where
4840 KeysFn: FnOnce(&mut StdRng, usize) -> ArrayRef,
4841 ValuesFn: FnOnce(&mut StdRng, usize) -> ArrayRef,
4842 {
4843 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4844 let entries_len = offsets.last().as_usize();
4845 let keys = gen_keys(rng, entries_len);
4846 let values = gen_values(rng, entries_len);
4847 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4848 let field = Arc::new(Field::new_map(
4849 "",
4850 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4851 Field::new(
4852 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4853 keys.data_type().clone(),
4854 false,
4855 ),
4856 Field::new(
4857 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4858 values.data_type().clone(),
4859 true,
4860 ),
4861 false,
4862 true,
4863 ));
4864 let DataType::Map(struct_field, _) = field.data_type() else {
4865 unreachable!();
4866 };
4867
4868 let DataType::Struct(fields) = struct_field.data_type() else {
4869 unreachable!();
4870 };
4871
4872 let entries = StructArray::new(fields.clone(), vec![keys, values], None);
4873
4874 let map_array = MapArray::new(struct_field.clone(), offsets, entries, Some(nulls), false);
4875
4876 assert_valid_map(&map_array);
4877
4878 map_array
4879 }
4880
4881 fn assert_valid_map(array: &MapArray) {
4891 let keys_arrow_row_converter =
4892 RowConverter::new(vec![SortField::new(array.key_type().clone())]).unwrap();
4893
4894 array.iter().enumerate().filter_map(|(index, entry)| entry.map(|entry| (index, Arc::clone(entry.column(0))))).for_each(|(entry_index, keys)| {
4895 let keys_as_rows = keys_arrow_row_converter.convert_columns(&[Arc::clone(&keys)]).expect("should be able to convert keys");
4896
4897 for i in 0..keys_as_rows.num_rows() {
4898 for j in (i + 1)..keys_as_rows.num_rows() {
4899 if keys_as_rows.row(i) == keys_as_rows.row(j) {
4900 let key_i = keys.slice(i, 1);
4901 let key_j = keys.slice(j, 1);
4902
4903 assert_ne!(keys_as_rows.row(i), keys_as_rows.row(j), "map keys should be unique, but key {i} and key {j} are equal in entry {entry_index}. key {i} value is {key_i:?} and key {j} value is {key_j:?}");
4904 }
4905 }
4906 }
4907 })
4908 }
4909
4910 fn generate_nulls(rng: &mut StdRng, len: usize) -> Option<NullBuffer> {
4911 Some(NullBuffer::from_iter(
4912 (0..len).map(|_| rng.random_bool(0.8)),
4913 ))
4914 }
4915
4916 fn change_underlying_null_values_for_primitive<T: ArrowPrimitiveType>(
4917 array: &PrimitiveArray<T>,
4918 ) -> PrimitiveArray<T> {
4919 let (dt, values, nulls) = array.clone().into_parts();
4920
4921 let new_values = ScalarBuffer::<T::Native>::from_iter(
4922 values
4923 .iter()
4924 .zip(nulls.as_ref().unwrap().iter())
4925 .map(|(val, is_valid)| {
4926 if is_valid {
4927 *val
4928 } else {
4929 val.add_wrapping(T::Native::usize_as(1))
4930 }
4931 }),
4932 );
4933
4934 PrimitiveArray::new(new_values, nulls).with_data_type(dt)
4935 }
4936
4937 fn change_underline_null_values_for_byte_array<T: ByteArrayType>(
4938 array: &GenericByteArray<T>,
4939 ) -> GenericByteArray<T> {
4940 let (offsets, values, nulls) = array.clone().into_parts();
4941
4942 let new_offsets = OffsetBuffer::<T::Offset>::from_lengths(
4943 offsets
4944 .lengths()
4945 .zip(nulls.as_ref().unwrap().iter())
4946 .map(|(len, is_valid)| if is_valid { len } else { len + 1 }),
4947 );
4948
4949 let mut new_bytes = Vec::<u8>::with_capacity(new_offsets[new_offsets.len() - 1].as_usize());
4950
4951 offsets
4952 .windows(2)
4953 .zip(nulls.as_ref().unwrap().iter())
4954 .for_each(|(start_and_end, is_valid)| {
4955 let start = start_and_end[0].as_usize();
4956 let end = start_and_end[1].as_usize();
4957 new_bytes.extend_from_slice(&values.as_slice()[start..end]);
4958
4959 if !is_valid {
4961 new_bytes.push(b'c');
4962 }
4963 });
4964
4965 GenericByteArray::<T>::new(new_offsets, Buffer::from_vec(new_bytes), nulls)
4966 }
4967
4968 fn change_underline_null_values_for_list_array<O: OffsetSizeTrait>(
4969 array: &GenericListArray<O>,
4970 ) -> GenericListArray<O> {
4971 let (field, offsets, values, nulls) = array.clone().into_parts();
4972
4973 let (new_values, new_offsets) = {
4974 let concat_values = offsets
4975 .windows(2)
4976 .zip(nulls.as_ref().unwrap().iter())
4977 .map(|(start_and_end, is_valid)| {
4978 let start = start_and_end[0].as_usize();
4979 let end = start_and_end[1].as_usize();
4980 if is_valid {
4981 return (start, end - start);
4982 }
4983
4984 if end == values.len() {
4986 (start, (end - start).saturating_sub(1))
4987 } else {
4988 (start, end - start + 1)
4989 }
4990 })
4991 .map(|(start, length)| values.slice(start, length))
4992 .collect::<Vec<_>>();
4993
4994 let new_offsets =
4995 OffsetBuffer::<O>::from_lengths(concat_values.iter().map(|s| s.len()));
4996
4997 let new_values = {
4998 let values = concat_values.iter().map(|a| a.as_ref()).collect::<Vec<_>>();
4999 arrow_select::concat::concat(&values).expect("should be able to concat")
5000 };
5001
5002 (new_values, new_offsets)
5003 };
5004
5005 GenericListArray::<O>::new(field, new_offsets, new_values, nulls)
5006 }
5007
5008 fn change_underline_null_values_for_map_array(array: &MapArray) -> MapArray {
5009 let (field, offsets, entries, nulls, ordered) = array.clone().into_parts();
5010 assert!(
5011 !ordered,
5012 "can't replace underlying null values for ordered map array as this can violate the ordering"
5013 );
5014
5015 let (new_entries, new_offsets) = {
5016 let concat_values = offsets
5017 .windows(2)
5018 .zip(nulls.as_ref().unwrap().iter())
5019 .map(|(start_and_end, is_valid)| {
5020 let start = start_and_end[0].as_usize();
5021 let end = start_and_end[1].as_usize();
5022 if is_valid {
5023 return (start, end - start);
5024 }
5025
5026 if end == entries.len() {
5028 (start, (end - start).saturating_sub(1))
5029 } else {
5030 (start, end - start + 1)
5032 }
5033 })
5034 .map(|(start, length)| entries.slice(start, length))
5035 .collect::<Vec<_>>();
5036
5037 let new_offsets = OffsetBuffer::from_lengths(concat_values.iter().map(|s| s.len()));
5038
5039 let new_values = {
5040 let values = concat_values
5041 .iter()
5042 .map(|a| a as &dyn Array)
5043 .collect::<Vec<_>>();
5044 arrow_select::concat::concat(&values).expect("should be able to concat")
5045 };
5046
5047 (new_values.as_struct().clone(), new_offsets)
5048 };
5049
5050 let new_map = MapArray::new(field, new_offsets, new_entries, nulls, ordered);
5051
5052 assert_valid_map(&new_map);
5053
5054 new_map
5055 }
5056
5057 fn change_underline_null_values(array: &ArrayRef) -> ArrayRef {
5058 if array.null_count() == 0 {
5059 return Arc::clone(array);
5060 }
5061
5062 downcast_primitive_array!(
5063 array => {
5064 let output = change_underlying_null_values_for_primitive(array);
5065
5066 Arc::new(output)
5067 }
5068
5069 DataType::Utf8 => {
5070 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i32>()))
5071 }
5072 DataType::LargeUtf8 => {
5073 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i64>()))
5074 }
5075 DataType::Binary => {
5076 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i32>()))
5077 }
5078 DataType::LargeBinary => {
5079 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i64>()))
5080 }
5081 DataType::List(_) => {
5082 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i32>()))
5083 }
5084 DataType::LargeList(_) => {
5085 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i64>()))
5086 }
5087 DataType::Map(_, _) => {
5088 Arc::new(change_underline_null_values_for_map_array(array.as_map()))
5089 }
5090 _ => {
5091 Arc::clone(array)
5092 }
5093 )
5094 }
5095
5096 fn generate_column(rng: &mut StdRng, len: usize) -> ArrayRef {
5097 match rng.random_range(0..24) {
5098 0 => Arc::new(generate_primitive_array::<Int32Type>(rng, len, 0.8)),
5099 1 => Arc::new(generate_primitive_array::<UInt32Type>(rng, len, 0.8)),
5100 2 => Arc::new(generate_primitive_array::<Int64Type>(rng, len, 0.8)),
5101 3 => Arc::new(generate_primitive_array::<UInt64Type>(rng, len, 0.8)),
5102 4 => Arc::new(generate_primitive_array::<Float32Type>(rng, len, 0.8)),
5103 5 => Arc::new(generate_primitive_array::<Float64Type>(rng, len, 0.8)),
5104 6 => Arc::new(generate_strings::<i32>(rng, len, 0.8)),
5105 7 => {
5106 let dict_values_len = rng.random_range(1..len);
5107 let strings = Arc::new(generate_strings::<i32>(rng, dict_values_len, 1.0));
5109 Arc::new(generate_dictionary::<Int64Type>(rng, strings, len, 0.8))
5110 }
5111 8 => {
5112 let dict_values_len = rng.random_range(1..len);
5113 let values = Arc::new(generate_primitive_array::<Int64Type>(
5115 rng,
5116 dict_values_len,
5117 1.0,
5118 ));
5119 Arc::new(generate_dictionary::<Int64Type>(rng, values, len, 0.8))
5120 }
5121 9 => Arc::new(generate_fixed_size_binary(rng, len, 0.8)),
5122 10 => Arc::new(generate_struct(rng, len, 0.8)),
5123 11 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5124 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5125 })),
5126 12 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5127 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5128 })),
5129 13 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5130 Arc::new(generate_struct(rng, values_len, 0.8))
5131 })),
5132 14 => Arc::new(generate_string_view(rng, len, 0.8)),
5133 15 => Arc::new(generate_byte_view(rng, len, 0.8)),
5134 16 => Arc::new(generate_fixed_stringview_column(len)),
5135 17 => Arc::new(
5136 generate_list(rng, len + 1000, 0.8, |rng, values_len| {
5137 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5138 })
5139 .slice(500, len),
5140 ),
5141 18 => Arc::new(generate_boolean_array(rng, len, 0.8)),
5142 19 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5143 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5144 })),
5145 20 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5146 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5147 })),
5148 21 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5149 Arc::new(generate_struct(rng, values_len, 0.8))
5150 })),
5151 22 => Arc::new(
5152 generate_list_view(rng, len + 1000, 0.8, |rng, values_len| {
5153 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5154 })
5155 .slice(500, len),
5156 ),
5157 23 => Arc::new(generate_map(
5158 rng,
5159 len,
5160 0.9,
5161 |rng, keys_len| {
5163 Arc::new(generate_all_unique_primitive_array::<Int64Type>(
5164 rng, keys_len,
5165 ))
5166 },
5167 |rng, values_len| Arc::new(generate_strings::<i32>(rng, values_len, 0.7)),
5168 )),
5169 _ => unreachable!(),
5170 }
5171 }
5172
5173 fn print_row(cols: &[SortColumn], row: usize) -> String {
5174 let t: Vec<_> = cols
5175 .iter()
5176 .map(|x| match x.values.is_valid(row) {
5177 true => {
5178 let opts = FormatOptions::default().with_null("NULL");
5179 let formatter = ArrayFormatter::try_new(x.values.as_ref(), &opts).unwrap();
5180 formatter.value(row).to_string()
5181 }
5182 false => "NULL".to_string(),
5183 })
5184 .collect();
5185 t.join(",")
5186 }
5187
5188 fn print_col_types(cols: &[SortColumn]) -> String {
5189 let t: Vec<_> = cols
5190 .iter()
5191 .map(|x| x.values.data_type().to_string())
5192 .collect();
5193 t.join(",")
5194 }
5195
5196 #[derive(Debug, PartialEq)]
5197 enum Nulls {
5198 AsIs,
5200
5201 Different,
5203
5204 None,
5206 }
5207
5208 #[test]
5209 #[cfg_attr(miri, ignore)] fn fuzz_test() {
5211 let mut rng = StdRng::seed_from_u64(42);
5212 for _ in 0..100 {
5213 for null_behavior in [Nulls::AsIs, Nulls::Different, Nulls::None] {
5214 let num_columns = rng.random_range(1..5);
5215 let len = rng.random_range(5..100);
5216 let mut arrays: Vec<_> = (0..num_columns)
5217 .map(|_| generate_column(&mut rng, len))
5218 .collect();
5219
5220 match null_behavior {
5221 Nulls::AsIs => {
5222 }
5224 Nulls::Different => {
5225 arrays = arrays
5227 .into_iter()
5228 .map(|a| replace_array_nulls(a, generate_nulls(&mut rng, len)))
5229 .collect()
5230 }
5231 Nulls::None => {
5232 arrays = arrays
5234 .into_iter()
5235 .map(|a| replace_array_nulls(a, None))
5236 .collect()
5237 }
5238 }
5239
5240 let options: Vec<_> = (0..num_columns)
5241 .map(|_| SortOptions {
5242 descending: rng.random_bool(0.5),
5243 nulls_first: rng.random_bool(0.5),
5244 })
5245 .collect();
5246
5247 let sort_columns: Vec<_> = options
5248 .iter()
5249 .zip(&arrays)
5250 .map(|(o, c)| SortColumn {
5251 values: Arc::clone(c),
5252 options: Some(*o),
5253 })
5254 .collect();
5255
5256 let comparator = LexicographicalComparator::try_new(&sort_columns).unwrap();
5257
5258 let columns: Vec<SortField> = options
5259 .into_iter()
5260 .zip(&arrays)
5261 .map(|(o, a)| SortField::new_with_options(a.data_type().clone(), o))
5262 .collect();
5263
5264 let converter = RowConverter::new(columns).unwrap();
5265 let rows = converter.convert_columns(&arrays).unwrap();
5266
5267 if !matches!(null_behavior, Nulls::None) {
5270 assert_same_rows_when_changing_input_underlying_null_values(
5271 &arrays, &converter, &rows,
5272 );
5273 }
5274
5275 for i in 0..len {
5276 for j in 0..len {
5277 let row_i = rows.row(i);
5278 let row_j = rows.row(j);
5279 let row_cmp = row_i.cmp(&row_j);
5280 let lex_cmp = comparator.compare(i, j);
5281 assert_eq!(
5282 row_cmp,
5283 lex_cmp,
5284 "({:?} vs {:?}) vs ({:?} vs {:?}) for types {}",
5285 print_row(&sort_columns, i),
5286 print_row(&sort_columns, j),
5287 row_i,
5288 row_j,
5289 print_col_types(&sort_columns)
5290 );
5291 }
5292 }
5293
5294 {
5296 let mut rows_iter = rows.iter();
5297 let mut rows_lengths_iter = rows.lengths();
5298 for (index, row) in rows_iter.by_ref().enumerate() {
5299 let len = rows_lengths_iter
5300 .next()
5301 .expect("Reached end of length iterator while still have rows");
5302 assert_eq!(
5303 row.data.len(),
5304 len,
5305 "Row length mismatch: {} vs {}",
5306 row.data.len(),
5307 len
5308 );
5309 assert_eq!(
5310 len,
5311 rows.row_len(index),
5312 "Row length mismatch at index {}: {} vs {}",
5313 index,
5314 len,
5315 rows.row_len(index)
5316 );
5317 }
5318
5319 assert_eq!(
5320 rows_lengths_iter.next(),
5321 None,
5322 "Length iterator did not reach end"
5323 );
5324 }
5325
5326 let back = converter.convert_rows(&rows).unwrap();
5329 for (actual, expected) in back.iter().zip(&arrays) {
5330 actual.to_data().validate_full().unwrap();
5331 dictionary_eq(actual, expected)
5332 }
5333
5334 let rows = rows.try_into_binary().expect("reasonable size");
5337 let parser = converter.parser();
5338 let back = converter
5339 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5340 .unwrap();
5341 for (actual, expected) in back.iter().zip(&arrays) {
5342 actual.to_data().validate_full().unwrap();
5343 dictionary_eq(actual, expected)
5344 }
5345
5346 let rows = converter.from_binary(rows);
5347 let back = converter.convert_rows(&rows).unwrap();
5348 for (actual, expected) in back.iter().zip(&arrays) {
5349 actual.to_data().validate_full().unwrap();
5350 dictionary_eq(actual, expected)
5351 }
5352 }
5353 }
5354 }
5355
5356 fn replace_array_nulls(array: ArrayRef, new_nulls: Option<NullBuffer>) -> ArrayRef {
5357 make_array(
5358 array
5359 .into_data()
5360 .into_builder()
5361 .nulls(new_nulls)
5363 .build()
5364 .unwrap(),
5365 )
5366 }
5367
5368 fn assert_same_rows_when_changing_input_underlying_null_values(
5369 arrays: &[ArrayRef],
5370 converter: &RowConverter,
5371 rows: &Rows,
5372 ) {
5373 let arrays_with_different_data_behind_nulls = arrays
5374 .iter()
5375 .map(|arr| change_underline_null_values(arr))
5376 .collect::<Vec<_>>();
5377
5378 if arrays
5380 .iter()
5381 .zip(arrays_with_different_data_behind_nulls.iter())
5382 .all(|(a, b)| Arc::ptr_eq(a, b))
5383 {
5384 return;
5385 }
5386
5387 let rows_with_different_nulls = converter
5388 .convert_columns(&arrays_with_different_data_behind_nulls)
5389 .unwrap();
5390
5391 assert_eq!(
5392 rows.iter().collect::<Vec<_>>(),
5393 rows_with_different_nulls.iter().collect::<Vec<_>>(),
5394 "Different underlying nulls should not output different rows"
5395 )
5396 }
5397
5398 #[test]
5399 fn test_clear() {
5400 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
5401 let mut rows = converter.empty_rows(3, 128);
5402
5403 let first = Int32Array::from(vec![None, Some(2), Some(4)]);
5404 let second = Int32Array::from(vec![Some(2), None, Some(4)]);
5405 let arrays = [Arc::new(first) as ArrayRef, Arc::new(second) as ArrayRef];
5406
5407 for array in &arrays {
5408 rows.clear();
5409 converter
5410 .append(&mut rows, std::slice::from_ref(array))
5411 .unwrap();
5412 let back = converter.convert_rows(&rows).unwrap();
5413 assert_eq!(&back[0], array);
5414 }
5415
5416 let mut rows_expected = converter.empty_rows(3, 128);
5417 converter.append(&mut rows_expected, &arrays[1..]).unwrap();
5418
5419 for (i, (actual, expected)) in rows.iter().zip(rows_expected.iter()).enumerate() {
5420 assert_eq!(
5421 actual, expected,
5422 "For row {i}: expected {expected:?}, actual: {actual:?}",
5423 );
5424 }
5425 }
5426
5427 #[test]
5428 fn test_append_codec_dictionary_binary() {
5429 use DataType::*;
5430 let converter = RowConverter::new(vec![SortField::new(Dictionary(
5432 Box::new(Int32),
5433 Box::new(Binary),
5434 ))])
5435 .unwrap();
5436 let mut rows = converter.empty_rows(4, 128);
5437
5438 let keys = Int32Array::from_iter_values([0, 1, 2, 3]);
5439 let values = BinaryArray::from(vec![
5440 Some(b"a".as_slice()),
5441 Some(b"b"),
5442 Some(b"c"),
5443 Some(b"d"),
5444 ]);
5445 let dict_array = DictionaryArray::new(keys, Arc::new(values));
5446
5447 rows.clear();
5448 let array = Arc::new(dict_array) as ArrayRef;
5449 converter
5450 .append(&mut rows, std::slice::from_ref(&array))
5451 .unwrap();
5452 let back = converter.convert_rows(&rows).unwrap();
5453
5454 dictionary_eq(&back[0], &array);
5455 }
5456
5457 #[test]
5458 fn test_list_prefix() {
5459 let mut a = ListBuilder::new(Int8Builder::new());
5460 a.append_value([None]);
5461 a.append_value([None, None]);
5462 let a = a.finish();
5463
5464 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
5465 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
5466 assert_eq!(rows.row(0).cmp(&rows.row(1)), Ordering::Less);
5467 }
5468
5469 #[test]
5470 fn test_utf8_validation_doesnt_affect_values_buffer_size() {
5471 fn assert_values_buffer_lens(col: ArrayRef) -> usize {
5472 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8View)]).unwrap();
5474
5475 let rows = converter.convert_columns(&[col]).unwrap();
5477 let converted = converter.convert_rows(&rows).unwrap();
5478 let unchecked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5479
5480 let rows = rows.try_into_binary().expect("reasonable size");
5482 let parser = converter.parser();
5483 let converted = converter
5484 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5485 .unwrap();
5486 let checked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5487 assert_eq!(unchecked_values_len, checked_values_len);
5489 checked_values_len
5490 }
5491
5492 let col = Arc::new(StringViewArray::from_iter([
5494 Some("hello"), None, Some("short"), Some("tiny"), ])) as ArrayRef;
5499
5500 let values_len = assert_values_buffer_lens(col);
5501 assert_eq!(values_len, 0);
5503
5504 let col = Arc::new(StringViewArray::from_iter([
5506 Some("1234567890123"), Some("12345678901234"), ])) as ArrayRef;
5509
5510 let values_len = assert_values_buffer_lens(col);
5511 assert_eq!(values_len, 13 + 14);
5512
5513 let col = Arc::new(StringViewArray::from_iter([
5515 Some("tiny"), Some("thisisexact13"), None,
5518 Some("short"), ])) as ArrayRef;
5520
5521 let values_len = assert_values_buffer_lens(col);
5522 assert_eq!(values_len, 13);
5524 }
5525
5526 #[test]
5527 fn test_sparse_union() {
5528 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5530 let str_array = StringArray::from(vec![None, Some("b"), None, Some("d"), None]);
5531
5532 let type_ids = vec![0, 1, 0, 1, 0].into();
5534
5535 let union_fields = [
5536 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5537 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5538 ]
5539 .into_iter()
5540 .collect();
5541
5542 let union_array = UnionArray::try_new(
5543 union_fields,
5544 type_ids,
5545 None,
5546 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5547 )
5548 .unwrap();
5549
5550 let union_type = union_array.data_type().clone();
5551 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5552
5553 let rows = converter
5554 .convert_columns(&[Arc::new(union_array.clone())])
5555 .unwrap();
5556
5557 let back = converter.convert_rows(&rows).unwrap();
5559 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5560
5561 assert_eq!(union_array.len(), back_union.len());
5562 for i in 0..union_array.len() {
5563 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5564 }
5565 }
5566
5567 #[test]
5568 fn test_sparse_union_with_nulls() {
5569 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5571 let str_array = StringArray::from(vec![None::<&str>; 5]);
5572
5573 let type_ids = vec![0, 1, 0, 1, 0].into();
5575
5576 let union_fields = [
5577 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5578 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5579 ]
5580 .into_iter()
5581 .collect();
5582
5583 let union_array = UnionArray::try_new(
5584 union_fields,
5585 type_ids,
5586 None,
5587 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5588 )
5589 .unwrap();
5590
5591 let union_type = union_array.data_type().clone();
5592 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5593
5594 let rows = converter
5595 .convert_columns(&[Arc::new(union_array.clone())])
5596 .unwrap();
5597
5598 let back = converter.convert_rows(&rows).unwrap();
5600 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5601
5602 assert_eq!(union_array.len(), back_union.len());
5603 for i in 0..union_array.len() {
5604 let expected_null = union_array.is_null(i);
5605 let actual_null = back_union.is_null(i);
5606 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5607 if !expected_null {
5608 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5609 }
5610 }
5611 }
5612
5613 #[test]
5614 fn test_dense_union() {
5615 let int_array = Int32Array::from(vec![1, 3, 5]);
5617 let str_array = StringArray::from(vec!["a", "b"]);
5618
5619 let type_ids = vec![0, 1, 0, 1, 0].into();
5620
5621 let offsets = vec![0, 0, 1, 1, 2].into();
5623
5624 let union_fields = [
5625 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5626 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5627 ]
5628 .into_iter()
5629 .collect();
5630
5631 let union_array = UnionArray::try_new(
5632 union_fields,
5633 type_ids,
5634 Some(offsets), vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5636 )
5637 .unwrap();
5638
5639 let union_type = union_array.data_type().clone();
5640 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5641
5642 let rows = converter
5643 .convert_columns(&[Arc::new(union_array.clone())])
5644 .unwrap();
5645
5646 let back = converter.convert_rows(&rows).unwrap();
5648 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5649
5650 assert_eq!(union_array.len(), back_union.len());
5651 for i in 0..union_array.len() {
5652 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5653 }
5654 }
5655
5656 #[test]
5657 fn test_dense_union_with_nulls() {
5658 let int_array = Int32Array::from(vec![Some(1), None, Some(5)]);
5660 let str_array = StringArray::from(vec![Some("a"), None]);
5661
5662 let type_ids = vec![0, 1, 0, 1, 0].into();
5664 let offsets = vec![0, 0, 1, 1, 2].into();
5665
5666 let union_fields = [
5667 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5668 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5669 ]
5670 .into_iter()
5671 .collect();
5672
5673 let union_array = UnionArray::try_new(
5674 union_fields,
5675 type_ids,
5676 Some(offsets),
5677 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5678 )
5679 .unwrap();
5680
5681 let union_type = union_array.data_type().clone();
5682 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5683
5684 let rows = converter
5685 .convert_columns(&[Arc::new(union_array.clone())])
5686 .unwrap();
5687
5688 let back = converter.convert_rows(&rows).unwrap();
5690 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5691
5692 assert_eq!(union_array.len(), back_union.len());
5693 for i in 0..union_array.len() {
5694 let expected_null = union_array.is_null(i);
5695 let actual_null = back_union.is_null(i);
5696 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5697 if !expected_null {
5698 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5699 }
5700 }
5701 }
5702
5703 #[test]
5704 fn test_union_ordering() {
5705 let int_array = Int32Array::from(vec![100, 5, 20]);
5706 let str_array = StringArray::from(vec!["z", "a"]);
5707
5708 let type_ids = vec![0, 1, 0, 1, 0].into();
5710 let offsets = vec![0, 0, 1, 1, 2].into();
5711
5712 let union_fields = [
5713 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5714 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5715 ]
5716 .into_iter()
5717 .collect();
5718
5719 let union_array = UnionArray::try_new(
5720 union_fields,
5721 type_ids,
5722 Some(offsets),
5723 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5724 )
5725 .unwrap();
5726
5727 let union_type = union_array.data_type().clone();
5728 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5729
5730 let rows = converter.convert_columns(&[Arc::new(union_array)]).unwrap();
5731
5732 assert!(rows.row(2) < rows.row(1));
5744
5745 assert!(rows.row(0) < rows.row(3));
5747
5748 assert!(rows.row(2) < rows.row(4));
5751 assert!(rows.row(4) < rows.row(0));
5753
5754 assert!(rows.row(3) < rows.row(1));
5757 }
5758
5759 #[test]
5760 fn test_row_converter_roundtrip_with_many_union_columns() {
5761 let fields1 = UnionFields::try_new(
5763 vec![0, 1],
5764 vec![
5765 Field::new("int", DataType::Int32, true),
5766 Field::new("string", DataType::Utf8, true),
5767 ],
5768 )
5769 .unwrap();
5770
5771 let int_array1 = Int32Array::from(vec![Some(67), None]);
5772 let string_array1 = StringArray::from(vec![None::<&str>, Some("hello")]);
5773 let type_ids1 = vec![0i8, 1].into();
5774
5775 let union_array1 = UnionArray::try_new(
5776 fields1.clone(),
5777 type_ids1,
5778 None,
5779 vec![
5780 Arc::new(int_array1) as ArrayRef,
5781 Arc::new(string_array1) as ArrayRef,
5782 ],
5783 )
5784 .unwrap();
5785
5786 let fields2 = UnionFields::try_new(
5788 vec![0, 1],
5789 vec![
5790 Field::new("int", DataType::Int32, true),
5791 Field::new("string", DataType::Utf8, true),
5792 ],
5793 )
5794 .unwrap();
5795
5796 let int_array2 = Int32Array::from(vec![Some(100), None]);
5797 let string_array2 = StringArray::from(vec![None::<&str>, Some("world")]);
5798 let type_ids2 = vec![0i8, 1].into();
5799
5800 let union_array2 = UnionArray::try_new(
5801 fields2.clone(),
5802 type_ids2,
5803 None,
5804 vec![
5805 Arc::new(int_array2) as ArrayRef,
5806 Arc::new(string_array2) as ArrayRef,
5807 ],
5808 )
5809 .unwrap();
5810
5811 let field1 = Field::new("col1", DataType::Union(fields1, UnionMode::Sparse), true);
5813 let field2 = Field::new("col2", DataType::Union(fields2, UnionMode::Sparse), true);
5814
5815 let sort_field1 = SortField::new(field1.data_type().clone());
5816 let sort_field2 = SortField::new(field2.data_type().clone());
5817
5818 let converter = RowConverter::new(vec![sort_field1, sort_field2]).unwrap();
5819
5820 let rows = converter
5821 .convert_columns(&[
5822 Arc::new(union_array1.clone()) as ArrayRef,
5823 Arc::new(union_array2.clone()) as ArrayRef,
5824 ])
5825 .unwrap();
5826
5827 let out = converter.convert_rows(&rows).unwrap();
5829
5830 let [col1, col2] = out.as_slice() else {
5831 panic!("expected 2 columns")
5832 };
5833
5834 let col1 = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5835 let col2 = col2.as_any().downcast_ref::<UnionArray>().unwrap();
5836
5837 for (expected, got) in [union_array1, union_array2].iter().zip([col1, col2]) {
5838 assert_eq!(expected.len(), got.len());
5839 assert_eq!(expected.type_ids(), got.type_ids());
5840
5841 for i in 0..expected.len() {
5842 assert_eq!(expected.value(i).as_ref(), got.value(i).as_ref());
5843 }
5844 }
5845 }
5846
5847 #[test]
5848 fn test_row_converter_roundtrip_with_one_union_column() {
5849 let fields = UnionFields::try_new(
5850 vec![0, 1],
5851 vec![
5852 Field::new("int", DataType::Int32, true),
5853 Field::new("string", DataType::Utf8, true),
5854 ],
5855 )
5856 .unwrap();
5857
5858 let int_array = Int32Array::from(vec![Some(67), None]);
5859 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5860 let type_ids = vec![0i8, 1].into();
5861
5862 let union_array = UnionArray::try_new(
5863 fields.clone(),
5864 type_ids,
5865 None,
5866 vec![
5867 Arc::new(int_array) as ArrayRef,
5868 Arc::new(string_array) as ArrayRef,
5869 ],
5870 )
5871 .unwrap();
5872
5873 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5874 let sort_field = SortField::new(field.data_type().clone());
5875 let converter = RowConverter::new(vec![sort_field]).unwrap();
5876
5877 let rows = converter
5878 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5879 .unwrap();
5880
5881 let out = converter.convert_rows(&rows).unwrap();
5883
5884 let [col1] = out.as_slice() else {
5885 panic!("expected 1 column")
5886 };
5887
5888 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5889 assert_eq!(col.len(), union_array.len());
5890 assert_eq!(col.type_ids(), union_array.type_ids());
5891
5892 for i in 0..col.len() {
5893 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5894 }
5895 }
5896
5897 #[test]
5898 fn test_row_converter_roundtrip_with_non_default_union_type_ids() {
5899 let fields = UnionFields::try_new(
5901 vec![70, 85],
5902 vec![
5903 Field::new("int", DataType::Int32, true),
5904 Field::new("string", DataType::Utf8, true),
5905 ],
5906 )
5907 .unwrap();
5908
5909 let int_array = Int32Array::from(vec![Some(67), None]);
5910 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5911 let type_ids = vec![70i8, 85].into();
5912
5913 let union_array = UnionArray::try_new(
5914 fields.clone(),
5915 type_ids,
5916 None,
5917 vec![
5918 Arc::new(int_array) as ArrayRef,
5919 Arc::new(string_array) as ArrayRef,
5920 ],
5921 )
5922 .unwrap();
5923
5924 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5925 let sort_field = SortField::new(field.data_type().clone());
5926 let converter = RowConverter::new(vec![sort_field]).unwrap();
5927
5928 let rows = converter
5929 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5930 .unwrap();
5931
5932 let out = converter.convert_rows(&rows).unwrap();
5934
5935 let [col1] = out.as_slice() else {
5936 panic!("expected 1 column")
5937 };
5938
5939 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5940 assert_eq!(col.len(), union_array.len());
5941 assert_eq!(col.type_ids(), union_array.type_ids());
5942
5943 for i in 0..col.len() {
5944 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5945 }
5946 }
5947
5948 #[test]
5949 fn rows_size_should_count_for_capacity() {
5950 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
5951
5952 let empty_rows_size_with_preallocate_rows_and_data = {
5953 let rows = row_converter.empty_rows(1000, 1000);
5954
5955 rows.size()
5956 };
5957 let empty_rows_size_with_preallocate_rows = {
5958 let rows = row_converter.empty_rows(1000, 0);
5959
5960 rows.size()
5961 };
5962 let empty_rows_size_with_preallocate_data = {
5963 let rows = row_converter.empty_rows(0, 1000);
5964
5965 rows.size()
5966 };
5967 let empty_rows_size_without_preallocate = {
5968 let rows = row_converter.empty_rows(0, 0);
5969
5970 rows.size()
5971 };
5972
5973 assert!(
5974 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_rows,
5975 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_rows}"
5976 );
5977 assert!(
5978 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_data,
5979 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_data}"
5980 );
5981 assert!(
5982 empty_rows_size_with_preallocate_rows > empty_rows_size_without_preallocate,
5983 "{empty_rows_size_with_preallocate_rows} should be larger than {empty_rows_size_without_preallocate}"
5984 );
5985 assert!(
5986 empty_rows_size_with_preallocate_data > empty_rows_size_without_preallocate,
5987 "{empty_rows_size_with_preallocate_data} should be larger than {empty_rows_size_without_preallocate}"
5988 );
5989 }
5990
5991 #[test]
5992 fn test_struct_no_child_fields() {
5993 fn run_test(array: ArrayRef) {
5994 let sort_fields = vec![SortField::new(array.data_type().clone())];
5995 let converter = RowConverter::new(sort_fields).unwrap();
5996 let r = converter.convert_columns(&[Arc::clone(&array)]).unwrap();
5997
5998 let back = converter.convert_rows(&r).unwrap();
5999 assert_eq!(back.len(), 1);
6000 assert_eq!(&back[0], &array);
6001 }
6002
6003 let s = Arc::new(StructArray::new_empty_fields(5, None)) as ArrayRef;
6004 run_test(s);
6005
6006 let s = Arc::new(StructArray::new_empty_fields(
6007 5,
6008 Some(vec![true, false, true, false, false].into()),
6009 )) as ArrayRef;
6010 run_test(s);
6011 }
6012
6013 #[test]
6014 fn reserve_should_increase_capacity_to_the_requested_size() {
6015 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6016 let mut empty_rows = row_converter.empty_rows(0, 0);
6017 empty_rows.reserve(50, 50);
6018 let before_size = empty_rows.size();
6019 empty_rows.reserve(50, 50);
6020 assert_eq!(
6021 empty_rows.size(),
6022 before_size,
6023 "Size should not change when reserving already reserved space"
6024 );
6025 empty_rows.reserve(10, 20);
6026 assert_eq!(
6027 empty_rows.size(),
6028 before_size,
6029 "Size should not change when already have space for the expected reserved data"
6030 );
6031
6032 empty_rows.reserve(100, 20);
6033 assert!(
6034 empty_rows.size() > before_size,
6035 "Size should increase when reserving more space than previously reserved"
6036 );
6037
6038 let before_size = empty_rows.size();
6039
6040 empty_rows.reserve(20, 100);
6041 assert!(
6042 empty_rows.size() > before_size,
6043 "Size should increase when reserving more space than previously reserved"
6044 );
6045 }
6046
6047 #[test]
6048 fn empty_rows_should_return_empty_lengths_iterator() {
6049 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6050 .unwrap()
6051 .empty_rows(0, 0);
6052 let mut lengths_iter = rows.lengths();
6053 assert_eq!(lengths_iter.next(), None);
6054 }
6055
6056 #[test]
6057 #[should_panic(expected = "row index out of bounds")]
6058 fn row_should_panic_on_overflowing_index() {
6059 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6060 .unwrap()
6061 .empty_rows(0, 0);
6062 rows.row(usize::MAX);
6063 }
6064
6065 #[test]
6066 #[should_panic(expected = "row index out of bounds")]
6067 fn row_len_should_panic_on_overflowing_index() {
6068 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6069 .unwrap()
6070 .empty_rows(0, 0);
6071 rows.row_len(usize::MAX);
6072 }
6073
6074 #[test]
6075 fn test_nested_null_list() {
6076 let null_array = Arc::new(NullArray::new(3));
6077 let list: ArrayRef = Arc::new(ListArray::new(
6079 Field::new_list_field(DataType::Null, true).into(),
6080 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6081 null_array,
6082 None,
6083 ));
6084
6085 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6086 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6087 let back = converter.convert_rows(&rows).unwrap();
6088
6089 assert_eq!(&list, &back[0]);
6090 }
6091
6092 #[test]
6094 fn test_double_nested_null_list() {
6095 let null_array = Arc::new(NullArray::new(1));
6096 let nested_field = Arc::new(Field::new_list_field(DataType::Null, true));
6098 let nested_list = Arc::new(ListArray::new(
6099 nested_field.clone(),
6100 OffsetBuffer::from_lengths(vec![1]),
6101 null_array,
6102 None,
6103 ));
6104 let list = Arc::new(ListArray::new(
6106 Field::new_list_field(DataType::List(nested_field), true).into(),
6107 OffsetBuffer::from_lengths(vec![1]),
6108 nested_list,
6109 None,
6110 )) as ArrayRef;
6111
6112 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6113 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6114 let back = converter.convert_rows(&rows).unwrap();
6115
6116 assert_eq!(&list, &back[0]);
6117 }
6118
6119 #[test]
6121 fn test_large_list_null() {
6122 let null_array = Arc::new(NullArray::new(3));
6123 let list: ArrayRef = Arc::new(LargeListArray::new(
6125 Field::new_list_field(DataType::Null, true).into(),
6126 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6127 null_array,
6128 None,
6129 ));
6130
6131 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6132 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6133 let back = converter.convert_rows(&rows).unwrap();
6134
6135 assert_eq!(&list, &back[0]);
6136 }
6137
6138 #[test]
6140 fn test_fixed_size_list_null() {
6141 let null_array = Arc::new(NullArray::new(6));
6142 let list: ArrayRef = Arc::new(FixedSizeListArray::new(
6144 Arc::new(Field::new_list_field(DataType::Null, true)),
6145 2,
6146 null_array,
6147 None,
6148 ));
6149
6150 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6151 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6152 let back = converter.convert_rows(&rows).unwrap();
6153
6154 assert_eq!(&list, &back[0]);
6155 }
6156
6157 #[test]
6159 fn test_fixed_size_list_of_dictionaries_round_trips() {
6160 let dict_dt = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
6163 let element_field = Arc::new(Field::new("item", dict_dt.clone(), true));
6164 let fsl_dt = DataType::FixedSizeList(Arc::clone(&element_field), 2);
6165
6166 let values = Arc::new(StringArray::from(vec!["a", "b"]));
6167 let keys = Int32Array::from(vec![0, 1]);
6168 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
6169 let fsl: ArrayRef = Arc::new(FixedSizeListArray::new(
6170 Arc::clone(&element_field),
6171 2,
6172 Arc::new(dict),
6173 None,
6174 ));
6175
6176 assert!(RowConverter::supports_fields(&[SortField::new(
6177 fsl_dt.clone()
6178 )]));
6179
6180 let converter = RowConverter::new(vec![SortField::new(fsl_dt.clone())]).unwrap();
6181 let rows = converter.convert_columns(&[Arc::clone(&fsl)]).unwrap();
6182
6183 let back = converter.convert_rows(&rows).unwrap();
6186 assert_eq!(back.len(), 1);
6187
6188 let out = back[0]
6192 .as_any()
6193 .downcast_ref::<FixedSizeListArray>()
6194 .expect("decoded array must be a FixedSizeListArray");
6195 assert_eq!(out.len(), 1);
6196 assert_eq!(out.value_length(), 2);
6197 assert_eq!(out.values().data_type(), &DataType::Utf8);
6201
6202 let values = out
6204 .values()
6205 .as_any()
6206 .downcast_ref::<StringArray>()
6207 .expect("child must be a StringArray after flattening");
6208 assert_eq!(values.value(0), "a");
6209 assert_eq!(values.value(1), "b");
6210 }
6211
6212 #[test]
6214 fn test_list_null_variations() {
6215 let null_array = Arc::new(NullArray::new(3));
6217 let list: ArrayRef = Arc::new(ListArray::new(
6218 Field::new_list_field(DataType::Null, true).into(),
6219 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6220 null_array,
6221 None,
6222 ));
6223
6224 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6225 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6226 let back = converter.convert_rows(&rows).unwrap();
6227 assert_eq!(&list, &back[0]);
6228
6229 let null_array = Arc::new(NullArray::new(3));
6231 let list: ArrayRef = Arc::new(ListArray::new(
6232 Field::new_list_field(DataType::Null, true).into(),
6233 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6234 null_array,
6235 Some(vec![true, false, true].into()),
6236 ));
6237
6238 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6239 let back = converter.convert_rows(&rows).unwrap();
6240 assert_eq!(&list, &back[0]);
6241
6242 let null_array = Arc::new(NullArray::new(0));
6244 let list: ArrayRef = Arc::new(ListArray::new(
6245 Field::new_list_field(DataType::Null, true).into(),
6246 OffsetBuffer::from_lengths(vec![]),
6247 null_array,
6248 None,
6249 ));
6250
6251 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6252 let back = converter.convert_rows(&rows).unwrap();
6253 assert_eq!(&list, &back[0]);
6254
6255 let null_array = Arc::new(NullArray::new(0));
6257 let list: ArrayRef = Arc::new(ListArray::new(
6258 Field::new_list_field(DataType::Null, true).into(),
6259 OffsetBuffer::from_lengths(vec![0, 0, 0]),
6260 null_array,
6261 None,
6262 ));
6263
6264 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6265 let back = converter.convert_rows(&rows).unwrap();
6266 assert_eq!(&list, &back[0]);
6267 }
6268
6269 #[test]
6271 fn test_list_null_descending() {
6272 let null_array = Arc::new(NullArray::new(3));
6273 let list: ArrayRef = Arc::new(ListArray::new(
6275 Field::new_list_field(DataType::Null, true).into(),
6276 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6277 null_array,
6278 None,
6279 ));
6280
6281 let options = SortOptions::default().with_descending(true);
6282 let field = SortField::new_with_options(list.data_type().clone(), options);
6283 let converter = RowConverter::new(vec![field]).unwrap();
6284 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6285 let back = converter.convert_rows(&rows).unwrap();
6286
6287 assert_eq!(&list, &back[0]);
6288 }
6289
6290 #[test]
6292 fn test_struct_with_null_field() {
6293 let null_array = Arc::new(NullArray::new(3));
6295 let int_array = Arc::new(Int32Array::from(vec![1, 2, 3]));
6296
6297 let struct_array: ArrayRef = Arc::new(StructArray::new(
6298 vec![
6299 Arc::new(Field::new("a", DataType::Null, true)),
6300 Arc::new(Field::new("b", DataType::Int32, true)),
6301 ]
6302 .into(),
6303 vec![null_array, int_array],
6304 Some(vec![true, true, false].into()), ));
6306
6307 let converter =
6308 RowConverter::new(vec![SortField::new(struct_array.data_type().clone())]).unwrap();
6309 let rows = converter
6310 .convert_columns(&[Arc::clone(&struct_array)])
6311 .unwrap();
6312 let back = converter.convert_rows(&rows).unwrap();
6313
6314 assert_eq!(&struct_array, &back[0]);
6315 }
6316
6317 #[test]
6319 fn test_nested_struct_with_null() {
6320 let inner_null = Arc::new(NullArray::new(2));
6322 let inner_struct = Arc::new(StructArray::new(
6323 vec![Arc::new(Field::new("x", DataType::Null, true))].into(),
6324 vec![inner_null],
6325 None,
6326 ));
6327
6328 let y_array = Arc::new(Int32Array::from(vec![10, 20]));
6330 let outer_struct: ArrayRef = Arc::new(StructArray::new(
6331 vec![
6332 Arc::new(Field::new("inner", inner_struct.data_type().clone(), true)),
6333 Arc::new(Field::new("y", DataType::Int32, true)),
6334 ]
6335 .into(),
6336 vec![inner_struct, y_array],
6337 None,
6338 ));
6339
6340 let converter =
6341 RowConverter::new(vec![SortField::new(outer_struct.data_type().clone())]).unwrap();
6342 let rows = converter
6343 .convert_columns(&[Arc::clone(&outer_struct)])
6344 .unwrap();
6345 let back = converter.convert_rows(&rows).unwrap();
6346
6347 assert_eq!(&outer_struct, &back[0]);
6348 }
6349
6350 #[test]
6352 fn test_map_null_variations() {
6353 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6355 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6356
6357 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6358 let entries_fields = vec![
6359 Arc::new(Field::new(
6360 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6361 DataType::Utf8,
6362 false,
6363 )),
6364 Arc::new(Field::new(
6365 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6366 DataType::Null,
6367 true,
6368 )),
6369 ];
6370 let struct_field = Arc::new(Field::new(
6371 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6372 DataType::Struct(entries_fields.clone().into()),
6373 false,
6374 ));
6375 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6376
6377 let map: ArrayRef = Arc::new(MapArray::new(
6378 struct_field.clone(),
6379 offsets,
6380 entries,
6381 None,
6382 false,
6383 ));
6384
6385 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6386 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6387 let back = converter.convert_rows(&rows).unwrap();
6388 assert_eq!(back.len(), 1);
6389 back[0].to_data().validate_full().unwrap();
6390 assert_eq!(&map, &back[0]);
6391
6392 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6394 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6395
6396 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6397 let entries_fields = vec![
6398 Arc::new(Field::new(
6399 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6400 DataType::Utf8,
6401 false,
6402 )),
6403 Arc::new(Field::new(
6404 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6405 DataType::Null,
6406 true,
6407 )),
6408 ];
6409 let struct_field = Arc::new(Field::new(
6410 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6411 DataType::Struct(entries_fields.clone().into()),
6412 false,
6413 ));
6414 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6415
6416 let map: ArrayRef = Arc::new(MapArray::new(
6417 struct_field.clone(),
6418 offsets,
6419 entries,
6420 Some(vec![true, false, true].into()),
6421 false,
6422 ));
6423
6424 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6425 let back = converter.convert_rows(&rows).unwrap();
6426 assert_eq!(back.len(), 1);
6427 back[0].to_data().validate_full().unwrap();
6428 assert_eq!(&map, &back[0]);
6429
6430 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6432 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6433
6434 let offsets = OffsetBuffer::new(vec![0i32].into());
6435 let entries_fields = vec![
6436 Arc::new(Field::new(
6437 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6438 DataType::Utf8,
6439 false,
6440 )),
6441 Arc::new(Field::new(
6442 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6443 DataType::Null,
6444 true,
6445 )),
6446 ];
6447 let struct_field = Arc::new(Field::new(
6448 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6449 DataType::Struct(entries_fields.clone().into()),
6450 false,
6451 ));
6452 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6453
6454 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6455
6456 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6457 let back = converter.convert_rows(&rows).unwrap();
6458 assert_eq!(back.len(), 1);
6459 back[0].to_data().validate_full().unwrap();
6460 assert_eq!(&map, &back[0]);
6461 }
6462
6463 #[test]
6465 fn test_map_null_descending() {
6466 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6468 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6469
6470 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6471 let entries_fields = vec![
6472 Arc::new(Field::new(
6473 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6474 DataType::Utf8,
6475 false,
6476 )),
6477 Arc::new(Field::new(
6478 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6479 DataType::Null,
6480 true,
6481 )),
6482 ];
6483 let struct_field = Arc::new(Field::new(
6484 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6485 DataType::Struct(entries_fields.clone().into()),
6486 false,
6487 ));
6488 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6489
6490 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6491
6492 let options = SortOptions::default().with_descending(true);
6493 let field = SortField::new_with_options(map.data_type().clone(), options);
6494 let converter = RowConverter::new(vec![field]).unwrap();
6495 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6496 let back = converter.convert_rows(&rows).unwrap();
6497 assert_eq!(back.len(), 1);
6498 back[0].to_data().validate_full().unwrap();
6499 assert_eq!(&map, &back[0]);
6500 }
6501
6502 #[test]
6504 fn test_map_null_all_empty() {
6505 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6506 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6507
6508 let offsets = OffsetBuffer::new(vec![0, 0, 0, 0].into());
6509 let entries_fields = vec![
6510 Arc::new(Field::new(
6511 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6512 DataType::Utf8,
6513 false,
6514 )),
6515 Arc::new(Field::new(
6516 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6517 DataType::Null,
6518 true,
6519 )),
6520 ];
6521 let struct_field = Arc::new(Field::new(
6522 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6523 DataType::Struct(entries_fields.clone().into()),
6524 false,
6525 ));
6526 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6527
6528 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6529
6530 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6531 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6532
6533 assert_eq!(rows.row(0), rows.row(1));
6535 assert_eq!(rows.row(1), rows.row(2));
6536
6537 let back = converter.convert_rows(&rows).unwrap();
6538 assert_eq!(back.len(), 1);
6539 back[0].to_data().validate_full().unwrap();
6540 assert_eq!(&map, &back[0]);
6541 }
6542
6543 #[test]
6545 fn test_nested_map_null() {
6546 let inner_keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6548 let inner_null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6549
6550 let inner_entries_fields = vec![
6551 Arc::new(Field::new(
6552 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6553 DataType::Utf8,
6554 false,
6555 )),
6556 Arc::new(Field::new(
6557 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6558 DataType::Null,
6559 true,
6560 )),
6561 ];
6562 let inner_struct_field = Arc::new(Field::new(
6563 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6564 DataType::Struct(inner_entries_fields.clone().into()),
6565 false,
6566 ));
6567 let inner_entries = StructArray::new(
6568 inner_entries_fields.clone().into(),
6569 vec![inner_keys, inner_null_values],
6570 None,
6571 );
6572
6573 let inner_map = Arc::new(MapArray::new(
6575 inner_struct_field.clone(),
6576 OffsetBuffer::new(vec![0, 1, 3].into()),
6577 inner_entries,
6578 None,
6579 false,
6580 )) as ArrayRef;
6581
6582 let outer_keys = Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef;
6584
6585 let inner_map_type = DataType::Map(inner_struct_field.clone(), false);
6586 let outer_entries_fields = vec![
6587 Arc::new(Field::new(
6588 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6589 DataType::Utf8,
6590 false,
6591 )),
6592 Arc::new(Field::new(
6593 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6594 inner_map_type,
6595 true,
6596 )),
6597 ];
6598 let outer_struct_field = Arc::new(Field::new(
6599 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6600 DataType::Struct(outer_entries_fields.clone().into()),
6601 false,
6602 ));
6603 let outer_entries = StructArray::new(
6604 outer_entries_fields.into(),
6605 vec![outer_keys, inner_map],
6606 None,
6607 );
6608
6609 let map: ArrayRef = Arc::new(MapArray::new(
6611 outer_struct_field,
6612 OffsetBuffer::new(vec![0, 1, 2].into()),
6613 outer_entries,
6614 None,
6615 false,
6616 ));
6617
6618 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6619 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6620 let back = converter.convert_rows(&rows).unwrap();
6621 assert_eq!(back.len(), 1);
6622 back[0].to_data().validate_full().unwrap();
6623 assert_eq!(&map, &back[0]);
6624 }
6625
6626 #[test]
6628 fn test_list_of_map_null() {
6629 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6631 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6632
6633 let entries_fields = vec![
6634 Arc::new(Field::new(
6635 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6636 DataType::Utf8,
6637 false,
6638 )),
6639 Arc::new(Field::new(
6640 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6641 DataType::Null,
6642 true,
6643 )),
6644 ];
6645 let struct_field = Arc::new(Field::new(
6646 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6647 DataType::Struct(entries_fields.clone().into()),
6648 false,
6649 ));
6650 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6651
6652 let map_array = Arc::new(MapArray::new(
6654 struct_field.clone(),
6655 OffsetBuffer::new(vec![0, 1, 1, 3].into()),
6656 entries,
6657 None,
6658 false,
6659 )) as ArrayRef;
6660
6661 let map_type = DataType::Map(struct_field, false);
6662 let list: ArrayRef = Arc::new(ListArray::new(
6664 Arc::new(Field::new_list_field(map_type, true)),
6665 OffsetBuffer::new(vec![0, 1, 3].into()),
6666 map_array,
6667 None,
6668 ));
6669
6670 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6671 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6672 let back = converter.convert_rows(&rows).unwrap();
6673 assert_eq!(&list, &back[0]);
6674 }
6675
6676 #[test]
6678 fn test_map_of_list_null() {
6679 let null_array = Arc::new(NullArray::new(3)) as ArrayRef;
6681 let list_array = Arc::new(ListArray::new(
6683 Arc::new(Field::new_list_field(DataType::Null, true)),
6684 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6685 null_array,
6686 None,
6687 )) as ArrayRef;
6688
6689 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6690
6691 let list_type = list_array.data_type().clone();
6692 let entries_fields = vec![
6693 Arc::new(Field::new(
6694 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6695 DataType::Utf8,
6696 false,
6697 )),
6698 Arc::new(Field::new(
6699 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6700 list_type,
6701 true,
6702 )),
6703 ];
6704 let struct_field = Arc::new(Field::new(
6705 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6706 DataType::Struct(entries_fields.clone().into()),
6707 false,
6708 ));
6709 let entries = StructArray::new(entries_fields.into(), vec![keys, list_array], None);
6710
6711 let map: ArrayRef = Arc::new(MapArray::new(
6713 struct_field,
6714 OffsetBuffer::new(vec![0, 3].into()),
6715 entries,
6716 None,
6717 false,
6718 ));
6719
6720 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6721 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6722 let back = converter.convert_rows(&rows).unwrap();
6723 assert_eq!(back.len(), 1);
6724 back[0].to_data().validate_full().unwrap();
6725 assert_eq!(&map, &back[0]);
6726 }
6727
6728 #[test]
6729 fn empty_row_iter_next_back() {
6730 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6731 .unwrap()
6732 .empty_rows(0, 0);
6733 let mut rows_iter = rows.iter();
6734 assert_eq!(rows_iter.next_back(), None);
6735 assert_eq!(rows_iter.next_back(), None);
6736 assert_eq!(rows_iter.next_back(), None);
6737 }
6738
6739 #[test]
6741 fn test_row_parser_skip_utf8_validation_roundtrip() {
6742 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
6743 let array = StringArray::from(vec!["arrow", "rust"]);
6744 let rows = converter.convert_columns(&[Arc::new(array) as _]).unwrap();
6745 let binary = rows.try_into_binary().expect("fits in i32 offsets");
6746
6747 let parser = unsafe { RowParser::with_skip_utf8_validate(Arc::clone(&converter.fields)) };
6749
6750 let decoded = converter
6751 .convert_rows(binary.iter().map(|b| parser.parse(b.unwrap())))
6752 .unwrap();
6753 let got: Vec<_> = decoded[0].as_string::<i32>().iter().flatten().collect();
6754 assert_eq!(got, vec!["arrow", "rust"]);
6755 }
6756
6757 #[test]
6758 fn row_iter_next_back() {
6759 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6760 let mut rng = StdRng::seed_from_u64(42);
6761 let array = generate_primitive_array::<UInt8Type>(&mut rng, 100, 0.8);
6762 let rows = row_converter.convert_columns(&[Arc::new(array)]).unwrap();
6763
6764 let mut rows_iter = rows.iter();
6765 let mut bytes: Vec<u8> = vec![];
6766
6767 while let Some(row) = rows_iter.next_back() {
6768 bytes.extend(row.data.iter().rev());
6769 }
6770
6771 bytes.reverse();
6772
6773 assert_eq!(
6774 bytes,
6775 &rows.buffer.as_slice()[..*rows.offsets.last().unwrap()]
6776 );
6777
6778 assert_eq!(rows_iter.next_back(), None);
6779 assert_eq!(rows_iter.next(), None);
6780 }
6781}