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 unsafe fn parser_skip_utf8_validation(&self) -> RowParser {
1264 unsafe { RowParser::with_skip_utf8_validate(Arc::clone(&self.fields)) }
1265 }
1266
1267 pub fn size(&self) -> usize {
1271 std::mem::size_of::<Self>()
1272 + self.fields.iter().map(|x| x.size()).sum::<usize>()
1273 + self.codecs.capacity() * std::mem::size_of::<Codec>()
1274 + self.codecs.iter().map(Codec::size).sum::<usize>()
1275 }
1276}
1277
1278#[derive(Debug)]
1280pub struct RowParser {
1281 config: RowConfig,
1282}
1283
1284impl RowParser {
1285 fn new(fields: Arc<[SortField]>) -> Self {
1286 Self {
1287 config: RowConfig {
1288 fields,
1289 validate_utf8: true,
1290 },
1291 }
1292 }
1293 unsafe fn with_skip_utf8_validate(fields: Arc<[SortField]>) -> Self {
1298 Self {
1299 config: RowConfig {
1300 fields,
1301 validate_utf8: false,
1302 },
1303 }
1304 }
1305
1306 pub fn parse<'a>(&'a self, bytes: &'a [u8]) -> Row<'a> {
1311 Row {
1312 data: bytes,
1313 config: &self.config,
1314 }
1315 }
1316}
1317
1318#[derive(Debug, Clone)]
1320struct RowConfig {
1321 fields: Arc<[SortField]>,
1323 validate_utf8: bool,
1325}
1326
1327#[derive(Debug, Clone)]
1331pub struct Rows {
1332 buffer: Vec<u8>,
1334 offsets: Vec<usize>,
1336 config: RowConfig,
1338}
1339
1340pub type RowLengthIter<'a> = Map<Windows<'a, usize>, fn(&'a [usize]) -> usize>;
1342
1343impl Rows {
1344 pub fn push(&mut self, row: Row<'_>) {
1346 assert!(
1347 Arc::ptr_eq(&row.config.fields, &self.config.fields),
1348 "row was not produced by this RowConverter"
1349 );
1350 self.config.validate_utf8 |= row.config.validate_utf8;
1351 self.buffer.extend_from_slice(row.data);
1352 self.offsets.push(self.buffer.len())
1353 }
1354
1355 pub fn reserve(&mut self, row_capacity: usize, data_capacity: usize) {
1357 self.buffer.reserve(data_capacity);
1358 self.offsets.reserve(row_capacity);
1359 }
1360
1361 pub fn row(&self, row: usize) -> Row<'_> {
1363 self.checked_row_end(row);
1364 unsafe { self.row_unchecked(row) }
1365 }
1366
1367 fn checked_row_end(&self, row: usize) -> usize {
1368 row.checked_add(1)
1369 .filter(|end| *end < self.offsets.len())
1370 .expect("row index out of bounds")
1371 }
1372
1373 pub unsafe fn row_unchecked(&self, index: usize) -> Row<'_> {
1378 let end = unsafe { self.offsets.get_unchecked(index + 1) };
1379 let start = unsafe { self.offsets.get_unchecked(index) };
1380 let data = unsafe { self.buffer.get_unchecked(*start..*end) };
1381 Row {
1382 data,
1383 config: &self.config,
1384 }
1385 }
1386
1387 pub fn row_len(&self, row: usize) -> usize {
1390 let end = self.checked_row_end(row);
1391
1392 self.offsets[end] - self.offsets[row]
1393 }
1394
1395 pub fn lengths(&self) -> RowLengthIter<'_> {
1397 self.offsets.windows(2).map(|w| w[1] - w[0])
1398 }
1399
1400 pub fn clear(&mut self) {
1402 self.offsets.truncate(1);
1403 self.buffer.clear();
1404 }
1405
1406 pub fn num_rows(&self) -> usize {
1408 self.offsets.len() - 1
1409 }
1410
1411 pub fn iter(&self) -> RowsIter<'_> {
1413 self.into_iter()
1414 }
1415
1416 pub fn size(&self) -> usize {
1420 std::mem::size_of::<Self>()
1422 + self.buffer.capacity()
1423 + self.offsets.capacity() * std::mem::size_of::<usize>()
1424 }
1425
1426 pub fn try_into_binary(self) -> Result<BinaryArray, ArrowError> {
1456 if self.buffer.len() > i32::MAX as usize {
1457 return Err(ArrowError::InvalidArgumentError(format!(
1458 "{}-byte rows buffer too long to convert into a i32-indexed BinaryArray",
1459 self.buffer.len()
1460 )));
1461 }
1462 let offsets_scalar = ScalarBuffer::from_iter(self.offsets.into_iter().map(i32::usize_as));
1464 let array = unsafe {
1466 BinaryArray::new_unchecked(
1467 OffsetBuffer::new_unchecked(offsets_scalar),
1468 Buffer::from_vec(self.buffer),
1469 None,
1470 )
1471 };
1472 Ok(array)
1473 }
1474}
1475
1476impl<'a> IntoIterator for &'a Rows {
1477 type Item = Row<'a>;
1478 type IntoIter = RowsIter<'a>;
1479
1480 fn into_iter(self) -> Self::IntoIter {
1481 RowsIter {
1482 rows: self,
1483 start: 0,
1484 end: self.num_rows(),
1485 }
1486 }
1487}
1488
1489#[derive(Debug)]
1491pub struct RowsIter<'a> {
1492 rows: &'a Rows,
1493 start: usize,
1494 end: usize,
1495}
1496
1497impl<'a> Iterator for RowsIter<'a> {
1498 type Item = Row<'a>;
1499
1500 fn next(&mut self) -> Option<Self::Item> {
1501 if self.end == self.start {
1502 return None;
1503 }
1504
1505 let row = unsafe { self.rows.row_unchecked(self.start) };
1507 self.start += 1;
1508 Some(row)
1509 }
1510
1511 fn size_hint(&self) -> (usize, Option<usize>) {
1512 let len = self.len();
1513 (len, Some(len))
1514 }
1515}
1516
1517impl ExactSizeIterator for RowsIter<'_> {
1518 fn len(&self) -> usize {
1519 self.end - self.start
1520 }
1521}
1522
1523impl DoubleEndedIterator for RowsIter<'_> {
1524 fn next_back(&mut self) -> Option<Self::Item> {
1525 if self.end == self.start {
1526 return None;
1527 }
1528
1529 self.end -= 1;
1530
1531 let row = unsafe { self.rows.row_unchecked(self.end) };
1534 Some(row)
1535 }
1536}
1537
1538#[derive(Debug, Copy, Clone)]
1547pub struct Row<'a> {
1548 data: &'a [u8],
1549 config: &'a RowConfig,
1550}
1551
1552impl<'a> Row<'a> {
1553 pub fn owned(&self) -> OwnedRow {
1555 OwnedRow {
1556 data: self.data.into(),
1557 config: self.config.clone(),
1558 }
1559 }
1560
1561 pub fn data(&self) -> &'a [u8] {
1563 self.data
1564 }
1565}
1566
1567impl PartialEq for Row<'_> {
1570 #[inline]
1571 fn eq(&self, other: &Self) -> bool {
1572 self.data.eq(other.data)
1573 }
1574}
1575
1576impl Eq for Row<'_> {}
1577
1578impl PartialOrd for Row<'_> {
1579 #[inline]
1580 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1581 Some(self.cmp(other))
1582 }
1583}
1584
1585impl Ord for Row<'_> {
1586 #[inline]
1587 fn cmp(&self, other: &Self) -> Ordering {
1588 self.data.cmp(other.data)
1589 }
1590}
1591
1592impl Hash for Row<'_> {
1593 #[inline]
1594 fn hash<H: Hasher>(&self, state: &mut H) {
1595 self.data.hash(state)
1596 }
1597}
1598
1599impl AsRef<[u8]> for Row<'_> {
1600 #[inline]
1601 fn as_ref(&self) -> &[u8] {
1602 self.data
1603 }
1604}
1605
1606#[derive(Debug, Clone)]
1610pub struct OwnedRow {
1611 data: Box<[u8]>,
1612 config: RowConfig,
1613}
1614
1615impl OwnedRow {
1616 pub fn row(&self) -> Row<'_> {
1620 Row {
1621 data: &self.data,
1622 config: &self.config,
1623 }
1624 }
1625}
1626
1627impl PartialEq for OwnedRow {
1630 #[inline]
1631 fn eq(&self, other: &Self) -> bool {
1632 self.row().eq(&other.row())
1633 }
1634}
1635
1636impl Eq for OwnedRow {}
1637
1638impl PartialOrd for OwnedRow {
1639 #[inline]
1640 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
1641 Some(self.cmp(other))
1642 }
1643}
1644
1645impl Ord for OwnedRow {
1646 #[inline]
1647 fn cmp(&self, other: &Self) -> Ordering {
1648 self.row().cmp(&other.row())
1649 }
1650}
1651
1652impl Hash for OwnedRow {
1653 #[inline]
1654 fn hash<H: Hasher>(&self, state: &mut H) {
1655 self.row().hash(state)
1656 }
1657}
1658
1659impl AsRef<[u8]> for OwnedRow {
1660 #[inline]
1661 fn as_ref(&self) -> &[u8] {
1662 &self.data
1663 }
1664}
1665
1666#[inline]
1668fn null_sentinel(options: SortOptions) -> u8 {
1669 match options.nulls_first {
1670 true => 0,
1671 false => 0xFF,
1672 }
1673}
1674
1675enum LengthTracker {
1677 Fixed { length: usize, num_rows: usize },
1679 Variable {
1681 fixed_length: usize,
1682 lengths: Vec<usize>,
1683 },
1684}
1685
1686impl LengthTracker {
1687 fn new(num_rows: usize) -> Self {
1688 Self::Fixed {
1689 length: 0,
1690 num_rows,
1691 }
1692 }
1693
1694 fn push_fixed(&mut self, new_length: usize) {
1696 match self {
1697 LengthTracker::Fixed { length, .. } => *length += new_length,
1698 LengthTracker::Variable { fixed_length, .. } => *fixed_length += new_length,
1699 }
1700 }
1701
1702 fn push_variable(&mut self, new_lengths: impl ExactSizeIterator<Item = usize>) {
1704 match self {
1705 LengthTracker::Fixed { length, .. } => {
1706 *self = LengthTracker::Variable {
1707 fixed_length: *length,
1708 lengths: new_lengths.collect(),
1709 }
1710 }
1711 LengthTracker::Variable { lengths, .. } => {
1712 assert_eq!(lengths.len(), new_lengths.len());
1713 lengths
1714 .iter_mut()
1715 .zip(new_lengths)
1716 .for_each(|(length, new_length)| *length += new_length);
1717 }
1718 }
1719 }
1720
1721 fn materialized(&mut self) -> &mut [usize] {
1723 if let LengthTracker::Fixed { length, num_rows } = *self {
1724 *self = LengthTracker::Variable {
1725 fixed_length: length,
1726 lengths: vec![0; num_rows],
1727 };
1728 }
1729
1730 match self {
1731 LengthTracker::Variable { lengths, .. } => lengths,
1732 LengthTracker::Fixed { .. } => unreachable!(),
1733 }
1734 }
1735
1736 fn extend_offsets(&self, initial_offset: usize, offsets: &mut Vec<usize>) -> usize {
1754 match self {
1755 LengthTracker::Fixed { length, num_rows } => {
1756 offsets.extend((0..*num_rows).map(|i| initial_offset + i * length));
1757
1758 initial_offset + num_rows * length
1759 }
1760 LengthTracker::Variable {
1761 fixed_length,
1762 lengths,
1763 } => {
1764 let mut acc = initial_offset;
1765
1766 offsets.extend(lengths.iter().map(|length| {
1767 let current = acc;
1768 acc += length + fixed_length;
1769 current
1770 }));
1771
1772 acc
1773 }
1774 }
1775 }
1776}
1777
1778fn row_lengths(cols: &[ArrayRef], encoders: &[Encoder]) -> LengthTracker {
1780 use fixed::FixedLengthEncoding;
1781
1782 let num_rows = cols.first().map(|x| x.len()).unwrap_or(0);
1783 let mut tracker = LengthTracker::new(num_rows);
1784
1785 for (array, encoder) in cols.iter().zip(encoders) {
1786 match encoder {
1787 Encoder::Stateless => {
1788 downcast_primitive_array! {
1789 array => tracker.push_fixed(fixed::encoded_len(array)),
1790 DataType::Null => tracker.push_fixed(2)
1791 DataType::Boolean => tracker.push_fixed(bool::ENCODED_LEN),
1792 DataType::Binary => push_generic_byte_array_lengths(&mut tracker, as_generic_binary_array::<i32>(array)),
1793 DataType::LargeBinary => push_generic_byte_array_lengths(&mut tracker, as_generic_binary_array::<i64>(array)),
1794 DataType::BinaryView => push_byte_view_array_lengths(&mut tracker, array.as_binary_view()),
1795 DataType::Utf8 => push_generic_byte_array_lengths(&mut tracker, array.as_string::<i32>()),
1796 DataType::LargeUtf8 => push_generic_byte_array_lengths(&mut tracker, array.as_string::<i64>()),
1797 DataType::Utf8View => push_byte_view_array_lengths(&mut tracker, array.as_string_view()),
1798 DataType::FixedSizeBinary(len) => {
1799 let len = len.to_usize().unwrap();
1800 tracker.push_fixed(1 + len)
1801 }
1802 _ => unimplemented!("unsupported data type: {}", array.data_type()),
1803 }
1804 }
1805 Encoder::Dictionary(values, null) => {
1806 downcast_dictionary_array! {
1807 array => {
1808 tracker.push_variable(
1809 array.keys().iter().map(|v| match v {
1810 Some(k) => values.row_len(k.as_usize()),
1811 None => null.data.len(),
1812 })
1813 )
1814 }
1815 _ => unreachable!(),
1816 }
1817 }
1818 Encoder::Struct(rows, null) => {
1819 let array = as_struct_array(array);
1820 if rows.num_rows() > 0 {
1821 tracker.push_variable((0..array.len()).map(|idx| match array.is_valid(idx) {
1823 true => 1 + rows.row_len(idx),
1824 false => 1 + null.data.len(),
1825 }));
1826 } else {
1827 tracker.push_variable((0..array.len()).map(|idx| match array.is_valid(idx) {
1829 true => 1,
1830 false => 1 + null.data.len(),
1831 }));
1832 }
1833 }
1834 Encoder::List(rows) => match array.data_type() {
1835 DataType::List(_) => {
1836 list::compute_lengths(tracker.materialized(), rows, as_list_array(array))
1837 }
1838 DataType::LargeList(_) => {
1839 list::compute_lengths(tracker.materialized(), rows, as_large_list_array(array))
1840 }
1841 DataType::ListView(_) => {
1842 let list_view = array.as_list_view::<i32>();
1843 let (min_offset, _) = compute_list_view_bounds(list_view);
1844 list::compute_lengths_list_view(
1845 tracker.materialized(),
1846 rows,
1847 list_view,
1848 min_offset,
1849 )
1850 }
1851 DataType::LargeListView(_) => {
1852 let list_view = array.as_list_view::<i64>();
1853 let (min_offset, _) = compute_list_view_bounds(list_view);
1854 list::compute_lengths_list_view(
1855 tracker.materialized(),
1856 rows,
1857 list_view,
1858 min_offset,
1859 )
1860 }
1861 DataType::FixedSizeList(_, _) => compute_lengths_fixed_size_list(
1862 &mut tracker,
1863 rows,
1864 as_fixed_size_list_array(array),
1865 ),
1866 _ => unreachable!(),
1867 },
1868 Encoder::Map(rows) => {
1869 list::compute_lengths(tracker.materialized(), rows, as_map_array(array))
1870 }
1871 Encoder::RunEndEncoded(rows) => match array.data_type() {
1872 DataType::RunEndEncoded(r, _) => match r.data_type() {
1873 DataType::Int16 => run::compute_lengths(
1874 tracker.materialized(),
1875 rows,
1876 array.as_run::<Int16Type>(),
1877 ),
1878 DataType::Int32 => run::compute_lengths(
1879 tracker.materialized(),
1880 rows,
1881 array.as_run::<Int32Type>(),
1882 ),
1883 DataType::Int64 => run::compute_lengths(
1884 tracker.materialized(),
1885 rows,
1886 array.as_run::<Int64Type>(),
1887 ),
1888 _ => unreachable!("Unsupported run end index type: {r:?}"),
1889 },
1890 _ => unreachable!(),
1891 },
1892 Encoder::Union {
1893 child_rows,
1894 field_to_type_ids,
1895 type_ids,
1896 offsets,
1897 } => {
1898 let union_array = array
1899 .as_any()
1900 .downcast_ref::<UnionArray>()
1901 .expect("expected UnionArray");
1902
1903 let mut type_id_to_field_idx = [0usize; 128];
1904 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
1905 type_id_to_field_idx[type_id as usize] = field_idx;
1906 }
1907
1908 let lengths = (0..union_array.len()).map(|i| {
1909 let type_id = type_ids[i];
1910 let field_idx = type_id_to_field_idx[type_id as usize];
1911 let child_row_i = offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i);
1912 let child_row_len = child_rows[field_idx].row_len(child_row_i);
1913
1914 1 + child_row_len
1916 });
1917
1918 tracker.push_variable(lengths);
1919 }
1920 }
1921 }
1922
1923 tracker
1924}
1925
1926fn push_generic_byte_array_lengths<T: ByteArrayType>(
1928 tracker: &mut LengthTracker,
1929 array: &GenericByteArray<T>,
1930) {
1931 if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
1932 tracker.push_variable(
1933 array
1934 .offsets()
1935 .lengths()
1936 .zip(nulls.iter())
1937 .map(|(length, is_valid)| if is_valid { Some(length) } else { None })
1938 .map(variable::padded_length),
1939 )
1940 } else {
1941 tracker.push_variable(
1942 array
1943 .offsets()
1944 .lengths()
1945 .map(variable::non_null_padded_length),
1946 )
1947 }
1948}
1949
1950fn push_byte_view_array_lengths<T: ByteViewType>(
1952 tracker: &mut LengthTracker,
1953 array: &GenericByteViewArray<T>,
1954) {
1955 if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
1956 tracker.push_variable(
1957 array
1958 .lengths()
1959 .zip(nulls.iter())
1960 .map(|(length, is_valid)| {
1961 if is_valid {
1962 Some(length as usize)
1963 } else {
1964 None
1965 }
1966 })
1967 .map(variable::padded_length),
1968 )
1969 } else {
1970 tracker.push_variable(
1971 array
1972 .lengths()
1973 .map(|len| variable::padded_length(Some(len as usize))),
1974 )
1975 }
1976}
1977
1978fn encode_column(
1980 data: &mut [u8],
1981 offsets: &mut [usize],
1982 column: &dyn Array,
1983 opts: SortOptions,
1984 encoder: &Encoder<'_>,
1985) {
1986 match encoder {
1987 Encoder::Stateless => {
1988 downcast_primitive_array! {
1989 column => {
1990 if let Some(nulls) = column.nulls().filter(|n| n.null_count() > 0){
1991 fixed::encode(data, offsets, column.values(), nulls, opts)
1992 } else {
1993 fixed::encode_not_null(data, offsets, column.values(), opts)
1994 }
1995 }
1996 DataType::Null => {
1997 for offset in offsets.iter_mut().skip(1) {
1998 variable::encode_null_value(&mut data[*offset..], opts);
1999 *offset += 2;
2000 }
2001 }
2002 DataType::Boolean => {
2003 if let Some(nulls) = column.nulls().filter(|n| n.null_count() > 0){
2004 fixed::encode_boolean(data, offsets, column.as_boolean().values(), nulls, opts)
2005 } else {
2006 fixed::encode_boolean_not_null(data, offsets, column.as_boolean().values(), opts)
2007 }
2008 }
2009 DataType::Binary => {
2010 variable::encode_generic_byte_array(data, offsets, as_generic_binary_array::<i32>(column), opts)
2011 }
2012 DataType::BinaryView => {
2013 variable::encode(data, offsets, column.as_binary_view().iter(), opts)
2014 }
2015 DataType::LargeBinary => {
2016 variable::encode_generic_byte_array(data, offsets, as_generic_binary_array::<i64>(column), opts)
2017 }
2018 DataType::Utf8 => variable::encode_generic_byte_array(
2019 data, offsets,
2020 column.as_string::<i32>(),
2021 opts,
2022 ),
2023 DataType::LargeUtf8 => variable::encode_generic_byte_array(
2024 data, offsets,
2025 column.as_string::<i64>(),
2026 opts,
2027 ),
2028 DataType::Utf8View => variable::encode(
2029 data, offsets,
2030 column.as_string_view().iter().map(|x| x.map(|x| x.as_bytes())),
2031 opts,
2032 ),
2033 DataType::FixedSizeBinary(_) => {
2034 let array = column.as_any().downcast_ref().unwrap();
2035 fixed::encode_fixed_size_binary(data, offsets, array, opts)
2036 }
2037 _ => unimplemented!("unsupported data type: {}", column.data_type()),
2038 }
2039 }
2040 Encoder::Dictionary(values, nulls) => {
2041 downcast_dictionary_array! {
2042 column => encode_dictionary_values(data, offsets, column, values, nulls),
2043 _ => unreachable!()
2044 }
2045 }
2046 Encoder::Struct(rows, null) => {
2047 fn struct_encode_helper<const NO_CHILD_FIELDS: bool>(
2048 array: &StructArray,
2049 offsets: &mut [usize],
2050 null_sentinel: u8,
2051 rows: &Rows,
2052 null: &Row<'_>,
2053 data: &mut [u8],
2054 ) {
2055 let empty_row = Row {
2056 data: &[],
2057 config: &rows.config,
2058 };
2059
2060 offsets
2061 .iter_mut()
2062 .skip(1)
2063 .enumerate()
2064 .for_each(|(idx, offset)| {
2065 let (row, sentinel) = match array.is_valid(idx) {
2066 true => (
2067 if NO_CHILD_FIELDS {
2068 empty_row
2069 } else {
2070 rows.row(idx)
2071 },
2072 0x01,
2073 ),
2074 false => (*null, null_sentinel),
2075 };
2076 let end_offset = *offset + 1 + row.as_ref().len();
2077 data[*offset] = sentinel;
2078 data[*offset + 1..end_offset].copy_from_slice(row.as_ref());
2079 *offset = end_offset;
2080 })
2081 }
2082
2083 let array = as_struct_array(column);
2084 let null_sentinel = null_sentinel(opts);
2085 if rows.num_rows() == 0 {
2086 struct_encode_helper::<true>(array, offsets, null_sentinel, rows, null, data);
2088 } else {
2089 struct_encode_helper::<false>(array, offsets, null_sentinel, rows, null, data);
2090 }
2091 }
2092 Encoder::List(rows) => match column.data_type() {
2093 DataType::List(_) => list::encode(data, offsets, rows, opts, as_list_array(column)),
2094 DataType::LargeList(_) => {
2095 list::encode(data, offsets, rows, opts, as_large_list_array(column))
2096 }
2097 DataType::ListView(_) => {
2098 let list_view = column.as_list_view::<i32>();
2099 let (min_offset, _) = compute_list_view_bounds(list_view);
2100 list::encode_list_view(data, offsets, rows, opts, list_view, min_offset)
2101 }
2102 DataType::LargeListView(_) => {
2103 let list_view = column.as_list_view::<i64>();
2104 let (min_offset, _) = compute_list_view_bounds(list_view);
2105 list::encode_list_view(data, offsets, rows, opts, list_view, min_offset)
2106 }
2107 DataType::FixedSizeList(_, _) => {
2108 encode_fixed_size_list(data, offsets, rows, opts, as_fixed_size_list_array(column))
2109 }
2110 _ => unreachable!(),
2111 },
2112 Encoder::Map(rows) => list::encode(data, offsets, rows, opts, as_map_array(column)),
2113 Encoder::RunEndEncoded(rows) => match column.data_type() {
2114 DataType::RunEndEncoded(r, _) => match r.data_type() {
2115 DataType::Int16 => {
2116 run::encode(data, offsets, rows, opts, column.as_run::<Int16Type>())
2117 }
2118 DataType::Int32 => {
2119 run::encode(data, offsets, rows, opts, column.as_run::<Int32Type>())
2120 }
2121 DataType::Int64 => {
2122 run::encode(data, offsets, rows, opts, column.as_run::<Int64Type>())
2123 }
2124 _ => unreachable!("Unsupported run end index type: {r:?}"),
2125 },
2126 _ => unreachable!(),
2127 },
2128 Encoder::Union {
2129 child_rows,
2130 field_to_type_ids,
2131 type_ids,
2132 offsets: offsets_buf,
2133 } => {
2134 let mut type_id_to_field_idx = [0usize; 128];
2135 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
2136 type_id_to_field_idx[type_id as usize] = field_idx;
2137 }
2138
2139 offsets
2140 .iter_mut()
2141 .skip(1)
2142 .enumerate()
2143 .for_each(|(i, offset)| {
2144 let type_id = type_ids[i];
2145 let field_idx = type_id_to_field_idx[type_id as usize];
2146
2147 let child_row_idx = offsets_buf.as_ref().map(|o| o[i] as usize).unwrap_or(i);
2148 let child_row = child_rows[field_idx].row(child_row_idx);
2149 let child_bytes = child_row.as_ref();
2150
2151 let type_id_byte = if opts.descending {
2152 !(type_id as u8)
2153 } else {
2154 type_id as u8
2155 };
2156 data[*offset] = type_id_byte;
2157
2158 let child_start = *offset + 1;
2159 let child_end = child_start + child_bytes.len();
2160 data[child_start..child_end].copy_from_slice(child_bytes);
2161
2162 *offset = child_end;
2163 });
2164 }
2165 }
2166}
2167
2168pub fn encode_dictionary_values<K: ArrowDictionaryKeyType>(
2170 data: &mut [u8],
2171 offsets: &mut [usize],
2172 column: &DictionaryArray<K>,
2173 values: &Rows,
2174 null: &Row<'_>,
2175) {
2176 for (offset, k) in offsets.iter_mut().skip(1).zip(column.keys()) {
2177 let row = match k {
2178 Some(k) => values.row(k.as_usize()).data,
2179 None => null.data,
2180 };
2181 let end_offset = *offset + row.len();
2182 data[*offset..end_offset].copy_from_slice(row);
2183 *offset = end_offset;
2184 }
2185}
2186
2187macro_rules! decode_primitive_helper {
2188 ($t:ty, $rows:ident, $data_type:ident, $options:ident) => {
2189 Arc::new(decode_primitive::<$t>($rows, $data_type, $options))
2190 };
2191}
2192
2193unsafe fn decode_column(
2199 field: &SortField,
2200 rows: &mut [&[u8]],
2201 codec: &Codec,
2202 validate_utf8: bool,
2203) -> Result<ArrayRef, ArrowError> {
2204 let options = field.options;
2205
2206 let array: ArrayRef = match codec {
2207 Codec::Stateless => {
2208 let data_type = field.data_type.clone();
2209 downcast_primitive! {
2210 data_type => (decode_primitive_helper, rows, data_type, options),
2211 DataType::Null => {
2212 variable::decode_null_value(rows, options);
2213 Arc::new(NullArray::new(rows.len()))
2214 }
2215 DataType::Boolean => Arc::new(decode_bool(rows, options)),
2216 DataType::Binary => Arc::new(decode_binary::<i32>(rows, options)),
2217 DataType::LargeBinary => Arc::new(decode_binary::<i64>(rows, options)),
2218 DataType::BinaryView => Arc::new(decode_binary_view(rows, options)),
2219 DataType::FixedSizeBinary(size) => Arc::new(decode_fixed_size_binary(rows, size, options)),
2220 DataType::Utf8 => Arc::new(unsafe{ decode_string::<i32>(rows, options, validate_utf8) }),
2221 DataType::LargeUtf8 => Arc::new(unsafe { decode_string::<i64>(rows, options, validate_utf8) }),
2222 DataType::Utf8View => Arc::new(unsafe { decode_string_view(rows, options, validate_utf8) }),
2223 _ => return Err(ArrowError::NotYetImplemented(format!("unsupported data type: {data_type}" )))
2224 }
2225 }
2226 Codec::Dictionary(converter, _) => {
2227 let cols = unsafe { converter.convert_raw(rows, validate_utf8) }?;
2228 cols.into_iter().next().unwrap()
2229 }
2230 Codec::Struct(converter, _) => {
2231 let nulls = fixed::decode_nulls(rows);
2232 rows.iter_mut().for_each(|row| *row = &row[1..]);
2233 let children = unsafe { converter.convert_raw(rows, validate_utf8) }?;
2234
2235 let corrected_fields: Vec<Field> = match &field.data_type {
2238 DataType::Struct(struct_fields) => struct_fields
2239 .iter()
2240 .zip(children.iter())
2241 .map(|(orig_field, child_array)| {
2242 orig_field
2243 .as_ref()
2244 .clone()
2245 .with_data_type(child_array.data_type().clone())
2246 })
2247 .collect(),
2248 _ => unreachable!("Only Struct types should be corrected here"),
2249 };
2250
2251 Arc::new(unsafe {
2252 StructArray::new_unchecked_with_length(
2253 corrected_fields.into(),
2254 children,
2255 nulls,
2256 rows.len(),
2257 )
2258 })
2259 }
2260 Codec::List(converter) => match &field.data_type {
2261 DataType::List(_) => Arc::new(unsafe {
2262 list::decode::<GenericListArray<i32>>(converter, rows, field, validate_utf8)
2263 }?),
2264 DataType::LargeList(_) => Arc::new(unsafe {
2265 list::decode::<GenericListArray<i64>>(converter, rows, field, validate_utf8)
2266 }?),
2267 DataType::ListView(_) => Arc::new(unsafe {
2268 list::decode_list_view::<i32>(converter, rows, field, validate_utf8)
2269 }?),
2270 DataType::LargeListView(_) => Arc::new(unsafe {
2271 list::decode_list_view::<i64>(converter, rows, field, validate_utf8)
2272 }?),
2273 DataType::FixedSizeList(_, value_length) => Arc::new(unsafe {
2274 list::decode_fixed_size_list(
2275 converter,
2276 rows,
2277 field,
2278 validate_utf8,
2279 value_length.as_usize(),
2280 )
2281 }?),
2282 _ => unreachable!(),
2283 },
2284 Codec::Map(converter) => {
2285 Arc::new(unsafe { list::decode::<MapArray>(converter, rows, field, validate_utf8) }?)
2286 }
2287 Codec::RunEndEncoded(converter) => match &field.data_type {
2288 DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() {
2289 DataType::Int16 => Arc::new(unsafe {
2290 run::decode::<Int16Type>(converter, rows, field, validate_utf8)
2291 }?),
2292 DataType::Int32 => Arc::new(unsafe {
2293 run::decode::<Int32Type>(converter, rows, field, validate_utf8)
2294 }?),
2295 DataType::Int64 => Arc::new(unsafe {
2296 run::decode::<Int64Type>(converter, rows, field, validate_utf8)
2297 }?),
2298 _ => unreachable!(),
2299 },
2300 _ => unreachable!(),
2301 },
2302 Codec::Union(converters, field_to_type_ids, null_rows) => {
2303 let len = rows.len();
2304
2305 let DataType::Union(union_fields, mode) = &field.data_type else {
2306 unreachable!()
2307 };
2308
2309 let mut type_id_to_field_idx = [0usize; 128];
2310 for (field_idx, &type_id) in field_to_type_ids.iter().enumerate() {
2311 type_id_to_field_idx[type_id as usize] = field_idx;
2312 }
2313
2314 let mut type_ids = Vec::with_capacity(len);
2315 let mut rows_by_field: Vec<Vec<(usize, &[u8])>> = vec![Vec::new(); converters.len()];
2316
2317 for (idx, row) in rows.iter_mut().enumerate() {
2318 let type_id_byte = {
2319 let id = row[0];
2320 if options.descending { !id } else { id }
2321 };
2322
2323 let type_id = type_id_byte as i8;
2324 type_ids.push(type_id);
2325
2326 let field_idx = type_id_to_field_idx[type_id as usize];
2327
2328 let child_row = &row[1..];
2329 rows_by_field[field_idx].push((idx, child_row));
2330 }
2331
2332 let mut child_arrays: Vec<ArrayRef> = Vec::with_capacity(converters.len());
2333 let mut offsets = (*mode == UnionMode::Dense).then(|| Vec::with_capacity(len));
2334
2335 for (field_idx, converter) in converters.iter().enumerate() {
2336 let field_rows = &rows_by_field[field_idx];
2337
2338 match &mode {
2339 UnionMode::Dense => {
2340 if field_rows.is_empty() {
2341 let (_, field) = union_fields.iter().nth(field_idx).unwrap();
2342 child_arrays.push(arrow_array::new_empty_array(field.data_type()));
2343 continue;
2344 }
2345
2346 let mut child_data = field_rows
2347 .iter()
2348 .map(|(_, bytes)| *bytes)
2349 .collect::<Vec<_>>();
2350
2351 let child_array =
2352 unsafe { converter.convert_raw(&mut child_data, validate_utf8) }?;
2353
2354 for ((row_idx, original_bytes), remaining_bytes) in
2356 field_rows.iter().zip(child_data)
2357 {
2358 let consumed_length = 1 + original_bytes.len() - remaining_bytes.len();
2359 rows[*row_idx] = &rows[*row_idx][consumed_length..];
2360 }
2361
2362 child_arrays.push(child_array.into_iter().next().unwrap());
2363 }
2364 UnionMode::Sparse => {
2365 let mut sparse_data: Vec<&[u8]> = Vec::with_capacity(len);
2366 let mut field_row_iter = field_rows.iter().peekable();
2367 let null_row_bytes: &[u8] = &null_rows[field_idx].data;
2368
2369 for idx in 0..len {
2370 if let Some((next_idx, bytes)) = field_row_iter.peek() {
2371 if *next_idx == idx {
2372 sparse_data.push(*bytes);
2373
2374 field_row_iter.next();
2375 continue;
2376 }
2377 }
2378 sparse_data.push(null_row_bytes);
2379 }
2380
2381 let child_array =
2382 unsafe { converter.convert_raw(&mut sparse_data, validate_utf8) }?;
2383
2384 for (row_idx, child_row) in field_rows.iter() {
2386 let remaining_len = sparse_data[*row_idx].len();
2387 let consumed_length = 1 + child_row.len() - remaining_len;
2388 rows[*row_idx] = &rows[*row_idx][consumed_length..];
2389 }
2390
2391 child_arrays.push(child_array.into_iter().next().unwrap());
2392 }
2393 }
2394 }
2395
2396 if let Some(ref mut offsets_vec) = offsets {
2398 let mut count = vec![0i32; converters.len()];
2399 for type_id in &type_ids {
2400 let field_idx = *type_id as usize;
2401 offsets_vec.push(count[field_idx]);
2402
2403 count[field_idx] += 1;
2404 }
2405 }
2406
2407 let type_ids_buffer = ScalarBuffer::from(type_ids);
2408 let offsets_buffer = offsets.map(ScalarBuffer::from);
2409
2410 let union_array = UnionArray::try_new(
2411 union_fields.clone(),
2412 type_ids_buffer,
2413 offsets_buffer,
2414 child_arrays,
2415 )?;
2416
2417 Arc::new(union_array)
2420 }
2421 };
2422 Ok(array)
2423}
2424
2425#[cfg(test)]
2426mod tests {
2427 use arrow_array::builder::*;
2428 use arrow_array::types::*;
2429 use arrow_array::*;
2430 use arrow_buffer::{Buffer, OffsetBuffer};
2431 use arrow_buffer::{NullBuffer, i256};
2432 use arrow_cast::display::{ArrayFormatter, FormatOptions};
2433 use arrow_ord::sort::{LexicographicalComparator, SortColumn};
2434 use rand::distr::uniform::SampleUniform;
2435 use rand::distr::{Distribution, StandardUniform};
2436 use rand::prelude::StdRng;
2437 use rand::{Rng, RngCore, SeedableRng};
2438
2439 use super::*;
2440
2441 fn all_sort_options() -> [SortOptions; 4] {
2442 [
2443 SortOptions {
2444 descending: false,
2445 nulls_first: false,
2446 },
2447 SortOptions {
2448 descending: false,
2449 nulls_first: true,
2450 },
2451 SortOptions {
2452 descending: true,
2453 nulls_first: false,
2454 },
2455 SortOptions {
2456 descending: true,
2457 nulls_first: true,
2458 },
2459 ]
2460 }
2461
2462 #[test]
2463 fn test_fixed_width() {
2464 let cols = [
2465 Arc::new(Int16Array::from_iter([
2466 Some(1),
2467 Some(2),
2468 None,
2469 Some(-5),
2470 Some(2),
2471 Some(2),
2472 Some(0),
2473 ])) as ArrayRef,
2474 Arc::new(Float32Array::from_iter([
2475 Some(1.3),
2476 Some(2.5),
2477 None,
2478 Some(4.),
2479 Some(0.1),
2480 Some(-4.),
2481 Some(-0.),
2482 ])) as ArrayRef,
2483 ];
2484
2485 let converter = RowConverter::new(vec![
2486 SortField::new(DataType::Int16),
2487 SortField::new(DataType::Float32),
2488 ])
2489 .unwrap();
2490 let rows = converter.convert_columns(&cols).unwrap();
2491
2492 assert_eq!(rows.offsets, &[0, 8, 16, 24, 32, 40, 48, 56]);
2493 assert_eq!(
2494 rows.buffer,
2495 &[
2496 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 ]
2511 );
2512
2513 assert!(rows.row(3) < rows.row(6));
2514 assert!(rows.row(0) < rows.row(1));
2515 assert!(rows.row(3) < rows.row(0));
2516 assert!(rows.row(4) < rows.row(1));
2517 assert!(rows.row(5) < rows.row(4));
2518
2519 let back = converter.convert_rows(&rows).unwrap();
2520 for (expected, actual) in cols.iter().zip(&back) {
2521 assert_eq!(expected, actual);
2522 }
2523 }
2524
2525 fn test_roundtrip(sort_option: SortOptions, col: ArrayRef) {
2526 let converter = RowConverter::new(vec![SortField::new_with_options(
2527 col.data_type().clone(),
2528 sort_option,
2529 )])
2530 .unwrap();
2531 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2532 let back = converter.convert_rows(&rows).unwrap();
2533 assert_eq!(back.len(), 1);
2534 assert_eq!(&back[0], &col);
2535 back[0].to_data().validate_full().unwrap();
2536 }
2537
2538 #[test]
2539 fn test_zero_width_fixed_size_binary_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 FixedSizeBinaryArray::try_new_with_len(0, Buffer::default(), nulls, 5).unwrap(),
2551 );
2552
2553 test_roundtrip(sort_option, col);
2554 }
2555 }
2556 }
2557
2558 #[test]
2559 fn test_zero_width_fixed_size_list_roundtrip() {
2560 for sort_option in all_sort_options() {
2561 for with_null in [true, false] {
2564 let nulls = if with_null {
2565 Some(NullBuffer::from(vec![true, false, true, false, true]))
2566 } else {
2567 None
2568 };
2569 let col: ArrayRef = Arc::new(
2570 FixedSizeListArray::try_new_with_length(
2571 Arc::new(Field::new("item", DataType::Boolean, false)),
2572 0,
2573 new_empty_array(&DataType::Boolean),
2574 nulls,
2575 5,
2576 )
2577 .unwrap(),
2578 );
2579
2580 test_roundtrip(sort_option, col);
2581 }
2582 }
2583 }
2584
2585 #[test]
2586 fn test_decimal32() {
2587 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal32(
2588 DECIMAL32_MAX_PRECISION,
2589 7,
2590 ))])
2591 .unwrap();
2592 let col = Arc::new(
2593 Decimal32Array::from_iter([
2594 None,
2595 Some(i32::MIN),
2596 Some(-13),
2597 Some(46_i32),
2598 Some(5456_i32),
2599 Some(i32::MAX),
2600 ])
2601 .with_precision_and_scale(9, 7)
2602 .unwrap(),
2603 ) as ArrayRef;
2604
2605 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2606 for i in 0..rows.num_rows() - 1 {
2607 assert!(rows.row(i) < rows.row(i + 1));
2608 }
2609
2610 let back = converter.convert_rows(&rows).unwrap();
2611 assert_eq!(back.len(), 1);
2612 assert_eq!(col.as_ref(), back[0].as_ref())
2613 }
2614
2615 #[test]
2616 fn test_decimal64() {
2617 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal64(
2618 DECIMAL64_MAX_PRECISION,
2619 7,
2620 ))])
2621 .unwrap();
2622 let col = Arc::new(
2623 Decimal64Array::from_iter([
2624 None,
2625 Some(i64::MIN),
2626 Some(-13),
2627 Some(46_i64),
2628 Some(5456_i64),
2629 Some(i64::MAX),
2630 ])
2631 .with_precision_and_scale(18, 7)
2632 .unwrap(),
2633 ) as ArrayRef;
2634
2635 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2636 for i in 0..rows.num_rows() - 1 {
2637 assert!(rows.row(i) < rows.row(i + 1));
2638 }
2639
2640 let back = converter.convert_rows(&rows).unwrap();
2641 assert_eq!(back.len(), 1);
2642 assert_eq!(col.as_ref(), back[0].as_ref())
2643 }
2644
2645 #[test]
2646 fn test_decimal128() {
2647 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal128(
2648 DECIMAL128_MAX_PRECISION,
2649 7,
2650 ))])
2651 .unwrap();
2652 let col = Arc::new(
2653 Decimal128Array::from_iter([
2654 None,
2655 Some(i128::MIN),
2656 Some(-13),
2657 Some(46_i128),
2658 Some(5456_i128),
2659 Some(i128::MAX),
2660 ])
2661 .with_precision_and_scale(38, 7)
2662 .unwrap(),
2663 ) as ArrayRef;
2664
2665 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2666 for i in 0..rows.num_rows() - 1 {
2667 assert!(rows.row(i) < rows.row(i + 1));
2668 }
2669
2670 let back = converter.convert_rows(&rows).unwrap();
2671 assert_eq!(back.len(), 1);
2672 assert_eq!(col.as_ref(), back[0].as_ref())
2673 }
2674
2675 #[test]
2676 fn test_decimal256() {
2677 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal256(
2678 DECIMAL256_MAX_PRECISION,
2679 7,
2680 ))])
2681 .unwrap();
2682 let col = Arc::new(
2683 Decimal256Array::from_iter([
2684 None,
2685 Some(i256::MIN),
2686 Some(i256::from_parts(0, -1)),
2687 Some(i256::from_parts(u128::MAX, -1)),
2688 Some(i256::from_parts(u128::MAX, 0)),
2689 Some(i256::from_parts(0, 46_i128)),
2690 Some(i256::from_parts(5, 46_i128)),
2691 Some(i256::MAX),
2692 ])
2693 .with_precision_and_scale(DECIMAL256_MAX_PRECISION, 7)
2694 .unwrap(),
2695 ) as ArrayRef;
2696
2697 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2698 for i in 0..rows.num_rows() - 1 {
2699 assert!(rows.row(i) < rows.row(i + 1));
2700 }
2701
2702 let back = converter.convert_rows(&rows).unwrap();
2703 assert_eq!(back.len(), 1);
2704 assert_eq!(col.as_ref(), back[0].as_ref())
2705 }
2706
2707 #[test]
2708 fn test_bool() {
2709 let converter = RowConverter::new(vec![SortField::new(DataType::Boolean)]).unwrap();
2710
2711 let col = Arc::new(BooleanArray::from_iter([None, Some(false), Some(true)])) as ArrayRef;
2712
2713 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2714 assert!(rows.row(2) > rows.row(1));
2715 assert!(rows.row(2) > rows.row(0));
2716 assert!(rows.row(1) > rows.row(0));
2717
2718 let cols = converter.convert_rows(&rows).unwrap();
2719 assert_eq!(&cols[0], &col);
2720
2721 let converter = RowConverter::new(vec![SortField::new_with_options(
2722 DataType::Boolean,
2723 SortOptions::default().desc().with_nulls_first(false),
2724 )])
2725 .unwrap();
2726
2727 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2728 assert!(rows.row(2) < rows.row(1));
2729 assert!(rows.row(2) < rows.row(0));
2730 assert!(rows.row(1) < rows.row(0));
2731 let cols = converter.convert_rows(&rows).unwrap();
2732 assert_eq!(&cols[0], &col);
2733 }
2734
2735 #[test]
2736 fn test_timezone() {
2737 let a =
2738 TimestampNanosecondArray::from(vec![1, 2, 3, 4, 5]).with_timezone("+01:00".to_string());
2739 let d = a.data_type().clone();
2740
2741 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
2742 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
2743 let back = converter.convert_rows(&rows).unwrap();
2744 assert_eq!(back.len(), 1);
2745 assert_eq!(back[0].data_type(), &d);
2746
2747 let mut a = PrimitiveDictionaryBuilder::<Int32Type, TimestampNanosecondType>::new();
2749 a.append(34).unwrap();
2750 a.append_null();
2751 a.append(345).unwrap();
2752
2753 let dict = a.finish();
2755 let values = TimestampNanosecondArray::from(dict.values().to_data());
2756 let dict_with_tz = dict.with_values(Arc::new(values.with_timezone("+02:00")));
2757 let v = DataType::Timestamp(TimeUnit::Nanosecond, Some("+02:00".into()));
2758 let d = DataType::Dictionary(Box::new(DataType::Int32), Box::new(v.clone()));
2759
2760 assert_eq!(dict_with_tz.data_type(), &d);
2761 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
2762 let rows = converter
2763 .convert_columns(&[Arc::new(dict_with_tz) as _])
2764 .unwrap();
2765 let back = converter.convert_rows(&rows).unwrap();
2766 assert_eq!(back.len(), 1);
2767 assert_eq!(back[0].data_type(), &v);
2768 }
2769
2770 #[test]
2771 fn test_null_encoding() {
2772 let col = Arc::new(NullArray::new(10));
2773 let converter = RowConverter::new(vec![SortField::new(DataType::Null)]).unwrap();
2774 let rows = converter.convert_columns(&[col]).unwrap();
2775 assert_eq!(rows.num_rows(), 10);
2776 assert_eq!(rows.row(1).data.len(), 2);
2778 }
2779
2780 #[test]
2781 fn test_variable_width() {
2782 let col = Arc::new(StringArray::from_iter([
2783 Some("hello"),
2784 Some("he"),
2785 None,
2786 Some("foo"),
2787 Some(""),
2788 ])) as ArrayRef;
2789
2790 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
2791 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2792
2793 assert!(rows.row(1) < rows.row(0));
2794 assert!(rows.row(2) < rows.row(4));
2795 assert!(rows.row(3) < rows.row(0));
2796 assert!(rows.row(3) < rows.row(1));
2797
2798 let cols = converter.convert_rows(&rows).unwrap();
2799 assert_eq!(&cols[0], &col);
2800
2801 let col = Arc::new(BinaryArray::from_iter([
2802 None,
2803 Some(vec![0_u8; 0]),
2804 Some(vec![0_u8; 6]),
2805 Some(vec![0_u8; variable::MINI_BLOCK_SIZE]),
2806 Some(vec![0_u8; variable::MINI_BLOCK_SIZE + 1]),
2807 Some(vec![0_u8; variable::BLOCK_SIZE]),
2808 Some(vec![0_u8; variable::BLOCK_SIZE + 1]),
2809 Some(vec![1_u8; 6]),
2810 Some(vec![1_u8; variable::MINI_BLOCK_SIZE]),
2811 Some(vec![1_u8; variable::MINI_BLOCK_SIZE + 1]),
2812 Some(vec![1_u8; variable::BLOCK_SIZE]),
2813 Some(vec![1_u8; variable::BLOCK_SIZE + 1]),
2814 Some(vec![0xFF_u8; 6]),
2815 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE]),
2816 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE + 1]),
2817 Some(vec![0xFF_u8; variable::BLOCK_SIZE]),
2818 Some(vec![0xFF_u8; variable::BLOCK_SIZE + 1]),
2819 ])) as ArrayRef;
2820
2821 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
2822 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2823
2824 for i in 0..rows.num_rows() {
2825 for j in i + 1..rows.num_rows() {
2826 assert!(
2827 rows.row(i) < rows.row(j),
2828 "{} < {} - {:?} < {:?}",
2829 i,
2830 j,
2831 rows.row(i),
2832 rows.row(j)
2833 );
2834 }
2835 }
2836
2837 let cols = converter.convert_rows(&rows).unwrap();
2838 assert_eq!(&cols[0], &col);
2839
2840 let converter = RowConverter::new(vec![SortField::new_with_options(
2841 DataType::Binary,
2842 SortOptions::default().desc().with_nulls_first(false),
2843 )])
2844 .unwrap();
2845 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2846
2847 for i in 0..rows.num_rows() {
2848 for j in i + 1..rows.num_rows() {
2849 assert!(
2850 rows.row(i) > rows.row(j),
2851 "{} > {} - {:?} > {:?}",
2852 i,
2853 j,
2854 rows.row(i),
2855 rows.row(j)
2856 );
2857 }
2858 }
2859
2860 let cols = converter.convert_rows(&rows).unwrap();
2861 assert_eq!(&cols[0], &col);
2862 }
2863
2864 fn dictionary_eq(a: &dyn Array, b: &dyn Array) {
2866 match b.data_type() {
2867 DataType::Dictionary(_, v) => {
2868 assert_eq!(a.data_type(), v.as_ref());
2869 let b = arrow_cast::cast(b, v).unwrap();
2870 assert_eq!(a, b.as_ref())
2871 }
2872 _ => assert_eq!(a, b),
2873 }
2874 }
2875
2876 #[test]
2877 fn test_string_dictionary() {
2878 let a = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2879 Some("foo"),
2880 Some("hello"),
2881 Some("he"),
2882 None,
2883 Some("hello"),
2884 Some(""),
2885 Some("hello"),
2886 Some("hello"),
2887 ])) as ArrayRef;
2888
2889 let field = SortField::new(a.data_type().clone());
2890 let converter = RowConverter::new(vec![field]).unwrap();
2891 let rows_a = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2892
2893 assert!(rows_a.row(3) < rows_a.row(5));
2894 assert!(rows_a.row(2) < rows_a.row(1));
2895 assert!(rows_a.row(0) < rows_a.row(1));
2896 assert!(rows_a.row(3) < rows_a.row(0));
2897
2898 assert_eq!(rows_a.row(1), rows_a.row(4));
2899 assert_eq!(rows_a.row(1), rows_a.row(6));
2900 assert_eq!(rows_a.row(1), rows_a.row(7));
2901
2902 let cols = converter.convert_rows(&rows_a).unwrap();
2903 dictionary_eq(&cols[0], &a);
2904
2905 let b = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2906 Some("hello"),
2907 None,
2908 Some("cupcakes"),
2909 ])) as ArrayRef;
2910
2911 let rows_b = converter.convert_columns(&[Arc::clone(&b)]).unwrap();
2912 assert_eq!(rows_a.row(1), rows_b.row(0));
2913 assert_eq!(rows_a.row(3), rows_b.row(1));
2914 assert!(rows_b.row(2) < rows_a.row(0));
2915
2916 let cols = converter.convert_rows(&rows_b).unwrap();
2917 dictionary_eq(&cols[0], &b);
2918
2919 let converter = RowConverter::new(vec![SortField::new_with_options(
2920 a.data_type().clone(),
2921 SortOptions::default().desc().with_nulls_first(false),
2922 )])
2923 .unwrap();
2924
2925 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2926 assert!(rows_c.row(3) > rows_c.row(5));
2927 assert!(rows_c.row(2) > rows_c.row(1));
2928 assert!(rows_c.row(0) > rows_c.row(1));
2929 assert!(rows_c.row(3) > rows_c.row(0));
2930
2931 let cols = converter.convert_rows(&rows_c).unwrap();
2932 dictionary_eq(&cols[0], &a);
2933
2934 let converter = RowConverter::new(vec![SortField::new_with_options(
2935 a.data_type().clone(),
2936 SortOptions::default().desc().with_nulls_first(true),
2937 )])
2938 .unwrap();
2939
2940 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2941 assert!(rows_c.row(3) < rows_c.row(5));
2942 assert!(rows_c.row(2) > rows_c.row(1));
2943 assert!(rows_c.row(0) > rows_c.row(1));
2944 assert!(rows_c.row(3) < rows_c.row(0));
2945
2946 let cols = converter.convert_rows(&rows_c).unwrap();
2947 dictionary_eq(&cols[0], &a);
2948 }
2949
2950 #[test]
2951 fn test_struct() {
2952 let a = Arc::new(Int32Array::from(vec![1, 1, 2, 2])) as ArrayRef;
2954 let a_f = Arc::new(Field::new("int", DataType::Int32, false));
2955 let u = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as ArrayRef;
2956 let u_f = Arc::new(Field::new("s", DataType::Utf8, false));
2957 let s1 = Arc::new(StructArray::from(vec![(a_f, a), (u_f, u)])) as ArrayRef;
2958
2959 let sort_fields = vec![SortField::new(s1.data_type().clone())];
2960 let converter = RowConverter::new(sort_fields).unwrap();
2961 let r1 = converter.convert_columns(&[Arc::clone(&s1)]).unwrap();
2962
2963 for (a, b) in r1.iter().zip(r1.iter().skip(1)) {
2964 assert!(a < b);
2965 }
2966
2967 let back = converter.convert_rows(&r1).unwrap();
2968 assert_eq!(back.len(), 1);
2969 assert_eq!(&back[0], &s1);
2970
2971 let data = s1
2973 .to_data()
2974 .into_builder()
2975 .null_bit_buffer(Some(Buffer::from_slice_ref([0b00001010])))
2976 .null_count(2)
2977 .build()
2978 .unwrap();
2979
2980 let s2 = Arc::new(StructArray::from(data)) as ArrayRef;
2981 let r2 = converter.convert_columns(&[Arc::clone(&s2)]).unwrap();
2982 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();
2988 assert_eq!(back.len(), 1);
2989 assert_eq!(&back[0], &s2);
2990
2991 back[0].to_data().validate_full().unwrap();
2992 }
2993
2994 #[test]
2995 fn test_dictionary_in_struct() {
2996 let builder = StringDictionaryBuilder::<Int32Type>::new();
2997 let mut struct_builder = StructBuilder::new(
2998 vec![Field::new_dictionary(
2999 "foo",
3000 DataType::Int32,
3001 DataType::Utf8,
3002 true,
3003 )],
3004 vec![Box::new(builder)],
3005 );
3006
3007 let dict_builder = struct_builder
3008 .field_builder::<StringDictionaryBuilder<Int32Type>>(0)
3009 .unwrap();
3010
3011 dict_builder.append_value("a");
3013 dict_builder.append_null();
3014 dict_builder.append_value("a");
3015 dict_builder.append_value("b");
3016
3017 for _ in 0..4 {
3018 struct_builder.append(true);
3019 }
3020
3021 let s = Arc::new(struct_builder.finish()) as ArrayRef;
3022 let sort_fields = vec![SortField::new(s.data_type().clone())];
3023 let converter = RowConverter::new(sort_fields).unwrap();
3024 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3025
3026 let back = converter.convert_rows(&r).unwrap();
3027 let [s2] = back.try_into().unwrap();
3028
3029 assert_ne!(&s.data_type(), &s2.data_type());
3032 s2.to_data().validate_full().unwrap();
3033
3034 let s1_struct = s.as_struct();
3038 let s1_0 = s1_struct.column(0);
3039 let s1_idx_0 = s1_0.as_dictionary::<Int32Type>();
3040 let keys = s1_idx_0.keys();
3041 let values = s1_idx_0.values().as_string::<i32>();
3042 let s2_struct = s2.as_struct();
3044 let s2_0 = s2_struct.column(0);
3045 let s2_idx_0 = s2_0.as_string::<i32>();
3046
3047 for i in 0..keys.len() {
3048 if keys.is_null(i) {
3049 assert!(s2_idx_0.is_null(i));
3050 } else {
3051 let dict_index = keys.value(i) as usize;
3052 assert_eq!(values.value(dict_index), s2_idx_0.value(i));
3053 }
3054 }
3055 }
3056
3057 #[test]
3058 fn test_dictionary_in_struct_empty() {
3059 let ty = DataType::Struct(
3060 vec![Field::new_dictionary(
3061 "foo",
3062 DataType::Int32,
3063 DataType::Int32,
3064 false,
3065 )]
3066 .into(),
3067 );
3068 let s = arrow_array::new_empty_array(&ty);
3069
3070 let sort_fields = vec![SortField::new(s.data_type().clone())];
3071 let converter = RowConverter::new(sort_fields).unwrap();
3072 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3073
3074 let back = converter.convert_rows(&r).unwrap();
3075 let [s2] = back.try_into().unwrap();
3076
3077 assert_ne!(&s.data_type(), &s2.data_type());
3080 s2.to_data().validate_full().unwrap();
3081 assert_eq!(s.len(), 0);
3082 assert_eq!(s2.len(), 0);
3083 }
3084
3085 #[test]
3086 fn test_list_of_string_dictionary() {
3087 let mut builder = ListBuilder::<StringDictionaryBuilder<Int32Type>>::default();
3088 builder.values().append("a").unwrap();
3090 builder.values().append("b").unwrap();
3091 builder.values().append("zero").unwrap();
3092 builder.values().append_null();
3093 builder.values().append("c").unwrap();
3094 builder.values().append("b").unwrap();
3095 builder.values().append("d").unwrap();
3096 builder.append(true);
3097 builder.append(false);
3099 builder.values().append("e").unwrap();
3101 builder.values().append("zero").unwrap();
3102 builder.values().append("a").unwrap();
3103 builder.append(true);
3104
3105 let a = Arc::new(builder.finish()) as ArrayRef;
3106 let data_type = a.data_type().clone();
3107
3108 let field = SortField::new(data_type.clone());
3109 let converter = RowConverter::new(vec![field]).unwrap();
3110 let rows = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
3111
3112 let back = converter.convert_rows(&rows).unwrap();
3113 assert_eq!(back.len(), 1);
3114 let [a2] = back.try_into().unwrap();
3115
3116 assert_ne!(&a.data_type(), &a2.data_type());
3119
3120 a2.to_data().validate_full().unwrap();
3121
3122 let a2_list = a2.as_list::<i32>();
3123 let a1_list = a.as_list::<i32>();
3124
3125 let a1_0 = a1_list.value(0);
3128 let a1_idx_0 = a1_0.as_dictionary::<Int32Type>();
3129 let keys = a1_idx_0.keys();
3130 let values = a1_idx_0.values().as_string::<i32>();
3131 let a2_0 = a2_list.value(0);
3132 let a2_idx_0 = a2_0.as_string::<i32>();
3133
3134 for i in 0..keys.len() {
3135 if keys.is_null(i) {
3136 assert!(a2_idx_0.is_null(i));
3137 } else {
3138 let dict_index = keys.value(i) as usize;
3139 assert_eq!(values.value(dict_index), a2_idx_0.value(i));
3140 }
3141 }
3142
3143 assert!(a1_list.is_null(1));
3145 assert!(a2_list.is_null(1));
3146
3147 let a1_2 = a1_list.value(2);
3149 let a1_idx_2 = a1_2.as_dictionary::<Int32Type>();
3150 let keys = a1_idx_2.keys();
3151 let values = a1_idx_2.values().as_string::<i32>();
3152 let a2_2 = a2_list.value(2);
3153 let a2_idx_2 = a2_2.as_string::<i32>();
3154
3155 for i in 0..keys.len() {
3156 if keys.is_null(i) {
3157 assert!(a2_idx_2.is_null(i));
3158 } else {
3159 let dict_index = keys.value(i) as usize;
3160 assert_eq!(values.value(dict_index), a2_idx_2.value(i));
3161 }
3162 }
3163 }
3164
3165 #[test]
3166 fn test_primitive_dictionary() {
3167 let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
3168 builder.append(2).unwrap();
3169 builder.append(3).unwrap();
3170 builder.append(0).unwrap();
3171 builder.append_null();
3172 builder.append(5).unwrap();
3173 builder.append(3).unwrap();
3174 builder.append(-1).unwrap();
3175
3176 let a = builder.finish();
3177 let data_type = a.data_type().clone();
3178 let columns = [Arc::new(a) as ArrayRef];
3179
3180 let field = SortField::new(data_type.clone());
3181 let converter = RowConverter::new(vec![field]).unwrap();
3182 let rows = converter.convert_columns(&columns).unwrap();
3183 assert!(rows.row(0) < rows.row(1));
3184 assert!(rows.row(2) < rows.row(0));
3185 assert!(rows.row(3) < rows.row(2));
3186 assert!(rows.row(6) < rows.row(2));
3187 assert!(rows.row(3) < rows.row(6));
3188
3189 let back = converter.convert_rows(&rows).unwrap();
3190 assert_eq!(back.len(), 1);
3191 back[0].to_data().validate_full().unwrap();
3192 }
3193
3194 #[test]
3195 fn test_dictionary_nulls() {
3196 let values = Int32Array::from_iter([Some(1), Some(-1), None, Some(4), None]).into_data();
3197 let keys =
3198 Int32Array::from_iter([Some(0), Some(0), Some(1), Some(2), Some(4), None]).into_data();
3199
3200 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3201 let data = keys
3202 .into_builder()
3203 .data_type(data_type.clone())
3204 .child_data(vec![values])
3205 .build()
3206 .unwrap();
3207
3208 let columns = [Arc::new(DictionaryArray::<Int32Type>::from(data)) as ArrayRef];
3209 let field = SortField::new(data_type.clone());
3210 let converter = RowConverter::new(vec![field]).unwrap();
3211 let rows = converter.convert_columns(&columns).unwrap();
3212
3213 assert_eq!(rows.row(0), rows.row(1));
3214 assert_eq!(rows.row(3), rows.row(4));
3215 assert_eq!(rows.row(4), rows.row(5));
3216 assert!(rows.row(3) < rows.row(0));
3217 }
3218
3219 #[test]
3220 fn test_from_binary_shared_buffer() {
3221 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3222 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3223 let rows = converter.convert_columns(&[array]).unwrap();
3224 let binary_rows = rows.try_into_binary().expect("known-small rows");
3225 let _binary_rows_shared_buffer = binary_rows.clone();
3226
3227 let parsed = converter.from_binary(binary_rows);
3228
3229 converter.convert_rows(parsed.iter()).unwrap();
3230 }
3231
3232 #[test]
3233 #[should_panic(expected = "Encountered non UTF-8 data")]
3234 fn test_invalid_utf8() {
3235 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3236 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3237 let rows = converter.convert_columns(&[array]).unwrap();
3238 let binary_row = rows.row(0);
3239
3240 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3241 let parser = converter.parser();
3242 let utf8_row = parser.parse(binary_row.as_ref());
3243
3244 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3245 }
3246
3247 #[test]
3248 #[should_panic(expected = "Encountered non UTF-8 data")]
3249 fn test_invalid_utf8_array() {
3250 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3251 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3252 let rows = converter.convert_columns(&[array]).unwrap();
3253 let binary_rows = rows.try_into_binary().expect("known-small rows");
3254
3255 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3256 let parsed = converter.from_binary(binary_rows);
3257
3258 converter.convert_rows(parsed.iter()).unwrap();
3259 }
3260
3261 #[test]
3262 #[should_panic(expected = "index out of bounds")]
3263 fn test_invalid_empty() {
3264 let binary_row: &[u8] = &[];
3265
3266 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3267 let parser = converter.parser();
3268 let utf8_row = parser.parse(binary_row.as_ref());
3269
3270 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3271 }
3272
3273 #[test]
3274 #[should_panic(expected = "index out of bounds")]
3275 fn test_invalid_empty_array() {
3276 let row: &[u8] = &[];
3277 let binary_rows = BinaryArray::from(vec![row]);
3278
3279 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3280 let parsed = converter.from_binary(binary_rows);
3281
3282 converter.convert_rows(parsed.iter()).unwrap();
3283 }
3284
3285 #[test]
3286 #[should_panic(expected = "index out of bounds")]
3287 fn test_invalid_truncated() {
3288 let binary_row: &[u8] = &[0x02];
3289
3290 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3291 let parser = converter.parser();
3292 let utf8_row = parser.parse(binary_row.as_ref());
3293
3294 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3295 }
3296
3297 #[test]
3298 #[should_panic(expected = "index out of bounds")]
3299 fn test_invalid_truncated_array() {
3300 let row: &[u8] = &[0x02];
3301 let binary_rows = BinaryArray::from(vec![row]);
3302
3303 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3304 let parsed = converter.from_binary(binary_rows);
3305
3306 converter.convert_rows(parsed.iter()).unwrap();
3307 }
3308
3309 #[test]
3310 #[should_panic(expected = "rows were not produced by this RowConverter")]
3311 fn test_different_converter() {
3312 let values = Arc::new(Int32Array::from_iter([Some(1), Some(-1)]));
3313 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3314 let rows = converter.convert_columns(&[values]).unwrap();
3315
3316 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3317 let _ = converter.convert_rows(&rows);
3318 }
3319
3320 fn test_single_list<O: OffsetSizeTrait>() {
3321 let mut builder = GenericListBuilder::<O, _>::new(Int32Builder::new());
3322 builder.values().append_value(32);
3323 builder.values().append_value(52);
3324 builder.values().append_value(32);
3325 builder.append(true);
3326 builder.values().append_value(32);
3327 builder.values().append_value(52);
3328 builder.values().append_value(12);
3329 builder.append(true);
3330 builder.values().append_value(32);
3331 builder.values().append_value(52);
3332 builder.append(true);
3333 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3336 builder.values().append_value(32);
3337 builder.values().append_null();
3338 builder.append(true);
3339 builder.append(true);
3340 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3343
3344 let list = Arc::new(builder.finish()) as ArrayRef;
3345 let d = list.data_type().clone();
3346
3347 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3348
3349 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3350 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();
3359 assert_eq!(back.len(), 1);
3360 back[0].to_data().validate_full().unwrap();
3361 assert_eq!(&back[0], &list);
3362
3363 let options = SortOptions::default().asc().with_nulls_first(false);
3364 let field = SortField::new_with_options(d.clone(), options);
3365 let converter = RowConverter::new(vec![field]).unwrap();
3366 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3367
3368 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();
3377 assert_eq!(back.len(), 1);
3378 back[0].to_data().validate_full().unwrap();
3379 assert_eq!(&back[0], &list);
3380
3381 let options = SortOptions::default().desc().with_nulls_first(false);
3382 let field = SortField::new_with_options(d.clone(), options);
3383 let converter = RowConverter::new(vec![field]).unwrap();
3384 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3385
3386 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();
3395 assert_eq!(back.len(), 1);
3396 back[0].to_data().validate_full().unwrap();
3397 assert_eq!(&back[0], &list);
3398
3399 let options = SortOptions::default().desc().with_nulls_first(true);
3400 let field = SortField::new_with_options(d, options);
3401 let converter = RowConverter::new(vec![field]).unwrap();
3402 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3403
3404 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();
3413 assert_eq!(back.len(), 1);
3414 back[0].to_data().validate_full().unwrap();
3415 assert_eq!(&back[0], &list);
3416
3417 let sliced_list = list.slice(1, 5);
3418 let rows_on_sliced_list = converter
3419 .convert_columns(&[Arc::clone(&sliced_list)])
3420 .unwrap();
3421
3422 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();
3429 assert_eq!(back.len(), 1);
3430 back[0].to_data().validate_full().unwrap();
3431 assert_eq!(&back[0], &sliced_list);
3432 }
3433
3434 fn test_nested_list<O: OffsetSizeTrait>() {
3435 let mut builder =
3436 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(Int32Builder::new()));
3437
3438 builder.values().values().append_value(1);
3439 builder.values().values().append_value(2);
3440 builder.values().append(true);
3441 builder.values().values().append_value(1);
3442 builder.values().values().append_null();
3443 builder.values().append(true);
3444 builder.append(true);
3445
3446 builder.values().values().append_value(1);
3447 builder.values().values().append_null();
3448 builder.values().append(true);
3449 builder.values().values().append_value(1);
3450 builder.values().values().append_null();
3451 builder.values().append(true);
3452 builder.append(true);
3453
3454 builder.values().values().append_value(1);
3455 builder.values().values().append_null();
3456 builder.values().append(true);
3457 builder.values().append(false);
3458 builder.append(true);
3459 builder.append(false);
3460
3461 builder.values().values().append_value(1);
3462 builder.values().values().append_value(2);
3463 builder.values().append(true);
3464 builder.append(true);
3465
3466 let list = Arc::new(builder.finish()) as ArrayRef;
3467 let d = list.data_type().clone();
3468
3469 let options = SortOptions::default().asc().with_nulls_first(true);
3477 let field = SortField::new_with_options(d.clone(), options);
3478 let converter = RowConverter::new(vec![field]).unwrap();
3479 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3480
3481 assert!(rows.row(0) > rows.row(1));
3482 assert!(rows.row(1) > rows.row(2));
3483 assert!(rows.row(2) > rows.row(3));
3484 assert!(rows.row(4) < rows.row(0));
3485 assert!(rows.row(4) > rows.row(1));
3486
3487 let back = converter.convert_rows(&rows).unwrap();
3488 assert_eq!(back.len(), 1);
3489 back[0].to_data().validate_full().unwrap();
3490 assert_eq!(&back[0], &list);
3491
3492 let options = SortOptions::default().desc().with_nulls_first(true);
3493 let field = SortField::new_with_options(d.clone(), options);
3494 let converter = RowConverter::new(vec![field]).unwrap();
3495 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3496
3497 assert!(rows.row(0) > rows.row(1));
3498 assert!(rows.row(1) > rows.row(2));
3499 assert!(rows.row(2) > rows.row(3));
3500 assert!(rows.row(4) > rows.row(0));
3501 assert!(rows.row(4) > rows.row(1));
3502
3503 let back = converter.convert_rows(&rows).unwrap();
3504 assert_eq!(back.len(), 1);
3505 back[0].to_data().validate_full().unwrap();
3506 assert_eq!(&back[0], &list);
3507
3508 let options = SortOptions::default().desc().with_nulls_first(false);
3509 let field = SortField::new_with_options(d, options);
3510 let converter = RowConverter::new(vec![field]).unwrap();
3511 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3512
3513 assert!(rows.row(0) < rows.row(1));
3514 assert!(rows.row(1) < rows.row(2));
3515 assert!(rows.row(2) < rows.row(3));
3516 assert!(rows.row(4) > rows.row(0));
3517 assert!(rows.row(4) < rows.row(1));
3518
3519 let back = converter.convert_rows(&rows).unwrap();
3520 assert_eq!(back.len(), 1);
3521 back[0].to_data().validate_full().unwrap();
3522 assert_eq!(&back[0], &list);
3523
3524 let sliced_list = list.slice(1, 3);
3525 let rows = converter
3526 .convert_columns(&[Arc::clone(&sliced_list)])
3527 .unwrap();
3528
3529 assert!(rows.row(0) < rows.row(1));
3530 assert!(rows.row(1) < rows.row(2));
3531
3532 let back = converter.convert_rows(&rows).unwrap();
3533 assert_eq!(back.len(), 1);
3534 back[0].to_data().validate_full().unwrap();
3535 assert_eq!(&back[0], &sliced_list);
3536 }
3537
3538 #[test]
3539 fn test_list() {
3540 test_single_list::<i32>();
3541 test_nested_list::<i32>();
3542 }
3543
3544 #[test]
3545 fn test_large_list() {
3546 test_single_list::<i64>();
3547 test_nested_list::<i64>();
3548 }
3549
3550 fn test_single_list_view<O: OffsetSizeTrait>() {
3551 let mut builder = GenericListViewBuilder::<O, _>::new(Int32Builder::new());
3552 builder.values().append_value(32);
3553 builder.values().append_value(52);
3554 builder.values().append_value(32);
3555 builder.append(true);
3556 builder.values().append_value(32);
3557 builder.values().append_value(52);
3558 builder.values().append_value(12);
3559 builder.append(true);
3560 builder.values().append_value(32);
3561 builder.values().append_value(52);
3562 builder.append(true);
3563 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3566 builder.values().append_value(32);
3567 builder.values().append_null();
3568 builder.append(true);
3569 builder.append(true);
3570 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3573
3574 let list = Arc::new(builder.finish()) as ArrayRef;
3575 let d = list.data_type().clone();
3576
3577 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3578
3579 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3580 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();
3589 assert_eq!(back.len(), 1);
3590 back[0].to_data().validate_full().unwrap();
3591
3592 let back_list_view = back[0]
3594 .as_any()
3595 .downcast_ref::<GenericListViewArray<O>>()
3596 .unwrap();
3597 let orig_list_view = list
3598 .as_any()
3599 .downcast_ref::<GenericListViewArray<O>>()
3600 .unwrap();
3601
3602 assert_eq!(back_list_view.len(), orig_list_view.len());
3603 for i in 0..back_list_view.len() {
3604 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3605 if back_list_view.is_valid(i) {
3606 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3607 }
3608 }
3609
3610 let options = SortOptions::default().asc().with_nulls_first(false);
3611 let field = SortField::new_with_options(d.clone(), options);
3612 let converter = RowConverter::new(vec![field]).unwrap();
3613 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3614
3615 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();
3624 assert_eq!(back.len(), 1);
3625 back[0].to_data().validate_full().unwrap();
3626
3627 let options = SortOptions::default().desc().with_nulls_first(false);
3628 let field = SortField::new_with_options(d.clone(), options);
3629 let converter = RowConverter::new(vec![field]).unwrap();
3630 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3631
3632 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();
3641 assert_eq!(back.len(), 1);
3642 back[0].to_data().validate_full().unwrap();
3643
3644 let options = SortOptions::default().desc().with_nulls_first(true);
3645 let field = SortField::new_with_options(d, options);
3646 let converter = RowConverter::new(vec![field]).unwrap();
3647 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3648
3649 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();
3658 assert_eq!(back.len(), 1);
3659 back[0].to_data().validate_full().unwrap();
3660
3661 let sliced_list = list.slice(1, 5);
3662 let rows_on_sliced_list = converter
3663 .convert_columns(&[Arc::clone(&sliced_list)])
3664 .unwrap();
3665
3666 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();
3673 assert_eq!(back.len(), 1);
3674 back[0].to_data().validate_full().unwrap();
3675 }
3676
3677 fn test_nested_list_view<O: OffsetSizeTrait>() {
3678 let mut builder = GenericListViewBuilder::<O, _>::new(GenericListViewBuilder::<O, _>::new(
3679 Int32Builder::new(),
3680 ));
3681
3682 builder.values().values().append_value(1);
3684 builder.values().values().append_value(2);
3685 builder.values().append(true);
3686 builder.values().values().append_value(1);
3687 builder.values().values().append_null();
3688 builder.values().append(true);
3689 builder.append(true);
3690
3691 builder.values().values().append_value(1);
3693 builder.values().values().append_null();
3694 builder.values().append(true);
3695 builder.values().values().append_value(1);
3696 builder.values().values().append_null();
3697 builder.values().append(true);
3698 builder.append(true);
3699
3700 builder.values().values().append_value(1);
3702 builder.values().values().append_null();
3703 builder.values().append(true);
3704 builder.values().append(false);
3705 builder.append(true);
3706
3707 builder.append(false);
3709
3710 builder.values().values().append_value(1);
3712 builder.values().values().append_value(2);
3713 builder.values().append(true);
3714 builder.append(true);
3715
3716 let list = Arc::new(builder.finish()) as ArrayRef;
3717 let d = list.data_type().clone();
3718
3719 let options = SortOptions::default().asc().with_nulls_first(true);
3727 let field = SortField::new_with_options(d.clone(), options);
3728 let converter = RowConverter::new(vec![field]).unwrap();
3729 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3730
3731 assert!(rows.row(0) > rows.row(1));
3732 assert!(rows.row(1) > rows.row(2));
3733 assert!(rows.row(2) > rows.row(3));
3734 assert!(rows.row(4) < rows.row(0));
3735 assert!(rows.row(4) > rows.row(1));
3736
3737 let back = converter.convert_rows(&rows).unwrap();
3738 assert_eq!(back.len(), 1);
3739 back[0].to_data().validate_full().unwrap();
3740
3741 let back_list_view = back[0]
3743 .as_any()
3744 .downcast_ref::<GenericListViewArray<O>>()
3745 .unwrap();
3746 let orig_list_view = list
3747 .as_any()
3748 .downcast_ref::<GenericListViewArray<O>>()
3749 .unwrap();
3750
3751 assert_eq!(back_list_view.len(), orig_list_view.len());
3752 for i in 0..back_list_view.len() {
3753 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3754 if back_list_view.is_valid(i) {
3755 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3756 }
3757 }
3758
3759 let options = SortOptions::default().desc().with_nulls_first(true);
3760 let field = SortField::new_with_options(d.clone(), options);
3761 let converter = RowConverter::new(vec![field]).unwrap();
3762 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3763
3764 assert!(rows.row(0) > rows.row(1));
3765 assert!(rows.row(1) > rows.row(2));
3766 assert!(rows.row(2) > rows.row(3));
3767 assert!(rows.row(4) > rows.row(0));
3768 assert!(rows.row(4) > rows.row(1));
3769
3770 let back = converter.convert_rows(&rows).unwrap();
3771 assert_eq!(back.len(), 1);
3772 back[0].to_data().validate_full().unwrap();
3773
3774 let back_list_view = back[0]
3776 .as_any()
3777 .downcast_ref::<GenericListViewArray<O>>()
3778 .unwrap();
3779
3780 assert_eq!(back_list_view.len(), orig_list_view.len());
3781 for i in 0..back_list_view.len() {
3782 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3783 if back_list_view.is_valid(i) {
3784 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3785 }
3786 }
3787
3788 let options = SortOptions::default().desc().with_nulls_first(false);
3789 let field = SortField::new_with_options(d.clone(), options);
3790 let converter = RowConverter::new(vec![field]).unwrap();
3791 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3792
3793 assert!(rows.row(0) < rows.row(1));
3794 assert!(rows.row(1) < rows.row(2));
3795 assert!(rows.row(2) < rows.row(3));
3796 assert!(rows.row(4) > rows.row(0));
3797 assert!(rows.row(4) < rows.row(1));
3798
3799 let back = converter.convert_rows(&rows).unwrap();
3800 assert_eq!(back.len(), 1);
3801 back[0].to_data().validate_full().unwrap();
3802
3803 let back_list_view = back[0]
3805 .as_any()
3806 .downcast_ref::<GenericListViewArray<O>>()
3807 .unwrap();
3808
3809 assert_eq!(back_list_view.len(), orig_list_view.len());
3810 for i in 0..back_list_view.len() {
3811 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3812 if back_list_view.is_valid(i) {
3813 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3814 }
3815 }
3816
3817 let sliced_list = list.slice(1, 3);
3818 let rows = converter
3819 .convert_columns(&[Arc::clone(&sliced_list)])
3820 .unwrap();
3821
3822 assert!(rows.row(0) < rows.row(1));
3823 assert!(rows.row(1) < rows.row(2));
3824
3825 let back = converter.convert_rows(&rows).unwrap();
3826 assert_eq!(back.len(), 1);
3827 back[0].to_data().validate_full().unwrap();
3828 }
3829
3830 #[test]
3831 fn test_list_view() {
3832 test_single_list_view::<i32>();
3833 test_nested_list_view::<i32>();
3834 }
3835
3836 #[test]
3837 fn test_large_list_view() {
3838 test_single_list_view::<i64>();
3839 test_nested_list_view::<i64>();
3840 }
3841
3842 fn test_list_view_with_shared_values<O: OffsetSizeTrait>() {
3843 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
3845 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
3846
3847 let offsets = ScalarBuffer::<O>::from(vec![
3855 O::from_usize(0).unwrap(),
3856 O::from_usize(0).unwrap(),
3857 O::from_usize(5).unwrap(),
3858 O::from_usize(2).unwrap(),
3859 O::from_usize(1).unwrap(),
3860 O::from_usize(2).unwrap(),
3861 ]);
3862 let sizes = ScalarBuffer::<O>::from(vec![
3863 O::from_usize(3).unwrap(),
3864 O::from_usize(3).unwrap(),
3865 O::from_usize(2).unwrap(),
3866 O::from_usize(2).unwrap(),
3867 O::from_usize(4).unwrap(),
3868 O::from_usize(1).unwrap(),
3869 ]);
3870
3871 let list_view: GenericListViewArray<O> =
3872 GenericListViewArray::try_new(field, offsets, sizes, Arc::new(values), None).unwrap();
3873
3874 let d = list_view.data_type().clone();
3875 let list = Arc::new(list_view) as ArrayRef;
3876
3877 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3878 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3879
3880 assert_eq!(rows.row(0), rows.row(1));
3882
3883 assert!(rows.row(0) < rows.row(2));
3885
3886 assert!(rows.row(3) > rows.row(0));
3888
3889 assert!(rows.row(4) > rows.row(0));
3891
3892 assert!(rows.row(5) < rows.row(3));
3894
3895 assert!(rows.row(5) > rows.row(4));
3897
3898 let back = converter.convert_rows(&rows).unwrap();
3900 assert_eq!(back.len(), 1);
3901 back[0].to_data().validate_full().unwrap();
3902
3903 let back_list_view = back[0]
3905 .as_any()
3906 .downcast_ref::<GenericListViewArray<O>>()
3907 .unwrap();
3908 let orig_list_view = list
3909 .as_any()
3910 .downcast_ref::<GenericListViewArray<O>>()
3911 .unwrap();
3912
3913 assert_eq!(back_list_view.len(), orig_list_view.len());
3914 for i in 0..back_list_view.len() {
3915 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3916 if back_list_view.is_valid(i) {
3917 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3918 }
3919 }
3920
3921 let options = SortOptions::default().desc();
3923 let field = SortField::new_with_options(d, options);
3924 let converter = RowConverter::new(vec![field]).unwrap();
3925 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3926
3927 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();
3933 assert_eq!(back.len(), 1);
3934 back[0].to_data().validate_full().unwrap();
3935 }
3936
3937 #[test]
3938 fn test_list_view_shared_values() {
3939 test_list_view_with_shared_values::<i32>();
3940 }
3941
3942 #[test]
3943 fn test_large_list_view_shared_values() {
3944 test_list_view_with_shared_values::<i64>();
3945 }
3946
3947 #[test]
3948 fn test_fixed_size_list() {
3949 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
3950 builder.values().append_value(32);
3951 builder.values().append_value(52);
3952 builder.values().append_value(32);
3953 builder.append(true);
3954 builder.values().append_value(32);
3955 builder.values().append_value(52);
3956 builder.values().append_value(12);
3957 builder.append(true);
3958 builder.values().append_value(32);
3959 builder.values().append_value(52);
3960 builder.values().append_null();
3961 builder.append(true);
3962 builder.values().append_value(32); builder.values().append_value(52); builder.values().append_value(13); builder.append(false);
3966 builder.values().append_value(32);
3967 builder.values().append_null();
3968 builder.values().append_null();
3969 builder.append(true);
3970 builder.values().append_null();
3971 builder.values().append_null();
3972 builder.values().append_null();
3973 builder.append(true);
3974 builder.values().append_value(17); builder.values().append_null(); builder.values().append_value(77); builder.append(false);
3978
3979 let list = Arc::new(builder.finish()) as ArrayRef;
3980 let d = list.data_type().clone();
3981
3982 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3984
3985 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3986 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();
3995 assert_eq!(back.len(), 1);
3996 back[0].to_data().validate_full().unwrap();
3997 assert_eq!(&back[0], &list);
3998
3999 let options = SortOptions::default().asc().with_nulls_first(false);
4001 let field = SortField::new_with_options(d.clone(), options);
4002 let converter = RowConverter::new(vec![field]).unwrap();
4003 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4004 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();
4013 assert_eq!(back.len(), 1);
4014 back[0].to_data().validate_full().unwrap();
4015 assert_eq!(&back[0], &list);
4016
4017 let options = SortOptions::default().desc().with_nulls_first(false);
4019 let field = SortField::new_with_options(d.clone(), options);
4020 let converter = RowConverter::new(vec![field]).unwrap();
4021 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4022 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();
4031 assert_eq!(back.len(), 1);
4032 back[0].to_data().validate_full().unwrap();
4033 assert_eq!(&back[0], &list);
4034
4035 let options = SortOptions::default().desc().with_nulls_first(true);
4037 let field = SortField::new_with_options(d, options);
4038 let converter = RowConverter::new(vec![field]).unwrap();
4039 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4040
4041 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();
4050 assert_eq!(back.len(), 1);
4051 back[0].to_data().validate_full().unwrap();
4052 assert_eq!(&back[0], &list);
4053
4054 let sliced_list = list.slice(1, 5);
4055 let rows_on_sliced_list = converter
4056 .convert_columns(&[Arc::clone(&sliced_list)])
4057 .unwrap();
4058
4059 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();
4065 assert_eq!(back.len(), 1);
4066 back[0].to_data().validate_full().unwrap();
4067 assert_eq!(&back[0], &sliced_list);
4068 }
4069
4070 #[test]
4071 fn test_two_fixed_size_lists() {
4072 let mut first = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4073 first.values().append_value(100);
4075 first.append(true);
4076 first.values().append_value(101);
4078 first.append(true);
4079 first.values().append_value(102);
4081 first.append(true);
4082 first.values().append_null();
4084 first.append(true);
4085 first.values().append_null(); first.append(false);
4088 let first = Arc::new(first.finish()) as ArrayRef;
4089 let first_type = first.data_type().clone();
4090
4091 let mut second = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4092 second.values().append_value(200);
4094 second.append(true);
4095 second.values().append_value(201);
4097 second.append(true);
4098 second.values().append_value(202);
4100 second.append(true);
4101 second.values().append_null();
4103 second.append(true);
4104 second.values().append_null(); second.append(false);
4107 let second = Arc::new(second.finish()) as ArrayRef;
4108 let second_type = second.data_type().clone();
4109
4110 let converter = RowConverter::new(vec![
4111 SortField::new(first_type.clone()),
4112 SortField::new(second_type.clone()),
4113 ])
4114 .unwrap();
4115
4116 let rows = converter
4117 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4118 .unwrap();
4119
4120 let back = converter.convert_rows(&rows).unwrap();
4121 assert_eq!(back.len(), 2);
4122 back[0].to_data().validate_full().unwrap();
4123 assert_eq!(&back[0], &first);
4124 back[1].to_data().validate_full().unwrap();
4125 assert_eq!(&back[1], &second);
4126 }
4127
4128 #[test]
4129 fn test_fixed_size_list_with_variable_width_content() {
4130 let mut first = FixedSizeListBuilder::new(
4131 StructBuilder::from_fields(
4132 vec![
4133 Field::new(
4134 "timestamp",
4135 DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))),
4136 false,
4137 ),
4138 Field::new("offset_minutes", DataType::Int16, false),
4139 Field::new("time_zone", DataType::Utf8, false),
4140 ],
4141 1,
4142 ),
4143 1,
4144 );
4145 first
4147 .values()
4148 .field_builder::<TimestampMicrosecondBuilder>(0)
4149 .unwrap()
4150 .append_null();
4151 first
4152 .values()
4153 .field_builder::<Int16Builder>(1)
4154 .unwrap()
4155 .append_null();
4156 first
4157 .values()
4158 .field_builder::<StringBuilder>(2)
4159 .unwrap()
4160 .append_null();
4161 first.values().append(false);
4162 first.append(false);
4163 first
4165 .values()
4166 .field_builder::<TimestampMicrosecondBuilder>(0)
4167 .unwrap()
4168 .append_null();
4169 first
4170 .values()
4171 .field_builder::<Int16Builder>(1)
4172 .unwrap()
4173 .append_null();
4174 first
4175 .values()
4176 .field_builder::<StringBuilder>(2)
4177 .unwrap()
4178 .append_null();
4179 first.values().append(false);
4180 first.append(true);
4181 first
4183 .values()
4184 .field_builder::<TimestampMicrosecondBuilder>(0)
4185 .unwrap()
4186 .append_value(0);
4187 first
4188 .values()
4189 .field_builder::<Int16Builder>(1)
4190 .unwrap()
4191 .append_value(0);
4192 first
4193 .values()
4194 .field_builder::<StringBuilder>(2)
4195 .unwrap()
4196 .append_value("UTC");
4197 first.values().append(true);
4198 first.append(true);
4199 first
4201 .values()
4202 .field_builder::<TimestampMicrosecondBuilder>(0)
4203 .unwrap()
4204 .append_value(1126351800123456);
4205 first
4206 .values()
4207 .field_builder::<Int16Builder>(1)
4208 .unwrap()
4209 .append_value(120);
4210 first
4211 .values()
4212 .field_builder::<StringBuilder>(2)
4213 .unwrap()
4214 .append_value("Europe/Warsaw");
4215 first.values().append(true);
4216 first.append(true);
4217 let first = Arc::new(first.finish()) as ArrayRef;
4218 let first_type = first.data_type().clone();
4219
4220 let mut second = StringBuilder::new();
4221 second.append_value("somewhere near");
4222 second.append_null();
4223 second.append_value("Greenwich");
4224 second.append_value("Warsaw");
4225 let second = Arc::new(second.finish()) as ArrayRef;
4226 let second_type = second.data_type().clone();
4227
4228 let converter = RowConverter::new(vec![
4229 SortField::new(first_type.clone()),
4230 SortField::new(second_type.clone()),
4231 ])
4232 .unwrap();
4233
4234 let rows = converter
4235 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4236 .unwrap();
4237
4238 let back = converter.convert_rows(&rows).unwrap();
4239 assert_eq!(back.len(), 2);
4240 back[0].to_data().validate_full().unwrap();
4241 assert_eq!(&back[0], &first);
4242 back[1].to_data().validate_full().unwrap();
4243 assert_eq!(&back[1], &second);
4244 }
4245
4246 #[test]
4247 fn test_single_map() {
4248 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4249 builder.keys().append_value("hello");
4251 builder.values().append_value(1);
4252 builder.keys().append_value("world");
4253 builder.values().append_value(2);
4254 builder.append(true).unwrap();
4255
4256 builder.keys().append_value("foo");
4258 builder.values().append_value(3);
4259 builder.append(true).unwrap();
4260
4261 builder.append(true).unwrap();
4263
4264 builder.keys().append_value("masked_key");
4266 builder.values().append_value(999);
4267 builder.append(false).unwrap();
4268
4269 builder.append(false).unwrap();
4271
4272 builder.keys().append_value("bar");
4274 builder.values().append_null();
4275 builder.append(true).unwrap();
4276
4277 builder.keys().append_value("other_masked");
4279 builder.values().append_value(0);
4280 builder.append(false).unwrap();
4281
4282 builder.keys().append_value("a");
4284 builder.values().append_value(10);
4285 builder.keys().append_value("b");
4286 builder.values().append_value(20);
4287 builder.keys().append_value("c");
4288 builder.values().append_value(30);
4289 builder.append(true).unwrap();
4290
4291 let map = Arc::new(builder.finish()) as ArrayRef;
4292 let d = map.data_type().clone();
4293
4294 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4295
4296 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4297
4298 assert_eq!(rows.row(3), rows.row(4));
4300 assert_eq!(rows.row(4), rows.row(6));
4301
4302 let back = converter.convert_rows(&rows).unwrap();
4303 assert_eq!(back.len(), 1);
4304 back[0].to_data().validate_full().unwrap();
4305 assert_eq!(&back[0], &map);
4306
4307 let sliced_map = map.slice(1, map.len() - 2);
4308 let rows_on_sliced = converter
4309 .convert_columns(&[Arc::clone(&sliced_map)])
4310 .unwrap();
4311
4312 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4313 assert_eq!(back.len(), 1);
4314 back[0].to_data().validate_full().unwrap();
4315 assert_eq!(&back[0], &sliced_map);
4316 }
4317
4318 #[test]
4319 fn two_maps_with_different_keys_order_should_sort_by_entry_order() {
4320 let map_1: ArrayRef =
4321 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4322 vec![Some(vec![("hello", Some(1)), ("world", Some(2))])],
4323 false,
4324 ));
4325 let map_2: ArrayRef =
4327 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4328 vec![Some(vec![("world", Some(2)), ("hello", Some(1))])],
4329 false,
4330 ));
4331
4332 let converter = RowConverter::new(vec![SortField::new(map_1.data_type().clone())]).unwrap();
4333
4334 let map_1_rows = converter.convert_columns(&[Arc::clone(&map_1)]).unwrap();
4335 let map_2_rows = converter.convert_columns(&[Arc::clone(&map_2)]).unwrap();
4336
4337 assert_ne!(map_1_rows.row(0), map_2_rows.row(0));
4338 assert!(map_1_rows.row(0) < map_2_rows.row(0));
4339
4340 let back_1 = converter.convert_rows(&map_1_rows).unwrap();
4341 let back_2 = converter.convert_rows(&map_2_rows).unwrap();
4342 assert_eq!(&back_1[0], &map_1);
4343 assert_eq!(&back_2[0], &map_2);
4344 }
4345
4346 #[test]
4347 fn test_nested_map() {
4348 let mut builder = MapBuilder::new(
4350 None,
4351 StringBuilder::new(),
4352 MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()),
4353 );
4354
4355 builder.keys().append_value("outer1");
4357 builder.values().keys().append_value("inner_a");
4358 builder.values().values().append_value(1);
4359 builder.values().keys().append_value("inner_b");
4360 builder.values().values().append_value(2);
4361 builder.values().append(true).unwrap();
4362 builder.keys().append_value("outer2");
4363 builder.values().keys().append_value("inner_c");
4364 builder.values().values().append_value(3);
4365 builder.values().append(true).unwrap();
4366 builder.append(true).unwrap();
4367
4368 builder.keys().append_value("x");
4370 builder.values().append(true).unwrap();
4371 builder.append(true).unwrap();
4372
4373 builder.keys().append_value("y");
4375 builder.values().keys().append_value("masked"); builder.values().values().append_value(0); builder.values().append(false).unwrap();
4378 builder.append(true).unwrap();
4379
4380 builder.keys().append_value("y");
4382 builder.values().append(false).unwrap(); builder.append(true).unwrap();
4384
4385 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();
4391
4392 builder.keys().append_value("masked_outer"); builder.values().append(false).unwrap(); builder.append(false).unwrap();
4396
4397 builder.append(false).unwrap(); builder.append(true).unwrap();
4402
4403 let map = Arc::new(builder.finish()) as ArrayRef;
4404 let d = map.data_type().clone();
4405
4406 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4407
4408 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4409
4410 let back = converter.convert_rows(&rows).unwrap();
4411 assert_eq!(back.len(), 1);
4412 back[0].to_data().validate_full().unwrap();
4413 assert_eq!(&back[0], &map);
4414
4415 let sliced_map = map.slice(1, 3);
4416 let rows_on_sliced = converter
4417 .convert_columns(&[Arc::clone(&sliced_map)])
4418 .unwrap();
4419
4420 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4421 assert_eq!(back.len(), 1);
4422 back[0].to_data().validate_full().unwrap();
4423 assert_eq!(&back[0], &sliced_map);
4424 }
4425
4426 #[test]
4427 fn test_single_map_with_non_nullable_values() {
4428 let value_field = Arc::new(Field::new(
4430 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4431 DataType::Int32,
4432 false,
4433 ));
4434 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4435 .with_values_field(value_field);
4436 builder.keys().append_value("a");
4438 builder.values().append_value(1);
4439 builder.keys().append_value("b");
4440 builder.values().append_value(2);
4441 builder.append(true).unwrap();
4442 builder.append(false).unwrap();
4444 builder.keys().append_value("c");
4446 builder.values().append_value(3);
4447 builder.append(true).unwrap();
4448 builder.append(true).unwrap();
4450 builder.keys().append_value("masked"); builder.values().append_value(0); builder.append(false).unwrap();
4454
4455 let map = Arc::new(builder.finish()) as ArrayRef;
4456 let d = map.data_type().clone();
4457
4458 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4459
4460 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4461
4462 let back = converter.convert_rows(&rows).unwrap();
4463 assert_eq!(back.len(), 1);
4464 back[0].to_data().validate_full().unwrap();
4465 assert_eq!(&back[0], &map);
4466 }
4467
4468 #[test]
4469 fn test_single_map_with_non_nullable_map_but_with_nullable_values() {
4470 let value_field = Arc::new(Field::new(
4472 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4473 DataType::Int32,
4474 true,
4475 ));
4476 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4477 .with_values_field(value_field);
4478
4479 builder.keys().append_value("a");
4481 builder.values().append_value(1);
4482 builder.keys().append_value("b");
4483 builder.values().append_null();
4484 builder.append(true).unwrap();
4485 builder.keys().append_value("c");
4487 builder.values().append_null();
4488 builder.keys().append_value("d");
4489 builder.values().append_null();
4490 builder.append(true).unwrap();
4491 builder.append(true).unwrap();
4493 builder.keys().append_value("e");
4495 builder.values().append_value(5);
4496 builder.append(true).unwrap();
4497
4498 let map = Arc::new(builder.finish()) as ArrayRef;
4499 let d = map.data_type().clone();
4500
4501 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4502
4503 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4504
4505 let back = converter.convert_rows(&rows).unwrap();
4506 assert_eq!(back.len(), 1);
4507 back[0].to_data().validate_full().unwrap();
4508 assert_eq!(&back[0], &map);
4509 }
4510
4511 #[test]
4512 fn test_map_all_nulls() {
4513 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4514 builder.keys().append_value("m1"); builder.values().append_value(1); builder.append(false).unwrap();
4518 builder.keys().append_value("m2"); builder.values().append_value(2); builder.append(false).unwrap();
4521
4522 builder.append(false).unwrap(); builder.keys().append_value("m3"); builder.values().append_value(3); builder.append(false).unwrap();
4527
4528 let map = Arc::new(builder.finish()) as ArrayRef;
4529 let d = map.data_type().clone();
4530
4531 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4532
4533 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4534
4535 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4537
4538 let back = converter.convert_rows(&rows).unwrap();
4539 assert_eq!(back.len(), 1);
4540 back[0].to_data().validate_full().unwrap();
4541 assert_eq!(&back[0], &map);
4542 }
4543
4544 #[test]
4545 fn test_map_all_empty() {
4546 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4547 builder.append(true).unwrap();
4549 builder.append(true).unwrap();
4550 builder.append(true).unwrap();
4551
4552 let map = Arc::new(builder.finish()) as ArrayRef;
4553 let d = map.data_type().clone();
4554
4555 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4556
4557 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4558
4559 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4561
4562 let back = converter.convert_rows(&rows).unwrap();
4563 assert_eq!(back.len(), 1);
4564 back[0].to_data().validate_full().unwrap();
4565 assert_eq!(&back[0], &map);
4566 }
4567
4568 #[test]
4569 fn test_map_empty_array() {
4570 let builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4572 let map = Arc::new(builder.finish_cloned()) as ArrayRef;
4573 let d = map.data_type().clone();
4574
4575 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4576
4577 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4578
4579 let back = converter.convert_rows(&rows).unwrap();
4580 assert_eq!(back.len(), 1);
4581 back[0].to_data().validate_full().unwrap();
4582 assert_eq!(&back[0], &map);
4583 }
4584
4585 fn generate_primitive_array<K>(
4586 rng: &mut impl RngCore,
4587 len: usize,
4588 valid_percent: f64,
4589 ) -> PrimitiveArray<K>
4590 where
4591 K: ArrowPrimitiveType,
4592 StandardUniform: Distribution<K::Native>,
4593 {
4594 (0..len)
4595 .map(|_| rng.random_bool(valid_percent).then(|| rng.random()))
4596 .collect()
4597 }
4598
4599 fn generate_all_unique_primitive_array<K>(
4600 rng: &mut impl RngCore,
4601 len: usize,
4602 ) -> PrimitiveArray<K>
4603 where
4604 K: ArrowPrimitiveType,
4605 K::Native: Hash + Eq,
4606 StandardUniform: Distribution<K::Native>,
4607 {
4608 let possible_number_of_values = 2i32.saturating_pow(size_of::<K::Native>() as u32 * 8);
4609 assert!(
4610 len <= possible_number_of_values as usize,
4611 "len {len} is larger than the number of possible values {possible_number_of_values}"
4612 );
4613
4614 let mut seen = std::collections::HashSet::new();
4615 (0..len)
4616 .map(|_| {
4617 let mut value;
4618 loop {
4619 value = rng.random();
4620
4621 if seen.insert(value) {
4622 break;
4623 }
4624 }
4625
4626 Some(value)
4627 })
4628 .collect()
4629 }
4630
4631 fn generate_boolean_array(
4632 rng: &mut impl RngCore,
4633 len: usize,
4634 valid_percent: f64,
4635 ) -> BooleanArray {
4636 (0..len)
4637 .map(|_| rng.random_bool(valid_percent).then(|| rng.random_bool(0.5)))
4638 .collect()
4639 }
4640
4641 fn generate_strings<O: OffsetSizeTrait>(
4642 rng: &mut impl RngCore,
4643 len: usize,
4644 valid_percent: f64,
4645 ) -> GenericStringArray<O> {
4646 (0..len)
4647 .map(|_| {
4648 rng.random_bool(valid_percent).then(|| {
4649 let len = rng.random_range(0..100);
4650 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4651 String::from_utf8(bytes).unwrap()
4652 })
4653 })
4654 .collect()
4655 }
4656
4657 fn generate_string_view(
4658 rng: &mut impl RngCore,
4659 len: usize,
4660 valid_percent: f64,
4661 ) -> StringViewArray {
4662 (0..len)
4663 .map(|_| {
4664 rng.random_bool(valid_percent).then(|| {
4665 let len = rng.random_range(0..100);
4666 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4667 String::from_utf8(bytes).unwrap()
4668 })
4669 })
4670 .collect()
4671 }
4672
4673 fn generate_byte_view(
4674 rng: &mut impl RngCore,
4675 len: usize,
4676 valid_percent: f64,
4677 ) -> BinaryViewArray {
4678 (0..len)
4679 .map(|_| {
4680 rng.random_bool(valid_percent).then(|| {
4681 let len = rng.random_range(0..100);
4682 let bytes: Vec<_> = (0..len).map(|_| rng.random_range(0..128)).collect();
4683 bytes
4684 })
4685 })
4686 .collect()
4687 }
4688
4689 fn generate_fixed_stringview_column(len: usize) -> StringViewArray {
4690 let edge_cases = vec![
4691 Some("bar".to_string()),
4692 Some("bar\0".to_string()),
4693 Some("LongerThan12Bytes".to_string()),
4694 Some("LongerThan12Bytez".to_string()),
4695 Some("LongerThan12Bytes\0".to_string()),
4696 Some("LongerThan12Byt".to_string()),
4697 Some("backend one".to_string()),
4698 Some("backend two".to_string()),
4699 Some("a".repeat(257)),
4700 Some("a".repeat(300)),
4701 ];
4702
4703 let mut values = Vec::with_capacity(len);
4705 for i in 0..len {
4706 values.push(
4707 edge_cases
4708 .get(i % edge_cases.len())
4709 .cloned()
4710 .unwrap_or(None),
4711 );
4712 }
4713
4714 StringViewArray::from(values)
4715 }
4716
4717 fn generate_dictionary<K>(
4718 rng: &mut impl RngCore,
4719 values: ArrayRef,
4720 len: usize,
4721 valid_percent: f64,
4722 ) -> DictionaryArray<K>
4723 where
4724 K: ArrowDictionaryKeyType,
4725 K::Native: SampleUniform,
4726 {
4727 let min_key = K::Native::from_usize(0).unwrap();
4728 let max_key = K::Native::from_usize(values.len()).unwrap();
4729 let keys: PrimitiveArray<K> = (0..len)
4730 .map(|_| {
4731 rng.random_bool(valid_percent)
4732 .then(|| rng.random_range(min_key..max_key))
4733 })
4734 .collect();
4735
4736 let data_type =
4737 DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
4738
4739 let data = keys
4740 .into_data()
4741 .into_builder()
4742 .data_type(data_type)
4743 .add_child_data(values.to_data())
4744 .build()
4745 .unwrap();
4746
4747 DictionaryArray::from(data)
4748 }
4749
4750 fn generate_fixed_size_binary(
4751 rng: &mut impl RngCore,
4752 len: usize,
4753 valid_percent: f64,
4754 ) -> FixedSizeBinaryArray {
4755 let width = rng.random_range(0..20);
4756 let mut builder = FixedSizeBinaryBuilder::new(width);
4757
4758 let mut b = vec![0; width as usize];
4759 for _ in 0..len {
4760 match rng.random_bool(valid_percent) {
4761 true => {
4762 b.iter_mut().for_each(|x| *x = rng.random());
4763 builder.append_value(&b).unwrap();
4764 }
4765 false => builder.append_null(),
4766 }
4767 }
4768
4769 builder.finish()
4770 }
4771
4772 fn generate_struct(rng: &mut impl RngCore, len: usize, valid_percent: f64) -> StructArray {
4773 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4774 let a = generate_primitive_array::<Int32Type>(rng, len, valid_percent);
4775 let b = generate_strings::<i32>(rng, len, valid_percent);
4776 let fields = Fields::from(vec![
4777 Field::new("a", DataType::Int32, true),
4778 Field::new("b", DataType::Utf8, true),
4779 ]);
4780 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
4781 StructArray::new(fields, values, Some(nulls))
4782 }
4783
4784 fn generate_list<R: RngCore, F>(
4785 rng: &mut R,
4786 len: usize,
4787 valid_percent: f64,
4788 values: F,
4789 ) -> ListArray
4790 where
4791 F: FnOnce(&mut R, usize) -> ArrayRef,
4792 {
4793 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4794 let values_len = offsets.last().unwrap().to_usize().unwrap();
4795 let values = values(rng, values_len);
4796 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4797 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4798 ListArray::new(field, offsets, values, Some(nulls))
4799 }
4800
4801 fn generate_list_view<F>(
4802 rng: &mut impl RngCore,
4803 len: usize,
4804 valid_percent: f64,
4805 values: F,
4806 ) -> ListViewArray
4807 where
4808 F: FnOnce(usize) -> ArrayRef,
4809 {
4810 let sizes: Vec<i32> = (0..len).map(|_| rng.random_range(0..10)).collect();
4812 let values_len: usize = sizes.iter().map(|s| *s as usize).sum::<usize>().max(1);
4813 let values = values(values_len);
4814
4815 let offsets: Vec<i32> = sizes
4817 .iter()
4818 .map(|&size| {
4819 if size == 0 {
4820 0
4821 } else {
4822 rng.random_range(0..=(values_len as i32 - size))
4823 }
4824 })
4825 .collect();
4826
4827 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4828 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4829 ListViewArray::new(
4830 field,
4831 ScalarBuffer::from(offsets),
4832 ScalarBuffer::from(sizes),
4833 values,
4834 Some(nulls),
4835 )
4836 }
4837
4838 fn generate_map<R: RngCore, KeysFn, ValuesFn>(
4839 rng: &mut R,
4840 len: usize,
4841 valid_percent: f64,
4842 gen_keys: KeysFn,
4843 gen_values: ValuesFn,
4844 ) -> MapArray
4845 where
4846 KeysFn: FnOnce(&mut R, usize) -> ArrayRef,
4847 ValuesFn: FnOnce(&mut R, usize) -> ArrayRef,
4848 {
4849 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4850 let entries_len = offsets.last().unwrap().to_usize().unwrap();
4851 let keys = gen_keys(rng, entries_len);
4852 let values = gen_values(rng, entries_len);
4853 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4854 let field = Arc::new(Field::new_map(
4855 "",
4856 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4857 Field::new(
4858 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4859 keys.data_type().clone(),
4860 false,
4861 ),
4862 Field::new(
4863 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4864 values.data_type().clone(),
4865 true,
4866 ),
4867 false,
4868 true,
4869 ));
4870 let DataType::Map(struct_field, _) = field.data_type() else {
4871 unreachable!();
4872 };
4873
4874 let DataType::Struct(fields) = struct_field.data_type() else {
4875 unreachable!();
4876 };
4877
4878 let entries = StructArray::new(fields.clone(), vec![keys, values], None);
4879
4880 let map_array = MapArray::new(struct_field.clone(), offsets, entries, Some(nulls), false);
4881
4882 assert_valid_map(&map_array);
4883
4884 map_array
4885 }
4886
4887 fn assert_valid_map(array: &MapArray) {
4897 let keys_arrow_row_converter =
4898 RowConverter::new(vec![SortField::new(array.key_type().clone())]).unwrap();
4899
4900 array.iter().enumerate().flat_map(|(index, entry)| entry.map(|entry| (index, Arc::clone(entry.column(0))))).for_each(|(entry_index, keys)| {
4901 let keys_as_rows = keys_arrow_row_converter.convert_columns(&[Arc::clone(&keys)]).expect("should be able to convert keys");
4902
4903 for i in 0..keys_as_rows.num_rows() {
4904 for j in (i + 1)..keys_as_rows.num_rows() {
4905 if keys_as_rows.row(i) == keys_as_rows.row(j) {
4906 let key_i = keys.slice(i, 1);
4907 let key_j = keys.slice(j, 1);
4908
4909 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:?}");
4910 }
4911 }
4912 }
4913 })
4914 }
4915
4916 fn generate_nulls(rng: &mut impl RngCore, len: usize) -> Option<NullBuffer> {
4917 Some(NullBuffer::from_iter(
4918 (0..len).map(|_| rng.random_bool(0.8)),
4919 ))
4920 }
4921
4922 fn change_underlying_null_values_for_primitive<T: ArrowPrimitiveType>(
4923 array: &PrimitiveArray<T>,
4924 ) -> PrimitiveArray<T> {
4925 let (dt, values, nulls) = array.clone().into_parts();
4926
4927 let new_values = ScalarBuffer::<T::Native>::from_iter(
4928 values
4929 .iter()
4930 .zip(nulls.as_ref().unwrap().iter())
4931 .map(|(val, is_valid)| {
4932 if is_valid {
4933 *val
4934 } else {
4935 val.add_wrapping(T::Native::usize_as(1))
4936 }
4937 }),
4938 );
4939
4940 PrimitiveArray::new(new_values, nulls).with_data_type(dt)
4941 }
4942
4943 fn change_underline_null_values_for_byte_array<T: ByteArrayType>(
4944 array: &GenericByteArray<T>,
4945 ) -> GenericByteArray<T> {
4946 let (offsets, values, nulls) = array.clone().into_parts();
4947
4948 let new_offsets = OffsetBuffer::<T::Offset>::from_lengths(
4949 offsets
4950 .lengths()
4951 .zip(nulls.as_ref().unwrap().iter())
4952 .map(|(len, is_valid)| if is_valid { len } else { len + 1 }),
4953 );
4954
4955 let mut new_bytes = Vec::<u8>::with_capacity(new_offsets[new_offsets.len() - 1].as_usize());
4956
4957 offsets
4958 .windows(2)
4959 .zip(nulls.as_ref().unwrap().iter())
4960 .for_each(|(start_and_end, is_valid)| {
4961 let start = start_and_end[0].as_usize();
4962 let end = start_and_end[1].as_usize();
4963 new_bytes.extend_from_slice(&values.as_slice()[start..end]);
4964
4965 if !is_valid {
4967 new_bytes.push(b'c');
4968 }
4969 });
4970
4971 GenericByteArray::<T>::new(new_offsets, Buffer::from_vec(new_bytes), nulls)
4972 }
4973
4974 fn change_underline_null_values_for_list_array<O: OffsetSizeTrait>(
4975 array: &GenericListArray<O>,
4976 ) -> GenericListArray<O> {
4977 let (field, offsets, values, nulls) = array.clone().into_parts();
4978
4979 let (new_values, new_offsets) = {
4980 let concat_values = offsets
4981 .windows(2)
4982 .zip(nulls.as_ref().unwrap().iter())
4983 .map(|(start_and_end, is_valid)| {
4984 let start = start_and_end[0].as_usize();
4985 let end = start_and_end[1].as_usize();
4986 if is_valid {
4987 return (start, end - start);
4988 }
4989
4990 if end == values.len() {
4992 (start, (end - start).saturating_sub(1))
4993 } else {
4994 (start, end - start + 1)
4995 }
4996 })
4997 .map(|(start, length)| values.slice(start, length))
4998 .collect::<Vec<_>>();
4999
5000 let new_offsets =
5001 OffsetBuffer::<O>::from_lengths(concat_values.iter().map(|s| s.len()));
5002
5003 let new_values = {
5004 let values = concat_values.iter().map(|a| a.as_ref()).collect::<Vec<_>>();
5005 arrow_select::concat::concat(&values).expect("should be able to concat")
5006 };
5007
5008 (new_values, new_offsets)
5009 };
5010
5011 GenericListArray::<O>::new(field, new_offsets, new_values, nulls)
5012 }
5013
5014 fn change_underline_null_values_for_map_array(array: &MapArray) -> MapArray {
5015 let (field, offsets, entries, nulls, ordered) = array.clone().into_parts();
5016 assert!(
5017 !ordered,
5018 "can't replace underlying null values for ordered map array as this can violate the ordering"
5019 );
5020
5021 let (new_entries, new_offsets) = {
5022 let concat_values = offsets
5023 .windows(2)
5024 .zip(nulls.as_ref().unwrap().iter())
5025 .map(|(start_and_end, is_valid)| {
5026 let start = start_and_end[0].as_usize();
5027 let end = start_and_end[1].as_usize();
5028 if is_valid {
5029 return (start, end - start);
5030 }
5031
5032 if end == entries.len() {
5034 (start, (end - start).saturating_sub(1))
5035 } else {
5036 (start, end - start + 1)
5038 }
5039 })
5040 .map(|(start, length)| entries.slice(start, length))
5041 .collect::<Vec<_>>();
5042
5043 let new_offsets = OffsetBuffer::from_lengths(concat_values.iter().map(|s| s.len()));
5044
5045 let new_values = {
5046 let values = concat_values
5047 .iter()
5048 .map(|a| a as &dyn Array)
5049 .collect::<Vec<_>>();
5050 arrow_select::concat::concat(&values).expect("should be able to concat")
5051 };
5052
5053 (new_values.as_struct().clone(), new_offsets)
5054 };
5055
5056 let new_map = MapArray::new(field, new_offsets, new_entries, nulls, ordered);
5057
5058 assert_valid_map(&new_map);
5059
5060 new_map
5061 }
5062
5063 fn change_underline_null_values(array: &ArrayRef) -> ArrayRef {
5064 if array.null_count() == 0 {
5065 return Arc::clone(array);
5066 }
5067
5068 downcast_primitive_array!(
5069 array => {
5070 let output = change_underlying_null_values_for_primitive(array);
5071
5072 Arc::new(output)
5073 }
5074
5075 DataType::Utf8 => {
5076 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i32>()))
5077 }
5078 DataType::LargeUtf8 => {
5079 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i64>()))
5080 }
5081 DataType::Binary => {
5082 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i32>()))
5083 }
5084 DataType::LargeBinary => {
5085 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i64>()))
5086 }
5087 DataType::List(_) => {
5088 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i32>()))
5089 }
5090 DataType::LargeList(_) => {
5091 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i64>()))
5092 }
5093 DataType::Map(_, _) => {
5094 Arc::new(change_underline_null_values_for_map_array(array.as_map()))
5095 }
5096 _ => {
5097 Arc::clone(array)
5098 }
5099 )
5100 }
5101
5102 fn generate_column(rng: &mut (impl RngCore + Clone), len: usize) -> ArrayRef {
5103 match rng.random_range(0..24) {
5104 0 => Arc::new(generate_primitive_array::<Int32Type>(rng, len, 0.8)),
5105 1 => Arc::new(generate_primitive_array::<UInt32Type>(rng, len, 0.8)),
5106 2 => Arc::new(generate_primitive_array::<Int64Type>(rng, len, 0.8)),
5107 3 => Arc::new(generate_primitive_array::<UInt64Type>(rng, len, 0.8)),
5108 4 => Arc::new(generate_primitive_array::<Float32Type>(rng, len, 0.8)),
5109 5 => Arc::new(generate_primitive_array::<Float64Type>(rng, len, 0.8)),
5110 6 => Arc::new(generate_strings::<i32>(rng, len, 0.8)),
5111 7 => {
5112 let dict_values_len = rng.random_range(1..len);
5113 let strings = Arc::new(generate_strings::<i32>(rng, dict_values_len, 1.0));
5115 Arc::new(generate_dictionary::<Int64Type>(rng, strings, len, 0.8))
5116 }
5117 8 => {
5118 let dict_values_len = rng.random_range(1..len);
5119 let values = Arc::new(generate_primitive_array::<Int64Type>(
5121 rng,
5122 dict_values_len,
5123 1.0,
5124 ));
5125 Arc::new(generate_dictionary::<Int64Type>(rng, values, len, 0.8))
5126 }
5127 9 => Arc::new(generate_fixed_size_binary(rng, len, 0.8)),
5128 10 => Arc::new(generate_struct(rng, len, 0.8)),
5129 11 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5130 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5131 })),
5132 12 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5133 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5134 })),
5135 13 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5136 Arc::new(generate_struct(rng, values_len, 0.8))
5137 })),
5138 14 => Arc::new(generate_string_view(rng, len, 0.8)),
5139 15 => Arc::new(generate_byte_view(rng, len, 0.8)),
5140 16 => Arc::new(generate_fixed_stringview_column(len)),
5141 17 => Arc::new(
5142 generate_list(&mut rng.clone(), len + 1000, 0.8, |rng, values_len| {
5143 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5144 })
5145 .slice(500, len),
5146 ),
5147 18 => Arc::new(generate_boolean_array(rng, len, 0.8)),
5148 19 => Arc::new(generate_list_view(
5149 &mut rng.clone(),
5150 len,
5151 0.8,
5152 |values_len| Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8)),
5153 )),
5154 20 => Arc::new(generate_list_view(
5155 &mut rng.clone(),
5156 len,
5157 0.8,
5158 |values_len| Arc::new(generate_strings::<i32>(rng, values_len, 0.8)),
5159 )),
5160 21 => Arc::new(generate_list_view(
5161 &mut rng.clone(),
5162 len,
5163 0.8,
5164 |values_len| Arc::new(generate_struct(rng, values_len, 0.8)),
5165 )),
5166 22 => Arc::new(
5167 generate_list_view(&mut rng.clone(), len + 1000, 0.8, |values_len| {
5168 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5169 })
5170 .slice(500, len),
5171 ),
5172 23 => Arc::new(generate_map(
5173 rng,
5174 len,
5175 0.9,
5176 |rng, keys_len| {
5178 Arc::new(generate_all_unique_primitive_array::<Int64Type>(
5179 rng, keys_len,
5180 ))
5181 },
5182 |rng, values_len| Arc::new(generate_strings::<i32>(rng, values_len, 0.7)),
5183 )),
5184 _ => unreachable!(),
5185 }
5186 }
5187
5188 fn print_row(cols: &[SortColumn], row: usize) -> String {
5189 let t: Vec<_> = cols
5190 .iter()
5191 .map(|x| match x.values.is_valid(row) {
5192 true => {
5193 let opts = FormatOptions::default().with_null("NULL");
5194 let formatter = ArrayFormatter::try_new(x.values.as_ref(), &opts).unwrap();
5195 formatter.value(row).to_string()
5196 }
5197 false => "NULL".to_string(),
5198 })
5199 .collect();
5200 t.join(",")
5201 }
5202
5203 fn print_col_types(cols: &[SortColumn]) -> String {
5204 let t: Vec<_> = cols
5205 .iter()
5206 .map(|x| x.values.data_type().to_string())
5207 .collect();
5208 t.join(",")
5209 }
5210
5211 #[derive(Debug, PartialEq)]
5212 enum Nulls {
5213 AsIs,
5215
5216 Different,
5218
5219 None,
5221 }
5222
5223 #[test]
5224 #[cfg_attr(miri, ignore)]
5225 fn fuzz_test() {
5226 let mut rng = StdRng::seed_from_u64(42);
5227 for _ in 0..100 {
5228 for null_behavior in [Nulls::AsIs, Nulls::Different, Nulls::None] {
5229 let num_columns = rng.random_range(1..5);
5230 let len = rng.random_range(5..100);
5231 let mut arrays: Vec<_> = (0..num_columns)
5232 .map(|_| generate_column(&mut rng, len))
5233 .collect();
5234
5235 match null_behavior {
5236 Nulls::AsIs => {
5237 }
5239 Nulls::Different => {
5240 arrays = arrays
5242 .into_iter()
5243 .map(|a| replace_array_nulls(a, generate_nulls(&mut rng, len)))
5244 .collect()
5245 }
5246 Nulls::None => {
5247 arrays = arrays
5249 .into_iter()
5250 .map(|a| replace_array_nulls(a, None))
5251 .collect()
5252 }
5253 }
5254
5255 let options: Vec<_> = (0..num_columns)
5256 .map(|_| SortOptions {
5257 descending: rng.random_bool(0.5),
5258 nulls_first: rng.random_bool(0.5),
5259 })
5260 .collect();
5261
5262 let sort_columns: Vec<_> = options
5263 .iter()
5264 .zip(&arrays)
5265 .map(|(o, c)| SortColumn {
5266 values: Arc::clone(c),
5267 options: Some(*o),
5268 })
5269 .collect();
5270
5271 let comparator = LexicographicalComparator::try_new(&sort_columns).unwrap();
5272
5273 let columns: Vec<SortField> = options
5274 .into_iter()
5275 .zip(&arrays)
5276 .map(|(o, a)| SortField::new_with_options(a.data_type().clone(), o))
5277 .collect();
5278
5279 let converter = RowConverter::new(columns).unwrap();
5280 let rows = converter.convert_columns(&arrays).unwrap();
5281
5282 if !matches!(null_behavior, Nulls::None) {
5285 assert_same_rows_when_changing_input_underlying_null_values(
5286 &arrays, &converter, &rows,
5287 );
5288 }
5289
5290 for i in 0..len {
5291 for j in 0..len {
5292 let row_i = rows.row(i);
5293 let row_j = rows.row(j);
5294 let row_cmp = row_i.cmp(&row_j);
5295 let lex_cmp = comparator.compare(i, j);
5296 assert_eq!(
5297 row_cmp,
5298 lex_cmp,
5299 "({:?} vs {:?}) vs ({:?} vs {:?}) for types {}",
5300 print_row(&sort_columns, i),
5301 print_row(&sort_columns, j),
5302 row_i,
5303 row_j,
5304 print_col_types(&sort_columns)
5305 );
5306 }
5307 }
5308
5309 {
5311 let mut rows_iter = rows.iter();
5312 let mut rows_lengths_iter = rows.lengths();
5313 for (index, row) in rows_iter.by_ref().enumerate() {
5314 let len = rows_lengths_iter
5315 .next()
5316 .expect("Reached end of length iterator while still have rows");
5317 assert_eq!(
5318 row.data.len(),
5319 len,
5320 "Row length mismatch: {} vs {}",
5321 row.data.len(),
5322 len
5323 );
5324 assert_eq!(
5325 len,
5326 rows.row_len(index),
5327 "Row length mismatch at index {}: {} vs {}",
5328 index,
5329 len,
5330 rows.row_len(index)
5331 );
5332 }
5333
5334 assert_eq!(
5335 rows_lengths_iter.next(),
5336 None,
5337 "Length iterator did not reach end"
5338 );
5339 }
5340
5341 let back = converter.convert_rows(&rows).unwrap();
5344 for (actual, expected) in back.iter().zip(&arrays) {
5345 actual.to_data().validate_full().unwrap();
5346 dictionary_eq(actual, expected)
5347 }
5348
5349 let rows = rows.try_into_binary().expect("reasonable size");
5352 let parser = converter.parser();
5353 let back = converter
5354 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5355 .unwrap();
5356 for (actual, expected) in back.iter().zip(&arrays) {
5357 actual.to_data().validate_full().unwrap();
5358 dictionary_eq(actual, expected)
5359 }
5360
5361 let rows = converter.from_binary(rows);
5362 let back = converter.convert_rows(&rows).unwrap();
5363 for (actual, expected) in back.iter().zip(&arrays) {
5364 actual.to_data().validate_full().unwrap();
5365 dictionary_eq(actual, expected)
5366 }
5367 }
5368 }
5369 }
5370
5371 fn replace_array_nulls(array: ArrayRef, new_nulls: Option<NullBuffer>) -> ArrayRef {
5372 make_array(
5373 array
5374 .into_data()
5375 .into_builder()
5376 .nulls(new_nulls)
5378 .build()
5379 .unwrap(),
5380 )
5381 }
5382
5383 fn assert_same_rows_when_changing_input_underlying_null_values(
5384 arrays: &[ArrayRef],
5385 converter: &RowConverter,
5386 rows: &Rows,
5387 ) {
5388 let arrays_with_different_data_behind_nulls = arrays
5389 .iter()
5390 .map(|arr| change_underline_null_values(arr))
5391 .collect::<Vec<_>>();
5392
5393 if arrays
5395 .iter()
5396 .zip(arrays_with_different_data_behind_nulls.iter())
5397 .all(|(a, b)| Arc::ptr_eq(a, b))
5398 {
5399 return;
5400 }
5401
5402 let rows_with_different_nulls = converter
5403 .convert_columns(&arrays_with_different_data_behind_nulls)
5404 .unwrap();
5405
5406 assert_eq!(
5407 rows.iter().collect::<Vec<_>>(),
5408 rows_with_different_nulls.iter().collect::<Vec<_>>(),
5409 "Different underlying nulls should not output different rows"
5410 )
5411 }
5412
5413 #[test]
5414 fn test_clear() {
5415 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
5416 let mut rows = converter.empty_rows(3, 128);
5417
5418 let first = Int32Array::from(vec![None, Some(2), Some(4)]);
5419 let second = Int32Array::from(vec![Some(2), None, Some(4)]);
5420 let arrays = [Arc::new(first) as ArrayRef, Arc::new(second) as ArrayRef];
5421
5422 for array in arrays.iter() {
5423 rows.clear();
5424 converter
5425 .append(&mut rows, std::slice::from_ref(array))
5426 .unwrap();
5427 let back = converter.convert_rows(&rows).unwrap();
5428 assert_eq!(&back[0], array);
5429 }
5430
5431 let mut rows_expected = converter.empty_rows(3, 128);
5432 converter.append(&mut rows_expected, &arrays[1..]).unwrap();
5433
5434 for (i, (actual, expected)) in rows.iter().zip(rows_expected.iter()).enumerate() {
5435 assert_eq!(
5436 actual, expected,
5437 "For row {i}: expected {expected:?}, actual: {actual:?}",
5438 );
5439 }
5440 }
5441
5442 #[test]
5443 fn test_append_codec_dictionary_binary() {
5444 use DataType::*;
5445 let converter = RowConverter::new(vec![SortField::new(Dictionary(
5447 Box::new(Int32),
5448 Box::new(Binary),
5449 ))])
5450 .unwrap();
5451 let mut rows = converter.empty_rows(4, 128);
5452
5453 let keys = Int32Array::from_iter_values([0, 1, 2, 3]);
5454 let values = BinaryArray::from(vec![
5455 Some("a".as_bytes()),
5456 Some(b"b"),
5457 Some(b"c"),
5458 Some(b"d"),
5459 ]);
5460 let dict_array = DictionaryArray::new(keys, Arc::new(values));
5461
5462 rows.clear();
5463 let array = Arc::new(dict_array) as ArrayRef;
5464 converter
5465 .append(&mut rows, std::slice::from_ref(&array))
5466 .unwrap();
5467 let back = converter.convert_rows(&rows).unwrap();
5468
5469 dictionary_eq(&back[0], &array);
5470 }
5471
5472 #[test]
5473 fn test_list_prefix() {
5474 let mut a = ListBuilder::new(Int8Builder::new());
5475 a.append_value([None]);
5476 a.append_value([None, None]);
5477 let a = a.finish();
5478
5479 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
5480 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
5481 assert_eq!(rows.row(0).cmp(&rows.row(1)), Ordering::Less);
5482 }
5483
5484 #[test]
5485 fn test_utf8_validation_doesnt_affect_values_buffer_size() {
5486 fn assert_values_buffer_lens(col: ArrayRef) -> usize {
5487 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8View)]).unwrap();
5489
5490 let rows = converter.convert_columns(&[col]).unwrap();
5492 let converted = converter.convert_rows(&rows).unwrap();
5493 let unchecked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5494
5495 let rows = rows.try_into_binary().expect("reasonable size");
5497 let parser = converter.parser();
5498 let converted = converter
5499 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5500 .unwrap();
5501 let checked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5502 assert_eq!(unchecked_values_len, checked_values_len);
5504 checked_values_len
5505 }
5506
5507 let col = Arc::new(StringViewArray::from_iter([
5509 Some("hello"), None, Some("short"), Some("tiny"), ])) as ArrayRef;
5514
5515 let values_len = assert_values_buffer_lens(col);
5516 assert_eq!(values_len, 0);
5518
5519 let col = Arc::new(StringViewArray::from_iter([
5521 Some("1234567890123"), Some("12345678901234"), ])) as ArrayRef;
5524
5525 let values_len = assert_values_buffer_lens(col);
5526 assert_eq!(values_len, 13 + 14);
5527
5528 let col = Arc::new(StringViewArray::from_iter([
5530 Some("tiny"), Some("thisisexact13"), None,
5533 Some("short"), ])) as ArrayRef;
5535
5536 let values_len = assert_values_buffer_lens(col);
5537 assert_eq!(values_len, 13);
5539 }
5540
5541 #[test]
5542 fn test_sparse_union() {
5543 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5545 let str_array = StringArray::from(vec![None, Some("b"), None, Some("d"), None]);
5546
5547 let type_ids = vec![0, 1, 0, 1, 0].into();
5549
5550 let union_fields = [
5551 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5552 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5553 ]
5554 .into_iter()
5555 .collect();
5556
5557 let union_array = UnionArray::try_new(
5558 union_fields,
5559 type_ids,
5560 None,
5561 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5562 )
5563 .unwrap();
5564
5565 let union_type = union_array.data_type().clone();
5566 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5567
5568 let rows = converter
5569 .convert_columns(&[Arc::new(union_array.clone())])
5570 .unwrap();
5571
5572 let back = converter.convert_rows(&rows).unwrap();
5574 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5575
5576 assert_eq!(union_array.len(), back_union.len());
5577 for i in 0..union_array.len() {
5578 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5579 }
5580 }
5581
5582 #[test]
5583 fn test_sparse_union_with_nulls() {
5584 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5586 let str_array = StringArray::from(vec![None::<&str>; 5]);
5587
5588 let type_ids = vec![0, 1, 0, 1, 0].into();
5590
5591 let union_fields = [
5592 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5593 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5594 ]
5595 .into_iter()
5596 .collect();
5597
5598 let union_array = UnionArray::try_new(
5599 union_fields,
5600 type_ids,
5601 None,
5602 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5603 )
5604 .unwrap();
5605
5606 let union_type = union_array.data_type().clone();
5607 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5608
5609 let rows = converter
5610 .convert_columns(&[Arc::new(union_array.clone())])
5611 .unwrap();
5612
5613 let back = converter.convert_rows(&rows).unwrap();
5615 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5616
5617 assert_eq!(union_array.len(), back_union.len());
5618 for i in 0..union_array.len() {
5619 let expected_null = union_array.is_null(i);
5620 let actual_null = back_union.is_null(i);
5621 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5622 if !expected_null {
5623 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5624 }
5625 }
5626 }
5627
5628 #[test]
5629 fn test_dense_union() {
5630 let int_array = Int32Array::from(vec![1, 3, 5]);
5632 let str_array = StringArray::from(vec!["a", "b"]);
5633
5634 let type_ids = vec![0, 1, 0, 1, 0].into();
5635
5636 let offsets = vec![0, 0, 1, 1, 2].into();
5638
5639 let union_fields = [
5640 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5641 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5642 ]
5643 .into_iter()
5644 .collect();
5645
5646 let union_array = UnionArray::try_new(
5647 union_fields,
5648 type_ids,
5649 Some(offsets), vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5651 )
5652 .unwrap();
5653
5654 let union_type = union_array.data_type().clone();
5655 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5656
5657 let rows = converter
5658 .convert_columns(&[Arc::new(union_array.clone())])
5659 .unwrap();
5660
5661 let back = converter.convert_rows(&rows).unwrap();
5663 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5664
5665 assert_eq!(union_array.len(), back_union.len());
5666 for i in 0..union_array.len() {
5667 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5668 }
5669 }
5670
5671 #[test]
5672 fn test_dense_union_with_nulls() {
5673 let int_array = Int32Array::from(vec![Some(1), None, Some(5)]);
5675 let str_array = StringArray::from(vec![Some("a"), None]);
5676
5677 let type_ids = vec![0, 1, 0, 1, 0].into();
5679 let offsets = vec![0, 0, 1, 1, 2].into();
5680
5681 let union_fields = [
5682 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5683 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5684 ]
5685 .into_iter()
5686 .collect();
5687
5688 let union_array = UnionArray::try_new(
5689 union_fields,
5690 type_ids,
5691 Some(offsets),
5692 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5693 )
5694 .unwrap();
5695
5696 let union_type = union_array.data_type().clone();
5697 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5698
5699 let rows = converter
5700 .convert_columns(&[Arc::new(union_array.clone())])
5701 .unwrap();
5702
5703 let back = converter.convert_rows(&rows).unwrap();
5705 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5706
5707 assert_eq!(union_array.len(), back_union.len());
5708 for i in 0..union_array.len() {
5709 let expected_null = union_array.is_null(i);
5710 let actual_null = back_union.is_null(i);
5711 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5712 if !expected_null {
5713 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5714 }
5715 }
5716 }
5717
5718 #[test]
5719 fn test_union_ordering() {
5720 let int_array = Int32Array::from(vec![100, 5, 20]);
5721 let str_array = StringArray::from(vec!["z", "a"]);
5722
5723 let type_ids = vec![0, 1, 0, 1, 0].into();
5725 let offsets = vec![0, 0, 1, 1, 2].into();
5726
5727 let union_fields = [
5728 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5729 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5730 ]
5731 .into_iter()
5732 .collect();
5733
5734 let union_array = UnionArray::try_new(
5735 union_fields,
5736 type_ids,
5737 Some(offsets),
5738 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5739 )
5740 .unwrap();
5741
5742 let union_type = union_array.data_type().clone();
5743 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5744
5745 let rows = converter.convert_columns(&[Arc::new(union_array)]).unwrap();
5746
5747 assert!(rows.row(2) < rows.row(1));
5759
5760 assert!(rows.row(0) < rows.row(3));
5762
5763 assert!(rows.row(2) < rows.row(4));
5766 assert!(rows.row(4) < rows.row(0));
5768
5769 assert!(rows.row(3) < rows.row(1));
5772 }
5773
5774 #[test]
5775 fn test_row_converter_roundtrip_with_many_union_columns() {
5776 let fields1 = UnionFields::try_new(
5778 vec![0, 1],
5779 vec![
5780 Field::new("int", DataType::Int32, true),
5781 Field::new("string", DataType::Utf8, true),
5782 ],
5783 )
5784 .unwrap();
5785
5786 let int_array1 = Int32Array::from(vec![Some(67), None]);
5787 let string_array1 = StringArray::from(vec![None::<&str>, Some("hello")]);
5788 let type_ids1 = vec![0i8, 1].into();
5789
5790 let union_array1 = UnionArray::try_new(
5791 fields1.clone(),
5792 type_ids1,
5793 None,
5794 vec![
5795 Arc::new(int_array1) as ArrayRef,
5796 Arc::new(string_array1) as ArrayRef,
5797 ],
5798 )
5799 .unwrap();
5800
5801 let fields2 = UnionFields::try_new(
5803 vec![0, 1],
5804 vec![
5805 Field::new("int", DataType::Int32, true),
5806 Field::new("string", DataType::Utf8, true),
5807 ],
5808 )
5809 .unwrap();
5810
5811 let int_array2 = Int32Array::from(vec![Some(100), None]);
5812 let string_array2 = StringArray::from(vec![None::<&str>, Some("world")]);
5813 let type_ids2 = vec![0i8, 1].into();
5814
5815 let union_array2 = UnionArray::try_new(
5816 fields2.clone(),
5817 type_ids2,
5818 None,
5819 vec![
5820 Arc::new(int_array2) as ArrayRef,
5821 Arc::new(string_array2) as ArrayRef,
5822 ],
5823 )
5824 .unwrap();
5825
5826 let field1 = Field::new("col1", DataType::Union(fields1, UnionMode::Sparse), true);
5828 let field2 = Field::new("col2", DataType::Union(fields2, UnionMode::Sparse), true);
5829
5830 let sort_field1 = SortField::new(field1.data_type().clone());
5831 let sort_field2 = SortField::new(field2.data_type().clone());
5832
5833 let converter = RowConverter::new(vec![sort_field1, sort_field2]).unwrap();
5834
5835 let rows = converter
5836 .convert_columns(&[
5837 Arc::new(union_array1.clone()) as ArrayRef,
5838 Arc::new(union_array2.clone()) as ArrayRef,
5839 ])
5840 .unwrap();
5841
5842 let out = converter.convert_rows(&rows).unwrap();
5844
5845 let [col1, col2] = out.as_slice() else {
5846 panic!("expected 2 columns")
5847 };
5848
5849 let col1 = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5850 let col2 = col2.as_any().downcast_ref::<UnionArray>().unwrap();
5851
5852 for (expected, got) in [union_array1, union_array2].iter().zip([col1, col2]) {
5853 assert_eq!(expected.len(), got.len());
5854 assert_eq!(expected.type_ids(), got.type_ids());
5855
5856 for i in 0..expected.len() {
5857 assert_eq!(expected.value(i).as_ref(), got.value(i).as_ref());
5858 }
5859 }
5860 }
5861
5862 #[test]
5863 fn test_row_converter_roundtrip_with_one_union_column() {
5864 let fields = UnionFields::try_new(
5865 vec![0, 1],
5866 vec![
5867 Field::new("int", DataType::Int32, true),
5868 Field::new("string", DataType::Utf8, true),
5869 ],
5870 )
5871 .unwrap();
5872
5873 let int_array = Int32Array::from(vec![Some(67), None]);
5874 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5875 let type_ids = vec![0i8, 1].into();
5876
5877 let union_array = UnionArray::try_new(
5878 fields.clone(),
5879 type_ids,
5880 None,
5881 vec![
5882 Arc::new(int_array) as ArrayRef,
5883 Arc::new(string_array) as ArrayRef,
5884 ],
5885 )
5886 .unwrap();
5887
5888 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5889 let sort_field = SortField::new(field.data_type().clone());
5890 let converter = RowConverter::new(vec![sort_field]).unwrap();
5891
5892 let rows = converter
5893 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5894 .unwrap();
5895
5896 let out = converter.convert_rows(&rows).unwrap();
5898
5899 let [col1] = out.as_slice() else {
5900 panic!("expected 1 column")
5901 };
5902
5903 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5904 assert_eq!(col.len(), union_array.len());
5905 assert_eq!(col.type_ids(), union_array.type_ids());
5906
5907 for i in 0..col.len() {
5908 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5909 }
5910 }
5911
5912 #[test]
5913 fn test_row_converter_roundtrip_with_non_default_union_type_ids() {
5914 let fields = UnionFields::try_new(
5916 vec![70, 85],
5917 vec![
5918 Field::new("int", DataType::Int32, true),
5919 Field::new("string", DataType::Utf8, true),
5920 ],
5921 )
5922 .unwrap();
5923
5924 let int_array = Int32Array::from(vec![Some(67), None]);
5925 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5926 let type_ids = vec![70i8, 85].into();
5927
5928 let union_array = UnionArray::try_new(
5929 fields.clone(),
5930 type_ids,
5931 None,
5932 vec![
5933 Arc::new(int_array) as ArrayRef,
5934 Arc::new(string_array) as ArrayRef,
5935 ],
5936 )
5937 .unwrap();
5938
5939 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5940 let sort_field = SortField::new(field.data_type().clone());
5941 let converter = RowConverter::new(vec![sort_field]).unwrap();
5942
5943 let rows = converter
5944 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5945 .unwrap();
5946
5947 let out = converter.convert_rows(&rows).unwrap();
5949
5950 let [col1] = out.as_slice() else {
5951 panic!("expected 1 column")
5952 };
5953
5954 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5955 assert_eq!(col.len(), union_array.len());
5956 assert_eq!(col.type_ids(), union_array.type_ids());
5957
5958 for i in 0..col.len() {
5959 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5960 }
5961 }
5962
5963 #[test]
5964 fn rows_size_should_count_for_capacity() {
5965 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
5966
5967 let empty_rows_size_with_preallocate_rows_and_data = {
5968 let rows = row_converter.empty_rows(1000, 1000);
5969
5970 rows.size()
5971 };
5972 let empty_rows_size_with_preallocate_rows = {
5973 let rows = row_converter.empty_rows(1000, 0);
5974
5975 rows.size()
5976 };
5977 let empty_rows_size_with_preallocate_data = {
5978 let rows = row_converter.empty_rows(0, 1000);
5979
5980 rows.size()
5981 };
5982 let empty_rows_size_without_preallocate = {
5983 let rows = row_converter.empty_rows(0, 0);
5984
5985 rows.size()
5986 };
5987
5988 assert!(
5989 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_rows,
5990 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_rows}"
5991 );
5992 assert!(
5993 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_data,
5994 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_data}"
5995 );
5996 assert!(
5997 empty_rows_size_with_preallocate_rows > empty_rows_size_without_preallocate,
5998 "{empty_rows_size_with_preallocate_rows} should be larger than {empty_rows_size_without_preallocate}"
5999 );
6000 assert!(
6001 empty_rows_size_with_preallocate_data > empty_rows_size_without_preallocate,
6002 "{empty_rows_size_with_preallocate_data} should be larger than {empty_rows_size_without_preallocate}"
6003 );
6004 }
6005
6006 #[test]
6007 fn test_struct_no_child_fields() {
6008 fn run_test(array: ArrayRef) {
6009 let sort_fields = vec![SortField::new(array.data_type().clone())];
6010 let converter = RowConverter::new(sort_fields).unwrap();
6011 let r = converter.convert_columns(&[Arc::clone(&array)]).unwrap();
6012
6013 let back = converter.convert_rows(&r).unwrap();
6014 assert_eq!(back.len(), 1);
6015 assert_eq!(&back[0], &array);
6016 }
6017
6018 let s = Arc::new(StructArray::new_empty_fields(5, None)) as ArrayRef;
6019 run_test(s);
6020
6021 let s = Arc::new(StructArray::new_empty_fields(
6022 5,
6023 Some(vec![true, false, true, false, false].into()),
6024 )) as ArrayRef;
6025 run_test(s);
6026 }
6027
6028 #[test]
6029 fn reserve_should_increase_capacity_to_the_requested_size() {
6030 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6031 let mut empty_rows = row_converter.empty_rows(0, 0);
6032 empty_rows.reserve(50, 50);
6033 let before_size = empty_rows.size();
6034 empty_rows.reserve(50, 50);
6035 assert_eq!(
6036 empty_rows.size(),
6037 before_size,
6038 "Size should not change when reserving already reserved space"
6039 );
6040 empty_rows.reserve(10, 20);
6041 assert_eq!(
6042 empty_rows.size(),
6043 before_size,
6044 "Size should not change when already have space for the expected reserved data"
6045 );
6046
6047 empty_rows.reserve(100, 20);
6048 assert!(
6049 empty_rows.size() > before_size,
6050 "Size should increase when reserving more space than previously reserved"
6051 );
6052
6053 let before_size = empty_rows.size();
6054
6055 empty_rows.reserve(20, 100);
6056 assert!(
6057 empty_rows.size() > before_size,
6058 "Size should increase when reserving more space than previously reserved"
6059 );
6060 }
6061
6062 #[test]
6063 fn empty_rows_should_return_empty_lengths_iterator() {
6064 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6065 .unwrap()
6066 .empty_rows(0, 0);
6067 let mut lengths_iter = rows.lengths();
6068 assert_eq!(lengths_iter.next(), None);
6069 }
6070
6071 #[test]
6072 #[should_panic(expected = "row index out of bounds")]
6073 fn row_should_panic_on_overflowing_index() {
6074 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6075 .unwrap()
6076 .empty_rows(0, 0);
6077 rows.row(usize::MAX);
6078 }
6079
6080 #[test]
6081 #[should_panic(expected = "row index out of bounds")]
6082 fn row_len_should_panic_on_overflowing_index() {
6083 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6084 .unwrap()
6085 .empty_rows(0, 0);
6086 rows.row_len(usize::MAX);
6087 }
6088
6089 #[test]
6090 fn test_nested_null_list() {
6091 let null_array = Arc::new(NullArray::new(3));
6092 let list: ArrayRef = Arc::new(ListArray::new(
6094 Field::new_list_field(DataType::Null, true).into(),
6095 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6096 null_array,
6097 None,
6098 ));
6099
6100 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6101 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6102 let back = converter.convert_rows(&rows).unwrap();
6103
6104 assert_eq!(&list, &back[0]);
6105 }
6106
6107 #[test]
6109 fn test_double_nested_null_list() {
6110 let null_array = Arc::new(NullArray::new(1));
6111 let nested_field = Arc::new(Field::new_list_field(DataType::Null, true));
6113 let nested_list = Arc::new(ListArray::new(
6114 nested_field.clone(),
6115 OffsetBuffer::from_lengths(vec![1]),
6116 null_array,
6117 None,
6118 ));
6119 let list = Arc::new(ListArray::new(
6121 Field::new_list_field(DataType::List(nested_field), true).into(),
6122 OffsetBuffer::from_lengths(vec![1]),
6123 nested_list,
6124 None,
6125 )) as ArrayRef;
6126
6127 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6128 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6129 let back = converter.convert_rows(&rows).unwrap();
6130
6131 assert_eq!(&list, &back[0]);
6132 }
6133
6134 #[test]
6136 fn test_large_list_null() {
6137 let null_array = Arc::new(NullArray::new(3));
6138 let list: ArrayRef = Arc::new(LargeListArray::new(
6140 Field::new_list_field(DataType::Null, true).into(),
6141 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6142 null_array,
6143 None,
6144 ));
6145
6146 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6147 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6148 let back = converter.convert_rows(&rows).unwrap();
6149
6150 assert_eq!(&list, &back[0]);
6151 }
6152
6153 #[test]
6155 fn test_fixed_size_list_null() {
6156 let null_array = Arc::new(NullArray::new(6));
6157 let list: ArrayRef = Arc::new(FixedSizeListArray::new(
6159 Arc::new(Field::new_list_field(DataType::Null, true)),
6160 2,
6161 null_array,
6162 None,
6163 ));
6164
6165 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6166 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6167 let back = converter.convert_rows(&rows).unwrap();
6168
6169 assert_eq!(&list, &back[0]);
6170 }
6171
6172 #[test]
6174 fn test_fixed_size_list_of_dictionaries_round_trips() {
6175 let dict_dt = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
6178 let element_field = Arc::new(Field::new("item", dict_dt.clone(), true));
6179 let fsl_dt = DataType::FixedSizeList(Arc::clone(&element_field), 2);
6180
6181 let values = Arc::new(StringArray::from(vec!["a", "b"]));
6182 let keys = Int32Array::from(vec![0, 1]);
6183 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
6184 let fsl: ArrayRef = Arc::new(FixedSizeListArray::new(
6185 Arc::clone(&element_field),
6186 2,
6187 Arc::new(dict),
6188 None,
6189 ));
6190
6191 assert!(RowConverter::supports_fields(&[SortField::new(
6192 fsl_dt.clone()
6193 )]));
6194
6195 let converter = RowConverter::new(vec![SortField::new(fsl_dt.clone())]).unwrap();
6196 let rows = converter.convert_columns(&[Arc::clone(&fsl)]).unwrap();
6197
6198 let back = converter.convert_rows(&rows).unwrap();
6201 assert_eq!(back.len(), 1);
6202
6203 let out = back[0]
6207 .as_any()
6208 .downcast_ref::<FixedSizeListArray>()
6209 .expect("decoded array must be a FixedSizeListArray");
6210 assert_eq!(out.len(), 1);
6211 assert_eq!(out.value_length(), 2);
6212 assert_eq!(out.values().data_type(), &DataType::Utf8);
6216
6217 let values = out
6219 .values()
6220 .as_any()
6221 .downcast_ref::<StringArray>()
6222 .expect("child must be a StringArray after flattening");
6223 assert_eq!(values.value(0), "a");
6224 assert_eq!(values.value(1), "b");
6225 }
6226
6227 #[test]
6229 fn test_list_null_variations() {
6230 let null_array = Arc::new(NullArray::new(3));
6232 let list: ArrayRef = Arc::new(ListArray::new(
6233 Field::new_list_field(DataType::Null, true).into(),
6234 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6235 null_array,
6236 None,
6237 ));
6238
6239 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6240 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6241 let back = converter.convert_rows(&rows).unwrap();
6242 assert_eq!(&list, &back[0]);
6243
6244 let null_array = Arc::new(NullArray::new(3));
6246 let list: ArrayRef = Arc::new(ListArray::new(
6247 Field::new_list_field(DataType::Null, true).into(),
6248 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6249 null_array,
6250 Some(vec![true, false, true].into()),
6251 ));
6252
6253 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6254 let back = converter.convert_rows(&rows).unwrap();
6255 assert_eq!(&list, &back[0]);
6256
6257 let null_array = Arc::new(NullArray::new(0));
6259 let list: ArrayRef = Arc::new(ListArray::new(
6260 Field::new_list_field(DataType::Null, true).into(),
6261 OffsetBuffer::from_lengths(vec![]),
6262 null_array,
6263 None,
6264 ));
6265
6266 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6267 let back = converter.convert_rows(&rows).unwrap();
6268 assert_eq!(&list, &back[0]);
6269
6270 let null_array = Arc::new(NullArray::new(0));
6272 let list: ArrayRef = Arc::new(ListArray::new(
6273 Field::new_list_field(DataType::Null, true).into(),
6274 OffsetBuffer::from_lengths(vec![0, 0, 0]),
6275 null_array,
6276 None,
6277 ));
6278
6279 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6280 let back = converter.convert_rows(&rows).unwrap();
6281 assert_eq!(&list, &back[0]);
6282 }
6283
6284 #[test]
6286 fn test_list_null_descending() {
6287 let null_array = Arc::new(NullArray::new(3));
6288 let list: ArrayRef = Arc::new(ListArray::new(
6290 Field::new_list_field(DataType::Null, true).into(),
6291 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6292 null_array,
6293 None,
6294 ));
6295
6296 let options = SortOptions::default().with_descending(true);
6297 let field = SortField::new_with_options(list.data_type().clone(), options);
6298 let converter = RowConverter::new(vec![field]).unwrap();
6299 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6300 let back = converter.convert_rows(&rows).unwrap();
6301
6302 assert_eq!(&list, &back[0]);
6303 }
6304
6305 #[test]
6307 fn test_struct_with_null_field() {
6308 let null_array = Arc::new(NullArray::new(3));
6310 let int_array = Arc::new(Int32Array::from(vec![1, 2, 3]));
6311
6312 let struct_array: ArrayRef = Arc::new(StructArray::new(
6313 vec![
6314 Arc::new(Field::new("a", DataType::Null, true)),
6315 Arc::new(Field::new("b", DataType::Int32, true)),
6316 ]
6317 .into(),
6318 vec![null_array, int_array],
6319 Some(vec![true, true, false].into()), ));
6321
6322 let converter =
6323 RowConverter::new(vec![SortField::new(struct_array.data_type().clone())]).unwrap();
6324 let rows = converter
6325 .convert_columns(&[Arc::clone(&struct_array)])
6326 .unwrap();
6327 let back = converter.convert_rows(&rows).unwrap();
6328
6329 assert_eq!(&struct_array, &back[0]);
6330 }
6331
6332 #[test]
6334 fn test_nested_struct_with_null() {
6335 let inner_null = Arc::new(NullArray::new(2));
6337 let inner_struct = Arc::new(StructArray::new(
6338 vec![Arc::new(Field::new("x", DataType::Null, true))].into(),
6339 vec![inner_null],
6340 None,
6341 ));
6342
6343 let y_array = Arc::new(Int32Array::from(vec![10, 20]));
6345 let outer_struct: ArrayRef = Arc::new(StructArray::new(
6346 vec![
6347 Arc::new(Field::new("inner", inner_struct.data_type().clone(), true)),
6348 Arc::new(Field::new("y", DataType::Int32, true)),
6349 ]
6350 .into(),
6351 vec![inner_struct, y_array],
6352 None,
6353 ));
6354
6355 let converter =
6356 RowConverter::new(vec![SortField::new(outer_struct.data_type().clone())]).unwrap();
6357 let rows = converter
6358 .convert_columns(&[Arc::clone(&outer_struct)])
6359 .unwrap();
6360 let back = converter.convert_rows(&rows).unwrap();
6361
6362 assert_eq!(&outer_struct, &back[0]);
6363 }
6364
6365 #[test]
6367 fn test_map_null_variations() {
6368 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6370 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6371
6372 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6373 let entries_fields = vec![
6374 Arc::new(Field::new(
6375 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6376 DataType::Utf8,
6377 false,
6378 )),
6379 Arc::new(Field::new(
6380 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6381 DataType::Null,
6382 true,
6383 )),
6384 ];
6385 let struct_field = Arc::new(Field::new(
6386 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6387 DataType::Struct(entries_fields.clone().into()),
6388 false,
6389 ));
6390 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6391
6392 let map: ArrayRef = Arc::new(MapArray::new(
6393 struct_field.clone(),
6394 offsets,
6395 entries,
6396 None,
6397 false,
6398 ));
6399
6400 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6401 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6402 let back = converter.convert_rows(&rows).unwrap();
6403 assert_eq!(back.len(), 1);
6404 back[0].to_data().validate_full().unwrap();
6405 assert_eq!(&map, &back[0]);
6406
6407 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6409 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6410
6411 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6412 let entries_fields = vec![
6413 Arc::new(Field::new(
6414 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6415 DataType::Utf8,
6416 false,
6417 )),
6418 Arc::new(Field::new(
6419 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6420 DataType::Null,
6421 true,
6422 )),
6423 ];
6424 let struct_field = Arc::new(Field::new(
6425 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6426 DataType::Struct(entries_fields.clone().into()),
6427 false,
6428 ));
6429 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6430
6431 let map: ArrayRef = Arc::new(MapArray::new(
6432 struct_field.clone(),
6433 offsets,
6434 entries,
6435 Some(vec![true, false, true].into()),
6436 false,
6437 ));
6438
6439 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6440 let back = converter.convert_rows(&rows).unwrap();
6441 assert_eq!(back.len(), 1);
6442 back[0].to_data().validate_full().unwrap();
6443 assert_eq!(&map, &back[0]);
6444
6445 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6447 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6448
6449 let offsets = OffsetBuffer::new(vec![0i32].into());
6450 let entries_fields = vec![
6451 Arc::new(Field::new(
6452 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6453 DataType::Utf8,
6454 false,
6455 )),
6456 Arc::new(Field::new(
6457 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6458 DataType::Null,
6459 true,
6460 )),
6461 ];
6462 let struct_field = Arc::new(Field::new(
6463 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6464 DataType::Struct(entries_fields.clone().into()),
6465 false,
6466 ));
6467 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6468
6469 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6470
6471 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6472 let back = converter.convert_rows(&rows).unwrap();
6473 assert_eq!(back.len(), 1);
6474 back[0].to_data().validate_full().unwrap();
6475 assert_eq!(&map, &back[0]);
6476 }
6477
6478 #[test]
6480 fn test_map_null_descending() {
6481 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6483 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6484
6485 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6486 let entries_fields = vec![
6487 Arc::new(Field::new(
6488 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6489 DataType::Utf8,
6490 false,
6491 )),
6492 Arc::new(Field::new(
6493 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6494 DataType::Null,
6495 true,
6496 )),
6497 ];
6498 let struct_field = Arc::new(Field::new(
6499 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6500 DataType::Struct(entries_fields.clone().into()),
6501 false,
6502 ));
6503 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6504
6505 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6506
6507 let options = SortOptions::default().with_descending(true);
6508 let field = SortField::new_with_options(map.data_type().clone(), options);
6509 let converter = RowConverter::new(vec![field]).unwrap();
6510 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6511 let back = converter.convert_rows(&rows).unwrap();
6512 assert_eq!(back.len(), 1);
6513 back[0].to_data().validate_full().unwrap();
6514 assert_eq!(&map, &back[0]);
6515 }
6516
6517 #[test]
6519 fn test_map_null_all_empty() {
6520 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6521 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6522
6523 let offsets = OffsetBuffer::new(vec![0, 0, 0, 0].into());
6524 let entries_fields = vec![
6525 Arc::new(Field::new(
6526 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6527 DataType::Utf8,
6528 false,
6529 )),
6530 Arc::new(Field::new(
6531 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6532 DataType::Null,
6533 true,
6534 )),
6535 ];
6536 let struct_field = Arc::new(Field::new(
6537 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
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: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6544
6545 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6546 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6547
6548 assert_eq!(rows.row(0), rows.row(1));
6550 assert_eq!(rows.row(1), rows.row(2));
6551
6552 let back = converter.convert_rows(&rows).unwrap();
6553 assert_eq!(back.len(), 1);
6554 back[0].to_data().validate_full().unwrap();
6555 assert_eq!(&map, &back[0]);
6556 }
6557
6558 #[test]
6560 fn test_nested_map_null() {
6561 let inner_keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6563 let inner_null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6564
6565 let inner_entries_fields = vec![
6566 Arc::new(Field::new(
6567 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6568 DataType::Utf8,
6569 false,
6570 )),
6571 Arc::new(Field::new(
6572 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6573 DataType::Null,
6574 true,
6575 )),
6576 ];
6577 let inner_struct_field = Arc::new(Field::new(
6578 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6579 DataType::Struct(inner_entries_fields.clone().into()),
6580 false,
6581 ));
6582 let inner_entries = StructArray::new(
6583 inner_entries_fields.clone().into(),
6584 vec![inner_keys, inner_null_values],
6585 None,
6586 );
6587
6588 let inner_map = Arc::new(MapArray::new(
6590 inner_struct_field.clone(),
6591 OffsetBuffer::new(vec![0, 1, 3].into()),
6592 inner_entries,
6593 None,
6594 false,
6595 )) as ArrayRef;
6596
6597 let outer_keys = Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef;
6599
6600 let inner_map_type = DataType::Map(inner_struct_field.clone(), false);
6601 let outer_entries_fields = vec![
6602 Arc::new(Field::new(
6603 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6604 DataType::Utf8,
6605 false,
6606 )),
6607 Arc::new(Field::new(
6608 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6609 inner_map_type,
6610 true,
6611 )),
6612 ];
6613 let outer_struct_field = Arc::new(Field::new(
6614 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6615 DataType::Struct(outer_entries_fields.clone().into()),
6616 false,
6617 ));
6618 let outer_entries = StructArray::new(
6619 outer_entries_fields.into(),
6620 vec![outer_keys, inner_map],
6621 None,
6622 );
6623
6624 let map: ArrayRef = Arc::new(MapArray::new(
6626 outer_struct_field,
6627 OffsetBuffer::new(vec![0, 1, 2].into()),
6628 outer_entries,
6629 None,
6630 false,
6631 ));
6632
6633 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6634 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6635 let back = converter.convert_rows(&rows).unwrap();
6636 assert_eq!(back.len(), 1);
6637 back[0].to_data().validate_full().unwrap();
6638 assert_eq!(&map, &back[0]);
6639 }
6640
6641 #[test]
6643 fn test_list_of_map_null() {
6644 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6646 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6647
6648 let entries_fields = vec![
6649 Arc::new(Field::new(
6650 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6651 DataType::Utf8,
6652 false,
6653 )),
6654 Arc::new(Field::new(
6655 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6656 DataType::Null,
6657 true,
6658 )),
6659 ];
6660 let struct_field = Arc::new(Field::new(
6661 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6662 DataType::Struct(entries_fields.clone().into()),
6663 false,
6664 ));
6665 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6666
6667 let map_array = Arc::new(MapArray::new(
6669 struct_field.clone(),
6670 OffsetBuffer::new(vec![0, 1, 1, 3].into()),
6671 entries,
6672 None,
6673 false,
6674 )) as ArrayRef;
6675
6676 let map_type = DataType::Map(struct_field, false);
6677 let list: ArrayRef = Arc::new(ListArray::new(
6679 Arc::new(Field::new_list_field(map_type, true)),
6680 OffsetBuffer::new(vec![0, 1, 3].into()),
6681 map_array,
6682 None,
6683 ));
6684
6685 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6686 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6687 let back = converter.convert_rows(&rows).unwrap();
6688 assert_eq!(&list, &back[0]);
6689 }
6690
6691 #[test]
6693 fn test_map_of_list_null() {
6694 let null_array = Arc::new(NullArray::new(3)) as ArrayRef;
6696 let list_array = Arc::new(ListArray::new(
6698 Arc::new(Field::new_list_field(DataType::Null, true)),
6699 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6700 null_array,
6701 None,
6702 )) as ArrayRef;
6703
6704 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6705
6706 let list_type = list_array.data_type().clone();
6707 let entries_fields = vec![
6708 Arc::new(Field::new(
6709 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6710 DataType::Utf8,
6711 false,
6712 )),
6713 Arc::new(Field::new(
6714 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6715 list_type,
6716 true,
6717 )),
6718 ];
6719 let struct_field = Arc::new(Field::new(
6720 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6721 DataType::Struct(entries_fields.clone().into()),
6722 false,
6723 ));
6724 let entries = StructArray::new(entries_fields.into(), vec![keys, list_array], None);
6725
6726 let map: ArrayRef = Arc::new(MapArray::new(
6728 struct_field,
6729 OffsetBuffer::new(vec![0, 3].into()),
6730 entries,
6731 None,
6732 false,
6733 ));
6734
6735 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6736 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6737 let back = converter.convert_rows(&rows).unwrap();
6738 assert_eq!(back.len(), 1);
6739 back[0].to_data().validate_full().unwrap();
6740 assert_eq!(&map, &back[0]);
6741 }
6742
6743 #[test]
6744 fn empty_row_iter_next_back() {
6745 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6746 .unwrap()
6747 .empty_rows(0, 0);
6748 let mut rows_iter = rows.iter();
6749 assert_eq!(rows_iter.next_back(), None);
6750 assert_eq!(rows_iter.next_back(), None);
6751 assert_eq!(rows_iter.next_back(), None);
6752 }
6753
6754 #[test]
6756 fn test_row_parser_skip_utf8_validation_roundtrip() {
6757 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
6758 let array = StringArray::from(vec!["arrow", "rust"]);
6759 let rows = converter.convert_columns(&[Arc::new(array) as _]).unwrap();
6760 let binary = rows.try_into_binary().expect("fits in i32 offsets");
6761
6762 let parser = unsafe { RowParser::with_skip_utf8_validate(Arc::clone(&converter.fields)) };
6764
6765 let decoded = converter
6766 .convert_rows(binary.iter().map(|b| parser.parse(b.unwrap())))
6767 .unwrap();
6768 let got: Vec<_> = decoded[0].as_string::<i32>().iter().flatten().collect();
6769 assert_eq!(got, vec!["arrow", "rust"]);
6770 }
6771
6772 #[test]
6773 fn row_iter_next_back() {
6774 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6775 let mut rng = StdRng::seed_from_u64(42);
6776 let array = generate_primitive_array::<UInt8Type>(&mut rng, 100, 0.8);
6777 let rows = row_converter.convert_columns(&[Arc::new(array)]).unwrap();
6778
6779 let mut rows_iter = rows.iter();
6780 let mut bytes: Vec<u8> = vec![];
6781
6782 while let Some(row) = rows_iter.next_back() {
6783 bytes.extend(row.data.iter().rev());
6784 }
6785
6786 bytes.reverse();
6787
6788 assert_eq!(
6789 bytes,
6790 &rows.buffer.as_slice()[..*rows.offsets.last().unwrap()]
6791 );
6792
6793 assert_eq!(rows_iter.next_back(), None);
6794 assert_eq!(rows_iter.next(), None);
6795 }
6796}