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, Clone)]
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 fn count(self) -> usize
1532 where
1533 Self: Sized,
1534 {
1535 self.len()
1536 }
1537
1538 #[inline]
1539 fn nth(&mut self, n: usize) -> Option<Self::Item> {
1540 match self.start.checked_add(n) {
1545 Some(new_offset) if new_offset < self.end => {
1547 self.start = new_offset;
1548 }
1549
1550 _ => {
1552 self.start = self.end;
1553 return None;
1554 }
1555 }
1556
1557 self.next()
1558 }
1559
1560 fn last(mut self) -> Option<Self::Item> {
1561 if self.start == self.end {
1563 return None;
1564 }
1565
1566 self.start = self.end - 1;
1568
1569 self.next()
1571 }
1572}
1573
1574impl ExactSizeIterator for RowsIter<'_> {
1575 fn len(&self) -> usize {
1576 self.end - self.start
1577 }
1578}
1579
1580impl DoubleEndedIterator for RowsIter<'_> {
1581 fn next_back(&mut self) -> Option<Self::Item> {
1582 if self.end == self.start {
1583 return None;
1584 }
1585
1586 self.end -= 1;
1587
1588 let row = unsafe { self.rows.row_unchecked(self.end) };
1591 Some(row)
1592 }
1593
1594 fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
1595 match self.end.checked_sub(n) {
1600 Some(new_offset) if self.start < new_offset => {
1602 self.end = new_offset;
1603 }
1604
1605 _ => {
1607 self.start = self.end;
1608 return None;
1609 }
1610 }
1611
1612 self.next_back()
1613 }
1614}
1615
1616#[derive(Debug, Copy, Clone)]
1625pub struct Row<'a> {
1626 data: &'a [u8],
1627 config: &'a RowConfig,
1628}
1629
1630impl<'a> Row<'a> {
1631 pub fn owned(&self) -> OwnedRow {
1633 OwnedRow {
1634 data: self.data.into(),
1635 config: self.config.clone(),
1636 }
1637 }
1638
1639 pub fn data(&self) -> &'a [u8] {
1641 self.data
1642 }
1643}
1644
1645impl PartialEq for Row<'_> {
1648 #[inline]
1649 fn eq(&self, other: &Self) -> bool {
1650 self.data.eq(other.data)
1651 }
1652}
1653
1654impl Eq for Row<'_> {}
1655
1656impl PartialOrd for Row<'_> {
1657 #[inline]
1658 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1659 Some(self.cmp(other))
1660 }
1661}
1662
1663impl Ord for Row<'_> {
1664 #[inline]
1665 fn cmp(&self, other: &Self) -> Ordering {
1666 self.data.cmp(other.data)
1667 }
1668}
1669
1670impl Hash for Row<'_> {
1671 #[inline]
1672 fn hash<H: Hasher>(&self, state: &mut H) {
1673 self.data.hash(state)
1674 }
1675}
1676
1677impl AsRef<[u8]> for Row<'_> {
1678 #[inline]
1679 fn as_ref(&self) -> &[u8] {
1680 self.data
1681 }
1682}
1683
1684#[derive(Debug, Clone)]
1688pub struct OwnedRow {
1689 data: Box<[u8]>,
1690 config: RowConfig,
1691}
1692
1693impl OwnedRow {
1694 pub fn row(&self) -> Row<'_> {
1698 Row {
1699 data: &self.data,
1700 config: &self.config,
1701 }
1702 }
1703}
1704
1705impl PartialEq for OwnedRow {
1708 #[inline]
1709 fn eq(&self, other: &Self) -> bool {
1710 self.row().eq(&other.row())
1711 }
1712}
1713
1714impl Eq for OwnedRow {}
1715
1716impl PartialOrd for OwnedRow {
1717 #[inline]
1718 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1719 Some(self.cmp(other))
1720 }
1721}
1722
1723impl Ord for OwnedRow {
1724 #[inline]
1725 fn cmp(&self, other: &Self) -> Ordering {
1726 self.row().cmp(&other.row())
1727 }
1728}
1729
1730impl Hash for OwnedRow {
1731 #[inline]
1732 fn hash<H: Hasher>(&self, state: &mut H) {
1733 self.row().hash(state)
1734 }
1735}
1736
1737impl AsRef<[u8]> for OwnedRow {
1738 #[inline]
1739 fn as_ref(&self) -> &[u8] {
1740 &self.data
1741 }
1742}
1743
1744#[inline]
1746fn null_sentinel(options: SortOptions) -> u8 {
1747 match options.nulls_first {
1748 true => 0,
1749 false => 0xFF,
1750 }
1751}
1752
1753enum LengthTracker {
1755 Fixed { length: usize, num_rows: usize },
1757 Variable {
1759 fixed_length: usize,
1760 lengths: Vec<usize>,
1761 },
1762}
1763
1764impl LengthTracker {
1765 fn new(num_rows: usize) -> Self {
1766 Self::Fixed {
1767 length: 0,
1768 num_rows,
1769 }
1770 }
1771
1772 fn push_fixed(&mut self, new_length: usize) {
1774 match self {
1775 LengthTracker::Fixed { length, .. } => *length += new_length,
1776 LengthTracker::Variable { fixed_length, .. } => *fixed_length += new_length,
1777 }
1778 }
1779
1780 fn push_variable(&mut self, new_lengths: impl ExactSizeIterator<Item = usize>) {
1782 match self {
1783 LengthTracker::Fixed { length, .. } => {
1784 *self = LengthTracker::Variable {
1785 fixed_length: *length,
1786 lengths: new_lengths.collect(),
1787 }
1788 }
1789 LengthTracker::Variable { lengths, .. } => {
1790 assert_eq!(lengths.len(), new_lengths.len());
1791 lengths
1792 .iter_mut()
1793 .zip(new_lengths)
1794 .for_each(|(length, new_length)| *length += new_length);
1795 }
1796 }
1797 }
1798
1799 fn materialized(&mut self) -> &mut [usize] {
1801 if let LengthTracker::Fixed { length, num_rows } = *self {
1802 *self = LengthTracker::Variable {
1803 fixed_length: length,
1804 lengths: vec![0; num_rows],
1805 };
1806 }
1807
1808 match self {
1809 LengthTracker::Variable { lengths, .. } => lengths,
1810 LengthTracker::Fixed { .. } => unreachable!(),
1811 }
1812 }
1813
1814 fn extend_offsets(&self, initial_offset: usize, offsets: &mut Vec<usize>) -> usize {
1832 match self {
1833 LengthTracker::Fixed { length, num_rows } => {
1834 offsets.extend((0..*num_rows).map(|i| initial_offset + i * length));
1835
1836 initial_offset + num_rows * length
1837 }
1838 LengthTracker::Variable {
1839 fixed_length,
1840 lengths,
1841 } => {
1842 let mut acc = initial_offset;
1843
1844 offsets.extend(lengths.iter().map(|length| {
1845 let current = acc;
1846 acc += length + fixed_length;
1847 current
1848 }));
1849
1850 acc
1851 }
1852 }
1853 }
1854}
1855
1856fn row_lengths(cols: &[ArrayRef], encoders: &[Encoder]) -> LengthTracker {
1858 use fixed::FixedLengthEncoding;
1859
1860 let num_rows = cols.first().map(|x| x.len()).unwrap_or(0);
1861 let mut tracker = LengthTracker::new(num_rows);
1862
1863 for (array, encoder) in cols.iter().zip(encoders) {
1864 match encoder {
1865 Encoder::Stateless => {
1866 downcast_primitive_array! {
1867 array => tracker.push_fixed(fixed::encoded_len(array)),
1868 DataType::Null => tracker.push_fixed(2)
1869 DataType::Boolean => tracker.push_fixed(bool::ENCODED_LEN),
1870 DataType::Binary => push_generic_byte_array_lengths(&mut tracker, as_generic_binary_array::<i32>(array)),
1871 DataType::LargeBinary => push_generic_byte_array_lengths(&mut tracker, as_generic_binary_array::<i64>(array)),
1872 DataType::BinaryView => push_byte_view_array_lengths(&mut tracker, array.as_binary_view()),
1873 DataType::Utf8 => push_generic_byte_array_lengths(&mut tracker, array.as_string::<i32>()),
1874 DataType::LargeUtf8 => push_generic_byte_array_lengths(&mut tracker, array.as_string::<i64>()),
1875 DataType::Utf8View => push_byte_view_array_lengths(&mut tracker, array.as_string_view()),
1876 DataType::FixedSizeBinary(len) => {
1877 let len = len.to_usize().unwrap();
1878 tracker.push_fixed(1 + len)
1879 }
1880 _ => unimplemented!("unsupported data type: {}", array.data_type()),
1881 }
1882 }
1883 Encoder::Dictionary(values, null) => {
1884 downcast_dictionary_array! {
1885 array => {
1886 tracker.push_variable(
1887 array.keys().iter().map(|v| match v {
1888 Some(k) => values.row_len(k.as_usize()),
1889 None => null.data.len(),
1890 })
1891 )
1892 }
1893 _ => unreachable!(),
1894 }
1895 }
1896 Encoder::Struct(rows, null) => {
1897 let array = as_struct_array(array);
1898 if rows.num_rows() > 0 {
1899 tracker.push_variable((0..array.len()).map(|idx| match array.is_valid(idx) {
1901 true => 1 + rows.row_len(idx),
1902 false => 1 + null.data.len(),
1903 }));
1904 } else {
1905 tracker.push_variable((0..array.len()).map(|idx| match array.is_valid(idx) {
1907 true => 1,
1908 false => 1 + null.data.len(),
1909 }));
1910 }
1911 }
1912 Encoder::List(rows) => match array.data_type() {
1913 DataType::List(_) => {
1914 list::compute_lengths(tracker.materialized(), rows, as_list_array(array))
1915 }
1916 DataType::LargeList(_) => {
1917 list::compute_lengths(tracker.materialized(), rows, as_large_list_array(array))
1918 }
1919 DataType::ListView(_) => {
1920 let list_view = array.as_list_view::<i32>();
1921 let (min_offset, _) = compute_list_view_bounds(list_view);
1922 list::compute_lengths_list_view(
1923 tracker.materialized(),
1924 rows,
1925 list_view,
1926 min_offset,
1927 )
1928 }
1929 DataType::LargeListView(_) => {
1930 let list_view = array.as_list_view::<i64>();
1931 let (min_offset, _) = compute_list_view_bounds(list_view);
1932 list::compute_lengths_list_view(
1933 tracker.materialized(),
1934 rows,
1935 list_view,
1936 min_offset,
1937 )
1938 }
1939 DataType::FixedSizeList(_, _) => compute_lengths_fixed_size_list(
1940 &mut tracker,
1941 rows,
1942 as_fixed_size_list_array(array),
1943 ),
1944 _ => unreachable!(),
1945 },
1946 Encoder::Map(rows) => {
1947 list::compute_lengths(tracker.materialized(), rows, as_map_array(array))
1948 }
1949 Encoder::RunEndEncoded(rows) => match array.data_type() {
1950 DataType::RunEndEncoded(r, _) => match r.data_type() {
1951 DataType::Int16 => run::compute_lengths(
1952 tracker.materialized(),
1953 rows,
1954 array.as_run::<Int16Type>(),
1955 ),
1956 DataType::Int32 => run::compute_lengths(
1957 tracker.materialized(),
1958 rows,
1959 array.as_run::<Int32Type>(),
1960 ),
1961 DataType::Int64 => run::compute_lengths(
1962 tracker.materialized(),
1963 rows,
1964 array.as_run::<Int64Type>(),
1965 ),
1966 _ => unreachable!("Unsupported run end index type: {r:?}"),
1967 },
1968 _ => unreachable!(),
1969 },
1970 Encoder::Union {
1971 child_rows,
1972 field_to_type_ids,
1973 type_ids,
1974 offsets,
1975 } => {
1976 let union_array = array
1977 .as_any()
1978 .downcast_ref::<UnionArray>()
1979 .expect("expected UnionArray");
1980
1981 let mut type_id_to_field_idx = [0usize; 128];
1982 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
1983 type_id_to_field_idx[type_id as usize] = field_idx;
1984 }
1985
1986 let lengths = (0..union_array.len()).map(|i| {
1987 let type_id = type_ids[i];
1988 let field_idx = type_id_to_field_idx[type_id as usize];
1989 let child_row_i = offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i);
1990 let child_row_len = child_rows[field_idx].row_len(child_row_i);
1991
1992 1 + child_row_len
1994 });
1995
1996 tracker.push_variable(lengths);
1997 }
1998 }
1999 }
2000
2001 tracker
2002}
2003
2004fn push_generic_byte_array_lengths<T: ByteArrayType>(
2006 tracker: &mut LengthTracker,
2007 array: &GenericByteArray<T>,
2008) {
2009 if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
2010 tracker.push_variable(
2011 array
2012 .offsets()
2013 .lengths()
2014 .zip(nulls.iter())
2015 .map(|(length, is_valid)| if is_valid { Some(length) } else { None })
2016 .map(variable::padded_length),
2017 )
2018 } else {
2019 tracker.push_variable(
2020 array
2021 .offsets()
2022 .lengths()
2023 .map(variable::non_null_padded_length),
2024 )
2025 }
2026}
2027
2028fn push_byte_view_array_lengths<T: ByteViewType>(
2030 tracker: &mut LengthTracker,
2031 array: &GenericByteViewArray<T>,
2032) {
2033 if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
2034 tracker.push_variable(
2035 array
2036 .lengths()
2037 .zip(nulls.iter())
2038 .map(|(length, is_valid)| {
2039 if is_valid {
2040 Some(length as usize)
2041 } else {
2042 None
2043 }
2044 })
2045 .map(variable::padded_length),
2046 )
2047 } else {
2048 tracker.push_variable(
2049 array
2050 .lengths()
2051 .map(|len| variable::padded_length(Some(len as usize))),
2052 )
2053 }
2054}
2055
2056fn encode_column(
2058 data: &mut [u8],
2059 offsets: &mut [usize],
2060 column: &dyn Array,
2061 opts: SortOptions,
2062 encoder: &Encoder<'_>,
2063) {
2064 match encoder {
2065 Encoder::Stateless => {
2066 downcast_primitive_array! {
2067 column => {
2068 if let Some(nulls) = column.nulls().filter(|n| n.null_count() > 0){
2069 fixed::encode(data, offsets, column.values(), nulls, opts)
2070 } else {
2071 fixed::encode_not_null(data, offsets, column.values(), opts)
2072 }
2073 }
2074 DataType::Null => {
2075 for offset in offsets.iter_mut().skip(1) {
2076 variable::encode_null_value(&mut data[*offset..], opts);
2077 *offset += 2;
2078 }
2079 }
2080 DataType::Boolean => {
2081 if let Some(nulls) = column.nulls().filter(|n| n.null_count() > 0){
2082 fixed::encode_boolean(data, offsets, column.as_boolean().values(), nulls, opts)
2083 } else {
2084 fixed::encode_boolean_not_null(data, offsets, column.as_boolean().values(), opts)
2085 }
2086 }
2087 DataType::Binary => {
2088 variable::encode_generic_byte_array(data, offsets, as_generic_binary_array::<i32>(column), opts)
2089 }
2090 DataType::BinaryView => {
2091 variable::encode(data, offsets, column.as_binary_view().iter(), opts)
2092 }
2093 DataType::LargeBinary => {
2094 variable::encode_generic_byte_array(data, offsets, as_generic_binary_array::<i64>(column), opts)
2095 }
2096 DataType::Utf8 => variable::encode_generic_byte_array(
2097 data, offsets,
2098 column.as_string::<i32>(),
2099 opts,
2100 ),
2101 DataType::LargeUtf8 => variable::encode_generic_byte_array(
2102 data, offsets,
2103 column.as_string::<i64>(),
2104 opts,
2105 ),
2106 DataType::Utf8View => variable::encode(
2107 data, offsets,
2108 column.as_string_view().iter().map(|x| x.map(|x| x.as_bytes())),
2109 opts,
2110 ),
2111 DataType::FixedSizeBinary(_) => {
2112 let array = column.as_any().downcast_ref().unwrap();
2113 fixed::encode_fixed_size_binary(data, offsets, array, opts)
2114 }
2115 _ => unimplemented!("unsupported data type: {}", column.data_type()),
2116 }
2117 }
2118 Encoder::Dictionary(values, nulls) => {
2119 downcast_dictionary_array! {
2120 column => encode_dictionary_values(data, offsets, column, values, nulls),
2121 _ => unreachable!()
2122 }
2123 }
2124 Encoder::Struct(rows, null) => {
2125 fn struct_encode_helper<const NO_CHILD_FIELDS: bool>(
2126 array: &StructArray,
2127 offsets: &mut [usize],
2128 null_sentinel: u8,
2129 rows: &Rows,
2130 null: &Row<'_>,
2131 data: &mut [u8],
2132 ) {
2133 let empty_row = Row {
2134 data: &[],
2135 config: &rows.config,
2136 };
2137
2138 offsets
2139 .iter_mut()
2140 .skip(1)
2141 .enumerate()
2142 .for_each(|(idx, offset)| {
2143 let (row, sentinel) = match array.is_valid(idx) {
2144 true => (
2145 if NO_CHILD_FIELDS {
2146 empty_row
2147 } else {
2148 rows.row(idx)
2149 },
2150 0x01,
2151 ),
2152 false => (*null, null_sentinel),
2153 };
2154 let end_offset = *offset + 1 + row.as_ref().len();
2155 data[*offset] = sentinel;
2156 data[*offset + 1..end_offset].copy_from_slice(row.as_ref());
2157 *offset = end_offset;
2158 })
2159 }
2160
2161 let array = as_struct_array(column);
2162 let null_sentinel = null_sentinel(opts);
2163 if rows.num_rows() == 0 {
2164 struct_encode_helper::<true>(array, offsets, null_sentinel, rows, null, data);
2166 } else {
2167 struct_encode_helper::<false>(array, offsets, null_sentinel, rows, null, data);
2168 }
2169 }
2170 Encoder::List(rows) => match column.data_type() {
2171 DataType::List(_) => list::encode(data, offsets, rows, opts, as_list_array(column)),
2172 DataType::LargeList(_) => {
2173 list::encode(data, offsets, rows, opts, as_large_list_array(column))
2174 }
2175 DataType::ListView(_) => {
2176 let list_view = column.as_list_view::<i32>();
2177 let (min_offset, _) = compute_list_view_bounds(list_view);
2178 list::encode_list_view(data, offsets, rows, opts, list_view, min_offset)
2179 }
2180 DataType::LargeListView(_) => {
2181 let list_view = column.as_list_view::<i64>();
2182 let (min_offset, _) = compute_list_view_bounds(list_view);
2183 list::encode_list_view(data, offsets, rows, opts, list_view, min_offset)
2184 }
2185 DataType::FixedSizeList(_, _) => {
2186 encode_fixed_size_list(data, offsets, rows, opts, as_fixed_size_list_array(column))
2187 }
2188 _ => unreachable!(),
2189 },
2190 Encoder::Map(rows) => list::encode(data, offsets, rows, opts, as_map_array(column)),
2191 Encoder::RunEndEncoded(rows) => match column.data_type() {
2192 DataType::RunEndEncoded(r, _) => match r.data_type() {
2193 DataType::Int16 => {
2194 run::encode(data, offsets, rows, opts, column.as_run::<Int16Type>())
2195 }
2196 DataType::Int32 => {
2197 run::encode(data, offsets, rows, opts, column.as_run::<Int32Type>())
2198 }
2199 DataType::Int64 => {
2200 run::encode(data, offsets, rows, opts, column.as_run::<Int64Type>())
2201 }
2202 _ => unreachable!("Unsupported run end index type: {r:?}"),
2203 },
2204 _ => unreachable!(),
2205 },
2206 Encoder::Union {
2207 child_rows,
2208 field_to_type_ids,
2209 type_ids,
2210 offsets: offsets_buf,
2211 } => {
2212 let mut type_id_to_field_idx = [0usize; 128];
2213 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
2214 type_id_to_field_idx[type_id as usize] = field_idx;
2215 }
2216
2217 offsets
2218 .iter_mut()
2219 .skip(1)
2220 .enumerate()
2221 .for_each(|(i, offset)| {
2222 let type_id = type_ids[i];
2223 let field_idx = type_id_to_field_idx[type_id as usize];
2224
2225 let child_row_idx = offsets_buf.as_ref().map(|o| o[i] as usize).unwrap_or(i);
2226 let child_row = child_rows[field_idx].row(child_row_idx);
2227 let child_bytes = child_row.as_ref();
2228
2229 let type_id_byte = if opts.descending {
2230 !(type_id as u8)
2231 } else {
2232 type_id as u8
2233 };
2234 data[*offset] = type_id_byte;
2235
2236 let child_start = *offset + 1;
2237 let child_end = child_start + child_bytes.len();
2238 data[child_start..child_end].copy_from_slice(child_bytes);
2239
2240 *offset = child_end;
2241 });
2242 }
2243 }
2244}
2245
2246pub fn encode_dictionary_values<K: ArrowDictionaryKeyType>(
2248 data: &mut [u8],
2249 offsets: &mut [usize],
2250 column: &DictionaryArray<K>,
2251 values: &Rows,
2252 null: &Row<'_>,
2253) {
2254 for (offset, k) in offsets.iter_mut().skip(1).zip(column.keys()) {
2255 let row = match k {
2256 Some(k) => values.row(k.as_usize()).data,
2257 None => null.data,
2258 };
2259 let end_offset = *offset + row.len();
2260 data[*offset..end_offset].copy_from_slice(row);
2261 *offset = end_offset;
2262 }
2263}
2264
2265macro_rules! decode_primitive_helper {
2266 ($t:ty, $rows:ident, $data_type:ident, $options:ident) => {
2267 Arc::new(decode_primitive::<$t>($rows, $data_type, $options))
2268 };
2269}
2270
2271unsafe fn decode_column(
2277 field: &SortField,
2278 rows: &mut [&[u8]],
2279 codec: &Codec,
2280 validate_utf8: bool,
2281) -> Result<ArrayRef, ArrowError> {
2282 let options = field.options;
2283
2284 let array: ArrayRef = match codec {
2285 Codec::Stateless => {
2286 let data_type = field.data_type.clone();
2287 downcast_primitive! {
2288 data_type => (decode_primitive_helper, rows, data_type, options),
2289 DataType::Null => {
2290 variable::decode_null_value(rows, options);
2291 Arc::new(NullArray::new(rows.len()))
2292 }
2293 DataType::Boolean => Arc::new(decode_bool(rows, options)),
2294 DataType::Binary => Arc::new(decode_binary::<i32>(rows, options)),
2295 DataType::LargeBinary => Arc::new(decode_binary::<i64>(rows, options)),
2296 DataType::BinaryView => Arc::new(decode_binary_view(rows, options)?),
2297 DataType::FixedSizeBinary(size) => Arc::new(decode_fixed_size_binary(rows, size, options)),
2298 DataType::Utf8 => Arc::new(unsafe{ decode_string::<i32>(rows, options, validate_utf8) }),
2299 DataType::LargeUtf8 => Arc::new(unsafe { decode_string::<i64>(rows, options, validate_utf8) }),
2300 DataType::Utf8View => Arc::new(unsafe { decode_string_view(rows, options, validate_utf8) }?),
2301 _ => return Err(ArrowError::NotYetImplemented(format!("unsupported data type: {data_type}" )))
2302 }
2303 }
2304 Codec::Dictionary(converter, _) => {
2305 let cols = unsafe { converter.convert_raw(rows, validate_utf8) }?;
2306 cols.into_iter().next().unwrap()
2307 }
2308 Codec::Struct(converter, _) => {
2309 let nulls = fixed::decode_nulls(rows);
2310 rows.iter_mut().for_each(|row| *row = &row[1..]);
2311 let children = unsafe { converter.convert_raw(rows, validate_utf8) }?;
2312
2313 let corrected_fields: Vec<Field> = match &field.data_type {
2316 DataType::Struct(struct_fields) => struct_fields
2317 .iter()
2318 .zip(children.iter())
2319 .map(|(orig_field, child_array)| {
2320 orig_field
2321 .as_ref()
2322 .clone()
2323 .with_data_type(child_array.data_type().clone())
2324 })
2325 .collect(),
2326 _ => unreachable!("Only Struct types should be corrected here"),
2327 };
2328
2329 Arc::new(unsafe {
2330 StructArray::new_unchecked_with_length(
2331 corrected_fields.into(),
2332 children,
2333 nulls,
2334 rows.len(),
2335 )
2336 })
2337 }
2338 Codec::List(converter) => match &field.data_type {
2339 DataType::List(_) => Arc::new(unsafe {
2340 list::decode::<GenericListArray<i32>>(converter, rows, field, validate_utf8)
2341 }?),
2342 DataType::LargeList(_) => Arc::new(unsafe {
2343 list::decode::<GenericListArray<i64>>(converter, rows, field, validate_utf8)
2344 }?),
2345 DataType::ListView(_) => Arc::new(unsafe {
2346 list::decode_list_view::<i32>(converter, rows, field, validate_utf8)
2347 }?),
2348 DataType::LargeListView(_) => Arc::new(unsafe {
2349 list::decode_list_view::<i64>(converter, rows, field, validate_utf8)
2350 }?),
2351 DataType::FixedSizeList(_, value_length) => Arc::new(unsafe {
2352 list::decode_fixed_size_list(
2353 converter,
2354 rows,
2355 field,
2356 validate_utf8,
2357 value_length.as_usize(),
2358 )
2359 }?),
2360 _ => unreachable!(),
2361 },
2362 Codec::Map(converter) => {
2363 Arc::new(unsafe { list::decode::<MapArray>(converter, rows, field, validate_utf8) }?)
2364 }
2365 Codec::RunEndEncoded(converter) => match &field.data_type {
2366 DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() {
2367 DataType::Int16 => Arc::new(unsafe {
2368 run::decode::<Int16Type>(converter, rows, field, validate_utf8)
2369 }?),
2370 DataType::Int32 => Arc::new(unsafe {
2371 run::decode::<Int32Type>(converter, rows, field, validate_utf8)
2372 }?),
2373 DataType::Int64 => Arc::new(unsafe {
2374 run::decode::<Int64Type>(converter, rows, field, validate_utf8)
2375 }?),
2376 _ => unreachable!(),
2377 },
2378 _ => unreachable!(),
2379 },
2380 Codec::Union(converters, field_to_type_ids, null_rows) => {
2381 let len = rows.len();
2382
2383 let DataType::Union(union_fields, mode) = &field.data_type else {
2384 unreachable!()
2385 };
2386
2387 let mut type_id_to_field_idx = [0usize; 128];
2388 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
2389 type_id_to_field_idx[type_id as usize] = field_idx;
2390 }
2391
2392 let mut type_ids = Vec::with_capacity(len);
2393 let mut rows_by_field: Vec<Vec<(usize, &[u8])>> = vec![Vec::new(); converters.len()];
2394
2395 for (idx, row) in rows.iter_mut().enumerate() {
2396 let type_id_byte = {
2397 let id = row[0];
2398 if options.descending { !id } else { id }
2399 };
2400
2401 let type_id = type_id_byte as i8;
2402 type_ids.push(type_id);
2403
2404 let field_idx = type_id_to_field_idx[type_id as usize];
2405
2406 let child_row = &row[1..];
2407 rows_by_field[field_idx].push((idx, child_row));
2408 }
2409
2410 let mut child_arrays: Vec<ArrayRef> = Vec::with_capacity(converters.len());
2411 let mut offsets = (*mode == UnionMode::Dense).then(|| Vec::with_capacity(len));
2412
2413 for (field_idx, converter) in converters.iter().enumerate() {
2414 let field_rows = &rows_by_field[field_idx];
2415
2416 match &mode {
2417 UnionMode::Dense => {
2418 if field_rows.is_empty() {
2419 let (_, field) = union_fields.iter().nth(field_idx).unwrap();
2420 child_arrays.push(arrow_array::new_empty_array(field.data_type()));
2421 continue;
2422 }
2423
2424 let mut child_data = field_rows
2425 .iter()
2426 .map(|(_, bytes)| *bytes)
2427 .collect::<Vec<_>>();
2428
2429 let child_array =
2430 unsafe { converter.convert_raw(&mut child_data, validate_utf8) }?;
2431
2432 for ((row_idx, original_bytes), remaining_bytes) in
2434 field_rows.iter().zip(child_data)
2435 {
2436 let consumed_length = 1 + original_bytes.len() - remaining_bytes.len();
2437 rows[*row_idx] = &rows[*row_idx][consumed_length..];
2438 }
2439
2440 child_arrays.push(child_array.into_iter().next().unwrap());
2441 }
2442 UnionMode::Sparse => {
2443 let mut sparse_data: Vec<&[u8]> = Vec::with_capacity(len);
2444 let mut field_row_iter = field_rows.iter().peekable();
2445 let null_row_bytes: &[u8] = &null_rows[field_idx].data;
2446
2447 for idx in 0..len {
2448 if let Some((next_idx, bytes)) = field_row_iter.peek()
2449 && *next_idx == idx
2450 {
2451 sparse_data.push(*bytes);
2452
2453 field_row_iter.next();
2454 continue;
2455 }
2456 sparse_data.push(null_row_bytes);
2457 }
2458
2459 let child_array =
2460 unsafe { converter.convert_raw(&mut sparse_data, validate_utf8) }?;
2461
2462 for (row_idx, child_row) in field_rows {
2464 let remaining_len = sparse_data[*row_idx].len();
2465 let consumed_length = 1 + child_row.len() - remaining_len;
2466 rows[*row_idx] = &rows[*row_idx][consumed_length..];
2467 }
2468
2469 child_arrays.push(child_array.into_iter().next().unwrap());
2470 }
2471 }
2472 }
2473
2474 if let Some(ref mut offsets_vec) = offsets {
2476 let mut count = vec![0i32; converters.len()];
2477 for type_id in &type_ids {
2478 let field_idx = type_id_to_field_idx[*type_id as usize];
2479 offsets_vec.push(count[field_idx]);
2480
2481 count[field_idx] += 1;
2482 }
2483 }
2484
2485 let type_ids_buffer = ScalarBuffer::from(type_ids);
2486 let offsets_buffer = offsets.map(ScalarBuffer::from);
2487
2488 let union_array = UnionArray::try_new(
2489 union_fields.clone(),
2490 type_ids_buffer,
2491 offsets_buffer,
2492 child_arrays,
2493 )?;
2494
2495 Arc::new(union_array)
2498 }
2499 };
2500 Ok(array)
2501}
2502
2503#[cfg(test)]
2504mod tests {
2505 use super::*;
2506 use arrow_array::builder::*;
2507 use arrow_array::types::*;
2508 use arrow_buffer::{Buffer, OffsetBuffer};
2509 use arrow_buffer::{NullBuffer, i256};
2510 use arrow_cast::display::{ArrayFormatter, FormatOptions};
2511 use arrow_ord::sort::{LexicographicalComparator, SortColumn};
2512 use rand::distr::uniform::SampleUniform;
2513 use rand::distr::{Distribution, StandardUniform};
2514 use rand::prelude::StdRng;
2515 use rand::{RngExt, SeedableRng};
2516 use std::fmt::Debug;
2517 use std::iter::Copied;
2518 use std::slice::Iter;
2519
2520 fn all_sort_options() -> [SortOptions; 4] {
2521 [
2522 SortOptions {
2523 descending: false,
2524 nulls_first: false,
2525 },
2526 SortOptions {
2527 descending: false,
2528 nulls_first: true,
2529 },
2530 SortOptions {
2531 descending: true,
2532 nulls_first: false,
2533 },
2534 SortOptions {
2535 descending: true,
2536 nulls_first: true,
2537 },
2538 ]
2539 }
2540
2541 #[test]
2542 fn test_fixed_width() {
2543 let cols = [
2544 Arc::new(Int16Array::from_iter([
2545 Some(1),
2546 Some(2),
2547 None,
2548 Some(-5),
2549 Some(2),
2550 Some(2),
2551 Some(0),
2552 ])) as ArrayRef,
2553 Arc::new(Float32Array::from_iter([
2554 Some(1.3),
2555 Some(2.5),
2556 None,
2557 Some(4.),
2558 Some(0.1),
2559 Some(-4.),
2560 Some(-0.),
2561 ])) as ArrayRef,
2562 ];
2563
2564 let converter = RowConverter::new(vec![
2565 SortField::new(DataType::Int16),
2566 SortField::new(DataType::Float32),
2567 ])
2568 .unwrap();
2569 let rows = converter.convert_columns(&cols).unwrap();
2570
2571 assert_eq!(rows.offsets, &[0, 8, 16, 24, 32, 40, 48, 56]);
2572 assert_eq!(
2573 rows.buffer,
2574 &[
2575 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 ]
2590 );
2591
2592 assert!(rows.row(3) < rows.row(6));
2593 assert!(rows.row(0) < rows.row(1));
2594 assert!(rows.row(3) < rows.row(0));
2595 assert!(rows.row(4) < rows.row(1));
2596 assert!(rows.row(5) < rows.row(4));
2597
2598 let back = converter.convert_rows(&rows).unwrap();
2599 for (expected, actual) in cols.iter().zip(&back) {
2600 assert_eq!(expected, actual);
2601 }
2602 }
2603
2604 fn test_roundtrip(sort_option: SortOptions, col: ArrayRef) {
2605 let converter = RowConverter::new(vec![SortField::new_with_options(
2606 col.data_type().clone(),
2607 sort_option,
2608 )])
2609 .unwrap();
2610 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2611 let back = converter.convert_rows(&rows).unwrap();
2612 assert_eq!(back.len(), 1);
2613 assert_eq!(&back[0], &col);
2614 back[0].to_data().validate_full().unwrap();
2615 }
2616
2617 #[test]
2618 fn test_zero_width_fixed_size_binary_roundtrip() {
2619 for sort_option in all_sort_options() {
2620 for with_null in [true, false] {
2623 let nulls = if with_null {
2624 Some(NullBuffer::from(vec![true, false, true, false, true]))
2625 } else {
2626 None
2627 };
2628 let col: ArrayRef = Arc::new(
2629 FixedSizeBinaryArray::try_new_with_len(0, Buffer::default(), nulls, 5).unwrap(),
2630 );
2631
2632 test_roundtrip(sort_option, col);
2633 }
2634 }
2635 }
2636
2637 #[test]
2638 fn test_zero_width_fixed_size_list_roundtrip() {
2639 for sort_option in all_sort_options() {
2640 for with_null in [true, false] {
2643 let nulls = if with_null {
2644 Some(NullBuffer::from(vec![true, false, true, false, true]))
2645 } else {
2646 None
2647 };
2648 let col: ArrayRef = Arc::new(
2649 FixedSizeListArray::try_new_with_length(
2650 Arc::new(Field::new("item", DataType::Boolean, false)),
2651 0,
2652 new_empty_array(&DataType::Boolean),
2653 nulls,
2654 5,
2655 )
2656 .unwrap(),
2657 );
2658
2659 test_roundtrip(sort_option, col);
2660 }
2661 }
2662 }
2663
2664 #[test]
2665 fn test_decimal32() {
2666 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal32(
2667 DECIMAL32_MAX_PRECISION,
2668 7,
2669 ))])
2670 .unwrap();
2671 let col = Arc::new(
2672 Decimal32Array::from_iter([
2673 None,
2674 Some(i32::MIN),
2675 Some(-13),
2676 Some(46_i32),
2677 Some(5456_i32),
2678 Some(i32::MAX),
2679 ])
2680 .with_precision_and_scale(9, 7)
2681 .unwrap(),
2682 ) as ArrayRef;
2683
2684 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2685 for i in 0..rows.num_rows() - 1 {
2686 assert!(rows.row(i) < rows.row(i + 1));
2687 }
2688
2689 let back = converter.convert_rows(&rows).unwrap();
2690 assert_eq!(back.len(), 1);
2691 assert_eq!(col.as_ref(), back[0].as_ref())
2692 }
2693
2694 #[test]
2695 fn test_decimal64() {
2696 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal64(
2697 DECIMAL64_MAX_PRECISION,
2698 7,
2699 ))])
2700 .unwrap();
2701 let col = Arc::new(
2702 Decimal64Array::from_iter([
2703 None,
2704 Some(i64::MIN),
2705 Some(-13),
2706 Some(46_i64),
2707 Some(5456_i64),
2708 Some(i64::MAX),
2709 ])
2710 .with_precision_and_scale(18, 7)
2711 .unwrap(),
2712 ) as ArrayRef;
2713
2714 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2715 for i in 0..rows.num_rows() - 1 {
2716 assert!(rows.row(i) < rows.row(i + 1));
2717 }
2718
2719 let back = converter.convert_rows(&rows).unwrap();
2720 assert_eq!(back.len(), 1);
2721 assert_eq!(col.as_ref(), back[0].as_ref())
2722 }
2723
2724 #[test]
2725 fn test_decimal128() {
2726 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal128(
2727 DECIMAL128_MAX_PRECISION,
2728 7,
2729 ))])
2730 .unwrap();
2731 let col = Arc::new(
2732 Decimal128Array::from_iter([
2733 None,
2734 Some(i128::MIN),
2735 Some(-13),
2736 Some(46_i128),
2737 Some(5456_i128),
2738 Some(i128::MAX),
2739 ])
2740 .with_precision_and_scale(38, 7)
2741 .unwrap(),
2742 ) as ArrayRef;
2743
2744 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2745 for i in 0..rows.num_rows() - 1 {
2746 assert!(rows.row(i) < rows.row(i + 1));
2747 }
2748
2749 let back = converter.convert_rows(&rows).unwrap();
2750 assert_eq!(back.len(), 1);
2751 assert_eq!(col.as_ref(), back[0].as_ref())
2752 }
2753
2754 #[test]
2755 fn test_decimal256() {
2756 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal256(
2757 DECIMAL256_MAX_PRECISION,
2758 7,
2759 ))])
2760 .unwrap();
2761 let col = Arc::new(
2762 Decimal256Array::from_iter([
2763 None,
2764 Some(i256::MIN),
2765 Some(i256::from_parts(0, -1)),
2766 Some(i256::from_parts(u128::MAX, -1)),
2767 Some(i256::from_parts(u128::MAX, 0)),
2768 Some(i256::from_parts(0, 46_i128)),
2769 Some(i256::from_parts(5, 46_i128)),
2770 Some(i256::MAX),
2771 ])
2772 .with_precision_and_scale(DECIMAL256_MAX_PRECISION, 7)
2773 .unwrap(),
2774 ) as ArrayRef;
2775
2776 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2777 for i in 0..rows.num_rows() - 1 {
2778 assert!(rows.row(i) < rows.row(i + 1));
2779 }
2780
2781 let back = converter.convert_rows(&rows).unwrap();
2782 assert_eq!(back.len(), 1);
2783 assert_eq!(col.as_ref(), back[0].as_ref())
2784 }
2785
2786 #[test]
2787 fn test_bool() {
2788 let converter = RowConverter::new(vec![SortField::new(DataType::Boolean)]).unwrap();
2789
2790 let col = Arc::new(BooleanArray::from_iter([None, Some(false), Some(true)])) as ArrayRef;
2791
2792 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2793 assert!(rows.row(2) > rows.row(1));
2794 assert!(rows.row(2) > rows.row(0));
2795 assert!(rows.row(1) > rows.row(0));
2796
2797 let cols = converter.convert_rows(&rows).unwrap();
2798 assert_eq!(&cols[0], &col);
2799
2800 let converter = RowConverter::new(vec![SortField::new_with_options(
2801 DataType::Boolean,
2802 SortOptions::default().desc().with_nulls_first(false),
2803 )])
2804 .unwrap();
2805
2806 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2807 assert!(rows.row(2) < rows.row(1));
2808 assert!(rows.row(2) < rows.row(0));
2809 assert!(rows.row(1) < rows.row(0));
2810 let cols = converter.convert_rows(&rows).unwrap();
2811 assert_eq!(&cols[0], &col);
2812 }
2813
2814 #[test]
2815 fn test_timezone() {
2816 let a =
2817 TimestampNanosecondArray::from(vec![1, 2, 3, 4, 5]).with_timezone("+01:00".to_string());
2818 let d = a.data_type().clone();
2819
2820 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
2821 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
2822 let back = converter.convert_rows(&rows).unwrap();
2823 assert_eq!(back.len(), 1);
2824 assert_eq!(back[0].data_type(), &d);
2825
2826 let mut a = PrimitiveDictionaryBuilder::<Int32Type, TimestampNanosecondType>::new();
2828 a.append(34).unwrap();
2829 a.append_null();
2830 a.append(345).unwrap();
2831
2832 let dict = a.finish();
2834 let values = TimestampNanosecondArray::from(dict.values().to_data());
2835 let dict_with_tz = dict.with_values(Arc::new(values.with_timezone("+02:00")));
2836 let v = DataType::Timestamp(TimeUnit::Nanosecond, Some("+02:00".into()));
2837 let d = DataType::Dictionary(Box::new(DataType::Int32), Box::new(v.clone()));
2838
2839 assert_eq!(dict_with_tz.data_type(), &d);
2840 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
2841 let rows = converter
2842 .convert_columns(&[Arc::new(dict_with_tz) as _])
2843 .unwrap();
2844 let back = converter.convert_rows(&rows).unwrap();
2845 assert_eq!(back.len(), 1);
2846 assert_eq!(back[0].data_type(), &v);
2847 }
2848
2849 #[test]
2850 fn test_null_encoding() {
2851 let col = Arc::new(NullArray::new(10));
2852 let converter = RowConverter::new(vec![SortField::new(DataType::Null)]).unwrap();
2853 let rows = converter.convert_columns(&[col]).unwrap();
2854 assert_eq!(rows.num_rows(), 10);
2855 assert_eq!(rows.row(1).data.len(), 2);
2857 }
2858
2859 #[test]
2860 fn test_variable_width() {
2861 let col = Arc::new(StringArray::from_iter([
2862 Some("hello"),
2863 Some("he"),
2864 None,
2865 Some("foo"),
2866 Some(""),
2867 ])) as ArrayRef;
2868
2869 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
2870 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2871
2872 assert!(rows.row(1) < rows.row(0));
2873 assert!(rows.row(2) < rows.row(4));
2874 assert!(rows.row(3) < rows.row(0));
2875 assert!(rows.row(3) < rows.row(1));
2876
2877 let cols = converter.convert_rows(&rows).unwrap();
2878 assert_eq!(&cols[0], &col);
2879
2880 let col = Arc::new(BinaryArray::from_iter([
2881 None,
2882 Some(vec![0_u8; 0]),
2883 Some(vec![0_u8; 6]),
2884 Some(vec![0_u8; variable::MINI_BLOCK_SIZE]),
2885 Some(vec![0_u8; variable::MINI_BLOCK_SIZE + 1]),
2886 Some(vec![0_u8; variable::BLOCK_SIZE]),
2887 Some(vec![0_u8; variable::BLOCK_SIZE + 1]),
2888 Some(vec![1_u8; 6]),
2889 Some(vec![1_u8; variable::MINI_BLOCK_SIZE]),
2890 Some(vec![1_u8; variable::MINI_BLOCK_SIZE + 1]),
2891 Some(vec![1_u8; variable::BLOCK_SIZE]),
2892 Some(vec![1_u8; variable::BLOCK_SIZE + 1]),
2893 Some(vec![0xFF_u8; 6]),
2894 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE]),
2895 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE + 1]),
2896 Some(vec![0xFF_u8; variable::BLOCK_SIZE]),
2897 Some(vec![0xFF_u8; variable::BLOCK_SIZE + 1]),
2898 ])) as ArrayRef;
2899
2900 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
2901 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2902
2903 for i in 0..rows.num_rows() {
2904 for j in i + 1..rows.num_rows() {
2905 assert!(
2906 rows.row(i) < rows.row(j),
2907 "{} < {} - {:?} < {:?}",
2908 i,
2909 j,
2910 rows.row(i),
2911 rows.row(j)
2912 );
2913 }
2914 }
2915
2916 let cols = converter.convert_rows(&rows).unwrap();
2917 assert_eq!(&cols[0], &col);
2918
2919 let converter = RowConverter::new(vec![SortField::new_with_options(
2920 DataType::Binary,
2921 SortOptions::default().desc().with_nulls_first(false),
2922 )])
2923 .unwrap();
2924 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2925
2926 for i in 0..rows.num_rows() {
2927 for j in i + 1..rows.num_rows() {
2928 assert!(
2929 rows.row(i) > rows.row(j),
2930 "{} > {} - {:?} > {:?}",
2931 i,
2932 j,
2933 rows.row(i),
2934 rows.row(j)
2935 );
2936 }
2937 }
2938
2939 let cols = converter.convert_rows(&rows).unwrap();
2940 assert_eq!(&cols[0], &col);
2941 }
2942
2943 fn dictionary_eq(a: &dyn Array, b: &dyn Array) {
2945 match b.data_type() {
2946 DataType::Dictionary(_, v) => {
2947 assert_eq!(a.data_type(), v.as_ref());
2948 let b = arrow_cast::cast(b, v).unwrap();
2949 assert_eq!(a, b.as_ref())
2950 }
2951 _ => assert_eq!(a, b),
2952 }
2953 }
2954
2955 #[test]
2956 fn test_string_dictionary() {
2957 let a = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2958 Some("foo"),
2959 Some("hello"),
2960 Some("he"),
2961 None,
2962 Some("hello"),
2963 Some(""),
2964 Some("hello"),
2965 Some("hello"),
2966 ])) as ArrayRef;
2967
2968 let field = SortField::new(a.data_type().clone());
2969 let converter = RowConverter::new(vec![field]).unwrap();
2970 let rows_a = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2971
2972 assert!(rows_a.row(3) < rows_a.row(5));
2973 assert!(rows_a.row(2) < rows_a.row(1));
2974 assert!(rows_a.row(0) < rows_a.row(1));
2975 assert!(rows_a.row(3) < rows_a.row(0));
2976
2977 assert_eq!(rows_a.row(1), rows_a.row(4));
2978 assert_eq!(rows_a.row(1), rows_a.row(6));
2979 assert_eq!(rows_a.row(1), rows_a.row(7));
2980
2981 let cols = converter.convert_rows(&rows_a).unwrap();
2982 dictionary_eq(&cols[0], &a);
2983
2984 let b = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2985 Some("hello"),
2986 None,
2987 Some("cupcakes"),
2988 ])) as ArrayRef;
2989
2990 let rows_b = converter.convert_columns(&[Arc::clone(&b)]).unwrap();
2991 assert_eq!(rows_a.row(1), rows_b.row(0));
2992 assert_eq!(rows_a.row(3), rows_b.row(1));
2993 assert!(rows_b.row(2) < rows_a.row(0));
2994
2995 let cols = converter.convert_rows(&rows_b).unwrap();
2996 dictionary_eq(&cols[0], &b);
2997
2998 let converter = RowConverter::new(vec![SortField::new_with_options(
2999 a.data_type().clone(),
3000 SortOptions::default().desc().with_nulls_first(false),
3001 )])
3002 .unwrap();
3003
3004 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
3005 assert!(rows_c.row(3) > rows_c.row(5));
3006 assert!(rows_c.row(2) > rows_c.row(1));
3007 assert!(rows_c.row(0) > rows_c.row(1));
3008 assert!(rows_c.row(3) > rows_c.row(0));
3009
3010 let cols = converter.convert_rows(&rows_c).unwrap();
3011 dictionary_eq(&cols[0], &a);
3012
3013 let converter = RowConverter::new(vec![SortField::new_with_options(
3014 a.data_type().clone(),
3015 SortOptions::default().desc().with_nulls_first(true),
3016 )])
3017 .unwrap();
3018
3019 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
3020 assert!(rows_c.row(3) < rows_c.row(5));
3021 assert!(rows_c.row(2) > rows_c.row(1));
3022 assert!(rows_c.row(0) > rows_c.row(1));
3023 assert!(rows_c.row(3) < rows_c.row(0));
3024
3025 let cols = converter.convert_rows(&rows_c).unwrap();
3026 dictionary_eq(&cols[0], &a);
3027 }
3028
3029 #[test]
3030 fn test_struct() {
3031 let a = Arc::new(Int32Array::from(vec![1, 1, 2, 2])) as ArrayRef;
3033 let a_f = Arc::new(Field::new("int", DataType::Int32, false));
3034 let u = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as ArrayRef;
3035 let u_f = Arc::new(Field::new("s", DataType::Utf8, false));
3036 let s1 = Arc::new(StructArray::from(vec![(a_f, a), (u_f, u)])) as ArrayRef;
3037
3038 let sort_fields = vec![SortField::new(s1.data_type().clone())];
3039 let converter = RowConverter::new(sort_fields).unwrap();
3040 let r1 = converter.convert_columns(&[Arc::clone(&s1)]).unwrap();
3041
3042 for (a, b) in r1.iter().zip(r1.iter().skip(1)) {
3043 assert!(a < b);
3044 }
3045
3046 let back = converter.convert_rows(&r1).unwrap();
3047 assert_eq!(back.len(), 1);
3048 assert_eq!(&back[0], &s1);
3049
3050 let data = s1
3052 .to_data()
3053 .into_builder()
3054 .null_bit_buffer(Some(Buffer::from_slice_ref([0b00001010])))
3055 .null_count(2)
3056 .build()
3057 .unwrap();
3058
3059 let s2 = Arc::new(StructArray::from(data)) as ArrayRef;
3060 let r2 = converter.convert_columns(&[Arc::clone(&s2)]).unwrap();
3061 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();
3067 assert_eq!(back.len(), 1);
3068 assert_eq!(&back[0], &s2);
3069
3070 back[0].to_data().validate_full().unwrap();
3071 }
3072
3073 #[test]
3074 fn test_dictionary_in_struct() {
3075 let builder = StringDictionaryBuilder::<Int32Type>::new();
3076 let mut struct_builder = StructBuilder::new(
3077 vec![Field::new_dictionary(
3078 "foo",
3079 DataType::Int32,
3080 DataType::Utf8,
3081 true,
3082 )],
3083 vec![Box::new(builder)],
3084 );
3085
3086 let dict_builder = struct_builder
3087 .field_builder::<StringDictionaryBuilder<Int32Type>>(0)
3088 .unwrap();
3089
3090 dict_builder.append_value("a");
3092 dict_builder.append_null();
3093 dict_builder.append_value("a");
3094 dict_builder.append_value("b");
3095
3096 for _ in 0..4 {
3097 struct_builder.append(true);
3098 }
3099
3100 let s = Arc::new(struct_builder.finish()) as ArrayRef;
3101 let sort_fields = vec![SortField::new(s.data_type().clone())];
3102 let converter = RowConverter::new(sort_fields).unwrap();
3103 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3104
3105 let back = converter.convert_rows(&r).unwrap();
3106 let [s2] = back.try_into().unwrap();
3107
3108 assert_ne!(&s.data_type(), &s2.data_type());
3111 s2.to_data().validate_full().unwrap();
3112
3113 let s1_struct = s.as_struct();
3117 let s1_0 = s1_struct.column(0);
3118 let s1_idx_0 = s1_0.as_dictionary::<Int32Type>();
3119 let keys = s1_idx_0.keys();
3120 let values = s1_idx_0.values().as_string::<i32>();
3121 let s2_struct = s2.as_struct();
3123 let s2_0 = s2_struct.column(0);
3124 let s2_idx_0 = s2_0.as_string::<i32>();
3125
3126 for i in 0..keys.len() {
3127 if keys.is_null(i) {
3128 assert!(s2_idx_0.is_null(i));
3129 } else {
3130 let dict_index = keys.value(i) as usize;
3131 assert_eq!(values.value(dict_index), s2_idx_0.value(i));
3132 }
3133 }
3134 }
3135
3136 #[test]
3137 fn test_dictionary_in_struct_empty() {
3138 let ty = DataType::Struct(
3139 vec![Field::new_dictionary(
3140 "foo",
3141 DataType::Int32,
3142 DataType::Int32,
3143 false,
3144 )]
3145 .into(),
3146 );
3147 let s = arrow_array::new_empty_array(&ty);
3148
3149 let sort_fields = vec![SortField::new(s.data_type().clone())];
3150 let converter = RowConverter::new(sort_fields).unwrap();
3151 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3152
3153 let back = converter.convert_rows(&r).unwrap();
3154 let [s2] = back.try_into().unwrap();
3155
3156 assert_ne!(&s.data_type(), &s2.data_type());
3159 s2.to_data().validate_full().unwrap();
3160 assert_eq!(s.len(), 0);
3161 assert_eq!(s2.len(), 0);
3162 }
3163
3164 #[test]
3165 fn test_list_of_string_dictionary() {
3166 let mut builder = ListBuilder::<StringDictionaryBuilder<Int32Type>>::default();
3167 builder.values().append("a").unwrap();
3169 builder.values().append("b").unwrap();
3170 builder.values().append("zero").unwrap();
3171 builder.values().append_null();
3172 builder.values().append("c").unwrap();
3173 builder.values().append("b").unwrap();
3174 builder.values().append("d").unwrap();
3175 builder.append(true);
3176 builder.append(false);
3178 builder.values().append("e").unwrap();
3180 builder.values().append("zero").unwrap();
3181 builder.values().append("a").unwrap();
3182 builder.append(true);
3183
3184 let a = Arc::new(builder.finish()) as ArrayRef;
3185 let data_type = a.data_type().clone();
3186
3187 let field = SortField::new(data_type.clone());
3188 let converter = RowConverter::new(vec![field]).unwrap();
3189 let rows = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
3190
3191 let back = converter.convert_rows(&rows).unwrap();
3192 assert_eq!(back.len(), 1);
3193 let [a2] = back.try_into().unwrap();
3194
3195 assert_ne!(&a.data_type(), &a2.data_type());
3198
3199 a2.to_data().validate_full().unwrap();
3200
3201 let a2_list = a2.as_list::<i32>();
3202 let a1_list = a.as_list::<i32>();
3203
3204 let a1_0 = a1_list.value(0);
3207 let a1_idx_0 = a1_0.as_dictionary::<Int32Type>();
3208 let keys = a1_idx_0.keys();
3209 let values = a1_idx_0.values().as_string::<i32>();
3210 let a2_0 = a2_list.value(0);
3211 let a2_idx_0 = a2_0.as_string::<i32>();
3212
3213 for i in 0..keys.len() {
3214 if keys.is_null(i) {
3215 assert!(a2_idx_0.is_null(i));
3216 } else {
3217 let dict_index = keys.value(i) as usize;
3218 assert_eq!(values.value(dict_index), a2_idx_0.value(i));
3219 }
3220 }
3221
3222 assert!(a1_list.is_null(1));
3224 assert!(a2_list.is_null(1));
3225
3226 let a1_2 = a1_list.value(2);
3228 let a1_idx_2 = a1_2.as_dictionary::<Int32Type>();
3229 let keys = a1_idx_2.keys();
3230 let values = a1_idx_2.values().as_string::<i32>();
3231 let a2_2 = a2_list.value(2);
3232 let a2_idx_2 = a2_2.as_string::<i32>();
3233
3234 for i in 0..keys.len() {
3235 if keys.is_null(i) {
3236 assert!(a2_idx_2.is_null(i));
3237 } else {
3238 let dict_index = keys.value(i) as usize;
3239 assert_eq!(values.value(dict_index), a2_idx_2.value(i));
3240 }
3241 }
3242 }
3243
3244 #[test]
3245 fn test_primitive_dictionary() {
3246 let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
3247 builder.append(2).unwrap();
3248 builder.append(3).unwrap();
3249 builder.append(0).unwrap();
3250 builder.append_null();
3251 builder.append(5).unwrap();
3252 builder.append(3).unwrap();
3253 builder.append(-1).unwrap();
3254
3255 let a = builder.finish();
3256 let data_type = a.data_type().clone();
3257 let columns = [Arc::new(a) as ArrayRef];
3258
3259 let field = SortField::new(data_type.clone());
3260 let converter = RowConverter::new(vec![field]).unwrap();
3261 let rows = converter.convert_columns(&columns).unwrap();
3262 assert!(rows.row(0) < rows.row(1));
3263 assert!(rows.row(2) < rows.row(0));
3264 assert!(rows.row(3) < rows.row(2));
3265 assert!(rows.row(6) < rows.row(2));
3266 assert!(rows.row(3) < rows.row(6));
3267
3268 let back = converter.convert_rows(&rows).unwrap();
3269 assert_eq!(back.len(), 1);
3270 back[0].to_data().validate_full().unwrap();
3271 }
3272
3273 #[test]
3274 fn test_dictionary_nulls() {
3275 let values = Int32Array::from_iter([Some(1), Some(-1), None, Some(4), None]).into_data();
3276 let keys =
3277 Int32Array::from_iter([Some(0), Some(0), Some(1), Some(2), Some(4), None]).into_data();
3278
3279 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3280 let data = keys
3281 .into_builder()
3282 .data_type(data_type.clone())
3283 .child_data(vec![values])
3284 .build()
3285 .unwrap();
3286
3287 let columns = [Arc::new(DictionaryArray::<Int32Type>::from(data)) as ArrayRef];
3288 let field = SortField::new(data_type.clone());
3289 let converter = RowConverter::new(vec![field]).unwrap();
3290 let rows = converter.convert_columns(&columns).unwrap();
3291
3292 assert_eq!(rows.row(0), rows.row(1));
3293 assert_eq!(rows.row(3), rows.row(4));
3294 assert_eq!(rows.row(4), rows.row(5));
3295 assert!(rows.row(3) < rows.row(0));
3296 }
3297
3298 #[test]
3299 fn test_from_binary_shared_buffer() {
3300 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3301 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3302 let rows = converter.convert_columns(&[array]).unwrap();
3303 let binary_rows = rows.try_into_binary().expect("known-small rows");
3304 let _binary_rows_shared_buffer = binary_rows.clone();
3305
3306 let parsed = converter.from_binary(binary_rows);
3307
3308 converter.convert_rows(parsed.iter()).unwrap();
3309 }
3310
3311 #[test]
3312 #[should_panic(expected = "Encountered non UTF-8 data")]
3313 fn test_invalid_utf8() {
3314 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3315 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3316 let rows = converter.convert_columns(&[array]).unwrap();
3317 let binary_row = rows.row(0);
3318
3319 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3320 let parser = converter.parser();
3321 let utf8_row = parser.parse(binary_row.as_ref());
3322
3323 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3324 }
3325
3326 #[test]
3327 #[should_panic(expected = "Encountered non UTF-8 data")]
3328 fn test_invalid_utf8_array() {
3329 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3330 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3331 let rows = converter.convert_columns(&[array]).unwrap();
3332 let binary_rows = rows.try_into_binary().expect("known-small rows");
3333
3334 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3335 let parsed = converter.from_binary(binary_rows);
3336
3337 converter.convert_rows(parsed.iter()).unwrap();
3338 }
3339
3340 #[test]
3341 #[should_panic(expected = "index out of bounds")]
3342 fn test_invalid_empty() {
3343 let binary_row: &[u8] = &[];
3344
3345 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3346 let parser = converter.parser();
3347 let utf8_row = parser.parse(binary_row.as_ref());
3348
3349 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3350 }
3351
3352 #[test]
3353 #[should_panic(expected = "index out of bounds")]
3354 fn test_invalid_empty_array() {
3355 let row: &[u8] = &[];
3356 let binary_rows = BinaryArray::from(vec![row]);
3357
3358 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3359 let parsed = converter.from_binary(binary_rows);
3360
3361 converter.convert_rows(parsed.iter()).unwrap();
3362 }
3363
3364 #[test]
3365 #[should_panic(expected = "index out of bounds")]
3366 fn test_invalid_truncated() {
3367 let binary_row: &[u8] = &[0x02];
3368
3369 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3370 let parser = converter.parser();
3371 let utf8_row = parser.parse(binary_row.as_ref());
3372
3373 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3374 }
3375
3376 #[test]
3377 #[should_panic(expected = "index out of bounds")]
3378 fn test_invalid_truncated_array() {
3379 let row: &[u8] = &[0x02];
3380 let binary_rows = BinaryArray::from(vec![row]);
3381
3382 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3383 let parsed = converter.from_binary(binary_rows);
3384
3385 converter.convert_rows(parsed.iter()).unwrap();
3386 }
3387
3388 #[test]
3389 #[should_panic(expected = "rows were not produced by this RowConverter")]
3390 fn test_different_converter() {
3391 let values = Arc::new(Int32Array::from_iter([Some(1), Some(-1)]));
3392 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3393 let rows = converter.convert_columns(&[values]).unwrap();
3394
3395 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3396 let _ = converter.convert_rows(&rows);
3397 }
3398
3399 fn test_single_list<O: OffsetSizeTrait>() {
3400 let mut builder = GenericListBuilder::<O, _>::new(Int32Builder::new());
3401 builder.values().append_value(32);
3402 builder.values().append_value(52);
3403 builder.values().append_value(32);
3404 builder.append(true);
3405 builder.values().append_value(32);
3406 builder.values().append_value(52);
3407 builder.values().append_value(12);
3408 builder.append(true);
3409 builder.values().append_value(32);
3410 builder.values().append_value(52);
3411 builder.append(true);
3412 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3415 builder.values().append_value(32);
3416 builder.values().append_null();
3417 builder.append(true);
3418 builder.append(true);
3419 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3422
3423 let list = Arc::new(builder.finish()) as ArrayRef;
3424 let d = list.data_type().clone();
3425
3426 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3427
3428 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3429 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();
3438 assert_eq!(back.len(), 1);
3439 back[0].to_data().validate_full().unwrap();
3440 assert_eq!(&back[0], &list);
3441
3442 let options = SortOptions::default().asc().with_nulls_first(false);
3443 let field = SortField::new_with_options(d.clone(), options);
3444 let converter = RowConverter::new(vec![field]).unwrap();
3445 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3446
3447 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();
3456 assert_eq!(back.len(), 1);
3457 back[0].to_data().validate_full().unwrap();
3458 assert_eq!(&back[0], &list);
3459
3460 let options = SortOptions::default().desc().with_nulls_first(false);
3461 let field = SortField::new_with_options(d.clone(), options);
3462 let converter = RowConverter::new(vec![field]).unwrap();
3463 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3464
3465 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();
3474 assert_eq!(back.len(), 1);
3475 back[0].to_data().validate_full().unwrap();
3476 assert_eq!(&back[0], &list);
3477
3478 let options = SortOptions::default().desc().with_nulls_first(true);
3479 let field = SortField::new_with_options(d, options);
3480 let converter = RowConverter::new(vec![field]).unwrap();
3481 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3482
3483 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();
3492 assert_eq!(back.len(), 1);
3493 back[0].to_data().validate_full().unwrap();
3494 assert_eq!(&back[0], &list);
3495
3496 let sliced_list = list.slice(1, 5);
3497 let rows_on_sliced_list = converter
3498 .convert_columns(&[Arc::clone(&sliced_list)])
3499 .unwrap();
3500
3501 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();
3508 assert_eq!(back.len(), 1);
3509 back[0].to_data().validate_full().unwrap();
3510 assert_eq!(&back[0], &sliced_list);
3511 }
3512
3513 fn test_nested_list<O: OffsetSizeTrait>() {
3514 let mut builder =
3515 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(Int32Builder::new()));
3516
3517 builder.values().values().append_value(1);
3518 builder.values().values().append_value(2);
3519 builder.values().append(true);
3520 builder.values().values().append_value(1);
3521 builder.values().values().append_null();
3522 builder.values().append(true);
3523 builder.append(true);
3524
3525 builder.values().values().append_value(1);
3526 builder.values().values().append_null();
3527 builder.values().append(true);
3528 builder.values().values().append_value(1);
3529 builder.values().values().append_null();
3530 builder.values().append(true);
3531 builder.append(true);
3532
3533 builder.values().values().append_value(1);
3534 builder.values().values().append_null();
3535 builder.values().append(true);
3536 builder.values().append(false);
3537 builder.append(true);
3538 builder.append(false);
3539
3540 builder.values().values().append_value(1);
3541 builder.values().values().append_value(2);
3542 builder.values().append(true);
3543 builder.append(true);
3544
3545 let list = Arc::new(builder.finish()) as ArrayRef;
3546 let d = list.data_type().clone();
3547
3548 let options = SortOptions::default().asc().with_nulls_first(true);
3556 let field = SortField::new_with_options(d.clone(), options);
3557 let converter = RowConverter::new(vec![field]).unwrap();
3558 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3559
3560 assert!(rows.row(0) > rows.row(1));
3561 assert!(rows.row(1) > rows.row(2));
3562 assert!(rows.row(2) > rows.row(3));
3563 assert!(rows.row(4) < rows.row(0));
3564 assert!(rows.row(4) > rows.row(1));
3565
3566 let back = converter.convert_rows(&rows).unwrap();
3567 assert_eq!(back.len(), 1);
3568 back[0].to_data().validate_full().unwrap();
3569 assert_eq!(&back[0], &list);
3570
3571 let options = SortOptions::default().desc().with_nulls_first(true);
3572 let field = SortField::new_with_options(d.clone(), options);
3573 let converter = RowConverter::new(vec![field]).unwrap();
3574 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3575
3576 assert!(rows.row(0) > rows.row(1));
3577 assert!(rows.row(1) > rows.row(2));
3578 assert!(rows.row(2) > rows.row(3));
3579 assert!(rows.row(4) > rows.row(0));
3580 assert!(rows.row(4) > rows.row(1));
3581
3582 let back = converter.convert_rows(&rows).unwrap();
3583 assert_eq!(back.len(), 1);
3584 back[0].to_data().validate_full().unwrap();
3585 assert_eq!(&back[0], &list);
3586
3587 let options = SortOptions::default().desc().with_nulls_first(false);
3588 let field = SortField::new_with_options(d, options);
3589 let converter = RowConverter::new(vec![field]).unwrap();
3590 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3591
3592 assert!(rows.row(0) < rows.row(1));
3593 assert!(rows.row(1) < rows.row(2));
3594 assert!(rows.row(2) < rows.row(3));
3595 assert!(rows.row(4) > rows.row(0));
3596 assert!(rows.row(4) < rows.row(1));
3597
3598 let back = converter.convert_rows(&rows).unwrap();
3599 assert_eq!(back.len(), 1);
3600 back[0].to_data().validate_full().unwrap();
3601 assert_eq!(&back[0], &list);
3602
3603 let sliced_list = list.slice(1, 3);
3604 let rows = converter
3605 .convert_columns(&[Arc::clone(&sliced_list)])
3606 .unwrap();
3607
3608 assert!(rows.row(0) < rows.row(1));
3609 assert!(rows.row(1) < rows.row(2));
3610
3611 let back = converter.convert_rows(&rows).unwrap();
3612 assert_eq!(back.len(), 1);
3613 back[0].to_data().validate_full().unwrap();
3614 assert_eq!(&back[0], &sliced_list);
3615 }
3616
3617 #[test]
3618 fn test_list() {
3619 test_single_list::<i32>();
3620 test_nested_list::<i32>();
3621 }
3622
3623 #[test]
3624 fn test_large_list() {
3625 test_single_list::<i64>();
3626 test_nested_list::<i64>();
3627 }
3628
3629 fn test_single_list_view<O: OffsetSizeTrait>() {
3630 let mut builder = GenericListViewBuilder::<O, _>::new(Int32Builder::new());
3631 builder.values().append_value(32);
3632 builder.values().append_value(52);
3633 builder.values().append_value(32);
3634 builder.append(true);
3635 builder.values().append_value(32);
3636 builder.values().append_value(52);
3637 builder.values().append_value(12);
3638 builder.append(true);
3639 builder.values().append_value(32);
3640 builder.values().append_value(52);
3641 builder.append(true);
3642 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3645 builder.values().append_value(32);
3646 builder.values().append_null();
3647 builder.append(true);
3648 builder.append(true);
3649 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3652
3653 let list = Arc::new(builder.finish()) as ArrayRef;
3654 let d = list.data_type().clone();
3655
3656 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3657
3658 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3659 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();
3668 assert_eq!(back.len(), 1);
3669 back[0].to_data().validate_full().unwrap();
3670
3671 let back_list_view = back[0]
3673 .as_any()
3674 .downcast_ref::<GenericListViewArray<O>>()
3675 .unwrap();
3676 let orig_list_view = list
3677 .as_any()
3678 .downcast_ref::<GenericListViewArray<O>>()
3679 .unwrap();
3680
3681 assert_eq!(back_list_view.len(), orig_list_view.len());
3682 for i in 0..back_list_view.len() {
3683 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3684 if back_list_view.is_valid(i) {
3685 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3686 }
3687 }
3688
3689 let options = SortOptions::default().asc().with_nulls_first(false);
3690 let field = SortField::new_with_options(d.clone(), options);
3691 let converter = RowConverter::new(vec![field]).unwrap();
3692 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3693
3694 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();
3703 assert_eq!(back.len(), 1);
3704 back[0].to_data().validate_full().unwrap();
3705
3706 let options = SortOptions::default().desc().with_nulls_first(false);
3707 let field = SortField::new_with_options(d.clone(), options);
3708 let converter = RowConverter::new(vec![field]).unwrap();
3709 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3710
3711 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();
3720 assert_eq!(back.len(), 1);
3721 back[0].to_data().validate_full().unwrap();
3722
3723 let options = SortOptions::default().desc().with_nulls_first(true);
3724 let field = SortField::new_with_options(d, options);
3725 let converter = RowConverter::new(vec![field]).unwrap();
3726 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3727
3728 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();
3737 assert_eq!(back.len(), 1);
3738 back[0].to_data().validate_full().unwrap();
3739
3740 let sliced_list = list.slice(1, 5);
3741 let rows_on_sliced_list = converter
3742 .convert_columns(&[Arc::clone(&sliced_list)])
3743 .unwrap();
3744
3745 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();
3752 assert_eq!(back.len(), 1);
3753 back[0].to_data().validate_full().unwrap();
3754 }
3755
3756 fn test_nested_list_view<O: OffsetSizeTrait>() {
3757 let mut builder = GenericListViewBuilder::<O, _>::new(GenericListViewBuilder::<O, _>::new(
3758 Int32Builder::new(),
3759 ));
3760
3761 builder.values().values().append_value(1);
3763 builder.values().values().append_value(2);
3764 builder.values().append(true);
3765 builder.values().values().append_value(1);
3766 builder.values().values().append_null();
3767 builder.values().append(true);
3768 builder.append(true);
3769
3770 builder.values().values().append_value(1);
3772 builder.values().values().append_null();
3773 builder.values().append(true);
3774 builder.values().values().append_value(1);
3775 builder.values().values().append_null();
3776 builder.values().append(true);
3777 builder.append(true);
3778
3779 builder.values().values().append_value(1);
3781 builder.values().values().append_null();
3782 builder.values().append(true);
3783 builder.values().append(false);
3784 builder.append(true);
3785
3786 builder.append(false);
3788
3789 builder.values().values().append_value(1);
3791 builder.values().values().append_value(2);
3792 builder.values().append(true);
3793 builder.append(true);
3794
3795 let list = Arc::new(builder.finish()) as ArrayRef;
3796 let d = list.data_type().clone();
3797
3798 let options = SortOptions::default().asc().with_nulls_first(true);
3806 let field = SortField::new_with_options(d.clone(), options);
3807 let converter = RowConverter::new(vec![field]).unwrap();
3808 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3809
3810 assert!(rows.row(0) > rows.row(1));
3811 assert!(rows.row(1) > rows.row(2));
3812 assert!(rows.row(2) > rows.row(3));
3813 assert!(rows.row(4) < rows.row(0));
3814 assert!(rows.row(4) > rows.row(1));
3815
3816 let back = converter.convert_rows(&rows).unwrap();
3817 assert_eq!(back.len(), 1);
3818 back[0].to_data().validate_full().unwrap();
3819
3820 let back_list_view = back[0]
3822 .as_any()
3823 .downcast_ref::<GenericListViewArray<O>>()
3824 .unwrap();
3825 let orig_list_view = list
3826 .as_any()
3827 .downcast_ref::<GenericListViewArray<O>>()
3828 .unwrap();
3829
3830 assert_eq!(back_list_view.len(), orig_list_view.len());
3831 for i in 0..back_list_view.len() {
3832 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3833 if back_list_view.is_valid(i) {
3834 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3835 }
3836 }
3837
3838 let options = SortOptions::default().desc().with_nulls_first(true);
3839 let field = SortField::new_with_options(d.clone(), options);
3840 let converter = RowConverter::new(vec![field]).unwrap();
3841 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3842
3843 assert!(rows.row(0) > rows.row(1));
3844 assert!(rows.row(1) > rows.row(2));
3845 assert!(rows.row(2) > rows.row(3));
3846 assert!(rows.row(4) > rows.row(0));
3847 assert!(rows.row(4) > rows.row(1));
3848
3849 let back = converter.convert_rows(&rows).unwrap();
3850 assert_eq!(back.len(), 1);
3851 back[0].to_data().validate_full().unwrap();
3852
3853 let back_list_view = back[0]
3855 .as_any()
3856 .downcast_ref::<GenericListViewArray<O>>()
3857 .unwrap();
3858
3859 assert_eq!(back_list_view.len(), orig_list_view.len());
3860 for i in 0..back_list_view.len() {
3861 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3862 if back_list_view.is_valid(i) {
3863 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3864 }
3865 }
3866
3867 let options = SortOptions::default().desc().with_nulls_first(false);
3868 let field = SortField::new_with_options(d.clone(), options);
3869 let converter = RowConverter::new(vec![field]).unwrap();
3870 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3871
3872 assert!(rows.row(0) < rows.row(1));
3873 assert!(rows.row(1) < rows.row(2));
3874 assert!(rows.row(2) < rows.row(3));
3875 assert!(rows.row(4) > rows.row(0));
3876 assert!(rows.row(4) < rows.row(1));
3877
3878 let back = converter.convert_rows(&rows).unwrap();
3879 assert_eq!(back.len(), 1);
3880 back[0].to_data().validate_full().unwrap();
3881
3882 let back_list_view = back[0]
3884 .as_any()
3885 .downcast_ref::<GenericListViewArray<O>>()
3886 .unwrap();
3887
3888 assert_eq!(back_list_view.len(), orig_list_view.len());
3889 for i in 0..back_list_view.len() {
3890 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3891 if back_list_view.is_valid(i) {
3892 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3893 }
3894 }
3895
3896 let sliced_list = list.slice(1, 3);
3897 let rows = converter
3898 .convert_columns(&[Arc::clone(&sliced_list)])
3899 .unwrap();
3900
3901 assert!(rows.row(0) < rows.row(1));
3902 assert!(rows.row(1) < rows.row(2));
3903
3904 let back = converter.convert_rows(&rows).unwrap();
3905 assert_eq!(back.len(), 1);
3906 back[0].to_data().validate_full().unwrap();
3907 }
3908
3909 #[test]
3910 fn test_list_view() {
3911 test_single_list_view::<i32>();
3912 test_nested_list_view::<i32>();
3913 }
3914
3915 #[test]
3916 fn test_large_list_view() {
3917 test_single_list_view::<i64>();
3918 test_nested_list_view::<i64>();
3919 }
3920
3921 fn test_list_view_with_shared_values<O: OffsetSizeTrait>() {
3922 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
3924 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
3925
3926 let offsets = ScalarBuffer::<O>::from(vec![
3934 O::from_usize(0).unwrap(),
3935 O::from_usize(0).unwrap(),
3936 O::from_usize(5).unwrap(),
3937 O::from_usize(2).unwrap(),
3938 O::from_usize(1).unwrap(),
3939 O::from_usize(2).unwrap(),
3940 ]);
3941 let sizes = ScalarBuffer::<O>::from(vec![
3942 O::from_usize(3).unwrap(),
3943 O::from_usize(3).unwrap(),
3944 O::from_usize(2).unwrap(),
3945 O::from_usize(2).unwrap(),
3946 O::from_usize(4).unwrap(),
3947 O::from_usize(1).unwrap(),
3948 ]);
3949
3950 let list_view: GenericListViewArray<O> =
3951 GenericListViewArray::try_new(field, offsets, sizes, Arc::new(values), None).unwrap();
3952
3953 let d = list_view.data_type().clone();
3954 let list = Arc::new(list_view) as ArrayRef;
3955
3956 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3957 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3958
3959 assert_eq!(rows.row(0), rows.row(1));
3961
3962 assert!(rows.row(0) < rows.row(2));
3964
3965 assert!(rows.row(3) > rows.row(0));
3967
3968 assert!(rows.row(4) > rows.row(0));
3970
3971 assert!(rows.row(5) < rows.row(3));
3973
3974 assert!(rows.row(5) > rows.row(4));
3976
3977 let back = converter.convert_rows(&rows).unwrap();
3979 assert_eq!(back.len(), 1);
3980 back[0].to_data().validate_full().unwrap();
3981
3982 let back_list_view = back[0]
3984 .as_any()
3985 .downcast_ref::<GenericListViewArray<O>>()
3986 .unwrap();
3987 let orig_list_view = list
3988 .as_any()
3989 .downcast_ref::<GenericListViewArray<O>>()
3990 .unwrap();
3991
3992 assert_eq!(back_list_view.len(), orig_list_view.len());
3993 for i in 0..back_list_view.len() {
3994 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3995 if back_list_view.is_valid(i) {
3996 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3997 }
3998 }
3999
4000 let options = SortOptions::default().desc();
4002 let field = SortField::new_with_options(d, options);
4003 let converter = RowConverter::new(vec![field]).unwrap();
4004 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4005
4006 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();
4012 assert_eq!(back.len(), 1);
4013 back[0].to_data().validate_full().unwrap();
4014 }
4015
4016 #[test]
4017 fn test_list_view_shared_values() {
4018 test_list_view_with_shared_values::<i32>();
4019 }
4020
4021 #[test]
4022 fn test_large_list_view_shared_values() {
4023 test_list_view_with_shared_values::<i64>();
4024 }
4025
4026 #[test]
4027 fn test_fixed_size_list() {
4028 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
4029 builder.values().append_value(32);
4030 builder.values().append_value(52);
4031 builder.values().append_value(32);
4032 builder.append(true);
4033 builder.values().append_value(32);
4034 builder.values().append_value(52);
4035 builder.values().append_value(12);
4036 builder.append(true);
4037 builder.values().append_value(32);
4038 builder.values().append_value(52);
4039 builder.values().append_null();
4040 builder.append(true);
4041 builder.values().append_value(32); builder.values().append_value(52); builder.values().append_value(13); builder.append(false);
4045 builder.values().append_value(32);
4046 builder.values().append_null();
4047 builder.values().append_null();
4048 builder.append(true);
4049 builder.values().append_null();
4050 builder.values().append_null();
4051 builder.values().append_null();
4052 builder.append(true);
4053 builder.values().append_value(17); builder.values().append_null(); builder.values().append_value(77); builder.append(false);
4057
4058 let list = Arc::new(builder.finish()) as ArrayRef;
4059 let d = list.data_type().clone();
4060
4061 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4063
4064 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4065 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();
4074 assert_eq!(back.len(), 1);
4075 back[0].to_data().validate_full().unwrap();
4076 assert_eq!(&back[0], &list);
4077
4078 let options = SortOptions::default().asc().with_nulls_first(false);
4080 let field = SortField::new_with_options(d.clone(), options);
4081 let converter = RowConverter::new(vec![field]).unwrap();
4082 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4083 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();
4092 assert_eq!(back.len(), 1);
4093 back[0].to_data().validate_full().unwrap();
4094 assert_eq!(&back[0], &list);
4095
4096 let options = SortOptions::default().desc().with_nulls_first(false);
4098 let field = SortField::new_with_options(d.clone(), options);
4099 let converter = RowConverter::new(vec![field]).unwrap();
4100 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4101 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();
4110 assert_eq!(back.len(), 1);
4111 back[0].to_data().validate_full().unwrap();
4112 assert_eq!(&back[0], &list);
4113
4114 let options = SortOptions::default().desc().with_nulls_first(true);
4116 let field = SortField::new_with_options(d, options);
4117 let converter = RowConverter::new(vec![field]).unwrap();
4118 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4119
4120 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();
4129 assert_eq!(back.len(), 1);
4130 back[0].to_data().validate_full().unwrap();
4131 assert_eq!(&back[0], &list);
4132
4133 let sliced_list = list.slice(1, 5);
4134 let rows_on_sliced_list = converter
4135 .convert_columns(&[Arc::clone(&sliced_list)])
4136 .unwrap();
4137
4138 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();
4144 assert_eq!(back.len(), 1);
4145 back[0].to_data().validate_full().unwrap();
4146 assert_eq!(&back[0], &sliced_list);
4147 }
4148
4149 #[test]
4150 fn test_two_fixed_size_lists() {
4151 let mut first = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4152 first.values().append_value(100);
4154 first.append(true);
4155 first.values().append_value(101);
4157 first.append(true);
4158 first.values().append_value(102);
4160 first.append(true);
4161 first.values().append_null();
4163 first.append(true);
4164 first.values().append_null(); first.append(false);
4167 let first = Arc::new(first.finish()) as ArrayRef;
4168 let first_type = first.data_type().clone();
4169
4170 let mut second = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4171 second.values().append_value(200);
4173 second.append(true);
4174 second.values().append_value(201);
4176 second.append(true);
4177 second.values().append_value(202);
4179 second.append(true);
4180 second.values().append_null();
4182 second.append(true);
4183 second.values().append_null(); second.append(false);
4186 let second = Arc::new(second.finish()) as ArrayRef;
4187 let second_type = second.data_type().clone();
4188
4189 let converter = RowConverter::new(vec![
4190 SortField::new(first_type.clone()),
4191 SortField::new(second_type.clone()),
4192 ])
4193 .unwrap();
4194
4195 let rows = converter
4196 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4197 .unwrap();
4198
4199 let back = converter.convert_rows(&rows).unwrap();
4200 assert_eq!(back.len(), 2);
4201 back[0].to_data().validate_full().unwrap();
4202 assert_eq!(&back[0], &first);
4203 back[1].to_data().validate_full().unwrap();
4204 assert_eq!(&back[1], &second);
4205 }
4206
4207 #[test]
4208 fn test_fixed_size_list_with_variable_width_content() {
4209 let mut first = FixedSizeListBuilder::new(
4210 StructBuilder::from_fields(
4211 vec![
4212 Field::new(
4213 "timestamp",
4214 DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))),
4215 false,
4216 ),
4217 Field::new("offset_minutes", DataType::Int16, false),
4218 Field::new("time_zone", DataType::Utf8, false),
4219 ],
4220 1,
4221 ),
4222 1,
4223 );
4224 first
4226 .values()
4227 .field_builder::<TimestampMicrosecondBuilder>(0)
4228 .unwrap()
4229 .append_null();
4230 first
4231 .values()
4232 .field_builder::<Int16Builder>(1)
4233 .unwrap()
4234 .append_null();
4235 first
4236 .values()
4237 .field_builder::<StringBuilder>(2)
4238 .unwrap()
4239 .append_null();
4240 first.values().append(false);
4241 first.append(false);
4242 first
4244 .values()
4245 .field_builder::<TimestampMicrosecondBuilder>(0)
4246 .unwrap()
4247 .append_null();
4248 first
4249 .values()
4250 .field_builder::<Int16Builder>(1)
4251 .unwrap()
4252 .append_null();
4253 first
4254 .values()
4255 .field_builder::<StringBuilder>(2)
4256 .unwrap()
4257 .append_null();
4258 first.values().append(false);
4259 first.append(true);
4260 first
4262 .values()
4263 .field_builder::<TimestampMicrosecondBuilder>(0)
4264 .unwrap()
4265 .append_value(0);
4266 first
4267 .values()
4268 .field_builder::<Int16Builder>(1)
4269 .unwrap()
4270 .append_value(0);
4271 first
4272 .values()
4273 .field_builder::<StringBuilder>(2)
4274 .unwrap()
4275 .append_value("UTC");
4276 first.values().append(true);
4277 first.append(true);
4278 first
4280 .values()
4281 .field_builder::<TimestampMicrosecondBuilder>(0)
4282 .unwrap()
4283 .append_value(1126351800123456);
4284 first
4285 .values()
4286 .field_builder::<Int16Builder>(1)
4287 .unwrap()
4288 .append_value(120);
4289 first
4290 .values()
4291 .field_builder::<StringBuilder>(2)
4292 .unwrap()
4293 .append_value("Europe/Warsaw");
4294 first.values().append(true);
4295 first.append(true);
4296 let first = Arc::new(first.finish()) as ArrayRef;
4297 let first_type = first.data_type().clone();
4298
4299 let mut second = StringBuilder::new();
4300 second.append_value("somewhere near");
4301 second.append_null();
4302 second.append_value("Greenwich");
4303 second.append_value("Warsaw");
4304 let second = Arc::new(second.finish()) as ArrayRef;
4305 let second_type = second.data_type().clone();
4306
4307 let converter = RowConverter::new(vec![
4308 SortField::new(first_type.clone()),
4309 SortField::new(second_type.clone()),
4310 ])
4311 .unwrap();
4312
4313 let rows = converter
4314 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4315 .unwrap();
4316
4317 let back = converter.convert_rows(&rows).unwrap();
4318 assert_eq!(back.len(), 2);
4319 back[0].to_data().validate_full().unwrap();
4320 assert_eq!(&back[0], &first);
4321 back[1].to_data().validate_full().unwrap();
4322 assert_eq!(&back[1], &second);
4323 }
4324
4325 #[test]
4326 fn test_single_map() {
4327 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4328 builder.keys().append_value("hello");
4330 builder.values().append_value(1);
4331 builder.keys().append_value("world");
4332 builder.values().append_value(2);
4333 builder.append(true).unwrap();
4334
4335 builder.keys().append_value("foo");
4337 builder.values().append_value(3);
4338 builder.append(true).unwrap();
4339
4340 builder.append(true).unwrap();
4342
4343 builder.keys().append_value("masked_key");
4345 builder.values().append_value(999);
4346 builder.append(false).unwrap();
4347
4348 builder.append(false).unwrap();
4350
4351 builder.keys().append_value("bar");
4353 builder.values().append_null();
4354 builder.append(true).unwrap();
4355
4356 builder.keys().append_value("other_masked");
4358 builder.values().append_value(0);
4359 builder.append(false).unwrap();
4360
4361 builder.keys().append_value("a");
4363 builder.values().append_value(10);
4364 builder.keys().append_value("b");
4365 builder.values().append_value(20);
4366 builder.keys().append_value("c");
4367 builder.values().append_value(30);
4368 builder.append(true).unwrap();
4369
4370 let map = Arc::new(builder.finish()) as ArrayRef;
4371 let d = map.data_type().clone();
4372
4373 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4374
4375 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4376
4377 assert_eq!(rows.row(3), rows.row(4));
4379 assert_eq!(rows.row(4), rows.row(6));
4380
4381 let back = converter.convert_rows(&rows).unwrap();
4382 assert_eq!(back.len(), 1);
4383 back[0].to_data().validate_full().unwrap();
4384 assert_eq!(&back[0], &map);
4385
4386 let sliced_map = map.slice(1, map.len() - 2);
4387 let rows_on_sliced = converter
4388 .convert_columns(&[Arc::clone(&sliced_map)])
4389 .unwrap();
4390
4391 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4392 assert_eq!(back.len(), 1);
4393 back[0].to_data().validate_full().unwrap();
4394 assert_eq!(&back[0], &sliced_map);
4395 }
4396
4397 #[test]
4398 fn two_maps_with_different_keys_order_should_sort_by_entry_order() {
4399 let map_1: ArrayRef =
4400 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4401 vec![Some(vec![("hello", Some(1)), ("world", Some(2))])],
4402 false,
4403 ));
4404 let map_2: ArrayRef =
4406 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4407 vec![Some(vec![("world", Some(2)), ("hello", Some(1))])],
4408 false,
4409 ));
4410
4411 let converter = RowConverter::new(vec![SortField::new(map_1.data_type().clone())]).unwrap();
4412
4413 let map_1_rows = converter.convert_columns(&[Arc::clone(&map_1)]).unwrap();
4414 let map_2_rows = converter.convert_columns(&[Arc::clone(&map_2)]).unwrap();
4415
4416 assert_ne!(map_1_rows.row(0), map_2_rows.row(0));
4417 assert!(map_1_rows.row(0) < map_2_rows.row(0));
4418
4419 let back_1 = converter.convert_rows(&map_1_rows).unwrap();
4420 let back_2 = converter.convert_rows(&map_2_rows).unwrap();
4421 assert_eq!(&back_1[0], &map_1);
4422 assert_eq!(&back_2[0], &map_2);
4423 }
4424
4425 #[test]
4426 fn test_nested_map() {
4427 let mut builder = MapBuilder::new(
4429 None,
4430 StringBuilder::new(),
4431 MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()),
4432 );
4433
4434 builder.keys().append_value("outer1");
4436 builder.values().keys().append_value("inner_a");
4437 builder.values().values().append_value(1);
4438 builder.values().keys().append_value("inner_b");
4439 builder.values().values().append_value(2);
4440 builder.values().append(true).unwrap();
4441 builder.keys().append_value("outer2");
4442 builder.values().keys().append_value("inner_c");
4443 builder.values().values().append_value(3);
4444 builder.values().append(true).unwrap();
4445 builder.append(true).unwrap();
4446
4447 builder.keys().append_value("x");
4449 builder.values().append(true).unwrap();
4450 builder.append(true).unwrap();
4451
4452 builder.keys().append_value("y");
4454 builder.values().keys().append_value("masked"); builder.values().values().append_value(0); builder.values().append(false).unwrap();
4457 builder.append(true).unwrap();
4458
4459 builder.keys().append_value("y");
4461 builder.values().append(false).unwrap(); builder.append(true).unwrap();
4463
4464 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();
4470
4471 builder.keys().append_value("masked_outer"); builder.values().append(false).unwrap(); builder.append(false).unwrap();
4475
4476 builder.append(false).unwrap(); builder.append(true).unwrap();
4481
4482 let map = Arc::new(builder.finish()) as ArrayRef;
4483 let d = map.data_type().clone();
4484
4485 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4486
4487 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4488
4489 let back = converter.convert_rows(&rows).unwrap();
4490 assert_eq!(back.len(), 1);
4491 back[0].to_data().validate_full().unwrap();
4492 assert_eq!(&back[0], &map);
4493
4494 let sliced_map = map.slice(1, 3);
4495 let rows_on_sliced = converter
4496 .convert_columns(&[Arc::clone(&sliced_map)])
4497 .unwrap();
4498
4499 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4500 assert_eq!(back.len(), 1);
4501 back[0].to_data().validate_full().unwrap();
4502 assert_eq!(&back[0], &sliced_map);
4503 }
4504
4505 #[test]
4506 fn test_single_map_with_non_nullable_values() {
4507 let value_field = Arc::new(Field::new(
4509 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4510 DataType::Int32,
4511 false,
4512 ));
4513 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4514 .with_values_field(value_field);
4515 builder.keys().append_value("a");
4517 builder.values().append_value(1);
4518 builder.keys().append_value("b");
4519 builder.values().append_value(2);
4520 builder.append(true).unwrap();
4521 builder.append(false).unwrap();
4523 builder.keys().append_value("c");
4525 builder.values().append_value(3);
4526 builder.append(true).unwrap();
4527 builder.append(true).unwrap();
4529 builder.keys().append_value("masked"); builder.values().append_value(0); builder.append(false).unwrap();
4533
4534 let map = Arc::new(builder.finish()) as ArrayRef;
4535 let d = map.data_type().clone();
4536
4537 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4538
4539 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4540
4541 let back = converter.convert_rows(&rows).unwrap();
4542 assert_eq!(back.len(), 1);
4543 back[0].to_data().validate_full().unwrap();
4544 assert_eq!(&back[0], &map);
4545 }
4546
4547 #[test]
4548 fn test_single_map_with_non_nullable_map_but_with_nullable_values() {
4549 let value_field = Arc::new(Field::new(
4551 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4552 DataType::Int32,
4553 true,
4554 ));
4555 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4556 .with_values_field(value_field);
4557
4558 builder.keys().append_value("a");
4560 builder.values().append_value(1);
4561 builder.keys().append_value("b");
4562 builder.values().append_null();
4563 builder.append(true).unwrap();
4564 builder.keys().append_value("c");
4566 builder.values().append_null();
4567 builder.keys().append_value("d");
4568 builder.values().append_null();
4569 builder.append(true).unwrap();
4570 builder.append(true).unwrap();
4572 builder.keys().append_value("e");
4574 builder.values().append_value(5);
4575 builder.append(true).unwrap();
4576
4577 let map = Arc::new(builder.finish()) as ArrayRef;
4578 let d = map.data_type().clone();
4579
4580 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4581
4582 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4583
4584 let back = converter.convert_rows(&rows).unwrap();
4585 assert_eq!(back.len(), 1);
4586 back[0].to_data().validate_full().unwrap();
4587 assert_eq!(&back[0], &map);
4588 }
4589
4590 #[test]
4591 fn test_map_all_nulls() {
4592 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4593 builder.keys().append_value("m1"); builder.values().append_value(1); builder.append(false).unwrap();
4597 builder.keys().append_value("m2"); builder.values().append_value(2); builder.append(false).unwrap();
4600
4601 builder.append(false).unwrap(); builder.keys().append_value("m3"); builder.values().append_value(3); builder.append(false).unwrap();
4606
4607 let map = Arc::new(builder.finish()) as ArrayRef;
4608 let d = map.data_type().clone();
4609
4610 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4611
4612 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4613
4614 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4616
4617 let back = converter.convert_rows(&rows).unwrap();
4618 assert_eq!(back.len(), 1);
4619 back[0].to_data().validate_full().unwrap();
4620 assert_eq!(&back[0], &map);
4621 }
4622
4623 #[test]
4624 fn test_map_all_empty() {
4625 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4626 builder.append(true).unwrap();
4628 builder.append(true).unwrap();
4629 builder.append(true).unwrap();
4630
4631 let map = Arc::new(builder.finish()) as ArrayRef;
4632 let d = map.data_type().clone();
4633
4634 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4635
4636 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4637
4638 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4640
4641 let back = converter.convert_rows(&rows).unwrap();
4642 assert_eq!(back.len(), 1);
4643 back[0].to_data().validate_full().unwrap();
4644 assert_eq!(&back[0], &map);
4645 }
4646
4647 #[test]
4648 fn test_map_empty_array() {
4649 let builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4651 let map = Arc::new(builder.finish_cloned()) as ArrayRef;
4652 let d = map.data_type().clone();
4653
4654 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4655
4656 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4657
4658 let back = converter.convert_rows(&rows).unwrap();
4659 assert_eq!(back.len(), 1);
4660 back[0].to_data().validate_full().unwrap();
4661 assert_eq!(&back[0], &map);
4662 }
4663
4664 fn generate_primitive_array<K>(
4665 rng: &mut StdRng,
4666 len: usize,
4667 valid_percent: f64,
4668 ) -> PrimitiveArray<K>
4669 where
4670 K: ArrowPrimitiveType,
4671 StandardUniform: Distribution<K::Native>,
4672 {
4673 (0..len)
4674 .map(|_| rng.random_bool(valid_percent).then(|| rng.random()))
4675 .collect()
4676 }
4677
4678 fn generate_all_unique_primitive_array<K>(rng: &mut StdRng, len: usize) -> PrimitiveArray<K>
4679 where
4680 K: ArrowPrimitiveType,
4681 K::Native: Hash + Eq,
4682 StandardUniform: Distribution<K::Native>,
4683 {
4684 let possible_number_of_values = 2i32.saturating_pow(size_of::<K::Native>() as u32 * 8);
4685 assert!(
4686 len <= possible_number_of_values as usize,
4687 "len {len} is larger than the number of possible values {possible_number_of_values}"
4688 );
4689
4690 let mut seen = std::collections::HashSet::new();
4691 (0..len)
4692 .map(|_| {
4693 let mut value;
4694 loop {
4695 value = rng.random();
4696
4697 if seen.insert(value) {
4698 break;
4699 }
4700 }
4701
4702 Some(value)
4703 })
4704 .collect()
4705 }
4706
4707 fn generate_boolean_array(rng: &mut StdRng, len: usize, valid_percent: f64) -> BooleanArray {
4708 (0..len)
4709 .map(|_| rng.random_bool(valid_percent).then(|| rng.random_bool(0.5)))
4710 .collect()
4711 }
4712
4713 fn generate_strings<O: OffsetSizeTrait>(
4714 rng: &mut StdRng,
4715 len: usize,
4716 valid_percent: f64,
4717 ) -> GenericStringArray<O> {
4718 (0..len)
4719 .map(|_| {
4720 rng.random_bool(valid_percent).then(|| {
4721 let len = rng.random_range(0..100);
4722 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4723 String::from_utf8(bytes).unwrap()
4724 })
4725 })
4726 .collect()
4727 }
4728
4729 fn generate_string_view(rng: &mut StdRng, len: usize, valid_percent: f64) -> StringViewArray {
4730 (0..len)
4731 .map(|_| {
4732 rng.random_bool(valid_percent).then(|| {
4733 let len = rng.random_range(0..100);
4734 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4735 String::from_utf8(bytes).unwrap()
4736 })
4737 })
4738 .collect()
4739 }
4740
4741 fn generate_byte_view(rng: &mut StdRng, len: usize, valid_percent: f64) -> BinaryViewArray {
4742 (0..len)
4743 .map(|_| {
4744 rng.random_bool(valid_percent).then(|| {
4745 let len = rng.random_range(0..100);
4746 let bytes: Vec<_> = (0..len).map(|_| rng.random_range(0..128)).collect();
4747 bytes
4748 })
4749 })
4750 .collect()
4751 }
4752
4753 fn generate_fixed_stringview_column(len: usize) -> StringViewArray {
4754 let edge_cases = vec![
4755 Some("bar".to_string()),
4756 Some("bar\0".to_string()),
4757 Some("LongerThan12Bytes".to_string()),
4758 Some("LongerThan12Bytez".to_string()),
4759 Some("LongerThan12Bytes\0".to_string()),
4760 Some("LongerThan12Byt".to_string()),
4761 Some("backend one".to_string()),
4762 Some("backend two".to_string()),
4763 Some("a".repeat(257)),
4764 Some("a".repeat(300)),
4765 ];
4766
4767 let mut values = Vec::with_capacity(len);
4769 for i in 0..len {
4770 values.push(
4771 edge_cases
4772 .get(i % edge_cases.len())
4773 .cloned()
4774 .unwrap_or(None),
4775 );
4776 }
4777
4778 StringViewArray::from(values)
4779 }
4780
4781 fn generate_dictionary<K>(
4782 rng: &mut StdRng,
4783 values: ArrayRef,
4784 len: usize,
4785 valid_percent: f64,
4786 ) -> DictionaryArray<K>
4787 where
4788 K: ArrowDictionaryKeyType,
4789 K::Native: SampleUniform,
4790 {
4791 let min_key = K::Native::from_usize(0).unwrap();
4792 let max_key = K::Native::from_usize(values.len()).unwrap();
4793 let keys: PrimitiveArray<K> = (0..len)
4794 .map(|_| {
4795 rng.random_bool(valid_percent)
4796 .then(|| rng.random_range(min_key..max_key))
4797 })
4798 .collect();
4799
4800 let data_type =
4801 DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
4802
4803 let data = keys
4804 .into_data()
4805 .into_builder()
4806 .data_type(data_type)
4807 .add_child_data(values.to_data())
4808 .build()
4809 .unwrap();
4810
4811 DictionaryArray::from(data)
4812 }
4813
4814 fn generate_fixed_size_binary(
4815 rng: &mut StdRng,
4816 len: usize,
4817 valid_percent: f64,
4818 ) -> FixedSizeBinaryArray {
4819 let width = rng.random_range(0..20);
4820 let mut builder = FixedSizeBinaryBuilder::new(width);
4821
4822 let mut b = vec![0; width as usize];
4823 for _ in 0..len {
4824 match rng.random_bool(valid_percent) {
4825 true => {
4826 b.iter_mut().for_each(|x| *x = rng.random());
4827 builder.append_value(&b).unwrap();
4828 }
4829 false => builder.append_null(),
4830 }
4831 }
4832
4833 builder.finish()
4834 }
4835
4836 fn generate_struct(rng: &mut StdRng, len: usize, valid_percent: f64) -> StructArray {
4837 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4838 let a = generate_primitive_array::<Int32Type>(rng, len, valid_percent);
4839 let b = generate_strings::<i32>(rng, len, valid_percent);
4840 let fields = Fields::from(vec![
4841 Field::new("a", DataType::Int32, true),
4842 Field::new("b", DataType::Utf8, true),
4843 ]);
4844 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
4845 StructArray::new(fields, values, Some(nulls))
4846 }
4847
4848 fn generate_list<F>(rng: &mut StdRng, len: usize, valid_percent: f64, values: F) -> ListArray
4849 where
4850 F: FnOnce(&mut StdRng, usize) -> ArrayRef,
4851 {
4852 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4853 let values_len = offsets.last().as_usize();
4854 let values = values(rng, values_len);
4855 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4856 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4857 ListArray::new(field, offsets, values, Some(nulls))
4858 }
4859
4860 fn generate_list_view<F>(
4861 rng: &mut StdRng,
4862 len: usize,
4863 valid_percent: f64,
4864 values: F,
4865 ) -> ListViewArray
4866 where
4867 F: FnOnce(&mut StdRng, usize) -> ArrayRef,
4868 {
4869 let sizes: Vec<i32> = (0..len).map(|_| rng.random_range(0..10)).collect();
4871 let values_len: usize = sizes.iter().map(|s| *s as usize).sum::<usize>().max(1);
4872 let values = values(rng, values_len);
4873
4874 let offsets: Vec<i32> = sizes
4876 .iter()
4877 .map(|&size| {
4878 if size == 0 {
4879 0
4880 } else {
4881 rng.random_range(0..=(values_len as i32 - size))
4882 }
4883 })
4884 .collect();
4885
4886 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4887 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4888 ListViewArray::new(
4889 field,
4890 ScalarBuffer::from(offsets),
4891 ScalarBuffer::from(sizes),
4892 values,
4893 Some(nulls),
4894 )
4895 }
4896
4897 fn generate_map<KeysFn, ValuesFn>(
4898 rng: &mut StdRng,
4899 len: usize,
4900 valid_percent: f64,
4901 gen_keys: KeysFn,
4902 gen_values: ValuesFn,
4903 ) -> MapArray
4904 where
4905 KeysFn: FnOnce(&mut StdRng, usize) -> ArrayRef,
4906 ValuesFn: FnOnce(&mut StdRng, usize) -> ArrayRef,
4907 {
4908 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4909 let entries_len = offsets.last().as_usize();
4910 let keys = gen_keys(rng, entries_len);
4911 let values = gen_values(rng, entries_len);
4912 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4913 let field = Arc::new(Field::new_map(
4914 "",
4915 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4916 Field::new(
4917 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4918 keys.data_type().clone(),
4919 false,
4920 ),
4921 Field::new(
4922 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4923 values.data_type().clone(),
4924 true,
4925 ),
4926 false,
4927 true,
4928 ));
4929 let DataType::Map(struct_field, _) = field.data_type() else {
4930 unreachable!();
4931 };
4932
4933 let DataType::Struct(fields) = struct_field.data_type() else {
4934 unreachable!();
4935 };
4936
4937 let entries = StructArray::new(fields.clone(), vec![keys, values], None);
4938
4939 let map_array = MapArray::new(struct_field.clone(), offsets, entries, Some(nulls), false);
4940
4941 assert_valid_map(&map_array);
4942
4943 map_array
4944 }
4945
4946 fn assert_valid_map(array: &MapArray) {
4956 let keys_arrow_row_converter =
4957 RowConverter::new(vec![SortField::new(array.key_type().clone())]).unwrap();
4958
4959 array.iter().enumerate().filter_map(|(index, entry)| entry.map(|entry| (index, Arc::clone(entry.column(0))))).for_each(|(entry_index, keys)| {
4960 let keys_as_rows = keys_arrow_row_converter.convert_columns(&[Arc::clone(&keys)]).expect("should be able to convert keys");
4961
4962 for i in 0..keys_as_rows.num_rows() {
4963 for j in (i + 1)..keys_as_rows.num_rows() {
4964 if keys_as_rows.row(i) == keys_as_rows.row(j) {
4965 let key_i = keys.slice(i, 1);
4966 let key_j = keys.slice(j, 1);
4967
4968 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:?}");
4969 }
4970 }
4971 }
4972 })
4973 }
4974
4975 fn generate_nulls(rng: &mut StdRng, len: usize) -> Option<NullBuffer> {
4976 Some(NullBuffer::from_iter(
4977 (0..len).map(|_| rng.random_bool(0.8)),
4978 ))
4979 }
4980
4981 fn change_underlying_null_values_for_primitive<T: ArrowPrimitiveType>(
4982 array: &PrimitiveArray<T>,
4983 ) -> PrimitiveArray<T> {
4984 let (dt, values, nulls) = array.clone().into_parts();
4985
4986 let new_values = ScalarBuffer::<T::Native>::from_iter(
4987 values
4988 .iter()
4989 .zip(nulls.as_ref().unwrap().iter())
4990 .map(|(val, is_valid)| {
4991 if is_valid {
4992 *val
4993 } else {
4994 val.add_wrapping(T::Native::usize_as(1))
4995 }
4996 }),
4997 );
4998
4999 PrimitiveArray::new(new_values, nulls).with_data_type(dt)
5000 }
5001
5002 fn change_underline_null_values_for_byte_array<T: ByteArrayType>(
5003 array: &GenericByteArray<T>,
5004 ) -> GenericByteArray<T> {
5005 let (offsets, values, nulls) = array.clone().into_parts();
5006
5007 let new_offsets = OffsetBuffer::<T::Offset>::from_lengths(
5008 offsets
5009 .lengths()
5010 .zip(nulls.as_ref().unwrap().iter())
5011 .map(|(len, is_valid)| if is_valid { len } else { len + 1 }),
5012 );
5013
5014 let mut new_bytes = Vec::<u8>::with_capacity(new_offsets[new_offsets.len() - 1].as_usize());
5015
5016 offsets
5017 .windows(2)
5018 .zip(nulls.as_ref().unwrap().iter())
5019 .for_each(|(start_and_end, is_valid)| {
5020 let start = start_and_end[0].as_usize();
5021 let end = start_and_end[1].as_usize();
5022 new_bytes.extend_from_slice(&values.as_slice()[start..end]);
5023
5024 if !is_valid {
5026 new_bytes.push(b'c');
5027 }
5028 });
5029
5030 GenericByteArray::<T>::new(new_offsets, Buffer::from_vec(new_bytes), nulls)
5031 }
5032
5033 fn change_underline_null_values_for_list_array<O: OffsetSizeTrait>(
5034 array: &GenericListArray<O>,
5035 ) -> GenericListArray<O> {
5036 let (field, offsets, values, nulls) = array.clone().into_parts();
5037
5038 let (new_values, new_offsets) = {
5039 let concat_values = offsets
5040 .windows(2)
5041 .zip(nulls.as_ref().unwrap().iter())
5042 .map(|(start_and_end, is_valid)| {
5043 let start = start_and_end[0].as_usize();
5044 let end = start_and_end[1].as_usize();
5045 if is_valid {
5046 return (start, end - start);
5047 }
5048
5049 if end == values.len() {
5051 (start, (end - start).saturating_sub(1))
5052 } else {
5053 (start, end - start + 1)
5054 }
5055 })
5056 .map(|(start, length)| values.slice(start, length))
5057 .collect::<Vec<_>>();
5058
5059 let new_offsets =
5060 OffsetBuffer::<O>::from_lengths(concat_values.iter().map(|s| s.len()));
5061
5062 let new_values = {
5063 let values = concat_values.iter().map(|a| a.as_ref()).collect::<Vec<_>>();
5064 arrow_select::concat::concat(&values).expect("should be able to concat")
5065 };
5066
5067 (new_values, new_offsets)
5068 };
5069
5070 GenericListArray::<O>::new(field, new_offsets, new_values, nulls)
5071 }
5072
5073 fn change_underline_null_values_for_map_array(array: &MapArray) -> MapArray {
5074 let (field, offsets, entries, nulls, ordered) = array.clone().into_parts();
5075 assert!(
5076 !ordered,
5077 "can't replace underlying null values for ordered map array as this can violate the ordering"
5078 );
5079
5080 let (new_entries, new_offsets) = {
5081 let concat_values = offsets
5082 .windows(2)
5083 .zip(nulls.as_ref().unwrap().iter())
5084 .map(|(start_and_end, is_valid)| {
5085 let start = start_and_end[0].as_usize();
5086 let end = start_and_end[1].as_usize();
5087 if is_valid {
5088 return (start, end - start);
5089 }
5090
5091 if end == entries.len() {
5093 (start, (end - start).saturating_sub(1))
5094 } else {
5095 (start, end - start + 1)
5097 }
5098 })
5099 .map(|(start, length)| entries.slice(start, length))
5100 .collect::<Vec<_>>();
5101
5102 let new_offsets = OffsetBuffer::from_lengths(concat_values.iter().map(|s| s.len()));
5103
5104 let new_values = {
5105 let values = concat_values
5106 .iter()
5107 .map(|a| a as &dyn Array)
5108 .collect::<Vec<_>>();
5109 arrow_select::concat::concat(&values).expect("should be able to concat")
5110 };
5111
5112 (new_values.as_struct().clone(), new_offsets)
5113 };
5114
5115 let new_map = MapArray::new(field, new_offsets, new_entries, nulls, ordered);
5116
5117 assert_valid_map(&new_map);
5118
5119 new_map
5120 }
5121
5122 fn change_underline_null_values(array: &ArrayRef) -> ArrayRef {
5123 if array.null_count() == 0 {
5124 return Arc::clone(array);
5125 }
5126
5127 downcast_primitive_array!(
5128 array => {
5129 let output = change_underlying_null_values_for_primitive(array);
5130
5131 Arc::new(output)
5132 }
5133
5134 DataType::Utf8 => {
5135 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i32>()))
5136 }
5137 DataType::LargeUtf8 => {
5138 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i64>()))
5139 }
5140 DataType::Binary => {
5141 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i32>()))
5142 }
5143 DataType::LargeBinary => {
5144 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i64>()))
5145 }
5146 DataType::List(_) => {
5147 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i32>()))
5148 }
5149 DataType::LargeList(_) => {
5150 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i64>()))
5151 }
5152 DataType::Map(_, _) => {
5153 Arc::new(change_underline_null_values_for_map_array(array.as_map()))
5154 }
5155 _ => {
5156 Arc::clone(array)
5157 }
5158 )
5159 }
5160
5161 fn generate_column(rng: &mut StdRng, len: usize) -> ArrayRef {
5162 match rng.random_range(0..24) {
5163 0 => Arc::new(generate_primitive_array::<Int32Type>(rng, len, 0.8)),
5164 1 => Arc::new(generate_primitive_array::<UInt32Type>(rng, len, 0.8)),
5165 2 => Arc::new(generate_primitive_array::<Int64Type>(rng, len, 0.8)),
5166 3 => Arc::new(generate_primitive_array::<UInt64Type>(rng, len, 0.8)),
5167 4 => Arc::new(generate_primitive_array::<Float32Type>(rng, len, 0.8)),
5168 5 => Arc::new(generate_primitive_array::<Float64Type>(rng, len, 0.8)),
5169 6 => Arc::new(generate_strings::<i32>(rng, len, 0.8)),
5170 7 => {
5171 let dict_values_len = rng.random_range(1..len);
5172 let strings = Arc::new(generate_strings::<i32>(rng, dict_values_len, 1.0));
5174 Arc::new(generate_dictionary::<Int64Type>(rng, strings, len, 0.8))
5175 }
5176 8 => {
5177 let dict_values_len = rng.random_range(1..len);
5178 let values = Arc::new(generate_primitive_array::<Int64Type>(
5180 rng,
5181 dict_values_len,
5182 1.0,
5183 ));
5184 Arc::new(generate_dictionary::<Int64Type>(rng, values, len, 0.8))
5185 }
5186 9 => Arc::new(generate_fixed_size_binary(rng, len, 0.8)),
5187 10 => Arc::new(generate_struct(rng, len, 0.8)),
5188 11 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5189 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5190 })),
5191 12 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5192 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5193 })),
5194 13 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5195 Arc::new(generate_struct(rng, values_len, 0.8))
5196 })),
5197 14 => Arc::new(generate_string_view(rng, len, 0.8)),
5198 15 => Arc::new(generate_byte_view(rng, len, 0.8)),
5199 16 => Arc::new(generate_fixed_stringview_column(len)),
5200 17 => Arc::new(
5201 generate_list(rng, len + 1000, 0.8, |rng, values_len| {
5202 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5203 })
5204 .slice(500, len),
5205 ),
5206 18 => Arc::new(generate_boolean_array(rng, len, 0.8)),
5207 19 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5208 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5209 })),
5210 20 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5211 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5212 })),
5213 21 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5214 Arc::new(generate_struct(rng, values_len, 0.8))
5215 })),
5216 22 => Arc::new(
5217 generate_list_view(rng, len + 1000, 0.8, |rng, values_len| {
5218 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5219 })
5220 .slice(500, len),
5221 ),
5222 23 => Arc::new(generate_map(
5223 rng,
5224 len,
5225 0.9,
5226 |rng, keys_len| {
5228 Arc::new(generate_all_unique_primitive_array::<Int64Type>(
5229 rng, keys_len,
5230 ))
5231 },
5232 |rng, values_len| Arc::new(generate_strings::<i32>(rng, values_len, 0.7)),
5233 )),
5234 _ => unreachable!(),
5235 }
5236 }
5237
5238 fn print_row(cols: &[SortColumn], row: usize) -> String {
5239 let t: Vec<_> = cols
5240 .iter()
5241 .map(|x| match x.values.is_valid(row) {
5242 true => {
5243 let opts = FormatOptions::default().with_null("NULL");
5244 let formatter = ArrayFormatter::try_new(x.values.as_ref(), &opts).unwrap();
5245 formatter.value(row).to_string()
5246 }
5247 false => "NULL".to_string(),
5248 })
5249 .collect();
5250 t.join(",")
5251 }
5252
5253 fn print_col_types(cols: &[SortColumn]) -> String {
5254 let t: Vec<_> = cols
5255 .iter()
5256 .map(|x| x.values.data_type().to_string())
5257 .collect();
5258 t.join(",")
5259 }
5260
5261 #[derive(Debug, PartialEq)]
5262 enum Nulls {
5263 AsIs,
5265
5266 Different,
5268
5269 None,
5271 }
5272
5273 #[test]
5274 #[cfg_attr(miri, ignore)] fn fuzz_test() {
5276 let mut rng = StdRng::seed_from_u64(42);
5277 for _ in 0..100 {
5278 for null_behavior in [Nulls::AsIs, Nulls::Different, Nulls::None] {
5279 let num_columns = rng.random_range(1..5);
5280 let len = rng.random_range(5..100);
5281 let mut arrays: Vec<_> = (0..num_columns)
5282 .map(|_| generate_column(&mut rng, len))
5283 .collect();
5284
5285 match null_behavior {
5286 Nulls::AsIs => {
5287 }
5289 Nulls::Different => {
5290 arrays = arrays
5292 .into_iter()
5293 .map(|a| replace_array_nulls(a, generate_nulls(&mut rng, len)))
5294 .collect()
5295 }
5296 Nulls::None => {
5297 arrays = arrays
5299 .into_iter()
5300 .map(|a| replace_array_nulls(a, None))
5301 .collect()
5302 }
5303 }
5304
5305 let options: Vec<_> = (0..num_columns)
5306 .map(|_| SortOptions {
5307 descending: rng.random_bool(0.5),
5308 nulls_first: rng.random_bool(0.5),
5309 })
5310 .collect();
5311
5312 let sort_columns: Vec<_> = options
5313 .iter()
5314 .zip(&arrays)
5315 .map(|(o, c)| SortColumn {
5316 values: Arc::clone(c),
5317 options: Some(*o),
5318 })
5319 .collect();
5320
5321 let comparator = LexicographicalComparator::try_new(&sort_columns).unwrap();
5322
5323 let columns: Vec<SortField> = options
5324 .into_iter()
5325 .zip(&arrays)
5326 .map(|(o, a)| SortField::new_with_options(a.data_type().clone(), o))
5327 .collect();
5328
5329 let converter = RowConverter::new(columns).unwrap();
5330 let rows = converter.convert_columns(&arrays).unwrap();
5331
5332 if !matches!(null_behavior, Nulls::None) {
5335 assert_same_rows_when_changing_input_underlying_null_values(
5336 &arrays, &converter, &rows,
5337 );
5338 }
5339
5340 for i in 0..len {
5341 for j in 0..len {
5342 let row_i = rows.row(i);
5343 let row_j = rows.row(j);
5344 let row_cmp = row_i.cmp(&row_j);
5345 let lex_cmp = comparator.compare(i, j);
5346 assert_eq!(
5347 row_cmp,
5348 lex_cmp,
5349 "({:?} vs {:?}) vs ({:?} vs {:?}) for types {}",
5350 print_row(&sort_columns, i),
5351 print_row(&sort_columns, j),
5352 row_i,
5353 row_j,
5354 print_col_types(&sort_columns)
5355 );
5356 }
5357 }
5358
5359 {
5361 let mut rows_iter = rows.iter();
5362 let mut rows_lengths_iter = rows.lengths();
5363 for (index, row) in rows_iter.by_ref().enumerate() {
5364 let len = rows_lengths_iter
5365 .next()
5366 .expect("Reached end of length iterator while still have rows");
5367 assert_eq!(
5368 row.data.len(),
5369 len,
5370 "Row length mismatch: {} vs {}",
5371 row.data.len(),
5372 len
5373 );
5374 assert_eq!(
5375 len,
5376 rows.row_len(index),
5377 "Row length mismatch at index {}: {} vs {}",
5378 index,
5379 len,
5380 rows.row_len(index)
5381 );
5382 }
5383
5384 assert_eq!(
5385 rows_lengths_iter.next(),
5386 None,
5387 "Length iterator did not reach end"
5388 );
5389 }
5390
5391 let back = converter.convert_rows(&rows).unwrap();
5394 for (actual, expected) in back.iter().zip(&arrays) {
5395 actual.to_data().validate_full().unwrap();
5396 dictionary_eq(actual, expected)
5397 }
5398
5399 let rows = rows.try_into_binary().expect("reasonable size");
5402 let parser = converter.parser();
5403 let back = converter
5404 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5405 .unwrap();
5406 for (actual, expected) in back.iter().zip(&arrays) {
5407 actual.to_data().validate_full().unwrap();
5408 dictionary_eq(actual, expected)
5409 }
5410
5411 let rows = converter.from_binary(rows);
5412 let back = converter.convert_rows(&rows).unwrap();
5413 for (actual, expected) in back.iter().zip(&arrays) {
5414 actual.to_data().validate_full().unwrap();
5415 dictionary_eq(actual, expected)
5416 }
5417 }
5418 }
5419 }
5420
5421 fn replace_array_nulls(array: ArrayRef, new_nulls: Option<NullBuffer>) -> ArrayRef {
5422 make_array(
5423 array
5424 .into_data()
5425 .into_builder()
5426 .nulls(new_nulls)
5428 .build()
5429 .unwrap(),
5430 )
5431 }
5432
5433 fn assert_same_rows_when_changing_input_underlying_null_values(
5434 arrays: &[ArrayRef],
5435 converter: &RowConverter,
5436 rows: &Rows,
5437 ) {
5438 let arrays_with_different_data_behind_nulls = arrays
5439 .iter()
5440 .map(|arr| change_underline_null_values(arr))
5441 .collect::<Vec<_>>();
5442
5443 if arrays
5445 .iter()
5446 .zip(arrays_with_different_data_behind_nulls.iter())
5447 .all(|(a, b)| Arc::ptr_eq(a, b))
5448 {
5449 return;
5450 }
5451
5452 let rows_with_different_nulls = converter
5453 .convert_columns(&arrays_with_different_data_behind_nulls)
5454 .unwrap();
5455
5456 assert_eq!(
5457 rows.iter().collect::<Vec<_>>(),
5458 rows_with_different_nulls.iter().collect::<Vec<_>>(),
5459 "Different underlying nulls should not output different rows"
5460 )
5461 }
5462
5463 #[test]
5464 fn test_clear() {
5465 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
5466 let mut rows = converter.empty_rows(3, 128);
5467
5468 let first = Int32Array::from(vec![None, Some(2), Some(4)]);
5469 let second = Int32Array::from(vec![Some(2), None, Some(4)]);
5470 let arrays = [Arc::new(first) as ArrayRef, Arc::new(second) as ArrayRef];
5471
5472 for array in &arrays {
5473 rows.clear();
5474 converter
5475 .append(&mut rows, std::slice::from_ref(array))
5476 .unwrap();
5477 let back = converter.convert_rows(&rows).unwrap();
5478 assert_eq!(&back[0], array);
5479 }
5480
5481 let mut rows_expected = converter.empty_rows(3, 128);
5482 converter.append(&mut rows_expected, &arrays[1..]).unwrap();
5483
5484 for (i, (actual, expected)) in rows.iter().zip(rows_expected.iter()).enumerate() {
5485 assert_eq!(
5486 actual, expected,
5487 "For row {i}: expected {expected:?}, actual: {actual:?}",
5488 );
5489 }
5490 }
5491
5492 #[test]
5493 fn test_append_codec_dictionary_binary() {
5494 use DataType::*;
5495 let converter = RowConverter::new(vec![SortField::new(Dictionary(
5497 Box::new(Int32),
5498 Box::new(Binary),
5499 ))])
5500 .unwrap();
5501 let mut rows = converter.empty_rows(4, 128);
5502
5503 let keys = Int32Array::from_iter_values([0, 1, 2, 3]);
5504 let values = BinaryArray::from(vec![
5505 Some(b"a".as_slice()),
5506 Some(b"b"),
5507 Some(b"c"),
5508 Some(b"d"),
5509 ]);
5510 let dict_array = DictionaryArray::new(keys, Arc::new(values));
5511
5512 rows.clear();
5513 let array = Arc::new(dict_array) as ArrayRef;
5514 converter
5515 .append(&mut rows, std::slice::from_ref(&array))
5516 .unwrap();
5517 let back = converter.convert_rows(&rows).unwrap();
5518
5519 dictionary_eq(&back[0], &array);
5520 }
5521
5522 #[test]
5523 fn test_list_prefix() {
5524 let mut a = ListBuilder::new(Int8Builder::new());
5525 a.append_value([None]);
5526 a.append_value([None, None]);
5527 let a = a.finish();
5528
5529 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
5530 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
5531 assert_eq!(rows.row(0).cmp(&rows.row(1)), Ordering::Less);
5532 }
5533
5534 #[test]
5535 fn test_utf8_validation_doesnt_affect_values_buffer_size() {
5536 fn assert_values_buffer_lens(col: ArrayRef) -> usize {
5537 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8View)]).unwrap();
5539
5540 let rows = converter.convert_columns(&[col]).unwrap();
5542 let converted = converter.convert_rows(&rows).unwrap();
5543 let unchecked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5544
5545 let rows = rows.try_into_binary().expect("reasonable size");
5547 let parser = converter.parser();
5548 let converted = converter
5549 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5550 .unwrap();
5551 let checked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5552 assert_eq!(unchecked_values_len, checked_values_len);
5554 checked_values_len
5555 }
5556
5557 let col = Arc::new(StringViewArray::from_iter([
5559 Some("hello"), None, Some("short"), Some("tiny"), ])) as ArrayRef;
5564
5565 let values_len = assert_values_buffer_lens(col);
5566 assert_eq!(values_len, 0);
5568
5569 let col = Arc::new(StringViewArray::from_iter([
5571 Some("1234567890123"), Some("12345678901234"), ])) as ArrayRef;
5574
5575 let values_len = assert_values_buffer_lens(col);
5576 assert_eq!(values_len, 13 + 14);
5577
5578 let col = Arc::new(StringViewArray::from_iter([
5580 Some("tiny"), Some("thisisexact13"), None,
5583 Some("short"), ])) as ArrayRef;
5585
5586 let values_len = assert_values_buffer_lens(col);
5587 assert_eq!(values_len, 13);
5589 }
5590
5591 #[test]
5592 fn test_sparse_union() {
5593 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5595 let str_array = StringArray::from(vec![None, Some("b"), None, Some("d"), None]);
5596
5597 let type_ids = vec![0, 1, 0, 1, 0].into();
5599
5600 let union_fields = [
5601 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5602 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5603 ]
5604 .into_iter()
5605 .collect();
5606
5607 let union_array = UnionArray::try_new(
5608 union_fields,
5609 type_ids,
5610 None,
5611 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5612 )
5613 .unwrap();
5614
5615 let union_type = union_array.data_type().clone();
5616 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5617
5618 let rows = converter
5619 .convert_columns(&[Arc::new(union_array.clone())])
5620 .unwrap();
5621
5622 let back = converter.convert_rows(&rows).unwrap();
5624 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5625
5626 assert_eq!(union_array.len(), back_union.len());
5627 for i in 0..union_array.len() {
5628 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5629 }
5630 }
5631
5632 #[test]
5633 fn test_sparse_union_with_nulls() {
5634 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5636 let str_array = StringArray::from(vec![None::<&str>; 5]);
5637
5638 let type_ids = vec![0, 1, 0, 1, 0].into();
5640
5641 let union_fields = [
5642 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5643 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5644 ]
5645 .into_iter()
5646 .collect();
5647
5648 let union_array = UnionArray::try_new(
5649 union_fields,
5650 type_ids,
5651 None,
5652 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5653 )
5654 .unwrap();
5655
5656 let union_type = union_array.data_type().clone();
5657 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5658
5659 let rows = converter
5660 .convert_columns(&[Arc::new(union_array.clone())])
5661 .unwrap();
5662
5663 let back = converter.convert_rows(&rows).unwrap();
5665 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5666
5667 assert_eq!(union_array.len(), back_union.len());
5668 for i in 0..union_array.len() {
5669 let expected_null = union_array.is_null(i);
5670 let actual_null = back_union.is_null(i);
5671 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5672 if !expected_null {
5673 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5674 }
5675 }
5676 }
5677
5678 #[test]
5679 fn test_dense_union() {
5680 let int_array = Int32Array::from(vec![1, 3, 5]);
5682 let str_array = StringArray::from(vec!["a", "b"]);
5683
5684 let type_ids = vec![0, 1, 0, 1, 0].into();
5685
5686 let offsets = vec![0, 0, 1, 1, 2].into();
5688
5689 let union_fields = [
5690 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5691 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5692 ]
5693 .into_iter()
5694 .collect();
5695
5696 let union_array = UnionArray::try_new(
5697 union_fields,
5698 type_ids,
5699 Some(offsets), vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5701 )
5702 .unwrap();
5703
5704 let union_type = union_array.data_type().clone();
5705 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5706
5707 let rows = converter
5708 .convert_columns(&[Arc::new(union_array.clone())])
5709 .unwrap();
5710
5711 let back = converter.convert_rows(&rows).unwrap();
5713 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5714
5715 assert_eq!(union_array.len(), back_union.len());
5716 for i in 0..union_array.len() {
5717 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5718 }
5719 }
5720
5721 #[test]
5722 fn test_dense_union_type_ids_not_below_field_count() {
5723 let union_fields = [
5725 (2, Arc::new(Field::new("int", DataType::Int32, false))),
5726 (7, Arc::new(Field::new("str", DataType::Utf8, false))),
5727 ]
5728 .into_iter()
5729 .collect();
5730
5731 let union_array: ArrayRef = Arc::new(
5732 UnionArray::try_new(
5733 union_fields,
5734 vec![2, 7, 2, 7].into(),
5735 Some(vec![0, 0, 1, 1].into()),
5736 vec![
5737 Arc::new(Int32Array::from(vec![1, 3])) as ArrayRef,
5738 Arc::new(StringArray::from(vec!["a", "b"])),
5739 ],
5740 )
5741 .unwrap(),
5742 );
5743
5744 let converter =
5745 RowConverter::new(vec![SortField::new(union_array.data_type().clone())]).unwrap();
5746 let rows = converter
5747 .convert_columns(std::slice::from_ref(&union_array))
5748 .unwrap();
5749 let back = converter.convert_rows(&rows).unwrap();
5750
5751 assert_eq!(back[0].to_data(), union_array.to_data());
5752 }
5753
5754 #[test]
5755 fn test_dense_union_with_nulls() {
5756 let int_array = Int32Array::from(vec![Some(1), None, Some(5)]);
5758 let str_array = StringArray::from(vec![Some("a"), None]);
5759
5760 let type_ids = vec![0, 1, 0, 1, 0].into();
5762 let offsets = vec![0, 0, 1, 1, 2].into();
5763
5764 let union_fields = [
5765 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5766 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5767 ]
5768 .into_iter()
5769 .collect();
5770
5771 let union_array = UnionArray::try_new(
5772 union_fields,
5773 type_ids,
5774 Some(offsets),
5775 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5776 )
5777 .unwrap();
5778
5779 let union_type = union_array.data_type().clone();
5780 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5781
5782 let rows = converter
5783 .convert_columns(&[Arc::new(union_array.clone())])
5784 .unwrap();
5785
5786 let back = converter.convert_rows(&rows).unwrap();
5788 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5789
5790 assert_eq!(union_array.len(), back_union.len());
5791 for i in 0..union_array.len() {
5792 let expected_null = union_array.is_null(i);
5793 let actual_null = back_union.is_null(i);
5794 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5795 if !expected_null {
5796 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5797 }
5798 }
5799 }
5800
5801 #[test]
5802 fn test_union_ordering() {
5803 let int_array = Int32Array::from(vec![100, 5, 20]);
5804 let str_array = StringArray::from(vec!["z", "a"]);
5805
5806 let type_ids = vec![0, 1, 0, 1, 0].into();
5808 let offsets = vec![0, 0, 1, 1, 2].into();
5809
5810 let union_fields = [
5811 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5812 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5813 ]
5814 .into_iter()
5815 .collect();
5816
5817 let union_array = UnionArray::try_new(
5818 union_fields,
5819 type_ids,
5820 Some(offsets),
5821 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5822 )
5823 .unwrap();
5824
5825 let union_type = union_array.data_type().clone();
5826 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5827
5828 let rows = converter.convert_columns(&[Arc::new(union_array)]).unwrap();
5829
5830 assert!(rows.row(2) < rows.row(1));
5842
5843 assert!(rows.row(0) < rows.row(3));
5845
5846 assert!(rows.row(2) < rows.row(4));
5849 assert!(rows.row(4) < rows.row(0));
5851
5852 assert!(rows.row(3) < rows.row(1));
5855 }
5856
5857 #[test]
5858 fn test_row_converter_roundtrip_with_many_union_columns() {
5859 let fields1 = UnionFields::try_new(
5861 vec![0, 1],
5862 vec![
5863 Field::new("int", DataType::Int32, true),
5864 Field::new("string", DataType::Utf8, true),
5865 ],
5866 )
5867 .unwrap();
5868
5869 let int_array1 = Int32Array::from(vec![Some(67), None]);
5870 let string_array1 = StringArray::from(vec![None::<&str>, Some("hello")]);
5871 let type_ids1 = vec![0i8, 1].into();
5872
5873 let union_array1 = UnionArray::try_new(
5874 fields1.clone(),
5875 type_ids1,
5876 None,
5877 vec![
5878 Arc::new(int_array1) as ArrayRef,
5879 Arc::new(string_array1) as ArrayRef,
5880 ],
5881 )
5882 .unwrap();
5883
5884 let fields2 = UnionFields::try_new(
5886 vec![0, 1],
5887 vec![
5888 Field::new("int", DataType::Int32, true),
5889 Field::new("string", DataType::Utf8, true),
5890 ],
5891 )
5892 .unwrap();
5893
5894 let int_array2 = Int32Array::from(vec![Some(100), None]);
5895 let string_array2 = StringArray::from(vec![None::<&str>, Some("world")]);
5896 let type_ids2 = vec![0i8, 1].into();
5897
5898 let union_array2 = UnionArray::try_new(
5899 fields2.clone(),
5900 type_ids2,
5901 None,
5902 vec![
5903 Arc::new(int_array2) as ArrayRef,
5904 Arc::new(string_array2) as ArrayRef,
5905 ],
5906 )
5907 .unwrap();
5908
5909 let field1 = Field::new("col1", DataType::Union(fields1, UnionMode::Sparse), true);
5911 let field2 = Field::new("col2", DataType::Union(fields2, UnionMode::Sparse), true);
5912
5913 let sort_field1 = SortField::new(field1.data_type().clone());
5914 let sort_field2 = SortField::new(field2.data_type().clone());
5915
5916 let converter = RowConverter::new(vec![sort_field1, sort_field2]).unwrap();
5917
5918 let rows = converter
5919 .convert_columns(&[
5920 Arc::new(union_array1.clone()) as ArrayRef,
5921 Arc::new(union_array2.clone()) as ArrayRef,
5922 ])
5923 .unwrap();
5924
5925 let out = converter.convert_rows(&rows).unwrap();
5927
5928 let [col1, col2] = out.as_slice() else {
5929 panic!("expected 2 columns")
5930 };
5931
5932 let col1 = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5933 let col2 = col2.as_any().downcast_ref::<UnionArray>().unwrap();
5934
5935 for (expected, got) in [union_array1, union_array2].iter().zip([col1, col2]) {
5936 assert_eq!(expected.len(), got.len());
5937 assert_eq!(expected.type_ids(), got.type_ids());
5938
5939 for i in 0..expected.len() {
5940 assert_eq!(expected.value(i).as_ref(), got.value(i).as_ref());
5941 }
5942 }
5943 }
5944
5945 #[test]
5946 fn test_row_converter_roundtrip_with_one_union_column() {
5947 let fields = UnionFields::try_new(
5948 vec![0, 1],
5949 vec![
5950 Field::new("int", DataType::Int32, true),
5951 Field::new("string", DataType::Utf8, true),
5952 ],
5953 )
5954 .unwrap();
5955
5956 let int_array = Int32Array::from(vec![Some(67), None]);
5957 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5958 let type_ids = vec![0i8, 1].into();
5959
5960 let union_array = UnionArray::try_new(
5961 fields.clone(),
5962 type_ids,
5963 None,
5964 vec![
5965 Arc::new(int_array) as ArrayRef,
5966 Arc::new(string_array) as ArrayRef,
5967 ],
5968 )
5969 .unwrap();
5970
5971 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5972 let sort_field = SortField::new(field.data_type().clone());
5973 let converter = RowConverter::new(vec![sort_field]).unwrap();
5974
5975 let rows = converter
5976 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5977 .unwrap();
5978
5979 let out = converter.convert_rows(&rows).unwrap();
5981
5982 let [col1] = out.as_slice() else {
5983 panic!("expected 1 column")
5984 };
5985
5986 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5987 assert_eq!(col.len(), union_array.len());
5988 assert_eq!(col.type_ids(), union_array.type_ids());
5989
5990 for i in 0..col.len() {
5991 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5992 }
5993 }
5994
5995 #[test]
5996 fn test_row_converter_roundtrip_with_non_default_union_type_ids() {
5997 let fields = UnionFields::try_new(
5999 vec![70, 85],
6000 vec![
6001 Field::new("int", DataType::Int32, true),
6002 Field::new("string", DataType::Utf8, true),
6003 ],
6004 )
6005 .unwrap();
6006
6007 let int_array = Int32Array::from(vec![Some(67), None]);
6008 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
6009 let type_ids = vec![70i8, 85].into();
6010
6011 let union_array = UnionArray::try_new(
6012 fields.clone(),
6013 type_ids,
6014 None,
6015 vec![
6016 Arc::new(int_array) as ArrayRef,
6017 Arc::new(string_array) as ArrayRef,
6018 ],
6019 )
6020 .unwrap();
6021
6022 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
6023 let sort_field = SortField::new(field.data_type().clone());
6024 let converter = RowConverter::new(vec![sort_field]).unwrap();
6025
6026 let rows = converter
6027 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
6028 .unwrap();
6029
6030 let out = converter.convert_rows(&rows).unwrap();
6032
6033 let [col1] = out.as_slice() else {
6034 panic!("expected 1 column")
6035 };
6036
6037 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
6038 assert_eq!(col.len(), union_array.len());
6039 assert_eq!(col.type_ids(), union_array.type_ids());
6040
6041 for i in 0..col.len() {
6042 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
6043 }
6044 }
6045
6046 #[test]
6047 fn rows_size_should_count_for_capacity() {
6048 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6049
6050 let empty_rows_size_with_preallocate_rows_and_data = {
6051 let rows = row_converter.empty_rows(1000, 1000);
6052
6053 rows.size()
6054 };
6055 let empty_rows_size_with_preallocate_rows = {
6056 let rows = row_converter.empty_rows(1000, 0);
6057
6058 rows.size()
6059 };
6060 let empty_rows_size_with_preallocate_data = {
6061 let rows = row_converter.empty_rows(0, 1000);
6062
6063 rows.size()
6064 };
6065 let empty_rows_size_without_preallocate = {
6066 let rows = row_converter.empty_rows(0, 0);
6067
6068 rows.size()
6069 };
6070
6071 assert!(
6072 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_rows,
6073 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_rows}"
6074 );
6075 assert!(
6076 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_data,
6077 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_data}"
6078 );
6079 assert!(
6080 empty_rows_size_with_preallocate_rows > empty_rows_size_without_preallocate,
6081 "{empty_rows_size_with_preallocate_rows} should be larger than {empty_rows_size_without_preallocate}"
6082 );
6083 assert!(
6084 empty_rows_size_with_preallocate_data > empty_rows_size_without_preallocate,
6085 "{empty_rows_size_with_preallocate_data} should be larger than {empty_rows_size_without_preallocate}"
6086 );
6087 }
6088
6089 #[test]
6090 fn test_struct_no_child_fields() {
6091 fn run_test(array: ArrayRef) {
6092 let sort_fields = vec![SortField::new(array.data_type().clone())];
6093 let converter = RowConverter::new(sort_fields).unwrap();
6094 let r = converter.convert_columns(&[Arc::clone(&array)]).unwrap();
6095
6096 let back = converter.convert_rows(&r).unwrap();
6097 assert_eq!(back.len(), 1);
6098 assert_eq!(&back[0], &array);
6099 }
6100
6101 let s = Arc::new(StructArray::new_empty_fields(5, None)) as ArrayRef;
6102 run_test(s);
6103
6104 let s = Arc::new(StructArray::new_empty_fields(
6105 5,
6106 Some(vec![true, false, true, false, false].into()),
6107 )) as ArrayRef;
6108 run_test(s);
6109 }
6110
6111 #[test]
6112 fn reserve_should_increase_capacity_to_the_requested_size() {
6113 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6114 let mut empty_rows = row_converter.empty_rows(0, 0);
6115 empty_rows.reserve(50, 50);
6116 let before_size = empty_rows.size();
6117 empty_rows.reserve(50, 50);
6118 assert_eq!(
6119 empty_rows.size(),
6120 before_size,
6121 "Size should not change when reserving already reserved space"
6122 );
6123 empty_rows.reserve(10, 20);
6124 assert_eq!(
6125 empty_rows.size(),
6126 before_size,
6127 "Size should not change when already have space for the expected reserved data"
6128 );
6129
6130 empty_rows.reserve(100, 20);
6131 assert!(
6132 empty_rows.size() > before_size,
6133 "Size should increase when reserving more space than previously reserved"
6134 );
6135
6136 let before_size = empty_rows.size();
6137
6138 empty_rows.reserve(20, 100);
6139 assert!(
6140 empty_rows.size() > before_size,
6141 "Size should increase when reserving more space than previously reserved"
6142 );
6143 }
6144
6145 #[test]
6146 fn empty_rows_should_return_empty_lengths_iterator() {
6147 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6148 .unwrap()
6149 .empty_rows(0, 0);
6150 let mut lengths_iter = rows.lengths();
6151 assert_eq!(lengths_iter.next(), None);
6152 }
6153
6154 #[test]
6155 #[should_panic(expected = "row index out of bounds")]
6156 fn row_should_panic_on_overflowing_index() {
6157 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6158 .unwrap()
6159 .empty_rows(0, 0);
6160 rows.row(usize::MAX);
6161 }
6162
6163 #[test]
6164 #[should_panic(expected = "row index out of bounds")]
6165 fn row_len_should_panic_on_overflowing_index() {
6166 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6167 .unwrap()
6168 .empty_rows(0, 0);
6169 rows.row_len(usize::MAX);
6170 }
6171
6172 #[test]
6173 fn test_nested_null_list() {
6174 let null_array = Arc::new(NullArray::new(3));
6175 let list: ArrayRef = Arc::new(ListArray::new(
6177 Field::new_list_field(DataType::Null, true).into(),
6178 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6179 null_array,
6180 None,
6181 ));
6182
6183 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6184 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6185 let back = converter.convert_rows(&rows).unwrap();
6186
6187 assert_eq!(&list, &back[0]);
6188 }
6189
6190 #[test]
6192 fn test_double_nested_null_list() {
6193 let null_array = Arc::new(NullArray::new(1));
6194 let nested_field = Arc::new(Field::new_list_field(DataType::Null, true));
6196 let nested_list = Arc::new(ListArray::new(
6197 nested_field.clone(),
6198 OffsetBuffer::from_lengths(vec![1]),
6199 null_array,
6200 None,
6201 ));
6202 let list = Arc::new(ListArray::new(
6204 Field::new_list_field(DataType::List(nested_field), true).into(),
6205 OffsetBuffer::from_lengths(vec![1]),
6206 nested_list,
6207 None,
6208 )) as ArrayRef;
6209
6210 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6211 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6212 let back = converter.convert_rows(&rows).unwrap();
6213
6214 assert_eq!(&list, &back[0]);
6215 }
6216
6217 #[test]
6219 fn test_large_list_null() {
6220 let null_array = Arc::new(NullArray::new(3));
6221 let list: ArrayRef = Arc::new(LargeListArray::new(
6223 Field::new_list_field(DataType::Null, true).into(),
6224 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6225 null_array,
6226 None,
6227 ));
6228
6229 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6230 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6231 let back = converter.convert_rows(&rows).unwrap();
6232
6233 assert_eq!(&list, &back[0]);
6234 }
6235
6236 #[test]
6238 fn test_fixed_size_list_null() {
6239 let null_array = Arc::new(NullArray::new(6));
6240 let list: ArrayRef = Arc::new(FixedSizeListArray::new(
6242 Arc::new(Field::new_list_field(DataType::Null, true)),
6243 2,
6244 null_array,
6245 None,
6246 ));
6247
6248 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6249 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6250 let back = converter.convert_rows(&rows).unwrap();
6251
6252 assert_eq!(&list, &back[0]);
6253 }
6254
6255 #[test]
6257 fn test_fixed_size_list_of_dictionaries_round_trips() {
6258 let dict_dt = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
6261 let element_field = Arc::new(Field::new("item", dict_dt.clone(), true));
6262 let fsl_dt = DataType::FixedSizeList(Arc::clone(&element_field), 2);
6263
6264 let values = Arc::new(StringArray::from(vec!["a", "b"]));
6265 let keys = Int32Array::from(vec![0, 1]);
6266 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
6267 let fsl: ArrayRef = Arc::new(FixedSizeListArray::new(
6268 Arc::clone(&element_field),
6269 2,
6270 Arc::new(dict),
6271 None,
6272 ));
6273
6274 assert!(RowConverter::supports_fields(&[SortField::new(
6275 fsl_dt.clone()
6276 )]));
6277
6278 let converter = RowConverter::new(vec![SortField::new(fsl_dt.clone())]).unwrap();
6279 let rows = converter.convert_columns(&[Arc::clone(&fsl)]).unwrap();
6280
6281 let back = converter.convert_rows(&rows).unwrap();
6284 assert_eq!(back.len(), 1);
6285
6286 let out = back[0]
6290 .as_any()
6291 .downcast_ref::<FixedSizeListArray>()
6292 .expect("decoded array must be a FixedSizeListArray");
6293 assert_eq!(out.len(), 1);
6294 assert_eq!(out.value_length(), 2);
6295 assert_eq!(out.values().data_type(), &DataType::Utf8);
6299
6300 let values = out
6302 .values()
6303 .as_any()
6304 .downcast_ref::<StringArray>()
6305 .expect("child must be a StringArray after flattening");
6306 assert_eq!(values.value(0), "a");
6307 assert_eq!(values.value(1), "b");
6308 }
6309
6310 #[test]
6312 fn test_list_null_variations() {
6313 let null_array = Arc::new(NullArray::new(3));
6315 let list: ArrayRef = Arc::new(ListArray::new(
6316 Field::new_list_field(DataType::Null, true).into(),
6317 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6318 null_array,
6319 None,
6320 ));
6321
6322 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6323 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6324 let back = converter.convert_rows(&rows).unwrap();
6325 assert_eq!(&list, &back[0]);
6326
6327 let null_array = Arc::new(NullArray::new(3));
6329 let list: ArrayRef = Arc::new(ListArray::new(
6330 Field::new_list_field(DataType::Null, true).into(),
6331 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6332 null_array,
6333 Some(vec![true, false, true].into()),
6334 ));
6335
6336 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6337 let back = converter.convert_rows(&rows).unwrap();
6338 assert_eq!(&list, &back[0]);
6339
6340 let null_array = Arc::new(NullArray::new(0));
6342 let list: ArrayRef = Arc::new(ListArray::new(
6343 Field::new_list_field(DataType::Null, true).into(),
6344 OffsetBuffer::from_lengths(vec![]),
6345 null_array,
6346 None,
6347 ));
6348
6349 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6350 let back = converter.convert_rows(&rows).unwrap();
6351 assert_eq!(&list, &back[0]);
6352
6353 let null_array = Arc::new(NullArray::new(0));
6355 let list: ArrayRef = Arc::new(ListArray::new(
6356 Field::new_list_field(DataType::Null, true).into(),
6357 OffsetBuffer::from_lengths(vec![0, 0, 0]),
6358 null_array,
6359 None,
6360 ));
6361
6362 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6363 let back = converter.convert_rows(&rows).unwrap();
6364 assert_eq!(&list, &back[0]);
6365 }
6366
6367 #[test]
6369 fn test_list_null_descending() {
6370 let null_array = Arc::new(NullArray::new(3));
6371 let list: ArrayRef = Arc::new(ListArray::new(
6373 Field::new_list_field(DataType::Null, true).into(),
6374 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6375 null_array,
6376 None,
6377 ));
6378
6379 let options = SortOptions::default().with_descending(true);
6380 let field = SortField::new_with_options(list.data_type().clone(), options);
6381 let converter = RowConverter::new(vec![field]).unwrap();
6382 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6383 let back = converter.convert_rows(&rows).unwrap();
6384
6385 assert_eq!(&list, &back[0]);
6386 }
6387
6388 #[test]
6390 fn test_struct_with_null_field() {
6391 let null_array = Arc::new(NullArray::new(3));
6393 let int_array = Arc::new(Int32Array::from(vec![1, 2, 3]));
6394
6395 let struct_array: ArrayRef = Arc::new(StructArray::new(
6396 vec![
6397 Arc::new(Field::new("a", DataType::Null, true)),
6398 Arc::new(Field::new("b", DataType::Int32, true)),
6399 ]
6400 .into(),
6401 vec![null_array, int_array],
6402 Some(vec![true, true, false].into()), ));
6404
6405 let converter =
6406 RowConverter::new(vec![SortField::new(struct_array.data_type().clone())]).unwrap();
6407 let rows = converter
6408 .convert_columns(&[Arc::clone(&struct_array)])
6409 .unwrap();
6410 let back = converter.convert_rows(&rows).unwrap();
6411
6412 assert_eq!(&struct_array, &back[0]);
6413 }
6414
6415 #[test]
6417 fn test_nested_struct_with_null() {
6418 let inner_null = Arc::new(NullArray::new(2));
6420 let inner_struct = Arc::new(StructArray::new(
6421 vec![Arc::new(Field::new("x", DataType::Null, true))].into(),
6422 vec![inner_null],
6423 None,
6424 ));
6425
6426 let y_array = Arc::new(Int32Array::from(vec![10, 20]));
6428 let outer_struct: ArrayRef = Arc::new(StructArray::new(
6429 vec![
6430 Arc::new(Field::new("inner", inner_struct.data_type().clone(), true)),
6431 Arc::new(Field::new("y", DataType::Int32, true)),
6432 ]
6433 .into(),
6434 vec![inner_struct, y_array],
6435 None,
6436 ));
6437
6438 let converter =
6439 RowConverter::new(vec![SortField::new(outer_struct.data_type().clone())]).unwrap();
6440 let rows = converter
6441 .convert_columns(&[Arc::clone(&outer_struct)])
6442 .unwrap();
6443 let back = converter.convert_rows(&rows).unwrap();
6444
6445 assert_eq!(&outer_struct, &back[0]);
6446 }
6447
6448 #[test]
6450 fn test_map_null_variations() {
6451 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6453 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6454
6455 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6456 let entries_fields = vec![
6457 Arc::new(Field::new(
6458 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6459 DataType::Utf8,
6460 false,
6461 )),
6462 Arc::new(Field::new(
6463 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6464 DataType::Null,
6465 true,
6466 )),
6467 ];
6468 let struct_field = Arc::new(Field::new(
6469 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6470 DataType::Struct(entries_fields.clone().into()),
6471 false,
6472 ));
6473 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6474
6475 let map: ArrayRef = Arc::new(MapArray::new(
6476 struct_field.clone(),
6477 offsets,
6478 entries,
6479 None,
6480 false,
6481 ));
6482
6483 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6484 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6485 let back = converter.convert_rows(&rows).unwrap();
6486 assert_eq!(back.len(), 1);
6487 back[0].to_data().validate_full().unwrap();
6488 assert_eq!(&map, &back[0]);
6489
6490 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6492 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6493
6494 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6495 let entries_fields = vec![
6496 Arc::new(Field::new(
6497 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6498 DataType::Utf8,
6499 false,
6500 )),
6501 Arc::new(Field::new(
6502 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6503 DataType::Null,
6504 true,
6505 )),
6506 ];
6507 let struct_field = Arc::new(Field::new(
6508 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6509 DataType::Struct(entries_fields.clone().into()),
6510 false,
6511 ));
6512 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6513
6514 let map: ArrayRef = Arc::new(MapArray::new(
6515 struct_field.clone(),
6516 offsets,
6517 entries,
6518 Some(vec![true, false, true].into()),
6519 false,
6520 ));
6521
6522 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6523 let back = converter.convert_rows(&rows).unwrap();
6524 assert_eq!(back.len(), 1);
6525 back[0].to_data().validate_full().unwrap();
6526 assert_eq!(&map, &back[0]);
6527
6528 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6530 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6531
6532 let offsets = OffsetBuffer::new(vec![0i32].into());
6533 let entries_fields = vec![
6534 Arc::new(Field::new(
6535 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6536 DataType::Utf8,
6537 false,
6538 )),
6539 Arc::new(Field::new(
6540 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6541 DataType::Null,
6542 true,
6543 )),
6544 ];
6545 let struct_field = Arc::new(Field::new(
6546 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6547 DataType::Struct(entries_fields.clone().into()),
6548 false,
6549 ));
6550 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6551
6552 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6553
6554 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6555 let back = converter.convert_rows(&rows).unwrap();
6556 assert_eq!(back.len(), 1);
6557 back[0].to_data().validate_full().unwrap();
6558 assert_eq!(&map, &back[0]);
6559 }
6560
6561 #[test]
6563 fn test_map_null_descending() {
6564 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6566 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6567
6568 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6569 let entries_fields = vec![
6570 Arc::new(Field::new(
6571 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6572 DataType::Utf8,
6573 false,
6574 )),
6575 Arc::new(Field::new(
6576 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6577 DataType::Null,
6578 true,
6579 )),
6580 ];
6581 let struct_field = Arc::new(Field::new(
6582 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6583 DataType::Struct(entries_fields.clone().into()),
6584 false,
6585 ));
6586 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6587
6588 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6589
6590 let options = SortOptions::default().with_descending(true);
6591 let field = SortField::new_with_options(map.data_type().clone(), options);
6592 let converter = RowConverter::new(vec![field]).unwrap();
6593 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6594 let back = converter.convert_rows(&rows).unwrap();
6595 assert_eq!(back.len(), 1);
6596 back[0].to_data().validate_full().unwrap();
6597 assert_eq!(&map, &back[0]);
6598 }
6599
6600 #[test]
6602 fn test_map_null_all_empty() {
6603 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6604 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6605
6606 let offsets = OffsetBuffer::new(vec![0, 0, 0, 0].into());
6607 let entries_fields = vec![
6608 Arc::new(Field::new(
6609 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6610 DataType::Utf8,
6611 false,
6612 )),
6613 Arc::new(Field::new(
6614 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6615 DataType::Null,
6616 true,
6617 )),
6618 ];
6619 let struct_field = Arc::new(Field::new(
6620 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6621 DataType::Struct(entries_fields.clone().into()),
6622 false,
6623 ));
6624 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6625
6626 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6627
6628 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6629 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6630
6631 assert_eq!(rows.row(0), rows.row(1));
6633 assert_eq!(rows.row(1), rows.row(2));
6634
6635 let back = converter.convert_rows(&rows).unwrap();
6636 assert_eq!(back.len(), 1);
6637 back[0].to_data().validate_full().unwrap();
6638 assert_eq!(&map, &back[0]);
6639 }
6640
6641 #[test]
6643 fn test_nested_map_null() {
6644 let inner_keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6646 let inner_null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6647
6648 let inner_entries_fields = vec![
6649 Arc::new(Field::new(
6650 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6651 DataType::Utf8,
6652 false,
6653 )),
6654 Arc::new(Field::new(
6655 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6656 DataType::Null,
6657 true,
6658 )),
6659 ];
6660 let inner_struct_field = Arc::new(Field::new(
6661 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6662 DataType::Struct(inner_entries_fields.clone().into()),
6663 false,
6664 ));
6665 let inner_entries = StructArray::new(
6666 inner_entries_fields.clone().into(),
6667 vec![inner_keys, inner_null_values],
6668 None,
6669 );
6670
6671 let inner_map = Arc::new(MapArray::new(
6673 inner_struct_field.clone(),
6674 OffsetBuffer::new(vec![0, 1, 3].into()),
6675 inner_entries,
6676 None,
6677 false,
6678 )) as ArrayRef;
6679
6680 let outer_keys = Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef;
6682
6683 let inner_map_type = DataType::Map(inner_struct_field.clone(), false);
6684 let outer_entries_fields = vec![
6685 Arc::new(Field::new(
6686 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6687 DataType::Utf8,
6688 false,
6689 )),
6690 Arc::new(Field::new(
6691 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6692 inner_map_type,
6693 true,
6694 )),
6695 ];
6696 let outer_struct_field = Arc::new(Field::new(
6697 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6698 DataType::Struct(outer_entries_fields.clone().into()),
6699 false,
6700 ));
6701 let outer_entries = StructArray::new(
6702 outer_entries_fields.into(),
6703 vec![outer_keys, inner_map],
6704 None,
6705 );
6706
6707 let map: ArrayRef = Arc::new(MapArray::new(
6709 outer_struct_field,
6710 OffsetBuffer::new(vec![0, 1, 2].into()),
6711 outer_entries,
6712 None,
6713 false,
6714 ));
6715
6716 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6717 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6718 let back = converter.convert_rows(&rows).unwrap();
6719 assert_eq!(back.len(), 1);
6720 back[0].to_data().validate_full().unwrap();
6721 assert_eq!(&map, &back[0]);
6722 }
6723
6724 #[test]
6726 fn test_list_of_map_null() {
6727 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6729 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6730
6731 let entries_fields = vec![
6732 Arc::new(Field::new(
6733 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6734 DataType::Utf8,
6735 false,
6736 )),
6737 Arc::new(Field::new(
6738 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6739 DataType::Null,
6740 true,
6741 )),
6742 ];
6743 let struct_field = Arc::new(Field::new(
6744 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6745 DataType::Struct(entries_fields.clone().into()),
6746 false,
6747 ));
6748 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6749
6750 let map_array = Arc::new(MapArray::new(
6752 struct_field.clone(),
6753 OffsetBuffer::new(vec![0, 1, 1, 3].into()),
6754 entries,
6755 None,
6756 false,
6757 )) as ArrayRef;
6758
6759 let map_type = DataType::Map(struct_field, false);
6760 let list: ArrayRef = Arc::new(ListArray::new(
6762 Arc::new(Field::new_list_field(map_type, true)),
6763 OffsetBuffer::new(vec![0, 1, 3].into()),
6764 map_array,
6765 None,
6766 ));
6767
6768 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6769 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6770 let back = converter.convert_rows(&rows).unwrap();
6771 assert_eq!(&list, &back[0]);
6772 }
6773
6774 #[test]
6776 fn test_map_of_list_null() {
6777 let null_array = Arc::new(NullArray::new(3)) as ArrayRef;
6779 let list_array = Arc::new(ListArray::new(
6781 Arc::new(Field::new_list_field(DataType::Null, true)),
6782 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6783 null_array,
6784 None,
6785 )) as ArrayRef;
6786
6787 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6788
6789 let list_type = list_array.data_type().clone();
6790 let entries_fields = vec![
6791 Arc::new(Field::new(
6792 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6793 DataType::Utf8,
6794 false,
6795 )),
6796 Arc::new(Field::new(
6797 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6798 list_type,
6799 true,
6800 )),
6801 ];
6802 let struct_field = Arc::new(Field::new(
6803 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6804 DataType::Struct(entries_fields.clone().into()),
6805 false,
6806 ));
6807 let entries = StructArray::new(entries_fields.into(), vec![keys, list_array], None);
6808
6809 let map: ArrayRef = Arc::new(MapArray::new(
6811 struct_field,
6812 OffsetBuffer::new(vec![0, 3].into()),
6813 entries,
6814 None,
6815 false,
6816 ));
6817
6818 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6819 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6820 let back = converter.convert_rows(&rows).unwrap();
6821 assert_eq!(back.len(), 1);
6822 back[0].to_data().validate_full().unwrap();
6823 assert_eq!(&map, &back[0]);
6824 }
6825
6826 #[test]
6827 fn empty_row_iter_next_back() {
6828 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6829 .unwrap()
6830 .empty_rows(0, 0);
6831 let mut rows_iter = rows.iter();
6832 assert_eq!(rows_iter.next_back(), None);
6833 assert_eq!(rows_iter.next_back(), None);
6834 assert_eq!(rows_iter.next_back(), None);
6835 }
6836
6837 #[test]
6839 fn test_row_parser_skip_utf8_validation_roundtrip() {
6840 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
6841 let array = StringArray::from(vec!["arrow", "rust"]);
6842 let rows = converter.convert_columns(&[Arc::new(array) as _]).unwrap();
6843 let binary = rows.try_into_binary().expect("fits in i32 offsets");
6844
6845 let parser = unsafe { RowParser::with_skip_utf8_validate(Arc::clone(&converter.fields)) };
6847
6848 let decoded = converter
6849 .convert_rows(binary.iter().map(|b| parser.parse(b.unwrap())))
6850 .unwrap();
6851 let got: Vec<_> = decoded[0].as_string::<i32>().iter().flatten().collect();
6852 assert_eq!(got, vec!["arrow", "rust"]);
6853 }
6854
6855 #[test]
6856 fn row_iter_next_back() {
6857 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6858 let mut rng = StdRng::seed_from_u64(42);
6859 let array = generate_primitive_array::<UInt8Type>(&mut rng, 100, 0.8);
6860 let rows = row_converter.convert_columns(&[Arc::new(array)]).unwrap();
6861
6862 let mut rows_iter = rows.iter();
6863 let mut bytes: Vec<u8> = vec![];
6864
6865 while let Some(row) = rows_iter.next_back() {
6866 bytes.extend(row.data.iter().rev());
6867 }
6868
6869 bytes.reverse();
6870
6871 assert_eq!(
6872 bytes,
6873 &rows.buffer.as_slice()[..*rows.offsets.last().unwrap()]
6874 );
6875
6876 assert_eq!(rows_iter.next_back(), None);
6877 assert_eq!(rows_iter.next(), None);
6878 }
6879
6880 trait SharedBetweenRowsIteratorAndSliceIter<'a>:
6881 ExactSizeIterator<Item = Row<'a>> + DoubleEndedIterator<Item = Row<'a>>
6882 {
6883 }
6884 impl<'a, T: ?Sized + ExactSizeIterator<Item = Row<'a>> + DoubleEndedIterator<Item = Row<'a>>>
6885 SharedBetweenRowsIteratorAndSliceIter<'a> for T
6886 {
6887 }
6888
6889 fn get_rows_iterator_cases() -> impl Iterator<Item = Rows> {
6890 let rows_converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
6891
6892 [0, 1, 6, 8, 100, 164]
6893 .map(|len| {
6894 let source = (0..).take(len).collect::<Vec<i32>>();
6895
6896 let source = Int32Array::from(source);
6897
6898 rows_converter.convert_columns(&[Arc::new(source)]).unwrap()
6899 })
6900 .into_iter()
6901 }
6902
6903 fn setup_and_assert(
6904 setup_iters: impl Fn(&mut dyn SharedBetweenRowsIteratorAndSliceIter),
6905 assert_fn: impl Fn(RowsIter, Copied<Iter<Row>>),
6906 ) {
6907 for rows in get_rows_iterator_cases() {
6908 let expected = (0..rows.num_rows())
6909 .map(|i| rows.row(i))
6910 .collect::<Vec<_>>();
6911 let mut expected_iter = expected.iter().copied();
6912
6913 let mut actual = rows.iter();
6914
6915 setup_iters(&mut actual);
6916 setup_iters(&mut expected_iter);
6917
6918 assert_fn(actual, expected_iter);
6919 }
6920 }
6921
6922 trait RowsIteratorOp {
6925 type Output<'a>: PartialEq + Debug;
6927
6928 const NAME: &'static str;
6930
6931 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(iter: T)
6934 -> Self::Output<'a>;
6935 }
6936
6937 fn assert_rows_iterator_cases<O: RowsIteratorOp>() {
6941 setup_and_assert(
6942 |_iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {},
6943 |actual, expected| {
6944 let current_iterator_values: Vec<Row> = expected.clone().collect();
6945 assert_eq!(
6946 O::get_value(actual),
6947 O::get_value(expected),
6948 "Failed on op {} for new iter (left actual, right expected) ({current_iterator_values:?})",
6949 O::NAME
6950 );
6951 },
6952 );
6953
6954 setup_and_assert(
6955 |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {
6956 iter.next();
6957 },
6958 |actual, expected| {
6959 let current_iterator_values: Vec<Row> = expected.clone().collect();
6960
6961 assert_eq!(
6962 O::get_value(actual),
6963 O::get_value(expected),
6964 "Failed on op {} for new iter after consuming 1 element from the start (left actual, right expected) ({current_iterator_values:?})",
6965 O::NAME
6966 );
6967 },
6968 );
6969
6970 setup_and_assert(
6971 |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {
6972 iter.next_back();
6973 },
6974 |actual, expected| {
6975 let current_iterator_values: Vec<Row> = expected.clone().collect();
6976
6977 assert_eq!(
6978 O::get_value(actual),
6979 O::get_value(expected),
6980 "Failed on op {} for new iter after consuming 1 element from the end (left actual, right expected) ({current_iterator_values:?})",
6981 O::NAME
6982 );
6983 },
6984 );
6985
6986 setup_and_assert(
6987 |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {
6988 iter.next();
6989 iter.next_back();
6990 },
6991 |actual, expected| {
6992 let current_iterator_values: Vec<Row> = expected.clone().collect();
6993
6994 assert_eq!(
6995 O::get_value(actual),
6996 O::get_value(expected),
6997 "Failed on op {} for new iter after consuming 1 element from start and end (left actual, right expected) ({current_iterator_values:?})",
6998 O::NAME
6999 );
7000 },
7001 );
7002
7003 setup_and_assert(
7004 |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {
7005 while iter.len() > 1 {
7006 iter.next();
7007 }
7008 },
7009 |actual, expected| {
7010 let current_iterator_values: Vec<Row> = expected.clone().collect();
7011
7012 assert_eq!(
7013 O::get_value(actual),
7014 O::get_value(expected),
7015 "Failed on op {} for new iter after consuming all from the start but 1 (left actual, right expected) ({current_iterator_values:?})",
7016 O::NAME
7017 );
7018 },
7019 );
7020
7021 setup_and_assert(
7022 |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {
7023 while iter.len() > 1 {
7024 iter.next_back();
7025 }
7026 },
7027 |actual, expected| {
7028 let current_iterator_values: Vec<Row> = expected.clone().collect();
7029
7030 assert_eq!(
7031 O::get_value(actual),
7032 O::get_value(expected),
7033 "Failed on op {} for new iter after consuming all from the end but 1 (left actual, right expected) ({current_iterator_values:?})",
7034 O::NAME
7035 );
7036 },
7037 );
7038
7039 setup_and_assert(
7040 |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {
7041 while iter.next().is_some() {}
7042 },
7043 |actual, expected| {
7044 let current_iterator_values: Vec<Row> = expected.clone().collect();
7045
7046 assert_eq!(
7047 O::get_value(actual),
7048 O::get_value(expected),
7049 "Failed on op {} for new iter after consuming all from the start (left actual, right expected) ({current_iterator_values:?})",
7050 O::NAME
7051 );
7052 },
7053 );
7054
7055 setup_and_assert(
7056 |iter: &mut dyn SharedBetweenRowsIteratorAndSliceIter| {
7057 while iter.next_back().is_some() {}
7058 },
7059 |actual, expected| {
7060 let current_iterator_values: Vec<Row> = expected.clone().collect();
7061
7062 assert_eq!(
7063 O::get_value(actual),
7064 O::get_value(expected),
7065 "Failed on op {} for new iter after consuming all from the end (left actual, right expected) ({current_iterator_values:?})",
7066 O::NAME
7067 );
7068 },
7069 );
7070 }
7071
7072 #[test]
7073 fn assert_rows_iterator_count() {
7074 struct CountOp;
7075
7076 impl RowsIteratorOp for CountOp {
7077 type Output<'a> = usize;
7078 const NAME: &'static str = "count";
7079
7080 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(
7081 iter: T,
7082 ) -> Self::Output<'a> {
7083 iter.count()
7084 }
7085 }
7086
7087 assert_rows_iterator_cases::<CountOp>()
7088 }
7089
7090 #[test]
7091 fn assert_rows_iterator_last() {
7092 struct LastOp;
7093
7094 impl RowsIteratorOp for LastOp {
7095 type Output<'a> = Option<Row<'a>>;
7096 const NAME: &'static str = "last";
7097
7098 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(
7099 iter: T,
7100 ) -> Self::Output<'a> {
7101 iter.last()
7102 }
7103 }
7104
7105 assert_rows_iterator_cases::<LastOp>()
7106 }
7107
7108 #[test]
7109 fn assert_rows_iterator_nth_0() {
7110 struct NthOp<const BACK: bool>;
7111
7112 impl<const BACK: bool> RowsIteratorOp for NthOp<BACK> {
7113 type Output<'a> = Option<Row<'a>>;
7114 const NAME: &'static str = if BACK { "nth_back(0)" } else { "nth(0)" };
7115
7116 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(
7117 mut iter: T,
7118 ) -> Self::Output<'a> {
7119 if BACK { iter.nth_back(0) } else { iter.nth(0) }
7120 }
7121 }
7122
7123 assert_rows_iterator_cases::<NthOp<false>>();
7124 assert_rows_iterator_cases::<NthOp<true>>();
7125 }
7126
7127 #[test]
7128 fn assert_rows_iterator_nth_1() {
7129 struct NthOp<const BACK: bool>;
7130
7131 impl<const BACK: bool> RowsIteratorOp for NthOp<BACK> {
7132 type Output<'a> = Option<Row<'a>>;
7133 const NAME: &'static str = if BACK { "nth_back(1)" } else { "nth(1)" };
7134
7135 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(
7136 mut iter: T,
7137 ) -> Self::Output<'a> {
7138 if BACK { iter.nth_back(1) } else { iter.nth(1) }
7139 }
7140 }
7141
7142 assert_rows_iterator_cases::<NthOp<false>>();
7143 assert_rows_iterator_cases::<NthOp<true>>();
7144 }
7145
7146 #[test]
7147 fn assert_rows_iterator_nth_after_end() {
7148 struct NthOp<const BACK: bool>;
7149
7150 impl<const BACK: bool> RowsIteratorOp for NthOp<BACK> {
7151 type Output<'a> = Option<Row<'a>>;
7152 const NAME: &'static str = if BACK {
7153 "nth_back(iter.len() + 1)"
7154 } else {
7155 "nth(iter.len() + 1)"
7156 };
7157
7158 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(
7159 mut iter: T,
7160 ) -> Self::Output<'a> {
7161 if BACK {
7162 iter.nth_back(iter.len() + 1)
7163 } else {
7164 iter.nth(iter.len() + 1)
7165 }
7166 }
7167 }
7168
7169 assert_rows_iterator_cases::<NthOp<false>>();
7170 assert_rows_iterator_cases::<NthOp<true>>();
7171 }
7172
7173 #[test]
7174 fn assert_rows_iterator_nth_len() {
7175 struct NthOp<const BACK: bool>;
7176
7177 impl<const BACK: bool> RowsIteratorOp for NthOp<BACK> {
7178 type Output<'a> = Option<Row<'a>>;
7179 const NAME: &'static str = if BACK {
7180 "nth_back(iter.len())"
7181 } else {
7182 "nth(iter.len())"
7183 };
7184
7185 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(
7186 mut iter: T,
7187 ) -> Self::Output<'a> {
7188 if BACK {
7189 iter.nth_back(iter.len())
7190 } else {
7191 iter.nth(iter.len())
7192 }
7193 }
7194 }
7195
7196 assert_rows_iterator_cases::<NthOp<false>>();
7197 assert_rows_iterator_cases::<NthOp<true>>();
7198 }
7199
7200 #[test]
7201 fn assert_rows_iterator_nth_last() {
7202 struct NthOp<const BACK: bool>;
7203
7204 impl<const BACK: bool> RowsIteratorOp for NthOp<BACK> {
7205 type Output<'a> = Option<Row<'a>>;
7206 const NAME: &'static str = if BACK {
7207 "nth_back(iter.len().saturating_sub(1))"
7208 } else {
7209 "nth(iter.len().saturating_sub(1))"
7210 };
7211
7212 fn get_value<'a, T: SharedBetweenRowsIteratorAndSliceIter<'a>>(
7213 mut iter: T,
7214 ) -> Self::Output<'a> {
7215 if BACK {
7216 iter.nth_back(iter.len().saturating_sub(1))
7217 } else {
7218 iter.nth(iter.len().saturating_sub(1))
7219 }
7220 }
7221 }
7222
7223 assert_rows_iterator_cases::<NthOp<false>>();
7224 assert_rows_iterator_cases::<NthOp<true>>();
7225 }
7226
7227 #[test]
7228 fn assert_rows_iterator_nth_and_reuse() {
7229 setup_and_assert(
7230 |_| {},
7231 |actual, expected| {
7232 {
7233 let mut actual = actual.clone();
7234 let mut expected = expected.clone();
7235 for _ in 0..expected.len() {
7236 #[expect(clippy::iter_nth_zero)]
7237 let actual_val = actual.nth(0);
7238 #[expect(clippy::iter_nth_zero)]
7239 let expected_val = expected.nth(0);
7240 assert_eq!(actual_val, expected_val, "Failed on nth(0)");
7241 }
7242 }
7243
7244 {
7245 let mut actual = actual.clone();
7246 let mut expected = expected.clone();
7247 for _ in 0..expected.len() {
7248 let actual_val = actual.nth(1);
7249 let expected_val = expected.nth(1);
7250 assert_eq!(actual_val, expected_val, "Failed on nth(1)");
7251 }
7252 }
7253
7254 {
7255 let mut actual = actual.clone();
7256 let mut expected = expected.clone();
7257 for _ in 0..expected.len() {
7258 let actual_val = actual.nth(2);
7259 let expected_val = expected.nth(2);
7260 assert_eq!(actual_val, expected_val, "Failed on nth(2)");
7261 }
7262 }
7263 },
7264 );
7265 }
7266
7267 #[test]
7268 fn assert_rows_iterator_nth_back_and_reuse() {
7269 setup_and_assert(
7270 |_| {},
7271 |actual, expected| {
7272 {
7273 let mut actual = actual.clone();
7274 let mut expected = expected.clone();
7275 for _ in 0..expected.len() {
7276 let actual_val = actual.nth_back(0);
7277 let expected_val = expected.nth_back(0);
7278 assert_eq!(actual_val, expected_val, "Failed on nth_back(0)");
7279 }
7280 }
7281
7282 {
7283 let mut actual = actual.clone();
7284 let mut expected = expected.clone();
7285 for _ in 0..expected.len() {
7286 let actual_val = actual.nth_back(1);
7287 let expected_val = expected.nth_back(1);
7288 assert_eq!(actual_val, expected_val, "Failed on nth_back(1)");
7289 }
7290 }
7291
7292 {
7293 let mut actual = actual.clone();
7294 let mut expected = expected.clone();
7295 for _ in 0..expected.len() {
7296 let actual_val = actual.nth_back(2);
7297 let expected_val = expected.nth_back(2);
7298 assert_eq!(actual_val, expected_val, "Failed on nth_back(2)");
7299 }
7300 }
7301 },
7302 );
7303 }
7304}