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("values", DataType::Int32, false));
4410 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4411 .with_values_field(value_field);
4412 builder.keys().append_value("a");
4414 builder.values().append_value(1);
4415 builder.keys().append_value("b");
4416 builder.values().append_value(2);
4417 builder.append(true).unwrap();
4418 builder.append(false).unwrap();
4420 builder.keys().append_value("c");
4422 builder.values().append_value(3);
4423 builder.append(true).unwrap();
4424 builder.append(true).unwrap();
4426 builder.keys().append_value("masked"); builder.values().append_value(0); builder.append(false).unwrap();
4430
4431 let map = Arc::new(builder.finish()) as ArrayRef;
4432 let d = map.data_type().clone();
4433
4434 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4435
4436 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4437
4438 let back = converter.convert_rows(&rows).unwrap();
4439 assert_eq!(back.len(), 1);
4440 back[0].to_data().validate_full().unwrap();
4441 assert_eq!(&back[0], &map);
4442 }
4443
4444 #[test]
4445 fn test_single_map_with_non_nullable_map_but_with_nullable_values() {
4446 let value_field = Arc::new(Field::new("values", DataType::Int32, true));
4448 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4449 .with_values_field(value_field);
4450
4451 builder.keys().append_value("a");
4453 builder.values().append_value(1);
4454 builder.keys().append_value("b");
4455 builder.values().append_null();
4456 builder.append(true).unwrap();
4457 builder.keys().append_value("c");
4459 builder.values().append_null();
4460 builder.keys().append_value("d");
4461 builder.values().append_null();
4462 builder.append(true).unwrap();
4463 builder.append(true).unwrap();
4465 builder.keys().append_value("e");
4467 builder.values().append_value(5);
4468 builder.append(true).unwrap();
4469
4470 let map = Arc::new(builder.finish()) as ArrayRef;
4471 let d = map.data_type().clone();
4472
4473 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4474
4475 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4476
4477 let back = converter.convert_rows(&rows).unwrap();
4478 assert_eq!(back.len(), 1);
4479 back[0].to_data().validate_full().unwrap();
4480 assert_eq!(&back[0], &map);
4481 }
4482
4483 #[test]
4484 fn test_map_all_nulls() {
4485 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4486 builder.keys().append_value("m1"); builder.values().append_value(1); builder.append(false).unwrap();
4490 builder.keys().append_value("m2"); builder.values().append_value(2); builder.append(false).unwrap();
4493
4494 builder.append(false).unwrap(); builder.keys().append_value("m3"); builder.values().append_value(3); builder.append(false).unwrap();
4499
4500 let map = Arc::new(builder.finish()) as ArrayRef;
4501 let d = map.data_type().clone();
4502
4503 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4504
4505 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4506
4507 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4509
4510 let back = converter.convert_rows(&rows).unwrap();
4511 assert_eq!(back.len(), 1);
4512 back[0].to_data().validate_full().unwrap();
4513 assert_eq!(&back[0], &map);
4514 }
4515
4516 #[test]
4517 fn test_map_all_empty() {
4518 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4519 builder.append(true).unwrap();
4521 builder.append(true).unwrap();
4522 builder.append(true).unwrap();
4523
4524 let map = Arc::new(builder.finish()) as ArrayRef;
4525 let d = map.data_type().clone();
4526
4527 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4528
4529 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4530
4531 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4533
4534 let back = converter.convert_rows(&rows).unwrap();
4535 assert_eq!(back.len(), 1);
4536 back[0].to_data().validate_full().unwrap();
4537 assert_eq!(&back[0], &map);
4538 }
4539
4540 #[test]
4541 fn test_map_empty_array() {
4542 let builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4544 let map = Arc::new(builder.finish_cloned()) as ArrayRef;
4545 let d = map.data_type().clone();
4546
4547 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4548
4549 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4550
4551 let back = converter.convert_rows(&rows).unwrap();
4552 assert_eq!(back.len(), 1);
4553 back[0].to_data().validate_full().unwrap();
4554 assert_eq!(&back[0], &map);
4555 }
4556
4557 fn generate_primitive_array<K>(
4558 rng: &mut impl RngCore,
4559 len: usize,
4560 valid_percent: f64,
4561 ) -> PrimitiveArray<K>
4562 where
4563 K: ArrowPrimitiveType,
4564 StandardUniform: Distribution<K::Native>,
4565 {
4566 (0..len)
4567 .map(|_| rng.random_bool(valid_percent).then(|| rng.random()))
4568 .collect()
4569 }
4570
4571 fn generate_all_unique_primitive_array<K>(
4572 rng: &mut impl RngCore,
4573 len: usize,
4574 ) -> PrimitiveArray<K>
4575 where
4576 K: ArrowPrimitiveType,
4577 K::Native: Hash + Eq,
4578 StandardUniform: Distribution<K::Native>,
4579 {
4580 let possible_number_of_values = 2i32.saturating_pow(size_of::<K::Native>() as u32 * 8);
4581 assert!(
4582 len <= possible_number_of_values as usize,
4583 "len {len} is larger than the number of possible values {possible_number_of_values}"
4584 );
4585
4586 let mut seen = std::collections::HashSet::new();
4587 (0..len)
4588 .map(|_| {
4589 let mut value;
4590 loop {
4591 value = rng.random();
4592
4593 if seen.insert(value) {
4594 break;
4595 }
4596 }
4597
4598 Some(value)
4599 })
4600 .collect()
4601 }
4602
4603 fn generate_boolean_array(
4604 rng: &mut impl RngCore,
4605 len: usize,
4606 valid_percent: f64,
4607 ) -> BooleanArray {
4608 (0..len)
4609 .map(|_| rng.random_bool(valid_percent).then(|| rng.random_bool(0.5)))
4610 .collect()
4611 }
4612
4613 fn generate_strings<O: OffsetSizeTrait>(
4614 rng: &mut impl RngCore,
4615 len: usize,
4616 valid_percent: f64,
4617 ) -> GenericStringArray<O> {
4618 (0..len)
4619 .map(|_| {
4620 rng.random_bool(valid_percent).then(|| {
4621 let len = rng.random_range(0..100);
4622 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4623 String::from_utf8(bytes).unwrap()
4624 })
4625 })
4626 .collect()
4627 }
4628
4629 fn generate_string_view(
4630 rng: &mut impl RngCore,
4631 len: usize,
4632 valid_percent: f64,
4633 ) -> StringViewArray {
4634 (0..len)
4635 .map(|_| {
4636 rng.random_bool(valid_percent).then(|| {
4637 let len = rng.random_range(0..100);
4638 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4639 String::from_utf8(bytes).unwrap()
4640 })
4641 })
4642 .collect()
4643 }
4644
4645 fn generate_byte_view(
4646 rng: &mut impl RngCore,
4647 len: usize,
4648 valid_percent: f64,
4649 ) -> BinaryViewArray {
4650 (0..len)
4651 .map(|_| {
4652 rng.random_bool(valid_percent).then(|| {
4653 let len = rng.random_range(0..100);
4654 let bytes: Vec<_> = (0..len).map(|_| rng.random_range(0..128)).collect();
4655 bytes
4656 })
4657 })
4658 .collect()
4659 }
4660
4661 fn generate_fixed_stringview_column(len: usize) -> StringViewArray {
4662 let edge_cases = vec![
4663 Some("bar".to_string()),
4664 Some("bar\0".to_string()),
4665 Some("LongerThan12Bytes".to_string()),
4666 Some("LongerThan12Bytez".to_string()),
4667 Some("LongerThan12Bytes\0".to_string()),
4668 Some("LongerThan12Byt".to_string()),
4669 Some("backend one".to_string()),
4670 Some("backend two".to_string()),
4671 Some("a".repeat(257)),
4672 Some("a".repeat(300)),
4673 ];
4674
4675 let mut values = Vec::with_capacity(len);
4677 for i in 0..len {
4678 values.push(
4679 edge_cases
4680 .get(i % edge_cases.len())
4681 .cloned()
4682 .unwrap_or(None),
4683 );
4684 }
4685
4686 StringViewArray::from(values)
4687 }
4688
4689 fn generate_dictionary<K>(
4690 rng: &mut impl RngCore,
4691 values: ArrayRef,
4692 len: usize,
4693 valid_percent: f64,
4694 ) -> DictionaryArray<K>
4695 where
4696 K: ArrowDictionaryKeyType,
4697 K::Native: SampleUniform,
4698 {
4699 let min_key = K::Native::from_usize(0).unwrap();
4700 let max_key = K::Native::from_usize(values.len()).unwrap();
4701 let keys: PrimitiveArray<K> = (0..len)
4702 .map(|_| {
4703 rng.random_bool(valid_percent)
4704 .then(|| rng.random_range(min_key..max_key))
4705 })
4706 .collect();
4707
4708 let data_type =
4709 DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
4710
4711 let data = keys
4712 .into_data()
4713 .into_builder()
4714 .data_type(data_type)
4715 .add_child_data(values.to_data())
4716 .build()
4717 .unwrap();
4718
4719 DictionaryArray::from(data)
4720 }
4721
4722 fn generate_fixed_size_binary(
4723 rng: &mut impl RngCore,
4724 len: usize,
4725 valid_percent: f64,
4726 ) -> FixedSizeBinaryArray {
4727 let width = rng.random_range(0..20);
4728 let mut builder = FixedSizeBinaryBuilder::new(width);
4729
4730 let mut b = vec![0; width as usize];
4731 for _ in 0..len {
4732 match rng.random_bool(valid_percent) {
4733 true => {
4734 b.iter_mut().for_each(|x| *x = rng.random());
4735 builder.append_value(&b).unwrap();
4736 }
4737 false => builder.append_null(),
4738 }
4739 }
4740
4741 builder.finish()
4742 }
4743
4744 fn generate_struct(rng: &mut impl RngCore, len: usize, valid_percent: f64) -> StructArray {
4745 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4746 let a = generate_primitive_array::<Int32Type>(rng, len, valid_percent);
4747 let b = generate_strings::<i32>(rng, len, valid_percent);
4748 let fields = Fields::from(vec![
4749 Field::new("a", DataType::Int32, true),
4750 Field::new("b", DataType::Utf8, true),
4751 ]);
4752 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
4753 StructArray::new(fields, values, Some(nulls))
4754 }
4755
4756 fn generate_list<R: RngCore, F>(
4757 rng: &mut R,
4758 len: usize,
4759 valid_percent: f64,
4760 values: F,
4761 ) -> ListArray
4762 where
4763 F: FnOnce(&mut R, usize) -> ArrayRef,
4764 {
4765 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4766 let values_len = offsets.last().unwrap().to_usize().unwrap();
4767 let values = values(rng, values_len);
4768 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4769 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4770 ListArray::new(field, offsets, values, Some(nulls))
4771 }
4772
4773 fn generate_list_view<F>(
4774 rng: &mut impl RngCore,
4775 len: usize,
4776 valid_percent: f64,
4777 values: F,
4778 ) -> ListViewArray
4779 where
4780 F: FnOnce(usize) -> ArrayRef,
4781 {
4782 let sizes: Vec<i32> = (0..len).map(|_| rng.random_range(0..10)).collect();
4784 let values_len: usize = sizes.iter().map(|s| *s as usize).sum::<usize>().max(1);
4785 let values = values(values_len);
4786
4787 let offsets: Vec<i32> = sizes
4789 .iter()
4790 .map(|&size| {
4791 if size == 0 {
4792 0
4793 } else {
4794 rng.random_range(0..=(values_len as i32 - size))
4795 }
4796 })
4797 .collect();
4798
4799 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4800 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4801 ListViewArray::new(
4802 field,
4803 ScalarBuffer::from(offsets),
4804 ScalarBuffer::from(sizes),
4805 values,
4806 Some(nulls),
4807 )
4808 }
4809
4810 fn generate_map<R: RngCore, KeysFn, ValuesFn>(
4811 rng: &mut R,
4812 len: usize,
4813 valid_percent: f64,
4814 gen_keys: KeysFn,
4815 gen_values: ValuesFn,
4816 ) -> MapArray
4817 where
4818 KeysFn: FnOnce(&mut R, usize) -> ArrayRef,
4819 ValuesFn: FnOnce(&mut R, usize) -> ArrayRef,
4820 {
4821 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4822 let entries_len = offsets.last().unwrap().to_usize().unwrap();
4823 let keys = gen_keys(rng, entries_len);
4824 let values = gen_values(rng, entries_len);
4825 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4826 let field = Arc::new(Field::new_map(
4827 "",
4828 "entries",
4829 Field::new("keys", keys.data_type().clone(), false),
4830 Field::new("values", values.data_type().clone(), true),
4831 false,
4832 true,
4833 ));
4834 let DataType::Map(struct_field, _) = field.data_type() else {
4835 unreachable!();
4836 };
4837
4838 let DataType::Struct(fields) = struct_field.data_type() else {
4839 unreachable!();
4840 };
4841
4842 let entries = StructArray::new(fields.clone(), vec![keys, values], None);
4843
4844 let map_array = MapArray::new(struct_field.clone(), offsets, entries, Some(nulls), false);
4845
4846 assert_valid_map(&map_array);
4847
4848 map_array
4849 }
4850
4851 fn assert_valid_map(array: &MapArray) {
4861 let keys_arrow_row_converter =
4862 RowConverter::new(vec![SortField::new(array.key_type().clone())]).unwrap();
4863
4864 array.iter().enumerate().flat_map(|(index, entry)| entry.map(|entry| (index, Arc::clone(entry.column(0))))).for_each(|(entry_index, keys)| {
4865 let keys_as_rows = keys_arrow_row_converter.convert_columns(&[Arc::clone(&keys)]).expect("should be able to convert keys");
4866
4867 for i in 0..keys_as_rows.num_rows() {
4868 for j in (i + 1)..keys_as_rows.num_rows() {
4869 if keys_as_rows.row(i) == keys_as_rows.row(j) {
4870 let key_i = keys.slice(i, 1);
4871 let key_j = keys.slice(j, 1);
4872
4873 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:?}");
4874 }
4875 }
4876 }
4877 })
4878 }
4879
4880 fn generate_nulls(rng: &mut impl RngCore, len: usize) -> Option<NullBuffer> {
4881 Some(NullBuffer::from_iter(
4882 (0..len).map(|_| rng.random_bool(0.8)),
4883 ))
4884 }
4885
4886 fn change_underlying_null_values_for_primitive<T: ArrowPrimitiveType>(
4887 array: &PrimitiveArray<T>,
4888 ) -> PrimitiveArray<T> {
4889 let (dt, values, nulls) = array.clone().into_parts();
4890
4891 let new_values = ScalarBuffer::<T::Native>::from_iter(
4892 values
4893 .iter()
4894 .zip(nulls.as_ref().unwrap().iter())
4895 .map(|(val, is_valid)| {
4896 if is_valid {
4897 *val
4898 } else {
4899 val.add_wrapping(T::Native::usize_as(1))
4900 }
4901 }),
4902 );
4903
4904 PrimitiveArray::new(new_values, nulls).with_data_type(dt)
4905 }
4906
4907 fn change_underline_null_values_for_byte_array<T: ByteArrayType>(
4908 array: &GenericByteArray<T>,
4909 ) -> GenericByteArray<T> {
4910 let (offsets, values, nulls) = array.clone().into_parts();
4911
4912 let new_offsets = OffsetBuffer::<T::Offset>::from_lengths(
4913 offsets
4914 .lengths()
4915 .zip(nulls.as_ref().unwrap().iter())
4916 .map(|(len, is_valid)| if is_valid { len } else { len + 1 }),
4917 );
4918
4919 let mut new_bytes = Vec::<u8>::with_capacity(new_offsets[new_offsets.len() - 1].as_usize());
4920
4921 offsets
4922 .windows(2)
4923 .zip(nulls.as_ref().unwrap().iter())
4924 .for_each(|(start_and_end, is_valid)| {
4925 let start = start_and_end[0].as_usize();
4926 let end = start_and_end[1].as_usize();
4927 new_bytes.extend_from_slice(&values.as_slice()[start..end]);
4928
4929 if !is_valid {
4931 new_bytes.push(b'c');
4932 }
4933 });
4934
4935 GenericByteArray::<T>::new(new_offsets, Buffer::from_vec(new_bytes), nulls)
4936 }
4937
4938 fn change_underline_null_values_for_list_array<O: OffsetSizeTrait>(
4939 array: &GenericListArray<O>,
4940 ) -> GenericListArray<O> {
4941 let (field, offsets, values, nulls) = array.clone().into_parts();
4942
4943 let (new_values, new_offsets) = {
4944 let concat_values = offsets
4945 .windows(2)
4946 .zip(nulls.as_ref().unwrap().iter())
4947 .map(|(start_and_end, is_valid)| {
4948 let start = start_and_end[0].as_usize();
4949 let end = start_and_end[1].as_usize();
4950 if is_valid {
4951 return (start, end - start);
4952 }
4953
4954 if end == values.len() {
4956 (start, (end - start).saturating_sub(1))
4957 } else {
4958 (start, end - start + 1)
4959 }
4960 })
4961 .map(|(start, length)| values.slice(start, length))
4962 .collect::<Vec<_>>();
4963
4964 let new_offsets =
4965 OffsetBuffer::<O>::from_lengths(concat_values.iter().map(|s| s.len()));
4966
4967 let new_values = {
4968 let values = concat_values.iter().map(|a| a.as_ref()).collect::<Vec<_>>();
4969 arrow_select::concat::concat(&values).expect("should be able to concat")
4970 };
4971
4972 (new_values, new_offsets)
4973 };
4974
4975 GenericListArray::<O>::new(field, new_offsets, new_values, nulls)
4976 }
4977
4978 fn change_underline_null_values_for_map_array(array: &MapArray) -> MapArray {
4979 let (field, offsets, entries, nulls, ordered) = array.clone().into_parts();
4980 assert!(
4981 !ordered,
4982 "can't replace underlying null values for ordered map array as this can violate the ordering"
4983 );
4984
4985 let (new_entries, new_offsets) = {
4986 let concat_values = offsets
4987 .windows(2)
4988 .zip(nulls.as_ref().unwrap().iter())
4989 .map(|(start_and_end, is_valid)| {
4990 let start = start_and_end[0].as_usize();
4991 let end = start_and_end[1].as_usize();
4992 if is_valid {
4993 return (start, end - start);
4994 }
4995
4996 if end == entries.len() {
4998 (start, (end - start).saturating_sub(1))
4999 } else {
5000 (start, end - start + 1)
5002 }
5003 })
5004 .map(|(start, length)| entries.slice(start, length))
5005 .collect::<Vec<_>>();
5006
5007 let new_offsets = OffsetBuffer::from_lengths(concat_values.iter().map(|s| s.len()));
5008
5009 let new_values = {
5010 let values = concat_values
5011 .iter()
5012 .map(|a| a as &dyn Array)
5013 .collect::<Vec<_>>();
5014 arrow_select::concat::concat(&values).expect("should be able to concat")
5015 };
5016
5017 (new_values.as_struct().clone(), new_offsets)
5018 };
5019
5020 let new_map = MapArray::new(field, new_offsets, new_entries, nulls, ordered);
5021
5022 assert_valid_map(&new_map);
5023
5024 new_map
5025 }
5026
5027 fn change_underline_null_values(array: &ArrayRef) -> ArrayRef {
5028 if array.null_count() == 0 {
5029 return Arc::clone(array);
5030 }
5031
5032 downcast_primitive_array!(
5033 array => {
5034 let output = change_underlying_null_values_for_primitive(array);
5035
5036 Arc::new(output)
5037 }
5038
5039 DataType::Utf8 => {
5040 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i32>()))
5041 }
5042 DataType::LargeUtf8 => {
5043 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i64>()))
5044 }
5045 DataType::Binary => {
5046 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i32>()))
5047 }
5048 DataType::LargeBinary => {
5049 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i64>()))
5050 }
5051 DataType::List(_) => {
5052 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i32>()))
5053 }
5054 DataType::LargeList(_) => {
5055 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i64>()))
5056 }
5057 DataType::Map(_, _) => {
5058 Arc::new(change_underline_null_values_for_map_array(array.as_map()))
5059 }
5060 _ => {
5061 Arc::clone(array)
5062 }
5063 )
5064 }
5065
5066 fn generate_column(rng: &mut (impl RngCore + Clone), len: usize) -> ArrayRef {
5067 match rng.random_range(0..24) {
5068 0 => Arc::new(generate_primitive_array::<Int32Type>(rng, len, 0.8)),
5069 1 => Arc::new(generate_primitive_array::<UInt32Type>(rng, len, 0.8)),
5070 2 => Arc::new(generate_primitive_array::<Int64Type>(rng, len, 0.8)),
5071 3 => Arc::new(generate_primitive_array::<UInt64Type>(rng, len, 0.8)),
5072 4 => Arc::new(generate_primitive_array::<Float32Type>(rng, len, 0.8)),
5073 5 => Arc::new(generate_primitive_array::<Float64Type>(rng, len, 0.8)),
5074 6 => Arc::new(generate_strings::<i32>(rng, len, 0.8)),
5075 7 => {
5076 let dict_values_len = rng.random_range(1..len);
5077 let strings = Arc::new(generate_strings::<i32>(rng, dict_values_len, 1.0));
5079 Arc::new(generate_dictionary::<Int64Type>(rng, strings, len, 0.8))
5080 }
5081 8 => {
5082 let dict_values_len = rng.random_range(1..len);
5083 let values = Arc::new(generate_primitive_array::<Int64Type>(
5085 rng,
5086 dict_values_len,
5087 1.0,
5088 ));
5089 Arc::new(generate_dictionary::<Int64Type>(rng, values, len, 0.8))
5090 }
5091 9 => Arc::new(generate_fixed_size_binary(rng, len, 0.8)),
5092 10 => Arc::new(generate_struct(rng, len, 0.8)),
5093 11 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5094 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5095 })),
5096 12 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5097 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5098 })),
5099 13 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5100 Arc::new(generate_struct(rng, values_len, 0.8))
5101 })),
5102 14 => Arc::new(generate_string_view(rng, len, 0.8)),
5103 15 => Arc::new(generate_byte_view(rng, len, 0.8)),
5104 16 => Arc::new(generate_fixed_stringview_column(len)),
5105 17 => Arc::new(
5106 generate_list(&mut rng.clone(), len + 1000, 0.8, |rng, values_len| {
5107 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5108 })
5109 .slice(500, len),
5110 ),
5111 18 => Arc::new(generate_boolean_array(rng, len, 0.8)),
5112 19 => Arc::new(generate_list_view(
5113 &mut rng.clone(),
5114 len,
5115 0.8,
5116 |values_len| Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8)),
5117 )),
5118 20 => Arc::new(generate_list_view(
5119 &mut rng.clone(),
5120 len,
5121 0.8,
5122 |values_len| Arc::new(generate_strings::<i32>(rng, values_len, 0.8)),
5123 )),
5124 21 => Arc::new(generate_list_view(
5125 &mut rng.clone(),
5126 len,
5127 0.8,
5128 |values_len| Arc::new(generate_struct(rng, values_len, 0.8)),
5129 )),
5130 22 => Arc::new(
5131 generate_list_view(&mut rng.clone(), len + 1000, 0.8, |values_len| {
5132 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5133 })
5134 .slice(500, len),
5135 ),
5136 23 => Arc::new(generate_map(
5137 rng,
5138 len,
5139 0.9,
5140 |rng, keys_len| {
5142 Arc::new(generate_all_unique_primitive_array::<Int64Type>(
5143 rng, keys_len,
5144 ))
5145 },
5146 |rng, values_len| Arc::new(generate_strings::<i32>(rng, values_len, 0.7)),
5147 )),
5148 _ => unreachable!(),
5149 }
5150 }
5151
5152 fn print_row(cols: &[SortColumn], row: usize) -> String {
5153 let t: Vec<_> = cols
5154 .iter()
5155 .map(|x| match x.values.is_valid(row) {
5156 true => {
5157 let opts = FormatOptions::default().with_null("NULL");
5158 let formatter = ArrayFormatter::try_new(x.values.as_ref(), &opts).unwrap();
5159 formatter.value(row).to_string()
5160 }
5161 false => "NULL".to_string(),
5162 })
5163 .collect();
5164 t.join(",")
5165 }
5166
5167 fn print_col_types(cols: &[SortColumn]) -> String {
5168 let t: Vec<_> = cols
5169 .iter()
5170 .map(|x| x.values.data_type().to_string())
5171 .collect();
5172 t.join(",")
5173 }
5174
5175 #[derive(Debug, PartialEq)]
5176 enum Nulls {
5177 AsIs,
5179
5180 Different,
5182
5183 None,
5185 }
5186
5187 #[test]
5188 #[cfg_attr(miri, ignore)]
5189 fn fuzz_test() {
5190 let mut rng = StdRng::seed_from_u64(42);
5191 for _ in 0..100 {
5192 for null_behavior in [Nulls::AsIs, Nulls::Different, Nulls::None] {
5193 let num_columns = rng.random_range(1..5);
5194 let len = rng.random_range(5..100);
5195 let mut arrays: Vec<_> = (0..num_columns)
5196 .map(|_| generate_column(&mut rng, len))
5197 .collect();
5198
5199 match null_behavior {
5200 Nulls::AsIs => {
5201 }
5203 Nulls::Different => {
5204 arrays = arrays
5206 .into_iter()
5207 .map(|a| replace_array_nulls(a, generate_nulls(&mut rng, len)))
5208 .collect()
5209 }
5210 Nulls::None => {
5211 arrays = arrays
5213 .into_iter()
5214 .map(|a| replace_array_nulls(a, None))
5215 .collect()
5216 }
5217 }
5218
5219 let options: Vec<_> = (0..num_columns)
5220 .map(|_| SortOptions {
5221 descending: rng.random_bool(0.5),
5222 nulls_first: rng.random_bool(0.5),
5223 })
5224 .collect();
5225
5226 let sort_columns: Vec<_> = options
5227 .iter()
5228 .zip(&arrays)
5229 .map(|(o, c)| SortColumn {
5230 values: Arc::clone(c),
5231 options: Some(*o),
5232 })
5233 .collect();
5234
5235 let comparator = LexicographicalComparator::try_new(&sort_columns).unwrap();
5236
5237 let columns: Vec<SortField> = options
5238 .into_iter()
5239 .zip(&arrays)
5240 .map(|(o, a)| SortField::new_with_options(a.data_type().clone(), o))
5241 .collect();
5242
5243 let converter = RowConverter::new(columns).unwrap();
5244 let rows = converter.convert_columns(&arrays).unwrap();
5245
5246 if !matches!(null_behavior, Nulls::None) {
5249 assert_same_rows_when_changing_input_underlying_null_values(
5250 &arrays, &converter, &rows,
5251 );
5252 }
5253
5254 for i in 0..len {
5255 for j in 0..len {
5256 let row_i = rows.row(i);
5257 let row_j = rows.row(j);
5258 let row_cmp = row_i.cmp(&row_j);
5259 let lex_cmp = comparator.compare(i, j);
5260 assert_eq!(
5261 row_cmp,
5262 lex_cmp,
5263 "({:?} vs {:?}) vs ({:?} vs {:?}) for types {}",
5264 print_row(&sort_columns, i),
5265 print_row(&sort_columns, j),
5266 row_i,
5267 row_j,
5268 print_col_types(&sort_columns)
5269 );
5270 }
5271 }
5272
5273 {
5275 let mut rows_iter = rows.iter();
5276 let mut rows_lengths_iter = rows.lengths();
5277 for (index, row) in rows_iter.by_ref().enumerate() {
5278 let len = rows_lengths_iter
5279 .next()
5280 .expect("Reached end of length iterator while still have rows");
5281 assert_eq!(
5282 row.data.len(),
5283 len,
5284 "Row length mismatch: {} vs {}",
5285 row.data.len(),
5286 len
5287 );
5288 assert_eq!(
5289 len,
5290 rows.row_len(index),
5291 "Row length mismatch at index {}: {} vs {}",
5292 index,
5293 len,
5294 rows.row_len(index)
5295 );
5296 }
5297
5298 assert_eq!(
5299 rows_lengths_iter.next(),
5300 None,
5301 "Length iterator did not reach end"
5302 );
5303 }
5304
5305 let back = converter.convert_rows(&rows).unwrap();
5308 for (actual, expected) in back.iter().zip(&arrays) {
5309 actual.to_data().validate_full().unwrap();
5310 dictionary_eq(actual, expected)
5311 }
5312
5313 let rows = rows.try_into_binary().expect("reasonable size");
5316 let parser = converter.parser();
5317 let back = converter
5318 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5319 .unwrap();
5320 for (actual, expected) in back.iter().zip(&arrays) {
5321 actual.to_data().validate_full().unwrap();
5322 dictionary_eq(actual, expected)
5323 }
5324
5325 let rows = converter.from_binary(rows);
5326 let back = converter.convert_rows(&rows).unwrap();
5327 for (actual, expected) in back.iter().zip(&arrays) {
5328 actual.to_data().validate_full().unwrap();
5329 dictionary_eq(actual, expected)
5330 }
5331 }
5332 }
5333 }
5334
5335 fn replace_array_nulls(array: ArrayRef, new_nulls: Option<NullBuffer>) -> ArrayRef {
5336 make_array(
5337 array
5338 .into_data()
5339 .into_builder()
5340 .nulls(new_nulls)
5342 .build()
5343 .unwrap(),
5344 )
5345 }
5346
5347 fn assert_same_rows_when_changing_input_underlying_null_values(
5348 arrays: &[ArrayRef],
5349 converter: &RowConverter,
5350 rows: &Rows,
5351 ) {
5352 let arrays_with_different_data_behind_nulls = arrays
5353 .iter()
5354 .map(|arr| change_underline_null_values(arr))
5355 .collect::<Vec<_>>();
5356
5357 if arrays
5359 .iter()
5360 .zip(arrays_with_different_data_behind_nulls.iter())
5361 .all(|(a, b)| Arc::ptr_eq(a, b))
5362 {
5363 return;
5364 }
5365
5366 let rows_with_different_nulls = converter
5367 .convert_columns(&arrays_with_different_data_behind_nulls)
5368 .unwrap();
5369
5370 assert_eq!(
5371 rows.iter().collect::<Vec<_>>(),
5372 rows_with_different_nulls.iter().collect::<Vec<_>>(),
5373 "Different underlying nulls should not output different rows"
5374 )
5375 }
5376
5377 #[test]
5378 fn test_clear() {
5379 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
5380 let mut rows = converter.empty_rows(3, 128);
5381
5382 let first = Int32Array::from(vec![None, Some(2), Some(4)]);
5383 let second = Int32Array::from(vec![Some(2), None, Some(4)]);
5384 let arrays = [Arc::new(first) as ArrayRef, Arc::new(second) as ArrayRef];
5385
5386 for array in arrays.iter() {
5387 rows.clear();
5388 converter
5389 .append(&mut rows, std::slice::from_ref(array))
5390 .unwrap();
5391 let back = converter.convert_rows(&rows).unwrap();
5392 assert_eq!(&back[0], array);
5393 }
5394
5395 let mut rows_expected = converter.empty_rows(3, 128);
5396 converter.append(&mut rows_expected, &arrays[1..]).unwrap();
5397
5398 for (i, (actual, expected)) in rows.iter().zip(rows_expected.iter()).enumerate() {
5399 assert_eq!(
5400 actual, expected,
5401 "For row {i}: expected {expected:?}, actual: {actual:?}",
5402 );
5403 }
5404 }
5405
5406 #[test]
5407 fn test_append_codec_dictionary_binary() {
5408 use DataType::*;
5409 let converter = RowConverter::new(vec![SortField::new(Dictionary(
5411 Box::new(Int32),
5412 Box::new(Binary),
5413 ))])
5414 .unwrap();
5415 let mut rows = converter.empty_rows(4, 128);
5416
5417 let keys = Int32Array::from_iter_values([0, 1, 2, 3]);
5418 let values = BinaryArray::from(vec![
5419 Some("a".as_bytes()),
5420 Some(b"b"),
5421 Some(b"c"),
5422 Some(b"d"),
5423 ]);
5424 let dict_array = DictionaryArray::new(keys, Arc::new(values));
5425
5426 rows.clear();
5427 let array = Arc::new(dict_array) as ArrayRef;
5428 converter
5429 .append(&mut rows, std::slice::from_ref(&array))
5430 .unwrap();
5431 let back = converter.convert_rows(&rows).unwrap();
5432
5433 dictionary_eq(&back[0], &array);
5434 }
5435
5436 #[test]
5437 fn test_list_prefix() {
5438 let mut a = ListBuilder::new(Int8Builder::new());
5439 a.append_value([None]);
5440 a.append_value([None, None]);
5441 let a = a.finish();
5442
5443 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
5444 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
5445 assert_eq!(rows.row(0).cmp(&rows.row(1)), Ordering::Less);
5446 }
5447
5448 #[test]
5449 fn test_values_buffer_smaller_when_utf8_validation_disabled() {
5450 fn get_values_buffer_len(col: ArrayRef) -> (usize, usize) {
5451 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8View)]).unwrap();
5453
5454 let rows = converter.convert_columns(&[col]).unwrap();
5456 let converted = converter.convert_rows(&rows).unwrap();
5457 let unchecked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5458
5459 let rows = rows.try_into_binary().expect("reasonable size");
5461 let parser = converter.parser();
5462 let converted = converter
5463 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5464 .unwrap();
5465 let checked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5466 (unchecked_values_len, checked_values_len)
5467 }
5468
5469 let col = Arc::new(StringViewArray::from_iter([
5471 Some("hello"), None, Some("short"), Some("tiny"), ])) as ArrayRef;
5476
5477 let (unchecked_values_len, checked_values_len) = get_values_buffer_len(col);
5478 assert_eq!(unchecked_values_len, 0);
5480 assert_eq!(checked_values_len, 14);
5482
5483 let col = Arc::new(StringViewArray::from_iter([
5485 Some("this is a very long string over 12 bytes"),
5486 Some("another long string to test the buffer"),
5487 ])) as ArrayRef;
5488
5489 let (unchecked_values_len, checked_values_len) = get_values_buffer_len(col);
5490 assert!(unchecked_values_len > 0);
5492 assert_eq!(unchecked_values_len, checked_values_len);
5493
5494 let col = Arc::new(StringViewArray::from_iter([
5496 Some("tiny"), Some("thisisexact13"), None,
5499 Some("short"), ])) as ArrayRef;
5501
5502 let (unchecked_values_len, checked_values_len) = get_values_buffer_len(col);
5503 assert_eq!(unchecked_values_len, 13);
5505 assert!(checked_values_len > unchecked_values_len);
5506 }
5507
5508 #[test]
5509 fn test_sparse_union() {
5510 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5512 let str_array = StringArray::from(vec![None, Some("b"), None, Some("d"), None]);
5513
5514 let type_ids = vec![0, 1, 0, 1, 0].into();
5516
5517 let union_fields = [
5518 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5519 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5520 ]
5521 .into_iter()
5522 .collect();
5523
5524 let union_array = UnionArray::try_new(
5525 union_fields,
5526 type_ids,
5527 None,
5528 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5529 )
5530 .unwrap();
5531
5532 let union_type = union_array.data_type().clone();
5533 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5534
5535 let rows = converter
5536 .convert_columns(&[Arc::new(union_array.clone())])
5537 .unwrap();
5538
5539 let back = converter.convert_rows(&rows).unwrap();
5541 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5542
5543 assert_eq!(union_array.len(), back_union.len());
5544 for i in 0..union_array.len() {
5545 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5546 }
5547 }
5548
5549 #[test]
5550 fn test_sparse_union_with_nulls() {
5551 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5553 let str_array = StringArray::from(vec![None::<&str>; 5]);
5554
5555 let type_ids = vec![0, 1, 0, 1, 0].into();
5557
5558 let union_fields = [
5559 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5560 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5561 ]
5562 .into_iter()
5563 .collect();
5564
5565 let union_array = UnionArray::try_new(
5566 union_fields,
5567 type_ids,
5568 None,
5569 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5570 )
5571 .unwrap();
5572
5573 let union_type = union_array.data_type().clone();
5574 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5575
5576 let rows = converter
5577 .convert_columns(&[Arc::new(union_array.clone())])
5578 .unwrap();
5579
5580 let back = converter.convert_rows(&rows).unwrap();
5582 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5583
5584 assert_eq!(union_array.len(), back_union.len());
5585 for i in 0..union_array.len() {
5586 let expected_null = union_array.is_null(i);
5587 let actual_null = back_union.is_null(i);
5588 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5589 if !expected_null {
5590 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5591 }
5592 }
5593 }
5594
5595 #[test]
5596 fn test_dense_union() {
5597 let int_array = Int32Array::from(vec![1, 3, 5]);
5599 let str_array = StringArray::from(vec!["a", "b"]);
5600
5601 let type_ids = vec![0, 1, 0, 1, 0].into();
5602
5603 let offsets = vec![0, 0, 1, 1, 2].into();
5605
5606 let union_fields = [
5607 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5608 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5609 ]
5610 .into_iter()
5611 .collect();
5612
5613 let union_array = UnionArray::try_new(
5614 union_fields,
5615 type_ids,
5616 Some(offsets), vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5618 )
5619 .unwrap();
5620
5621 let union_type = union_array.data_type().clone();
5622 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5623
5624 let rows = converter
5625 .convert_columns(&[Arc::new(union_array.clone())])
5626 .unwrap();
5627
5628 let back = converter.convert_rows(&rows).unwrap();
5630 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5631
5632 assert_eq!(union_array.len(), back_union.len());
5633 for i in 0..union_array.len() {
5634 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5635 }
5636 }
5637
5638 #[test]
5639 fn test_dense_union_with_nulls() {
5640 let int_array = Int32Array::from(vec![Some(1), None, Some(5)]);
5642 let str_array = StringArray::from(vec![Some("a"), None]);
5643
5644 let type_ids = vec![0, 1, 0, 1, 0].into();
5646 let offsets = vec![0, 0, 1, 1, 2].into();
5647
5648 let union_fields = [
5649 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5650 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5651 ]
5652 .into_iter()
5653 .collect();
5654
5655 let union_array = UnionArray::try_new(
5656 union_fields,
5657 type_ids,
5658 Some(offsets),
5659 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5660 )
5661 .unwrap();
5662
5663 let union_type = union_array.data_type().clone();
5664 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5665
5666 let rows = converter
5667 .convert_columns(&[Arc::new(union_array.clone())])
5668 .unwrap();
5669
5670 let back = converter.convert_rows(&rows).unwrap();
5672 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5673
5674 assert_eq!(union_array.len(), back_union.len());
5675 for i in 0..union_array.len() {
5676 let expected_null = union_array.is_null(i);
5677 let actual_null = back_union.is_null(i);
5678 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5679 if !expected_null {
5680 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5681 }
5682 }
5683 }
5684
5685 #[test]
5686 fn test_union_ordering() {
5687 let int_array = Int32Array::from(vec![100, 5, 20]);
5688 let str_array = StringArray::from(vec!["z", "a"]);
5689
5690 let type_ids = vec![0, 1, 0, 1, 0].into();
5692 let offsets = vec![0, 0, 1, 1, 2].into();
5693
5694 let union_fields = [
5695 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5696 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5697 ]
5698 .into_iter()
5699 .collect();
5700
5701 let union_array = UnionArray::try_new(
5702 union_fields,
5703 type_ids,
5704 Some(offsets),
5705 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5706 )
5707 .unwrap();
5708
5709 let union_type = union_array.data_type().clone();
5710 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5711
5712 let rows = converter.convert_columns(&[Arc::new(union_array)]).unwrap();
5713
5714 assert!(rows.row(2) < rows.row(1));
5726
5727 assert!(rows.row(0) < rows.row(3));
5729
5730 assert!(rows.row(2) < rows.row(4));
5733 assert!(rows.row(4) < rows.row(0));
5735
5736 assert!(rows.row(3) < rows.row(1));
5739 }
5740
5741 #[test]
5742 fn test_row_converter_roundtrip_with_many_union_columns() {
5743 let fields1 = UnionFields::try_new(
5745 vec![0, 1],
5746 vec![
5747 Field::new("int", DataType::Int32, true),
5748 Field::new("string", DataType::Utf8, true),
5749 ],
5750 )
5751 .unwrap();
5752
5753 let int_array1 = Int32Array::from(vec![Some(67), None]);
5754 let string_array1 = StringArray::from(vec![None::<&str>, Some("hello")]);
5755 let type_ids1 = vec![0i8, 1].into();
5756
5757 let union_array1 = UnionArray::try_new(
5758 fields1.clone(),
5759 type_ids1,
5760 None,
5761 vec![
5762 Arc::new(int_array1) as ArrayRef,
5763 Arc::new(string_array1) as ArrayRef,
5764 ],
5765 )
5766 .unwrap();
5767
5768 let fields2 = UnionFields::try_new(
5770 vec![0, 1],
5771 vec![
5772 Field::new("int", DataType::Int32, true),
5773 Field::new("string", DataType::Utf8, true),
5774 ],
5775 )
5776 .unwrap();
5777
5778 let int_array2 = Int32Array::from(vec![Some(100), None]);
5779 let string_array2 = StringArray::from(vec![None::<&str>, Some("world")]);
5780 let type_ids2 = vec![0i8, 1].into();
5781
5782 let union_array2 = UnionArray::try_new(
5783 fields2.clone(),
5784 type_ids2,
5785 None,
5786 vec![
5787 Arc::new(int_array2) as ArrayRef,
5788 Arc::new(string_array2) as ArrayRef,
5789 ],
5790 )
5791 .unwrap();
5792
5793 let field1 = Field::new("col1", DataType::Union(fields1, UnionMode::Sparse), true);
5795 let field2 = Field::new("col2", DataType::Union(fields2, UnionMode::Sparse), true);
5796
5797 let sort_field1 = SortField::new(field1.data_type().clone());
5798 let sort_field2 = SortField::new(field2.data_type().clone());
5799
5800 let converter = RowConverter::new(vec![sort_field1, sort_field2]).unwrap();
5801
5802 let rows = converter
5803 .convert_columns(&[
5804 Arc::new(union_array1.clone()) as ArrayRef,
5805 Arc::new(union_array2.clone()) as ArrayRef,
5806 ])
5807 .unwrap();
5808
5809 let out = converter.convert_rows(&rows).unwrap();
5811
5812 let [col1, col2] = out.as_slice() else {
5813 panic!("expected 2 columns")
5814 };
5815
5816 let col1 = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5817 let col2 = col2.as_any().downcast_ref::<UnionArray>().unwrap();
5818
5819 for (expected, got) in [union_array1, union_array2].iter().zip([col1, col2]) {
5820 assert_eq!(expected.len(), got.len());
5821 assert_eq!(expected.type_ids(), got.type_ids());
5822
5823 for i in 0..expected.len() {
5824 assert_eq!(expected.value(i).as_ref(), got.value(i).as_ref());
5825 }
5826 }
5827 }
5828
5829 #[test]
5830 fn test_row_converter_roundtrip_with_one_union_column() {
5831 let fields = UnionFields::try_new(
5832 vec![0, 1],
5833 vec![
5834 Field::new("int", DataType::Int32, true),
5835 Field::new("string", DataType::Utf8, true),
5836 ],
5837 )
5838 .unwrap();
5839
5840 let int_array = Int32Array::from(vec![Some(67), None]);
5841 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5842 let type_ids = vec![0i8, 1].into();
5843
5844 let union_array = UnionArray::try_new(
5845 fields.clone(),
5846 type_ids,
5847 None,
5848 vec![
5849 Arc::new(int_array) as ArrayRef,
5850 Arc::new(string_array) as ArrayRef,
5851 ],
5852 )
5853 .unwrap();
5854
5855 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5856 let sort_field = SortField::new(field.data_type().clone());
5857 let converter = RowConverter::new(vec![sort_field]).unwrap();
5858
5859 let rows = converter
5860 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5861 .unwrap();
5862
5863 let out = converter.convert_rows(&rows).unwrap();
5865
5866 let [col1] = out.as_slice() else {
5867 panic!("expected 1 column")
5868 };
5869
5870 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5871 assert_eq!(col.len(), union_array.len());
5872 assert_eq!(col.type_ids(), union_array.type_ids());
5873
5874 for i in 0..col.len() {
5875 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5876 }
5877 }
5878
5879 #[test]
5880 fn test_row_converter_roundtrip_with_non_default_union_type_ids() {
5881 let fields = UnionFields::try_new(
5883 vec![70, 85],
5884 vec![
5885 Field::new("int", DataType::Int32, true),
5886 Field::new("string", DataType::Utf8, true),
5887 ],
5888 )
5889 .unwrap();
5890
5891 let int_array = Int32Array::from(vec![Some(67), None]);
5892 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5893 let type_ids = vec![70i8, 85].into();
5894
5895 let union_array = UnionArray::try_new(
5896 fields.clone(),
5897 type_ids,
5898 None,
5899 vec![
5900 Arc::new(int_array) as ArrayRef,
5901 Arc::new(string_array) as ArrayRef,
5902 ],
5903 )
5904 .unwrap();
5905
5906 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5907 let sort_field = SortField::new(field.data_type().clone());
5908 let converter = RowConverter::new(vec![sort_field]).unwrap();
5909
5910 let rows = converter
5911 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5912 .unwrap();
5913
5914 let out = converter.convert_rows(&rows).unwrap();
5916
5917 let [col1] = out.as_slice() else {
5918 panic!("expected 1 column")
5919 };
5920
5921 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5922 assert_eq!(col.len(), union_array.len());
5923 assert_eq!(col.type_ids(), union_array.type_ids());
5924
5925 for i in 0..col.len() {
5926 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5927 }
5928 }
5929
5930 #[test]
5931 fn rows_size_should_count_for_capacity() {
5932 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
5933
5934 let empty_rows_size_with_preallocate_rows_and_data = {
5935 let rows = row_converter.empty_rows(1000, 1000);
5936
5937 rows.size()
5938 };
5939 let empty_rows_size_with_preallocate_rows = {
5940 let rows = row_converter.empty_rows(1000, 0);
5941
5942 rows.size()
5943 };
5944 let empty_rows_size_with_preallocate_data = {
5945 let rows = row_converter.empty_rows(0, 1000);
5946
5947 rows.size()
5948 };
5949 let empty_rows_size_without_preallocate = {
5950 let rows = row_converter.empty_rows(0, 0);
5951
5952 rows.size()
5953 };
5954
5955 assert!(
5956 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_rows,
5957 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_rows}"
5958 );
5959 assert!(
5960 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_data,
5961 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_data}"
5962 );
5963 assert!(
5964 empty_rows_size_with_preallocate_rows > empty_rows_size_without_preallocate,
5965 "{empty_rows_size_with_preallocate_rows} should be larger than {empty_rows_size_without_preallocate}"
5966 );
5967 assert!(
5968 empty_rows_size_with_preallocate_data > empty_rows_size_without_preallocate,
5969 "{empty_rows_size_with_preallocate_data} should be larger than {empty_rows_size_without_preallocate}"
5970 );
5971 }
5972
5973 #[test]
5974 fn test_struct_no_child_fields() {
5975 fn run_test(array: ArrayRef) {
5976 let sort_fields = vec![SortField::new(array.data_type().clone())];
5977 let converter = RowConverter::new(sort_fields).unwrap();
5978 let r = converter.convert_columns(&[Arc::clone(&array)]).unwrap();
5979
5980 let back = converter.convert_rows(&r).unwrap();
5981 assert_eq!(back.len(), 1);
5982 assert_eq!(&back[0], &array);
5983 }
5984
5985 let s = Arc::new(StructArray::new_empty_fields(5, None)) as ArrayRef;
5986 run_test(s);
5987
5988 let s = Arc::new(StructArray::new_empty_fields(
5989 5,
5990 Some(vec![true, false, true, false, false].into()),
5991 )) as ArrayRef;
5992 run_test(s);
5993 }
5994
5995 #[test]
5996 fn reserve_should_increase_capacity_to_the_requested_size() {
5997 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
5998 let mut empty_rows = row_converter.empty_rows(0, 0);
5999 empty_rows.reserve(50, 50);
6000 let before_size = empty_rows.size();
6001 empty_rows.reserve(50, 50);
6002 assert_eq!(
6003 empty_rows.size(),
6004 before_size,
6005 "Size should not change when reserving already reserved space"
6006 );
6007 empty_rows.reserve(10, 20);
6008 assert_eq!(
6009 empty_rows.size(),
6010 before_size,
6011 "Size should not change when already have space for the expected reserved data"
6012 );
6013
6014 empty_rows.reserve(100, 20);
6015 assert!(
6016 empty_rows.size() > before_size,
6017 "Size should increase when reserving more space than previously reserved"
6018 );
6019
6020 let before_size = empty_rows.size();
6021
6022 empty_rows.reserve(20, 100);
6023 assert!(
6024 empty_rows.size() > before_size,
6025 "Size should increase when reserving more space than previously reserved"
6026 );
6027 }
6028
6029 #[test]
6030 fn empty_rows_should_return_empty_lengths_iterator() {
6031 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6032 .unwrap()
6033 .empty_rows(0, 0);
6034 let mut lengths_iter = rows.lengths();
6035 assert_eq!(lengths_iter.next(), None);
6036 }
6037
6038 #[test]
6039 #[should_panic(expected = "row index out of bounds")]
6040 fn row_should_panic_on_overflowing_index() {
6041 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6042 .unwrap()
6043 .empty_rows(0, 0);
6044 rows.row(usize::MAX);
6045 }
6046
6047 #[test]
6048 #[should_panic(expected = "row index out of bounds")]
6049 fn row_len_should_panic_on_overflowing_index() {
6050 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6051 .unwrap()
6052 .empty_rows(0, 0);
6053 rows.row_len(usize::MAX);
6054 }
6055
6056 #[test]
6057 fn test_nested_null_list() {
6058 let null_array = Arc::new(NullArray::new(3));
6059 let list: ArrayRef = Arc::new(ListArray::new(
6061 Field::new_list_field(DataType::Null, true).into(),
6062 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6063 null_array,
6064 None,
6065 ));
6066
6067 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6068 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6069 let back = converter.convert_rows(&rows).unwrap();
6070
6071 assert_eq!(&list, &back[0]);
6072 }
6073
6074 #[test]
6076 fn test_double_nested_null_list() {
6077 let null_array = Arc::new(NullArray::new(1));
6078 let nested_field = Arc::new(Field::new_list_field(DataType::Null, true));
6080 let nested_list = Arc::new(ListArray::new(
6081 nested_field.clone(),
6082 OffsetBuffer::from_lengths(vec![1]),
6083 null_array,
6084 None,
6085 ));
6086 let list = Arc::new(ListArray::new(
6088 Field::new_list_field(DataType::List(nested_field), true).into(),
6089 OffsetBuffer::from_lengths(vec![1]),
6090 nested_list,
6091 None,
6092 )) as ArrayRef;
6093
6094 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6095 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6096 let back = converter.convert_rows(&rows).unwrap();
6097
6098 assert_eq!(&list, &back[0]);
6099 }
6100
6101 #[test]
6103 fn test_large_list_null() {
6104 let null_array = Arc::new(NullArray::new(3));
6105 let list: ArrayRef = Arc::new(LargeListArray::new(
6107 Field::new_list_field(DataType::Null, true).into(),
6108 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6109 null_array,
6110 None,
6111 ));
6112
6113 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6114 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6115 let back = converter.convert_rows(&rows).unwrap();
6116
6117 assert_eq!(&list, &back[0]);
6118 }
6119
6120 #[test]
6122 fn test_fixed_size_list_null() {
6123 let null_array = Arc::new(NullArray::new(6));
6124 let list: ArrayRef = Arc::new(FixedSizeListArray::new(
6126 Arc::new(Field::new_list_field(DataType::Null, true)),
6127 2,
6128 null_array,
6129 None,
6130 ));
6131
6132 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6133 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6134 let back = converter.convert_rows(&rows).unwrap();
6135
6136 assert_eq!(&list, &back[0]);
6137 }
6138
6139 #[test]
6141 fn test_list_null_variations() {
6142 let null_array = Arc::new(NullArray::new(3));
6144 let list: ArrayRef = Arc::new(ListArray::new(
6145 Field::new_list_field(DataType::Null, true).into(),
6146 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6147 null_array,
6148 None,
6149 ));
6150
6151 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6152 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6153 let back = converter.convert_rows(&rows).unwrap();
6154 assert_eq!(&list, &back[0]);
6155
6156 let null_array = Arc::new(NullArray::new(3));
6158 let list: ArrayRef = Arc::new(ListArray::new(
6159 Field::new_list_field(DataType::Null, true).into(),
6160 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6161 null_array,
6162 Some(vec![true, false, true].into()),
6163 ));
6164
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(0));
6171 let list: ArrayRef = Arc::new(ListArray::new(
6172 Field::new_list_field(DataType::Null, true).into(),
6173 OffsetBuffer::from_lengths(vec![]),
6174 null_array,
6175 None,
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![0, 0, 0]),
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
6196 #[test]
6198 fn test_list_null_descending() {
6199 let null_array = Arc::new(NullArray::new(3));
6200 let list: ArrayRef = Arc::new(ListArray::new(
6202 Field::new_list_field(DataType::Null, true).into(),
6203 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6204 null_array,
6205 None,
6206 ));
6207
6208 let options = SortOptions::default().with_descending(true);
6209 let field = SortField::new_with_options(list.data_type().clone(), options);
6210 let converter = RowConverter::new(vec![field]).unwrap();
6211 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6212 let back = converter.convert_rows(&rows).unwrap();
6213
6214 assert_eq!(&list, &back[0]);
6215 }
6216
6217 #[test]
6219 fn test_struct_with_null_field() {
6220 let null_array = Arc::new(NullArray::new(3));
6222 let int_array = Arc::new(Int32Array::from(vec![1, 2, 3]));
6223
6224 let struct_array: ArrayRef = Arc::new(StructArray::new(
6225 vec![
6226 Arc::new(Field::new("a", DataType::Null, true)),
6227 Arc::new(Field::new("b", DataType::Int32, true)),
6228 ]
6229 .into(),
6230 vec![null_array, int_array],
6231 Some(vec![true, true, false].into()), ));
6233
6234 let converter =
6235 RowConverter::new(vec![SortField::new(struct_array.data_type().clone())]).unwrap();
6236 let rows = converter
6237 .convert_columns(&[Arc::clone(&struct_array)])
6238 .unwrap();
6239 let back = converter.convert_rows(&rows).unwrap();
6240
6241 assert_eq!(&struct_array, &back[0]);
6242 }
6243
6244 #[test]
6246 fn test_nested_struct_with_null() {
6247 let inner_null = Arc::new(NullArray::new(2));
6249 let inner_struct = Arc::new(StructArray::new(
6250 vec![Arc::new(Field::new("x", DataType::Null, true))].into(),
6251 vec![inner_null],
6252 None,
6253 ));
6254
6255 let y_array = Arc::new(Int32Array::from(vec![10, 20]));
6257 let outer_struct: ArrayRef = Arc::new(StructArray::new(
6258 vec![
6259 Arc::new(Field::new("inner", inner_struct.data_type().clone(), true)),
6260 Arc::new(Field::new("y", DataType::Int32, true)),
6261 ]
6262 .into(),
6263 vec![inner_struct, y_array],
6264 None,
6265 ));
6266
6267 let converter =
6268 RowConverter::new(vec![SortField::new(outer_struct.data_type().clone())]).unwrap();
6269 let rows = converter
6270 .convert_columns(&[Arc::clone(&outer_struct)])
6271 .unwrap();
6272 let back = converter.convert_rows(&rows).unwrap();
6273
6274 assert_eq!(&outer_struct, &back[0]);
6275 }
6276
6277 #[test]
6279 fn test_map_null_variations() {
6280 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6282 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6283
6284 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6285 let entries_fields = vec![
6286 Arc::new(Field::new("keys", DataType::Utf8, false)),
6287 Arc::new(Field::new("values", DataType::Null, true)),
6288 ];
6289 let struct_field = Arc::new(Field::new(
6290 "entries",
6291 DataType::Struct(entries_fields.clone().into()),
6292 false,
6293 ));
6294 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6295
6296 let map: ArrayRef = Arc::new(MapArray::new(
6297 struct_field.clone(),
6298 offsets,
6299 entries,
6300 None,
6301 false,
6302 ));
6303
6304 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6305 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6306 let back = converter.convert_rows(&rows).unwrap();
6307 assert_eq!(back.len(), 1);
6308 back[0].to_data().validate_full().unwrap();
6309 assert_eq!(&map, &back[0]);
6310
6311 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6313 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6314
6315 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6316 let entries_fields = vec![
6317 Arc::new(Field::new("keys", DataType::Utf8, false)),
6318 Arc::new(Field::new("values", DataType::Null, true)),
6319 ];
6320 let struct_field = Arc::new(Field::new(
6321 "entries",
6322 DataType::Struct(entries_fields.clone().into()),
6323 false,
6324 ));
6325 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6326
6327 let map: ArrayRef = Arc::new(MapArray::new(
6328 struct_field.clone(),
6329 offsets,
6330 entries,
6331 Some(vec![true, false, true].into()),
6332 false,
6333 ));
6334
6335 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6336 let back = converter.convert_rows(&rows).unwrap();
6337 assert_eq!(back.len(), 1);
6338 back[0].to_data().validate_full().unwrap();
6339 assert_eq!(&map, &back[0]);
6340
6341 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6343 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6344
6345 let offsets = OffsetBuffer::new(vec![0i32].into());
6346 let entries_fields = vec![
6347 Arc::new(Field::new("keys", DataType::Utf8, false)),
6348 Arc::new(Field::new("values", DataType::Null, true)),
6349 ];
6350 let struct_field = Arc::new(Field::new(
6351 "entries",
6352 DataType::Struct(entries_fields.clone().into()),
6353 false,
6354 ));
6355 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6356
6357 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6358
6359 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6360 let back = converter.convert_rows(&rows).unwrap();
6361 assert_eq!(back.len(), 1);
6362 back[0].to_data().validate_full().unwrap();
6363 assert_eq!(&map, &back[0]);
6364 }
6365
6366 #[test]
6368 fn test_map_null_descending() {
6369 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6371 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6372
6373 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6374 let entries_fields = vec![
6375 Arc::new(Field::new("keys", DataType::Utf8, false)),
6376 Arc::new(Field::new("values", DataType::Null, true)),
6377 ];
6378 let struct_field = Arc::new(Field::new(
6379 "entries",
6380 DataType::Struct(entries_fields.clone().into()),
6381 false,
6382 ));
6383 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6384
6385 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6386
6387 let options = SortOptions::default().with_descending(true);
6388 let field = SortField::new_with_options(map.data_type().clone(), options);
6389 let converter = RowConverter::new(vec![field]).unwrap();
6390 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6391 let back = converter.convert_rows(&rows).unwrap();
6392 assert_eq!(back.len(), 1);
6393 back[0].to_data().validate_full().unwrap();
6394 assert_eq!(&map, &back[0]);
6395 }
6396
6397 #[test]
6399 fn test_map_null_keys_and_null_values() {
6400 let null_keys = Arc::new(NullArray::new(3)) as ArrayRef;
6401 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6402
6403 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6404 let entries_fields = vec![
6405 Arc::new(Field::new("keys", DataType::Null, true)),
6406 Arc::new(Field::new("values", DataType::Null, true)),
6407 ];
6408 let struct_field = Arc::new(Field::new(
6409 "entries",
6410 DataType::Struct(entries_fields.clone().into()),
6411 false,
6412 ));
6413 let entries = StructArray::new(entries_fields.into(), vec![null_keys, null_values], None);
6414
6415 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6416
6417 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6418 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6419 let back = converter.convert_rows(&rows).unwrap();
6420 assert_eq!(back.len(), 1);
6421 back[0].to_data().validate_full().unwrap();
6422 assert_eq!(&map, &back[0]);
6423 }
6424
6425 #[test]
6427 fn test_map_null_all_empty() {
6428 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6429 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6430
6431 let offsets = OffsetBuffer::new(vec![0, 0, 0, 0].into());
6432 let entries_fields = vec![
6433 Arc::new(Field::new("keys", DataType::Utf8, false)),
6434 Arc::new(Field::new("values", DataType::Null, true)),
6435 ];
6436 let struct_field = Arc::new(Field::new(
6437 "entries",
6438 DataType::Struct(entries_fields.clone().into()),
6439 false,
6440 ));
6441 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6442
6443 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6444
6445 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6446 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6447
6448 assert_eq!(rows.row(0), rows.row(1));
6450 assert_eq!(rows.row(1), rows.row(2));
6451
6452 let back = converter.convert_rows(&rows).unwrap();
6453 assert_eq!(back.len(), 1);
6454 back[0].to_data().validate_full().unwrap();
6455 assert_eq!(&map, &back[0]);
6456 }
6457
6458 #[test]
6460 fn test_nested_map_null() {
6461 let inner_keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6463 let inner_null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6464
6465 let inner_entries_fields = vec![
6466 Arc::new(Field::new("keys", DataType::Utf8, false)),
6467 Arc::new(Field::new("values", DataType::Null, true)),
6468 ];
6469 let inner_struct_field = Arc::new(Field::new(
6470 "entries",
6471 DataType::Struct(inner_entries_fields.clone().into()),
6472 false,
6473 ));
6474 let inner_entries = StructArray::new(
6475 inner_entries_fields.clone().into(),
6476 vec![inner_keys, inner_null_values],
6477 None,
6478 );
6479
6480 let inner_map = Arc::new(MapArray::new(
6482 inner_struct_field.clone(),
6483 OffsetBuffer::new(vec![0, 1, 3].into()),
6484 inner_entries,
6485 None,
6486 false,
6487 )) as ArrayRef;
6488
6489 let outer_keys = Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef;
6491
6492 let inner_map_type = DataType::Map(inner_struct_field.clone(), false);
6493 let outer_entries_fields = vec![
6494 Arc::new(Field::new("keys", DataType::Utf8, false)),
6495 Arc::new(Field::new("values", inner_map_type, true)),
6496 ];
6497 let outer_struct_field = Arc::new(Field::new(
6498 "entries",
6499 DataType::Struct(outer_entries_fields.clone().into()),
6500 false,
6501 ));
6502 let outer_entries = StructArray::new(
6503 outer_entries_fields.into(),
6504 vec![outer_keys, inner_map],
6505 None,
6506 );
6507
6508 let map: ArrayRef = Arc::new(MapArray::new(
6510 outer_struct_field,
6511 OffsetBuffer::new(vec![0, 1, 2].into()),
6512 outer_entries,
6513 None,
6514 false,
6515 ));
6516
6517 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6518 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6519 let back = converter.convert_rows(&rows).unwrap();
6520 assert_eq!(back.len(), 1);
6521 back[0].to_data().validate_full().unwrap();
6522 assert_eq!(&map, &back[0]);
6523 }
6524
6525 #[test]
6527 fn test_list_of_map_null() {
6528 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6530 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6531
6532 let entries_fields = vec![
6533 Arc::new(Field::new("keys", DataType::Utf8, false)),
6534 Arc::new(Field::new("values", DataType::Null, true)),
6535 ];
6536 let struct_field = Arc::new(Field::new(
6537 "entries",
6538 DataType::Struct(entries_fields.clone().into()),
6539 false,
6540 ));
6541 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6542
6543 let map_array = Arc::new(MapArray::new(
6545 struct_field.clone(),
6546 OffsetBuffer::new(vec![0, 1, 1, 3].into()),
6547 entries,
6548 None,
6549 false,
6550 )) as ArrayRef;
6551
6552 let map_type = DataType::Map(struct_field, false);
6553 let list: ArrayRef = Arc::new(ListArray::new(
6555 Arc::new(Field::new_list_field(map_type, true)),
6556 OffsetBuffer::new(vec![0, 1, 3].into()),
6557 map_array,
6558 None,
6559 ));
6560
6561 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6562 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6563 let back = converter.convert_rows(&rows).unwrap();
6564 assert_eq!(&list, &back[0]);
6565 }
6566
6567 #[test]
6569 fn test_map_of_list_null() {
6570 let null_array = Arc::new(NullArray::new(3)) as ArrayRef;
6572 let list_array = Arc::new(ListArray::new(
6574 Arc::new(Field::new_list_field(DataType::Null, true)),
6575 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6576 null_array,
6577 None,
6578 )) as ArrayRef;
6579
6580 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6581
6582 let list_type = list_array.data_type().clone();
6583 let entries_fields = vec![
6584 Arc::new(Field::new("keys", DataType::Utf8, false)),
6585 Arc::new(Field::new("values", list_type, true)),
6586 ];
6587 let struct_field = Arc::new(Field::new(
6588 "entries",
6589 DataType::Struct(entries_fields.clone().into()),
6590 false,
6591 ));
6592 let entries = StructArray::new(entries_fields.into(), vec![keys, list_array], None);
6593
6594 let map: ArrayRef = Arc::new(MapArray::new(
6596 struct_field,
6597 OffsetBuffer::new(vec![0, 3].into()),
6598 entries,
6599 None,
6600 false,
6601 ));
6602
6603 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6604 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6605 let back = converter.convert_rows(&rows).unwrap();
6606 assert_eq!(back.len(), 1);
6607 back[0].to_data().validate_full().unwrap();
6608 assert_eq!(&map, &back[0]);
6609 }
6610
6611 #[test]
6612 fn empty_row_iter_next_back() {
6613 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6614 .unwrap()
6615 .empty_rows(0, 0);
6616 let mut rows_iter = rows.iter();
6617 assert_eq!(rows_iter.next_back(), None);
6618 assert_eq!(rows_iter.next_back(), None);
6619 assert_eq!(rows_iter.next_back(), None);
6620 }
6621
6622 #[test]
6623 fn row_iter_next_back() {
6624 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6625 let mut rng = StdRng::seed_from_u64(42);
6626 let array = generate_primitive_array::<UInt8Type>(&mut rng, 100, 0.8);
6627 let rows = row_converter.convert_columns(&[Arc::new(array)]).unwrap();
6628
6629 let mut rows_iter = rows.iter();
6630 let mut bytes: Vec<u8> = vec![];
6631
6632 while let Some(row) = rows_iter.next_back() {
6633 bytes.extend(row.data.iter().rev());
6634 }
6635
6636 bytes.reverse();
6637
6638 assert_eq!(
6639 bytes,
6640 &rows.buffer.as_slice()[..*rows.offsets.last().unwrap()]
6641 );
6642
6643 assert_eq!(rows_iter.next_back(), None);
6644 assert_eq!(rows_iter.next(), None);
6645 }
6646}