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 column in columns.iter().skip(1) {
1055 if column.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 column.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_buffer::{Buffer, OffsetBuffer};
2438 use arrow_buffer::{NullBuffer, i256};
2439 use arrow_cast::display::{ArrayFormatter, FormatOptions};
2440 use arrow_ord::sort::{LexicographicalComparator, SortColumn};
2441 use rand::distr::uniform::SampleUniform;
2442 use rand::distr::{Distribution, StandardUniform};
2443 use rand::prelude::StdRng;
2444 use rand::{RngExt, SeedableRng};
2445
2446 use super::*;
2447
2448 fn all_sort_options() -> [SortOptions; 4] {
2449 [
2450 SortOptions {
2451 descending: false,
2452 nulls_first: false,
2453 },
2454 SortOptions {
2455 descending: false,
2456 nulls_first: true,
2457 },
2458 SortOptions {
2459 descending: true,
2460 nulls_first: false,
2461 },
2462 SortOptions {
2463 descending: true,
2464 nulls_first: true,
2465 },
2466 ]
2467 }
2468
2469 #[test]
2470 fn test_fixed_width() {
2471 let cols = [
2472 Arc::new(Int16Array::from_iter([
2473 Some(1),
2474 Some(2),
2475 None,
2476 Some(-5),
2477 Some(2),
2478 Some(2),
2479 Some(0),
2480 ])) as ArrayRef,
2481 Arc::new(Float32Array::from_iter([
2482 Some(1.3),
2483 Some(2.5),
2484 None,
2485 Some(4.),
2486 Some(0.1),
2487 Some(-4.),
2488 Some(-0.),
2489 ])) as ArrayRef,
2490 ];
2491
2492 let converter = RowConverter::new(vec![
2493 SortField::new(DataType::Int16),
2494 SortField::new(DataType::Float32),
2495 ])
2496 .unwrap();
2497 let rows = converter.convert_columns(&cols).unwrap();
2498
2499 assert_eq!(rows.offsets, &[0, 8, 16, 24, 32, 40, 48, 56]);
2500 assert_eq!(
2501 rows.buffer,
2502 &[
2503 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 ]
2518 );
2519
2520 assert!(rows.row(3) < rows.row(6));
2521 assert!(rows.row(0) < rows.row(1));
2522 assert!(rows.row(3) < rows.row(0));
2523 assert!(rows.row(4) < rows.row(1));
2524 assert!(rows.row(5) < rows.row(4));
2525
2526 let back = converter.convert_rows(&rows).unwrap();
2527 for (expected, actual) in cols.iter().zip(&back) {
2528 assert_eq!(expected, actual);
2529 }
2530 }
2531
2532 fn test_roundtrip(sort_option: SortOptions, col: ArrayRef) {
2533 let converter = RowConverter::new(vec![SortField::new_with_options(
2534 col.data_type().clone(),
2535 sort_option,
2536 )])
2537 .unwrap();
2538 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2539 let back = converter.convert_rows(&rows).unwrap();
2540 assert_eq!(back.len(), 1);
2541 assert_eq!(&back[0], &col);
2542 back[0].to_data().validate_full().unwrap();
2543 }
2544
2545 #[test]
2546 fn test_zero_width_fixed_size_binary_roundtrip() {
2547 for sort_option in all_sort_options() {
2548 for with_null in [true, false] {
2551 let nulls = if with_null {
2552 Some(NullBuffer::from(vec![true, false, true, false, true]))
2553 } else {
2554 None
2555 };
2556 let col: ArrayRef = Arc::new(
2557 FixedSizeBinaryArray::try_new_with_len(0, Buffer::default(), nulls, 5).unwrap(),
2558 );
2559
2560 test_roundtrip(sort_option, col);
2561 }
2562 }
2563 }
2564
2565 #[test]
2566 fn test_zero_width_fixed_size_list_roundtrip() {
2567 for sort_option in all_sort_options() {
2568 for with_null in [true, false] {
2571 let nulls = if with_null {
2572 Some(NullBuffer::from(vec![true, false, true, false, true]))
2573 } else {
2574 None
2575 };
2576 let col: ArrayRef = Arc::new(
2577 FixedSizeListArray::try_new_with_length(
2578 Arc::new(Field::new("item", DataType::Boolean, false)),
2579 0,
2580 new_empty_array(&DataType::Boolean),
2581 nulls,
2582 5,
2583 )
2584 .unwrap(),
2585 );
2586
2587 test_roundtrip(sort_option, col);
2588 }
2589 }
2590 }
2591
2592 #[test]
2593 fn test_decimal32() {
2594 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal32(
2595 DECIMAL32_MAX_PRECISION,
2596 7,
2597 ))])
2598 .unwrap();
2599 let col = Arc::new(
2600 Decimal32Array::from_iter([
2601 None,
2602 Some(i32::MIN),
2603 Some(-13),
2604 Some(46_i32),
2605 Some(5456_i32),
2606 Some(i32::MAX),
2607 ])
2608 .with_precision_and_scale(9, 7)
2609 .unwrap(),
2610 ) as ArrayRef;
2611
2612 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2613 for i in 0..rows.num_rows() - 1 {
2614 assert!(rows.row(i) < rows.row(i + 1));
2615 }
2616
2617 let back = converter.convert_rows(&rows).unwrap();
2618 assert_eq!(back.len(), 1);
2619 assert_eq!(col.as_ref(), back[0].as_ref())
2620 }
2621
2622 #[test]
2623 fn test_decimal64() {
2624 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal64(
2625 DECIMAL64_MAX_PRECISION,
2626 7,
2627 ))])
2628 .unwrap();
2629 let col = Arc::new(
2630 Decimal64Array::from_iter([
2631 None,
2632 Some(i64::MIN),
2633 Some(-13),
2634 Some(46_i64),
2635 Some(5456_i64),
2636 Some(i64::MAX),
2637 ])
2638 .with_precision_and_scale(18, 7)
2639 .unwrap(),
2640 ) as ArrayRef;
2641
2642 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2643 for i in 0..rows.num_rows() - 1 {
2644 assert!(rows.row(i) < rows.row(i + 1));
2645 }
2646
2647 let back = converter.convert_rows(&rows).unwrap();
2648 assert_eq!(back.len(), 1);
2649 assert_eq!(col.as_ref(), back[0].as_ref())
2650 }
2651
2652 #[test]
2653 fn test_decimal128() {
2654 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal128(
2655 DECIMAL128_MAX_PRECISION,
2656 7,
2657 ))])
2658 .unwrap();
2659 let col = Arc::new(
2660 Decimal128Array::from_iter([
2661 None,
2662 Some(i128::MIN),
2663 Some(-13),
2664 Some(46_i128),
2665 Some(5456_i128),
2666 Some(i128::MAX),
2667 ])
2668 .with_precision_and_scale(38, 7)
2669 .unwrap(),
2670 ) as ArrayRef;
2671
2672 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2673 for i in 0..rows.num_rows() - 1 {
2674 assert!(rows.row(i) < rows.row(i + 1));
2675 }
2676
2677 let back = converter.convert_rows(&rows).unwrap();
2678 assert_eq!(back.len(), 1);
2679 assert_eq!(col.as_ref(), back[0].as_ref())
2680 }
2681
2682 #[test]
2683 fn test_decimal256() {
2684 let converter = RowConverter::new(vec![SortField::new(DataType::Decimal256(
2685 DECIMAL256_MAX_PRECISION,
2686 7,
2687 ))])
2688 .unwrap();
2689 let col = Arc::new(
2690 Decimal256Array::from_iter([
2691 None,
2692 Some(i256::MIN),
2693 Some(i256::from_parts(0, -1)),
2694 Some(i256::from_parts(u128::MAX, -1)),
2695 Some(i256::from_parts(u128::MAX, 0)),
2696 Some(i256::from_parts(0, 46_i128)),
2697 Some(i256::from_parts(5, 46_i128)),
2698 Some(i256::MAX),
2699 ])
2700 .with_precision_and_scale(DECIMAL256_MAX_PRECISION, 7)
2701 .unwrap(),
2702 ) as ArrayRef;
2703
2704 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2705 for i in 0..rows.num_rows() - 1 {
2706 assert!(rows.row(i) < rows.row(i + 1));
2707 }
2708
2709 let back = converter.convert_rows(&rows).unwrap();
2710 assert_eq!(back.len(), 1);
2711 assert_eq!(col.as_ref(), back[0].as_ref())
2712 }
2713
2714 #[test]
2715 fn test_bool() {
2716 let converter = RowConverter::new(vec![SortField::new(DataType::Boolean)]).unwrap();
2717
2718 let col = Arc::new(BooleanArray::from_iter([None, Some(false), Some(true)])) as ArrayRef;
2719
2720 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2721 assert!(rows.row(2) > rows.row(1));
2722 assert!(rows.row(2) > rows.row(0));
2723 assert!(rows.row(1) > rows.row(0));
2724
2725 let cols = converter.convert_rows(&rows).unwrap();
2726 assert_eq!(&cols[0], &col);
2727
2728 let converter = RowConverter::new(vec![SortField::new_with_options(
2729 DataType::Boolean,
2730 SortOptions::default().desc().with_nulls_first(false),
2731 )])
2732 .unwrap();
2733
2734 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2735 assert!(rows.row(2) < rows.row(1));
2736 assert!(rows.row(2) < rows.row(0));
2737 assert!(rows.row(1) < rows.row(0));
2738 let cols = converter.convert_rows(&rows).unwrap();
2739 assert_eq!(&cols[0], &col);
2740 }
2741
2742 #[test]
2743 fn test_timezone() {
2744 let a =
2745 TimestampNanosecondArray::from(vec![1, 2, 3, 4, 5]).with_timezone("+01:00".to_string());
2746 let d = a.data_type().clone();
2747
2748 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
2749 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
2750 let back = converter.convert_rows(&rows).unwrap();
2751 assert_eq!(back.len(), 1);
2752 assert_eq!(back[0].data_type(), &d);
2753
2754 let mut a = PrimitiveDictionaryBuilder::<Int32Type, TimestampNanosecondType>::new();
2756 a.append(34).unwrap();
2757 a.append_null();
2758 a.append(345).unwrap();
2759
2760 let dict = a.finish();
2762 let values = TimestampNanosecondArray::from(dict.values().to_data());
2763 let dict_with_tz = dict.with_values(Arc::new(values.with_timezone("+02:00")));
2764 let v = DataType::Timestamp(TimeUnit::Nanosecond, Some("+02:00".into()));
2765 let d = DataType::Dictionary(Box::new(DataType::Int32), Box::new(v.clone()));
2766
2767 assert_eq!(dict_with_tz.data_type(), &d);
2768 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
2769 let rows = converter
2770 .convert_columns(&[Arc::new(dict_with_tz) as _])
2771 .unwrap();
2772 let back = converter.convert_rows(&rows).unwrap();
2773 assert_eq!(back.len(), 1);
2774 assert_eq!(back[0].data_type(), &v);
2775 }
2776
2777 #[test]
2778 fn test_null_encoding() {
2779 let col = Arc::new(NullArray::new(10));
2780 let converter = RowConverter::new(vec![SortField::new(DataType::Null)]).unwrap();
2781 let rows = converter.convert_columns(&[col]).unwrap();
2782 assert_eq!(rows.num_rows(), 10);
2783 assert_eq!(rows.row(1).data.len(), 2);
2785 }
2786
2787 #[test]
2788 fn test_variable_width() {
2789 let col = Arc::new(StringArray::from_iter([
2790 Some("hello"),
2791 Some("he"),
2792 None,
2793 Some("foo"),
2794 Some(""),
2795 ])) as ArrayRef;
2796
2797 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
2798 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2799
2800 assert!(rows.row(1) < rows.row(0));
2801 assert!(rows.row(2) < rows.row(4));
2802 assert!(rows.row(3) < rows.row(0));
2803 assert!(rows.row(3) < rows.row(1));
2804
2805 let cols = converter.convert_rows(&rows).unwrap();
2806 assert_eq!(&cols[0], &col);
2807
2808 let col = Arc::new(BinaryArray::from_iter([
2809 None,
2810 Some(vec![0_u8; 0]),
2811 Some(vec![0_u8; 6]),
2812 Some(vec![0_u8; variable::MINI_BLOCK_SIZE]),
2813 Some(vec![0_u8; variable::MINI_BLOCK_SIZE + 1]),
2814 Some(vec![0_u8; variable::BLOCK_SIZE]),
2815 Some(vec![0_u8; variable::BLOCK_SIZE + 1]),
2816 Some(vec![1_u8; 6]),
2817 Some(vec![1_u8; variable::MINI_BLOCK_SIZE]),
2818 Some(vec![1_u8; variable::MINI_BLOCK_SIZE + 1]),
2819 Some(vec![1_u8; variable::BLOCK_SIZE]),
2820 Some(vec![1_u8; variable::BLOCK_SIZE + 1]),
2821 Some(vec![0xFF_u8; 6]),
2822 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE]),
2823 Some(vec![0xFF_u8; variable::MINI_BLOCK_SIZE + 1]),
2824 Some(vec![0xFF_u8; variable::BLOCK_SIZE]),
2825 Some(vec![0xFF_u8; variable::BLOCK_SIZE + 1]),
2826 ])) as ArrayRef;
2827
2828 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
2829 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2830
2831 for i in 0..rows.num_rows() {
2832 for j in i + 1..rows.num_rows() {
2833 assert!(
2834 rows.row(i) < rows.row(j),
2835 "{} < {} - {:?} < {:?}",
2836 i,
2837 j,
2838 rows.row(i),
2839 rows.row(j)
2840 );
2841 }
2842 }
2843
2844 let cols = converter.convert_rows(&rows).unwrap();
2845 assert_eq!(&cols[0], &col);
2846
2847 let converter = RowConverter::new(vec![SortField::new_with_options(
2848 DataType::Binary,
2849 SortOptions::default().desc().with_nulls_first(false),
2850 )])
2851 .unwrap();
2852 let rows = converter.convert_columns(&[Arc::clone(&col)]).unwrap();
2853
2854 for i in 0..rows.num_rows() {
2855 for j in i + 1..rows.num_rows() {
2856 assert!(
2857 rows.row(i) > rows.row(j),
2858 "{} > {} - {:?} > {:?}",
2859 i,
2860 j,
2861 rows.row(i),
2862 rows.row(j)
2863 );
2864 }
2865 }
2866
2867 let cols = converter.convert_rows(&rows).unwrap();
2868 assert_eq!(&cols[0], &col);
2869 }
2870
2871 fn dictionary_eq(a: &dyn Array, b: &dyn Array) {
2873 match b.data_type() {
2874 DataType::Dictionary(_, v) => {
2875 assert_eq!(a.data_type(), v.as_ref());
2876 let b = arrow_cast::cast(b, v).unwrap();
2877 assert_eq!(a, b.as_ref())
2878 }
2879 _ => assert_eq!(a, b),
2880 }
2881 }
2882
2883 #[test]
2884 fn test_string_dictionary() {
2885 let a = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2886 Some("foo"),
2887 Some("hello"),
2888 Some("he"),
2889 None,
2890 Some("hello"),
2891 Some(""),
2892 Some("hello"),
2893 Some("hello"),
2894 ])) as ArrayRef;
2895
2896 let field = SortField::new(a.data_type().clone());
2897 let converter = RowConverter::new(vec![field]).unwrap();
2898 let rows_a = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2899
2900 assert!(rows_a.row(3) < rows_a.row(5));
2901 assert!(rows_a.row(2) < rows_a.row(1));
2902 assert!(rows_a.row(0) < rows_a.row(1));
2903 assert!(rows_a.row(3) < rows_a.row(0));
2904
2905 assert_eq!(rows_a.row(1), rows_a.row(4));
2906 assert_eq!(rows_a.row(1), rows_a.row(6));
2907 assert_eq!(rows_a.row(1), rows_a.row(7));
2908
2909 let cols = converter.convert_rows(&rows_a).unwrap();
2910 dictionary_eq(&cols[0], &a);
2911
2912 let b = Arc::new(DictionaryArray::<Int32Type>::from_iter([
2913 Some("hello"),
2914 None,
2915 Some("cupcakes"),
2916 ])) as ArrayRef;
2917
2918 let rows_b = converter.convert_columns(&[Arc::clone(&b)]).unwrap();
2919 assert_eq!(rows_a.row(1), rows_b.row(0));
2920 assert_eq!(rows_a.row(3), rows_b.row(1));
2921 assert!(rows_b.row(2) < rows_a.row(0));
2922
2923 let cols = converter.convert_rows(&rows_b).unwrap();
2924 dictionary_eq(&cols[0], &b);
2925
2926 let converter = RowConverter::new(vec![SortField::new_with_options(
2927 a.data_type().clone(),
2928 SortOptions::default().desc().with_nulls_first(false),
2929 )])
2930 .unwrap();
2931
2932 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2933 assert!(rows_c.row(3) > rows_c.row(5));
2934 assert!(rows_c.row(2) > rows_c.row(1));
2935 assert!(rows_c.row(0) > rows_c.row(1));
2936 assert!(rows_c.row(3) > rows_c.row(0));
2937
2938 let cols = converter.convert_rows(&rows_c).unwrap();
2939 dictionary_eq(&cols[0], &a);
2940
2941 let converter = RowConverter::new(vec![SortField::new_with_options(
2942 a.data_type().clone(),
2943 SortOptions::default().desc().with_nulls_first(true),
2944 )])
2945 .unwrap();
2946
2947 let rows_c = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
2948 assert!(rows_c.row(3) < rows_c.row(5));
2949 assert!(rows_c.row(2) > rows_c.row(1));
2950 assert!(rows_c.row(0) > rows_c.row(1));
2951 assert!(rows_c.row(3) < rows_c.row(0));
2952
2953 let cols = converter.convert_rows(&rows_c).unwrap();
2954 dictionary_eq(&cols[0], &a);
2955 }
2956
2957 #[test]
2958 fn test_struct() {
2959 let a = Arc::new(Int32Array::from(vec![1, 1, 2, 2])) as ArrayRef;
2961 let a_f = Arc::new(Field::new("int", DataType::Int32, false));
2962 let u = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as ArrayRef;
2963 let u_f = Arc::new(Field::new("s", DataType::Utf8, false));
2964 let s1 = Arc::new(StructArray::from(vec![(a_f, a), (u_f, u)])) as ArrayRef;
2965
2966 let sort_fields = vec![SortField::new(s1.data_type().clone())];
2967 let converter = RowConverter::new(sort_fields).unwrap();
2968 let r1 = converter.convert_columns(&[Arc::clone(&s1)]).unwrap();
2969
2970 for (a, b) in r1.iter().zip(r1.iter().skip(1)) {
2971 assert!(a < b);
2972 }
2973
2974 let back = converter.convert_rows(&r1).unwrap();
2975 assert_eq!(back.len(), 1);
2976 assert_eq!(&back[0], &s1);
2977
2978 let data = s1
2980 .to_data()
2981 .into_builder()
2982 .null_bit_buffer(Some(Buffer::from_slice_ref([0b00001010])))
2983 .null_count(2)
2984 .build()
2985 .unwrap();
2986
2987 let s2 = Arc::new(StructArray::from(data)) as ArrayRef;
2988 let r2 = converter.convert_columns(&[Arc::clone(&s2)]).unwrap();
2989 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();
2995 assert_eq!(back.len(), 1);
2996 assert_eq!(&back[0], &s2);
2997
2998 back[0].to_data().validate_full().unwrap();
2999 }
3000
3001 #[test]
3002 fn test_dictionary_in_struct() {
3003 let builder = StringDictionaryBuilder::<Int32Type>::new();
3004 let mut struct_builder = StructBuilder::new(
3005 vec![Field::new_dictionary(
3006 "foo",
3007 DataType::Int32,
3008 DataType::Utf8,
3009 true,
3010 )],
3011 vec![Box::new(builder)],
3012 );
3013
3014 let dict_builder = struct_builder
3015 .field_builder::<StringDictionaryBuilder<Int32Type>>(0)
3016 .unwrap();
3017
3018 dict_builder.append_value("a");
3020 dict_builder.append_null();
3021 dict_builder.append_value("a");
3022 dict_builder.append_value("b");
3023
3024 for _ in 0..4 {
3025 struct_builder.append(true);
3026 }
3027
3028 let s = Arc::new(struct_builder.finish()) as ArrayRef;
3029 let sort_fields = vec![SortField::new(s.data_type().clone())];
3030 let converter = RowConverter::new(sort_fields).unwrap();
3031 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3032
3033 let back = converter.convert_rows(&r).unwrap();
3034 let [s2] = back.try_into().unwrap();
3035
3036 assert_ne!(&s.data_type(), &s2.data_type());
3039 s2.to_data().validate_full().unwrap();
3040
3041 let s1_struct = s.as_struct();
3045 let s1_0 = s1_struct.column(0);
3046 let s1_idx_0 = s1_0.as_dictionary::<Int32Type>();
3047 let keys = s1_idx_0.keys();
3048 let values = s1_idx_0.values().as_string::<i32>();
3049 let s2_struct = s2.as_struct();
3051 let s2_0 = s2_struct.column(0);
3052 let s2_idx_0 = s2_0.as_string::<i32>();
3053
3054 for i in 0..keys.len() {
3055 if keys.is_null(i) {
3056 assert!(s2_idx_0.is_null(i));
3057 } else {
3058 let dict_index = keys.value(i) as usize;
3059 assert_eq!(values.value(dict_index), s2_idx_0.value(i));
3060 }
3061 }
3062 }
3063
3064 #[test]
3065 fn test_dictionary_in_struct_empty() {
3066 let ty = DataType::Struct(
3067 vec![Field::new_dictionary(
3068 "foo",
3069 DataType::Int32,
3070 DataType::Int32,
3071 false,
3072 )]
3073 .into(),
3074 );
3075 let s = arrow_array::new_empty_array(&ty);
3076
3077 let sort_fields = vec![SortField::new(s.data_type().clone())];
3078 let converter = RowConverter::new(sort_fields).unwrap();
3079 let r = converter.convert_columns(&[Arc::clone(&s)]).unwrap();
3080
3081 let back = converter.convert_rows(&r).unwrap();
3082 let [s2] = back.try_into().unwrap();
3083
3084 assert_ne!(&s.data_type(), &s2.data_type());
3087 s2.to_data().validate_full().unwrap();
3088 assert_eq!(s.len(), 0);
3089 assert_eq!(s2.len(), 0);
3090 }
3091
3092 #[test]
3093 fn test_list_of_string_dictionary() {
3094 let mut builder = ListBuilder::<StringDictionaryBuilder<Int32Type>>::default();
3095 builder.values().append("a").unwrap();
3097 builder.values().append("b").unwrap();
3098 builder.values().append("zero").unwrap();
3099 builder.values().append_null();
3100 builder.values().append("c").unwrap();
3101 builder.values().append("b").unwrap();
3102 builder.values().append("d").unwrap();
3103 builder.append(true);
3104 builder.append(false);
3106 builder.values().append("e").unwrap();
3108 builder.values().append("zero").unwrap();
3109 builder.values().append("a").unwrap();
3110 builder.append(true);
3111
3112 let a = Arc::new(builder.finish()) as ArrayRef;
3113 let data_type = a.data_type().clone();
3114
3115 let field = SortField::new(data_type.clone());
3116 let converter = RowConverter::new(vec![field]).unwrap();
3117 let rows = converter.convert_columns(&[Arc::clone(&a)]).unwrap();
3118
3119 let back = converter.convert_rows(&rows).unwrap();
3120 assert_eq!(back.len(), 1);
3121 let [a2] = back.try_into().unwrap();
3122
3123 assert_ne!(&a.data_type(), &a2.data_type());
3126
3127 a2.to_data().validate_full().unwrap();
3128
3129 let a2_list = a2.as_list::<i32>();
3130 let a1_list = a.as_list::<i32>();
3131
3132 let a1_0 = a1_list.value(0);
3135 let a1_idx_0 = a1_0.as_dictionary::<Int32Type>();
3136 let keys = a1_idx_0.keys();
3137 let values = a1_idx_0.values().as_string::<i32>();
3138 let a2_0 = a2_list.value(0);
3139 let a2_idx_0 = a2_0.as_string::<i32>();
3140
3141 for i in 0..keys.len() {
3142 if keys.is_null(i) {
3143 assert!(a2_idx_0.is_null(i));
3144 } else {
3145 let dict_index = keys.value(i) as usize;
3146 assert_eq!(values.value(dict_index), a2_idx_0.value(i));
3147 }
3148 }
3149
3150 assert!(a1_list.is_null(1));
3152 assert!(a2_list.is_null(1));
3153
3154 let a1_2 = a1_list.value(2);
3156 let a1_idx_2 = a1_2.as_dictionary::<Int32Type>();
3157 let keys = a1_idx_2.keys();
3158 let values = a1_idx_2.values().as_string::<i32>();
3159 let a2_2 = a2_list.value(2);
3160 let a2_idx_2 = a2_2.as_string::<i32>();
3161
3162 for i in 0..keys.len() {
3163 if keys.is_null(i) {
3164 assert!(a2_idx_2.is_null(i));
3165 } else {
3166 let dict_index = keys.value(i) as usize;
3167 assert_eq!(values.value(dict_index), a2_idx_2.value(i));
3168 }
3169 }
3170 }
3171
3172 #[test]
3173 fn test_primitive_dictionary() {
3174 let mut builder = PrimitiveDictionaryBuilder::<Int32Type, Int32Type>::new();
3175 builder.append(2).unwrap();
3176 builder.append(3).unwrap();
3177 builder.append(0).unwrap();
3178 builder.append_null();
3179 builder.append(5).unwrap();
3180 builder.append(3).unwrap();
3181 builder.append(-1).unwrap();
3182
3183 let a = builder.finish();
3184 let data_type = a.data_type().clone();
3185 let columns = [Arc::new(a) as ArrayRef];
3186
3187 let field = SortField::new(data_type.clone());
3188 let converter = RowConverter::new(vec![field]).unwrap();
3189 let rows = converter.convert_columns(&columns).unwrap();
3190 assert!(rows.row(0) < rows.row(1));
3191 assert!(rows.row(2) < rows.row(0));
3192 assert!(rows.row(3) < rows.row(2));
3193 assert!(rows.row(6) < rows.row(2));
3194 assert!(rows.row(3) < rows.row(6));
3195
3196 let back = converter.convert_rows(&rows).unwrap();
3197 assert_eq!(back.len(), 1);
3198 back[0].to_data().validate_full().unwrap();
3199 }
3200
3201 #[test]
3202 fn test_dictionary_nulls() {
3203 let values = Int32Array::from_iter([Some(1), Some(-1), None, Some(4), None]).into_data();
3204 let keys =
3205 Int32Array::from_iter([Some(0), Some(0), Some(1), Some(2), Some(4), None]).into_data();
3206
3207 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3208 let data = keys
3209 .into_builder()
3210 .data_type(data_type.clone())
3211 .child_data(vec![values])
3212 .build()
3213 .unwrap();
3214
3215 let columns = [Arc::new(DictionaryArray::<Int32Type>::from(data)) as ArrayRef];
3216 let field = SortField::new(data_type.clone());
3217 let converter = RowConverter::new(vec![field]).unwrap();
3218 let rows = converter.convert_columns(&columns).unwrap();
3219
3220 assert_eq!(rows.row(0), rows.row(1));
3221 assert_eq!(rows.row(3), rows.row(4));
3222 assert_eq!(rows.row(4), rows.row(5));
3223 assert!(rows.row(3) < rows.row(0));
3224 }
3225
3226 #[test]
3227 fn test_from_binary_shared_buffer() {
3228 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3229 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3230 let rows = converter.convert_columns(&[array]).unwrap();
3231 let binary_rows = rows.try_into_binary().expect("known-small rows");
3232 let _binary_rows_shared_buffer = binary_rows.clone();
3233
3234 let parsed = converter.from_binary(binary_rows);
3235
3236 converter.convert_rows(parsed.iter()).unwrap();
3237 }
3238
3239 #[test]
3240 #[should_panic(expected = "Encountered non UTF-8 data")]
3241 fn test_invalid_utf8() {
3242 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3243 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3244 let rows = converter.convert_columns(&[array]).unwrap();
3245 let binary_row = rows.row(0);
3246
3247 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3248 let parser = converter.parser();
3249 let utf8_row = parser.parse(binary_row.as_ref());
3250
3251 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3252 }
3253
3254 #[test]
3255 #[should_panic(expected = "Encountered non UTF-8 data")]
3256 fn test_invalid_utf8_array() {
3257 let converter = RowConverter::new(vec![SortField::new(DataType::Binary)]).unwrap();
3258 let array = Arc::new(BinaryArray::from_iter_values([&[0xFF]])) as _;
3259 let rows = converter.convert_columns(&[array]).unwrap();
3260 let binary_rows = rows.try_into_binary().expect("known-small rows");
3261
3262 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3263 let parsed = converter.from_binary(binary_rows);
3264
3265 converter.convert_rows(parsed.iter()).unwrap();
3266 }
3267
3268 #[test]
3269 #[should_panic(expected = "index out of bounds")]
3270 fn test_invalid_empty() {
3271 let binary_row: &[u8] = &[];
3272
3273 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3274 let parser = converter.parser();
3275 let utf8_row = parser.parse(binary_row.as_ref());
3276
3277 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3278 }
3279
3280 #[test]
3281 #[should_panic(expected = "index out of bounds")]
3282 fn test_invalid_empty_array() {
3283 let row: &[u8] = &[];
3284 let binary_rows = BinaryArray::from(vec![row]);
3285
3286 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3287 let parsed = converter.from_binary(binary_rows);
3288
3289 converter.convert_rows(parsed.iter()).unwrap();
3290 }
3291
3292 #[test]
3293 #[should_panic(expected = "index out of bounds")]
3294 fn test_invalid_truncated() {
3295 let binary_row: &[u8] = &[0x02];
3296
3297 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3298 let parser = converter.parser();
3299 let utf8_row = parser.parse(binary_row.as_ref());
3300
3301 converter.convert_rows(std::iter::once(utf8_row)).unwrap();
3302 }
3303
3304 #[test]
3305 #[should_panic(expected = "index out of bounds")]
3306 fn test_invalid_truncated_array() {
3307 let row: &[u8] = &[0x02];
3308 let binary_rows = BinaryArray::from(vec![row]);
3309
3310 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
3311 let parsed = converter.from_binary(binary_rows);
3312
3313 converter.convert_rows(parsed.iter()).unwrap();
3314 }
3315
3316 #[test]
3317 #[should_panic(expected = "rows were not produced by this RowConverter")]
3318 fn test_different_converter() {
3319 let values = Arc::new(Int32Array::from_iter([Some(1), Some(-1)]));
3320 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3321 let rows = converter.convert_columns(&[values]).unwrap();
3322
3323 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
3324 let _ = converter.convert_rows(&rows);
3325 }
3326
3327 fn test_single_list<O: OffsetSizeTrait>() {
3328 let mut builder = GenericListBuilder::<O, _>::new(Int32Builder::new());
3329 builder.values().append_value(32);
3330 builder.values().append_value(52);
3331 builder.values().append_value(32);
3332 builder.append(true);
3333 builder.values().append_value(32);
3334 builder.values().append_value(52);
3335 builder.values().append_value(12);
3336 builder.append(true);
3337 builder.values().append_value(32);
3338 builder.values().append_value(52);
3339 builder.append(true);
3340 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3343 builder.values().append_value(32);
3344 builder.values().append_null();
3345 builder.append(true);
3346 builder.append(true);
3347 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3350
3351 let list = Arc::new(builder.finish()) as ArrayRef;
3352 let d = list.data_type().clone();
3353
3354 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3355
3356 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3357 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();
3366 assert_eq!(back.len(), 1);
3367 back[0].to_data().validate_full().unwrap();
3368 assert_eq!(&back[0], &list);
3369
3370 let options = SortOptions::default().asc().with_nulls_first(false);
3371 let field = SortField::new_with_options(d.clone(), options);
3372 let converter = RowConverter::new(vec![field]).unwrap();
3373 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3374
3375 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();
3384 assert_eq!(back.len(), 1);
3385 back[0].to_data().validate_full().unwrap();
3386 assert_eq!(&back[0], &list);
3387
3388 let options = SortOptions::default().desc().with_nulls_first(false);
3389 let field = SortField::new_with_options(d.clone(), options);
3390 let converter = RowConverter::new(vec![field]).unwrap();
3391 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3392
3393 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();
3402 assert_eq!(back.len(), 1);
3403 back[0].to_data().validate_full().unwrap();
3404 assert_eq!(&back[0], &list);
3405
3406 let options = SortOptions::default().desc().with_nulls_first(true);
3407 let field = SortField::new_with_options(d, options);
3408 let converter = RowConverter::new(vec![field]).unwrap();
3409 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3410
3411 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();
3420 assert_eq!(back.len(), 1);
3421 back[0].to_data().validate_full().unwrap();
3422 assert_eq!(&back[0], &list);
3423
3424 let sliced_list = list.slice(1, 5);
3425 let rows_on_sliced_list = converter
3426 .convert_columns(&[Arc::clone(&sliced_list)])
3427 .unwrap();
3428
3429 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();
3436 assert_eq!(back.len(), 1);
3437 back[0].to_data().validate_full().unwrap();
3438 assert_eq!(&back[0], &sliced_list);
3439 }
3440
3441 fn test_nested_list<O: OffsetSizeTrait>() {
3442 let mut builder =
3443 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(Int32Builder::new()));
3444
3445 builder.values().values().append_value(1);
3446 builder.values().values().append_value(2);
3447 builder.values().append(true);
3448 builder.values().values().append_value(1);
3449 builder.values().values().append_null();
3450 builder.values().append(true);
3451 builder.append(true);
3452
3453 builder.values().values().append_value(1);
3454 builder.values().values().append_null();
3455 builder.values().append(true);
3456 builder.values().values().append_value(1);
3457 builder.values().values().append_null();
3458 builder.values().append(true);
3459 builder.append(true);
3460
3461 builder.values().values().append_value(1);
3462 builder.values().values().append_null();
3463 builder.values().append(true);
3464 builder.values().append(false);
3465 builder.append(true);
3466 builder.append(false);
3467
3468 builder.values().values().append_value(1);
3469 builder.values().values().append_value(2);
3470 builder.values().append(true);
3471 builder.append(true);
3472
3473 let list = Arc::new(builder.finish()) as ArrayRef;
3474 let d = list.data_type().clone();
3475
3476 let options = SortOptions::default().asc().with_nulls_first(true);
3484 let field = SortField::new_with_options(d.clone(), options);
3485 let converter = RowConverter::new(vec![field]).unwrap();
3486 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3487
3488 assert!(rows.row(0) > rows.row(1));
3489 assert!(rows.row(1) > rows.row(2));
3490 assert!(rows.row(2) > rows.row(3));
3491 assert!(rows.row(4) < rows.row(0));
3492 assert!(rows.row(4) > rows.row(1));
3493
3494 let back = converter.convert_rows(&rows).unwrap();
3495 assert_eq!(back.len(), 1);
3496 back[0].to_data().validate_full().unwrap();
3497 assert_eq!(&back[0], &list);
3498
3499 let options = SortOptions::default().desc().with_nulls_first(true);
3500 let field = SortField::new_with_options(d.clone(), options);
3501 let converter = RowConverter::new(vec![field]).unwrap();
3502 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3503
3504 assert!(rows.row(0) > rows.row(1));
3505 assert!(rows.row(1) > rows.row(2));
3506 assert!(rows.row(2) > rows.row(3));
3507 assert!(rows.row(4) > rows.row(0));
3508 assert!(rows.row(4) > rows.row(1));
3509
3510 let back = converter.convert_rows(&rows).unwrap();
3511 assert_eq!(back.len(), 1);
3512 back[0].to_data().validate_full().unwrap();
3513 assert_eq!(&back[0], &list);
3514
3515 let options = SortOptions::default().desc().with_nulls_first(false);
3516 let field = SortField::new_with_options(d, options);
3517 let converter = RowConverter::new(vec![field]).unwrap();
3518 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3519
3520 assert!(rows.row(0) < rows.row(1));
3521 assert!(rows.row(1) < rows.row(2));
3522 assert!(rows.row(2) < rows.row(3));
3523 assert!(rows.row(4) > rows.row(0));
3524 assert!(rows.row(4) < rows.row(1));
3525
3526 let back = converter.convert_rows(&rows).unwrap();
3527 assert_eq!(back.len(), 1);
3528 back[0].to_data().validate_full().unwrap();
3529 assert_eq!(&back[0], &list);
3530
3531 let sliced_list = list.slice(1, 3);
3532 let rows = converter
3533 .convert_columns(&[Arc::clone(&sliced_list)])
3534 .unwrap();
3535
3536 assert!(rows.row(0) < rows.row(1));
3537 assert!(rows.row(1) < rows.row(2));
3538
3539 let back = converter.convert_rows(&rows).unwrap();
3540 assert_eq!(back.len(), 1);
3541 back[0].to_data().validate_full().unwrap();
3542 assert_eq!(&back[0], &sliced_list);
3543 }
3544
3545 #[test]
3546 fn test_list() {
3547 test_single_list::<i32>();
3548 test_nested_list::<i32>();
3549 }
3550
3551 #[test]
3552 fn test_large_list() {
3553 test_single_list::<i64>();
3554 test_nested_list::<i64>();
3555 }
3556
3557 fn test_single_list_view<O: OffsetSizeTrait>() {
3558 let mut builder = GenericListViewBuilder::<O, _>::new(Int32Builder::new());
3559 builder.values().append_value(32);
3560 builder.values().append_value(52);
3561 builder.values().append_value(32);
3562 builder.append(true);
3563 builder.values().append_value(32);
3564 builder.values().append_value(52);
3565 builder.values().append_value(12);
3566 builder.append(true);
3567 builder.values().append_value(32);
3568 builder.values().append_value(52);
3569 builder.append(true);
3570 builder.values().append_value(32); builder.values().append_value(52); builder.append(false);
3573 builder.values().append_value(32);
3574 builder.values().append_null();
3575 builder.append(true);
3576 builder.append(true);
3577 builder.values().append_value(17); builder.values().append_null(); builder.append(false);
3580
3581 let list = Arc::new(builder.finish()) as ArrayRef;
3582 let d = list.data_type().clone();
3583
3584 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3585
3586 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3587 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();
3596 assert_eq!(back.len(), 1);
3597 back[0].to_data().validate_full().unwrap();
3598
3599 let back_list_view = back[0]
3601 .as_any()
3602 .downcast_ref::<GenericListViewArray<O>>()
3603 .unwrap();
3604 let orig_list_view = list
3605 .as_any()
3606 .downcast_ref::<GenericListViewArray<O>>()
3607 .unwrap();
3608
3609 assert_eq!(back_list_view.len(), orig_list_view.len());
3610 for i in 0..back_list_view.len() {
3611 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3612 if back_list_view.is_valid(i) {
3613 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3614 }
3615 }
3616
3617 let options = SortOptions::default().asc().with_nulls_first(false);
3618 let field = SortField::new_with_options(d.clone(), options);
3619 let converter = RowConverter::new(vec![field]).unwrap();
3620 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3621
3622 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();
3631 assert_eq!(back.len(), 1);
3632 back[0].to_data().validate_full().unwrap();
3633
3634 let options = SortOptions::default().desc().with_nulls_first(false);
3635 let field = SortField::new_with_options(d.clone(), options);
3636 let converter = RowConverter::new(vec![field]).unwrap();
3637 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3638
3639 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();
3648 assert_eq!(back.len(), 1);
3649 back[0].to_data().validate_full().unwrap();
3650
3651 let options = SortOptions::default().desc().with_nulls_first(true);
3652 let field = SortField::new_with_options(d, options);
3653 let converter = RowConverter::new(vec![field]).unwrap();
3654 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3655
3656 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();
3665 assert_eq!(back.len(), 1);
3666 back[0].to_data().validate_full().unwrap();
3667
3668 let sliced_list = list.slice(1, 5);
3669 let rows_on_sliced_list = converter
3670 .convert_columns(&[Arc::clone(&sliced_list)])
3671 .unwrap();
3672
3673 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();
3680 assert_eq!(back.len(), 1);
3681 back[0].to_data().validate_full().unwrap();
3682 }
3683
3684 fn test_nested_list_view<O: OffsetSizeTrait>() {
3685 let mut builder = GenericListViewBuilder::<O, _>::new(GenericListViewBuilder::<O, _>::new(
3686 Int32Builder::new(),
3687 ));
3688
3689 builder.values().values().append_value(1);
3691 builder.values().values().append_value(2);
3692 builder.values().append(true);
3693 builder.values().values().append_value(1);
3694 builder.values().values().append_null();
3695 builder.values().append(true);
3696 builder.append(true);
3697
3698 builder.values().values().append_value(1);
3700 builder.values().values().append_null();
3701 builder.values().append(true);
3702 builder.values().values().append_value(1);
3703 builder.values().values().append_null();
3704 builder.values().append(true);
3705 builder.append(true);
3706
3707 builder.values().values().append_value(1);
3709 builder.values().values().append_null();
3710 builder.values().append(true);
3711 builder.values().append(false);
3712 builder.append(true);
3713
3714 builder.append(false);
3716
3717 builder.values().values().append_value(1);
3719 builder.values().values().append_value(2);
3720 builder.values().append(true);
3721 builder.append(true);
3722
3723 let list = Arc::new(builder.finish()) as ArrayRef;
3724 let d = list.data_type().clone();
3725
3726 let options = SortOptions::default().asc().with_nulls_first(true);
3734 let field = SortField::new_with_options(d.clone(), options);
3735 let converter = RowConverter::new(vec![field]).unwrap();
3736 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3737
3738 assert!(rows.row(0) > rows.row(1));
3739 assert!(rows.row(1) > rows.row(2));
3740 assert!(rows.row(2) > rows.row(3));
3741 assert!(rows.row(4) < rows.row(0));
3742 assert!(rows.row(4) > rows.row(1));
3743
3744 let back = converter.convert_rows(&rows).unwrap();
3745 assert_eq!(back.len(), 1);
3746 back[0].to_data().validate_full().unwrap();
3747
3748 let back_list_view = back[0]
3750 .as_any()
3751 .downcast_ref::<GenericListViewArray<O>>()
3752 .unwrap();
3753 let orig_list_view = list
3754 .as_any()
3755 .downcast_ref::<GenericListViewArray<O>>()
3756 .unwrap();
3757
3758 assert_eq!(back_list_view.len(), orig_list_view.len());
3759 for i in 0..back_list_view.len() {
3760 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3761 if back_list_view.is_valid(i) {
3762 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3763 }
3764 }
3765
3766 let options = SortOptions::default().desc().with_nulls_first(true);
3767 let field = SortField::new_with_options(d.clone(), options);
3768 let converter = RowConverter::new(vec![field]).unwrap();
3769 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3770
3771 assert!(rows.row(0) > rows.row(1));
3772 assert!(rows.row(1) > rows.row(2));
3773 assert!(rows.row(2) > rows.row(3));
3774 assert!(rows.row(4) > rows.row(0));
3775 assert!(rows.row(4) > rows.row(1));
3776
3777 let back = converter.convert_rows(&rows).unwrap();
3778 assert_eq!(back.len(), 1);
3779 back[0].to_data().validate_full().unwrap();
3780
3781 let back_list_view = back[0]
3783 .as_any()
3784 .downcast_ref::<GenericListViewArray<O>>()
3785 .unwrap();
3786
3787 assert_eq!(back_list_view.len(), orig_list_view.len());
3788 for i in 0..back_list_view.len() {
3789 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3790 if back_list_view.is_valid(i) {
3791 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3792 }
3793 }
3794
3795 let options = SortOptions::default().desc().with_nulls_first(false);
3796 let field = SortField::new_with_options(d.clone(), options);
3797 let converter = RowConverter::new(vec![field]).unwrap();
3798 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3799
3800 assert!(rows.row(0) < rows.row(1));
3801 assert!(rows.row(1) < rows.row(2));
3802 assert!(rows.row(2) < rows.row(3));
3803 assert!(rows.row(4) > rows.row(0));
3804 assert!(rows.row(4) < rows.row(1));
3805
3806 let back = converter.convert_rows(&rows).unwrap();
3807 assert_eq!(back.len(), 1);
3808 back[0].to_data().validate_full().unwrap();
3809
3810 let back_list_view = back[0]
3812 .as_any()
3813 .downcast_ref::<GenericListViewArray<O>>()
3814 .unwrap();
3815
3816 assert_eq!(back_list_view.len(), orig_list_view.len());
3817 for i in 0..back_list_view.len() {
3818 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3819 if back_list_view.is_valid(i) {
3820 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3821 }
3822 }
3823
3824 let sliced_list = list.slice(1, 3);
3825 let rows = converter
3826 .convert_columns(&[Arc::clone(&sliced_list)])
3827 .unwrap();
3828
3829 assert!(rows.row(0) < rows.row(1));
3830 assert!(rows.row(1) < rows.row(2));
3831
3832 let back = converter.convert_rows(&rows).unwrap();
3833 assert_eq!(back.len(), 1);
3834 back[0].to_data().validate_full().unwrap();
3835 }
3836
3837 #[test]
3838 fn test_list_view() {
3839 test_single_list_view::<i32>();
3840 test_nested_list_view::<i32>();
3841 }
3842
3843 #[test]
3844 fn test_large_list_view() {
3845 test_single_list_view::<i64>();
3846 test_nested_list_view::<i64>();
3847 }
3848
3849 fn test_list_view_with_shared_values<O: OffsetSizeTrait>() {
3850 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]);
3852 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
3853
3854 let offsets = ScalarBuffer::<O>::from(vec![
3862 O::from_usize(0).unwrap(),
3863 O::from_usize(0).unwrap(),
3864 O::from_usize(5).unwrap(),
3865 O::from_usize(2).unwrap(),
3866 O::from_usize(1).unwrap(),
3867 O::from_usize(2).unwrap(),
3868 ]);
3869 let sizes = ScalarBuffer::<O>::from(vec![
3870 O::from_usize(3).unwrap(),
3871 O::from_usize(3).unwrap(),
3872 O::from_usize(2).unwrap(),
3873 O::from_usize(2).unwrap(),
3874 O::from_usize(4).unwrap(),
3875 O::from_usize(1).unwrap(),
3876 ]);
3877
3878 let list_view: GenericListViewArray<O> =
3879 GenericListViewArray::try_new(field, offsets, sizes, Arc::new(values), None).unwrap();
3880
3881 let d = list_view.data_type().clone();
3882 let list = Arc::new(list_view) as ArrayRef;
3883
3884 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3885 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3886
3887 assert_eq!(rows.row(0), rows.row(1));
3889
3890 assert!(rows.row(0) < rows.row(2));
3892
3893 assert!(rows.row(3) > rows.row(0));
3895
3896 assert!(rows.row(4) > rows.row(0));
3898
3899 assert!(rows.row(5) < rows.row(3));
3901
3902 assert!(rows.row(5) > rows.row(4));
3904
3905 let back = converter.convert_rows(&rows).unwrap();
3907 assert_eq!(back.len(), 1);
3908 back[0].to_data().validate_full().unwrap();
3909
3910 let back_list_view = back[0]
3912 .as_any()
3913 .downcast_ref::<GenericListViewArray<O>>()
3914 .unwrap();
3915 let orig_list_view = list
3916 .as_any()
3917 .downcast_ref::<GenericListViewArray<O>>()
3918 .unwrap();
3919
3920 assert_eq!(back_list_view.len(), orig_list_view.len());
3921 for i in 0..back_list_view.len() {
3922 assert_eq!(back_list_view.is_valid(i), orig_list_view.is_valid(i));
3923 if back_list_view.is_valid(i) {
3924 assert_eq!(&back_list_view.value(i), &orig_list_view.value(i));
3925 }
3926 }
3927
3928 let options = SortOptions::default().desc();
3930 let field = SortField::new_with_options(d, options);
3931 let converter = RowConverter::new(vec![field]).unwrap();
3932 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3933
3934 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();
3940 assert_eq!(back.len(), 1);
3941 back[0].to_data().validate_full().unwrap();
3942 }
3943
3944 #[test]
3945 fn test_list_view_shared_values() {
3946 test_list_view_with_shared_values::<i32>();
3947 }
3948
3949 #[test]
3950 fn test_large_list_view_shared_values() {
3951 test_list_view_with_shared_values::<i64>();
3952 }
3953
3954 #[test]
3955 fn test_fixed_size_list() {
3956 let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 3);
3957 builder.values().append_value(32);
3958 builder.values().append_value(52);
3959 builder.values().append_value(32);
3960 builder.append(true);
3961 builder.values().append_value(32);
3962 builder.values().append_value(52);
3963 builder.values().append_value(12);
3964 builder.append(true);
3965 builder.values().append_value(32);
3966 builder.values().append_value(52);
3967 builder.values().append_null();
3968 builder.append(true);
3969 builder.values().append_value(32); builder.values().append_value(52); builder.values().append_value(13); builder.append(false);
3973 builder.values().append_value(32);
3974 builder.values().append_null();
3975 builder.values().append_null();
3976 builder.append(true);
3977 builder.values().append_null();
3978 builder.values().append_null();
3979 builder.values().append_null();
3980 builder.append(true);
3981 builder.values().append_value(17); builder.values().append_null(); builder.values().append_value(77); builder.append(false);
3985
3986 let list = Arc::new(builder.finish()) as ArrayRef;
3987 let d = list.data_type().clone();
3988
3989 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
3991
3992 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
3993 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();
4002 assert_eq!(back.len(), 1);
4003 back[0].to_data().validate_full().unwrap();
4004 assert_eq!(&back[0], &list);
4005
4006 let options = SortOptions::default().asc().with_nulls_first(false);
4008 let field = SortField::new_with_options(d.clone(), options);
4009 let converter = RowConverter::new(vec![field]).unwrap();
4010 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4011 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();
4020 assert_eq!(back.len(), 1);
4021 back[0].to_data().validate_full().unwrap();
4022 assert_eq!(&back[0], &list);
4023
4024 let options = SortOptions::default().desc().with_nulls_first(false);
4026 let field = SortField::new_with_options(d.clone(), options);
4027 let converter = RowConverter::new(vec![field]).unwrap();
4028 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4029 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();
4038 assert_eq!(back.len(), 1);
4039 back[0].to_data().validate_full().unwrap();
4040 assert_eq!(&back[0], &list);
4041
4042 let options = SortOptions::default().desc().with_nulls_first(true);
4044 let field = SortField::new_with_options(d, options);
4045 let converter = RowConverter::new(vec![field]).unwrap();
4046 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
4047
4048 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();
4057 assert_eq!(back.len(), 1);
4058 back[0].to_data().validate_full().unwrap();
4059 assert_eq!(&back[0], &list);
4060
4061 let sliced_list = list.slice(1, 5);
4062 let rows_on_sliced_list = converter
4063 .convert_columns(&[Arc::clone(&sliced_list)])
4064 .unwrap();
4065
4066 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();
4072 assert_eq!(back.len(), 1);
4073 back[0].to_data().validate_full().unwrap();
4074 assert_eq!(&back[0], &sliced_list);
4075 }
4076
4077 #[test]
4078 fn test_two_fixed_size_lists() {
4079 let mut first = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4080 first.values().append_value(100);
4082 first.append(true);
4083 first.values().append_value(101);
4085 first.append(true);
4086 first.values().append_value(102);
4088 first.append(true);
4089 first.values().append_null();
4091 first.append(true);
4092 first.values().append_null(); first.append(false);
4095 let first = Arc::new(first.finish()) as ArrayRef;
4096 let first_type = first.data_type().clone();
4097
4098 let mut second = FixedSizeListBuilder::new(UInt8Builder::new(), 1);
4099 second.values().append_value(200);
4101 second.append(true);
4102 second.values().append_value(201);
4104 second.append(true);
4105 second.values().append_value(202);
4107 second.append(true);
4108 second.values().append_null();
4110 second.append(true);
4111 second.values().append_null(); second.append(false);
4114 let second = Arc::new(second.finish()) as ArrayRef;
4115 let second_type = second.data_type().clone();
4116
4117 let converter = RowConverter::new(vec![
4118 SortField::new(first_type.clone()),
4119 SortField::new(second_type.clone()),
4120 ])
4121 .unwrap();
4122
4123 let rows = converter
4124 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4125 .unwrap();
4126
4127 let back = converter.convert_rows(&rows).unwrap();
4128 assert_eq!(back.len(), 2);
4129 back[0].to_data().validate_full().unwrap();
4130 assert_eq!(&back[0], &first);
4131 back[1].to_data().validate_full().unwrap();
4132 assert_eq!(&back[1], &second);
4133 }
4134
4135 #[test]
4136 fn test_fixed_size_list_with_variable_width_content() {
4137 let mut first = FixedSizeListBuilder::new(
4138 StructBuilder::from_fields(
4139 vec![
4140 Field::new(
4141 "timestamp",
4142 DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("UTC"))),
4143 false,
4144 ),
4145 Field::new("offset_minutes", DataType::Int16, false),
4146 Field::new("time_zone", DataType::Utf8, false),
4147 ],
4148 1,
4149 ),
4150 1,
4151 );
4152 first
4154 .values()
4155 .field_builder::<TimestampMicrosecondBuilder>(0)
4156 .unwrap()
4157 .append_null();
4158 first
4159 .values()
4160 .field_builder::<Int16Builder>(1)
4161 .unwrap()
4162 .append_null();
4163 first
4164 .values()
4165 .field_builder::<StringBuilder>(2)
4166 .unwrap()
4167 .append_null();
4168 first.values().append(false);
4169 first.append(false);
4170 first
4172 .values()
4173 .field_builder::<TimestampMicrosecondBuilder>(0)
4174 .unwrap()
4175 .append_null();
4176 first
4177 .values()
4178 .field_builder::<Int16Builder>(1)
4179 .unwrap()
4180 .append_null();
4181 first
4182 .values()
4183 .field_builder::<StringBuilder>(2)
4184 .unwrap()
4185 .append_null();
4186 first.values().append(false);
4187 first.append(true);
4188 first
4190 .values()
4191 .field_builder::<TimestampMicrosecondBuilder>(0)
4192 .unwrap()
4193 .append_value(0);
4194 first
4195 .values()
4196 .field_builder::<Int16Builder>(1)
4197 .unwrap()
4198 .append_value(0);
4199 first
4200 .values()
4201 .field_builder::<StringBuilder>(2)
4202 .unwrap()
4203 .append_value("UTC");
4204 first.values().append(true);
4205 first.append(true);
4206 first
4208 .values()
4209 .field_builder::<TimestampMicrosecondBuilder>(0)
4210 .unwrap()
4211 .append_value(1126351800123456);
4212 first
4213 .values()
4214 .field_builder::<Int16Builder>(1)
4215 .unwrap()
4216 .append_value(120);
4217 first
4218 .values()
4219 .field_builder::<StringBuilder>(2)
4220 .unwrap()
4221 .append_value("Europe/Warsaw");
4222 first.values().append(true);
4223 first.append(true);
4224 let first = Arc::new(first.finish()) as ArrayRef;
4225 let first_type = first.data_type().clone();
4226
4227 let mut second = StringBuilder::new();
4228 second.append_value("somewhere near");
4229 second.append_null();
4230 second.append_value("Greenwich");
4231 second.append_value("Warsaw");
4232 let second = Arc::new(second.finish()) as ArrayRef;
4233 let second_type = second.data_type().clone();
4234
4235 let converter = RowConverter::new(vec![
4236 SortField::new(first_type.clone()),
4237 SortField::new(second_type.clone()),
4238 ])
4239 .unwrap();
4240
4241 let rows = converter
4242 .convert_columns(&[Arc::clone(&first), Arc::clone(&second)])
4243 .unwrap();
4244
4245 let back = converter.convert_rows(&rows).unwrap();
4246 assert_eq!(back.len(), 2);
4247 back[0].to_data().validate_full().unwrap();
4248 assert_eq!(&back[0], &first);
4249 back[1].to_data().validate_full().unwrap();
4250 assert_eq!(&back[1], &second);
4251 }
4252
4253 #[test]
4254 fn test_single_map() {
4255 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4256 builder.keys().append_value("hello");
4258 builder.values().append_value(1);
4259 builder.keys().append_value("world");
4260 builder.values().append_value(2);
4261 builder.append(true).unwrap();
4262
4263 builder.keys().append_value("foo");
4265 builder.values().append_value(3);
4266 builder.append(true).unwrap();
4267
4268 builder.append(true).unwrap();
4270
4271 builder.keys().append_value("masked_key");
4273 builder.values().append_value(999);
4274 builder.append(false).unwrap();
4275
4276 builder.append(false).unwrap();
4278
4279 builder.keys().append_value("bar");
4281 builder.values().append_null();
4282 builder.append(true).unwrap();
4283
4284 builder.keys().append_value("other_masked");
4286 builder.values().append_value(0);
4287 builder.append(false).unwrap();
4288
4289 builder.keys().append_value("a");
4291 builder.values().append_value(10);
4292 builder.keys().append_value("b");
4293 builder.values().append_value(20);
4294 builder.keys().append_value("c");
4295 builder.values().append_value(30);
4296 builder.append(true).unwrap();
4297
4298 let map = Arc::new(builder.finish()) as ArrayRef;
4299 let d = map.data_type().clone();
4300
4301 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4302
4303 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4304
4305 assert_eq!(rows.row(3), rows.row(4));
4307 assert_eq!(rows.row(4), rows.row(6));
4308
4309 let back = converter.convert_rows(&rows).unwrap();
4310 assert_eq!(back.len(), 1);
4311 back[0].to_data().validate_full().unwrap();
4312 assert_eq!(&back[0], &map);
4313
4314 let sliced_map = map.slice(1, map.len() - 2);
4315 let rows_on_sliced = converter
4316 .convert_columns(&[Arc::clone(&sliced_map)])
4317 .unwrap();
4318
4319 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4320 assert_eq!(back.len(), 1);
4321 back[0].to_data().validate_full().unwrap();
4322 assert_eq!(&back[0], &sliced_map);
4323 }
4324
4325 #[test]
4326 fn two_maps_with_different_keys_order_should_sort_by_entry_order() {
4327 let map_1: ArrayRef =
4328 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4329 vec![Some(vec![("hello", Some(1)), ("world", Some(2))])],
4330 false,
4331 ));
4332 let map_2: ArrayRef =
4334 Arc::new(MapArray::from_vec_of_maps::<StringArray, Int32Array, _, _>(
4335 vec![Some(vec![("world", Some(2)), ("hello", Some(1))])],
4336 false,
4337 ));
4338
4339 let converter = RowConverter::new(vec![SortField::new(map_1.data_type().clone())]).unwrap();
4340
4341 let map_1_rows = converter.convert_columns(&[Arc::clone(&map_1)]).unwrap();
4342 let map_2_rows = converter.convert_columns(&[Arc::clone(&map_2)]).unwrap();
4343
4344 assert_ne!(map_1_rows.row(0), map_2_rows.row(0));
4345 assert!(map_1_rows.row(0) < map_2_rows.row(0));
4346
4347 let back_1 = converter.convert_rows(&map_1_rows).unwrap();
4348 let back_2 = converter.convert_rows(&map_2_rows).unwrap();
4349 assert_eq!(&back_1[0], &map_1);
4350 assert_eq!(&back_2[0], &map_2);
4351 }
4352
4353 #[test]
4354 fn test_nested_map() {
4355 let mut builder = MapBuilder::new(
4357 None,
4358 StringBuilder::new(),
4359 MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()),
4360 );
4361
4362 builder.keys().append_value("outer1");
4364 builder.values().keys().append_value("inner_a");
4365 builder.values().values().append_value(1);
4366 builder.values().keys().append_value("inner_b");
4367 builder.values().values().append_value(2);
4368 builder.values().append(true).unwrap();
4369 builder.keys().append_value("outer2");
4370 builder.values().keys().append_value("inner_c");
4371 builder.values().values().append_value(3);
4372 builder.values().append(true).unwrap();
4373 builder.append(true).unwrap();
4374
4375 builder.keys().append_value("x");
4377 builder.values().append(true).unwrap();
4378 builder.append(true).unwrap();
4379
4380 builder.keys().append_value("y");
4382 builder.values().keys().append_value("masked"); builder.values().values().append_value(0); builder.values().append(false).unwrap();
4385 builder.append(true).unwrap();
4386
4387 builder.keys().append_value("y");
4389 builder.values().append(false).unwrap(); builder.append(true).unwrap();
4391
4392 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();
4398
4399 builder.keys().append_value("masked_outer"); builder.values().append(false).unwrap(); builder.append(false).unwrap();
4403
4404 builder.append(false).unwrap(); builder.append(true).unwrap();
4409
4410 let map = Arc::new(builder.finish()) as ArrayRef;
4411 let d = map.data_type().clone();
4412
4413 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4414
4415 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4416
4417 let back = converter.convert_rows(&rows).unwrap();
4418 assert_eq!(back.len(), 1);
4419 back[0].to_data().validate_full().unwrap();
4420 assert_eq!(&back[0], &map);
4421
4422 let sliced_map = map.slice(1, 3);
4423 let rows_on_sliced = converter
4424 .convert_columns(&[Arc::clone(&sliced_map)])
4425 .unwrap();
4426
4427 let back = converter.convert_rows(&rows_on_sliced).unwrap();
4428 assert_eq!(back.len(), 1);
4429 back[0].to_data().validate_full().unwrap();
4430 assert_eq!(&back[0], &sliced_map);
4431 }
4432
4433 #[test]
4434 fn test_single_map_with_non_nullable_values() {
4435 let value_field = Arc::new(Field::new(
4437 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4438 DataType::Int32,
4439 false,
4440 ));
4441 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4442 .with_values_field(value_field);
4443 builder.keys().append_value("a");
4445 builder.values().append_value(1);
4446 builder.keys().append_value("b");
4447 builder.values().append_value(2);
4448 builder.append(true).unwrap();
4449 builder.append(false).unwrap();
4451 builder.keys().append_value("c");
4453 builder.values().append_value(3);
4454 builder.append(true).unwrap();
4455 builder.append(true).unwrap();
4457 builder.keys().append_value("masked"); builder.values().append_value(0); builder.append(false).unwrap();
4461
4462 let map = Arc::new(builder.finish()) as ArrayRef;
4463 let d = map.data_type().clone();
4464
4465 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4466
4467 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4468
4469 let back = converter.convert_rows(&rows).unwrap();
4470 assert_eq!(back.len(), 1);
4471 back[0].to_data().validate_full().unwrap();
4472 assert_eq!(&back[0], &map);
4473 }
4474
4475 #[test]
4476 fn test_single_map_with_non_nullable_map_but_with_nullable_values() {
4477 let value_field = Arc::new(Field::new(
4479 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4480 DataType::Int32,
4481 true,
4482 ));
4483 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new())
4484 .with_values_field(value_field);
4485
4486 builder.keys().append_value("a");
4488 builder.values().append_value(1);
4489 builder.keys().append_value("b");
4490 builder.values().append_null();
4491 builder.append(true).unwrap();
4492 builder.keys().append_value("c");
4494 builder.values().append_null();
4495 builder.keys().append_value("d");
4496 builder.values().append_null();
4497 builder.append(true).unwrap();
4498 builder.append(true).unwrap();
4500 builder.keys().append_value("e");
4502 builder.values().append_value(5);
4503 builder.append(true).unwrap();
4504
4505 let map = Arc::new(builder.finish()) as ArrayRef;
4506 let d = map.data_type().clone();
4507
4508 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4509
4510 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4511
4512 let back = converter.convert_rows(&rows).unwrap();
4513 assert_eq!(back.len(), 1);
4514 back[0].to_data().validate_full().unwrap();
4515 assert_eq!(&back[0], &map);
4516 }
4517
4518 #[test]
4519 fn test_map_all_nulls() {
4520 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4521 builder.keys().append_value("m1"); builder.values().append_value(1); builder.append(false).unwrap();
4525 builder.keys().append_value("m2"); builder.values().append_value(2); builder.append(false).unwrap();
4528
4529 builder.append(false).unwrap(); builder.keys().append_value("m3"); builder.values().append_value(3); builder.append(false).unwrap();
4534
4535 let map = Arc::new(builder.finish()) as ArrayRef;
4536 let d = map.data_type().clone();
4537
4538 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4539
4540 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4541
4542 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4544
4545 let back = converter.convert_rows(&rows).unwrap();
4546 assert_eq!(back.len(), 1);
4547 back[0].to_data().validate_full().unwrap();
4548 assert_eq!(&back[0], &map);
4549 }
4550
4551 #[test]
4552 fn test_map_all_empty() {
4553 let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4554 builder.append(true).unwrap();
4556 builder.append(true).unwrap();
4557 builder.append(true).unwrap();
4558
4559 let map = Arc::new(builder.finish()) as ArrayRef;
4560 let d = map.data_type().clone();
4561
4562 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4563
4564 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4565
4566 rows.iter().for_each(|row| assert_eq!(row, rows.row(0)));
4568
4569 let back = converter.convert_rows(&rows).unwrap();
4570 assert_eq!(back.len(), 1);
4571 back[0].to_data().validate_full().unwrap();
4572 assert_eq!(&back[0], &map);
4573 }
4574
4575 #[test]
4576 fn test_map_empty_array() {
4577 let builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
4579 let map = Arc::new(builder.finish_cloned()) as ArrayRef;
4580 let d = map.data_type().clone();
4581
4582 let converter = RowConverter::new(vec![SortField::new(d.clone())]).unwrap();
4583
4584 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
4585
4586 let back = converter.convert_rows(&rows).unwrap();
4587 assert_eq!(back.len(), 1);
4588 back[0].to_data().validate_full().unwrap();
4589 assert_eq!(&back[0], &map);
4590 }
4591
4592 fn generate_primitive_array<K>(
4593 rng: &mut StdRng,
4594 len: usize,
4595 valid_percent: f64,
4596 ) -> PrimitiveArray<K>
4597 where
4598 K: ArrowPrimitiveType,
4599 StandardUniform: Distribution<K::Native>,
4600 {
4601 (0..len)
4602 .map(|_| rng.random_bool(valid_percent).then(|| rng.random()))
4603 .collect()
4604 }
4605
4606 fn generate_all_unique_primitive_array<K>(rng: &mut StdRng, len: usize) -> PrimitiveArray<K>
4607 where
4608 K: ArrowPrimitiveType,
4609 K::Native: Hash + Eq,
4610 StandardUniform: Distribution<K::Native>,
4611 {
4612 let possible_number_of_values = 2i32.saturating_pow(size_of::<K::Native>() as u32 * 8);
4613 assert!(
4614 len <= possible_number_of_values as usize,
4615 "len {len} is larger than the number of possible values {possible_number_of_values}"
4616 );
4617
4618 let mut seen = std::collections::HashSet::new();
4619 (0..len)
4620 .map(|_| {
4621 let mut value;
4622 loop {
4623 value = rng.random();
4624
4625 if seen.insert(value) {
4626 break;
4627 }
4628 }
4629
4630 Some(value)
4631 })
4632 .collect()
4633 }
4634
4635 fn generate_boolean_array(rng: &mut StdRng, len: usize, valid_percent: f64) -> BooleanArray {
4636 (0..len)
4637 .map(|_| rng.random_bool(valid_percent).then(|| rng.random_bool(0.5)))
4638 .collect()
4639 }
4640
4641 fn generate_strings<O: OffsetSizeTrait>(
4642 rng: &mut StdRng,
4643 len: usize,
4644 valid_percent: f64,
4645 ) -> GenericStringArray<O> {
4646 (0..len)
4647 .map(|_| {
4648 rng.random_bool(valid_percent).then(|| {
4649 let len = rng.random_range(0..100);
4650 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4651 String::from_utf8(bytes).unwrap()
4652 })
4653 })
4654 .collect()
4655 }
4656
4657 fn generate_string_view(rng: &mut StdRng, len: usize, valid_percent: f64) -> StringViewArray {
4658 (0..len)
4659 .map(|_| {
4660 rng.random_bool(valid_percent).then(|| {
4661 let len = rng.random_range(0..100);
4662 let bytes = (0..len).map(|_| rng.random_range(0..128)).collect();
4663 String::from_utf8(bytes).unwrap()
4664 })
4665 })
4666 .collect()
4667 }
4668
4669 fn generate_byte_view(rng: &mut StdRng, len: usize, valid_percent: f64) -> BinaryViewArray {
4670 (0..len)
4671 .map(|_| {
4672 rng.random_bool(valid_percent).then(|| {
4673 let len = rng.random_range(0..100);
4674 let bytes: Vec<_> = (0..len).map(|_| rng.random_range(0..128)).collect();
4675 bytes
4676 })
4677 })
4678 .collect()
4679 }
4680
4681 fn generate_fixed_stringview_column(len: usize) -> StringViewArray {
4682 let edge_cases = vec![
4683 Some("bar".to_string()),
4684 Some("bar\0".to_string()),
4685 Some("LongerThan12Bytes".to_string()),
4686 Some("LongerThan12Bytez".to_string()),
4687 Some("LongerThan12Bytes\0".to_string()),
4688 Some("LongerThan12Byt".to_string()),
4689 Some("backend one".to_string()),
4690 Some("backend two".to_string()),
4691 Some("a".repeat(257)),
4692 Some("a".repeat(300)),
4693 ];
4694
4695 let mut values = Vec::with_capacity(len);
4697 for i in 0..len {
4698 values.push(
4699 edge_cases
4700 .get(i % edge_cases.len())
4701 .cloned()
4702 .unwrap_or(None),
4703 );
4704 }
4705
4706 StringViewArray::from(values)
4707 }
4708
4709 fn generate_dictionary<K>(
4710 rng: &mut StdRng,
4711 values: ArrayRef,
4712 len: usize,
4713 valid_percent: f64,
4714 ) -> DictionaryArray<K>
4715 where
4716 K: ArrowDictionaryKeyType,
4717 K::Native: SampleUniform,
4718 {
4719 let min_key = K::Native::from_usize(0).unwrap();
4720 let max_key = K::Native::from_usize(values.len()).unwrap();
4721 let keys: PrimitiveArray<K> = (0..len)
4722 .map(|_| {
4723 rng.random_bool(valid_percent)
4724 .then(|| rng.random_range(min_key..max_key))
4725 })
4726 .collect();
4727
4728 let data_type =
4729 DataType::Dictionary(Box::new(K::DATA_TYPE), Box::new(values.data_type().clone()));
4730
4731 let data = keys
4732 .into_data()
4733 .into_builder()
4734 .data_type(data_type)
4735 .add_child_data(values.to_data())
4736 .build()
4737 .unwrap();
4738
4739 DictionaryArray::from(data)
4740 }
4741
4742 fn generate_fixed_size_binary(
4743 rng: &mut StdRng,
4744 len: usize,
4745 valid_percent: f64,
4746 ) -> FixedSizeBinaryArray {
4747 let width = rng.random_range(0..20);
4748 let mut builder = FixedSizeBinaryBuilder::new(width);
4749
4750 let mut b = vec![0; width as usize];
4751 for _ in 0..len {
4752 match rng.random_bool(valid_percent) {
4753 true => {
4754 b.iter_mut().for_each(|x| *x = rng.random());
4755 builder.append_value(&b).unwrap();
4756 }
4757 false => builder.append_null(),
4758 }
4759 }
4760
4761 builder.finish()
4762 }
4763
4764 fn generate_struct(rng: &mut StdRng, len: usize, valid_percent: f64) -> StructArray {
4765 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4766 let a = generate_primitive_array::<Int32Type>(rng, len, valid_percent);
4767 let b = generate_strings::<i32>(rng, len, valid_percent);
4768 let fields = Fields::from(vec![
4769 Field::new("a", DataType::Int32, true),
4770 Field::new("b", DataType::Utf8, true),
4771 ]);
4772 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
4773 StructArray::new(fields, values, Some(nulls))
4774 }
4775
4776 fn generate_list<F>(rng: &mut StdRng, len: usize, valid_percent: f64, values: F) -> ListArray
4777 where
4778 F: FnOnce(&mut StdRng, usize) -> ArrayRef,
4779 {
4780 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4781 let values_len = offsets.last().as_usize();
4782 let values = values(rng, values_len);
4783 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4784 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4785 ListArray::new(field, offsets, values, Some(nulls))
4786 }
4787
4788 fn generate_list_view<F>(
4789 rng: &mut StdRng,
4790 len: usize,
4791 valid_percent: f64,
4792 values: F,
4793 ) -> ListViewArray
4794 where
4795 F: FnOnce(&mut StdRng, usize) -> ArrayRef,
4796 {
4797 let sizes: Vec<i32> = (0..len).map(|_| rng.random_range(0..10)).collect();
4799 let values_len: usize = sizes.iter().map(|s| *s as usize).sum::<usize>().max(1);
4800 let values = values(rng, values_len);
4801
4802 let offsets: Vec<i32> = sizes
4804 .iter()
4805 .map(|&size| {
4806 if size == 0 {
4807 0
4808 } else {
4809 rng.random_range(0..=(values_len as i32 - size))
4810 }
4811 })
4812 .collect();
4813
4814 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4815 let field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
4816 ListViewArray::new(
4817 field,
4818 ScalarBuffer::from(offsets),
4819 ScalarBuffer::from(sizes),
4820 values,
4821 Some(nulls),
4822 )
4823 }
4824
4825 fn generate_map<KeysFn, ValuesFn>(
4826 rng: &mut StdRng,
4827 len: usize,
4828 valid_percent: f64,
4829 gen_keys: KeysFn,
4830 gen_values: ValuesFn,
4831 ) -> MapArray
4832 where
4833 KeysFn: FnOnce(&mut StdRng, usize) -> ArrayRef,
4834 ValuesFn: FnOnce(&mut StdRng, usize) -> ArrayRef,
4835 {
4836 let offsets = OffsetBuffer::<i32>::from_lengths((0..len).map(|_| rng.random_range(0..10)));
4837 let entries_len = offsets.last().as_usize();
4838 let keys = gen_keys(rng, entries_len);
4839 let values = gen_values(rng, entries_len);
4840 let nulls = NullBuffer::from_iter((0..len).map(|_| rng.random_bool(valid_percent)));
4841 let field = Arc::new(Field::new_map(
4842 "",
4843 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4844 Field::new(
4845 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4846 keys.data_type().clone(),
4847 false,
4848 ),
4849 Field::new(
4850 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4851 values.data_type().clone(),
4852 true,
4853 ),
4854 false,
4855 true,
4856 ));
4857 let DataType::Map(struct_field, _) = field.data_type() else {
4858 unreachable!();
4859 };
4860
4861 let DataType::Struct(fields) = struct_field.data_type() else {
4862 unreachable!();
4863 };
4864
4865 let entries = StructArray::new(fields.clone(), vec![keys, values], None);
4866
4867 let map_array = MapArray::new(struct_field.clone(), offsets, entries, Some(nulls), false);
4868
4869 assert_valid_map(&map_array);
4870
4871 map_array
4872 }
4873
4874 fn assert_valid_map(array: &MapArray) {
4884 let keys_arrow_row_converter =
4885 RowConverter::new(vec![SortField::new(array.key_type().clone())]).unwrap();
4886
4887 array.iter().enumerate().filter_map(|(index, entry)| entry.map(|entry| (index, Arc::clone(entry.column(0))))).for_each(|(entry_index, keys)| {
4888 let keys_as_rows = keys_arrow_row_converter.convert_columns(&[Arc::clone(&keys)]).expect("should be able to convert keys");
4889
4890 for i in 0..keys_as_rows.num_rows() {
4891 for j in (i + 1)..keys_as_rows.num_rows() {
4892 if keys_as_rows.row(i) == keys_as_rows.row(j) {
4893 let key_i = keys.slice(i, 1);
4894 let key_j = keys.slice(j, 1);
4895
4896 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:?}");
4897 }
4898 }
4899 }
4900 })
4901 }
4902
4903 fn generate_nulls(rng: &mut StdRng, len: usize) -> Option<NullBuffer> {
4904 Some(NullBuffer::from_iter(
4905 (0..len).map(|_| rng.random_bool(0.8)),
4906 ))
4907 }
4908
4909 fn change_underlying_null_values_for_primitive<T: ArrowPrimitiveType>(
4910 array: &PrimitiveArray<T>,
4911 ) -> PrimitiveArray<T> {
4912 let (dt, values, nulls) = array.clone().into_parts();
4913
4914 let new_values = ScalarBuffer::<T::Native>::from_iter(
4915 values
4916 .iter()
4917 .zip(nulls.as_ref().unwrap().iter())
4918 .map(|(val, is_valid)| {
4919 if is_valid {
4920 *val
4921 } else {
4922 val.add_wrapping(T::Native::usize_as(1))
4923 }
4924 }),
4925 );
4926
4927 PrimitiveArray::new(new_values, nulls).with_data_type(dt)
4928 }
4929
4930 fn change_underline_null_values_for_byte_array<T: ByteArrayType>(
4931 array: &GenericByteArray<T>,
4932 ) -> GenericByteArray<T> {
4933 let (offsets, values, nulls) = array.clone().into_parts();
4934
4935 let new_offsets = OffsetBuffer::<T::Offset>::from_lengths(
4936 offsets
4937 .lengths()
4938 .zip(nulls.as_ref().unwrap().iter())
4939 .map(|(len, is_valid)| if is_valid { len } else { len + 1 }),
4940 );
4941
4942 let mut new_bytes = Vec::<u8>::with_capacity(new_offsets[new_offsets.len() - 1].as_usize());
4943
4944 offsets
4945 .windows(2)
4946 .zip(nulls.as_ref().unwrap().iter())
4947 .for_each(|(start_and_end, is_valid)| {
4948 let start = start_and_end[0].as_usize();
4949 let end = start_and_end[1].as_usize();
4950 new_bytes.extend_from_slice(&values.as_slice()[start..end]);
4951
4952 if !is_valid {
4954 new_bytes.push(b'c');
4955 }
4956 });
4957
4958 GenericByteArray::<T>::new(new_offsets, Buffer::from_vec(new_bytes), nulls)
4959 }
4960
4961 fn change_underline_null_values_for_list_array<O: OffsetSizeTrait>(
4962 array: &GenericListArray<O>,
4963 ) -> GenericListArray<O> {
4964 let (field, offsets, values, nulls) = array.clone().into_parts();
4965
4966 let (new_values, new_offsets) = {
4967 let concat_values = offsets
4968 .windows(2)
4969 .zip(nulls.as_ref().unwrap().iter())
4970 .map(|(start_and_end, is_valid)| {
4971 let start = start_and_end[0].as_usize();
4972 let end = start_and_end[1].as_usize();
4973 if is_valid {
4974 return (start, end - start);
4975 }
4976
4977 if end == values.len() {
4979 (start, (end - start).saturating_sub(1))
4980 } else {
4981 (start, end - start + 1)
4982 }
4983 })
4984 .map(|(start, length)| values.slice(start, length))
4985 .collect::<Vec<_>>();
4986
4987 let new_offsets =
4988 OffsetBuffer::<O>::from_lengths(concat_values.iter().map(|s| s.len()));
4989
4990 let new_values = {
4991 let values = concat_values.iter().map(|a| a.as_ref()).collect::<Vec<_>>();
4992 arrow_select::concat::concat(&values).expect("should be able to concat")
4993 };
4994
4995 (new_values, new_offsets)
4996 };
4997
4998 GenericListArray::<O>::new(field, new_offsets, new_values, nulls)
4999 }
5000
5001 fn change_underline_null_values_for_map_array(array: &MapArray) -> MapArray {
5002 let (field, offsets, entries, nulls, ordered) = array.clone().into_parts();
5003 assert!(
5004 !ordered,
5005 "can't replace underlying null values for ordered map array as this can violate the ordering"
5006 );
5007
5008 let (new_entries, new_offsets) = {
5009 let concat_values = offsets
5010 .windows(2)
5011 .zip(nulls.as_ref().unwrap().iter())
5012 .map(|(start_and_end, is_valid)| {
5013 let start = start_and_end[0].as_usize();
5014 let end = start_and_end[1].as_usize();
5015 if is_valid {
5016 return (start, end - start);
5017 }
5018
5019 if end == entries.len() {
5021 (start, (end - start).saturating_sub(1))
5022 } else {
5023 (start, end - start + 1)
5025 }
5026 })
5027 .map(|(start, length)| entries.slice(start, length))
5028 .collect::<Vec<_>>();
5029
5030 let new_offsets = OffsetBuffer::from_lengths(concat_values.iter().map(|s| s.len()));
5031
5032 let new_values = {
5033 let values = concat_values
5034 .iter()
5035 .map(|a| a as &dyn Array)
5036 .collect::<Vec<_>>();
5037 arrow_select::concat::concat(&values).expect("should be able to concat")
5038 };
5039
5040 (new_values.as_struct().clone(), new_offsets)
5041 };
5042
5043 let new_map = MapArray::new(field, new_offsets, new_entries, nulls, ordered);
5044
5045 assert_valid_map(&new_map);
5046
5047 new_map
5048 }
5049
5050 fn change_underline_null_values(array: &ArrayRef) -> ArrayRef {
5051 if array.null_count() == 0 {
5052 return Arc::clone(array);
5053 }
5054
5055 downcast_primitive_array!(
5056 array => {
5057 let output = change_underlying_null_values_for_primitive(array);
5058
5059 Arc::new(output)
5060 }
5061
5062 DataType::Utf8 => {
5063 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i32>()))
5064 }
5065 DataType::LargeUtf8 => {
5066 Arc::new(change_underline_null_values_for_byte_array(array.as_string::<i64>()))
5067 }
5068 DataType::Binary => {
5069 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i32>()))
5070 }
5071 DataType::LargeBinary => {
5072 Arc::new(change_underline_null_values_for_byte_array(array.as_binary::<i64>()))
5073 }
5074 DataType::List(_) => {
5075 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i32>()))
5076 }
5077 DataType::LargeList(_) => {
5078 Arc::new(change_underline_null_values_for_list_array(array.as_list::<i64>()))
5079 }
5080 DataType::Map(_, _) => {
5081 Arc::new(change_underline_null_values_for_map_array(array.as_map()))
5082 }
5083 _ => {
5084 Arc::clone(array)
5085 }
5086 )
5087 }
5088
5089 fn generate_column(rng: &mut StdRng, len: usize) -> ArrayRef {
5090 match rng.random_range(0..24) {
5091 0 => Arc::new(generate_primitive_array::<Int32Type>(rng, len, 0.8)),
5092 1 => Arc::new(generate_primitive_array::<UInt32Type>(rng, len, 0.8)),
5093 2 => Arc::new(generate_primitive_array::<Int64Type>(rng, len, 0.8)),
5094 3 => Arc::new(generate_primitive_array::<UInt64Type>(rng, len, 0.8)),
5095 4 => Arc::new(generate_primitive_array::<Float32Type>(rng, len, 0.8)),
5096 5 => Arc::new(generate_primitive_array::<Float64Type>(rng, len, 0.8)),
5097 6 => Arc::new(generate_strings::<i32>(rng, len, 0.8)),
5098 7 => {
5099 let dict_values_len = rng.random_range(1..len);
5100 let strings = Arc::new(generate_strings::<i32>(rng, dict_values_len, 1.0));
5102 Arc::new(generate_dictionary::<Int64Type>(rng, strings, len, 0.8))
5103 }
5104 8 => {
5105 let dict_values_len = rng.random_range(1..len);
5106 let values = Arc::new(generate_primitive_array::<Int64Type>(
5108 rng,
5109 dict_values_len,
5110 1.0,
5111 ));
5112 Arc::new(generate_dictionary::<Int64Type>(rng, values, len, 0.8))
5113 }
5114 9 => Arc::new(generate_fixed_size_binary(rng, len, 0.8)),
5115 10 => Arc::new(generate_struct(rng, len, 0.8)),
5116 11 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5117 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5118 })),
5119 12 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5120 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5121 })),
5122 13 => Arc::new(generate_list(rng, len, 0.8, |rng, values_len| {
5123 Arc::new(generate_struct(rng, values_len, 0.8))
5124 })),
5125 14 => Arc::new(generate_string_view(rng, len, 0.8)),
5126 15 => Arc::new(generate_byte_view(rng, len, 0.8)),
5127 16 => Arc::new(generate_fixed_stringview_column(len)),
5128 17 => Arc::new(
5129 generate_list(rng, len + 1000, 0.8, |rng, values_len| {
5130 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5131 })
5132 .slice(500, len),
5133 ),
5134 18 => Arc::new(generate_boolean_array(rng, len, 0.8)),
5135 19 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5136 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5137 })),
5138 20 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5139 Arc::new(generate_strings::<i32>(rng, values_len, 0.8))
5140 })),
5141 21 => Arc::new(generate_list_view(rng, len, 0.8, |rng, values_len| {
5142 Arc::new(generate_struct(rng, values_len, 0.8))
5143 })),
5144 22 => Arc::new(
5145 generate_list_view(rng, len + 1000, 0.8, |rng, values_len| {
5146 Arc::new(generate_primitive_array::<Int64Type>(rng, values_len, 0.8))
5147 })
5148 .slice(500, len),
5149 ),
5150 23 => Arc::new(generate_map(
5151 rng,
5152 len,
5153 0.9,
5154 |rng, keys_len| {
5156 Arc::new(generate_all_unique_primitive_array::<Int64Type>(
5157 rng, keys_len,
5158 ))
5159 },
5160 |rng, values_len| Arc::new(generate_strings::<i32>(rng, values_len, 0.7)),
5161 )),
5162 _ => unreachable!(),
5163 }
5164 }
5165
5166 fn print_row(cols: &[SortColumn], row: usize) -> String {
5167 let t: Vec<_> = cols
5168 .iter()
5169 .map(|x| match x.values.is_valid(row) {
5170 true => {
5171 let opts = FormatOptions::default().with_null("NULL");
5172 let formatter = ArrayFormatter::try_new(x.values.as_ref(), &opts).unwrap();
5173 formatter.value(row).to_string()
5174 }
5175 false => "NULL".to_string(),
5176 })
5177 .collect();
5178 t.join(",")
5179 }
5180
5181 fn print_col_types(cols: &[SortColumn]) -> String {
5182 let t: Vec<_> = cols
5183 .iter()
5184 .map(|x| x.values.data_type().to_string())
5185 .collect();
5186 t.join(",")
5187 }
5188
5189 #[derive(Debug, PartialEq)]
5190 enum Nulls {
5191 AsIs,
5193
5194 Different,
5196
5197 None,
5199 }
5200
5201 #[test]
5202 #[cfg_attr(miri, ignore)] fn fuzz_test() {
5204 let mut rng = StdRng::seed_from_u64(42);
5205 for _ in 0..100 {
5206 for null_behavior in [Nulls::AsIs, Nulls::Different, Nulls::None] {
5207 let num_columns = rng.random_range(1..5);
5208 let len = rng.random_range(5..100);
5209 let mut arrays: Vec<_> = (0..num_columns)
5210 .map(|_| generate_column(&mut rng, len))
5211 .collect();
5212
5213 match null_behavior {
5214 Nulls::AsIs => {
5215 }
5217 Nulls::Different => {
5218 arrays = arrays
5220 .into_iter()
5221 .map(|a| replace_array_nulls(a, generate_nulls(&mut rng, len)))
5222 .collect()
5223 }
5224 Nulls::None => {
5225 arrays = arrays
5227 .into_iter()
5228 .map(|a| replace_array_nulls(a, None))
5229 .collect()
5230 }
5231 }
5232
5233 let options: Vec<_> = (0..num_columns)
5234 .map(|_| SortOptions {
5235 descending: rng.random_bool(0.5),
5236 nulls_first: rng.random_bool(0.5),
5237 })
5238 .collect();
5239
5240 let sort_columns: Vec<_> = options
5241 .iter()
5242 .zip(&arrays)
5243 .map(|(o, c)| SortColumn {
5244 values: Arc::clone(c),
5245 options: Some(*o),
5246 })
5247 .collect();
5248
5249 let comparator = LexicographicalComparator::try_new(&sort_columns).unwrap();
5250
5251 let columns: Vec<SortField> = options
5252 .into_iter()
5253 .zip(&arrays)
5254 .map(|(o, a)| SortField::new_with_options(a.data_type().clone(), o))
5255 .collect();
5256
5257 let converter = RowConverter::new(columns).unwrap();
5258 let rows = converter.convert_columns(&arrays).unwrap();
5259
5260 if !matches!(null_behavior, Nulls::None) {
5263 assert_same_rows_when_changing_input_underlying_null_values(
5264 &arrays, &converter, &rows,
5265 );
5266 }
5267
5268 for i in 0..len {
5269 for j in 0..len {
5270 let row_i = rows.row(i);
5271 let row_j = rows.row(j);
5272 let row_cmp = row_i.cmp(&row_j);
5273 let lex_cmp = comparator.compare(i, j);
5274 assert_eq!(
5275 row_cmp,
5276 lex_cmp,
5277 "({:?} vs {:?}) vs ({:?} vs {:?}) for types {}",
5278 print_row(&sort_columns, i),
5279 print_row(&sort_columns, j),
5280 row_i,
5281 row_j,
5282 print_col_types(&sort_columns)
5283 );
5284 }
5285 }
5286
5287 {
5289 let mut rows_iter = rows.iter();
5290 let mut rows_lengths_iter = rows.lengths();
5291 for (index, row) in rows_iter.by_ref().enumerate() {
5292 let len = rows_lengths_iter
5293 .next()
5294 .expect("Reached end of length iterator while still have rows");
5295 assert_eq!(
5296 row.data.len(),
5297 len,
5298 "Row length mismatch: {} vs {}",
5299 row.data.len(),
5300 len
5301 );
5302 assert_eq!(
5303 len,
5304 rows.row_len(index),
5305 "Row length mismatch at index {}: {} vs {}",
5306 index,
5307 len,
5308 rows.row_len(index)
5309 );
5310 }
5311
5312 assert_eq!(
5313 rows_lengths_iter.next(),
5314 None,
5315 "Length iterator did not reach end"
5316 );
5317 }
5318
5319 let back = converter.convert_rows(&rows).unwrap();
5322 for (actual, expected) in back.iter().zip(&arrays) {
5323 actual.to_data().validate_full().unwrap();
5324 dictionary_eq(actual, expected)
5325 }
5326
5327 let rows = rows.try_into_binary().expect("reasonable size");
5330 let parser = converter.parser();
5331 let back = converter
5332 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5333 .unwrap();
5334 for (actual, expected) in back.iter().zip(&arrays) {
5335 actual.to_data().validate_full().unwrap();
5336 dictionary_eq(actual, expected)
5337 }
5338
5339 let rows = converter.from_binary(rows);
5340 let back = converter.convert_rows(&rows).unwrap();
5341 for (actual, expected) in back.iter().zip(&arrays) {
5342 actual.to_data().validate_full().unwrap();
5343 dictionary_eq(actual, expected)
5344 }
5345 }
5346 }
5347 }
5348
5349 fn replace_array_nulls(array: ArrayRef, new_nulls: Option<NullBuffer>) -> ArrayRef {
5350 make_array(
5351 array
5352 .into_data()
5353 .into_builder()
5354 .nulls(new_nulls)
5356 .build()
5357 .unwrap(),
5358 )
5359 }
5360
5361 fn assert_same_rows_when_changing_input_underlying_null_values(
5362 arrays: &[ArrayRef],
5363 converter: &RowConverter,
5364 rows: &Rows,
5365 ) {
5366 let arrays_with_different_data_behind_nulls = arrays
5367 .iter()
5368 .map(|arr| change_underline_null_values(arr))
5369 .collect::<Vec<_>>();
5370
5371 if arrays
5373 .iter()
5374 .zip(arrays_with_different_data_behind_nulls.iter())
5375 .all(|(a, b)| Arc::ptr_eq(a, b))
5376 {
5377 return;
5378 }
5379
5380 let rows_with_different_nulls = converter
5381 .convert_columns(&arrays_with_different_data_behind_nulls)
5382 .unwrap();
5383
5384 assert_eq!(
5385 rows.iter().collect::<Vec<_>>(),
5386 rows_with_different_nulls.iter().collect::<Vec<_>>(),
5387 "Different underlying nulls should not output different rows"
5388 )
5389 }
5390
5391 #[test]
5392 fn test_clear() {
5393 let converter = RowConverter::new(vec![SortField::new(DataType::Int32)]).unwrap();
5394 let mut rows = converter.empty_rows(3, 128);
5395
5396 let first = Int32Array::from(vec![None, Some(2), Some(4)]);
5397 let second = Int32Array::from(vec![Some(2), None, Some(4)]);
5398 let arrays = [Arc::new(first) as ArrayRef, Arc::new(second) as ArrayRef];
5399
5400 for array in &arrays {
5401 rows.clear();
5402 converter
5403 .append(&mut rows, std::slice::from_ref(array))
5404 .unwrap();
5405 let back = converter.convert_rows(&rows).unwrap();
5406 assert_eq!(&back[0], array);
5407 }
5408
5409 let mut rows_expected = converter.empty_rows(3, 128);
5410 converter.append(&mut rows_expected, &arrays[1..]).unwrap();
5411
5412 for (i, (actual, expected)) in rows.iter().zip(rows_expected.iter()).enumerate() {
5413 assert_eq!(
5414 actual, expected,
5415 "For row {i}: expected {expected:?}, actual: {actual:?}",
5416 );
5417 }
5418 }
5419
5420 #[test]
5421 fn test_append_codec_dictionary_binary() {
5422 use DataType::*;
5423 let converter = RowConverter::new(vec![SortField::new(Dictionary(
5425 Box::new(Int32),
5426 Box::new(Binary),
5427 ))])
5428 .unwrap();
5429 let mut rows = converter.empty_rows(4, 128);
5430
5431 let keys = Int32Array::from_iter_values([0, 1, 2, 3]);
5432 let values = BinaryArray::from(vec![
5433 Some(b"a".as_slice()),
5434 Some(b"b"),
5435 Some(b"c"),
5436 Some(b"d"),
5437 ]);
5438 let dict_array = DictionaryArray::new(keys, Arc::new(values));
5439
5440 rows.clear();
5441 let array = Arc::new(dict_array) as ArrayRef;
5442 converter
5443 .append(&mut rows, std::slice::from_ref(&array))
5444 .unwrap();
5445 let back = converter.convert_rows(&rows).unwrap();
5446
5447 dictionary_eq(&back[0], &array);
5448 }
5449
5450 #[test]
5451 fn test_list_prefix() {
5452 let mut a = ListBuilder::new(Int8Builder::new());
5453 a.append_value([None]);
5454 a.append_value([None, None]);
5455 let a = a.finish();
5456
5457 let converter = RowConverter::new(vec![SortField::new(a.data_type().clone())]).unwrap();
5458 let rows = converter.convert_columns(&[Arc::new(a) as _]).unwrap();
5459 assert_eq!(rows.row(0).cmp(&rows.row(1)), Ordering::Less);
5460 }
5461
5462 #[test]
5463 fn test_utf8_validation_doesnt_affect_values_buffer_size() {
5464 fn assert_values_buffer_lens(col: ArrayRef) -> usize {
5465 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8View)]).unwrap();
5467
5468 let rows = converter.convert_columns(&[col]).unwrap();
5470 let converted = converter.convert_rows(&rows).unwrap();
5471 let unchecked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5472
5473 let rows = rows.try_into_binary().expect("reasonable size");
5475 let parser = converter.parser();
5476 let converted = converter
5477 .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
5478 .unwrap();
5479 let checked_values_len = converted[0].as_string_view().data_buffers()[0].len();
5480 assert_eq!(unchecked_values_len, checked_values_len);
5482 checked_values_len
5483 }
5484
5485 let col = Arc::new(StringViewArray::from_iter([
5487 Some("hello"), None, Some("short"), Some("tiny"), ])) as ArrayRef;
5492
5493 let values_len = assert_values_buffer_lens(col);
5494 assert_eq!(values_len, 0);
5496
5497 let col = Arc::new(StringViewArray::from_iter([
5499 Some("1234567890123"), Some("12345678901234"), ])) as ArrayRef;
5502
5503 let values_len = assert_values_buffer_lens(col);
5504 assert_eq!(values_len, 13 + 14);
5505
5506 let col = Arc::new(StringViewArray::from_iter([
5508 Some("tiny"), Some("thisisexact13"), None,
5511 Some("short"), ])) as ArrayRef;
5513
5514 let values_len = assert_values_buffer_lens(col);
5515 assert_eq!(values_len, 13);
5517 }
5518
5519 #[test]
5520 fn test_sparse_union() {
5521 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5523 let str_array = StringArray::from(vec![None, Some("b"), None, Some("d"), None]);
5524
5525 let type_ids = vec![0, 1, 0, 1, 0].into();
5527
5528 let union_fields = [
5529 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5530 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5531 ]
5532 .into_iter()
5533 .collect();
5534
5535 let union_array = UnionArray::try_new(
5536 union_fields,
5537 type_ids,
5538 None,
5539 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5540 )
5541 .unwrap();
5542
5543 let union_type = union_array.data_type().clone();
5544 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5545
5546 let rows = converter
5547 .convert_columns(&[Arc::new(union_array.clone())])
5548 .unwrap();
5549
5550 let back = converter.convert_rows(&rows).unwrap();
5552 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5553
5554 assert_eq!(union_array.len(), back_union.len());
5555 for i in 0..union_array.len() {
5556 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5557 }
5558 }
5559
5560 #[test]
5561 fn test_sparse_union_with_nulls() {
5562 let int_array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
5564 let str_array = StringArray::from(vec![None::<&str>; 5]);
5565
5566 let type_ids = vec![0, 1, 0, 1, 0].into();
5568
5569 let union_fields = [
5570 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5571 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5572 ]
5573 .into_iter()
5574 .collect();
5575
5576 let union_array = UnionArray::try_new(
5577 union_fields,
5578 type_ids,
5579 None,
5580 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5581 )
5582 .unwrap();
5583
5584 let union_type = union_array.data_type().clone();
5585 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5586
5587 let rows = converter
5588 .convert_columns(&[Arc::new(union_array.clone())])
5589 .unwrap();
5590
5591 let back = converter.convert_rows(&rows).unwrap();
5593 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5594
5595 assert_eq!(union_array.len(), back_union.len());
5596 for i in 0..union_array.len() {
5597 let expected_null = union_array.is_null(i);
5598 let actual_null = back_union.is_null(i);
5599 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5600 if !expected_null {
5601 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5602 }
5603 }
5604 }
5605
5606 #[test]
5607 fn test_dense_union() {
5608 let int_array = Int32Array::from(vec![1, 3, 5]);
5610 let str_array = StringArray::from(vec!["a", "b"]);
5611
5612 let type_ids = vec![0, 1, 0, 1, 0].into();
5613
5614 let offsets = vec![0, 0, 1, 1, 2].into();
5616
5617 let union_fields = [
5618 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5619 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5620 ]
5621 .into_iter()
5622 .collect();
5623
5624 let union_array = UnionArray::try_new(
5625 union_fields,
5626 type_ids,
5627 Some(offsets), vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5629 )
5630 .unwrap();
5631
5632 let union_type = union_array.data_type().clone();
5633 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5634
5635 let rows = converter
5636 .convert_columns(&[Arc::new(union_array.clone())])
5637 .unwrap();
5638
5639 let back = converter.convert_rows(&rows).unwrap();
5641 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5642
5643 assert_eq!(union_array.len(), back_union.len());
5644 for i in 0..union_array.len() {
5645 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5646 }
5647 }
5648
5649 #[test]
5650 fn test_dense_union_with_nulls() {
5651 let int_array = Int32Array::from(vec![Some(1), None, Some(5)]);
5653 let str_array = StringArray::from(vec![Some("a"), None]);
5654
5655 let type_ids = vec![0, 1, 0, 1, 0].into();
5657 let offsets = vec![0, 0, 1, 1, 2].into();
5658
5659 let union_fields = [
5660 (0, Arc::new(Field::new("int", DataType::Int32, true))),
5661 (1, Arc::new(Field::new("str", DataType::Utf8, true))),
5662 ]
5663 .into_iter()
5664 .collect();
5665
5666 let union_array = UnionArray::try_new(
5667 union_fields,
5668 type_ids,
5669 Some(offsets),
5670 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5671 )
5672 .unwrap();
5673
5674 let union_type = union_array.data_type().clone();
5675 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5676
5677 let rows = converter
5678 .convert_columns(&[Arc::new(union_array.clone())])
5679 .unwrap();
5680
5681 let back = converter.convert_rows(&rows).unwrap();
5683 let back_union = back[0].as_any().downcast_ref::<UnionArray>().unwrap();
5684
5685 assert_eq!(union_array.len(), back_union.len());
5686 for i in 0..union_array.len() {
5687 let expected_null = union_array.is_null(i);
5688 let actual_null = back_union.is_null(i);
5689 assert_eq!(expected_null, actual_null, "Null mismatch at index {i}");
5690 if !expected_null {
5691 assert_eq!(union_array.type_id(i), back_union.type_id(i));
5692 }
5693 }
5694 }
5695
5696 #[test]
5697 fn test_union_ordering() {
5698 let int_array = Int32Array::from(vec![100, 5, 20]);
5699 let str_array = StringArray::from(vec!["z", "a"]);
5700
5701 let type_ids = vec![0, 1, 0, 1, 0].into();
5703 let offsets = vec![0, 0, 1, 1, 2].into();
5704
5705 let union_fields = [
5706 (0, Arc::new(Field::new("int", DataType::Int32, false))),
5707 (1, Arc::new(Field::new("str", DataType::Utf8, false))),
5708 ]
5709 .into_iter()
5710 .collect();
5711
5712 let union_array = UnionArray::try_new(
5713 union_fields,
5714 type_ids,
5715 Some(offsets),
5716 vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)],
5717 )
5718 .unwrap();
5719
5720 let union_type = union_array.data_type().clone();
5721 let converter = RowConverter::new(vec![SortField::new(union_type)]).unwrap();
5722
5723 let rows = converter.convert_columns(&[Arc::new(union_array)]).unwrap();
5724
5725 assert!(rows.row(2) < rows.row(1));
5737
5738 assert!(rows.row(0) < rows.row(3));
5740
5741 assert!(rows.row(2) < rows.row(4));
5744 assert!(rows.row(4) < rows.row(0));
5746
5747 assert!(rows.row(3) < rows.row(1));
5750 }
5751
5752 #[test]
5753 fn test_row_converter_roundtrip_with_many_union_columns() {
5754 let fields1 = UnionFields::try_new(
5756 vec![0, 1],
5757 vec![
5758 Field::new("int", DataType::Int32, true),
5759 Field::new("string", DataType::Utf8, true),
5760 ],
5761 )
5762 .unwrap();
5763
5764 let int_array1 = Int32Array::from(vec![Some(67), None]);
5765 let string_array1 = StringArray::from(vec![None::<&str>, Some("hello")]);
5766 let type_ids1 = vec![0i8, 1].into();
5767
5768 let union_array1 = UnionArray::try_new(
5769 fields1.clone(),
5770 type_ids1,
5771 None,
5772 vec![
5773 Arc::new(int_array1) as ArrayRef,
5774 Arc::new(string_array1) as ArrayRef,
5775 ],
5776 )
5777 .unwrap();
5778
5779 let fields2 = UnionFields::try_new(
5781 vec![0, 1],
5782 vec![
5783 Field::new("int", DataType::Int32, true),
5784 Field::new("string", DataType::Utf8, true),
5785 ],
5786 )
5787 .unwrap();
5788
5789 let int_array2 = Int32Array::from(vec![Some(100), None]);
5790 let string_array2 = StringArray::from(vec![None::<&str>, Some("world")]);
5791 let type_ids2 = vec![0i8, 1].into();
5792
5793 let union_array2 = UnionArray::try_new(
5794 fields2.clone(),
5795 type_ids2,
5796 None,
5797 vec![
5798 Arc::new(int_array2) as ArrayRef,
5799 Arc::new(string_array2) as ArrayRef,
5800 ],
5801 )
5802 .unwrap();
5803
5804 let field1 = Field::new("col1", DataType::Union(fields1, UnionMode::Sparse), true);
5806 let field2 = Field::new("col2", DataType::Union(fields2, UnionMode::Sparse), true);
5807
5808 let sort_field1 = SortField::new(field1.data_type().clone());
5809 let sort_field2 = SortField::new(field2.data_type().clone());
5810
5811 let converter = RowConverter::new(vec![sort_field1, sort_field2]).unwrap();
5812
5813 let rows = converter
5814 .convert_columns(&[
5815 Arc::new(union_array1.clone()) as ArrayRef,
5816 Arc::new(union_array2.clone()) as ArrayRef,
5817 ])
5818 .unwrap();
5819
5820 let out = converter.convert_rows(&rows).unwrap();
5822
5823 let [col1, col2] = out.as_slice() else {
5824 panic!("expected 2 columns")
5825 };
5826
5827 let col1 = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5828 let col2 = col2.as_any().downcast_ref::<UnionArray>().unwrap();
5829
5830 for (expected, got) in [union_array1, union_array2].iter().zip([col1, col2]) {
5831 assert_eq!(expected.len(), got.len());
5832 assert_eq!(expected.type_ids(), got.type_ids());
5833
5834 for i in 0..expected.len() {
5835 assert_eq!(expected.value(i).as_ref(), got.value(i).as_ref());
5836 }
5837 }
5838 }
5839
5840 #[test]
5841 fn test_row_converter_roundtrip_with_one_union_column() {
5842 let fields = UnionFields::try_new(
5843 vec![0, 1],
5844 vec![
5845 Field::new("int", DataType::Int32, true),
5846 Field::new("string", DataType::Utf8, true),
5847 ],
5848 )
5849 .unwrap();
5850
5851 let int_array = Int32Array::from(vec![Some(67), None]);
5852 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5853 let type_ids = vec![0i8, 1].into();
5854
5855 let union_array = UnionArray::try_new(
5856 fields.clone(),
5857 type_ids,
5858 None,
5859 vec![
5860 Arc::new(int_array) as ArrayRef,
5861 Arc::new(string_array) as ArrayRef,
5862 ],
5863 )
5864 .unwrap();
5865
5866 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5867 let sort_field = SortField::new(field.data_type().clone());
5868 let converter = RowConverter::new(vec![sort_field]).unwrap();
5869
5870 let rows = converter
5871 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5872 .unwrap();
5873
5874 let out = converter.convert_rows(&rows).unwrap();
5876
5877 let [col1] = out.as_slice() else {
5878 panic!("expected 1 column")
5879 };
5880
5881 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5882 assert_eq!(col.len(), union_array.len());
5883 assert_eq!(col.type_ids(), union_array.type_ids());
5884
5885 for i in 0..col.len() {
5886 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5887 }
5888 }
5889
5890 #[test]
5891 fn test_row_converter_roundtrip_with_non_default_union_type_ids() {
5892 let fields = UnionFields::try_new(
5894 vec![70, 85],
5895 vec![
5896 Field::new("int", DataType::Int32, true),
5897 Field::new("string", DataType::Utf8, true),
5898 ],
5899 )
5900 .unwrap();
5901
5902 let int_array = Int32Array::from(vec![Some(67), None]);
5903 let string_array = StringArray::from(vec![None::<&str>, Some("hello")]);
5904 let type_ids = vec![70i8, 85].into();
5905
5906 let union_array = UnionArray::try_new(
5907 fields.clone(),
5908 type_ids,
5909 None,
5910 vec![
5911 Arc::new(int_array) as ArrayRef,
5912 Arc::new(string_array) as ArrayRef,
5913 ],
5914 )
5915 .unwrap();
5916
5917 let field = Field::new("col", DataType::Union(fields, UnionMode::Sparse), true);
5918 let sort_field = SortField::new(field.data_type().clone());
5919 let converter = RowConverter::new(vec![sort_field]).unwrap();
5920
5921 let rows = converter
5922 .convert_columns(&[Arc::new(union_array.clone()) as ArrayRef])
5923 .unwrap();
5924
5925 let out = converter.convert_rows(&rows).unwrap();
5927
5928 let [col1] = out.as_slice() else {
5929 panic!("expected 1 column")
5930 };
5931
5932 let col = col1.as_any().downcast_ref::<UnionArray>().unwrap();
5933 assert_eq!(col.len(), union_array.len());
5934 assert_eq!(col.type_ids(), union_array.type_ids());
5935
5936 for i in 0..col.len() {
5937 assert_eq!(col.value(i).as_ref(), union_array.value(i).as_ref());
5938 }
5939 }
5940
5941 #[test]
5942 fn rows_size_should_count_for_capacity() {
5943 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
5944
5945 let empty_rows_size_with_preallocate_rows_and_data = {
5946 let rows = row_converter.empty_rows(1000, 1000);
5947
5948 rows.size()
5949 };
5950 let empty_rows_size_with_preallocate_rows = {
5951 let rows = row_converter.empty_rows(1000, 0);
5952
5953 rows.size()
5954 };
5955 let empty_rows_size_with_preallocate_data = {
5956 let rows = row_converter.empty_rows(0, 1000);
5957
5958 rows.size()
5959 };
5960 let empty_rows_size_without_preallocate = {
5961 let rows = row_converter.empty_rows(0, 0);
5962
5963 rows.size()
5964 };
5965
5966 assert!(
5967 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_rows,
5968 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_rows}"
5969 );
5970 assert!(
5971 empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_data,
5972 "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_data}"
5973 );
5974 assert!(
5975 empty_rows_size_with_preallocate_rows > empty_rows_size_without_preallocate,
5976 "{empty_rows_size_with_preallocate_rows} should be larger than {empty_rows_size_without_preallocate}"
5977 );
5978 assert!(
5979 empty_rows_size_with_preallocate_data > empty_rows_size_without_preallocate,
5980 "{empty_rows_size_with_preallocate_data} should be larger than {empty_rows_size_without_preallocate}"
5981 );
5982 }
5983
5984 #[test]
5985 fn test_struct_no_child_fields() {
5986 fn run_test(array: ArrayRef) {
5987 let sort_fields = vec![SortField::new(array.data_type().clone())];
5988 let converter = RowConverter::new(sort_fields).unwrap();
5989 let r = converter.convert_columns(&[Arc::clone(&array)]).unwrap();
5990
5991 let back = converter.convert_rows(&r).unwrap();
5992 assert_eq!(back.len(), 1);
5993 assert_eq!(&back[0], &array);
5994 }
5995
5996 let s = Arc::new(StructArray::new_empty_fields(5, None)) as ArrayRef;
5997 run_test(s);
5998
5999 let s = Arc::new(StructArray::new_empty_fields(
6000 5,
6001 Some(vec![true, false, true, false, false].into()),
6002 )) as ArrayRef;
6003 run_test(s);
6004 }
6005
6006 #[test]
6007 fn reserve_should_increase_capacity_to_the_requested_size() {
6008 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6009 let mut empty_rows = row_converter.empty_rows(0, 0);
6010 empty_rows.reserve(50, 50);
6011 let before_size = empty_rows.size();
6012 empty_rows.reserve(50, 50);
6013 assert_eq!(
6014 empty_rows.size(),
6015 before_size,
6016 "Size should not change when reserving already reserved space"
6017 );
6018 empty_rows.reserve(10, 20);
6019 assert_eq!(
6020 empty_rows.size(),
6021 before_size,
6022 "Size should not change when already have space for the expected reserved data"
6023 );
6024
6025 empty_rows.reserve(100, 20);
6026 assert!(
6027 empty_rows.size() > before_size,
6028 "Size should increase when reserving more space than previously reserved"
6029 );
6030
6031 let before_size = empty_rows.size();
6032
6033 empty_rows.reserve(20, 100);
6034 assert!(
6035 empty_rows.size() > before_size,
6036 "Size should increase when reserving more space than previously reserved"
6037 );
6038 }
6039
6040 #[test]
6041 fn empty_rows_should_return_empty_lengths_iterator() {
6042 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6043 .unwrap()
6044 .empty_rows(0, 0);
6045 let mut lengths_iter = rows.lengths();
6046 assert_eq!(lengths_iter.next(), None);
6047 }
6048
6049 #[test]
6050 #[should_panic(expected = "row index out of bounds")]
6051 fn row_should_panic_on_overflowing_index() {
6052 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6053 .unwrap()
6054 .empty_rows(0, 0);
6055 rows.row(usize::MAX);
6056 }
6057
6058 #[test]
6059 #[should_panic(expected = "row index out of bounds")]
6060 fn row_len_should_panic_on_overflowing_index() {
6061 let rows = RowConverter::new(vec![SortField::new(DataType::Int32)])
6062 .unwrap()
6063 .empty_rows(0, 0);
6064 rows.row_len(usize::MAX);
6065 }
6066
6067 #[test]
6068 fn test_nested_null_list() {
6069 let null_array = Arc::new(NullArray::new(3));
6070 let list: ArrayRef = Arc::new(ListArray::new(
6072 Field::new_list_field(DataType::Null, true).into(),
6073 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6074 null_array,
6075 None,
6076 ));
6077
6078 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6079 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6080 let back = converter.convert_rows(&rows).unwrap();
6081
6082 assert_eq!(&list, &back[0]);
6083 }
6084
6085 #[test]
6087 fn test_double_nested_null_list() {
6088 let null_array = Arc::new(NullArray::new(1));
6089 let nested_field = Arc::new(Field::new_list_field(DataType::Null, true));
6091 let nested_list = Arc::new(ListArray::new(
6092 nested_field.clone(),
6093 OffsetBuffer::from_lengths(vec![1]),
6094 null_array,
6095 None,
6096 ));
6097 let list = Arc::new(ListArray::new(
6099 Field::new_list_field(DataType::List(nested_field), true).into(),
6100 OffsetBuffer::from_lengths(vec![1]),
6101 nested_list,
6102 None,
6103 )) as ArrayRef;
6104
6105 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6106 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6107 let back = converter.convert_rows(&rows).unwrap();
6108
6109 assert_eq!(&list, &back[0]);
6110 }
6111
6112 #[test]
6114 fn test_large_list_null() {
6115 let null_array = Arc::new(NullArray::new(3));
6116 let list: ArrayRef = Arc::new(LargeListArray::new(
6118 Field::new_list_field(DataType::Null, true).into(),
6119 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6120 null_array,
6121 None,
6122 ));
6123
6124 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6125 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6126 let back = converter.convert_rows(&rows).unwrap();
6127
6128 assert_eq!(&list, &back[0]);
6129 }
6130
6131 #[test]
6133 fn test_fixed_size_list_null() {
6134 let null_array = Arc::new(NullArray::new(6));
6135 let list: ArrayRef = Arc::new(FixedSizeListArray::new(
6137 Arc::new(Field::new_list_field(DataType::Null, true)),
6138 2,
6139 null_array,
6140 None,
6141 ));
6142
6143 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6144 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6145 let back = converter.convert_rows(&rows).unwrap();
6146
6147 assert_eq!(&list, &back[0]);
6148 }
6149
6150 #[test]
6152 fn test_fixed_size_list_of_dictionaries_round_trips() {
6153 let dict_dt = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
6156 let element_field = Arc::new(Field::new("item", dict_dt.clone(), true));
6157 let fsl_dt = DataType::FixedSizeList(Arc::clone(&element_field), 2);
6158
6159 let values = Arc::new(StringArray::from(vec!["a", "b"]));
6160 let keys = Int32Array::from(vec![0, 1]);
6161 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
6162 let fsl: ArrayRef = Arc::new(FixedSizeListArray::new(
6163 Arc::clone(&element_field),
6164 2,
6165 Arc::new(dict),
6166 None,
6167 ));
6168
6169 assert!(RowConverter::supports_fields(&[SortField::new(
6170 fsl_dt.clone()
6171 )]));
6172
6173 let converter = RowConverter::new(vec![SortField::new(fsl_dt.clone())]).unwrap();
6174 let rows = converter.convert_columns(&[Arc::clone(&fsl)]).unwrap();
6175
6176 let back = converter.convert_rows(&rows).unwrap();
6179 assert_eq!(back.len(), 1);
6180
6181 let out = back[0]
6185 .as_any()
6186 .downcast_ref::<FixedSizeListArray>()
6187 .expect("decoded array must be a FixedSizeListArray");
6188 assert_eq!(out.len(), 1);
6189 assert_eq!(out.value_length(), 2);
6190 assert_eq!(out.values().data_type(), &DataType::Utf8);
6194
6195 let values = out
6197 .values()
6198 .as_any()
6199 .downcast_ref::<StringArray>()
6200 .expect("child must be a StringArray after flattening");
6201 assert_eq!(values.value(0), "a");
6202 assert_eq!(values.value(1), "b");
6203 }
6204
6205 #[test]
6207 fn test_list_null_variations() {
6208 let null_array = Arc::new(NullArray::new(3));
6210 let list: ArrayRef = Arc::new(ListArray::new(
6211 Field::new_list_field(DataType::Null, true).into(),
6212 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6213 null_array,
6214 None,
6215 ));
6216
6217 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6218 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6219 let back = converter.convert_rows(&rows).unwrap();
6220 assert_eq!(&list, &back[0]);
6221
6222 let null_array = Arc::new(NullArray::new(3));
6224 let list: ArrayRef = Arc::new(ListArray::new(
6225 Field::new_list_field(DataType::Null, true).into(),
6226 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6227 null_array,
6228 Some(vec![true, false, true].into()),
6229 ));
6230
6231 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6232 let back = converter.convert_rows(&rows).unwrap();
6233 assert_eq!(&list, &back[0]);
6234
6235 let null_array = Arc::new(NullArray::new(0));
6237 let list: ArrayRef = Arc::new(ListArray::new(
6238 Field::new_list_field(DataType::Null, true).into(),
6239 OffsetBuffer::from_lengths(vec![]),
6240 null_array,
6241 None,
6242 ));
6243
6244 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6245 let back = converter.convert_rows(&rows).unwrap();
6246 assert_eq!(&list, &back[0]);
6247
6248 let null_array = Arc::new(NullArray::new(0));
6250 let list: ArrayRef = Arc::new(ListArray::new(
6251 Field::new_list_field(DataType::Null, true).into(),
6252 OffsetBuffer::from_lengths(vec![0, 0, 0]),
6253 null_array,
6254 None,
6255 ));
6256
6257 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6258 let back = converter.convert_rows(&rows).unwrap();
6259 assert_eq!(&list, &back[0]);
6260 }
6261
6262 #[test]
6264 fn test_list_null_descending() {
6265 let null_array = Arc::new(NullArray::new(3));
6266 let list: ArrayRef = Arc::new(ListArray::new(
6268 Field::new_list_field(DataType::Null, true).into(),
6269 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6270 null_array,
6271 None,
6272 ));
6273
6274 let options = SortOptions::default().with_descending(true);
6275 let field = SortField::new_with_options(list.data_type().clone(), options);
6276 let converter = RowConverter::new(vec![field]).unwrap();
6277 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6278 let back = converter.convert_rows(&rows).unwrap();
6279
6280 assert_eq!(&list, &back[0]);
6281 }
6282
6283 #[test]
6285 fn test_struct_with_null_field() {
6286 let null_array = Arc::new(NullArray::new(3));
6288 let int_array = Arc::new(Int32Array::from(vec![1, 2, 3]));
6289
6290 let struct_array: ArrayRef = Arc::new(StructArray::new(
6291 vec![
6292 Arc::new(Field::new("a", DataType::Null, true)),
6293 Arc::new(Field::new("b", DataType::Int32, true)),
6294 ]
6295 .into(),
6296 vec![null_array, int_array],
6297 Some(vec![true, true, false].into()), ));
6299
6300 let converter =
6301 RowConverter::new(vec![SortField::new(struct_array.data_type().clone())]).unwrap();
6302 let rows = converter
6303 .convert_columns(&[Arc::clone(&struct_array)])
6304 .unwrap();
6305 let back = converter.convert_rows(&rows).unwrap();
6306
6307 assert_eq!(&struct_array, &back[0]);
6308 }
6309
6310 #[test]
6312 fn test_nested_struct_with_null() {
6313 let inner_null = Arc::new(NullArray::new(2));
6315 let inner_struct = Arc::new(StructArray::new(
6316 vec![Arc::new(Field::new("x", DataType::Null, true))].into(),
6317 vec![inner_null],
6318 None,
6319 ));
6320
6321 let y_array = Arc::new(Int32Array::from(vec![10, 20]));
6323 let outer_struct: ArrayRef = Arc::new(StructArray::new(
6324 vec![
6325 Arc::new(Field::new("inner", inner_struct.data_type().clone(), true)),
6326 Arc::new(Field::new("y", DataType::Int32, true)),
6327 ]
6328 .into(),
6329 vec![inner_struct, y_array],
6330 None,
6331 ));
6332
6333 let converter =
6334 RowConverter::new(vec![SortField::new(outer_struct.data_type().clone())]).unwrap();
6335 let rows = converter
6336 .convert_columns(&[Arc::clone(&outer_struct)])
6337 .unwrap();
6338 let back = converter.convert_rows(&rows).unwrap();
6339
6340 assert_eq!(&outer_struct, &back[0]);
6341 }
6342
6343 #[test]
6345 fn test_map_null_variations() {
6346 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6348 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6349
6350 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6351 let entries_fields = vec![
6352 Arc::new(Field::new(
6353 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6354 DataType::Utf8,
6355 false,
6356 )),
6357 Arc::new(Field::new(
6358 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6359 DataType::Null,
6360 true,
6361 )),
6362 ];
6363 let struct_field = Arc::new(Field::new(
6364 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6365 DataType::Struct(entries_fields.clone().into()),
6366 false,
6367 ));
6368 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6369
6370 let map: ArrayRef = Arc::new(MapArray::new(
6371 struct_field.clone(),
6372 offsets,
6373 entries,
6374 None,
6375 false,
6376 ));
6377
6378 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6379 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6380 let back = converter.convert_rows(&rows).unwrap();
6381 assert_eq!(back.len(), 1);
6382 back[0].to_data().validate_full().unwrap();
6383 assert_eq!(&map, &back[0]);
6384
6385 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6387 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6388
6389 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6390 let entries_fields = vec![
6391 Arc::new(Field::new(
6392 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6393 DataType::Utf8,
6394 false,
6395 )),
6396 Arc::new(Field::new(
6397 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6398 DataType::Null,
6399 true,
6400 )),
6401 ];
6402 let struct_field = Arc::new(Field::new(
6403 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6404 DataType::Struct(entries_fields.clone().into()),
6405 false,
6406 ));
6407 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6408
6409 let map: ArrayRef = Arc::new(MapArray::new(
6410 struct_field.clone(),
6411 offsets,
6412 entries,
6413 Some(vec![true, false, true].into()),
6414 false,
6415 ));
6416
6417 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6418 let back = converter.convert_rows(&rows).unwrap();
6419 assert_eq!(back.len(), 1);
6420 back[0].to_data().validate_full().unwrap();
6421 assert_eq!(&map, &back[0]);
6422
6423 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6425 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6426
6427 let offsets = OffsetBuffer::new(vec![0i32].into());
6428 let entries_fields = vec![
6429 Arc::new(Field::new(
6430 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6431 DataType::Utf8,
6432 false,
6433 )),
6434 Arc::new(Field::new(
6435 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6436 DataType::Null,
6437 true,
6438 )),
6439 ];
6440 let struct_field = Arc::new(Field::new(
6441 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6442 DataType::Struct(entries_fields.clone().into()),
6443 false,
6444 ));
6445 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6446
6447 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6448
6449 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6450 let back = converter.convert_rows(&rows).unwrap();
6451 assert_eq!(back.len(), 1);
6452 back[0].to_data().validate_full().unwrap();
6453 assert_eq!(&map, &back[0]);
6454 }
6455
6456 #[test]
6458 fn test_map_null_descending() {
6459 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6461 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6462
6463 let offsets = OffsetBuffer::new(vec![0, 1, 1, 3].into());
6464 let entries_fields = vec![
6465 Arc::new(Field::new(
6466 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6467 DataType::Utf8,
6468 false,
6469 )),
6470 Arc::new(Field::new(
6471 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6472 DataType::Null,
6473 true,
6474 )),
6475 ];
6476 let struct_field = Arc::new(Field::new(
6477 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6478 DataType::Struct(entries_fields.clone().into()),
6479 false,
6480 ));
6481 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6482
6483 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6484
6485 let options = SortOptions::default().with_descending(true);
6486 let field = SortField::new_with_options(map.data_type().clone(), options);
6487 let converter = RowConverter::new(vec![field]).unwrap();
6488 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6489 let back = converter.convert_rows(&rows).unwrap();
6490 assert_eq!(back.len(), 1);
6491 back[0].to_data().validate_full().unwrap();
6492 assert_eq!(&map, &back[0]);
6493 }
6494
6495 #[test]
6497 fn test_map_null_all_empty() {
6498 let keys = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef;
6499 let null_values = Arc::new(NullArray::new(0)) as ArrayRef;
6500
6501 let offsets = OffsetBuffer::new(vec![0, 0, 0, 0].into());
6502 let entries_fields = vec![
6503 Arc::new(Field::new(
6504 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6505 DataType::Utf8,
6506 false,
6507 )),
6508 Arc::new(Field::new(
6509 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6510 DataType::Null,
6511 true,
6512 )),
6513 ];
6514 let struct_field = Arc::new(Field::new(
6515 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6516 DataType::Struct(entries_fields.clone().into()),
6517 false,
6518 ));
6519 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6520
6521 let map: ArrayRef = Arc::new(MapArray::new(struct_field, offsets, entries, None, false));
6522
6523 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6524 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6525
6526 assert_eq!(rows.row(0), rows.row(1));
6528 assert_eq!(rows.row(1), rows.row(2));
6529
6530 let back = converter.convert_rows(&rows).unwrap();
6531 assert_eq!(back.len(), 1);
6532 back[0].to_data().validate_full().unwrap();
6533 assert_eq!(&map, &back[0]);
6534 }
6535
6536 #[test]
6538 fn test_nested_map_null() {
6539 let inner_keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6541 let inner_null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6542
6543 let inner_entries_fields = vec![
6544 Arc::new(Field::new(
6545 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6546 DataType::Utf8,
6547 false,
6548 )),
6549 Arc::new(Field::new(
6550 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6551 DataType::Null,
6552 true,
6553 )),
6554 ];
6555 let inner_struct_field = Arc::new(Field::new(
6556 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6557 DataType::Struct(inner_entries_fields.clone().into()),
6558 false,
6559 ));
6560 let inner_entries = StructArray::new(
6561 inner_entries_fields.clone().into(),
6562 vec![inner_keys, inner_null_values],
6563 None,
6564 );
6565
6566 let inner_map = Arc::new(MapArray::new(
6568 inner_struct_field.clone(),
6569 OffsetBuffer::new(vec![0, 1, 3].into()),
6570 inner_entries,
6571 None,
6572 false,
6573 )) as ArrayRef;
6574
6575 let outer_keys = Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef;
6577
6578 let inner_map_type = DataType::Map(inner_struct_field.clone(), false);
6579 let outer_entries_fields = vec![
6580 Arc::new(Field::new(
6581 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6582 DataType::Utf8,
6583 false,
6584 )),
6585 Arc::new(Field::new(
6586 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6587 inner_map_type,
6588 true,
6589 )),
6590 ];
6591 let outer_struct_field = Arc::new(Field::new(
6592 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6593 DataType::Struct(outer_entries_fields.clone().into()),
6594 false,
6595 ));
6596 let outer_entries = StructArray::new(
6597 outer_entries_fields.into(),
6598 vec![outer_keys, inner_map],
6599 None,
6600 );
6601
6602 let map: ArrayRef = Arc::new(MapArray::new(
6604 outer_struct_field,
6605 OffsetBuffer::new(vec![0, 1, 2].into()),
6606 outer_entries,
6607 None,
6608 false,
6609 ));
6610
6611 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6612 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6613 let back = converter.convert_rows(&rows).unwrap();
6614 assert_eq!(back.len(), 1);
6615 back[0].to_data().validate_full().unwrap();
6616 assert_eq!(&map, &back[0]);
6617 }
6618
6619 #[test]
6621 fn test_list_of_map_null() {
6622 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6624 let null_values = Arc::new(NullArray::new(3)) as ArrayRef;
6625
6626 let entries_fields = vec![
6627 Arc::new(Field::new(
6628 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6629 DataType::Utf8,
6630 false,
6631 )),
6632 Arc::new(Field::new(
6633 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6634 DataType::Null,
6635 true,
6636 )),
6637 ];
6638 let struct_field = Arc::new(Field::new(
6639 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6640 DataType::Struct(entries_fields.clone().into()),
6641 false,
6642 ));
6643 let entries = StructArray::new(entries_fields.into(), vec![keys, null_values], None);
6644
6645 let map_array = Arc::new(MapArray::new(
6647 struct_field.clone(),
6648 OffsetBuffer::new(vec![0, 1, 1, 3].into()),
6649 entries,
6650 None,
6651 false,
6652 )) as ArrayRef;
6653
6654 let map_type = DataType::Map(struct_field, false);
6655 let list: ArrayRef = Arc::new(ListArray::new(
6657 Arc::new(Field::new_list_field(map_type, true)),
6658 OffsetBuffer::new(vec![0, 1, 3].into()),
6659 map_array,
6660 None,
6661 ));
6662
6663 let converter = RowConverter::new(vec![SortField::new(list.data_type().clone())]).unwrap();
6664 let rows = converter.convert_columns(&[Arc::clone(&list)]).unwrap();
6665 let back = converter.convert_rows(&rows).unwrap();
6666 assert_eq!(&list, &back[0]);
6667 }
6668
6669 #[test]
6671 fn test_map_of_list_null() {
6672 let null_array = Arc::new(NullArray::new(3)) as ArrayRef;
6674 let list_array = Arc::new(ListArray::new(
6676 Arc::new(Field::new_list_field(DataType::Null, true)),
6677 OffsetBuffer::from_lengths(vec![1, 0, 2]),
6678 null_array,
6679 None,
6680 )) as ArrayRef;
6681
6682 let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
6683
6684 let list_type = list_array.data_type().clone();
6685 let entries_fields = vec![
6686 Arc::new(Field::new(
6687 Field::MAP_KEY_FIELD_DEFAULT_NAME,
6688 DataType::Utf8,
6689 false,
6690 )),
6691 Arc::new(Field::new(
6692 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6693 list_type,
6694 true,
6695 )),
6696 ];
6697 let struct_field = Arc::new(Field::new(
6698 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6699 DataType::Struct(entries_fields.clone().into()),
6700 false,
6701 ));
6702 let entries = StructArray::new(entries_fields.into(), vec![keys, list_array], None);
6703
6704 let map: ArrayRef = Arc::new(MapArray::new(
6706 struct_field,
6707 OffsetBuffer::new(vec![0, 3].into()),
6708 entries,
6709 None,
6710 false,
6711 ));
6712
6713 let converter = RowConverter::new(vec![SortField::new(map.data_type().clone())]).unwrap();
6714 let rows = converter.convert_columns(&[Arc::clone(&map)]).unwrap();
6715 let back = converter.convert_rows(&rows).unwrap();
6716 assert_eq!(back.len(), 1);
6717 back[0].to_data().validate_full().unwrap();
6718 assert_eq!(&map, &back[0]);
6719 }
6720
6721 #[test]
6722 fn empty_row_iter_next_back() {
6723 let rows = RowConverter::new(vec![SortField::new(DataType::UInt8)])
6724 .unwrap()
6725 .empty_rows(0, 0);
6726 let mut rows_iter = rows.iter();
6727 assert_eq!(rows_iter.next_back(), None);
6728 assert_eq!(rows_iter.next_back(), None);
6729 assert_eq!(rows_iter.next_back(), None);
6730 }
6731
6732 #[test]
6734 fn test_row_parser_skip_utf8_validation_roundtrip() {
6735 let converter = RowConverter::new(vec![SortField::new(DataType::Utf8)]).unwrap();
6736 let array = StringArray::from(vec!["arrow", "rust"]);
6737 let rows = converter.convert_columns(&[Arc::new(array) as _]).unwrap();
6738 let binary = rows.try_into_binary().expect("fits in i32 offsets");
6739
6740 let parser = unsafe { RowParser::with_skip_utf8_validate(Arc::clone(&converter.fields)) };
6742
6743 let decoded = converter
6744 .convert_rows(binary.iter().map(|b| parser.parse(b.unwrap())))
6745 .unwrap();
6746 let got: Vec<_> = decoded[0].as_string::<i32>().iter().flatten().collect();
6747 assert_eq!(got, vec!["arrow", "rust"]);
6748 }
6749
6750 #[test]
6751 fn row_iter_next_back() {
6752 let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap();
6753 let mut rng = StdRng::seed_from_u64(42);
6754 let array = generate_primitive_array::<UInt8Type>(&mut rng, 100, 0.8);
6755 let rows = row_converter.convert_columns(&[Arc::new(array)]).unwrap();
6756
6757 let mut rows_iter = rows.iter();
6758 let mut bytes: Vec<u8> = vec![];
6759
6760 while let Some(row) = rows_iter.next_back() {
6761 bytes.extend(row.data.iter().rev());
6762 }
6763
6764 bytes.reverse();
6765
6766 assert_eq!(
6767 bytes,
6768 &rows.buffer.as_slice()[..*rows.offsets.last().unwrap()]
6769 );
6770
6771 assert_eq!(rows_iter.next_back(), None);
6772 assert_eq!(rows_iter.next(), None);
6773 }
6774}