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