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