1use std::borrow::Cow;
34
35use crate::query::{ColumnInfo, QueryRow};
36use crate::schema::CqlType;
37use crate::types::{DataType, Value};
38use crate::util::value_fmt::ValueFormatter;
39use arrow::array::{
40 ArrayRef, BinaryArray, BooleanArray, Date32Array, Float32Array, Float64Array, Int16Array,
41 Int32Array, Int64Array, Int8Array, ListArray, MapArray, StringArray, StructArray,
42 Time64NanosecondArray, TimestampMillisecondArray,
43};
44use arrow::buffer::{NullBuffer, OffsetBuffer};
45use arrow::datatypes::{DataType as ArrowDataType, Field, Fields, Schema, TimeUnit};
46use arrow::record_batch::RecordBatch;
47use std::collections::HashMap;
48use std::sync::Arc;
49use thiserror::Error;
50
51pub(crate) const DECIMAL_FIXED_SCALE: i32 = 9;
61
62pub(crate) const DECIMAL_MAX_PRECISION: u8 = 38;
64
65pub(crate) const ARROW_EXTENSION_NAME_KEY: &str = "ARROW:extension:name";
69pub(crate) const ARROW_UUID_EXTENSION_NAME: &str = "arrow.uuid";
70
71#[derive(Debug, Error)]
81pub enum ArrowConvertError {
82 #[error("Arrow error: {0}")]
84 Arrow(#[from] arrow::error::ArrowError),
85 #[error("{0}")]
87 InvalidValue(String),
88}
89
90#[inline]
102fn checked_offset(len: usize) -> Result<i32, ArrowConvertError> {
103 i32::try_from(len).map_err(|_| {
104 ArrowConvertError::InvalidValue(format!(
105 "collection offset {} exceeds i32::MAX ({}); Arrow List/Map offsets \
106 are 32-bit — split the row group (fewer rows) or export via LargeList",
107 len,
108 i32::MAX
109 ))
110 })
111}
112
113#[inline]
128fn checked_value_bytes(total_bytes: usize) -> Result<(), ArrowConvertError> {
129 if total_bytes > i32::MAX as usize {
130 return Err(ArrowConvertError::InvalidValue(format!(
131 "cumulative Utf8/Binary byte length {} exceeds i32::MAX ({}); Arrow \
132 StringArray/BinaryArray value offsets are 32-bit — reduce the batch \
133 row count (byte-bounded batching) or export via LargeUtf8/LargeBinary",
134 total_bytes,
135 i32::MAX
136 )));
137 }
138 Ok(())
139}
140
141#[inline]
149fn checked_string_offsets<S: AsRef<str>>(values: &[Option<S>]) -> Result<(), ArrowConvertError> {
150 let total = values
151 .iter()
152 .flatten()
153 .fold(0usize, |acc, s| acc.saturating_add(s.as_ref().len()));
154 checked_value_bytes(total)
155}
156
157#[inline]
164fn checked_binary_offsets<B: AsRef<[u8]>>(values: &[Option<B>]) -> Result<(), ArrowConvertError> {
165 let total = values
166 .iter()
167 .flatten()
168 .fold(0usize, |acc, b| acc.saturating_add(b.as_ref().len()));
169 checked_value_bytes(total)
170}
171
172pub fn build_arrow_schema(columns: &[ColumnInfo]) -> Result<Schema, ArrowConvertError> {
184 let fields: Vec<Field> = columns.iter().map(column_to_field).collect();
185 Ok(Schema::new(fields))
186}
187
188pub fn rows_to_record_batch(
198 columns: &[ColumnInfo],
199 rows: &[QueryRow],
200) -> Result<RecordBatch, ArrowConvertError> {
201 let schema = build_arrow_schema(columns)?;
202 let arrays = convert_to_arrays(columns, rows)?;
203 let batch = RecordBatch::try_new(Arc::new(schema), arrays)?;
204 Ok(batch)
205}
206
207pub(crate) fn bigint_to_i128(n: &num_bigint::BigInt) -> Result<i128, ArrowConvertError> {
218 let tc_bytes = n.to_signed_bytes_be();
219 if tc_bytes.len() > 16 {
220 return Err(ArrowConvertError::InvalidValue(
221 "BigInt value requires more than 16 bytes; cannot fit in i128".to_string(),
222 ));
223 }
224 let pad: u8 = if n.sign() == num_bigint::Sign::Minus {
226 0xFF
227 } else {
228 0x00
229 };
230 let mut buf = [pad; 16];
231 buf[16 - tc_bytes.len()..].copy_from_slice(&tc_bytes);
233 Ok(i128::from_be_bytes(buf))
234}
235
236pub(crate) fn column_to_field(col: &ColumnInfo) -> Field {
246 if let Some(cql_type) = &col.cql_type {
247 if let Some(field) = cql_type_to_arrow_field(&col.name, cql_type, col.nullable) {
248 return field;
249 }
250 }
251 let arrow_type = data_type_to_arrow(&col.data_type);
252 Field::new(&col.name, arrow_type, col.nullable)
253}
254
255pub(crate) fn cql_type_to_arrow_field(
264 name: &str,
265 cql_type: &CqlType,
266 nullable: bool,
267) -> Option<Field> {
268 match cql_type {
269 CqlType::Date => Some(Field::new(name, ArrowDataType::Date32, nullable)),
270 CqlType::Time => Some(Field::new(
271 name,
272 ArrowDataType::Time64(TimeUnit::Nanosecond),
273 nullable,
274 )),
275 CqlType::Decimal => Some(Field::new(
276 name,
277 ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8),
278 nullable,
279 )),
280 CqlType::Varint => {
281 Some(Field::new(
285 name,
286 ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, 0),
287 nullable,
288 ))
289 }
290 CqlType::Duration => {
291 Some(Field::new(name, ArrowDataType::Utf8, nullable))
303 }
304 CqlType::Uuid | CqlType::TimeUuid => {
305 let mut meta = HashMap::new();
308 meta.insert(
309 ARROW_EXTENSION_NAME_KEY.to_string(),
310 ARROW_UUID_EXTENSION_NAME.to_string(),
311 );
312 Some(Field::new(name, ArrowDataType::FixedSizeBinary(16), nullable).with_metadata(meta))
313 }
314 CqlType::Inet => Some(Field::new(name, ArrowDataType::Utf8, nullable)),
315 CqlType::Counter => Some(Field::new(name, ArrowDataType::Int64, nullable)),
316 CqlType::List(inner) | CqlType::Set(inner) => {
319 let item_type = cql_type_to_arrow_data_type(inner);
320 let item_field = Arc::new(Field::new("item", item_type, true));
321 Some(Field::new(name, ArrowDataType::List(item_field), nullable))
322 }
323 CqlType::Frozen(inner) => cql_type_to_arrow_field(name, inner, nullable),
325 CqlType::Map(key_type, val_type) => {
329 let key_arrow = cql_type_to_arrow_data_type(key_type);
330 let val_arrow = cql_type_to_arrow_data_type(val_type);
331 let entries_field = Arc::new(Field::new(
332 "entries",
333 ArrowDataType::Struct(Fields::from(vec![
334 Field::new("key", key_arrow, false),
335 Field::new("value", val_arrow, true),
336 ])),
337 false,
338 ));
339 Some(Field::new(
340 name,
341 ArrowDataType::Map(entries_field, false),
342 nullable,
343 ))
344 }
345 CqlType::Tuple(element_types) => {
348 if element_types.is_empty() {
349 return Some(Field::new(name, ArrowDataType::Utf8, nullable));
350 }
351 let struct_type = cql_type_to_arrow_data_type(cql_type);
352 Some(Field::new(name, struct_type, nullable))
353 }
354 CqlType::Udt(_udt_name, udt_fields) => {
357 if udt_fields.is_empty() {
358 return Some(Field::new(name, ArrowDataType::Utf8, nullable));
359 }
360 let struct_type = cql_type_to_arrow_data_type(cql_type);
361 Some(Field::new(name, struct_type, nullable))
362 }
363 _ => None,
366 }
367}
368
369pub(crate) fn cql_type_to_arrow_data_type(cql_type: &CqlType) -> ArrowDataType {
386 match cql_type {
387 CqlType::Boolean => ArrowDataType::Boolean,
389 CqlType::TinyInt => ArrowDataType::Int8,
390 CqlType::SmallInt => ArrowDataType::Int16,
391 CqlType::Int => ArrowDataType::Int32,
392 CqlType::BigInt => ArrowDataType::Int64,
393 CqlType::Counter => ArrowDataType::Int64,
394 CqlType::Float => ArrowDataType::Float32,
395 CqlType::Double => ArrowDataType::Float64,
396 CqlType::Text | CqlType::Ascii | CqlType::Varchar => ArrowDataType::Utf8,
397 CqlType::Blob => ArrowDataType::Binary,
398 CqlType::Timestamp => ArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
399 CqlType::Date => ArrowDataType::Date32,
400 CqlType::Time => ArrowDataType::Time64(TimeUnit::Nanosecond),
401 CqlType::Decimal => {
402 ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8)
403 }
404 CqlType::Varint => ArrowDataType::Decimal128(DECIMAL_MAX_PRECISION, 0),
405 CqlType::Duration => ArrowDataType::Utf8,
407 CqlType::Uuid | CqlType::TimeUuid => ArrowDataType::FixedSizeBinary(16),
408 CqlType::Inet => ArrowDataType::Utf8,
410 CqlType::List(inner) | CqlType::Set(inner) => {
413 let item_type = cql_type_to_arrow_data_type(inner);
414 ArrowDataType::List(Arc::new(Field::new("item", item_type, true)))
415 }
416 CqlType::Frozen(inner) => cql_type_to_arrow_data_type(inner),
418 CqlType::Map(key_type, val_type) => {
422 let key_arrow = cql_type_to_arrow_data_type(key_type);
423 let val_arrow = cql_type_to_arrow_data_type(val_type);
424 ArrowDataType::Map(
425 Arc::new(Field::new(
426 "entries",
427 ArrowDataType::Struct(Fields::from(vec![
428 Field::new("key", key_arrow, false),
429 Field::new("value", val_arrow, true),
430 ])),
431 false,
432 )),
433 false,
434 )
435 }
436 CqlType::Tuple(element_types) => {
439 if element_types.is_empty() {
440 return ArrowDataType::Utf8;
441 }
442 let struct_fields: Vec<Field> = element_types
443 .iter()
444 .enumerate()
445 .map(|(i, t)| {
446 Field::new(
447 format!("field_{i}"),
448 cql_type_to_arrow_data_type(t),
449 true, )
451 })
452 .collect();
453 ArrowDataType::Struct(Fields::from(struct_fields))
454 }
455 CqlType::Udt(_udt_name, udt_fields) => {
458 if udt_fields.is_empty() {
459 return ArrowDataType::Utf8;
460 }
461 let struct_fields: Vec<Field> = udt_fields
462 .iter()
463 .map(|(field_name, field_type)| {
464 Field::new(
465 field_name.as_str(),
466 cql_type_to_arrow_data_type(field_type),
467 true, )
469 })
470 .collect();
471 ArrowDataType::Struct(Fields::from(struct_fields))
472 }
473 CqlType::Custom(_) => ArrowDataType::Utf8,
475 }
476}
477
478pub(crate) fn build_typed_value_array(
491 cql_type: &CqlType,
492 values: &[Option<&Value>],
493) -> Result<ArrayRef, ArrowConvertError> {
494 let effective_type = unwrap_frozen_type(cql_type);
496
497 match effective_type {
498 CqlType::Boolean => {
502 let arr: Vec<Option<bool>> = values
503 .iter()
504 .filter_map(|opt| {
505 let v = unwrap_frozen_value(*opt)?;
506 Some(match v {
507 Value::Boolean(b) => Ok(Some(*b)),
508 Value::Null => Ok(None),
509 other => Err(ArrowConvertError::InvalidValue(format!(
510 "expected Boolean value in element, got {:?}",
511 other
512 ))),
513 })
514 })
515 .collect::<Result<Vec<Option<bool>>, ArrowConvertError>>()?;
516 Ok(Arc::new(BooleanArray::from(arr)))
517 }
518 CqlType::TinyInt => {
519 let arr: Vec<Option<i8>> = values
520 .iter()
521 .filter_map(|opt| {
522 let v = unwrap_frozen_value(*opt)?;
523 Some(match v {
524 Value::TinyInt(i) => Ok(Some(*i)),
525 Value::Null => Ok(None),
526 other => Err(ArrowConvertError::InvalidValue(format!(
527 "expected TinyInt value in element, got {:?}",
528 other
529 ))),
530 })
531 })
532 .collect::<Result<Vec<Option<i8>>, ArrowConvertError>>()?;
533 Ok(Arc::new(Int8Array::from(arr)))
534 }
535 CqlType::SmallInt => {
536 let arr: Vec<Option<i16>> = values
537 .iter()
538 .filter_map(|opt| {
539 let v = unwrap_frozen_value(*opt)?;
540 Some(match v {
541 Value::SmallInt(i) => Ok(Some(*i)),
542 Value::Null => Ok(None),
543 other => Err(ArrowConvertError::InvalidValue(format!(
544 "expected SmallInt value in element, got {:?}",
545 other
546 ))),
547 })
548 })
549 .collect::<Result<Vec<Option<i16>>, ArrowConvertError>>()?;
550 Ok(Arc::new(Int16Array::from(arr)))
551 }
552 CqlType::Int => {
553 let arr: Vec<Option<i32>> = values
554 .iter()
555 .filter_map(|opt| {
556 let v = unwrap_frozen_value(*opt)?;
557 Some(match v {
558 Value::Integer(i) => Ok(Some(*i)),
559 Value::Null => Ok(None),
560 other => Err(ArrowConvertError::InvalidValue(format!(
561 "expected Int value in element, got {:?}",
562 other
563 ))),
564 })
565 })
566 .collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
567 Ok(Arc::new(Int32Array::from(arr)))
568 }
569 CqlType::BigInt => {
570 let arr: Vec<Option<i64>> = values
571 .iter()
572 .filter_map(|opt| {
573 let v = unwrap_frozen_value(*opt)?;
574 Some(match v {
575 Value::BigInt(i) => Ok(Some(*i)),
576 Value::Null => Ok(None),
577 other => Err(ArrowConvertError::InvalidValue(format!(
578 "expected BigInt value in element, got {:?}",
579 other
580 ))),
581 })
582 })
583 .collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
584 Ok(Arc::new(Int64Array::from(arr)))
585 }
586 CqlType::Counter => {
587 let arr: Vec<Option<i64>> = values
588 .iter()
589 .filter_map(|opt| {
590 let v = unwrap_frozen_value(*opt)?;
591 Some(match v {
592 Value::Counter(c) => Ok(Some(*c)),
593 Value::BigInt(i) => Ok(Some(*i)),
594 Value::Null => Ok(None),
595 other => Err(ArrowConvertError::InvalidValue(format!(
596 "expected Counter value in element, got {:?}",
597 other
598 ))),
599 })
600 })
601 .collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
602 Ok(Arc::new(Int64Array::from(arr)))
603 }
604 CqlType::Float => {
605 let arr: Vec<Option<f32>> = values
606 .iter()
607 .filter_map(|opt| {
608 let v = unwrap_frozen_value(*opt)?;
609 Some(match v {
610 Value::Float32(f) => Ok(Some(*f)),
611 Value::Float(f) => Ok(Some(*f as f32)),
616 Value::Null => Ok(None),
617 other => Err(ArrowConvertError::InvalidValue(format!(
618 "expected Float value in element, got {:?}",
619 other
620 ))),
621 })
622 })
623 .collect::<Result<Vec<Option<f32>>, ArrowConvertError>>()?;
624 Ok(Arc::new(Float32Array::from(arr)))
625 }
626 CqlType::Double => {
627 let arr: Vec<Option<f64>> = values
628 .iter()
629 .filter_map(|opt| {
630 let v = unwrap_frozen_value(*opt)?;
631 Some(match v {
632 Value::Float(f) => Ok(Some(*f)),
633 Value::Float32(f) => Ok(Some(*f as f64)),
634 Value::Null => Ok(None),
635 other => Err(ArrowConvertError::InvalidValue(format!(
636 "expected Double value in element, got {:?}",
637 other
638 ))),
639 })
640 })
641 .collect::<Result<Vec<Option<f64>>, ArrowConvertError>>()?;
642 Ok(Arc::new(Float64Array::from(arr)))
643 }
644 CqlType::Text | CqlType::Ascii | CqlType::Varchar => {
645 let refs: Vec<Option<&str>> = values
650 .iter()
651 .filter_map(|opt| {
652 let v = unwrap_frozen_value(*opt)?;
653 Some(match v {
654 Value::Text(s) => std::str::from_utf8(s).map(Some).map_err(|e| {
655 ArrowConvertError::InvalidValue(format!("invalid UTF-8 in text: {e}"))
656 }),
657 Value::Null => Ok(None),
658 other => Err(ArrowConvertError::InvalidValue(format!(
659 "expected Text value in element, got {:?}",
660 other
661 ))),
662 })
663 })
664 .collect::<Result<Vec<Option<&str>>, ArrowConvertError>>()?;
665 checked_string_offsets(&refs)?;
666 Ok(Arc::new(StringArray::from(refs)))
667 }
668 CqlType::Blob => {
669 let refs: Vec<Option<&[u8]>> = values
672 .iter()
673 .filter_map(|opt| {
674 let v = unwrap_frozen_value(*opt)?;
675 Some(match v {
676 Value::Blob(b) => Ok(Some(b.as_ref())),
677 Value::Null => Ok(None),
678 other => Err(ArrowConvertError::InvalidValue(format!(
679 "expected Blob value in element, got {:?}",
680 other
681 ))),
682 })
683 })
684 .collect::<Result<Vec<Option<&[u8]>>, ArrowConvertError>>()?;
685 checked_binary_offsets(&refs)?;
686 Ok(Arc::new(BinaryArray::from(refs)))
687 }
688 CqlType::Timestamp => {
689 let arr: Vec<Option<i64>> = values
690 .iter()
691 .filter_map(|opt| {
692 let v = unwrap_frozen_value(*opt)?;
693 Some(match v {
694 Value::Timestamp(ts) => Ok(Some(*ts)),
695 Value::Null => Ok(None),
696 other => Err(ArrowConvertError::InvalidValue(format!(
697 "expected Timestamp value in element, got {:?}",
698 other
699 ))),
700 })
701 })
702 .collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
703 Ok(Arc::new(
704 TimestampMillisecondArray::from(arr).with_timezone("UTC"),
705 ))
706 }
707 CqlType::Date => {
708 let arr: Vec<Option<i32>> = values
709 .iter()
710 .filter_map(|opt| {
711 let v = unwrap_frozen_value(*opt)?;
712 Some(match v {
713 Value::Date(d) => Ok(Some(*d)),
714 Value::Null => Ok(None),
715 other => Err(ArrowConvertError::InvalidValue(format!(
716 "expected Date value in element, got {:?}",
717 other
718 ))),
719 })
720 })
721 .collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
722 Ok(Arc::new(Date32Array::from(arr)))
723 }
724 CqlType::Time => {
725 let arr: Vec<Option<i64>> = values
726 .iter()
727 .filter_map(|opt| {
728 let v = unwrap_frozen_value(*opt)?;
729 Some(match v {
730 Value::Time(t) => Ok(Some(*t)),
731 Value::Null => Ok(None),
732 other => Err(ArrowConvertError::InvalidValue(format!(
733 "expected Time value in element, got {:?}",
734 other
735 ))),
736 })
737 })
738 .collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
739 Ok(Arc::new(Time64NanosecondArray::from(arr)))
740 }
741 CqlType::Decimal => {
742 let mut builder = arrow::array::Decimal128Builder::with_capacity(values.len())
743 .with_precision_and_scale(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8)?;
744 for opt in values {
745 let v = unwrap_frozen_value(*opt);
746 match v {
747 Some(Value::Decimal { scale, unscaled }) => {
748 let rescaled = rescale_decimal(*scale, unscaled)?;
749 builder.append_value(rescaled);
750 }
751 Some(Value::Null) | None => builder.append_null(),
752 Some(other) => {
753 return Err(ArrowConvertError::InvalidValue(format!(
754 "expected Decimal value in element, got {:?}",
755 other
756 )));
757 }
758 }
759 }
760 Ok(Arc::new(builder.finish()))
761 }
762 CqlType::Varint => {
763 use num_bigint::BigInt;
764 let mut builder = arrow::array::Decimal128Builder::with_capacity(values.len())
765 .with_precision_and_scale(DECIMAL_MAX_PRECISION, 0)?;
766 for opt in values {
767 let v = unwrap_frozen_value(*opt);
768 match v {
769 Some(Value::Varint(bytes)) => {
770 if bytes.is_empty() {
771 builder.append_value(0);
772 } else {
773 let bigint = BigInt::from_signed_bytes_be(bytes);
774 let max_abs = BigInt::from(10i64).pow(38u32) - BigInt::from(1i64);
775 let abs_val = if bigint.sign() == num_bigint::Sign::Minus {
776 -bigint.clone()
777 } else {
778 bigint.clone()
779 };
780 if abs_val > max_abs {
781 return Err(ArrowConvertError::InvalidValue(
782 "varint element exceeds Decimal128(38, 0) range".to_string(),
783 ));
784 }
785 let i128_val = bigint_to_i128(&bigint)?;
786 builder.append_value(i128_val);
787 }
788 }
789 Some(Value::Null) | None => builder.append_null(),
790 Some(other) => {
791 return Err(ArrowConvertError::InvalidValue(format!(
792 "expected Varint value in element, got {:?}",
793 other
794 )));
795 }
796 }
797 }
798 Ok(Arc::new(builder.finish()))
799 }
800 CqlType::Duration => {
801 let arr: Vec<Option<String>> = values
803 .iter()
804 .filter_map(|opt| {
805 let v = unwrap_frozen_value(*opt)?;
806 Some(match v {
807 Value::Duration { .. } => Ok(Some(ValueFormatter::format_value(v))),
808 Value::Null => Ok(None),
809 other => Err(ArrowConvertError::InvalidValue(format!(
810 "expected Duration value in element, got {:?}",
811 other
812 ))),
813 })
814 })
815 .collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
816 checked_string_offsets(&arr)?;
817 Ok(Arc::new(StringArray::from(arr)))
818 }
819 CqlType::Uuid | CqlType::TimeUuid => {
820 let mut builder = arrow::array::FixedSizeBinaryBuilder::with_capacity(values.len(), 16);
821 for opt in values {
822 let v = unwrap_frozen_value(*opt);
823 match v {
824 Some(Value::Uuid(bytes)) => builder.append_value(bytes)?,
825 Some(Value::Null) | None => builder.append_null(),
826 Some(other) => {
827 return Err(ArrowConvertError::InvalidValue(format!(
828 "expected Uuid value in element, got {:?}",
829 other
830 )));
831 }
832 }
833 }
834 Ok(Arc::new(builder.finish()))
835 }
836 CqlType::Inet => {
837 let arr: Vec<Option<String>> = values
838 .iter()
839 .filter_map(|opt| {
840 let v = unwrap_frozen_value(*opt)?;
841 Some(match v {
842 inet @ Value::Inet(_) => Ok(Some(ValueFormatter::format_value(inet))),
845 Value::Null => Ok(None),
846 other => Err(ArrowConvertError::InvalidValue(format!(
847 "expected Inet value in element, got {:?}",
848 other
849 ))),
850 })
851 })
852 .collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
853 checked_string_offsets(&arr)?;
854 Ok(Arc::new(StringArray::from(arr)))
855 }
856 CqlType::List(inner) | CqlType::Set(inner) => {
861 let element_type = cql_type_to_arrow_data_type(inner);
862 let item_field = Arc::new(Field::new("item", element_type, true));
863
864 let mut offsets: Vec<i32> = vec![0];
867 let mut flat_elements: Vec<Option<&Value>> = Vec::new();
868 let mut null_bitmap: Vec<bool> = Vec::new();
869
870 for opt in values {
871 let v = unwrap_frozen_value(*opt);
872 match v {
873 Some(Value::List(items)) | Some(Value::Set(items)) => {
874 null_bitmap.push(true);
875 for item in items {
876 flat_elements.push(Some(item));
877 }
878 offsets.push(checked_offset(flat_elements.len())?);
879 }
880 Some(Value::Null) | None => {
881 null_bitmap.push(false);
882 offsets.push(checked_offset(flat_elements.len())?);
883 }
884 Some(other) => {
885 return Err(ArrowConvertError::InvalidValue(format!(
886 "expected List/Set value, got {:?}",
887 other
888 )));
889 }
890 }
891 }
892
893 let elements_array = build_typed_value_array(inner, &flat_elements)?;
895
896 let offset_buffer = OffsetBuffer::new(offsets.into());
897 let null_buffer = NullBuffer::from(null_bitmap);
898
899 Ok(Arc::new(ListArray::new(
900 item_field,
901 offset_buffer,
902 elements_array,
903 Some(null_buffer),
904 )))
905 }
906 CqlType::Frozen(inner) => build_typed_value_array(inner, values),
909 CqlType::Map(key_type, val_type) => {
924 let key_arrow = cql_type_to_arrow_data_type(key_type);
925 let val_arrow = cql_type_to_arrow_data_type(val_type);
926
927 let mut offsets: Vec<i32> = vec![0];
928 let mut flat_keys: Vec<Option<&Value>> = Vec::new();
929 let mut flat_vals: Vec<Option<&Value>> = Vec::new();
930 let mut null_bitmap: Vec<bool> = Vec::new();
931
932 for opt in values {
933 let v = unwrap_frozen_value(*opt);
934 match v {
935 Some(Value::Map(pairs)) => {
936 null_bitmap.push(true);
937 for (k, val) in pairs {
938 if matches!(k, Value::Null) {
940 return Err(ArrowConvertError::InvalidValue(
941 "null key in map is not allowed in Arrow MapArray".to_string(),
942 ));
943 }
944 flat_keys.push(Some(k));
945 flat_vals.push(Some(val));
946 }
947 offsets.push(checked_offset(flat_keys.len())?);
948 }
949 Some(Value::Null) | None => {
950 null_bitmap.push(false);
951 offsets.push(checked_offset(flat_keys.len())?);
952 }
953 Some(other) => {
954 return Err(ArrowConvertError::InvalidValue(format!(
955 "expected Map value, got {:?}",
956 other
957 )));
958 }
959 }
960 }
961
962 let key_array = build_typed_value_array(key_type, &flat_keys)?;
964 let val_array = build_typed_value_array(val_type, &flat_vals)?;
965
966 let struct_fields = Fields::from(vec![
968 Field::new("key", key_arrow, false),
969 Field::new("value", val_arrow, true),
970 ]);
971 let entries_array =
972 StructArray::new(struct_fields.clone(), vec![key_array, val_array], None);
973
974 let map_field = Arc::new(Field::new(
975 "entries",
976 ArrowDataType::Struct(struct_fields),
977 false,
978 ));
979 let offset_buffer = OffsetBuffer::new(offsets.into());
980 let null_buffer = NullBuffer::from(null_bitmap);
981
982 Ok(Arc::new(MapArray::new(
983 map_field,
984 offset_buffer,
985 entries_array,
986 Some(null_buffer),
987 false,
988 )))
989 }
990 CqlType::Tuple(element_types) => {
1001 if element_types.is_empty() {
1002 let arr: Vec<Option<String>> = values
1005 .iter()
1006 .map(|opt| match unwrap_frozen_value(*opt) {
1007 Some(Value::Null) | None => Ok(None),
1008 Some(v @ Value::Tuple(_)) => Ok(Some(ValueFormatter::format_value(v))),
1009 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1010 "expected Tuple value, got {:?}",
1011 other
1012 ))),
1013 })
1014 .collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
1015 checked_string_offsets(&arr)?;
1016 return Ok(Arc::new(StringArray::from(arr)));
1017 }
1018
1019 let n_rows = values.len();
1020 let n_fields = element_types.len();
1021
1022 let unwrapped: Vec<Option<&Value>> =
1024 values.iter().map(|opt| unwrap_frozen_value(*opt)).collect();
1025
1026 for v in unwrapped.iter() {
1030 match v {
1031 Some(Value::Tuple(_)) | Some(Value::Null) | None => {}
1032 Some(other) => {
1033 return Err(ArrowConvertError::InvalidValue(format!(
1034 "expected Tuple value, got {:?}",
1035 other
1036 )));
1037 }
1038 }
1039 }
1040
1041 let null_bitmap: Vec<bool> = unwrapped
1043 .iter()
1044 .map(|v| !matches!(v, Some(Value::Null) | None))
1045 .collect();
1046
1047 let null_sentinel = Value::Null;
1058 let mut child_arrays: Vec<ArrayRef> = Vec::with_capacity(n_fields);
1059 for (field_idx, element_type) in element_types.iter().enumerate() {
1060 let child_values: Vec<Option<&Value>> = (0..n_rows)
1061 .map(|row_idx| {
1062 match unwrapped[row_idx] {
1063 Some(Value::Tuple(items)) => {
1064 Some(
1066 items
1067 .get(field_idx)
1068 .map(|v| v as &Value)
1069 .unwrap_or(&null_sentinel),
1070 )
1071 }
1072 _ => Some(&null_sentinel),
1075 }
1076 })
1077 .collect();
1078 let child_arr = build_typed_value_array(element_type, &child_values)?;
1079 child_arrays.push(child_arr);
1080 }
1081
1082 let struct_fields: Fields = Fields::from(
1083 element_types
1084 .iter()
1085 .enumerate()
1086 .map(|(i, t)| {
1087 Field::new(format!("field_{i}"), cql_type_to_arrow_data_type(t), true)
1088 })
1089 .collect::<Vec<_>>(),
1090 );
1091
1092 let null_buffer = NullBuffer::from(null_bitmap);
1093 Ok(Arc::new(StructArray::new(
1094 struct_fields,
1095 child_arrays,
1096 Some(null_buffer),
1097 )))
1098 }
1099 CqlType::Udt(_udt_name, udt_fields) => {
1110 if udt_fields.is_empty() {
1111 let arr: Vec<Option<String>> = values
1115 .iter()
1116 .map(|opt| match unwrap_frozen_value(*opt) {
1117 Some(Value::Null) | None => Ok(None),
1118 Some(v @ Value::Udt(_)) => Ok(Some(ValueFormatter::format_value(v))),
1119 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1120 "expected Udt value, got {:?}",
1121 other
1122 ))),
1123 })
1124 .collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
1125 checked_string_offsets(&arr)?;
1126 return Ok(Arc::new(StringArray::from(arr)));
1127 }
1128
1129 let n_rows = values.len();
1130
1131 let unwrapped: Vec<Option<&Value>> =
1133 values.iter().map(|opt| unwrap_frozen_value(*opt)).collect();
1134
1135 for v in unwrapped.iter() {
1139 match v {
1140 Some(Value::Udt(_)) | Some(Value::Null) | None => {}
1141 Some(other) => {
1142 return Err(ArrowConvertError::InvalidValue(format!(
1143 "expected Udt value, got {:?}",
1144 other
1145 )));
1146 }
1147 }
1148 }
1149
1150 let null_bitmap: Vec<bool> = unwrapped
1152 .iter()
1153 .map(|v| !matches!(v, Some(Value::Null) | None))
1154 .collect();
1155
1156 let null_sentinel = Value::Null;
1162 let mut child_arrays: Vec<ArrayRef> = Vec::with_capacity(udt_fields.len());
1163 for (field_name, field_type) in udt_fields.iter() {
1164 let child_values: Vec<Option<&Value>> = (0..n_rows)
1165 .map(|row_idx| match unwrapped[row_idx] {
1166 Some(Value::Udt(udt_val)) => {
1167 Some(
1169 udt_val
1170 .fields
1171 .iter()
1172 .find(|f| &f.name == field_name)
1173 .and_then(|f| f.value.as_ref().map(|v| v as &Value))
1174 .unwrap_or(&null_sentinel),
1175 )
1176 }
1177 _ => Some(&null_sentinel),
1180 })
1181 .collect();
1182 let child_arr = build_typed_value_array(field_type, &child_values)?;
1183 child_arrays.push(child_arr);
1184 }
1185
1186 let struct_fields: Fields = Fields::from(
1187 udt_fields
1188 .iter()
1189 .map(|(field_name, field_type)| {
1190 Field::new(
1191 field_name.as_str(),
1192 cql_type_to_arrow_data_type(field_type),
1193 true,
1194 )
1195 })
1196 .collect::<Vec<_>>(),
1197 );
1198
1199 let null_buffer = NullBuffer::from(null_bitmap);
1200 Ok(Arc::new(StructArray::new(
1201 struct_fields,
1202 child_arrays,
1203 Some(null_buffer),
1204 )))
1205 }
1206 CqlType::Custom(_) => {
1207 let arr: Vec<Option<String>> = values
1208 .iter()
1209 .map(|opt| match opt {
1210 Some(Value::Null) | None => None,
1211 Some(v) => Some(ValueFormatter::format_value(v)),
1212 })
1213 .collect();
1214 checked_string_offsets(&arr)?;
1215 Ok(Arc::new(StringArray::from(arr)))
1216 }
1217 }
1218}
1219
1220pub(crate) fn unwrap_frozen_type(cql_type: &CqlType) -> &CqlType {
1225 let mut t = cql_type;
1226 while let CqlType::Frozen(inner) = t {
1227 t = inner.as_ref();
1228 }
1229 t
1230}
1231
1232pub(crate) fn unwrap_frozen_value(v: Option<&Value>) -> Option<&Value> {
1237 match v {
1238 Some(Value::Frozen(inner)) => Some(inner.as_ref()),
1239 other => other,
1240 }
1241}
1242
1243pub(crate) fn data_type_to_arrow(data_type: &DataType) -> ArrowDataType {
1247 match data_type {
1248 DataType::Null => ArrowDataType::Null,
1249 DataType::Boolean => ArrowDataType::Boolean,
1250 DataType::TinyInt => ArrowDataType::Int8,
1251 DataType::SmallInt => ArrowDataType::Int16,
1252 DataType::Integer => ArrowDataType::Int32,
1253 DataType::BigInt => ArrowDataType::Int64,
1254 DataType::Float32 => ArrowDataType::Float32,
1255 DataType::Float => ArrowDataType::Float64,
1256 DataType::Text => ArrowDataType::Utf8,
1257 DataType::Blob => ArrowDataType::Binary,
1258 DataType::Timestamp => ArrowDataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
1259 DataType::Uuid => ArrowDataType::FixedSizeBinary(16),
1260 DataType::Json => ArrowDataType::Utf8,
1261 DataType::List => {
1262 ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Utf8, true)))
1263 }
1264 DataType::Set => {
1265 ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Utf8, true)))
1266 }
1267 DataType::Map => ArrowDataType::Map(
1268 Arc::new(Field::new(
1269 "entries",
1270 ArrowDataType::Struct(Fields::from(vec![
1271 Field::new("key", ArrowDataType::Utf8, false),
1272 Field::new("value", ArrowDataType::Utf8, true),
1273 ])),
1274 false,
1275 )),
1276 false,
1277 ),
1278 DataType::Tuple => ArrowDataType::Utf8, DataType::Udt => ArrowDataType::Utf8, DataType::Frozen => ArrowDataType::Utf8,
1281 DataType::Tombstone => ArrowDataType::Utf8,
1282 }
1283}
1284
1285pub(crate) fn convert_to_arrays(
1297 columns: &[ColumnInfo],
1298 rows: &[QueryRow],
1299) -> Result<Vec<ArrayRef>, ArrowConvertError> {
1300 let columnar = super::arrow_columnar::transpose_columns(columns, rows);
1301 columns
1302 .iter()
1303 .zip(columnar.iter())
1304 .map(|(col, cells)| convert_column_to_array(col, cells))
1305 .collect()
1306}
1307
1308pub(crate) fn convert_column_to_array(
1323 col: &ColumnInfo,
1324 cells: Cells,
1325) -> Result<ArrayRef, ArrowConvertError> {
1326 if let Some(cql_type) = &col.cql_type {
1328 let effective = unwrap_frozen_type(cql_type);
1330 match effective {
1331 CqlType::Date => return build_date32_array(col, cells),
1332 CqlType::Time => return build_time64_ns_array(col, cells),
1333 CqlType::Decimal => return build_decimal128_array(col, cells),
1334 CqlType::Varint => return build_varint_as_decimal128_array(col, cells),
1335 CqlType::Duration => return build_duration_utf8_array(col, cells),
1336 CqlType::Uuid | CqlType::TimeUuid => return build_uuid_fixed_binary_array(col, cells),
1337 CqlType::Inet => return build_inet_utf8_array(col, cells),
1338 CqlType::Counter => return build_int64_array(col, cells),
1339 CqlType::List(_)
1341 | CqlType::Set(_)
1342 | CqlType::Map(_, _)
1343 | CqlType::Tuple(_)
1344 | CqlType::Udt(_, _) => {
1345 return build_typed_value_array(cql_type, cells);
1346 }
1347 _ => {}
1349 }
1350 }
1351
1352 match &col.data_type {
1354 DataType::Boolean => build_boolean_array(col, cells),
1355 DataType::TinyInt => build_int8_array(col, cells),
1356 DataType::SmallInt => build_int16_array(col, cells),
1357 DataType::Integer => build_int32_array(col, cells),
1358 DataType::BigInt => build_int64_array(col, cells),
1359 DataType::Float32 => build_float32_array(col, cells),
1360 DataType::Float => build_float64_array(col, cells),
1361 DataType::Text | DataType::Json => build_string_array(col, cells),
1362 DataType::Blob => build_binary_array(col, cells),
1363 DataType::Timestamp => build_timestamp_array(col, cells),
1364 DataType::Uuid => build_uuid_array(col, cells),
1365 DataType::List | DataType::Set => build_list_array(col, cells),
1366 DataType::Map => build_map_array(col, cells),
1367 DataType::Tuple
1368 | DataType::Udt
1369 | DataType::Frozen
1370 | DataType::Tombstone
1371 | DataType::Null => {
1372 build_string_array(col, cells) }
1374 }
1375}
1376
1377pub(crate) use super::arrow_decimal::rescale_decimal;
1385
1386type Cells<'a> = &'a [Option<&'a Value>];
1394
1395fn build_boolean_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1396 let values: Vec<Option<bool>> = cells
1397 .iter()
1398 .map(|cell| match unwrap_frozen_value(*cell) {
1399 None => Ok(None),
1400 Some(Value::Boolean(b)) => Ok(Some(*b)),
1401 Some(Value::Null) => Ok(None),
1402 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1403 "column '{}': expected Boolean value, got {:?}",
1404 col.name, other
1405 ))),
1406 })
1407 .collect::<Result<Vec<Option<bool>>, ArrowConvertError>>()?;
1408 Ok(Arc::new(BooleanArray::from(values)))
1409}
1410
1411fn build_int8_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1412 let values: Vec<Option<i8>> = cells
1413 .iter()
1414 .map(|cell| match unwrap_frozen_value(*cell) {
1415 None => Ok(None),
1416 Some(Value::TinyInt(i)) => Ok(Some(*i)),
1417 Some(Value::Null) => Ok(None),
1418 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1419 "column '{}': expected TinyInt value, got {:?}",
1420 col.name, other
1421 ))),
1422 })
1423 .collect::<Result<Vec<Option<i8>>, ArrowConvertError>>()?;
1424 Ok(Arc::new(Int8Array::from(values)))
1425}
1426
1427fn build_int16_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1428 let values: Vec<Option<i16>> = cells
1429 .iter()
1430 .map(|cell| match unwrap_frozen_value(*cell) {
1431 None => Ok(None),
1432 Some(Value::SmallInt(i)) => Ok(Some(*i)),
1433 Some(Value::Null) => Ok(None),
1434 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1435 "column '{}': expected SmallInt value, got {:?}",
1436 col.name, other
1437 ))),
1438 })
1439 .collect::<Result<Vec<Option<i16>>, ArrowConvertError>>()?;
1440 Ok(Arc::new(Int16Array::from(values)))
1441}
1442
1443fn build_int32_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1444 let allow_compat = col.cql_type.is_none();
1449 let values: Vec<Option<i32>> = cells
1450 .iter()
1451 .map(|cell| match unwrap_frozen_value(*cell) {
1452 None => Ok(None),
1453 Some(Value::Integer(i)) => Ok(Some(*i)),
1454 Some(Value::Date(d)) if allow_compat => Ok(Some(*d)), Some(Value::Null) => Ok(None),
1456 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1457 "column '{}': expected Int value, got {:?}",
1458 col.name, other
1459 ))),
1460 })
1461 .collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
1462 Ok(Arc::new(Int32Array::from(values)))
1463}
1464
1465fn build_int64_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1466 let effective = col.cql_type.as_ref().map(unwrap_frozen_type);
1472 let allow_counter = matches!(effective, None | Some(CqlType::Counter));
1473 let allow_compat = effective.is_none();
1474 let values: Vec<Option<i64>> = cells
1475 .iter()
1476 .map(|cell| match unwrap_frozen_value(*cell) {
1477 None => Ok(None),
1478 Some(Value::BigInt(i)) => Ok(Some(*i)),
1479 Some(Value::Counter(c)) if allow_counter => Ok(Some(*c)),
1480 Some(Value::Time(t)) if allow_compat => Ok(Some(*t)), Some(Value::Null) => Ok(None),
1482 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1483 "column '{}': expected BigInt value, got {:?}",
1484 col.name, other
1485 ))),
1486 })
1487 .collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
1488 Ok(Arc::new(Int64Array::from(values)))
1489}
1490
1491fn build_float32_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1492 let values: Vec<Option<f32>> = cells
1493 .iter()
1494 .map(|cell| match unwrap_frozen_value(*cell) {
1495 None => Ok(None),
1496 Some(Value::Float32(f)) => Ok(Some(*f)),
1497 Some(Value::Float(f)) => Ok(Some(*f as f32)),
1501 Some(Value::Null) => Ok(None),
1502 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1503 "column '{}': expected Float value, got {:?}",
1504 col.name, other
1505 ))),
1506 })
1507 .collect::<Result<Vec<Option<f32>>, ArrowConvertError>>()?;
1508 Ok(Arc::new(Float32Array::from(values)))
1509}
1510
1511fn build_float64_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1512 let values: Vec<Option<f64>> = cells
1513 .iter()
1514 .map(|cell| match unwrap_frozen_value(*cell) {
1515 None => Ok(None),
1516 Some(Value::Float(f)) => Ok(Some(*f)),
1517 Some(Value::Float32(f)) => Ok(Some(*f as f64)),
1518 Some(Value::Null) => Ok(None),
1519 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1520 "column '{}': expected Double value, got {:?}",
1521 col.name, other
1522 ))),
1523 })
1524 .collect::<Result<Vec<Option<f64>>, ArrowConvertError>>()?;
1525 Ok(Arc::new(Float64Array::from(values)))
1526}
1527
1528fn build_string_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1529 let strict_text = matches!(
1535 col.cql_type.as_ref().map(unwrap_frozen_type),
1536 Some(CqlType::Text | CqlType::Ascii | CqlType::Varchar)
1537 );
1538 if strict_text {
1539 let refs: Vec<Option<&str>> = cells
1544 .iter()
1545 .map(|cell| match unwrap_frozen_value(*cell) {
1546 None | Some(Value::Null) => Ok(None),
1547 Some(Value::Text(s)) => std::str::from_utf8(s).map(Some).map_err(|e| {
1548 ArrowConvertError::InvalidValue(format!("invalid UTF-8 in text: {e}"))
1549 }),
1550 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1551 "column '{}': expected Text value, got {:?}",
1552 col.name, other
1553 ))),
1554 })
1555 .collect::<Result<Vec<Option<&str>>, ArrowConvertError>>()?;
1556 checked_string_offsets(&refs)?;
1557 return Ok(Arc::new(StringArray::from(refs)));
1558 }
1559 let values: Vec<Option<Cow<str>>> = cells
1565 .iter()
1566 .map(|cell| match unwrap_frozen_value(*cell) {
1567 None => Ok(None),
1568 Some(Value::Null) => Ok(None),
1569 Some(Value::Text(s)) => std::str::from_utf8(s)
1570 .map(|st| Some(Cow::Borrowed(st)))
1571 .map_err(|e| {
1572 ArrowConvertError::InvalidValue(format!("invalid UTF-8 in text: {e}"))
1573 }),
1574 Some(Value::Json(j)) => Ok(Some(Cow::Owned(j.to_string()))),
1575 Some(other) => Ok(Some(Cow::Owned(ValueFormatter::format_value(other)))),
1576 })
1577 .collect::<Result<Vec<Option<Cow<str>>>, ArrowConvertError>>()?;
1578 checked_string_offsets(&values)?;
1579 Ok(Arc::new(StringArray::from_iter(
1580 values.iter().map(|v| v.as_deref()),
1581 )))
1582}
1583
1584fn build_binary_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1585 let values: Vec<Option<&[u8]>> = cells
1586 .iter()
1587 .map(|cell| match unwrap_frozen_value(*cell) {
1588 None => Ok(None),
1589 Some(Value::Blob(b)) => Ok(Some(b.as_ref())),
1590 Some(Value::Null) => Ok(None),
1591 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1592 "column '{}': expected Blob value, got {:?}",
1593 col.name, other
1594 ))),
1595 })
1596 .collect::<Result<Vec<Option<&[u8]>>, ArrowConvertError>>()?;
1597 checked_binary_offsets(&values)?;
1598 Ok(Arc::new(BinaryArray::from(values)))
1599}
1600
1601fn build_timestamp_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1602 let values: Vec<Option<i64>> = cells
1603 .iter()
1604 .map(|cell| match unwrap_frozen_value(*cell) {
1605 None => Ok(None),
1606 Some(Value::Timestamp(ts)) => Ok(Some(*ts)),
1607 Some(Value::Null) => Ok(None),
1608 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1609 "column '{}': expected Timestamp value, got {:?}",
1610 col.name, other
1611 ))),
1612 })
1613 .collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
1614 Ok(Arc::new(
1615 TimestampMillisecondArray::from(values).with_timezone("UTC"),
1616 ))
1617}
1618
1619fn build_uuid_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1620 let mut builder = arrow::array::FixedSizeBinaryBuilder::with_capacity(cells.len(), 16);
1623 for cell in cells {
1624 match unwrap_frozen_value(*cell) {
1625 None | Some(Value::Null) => builder.append_null(),
1626 Some(Value::Uuid(uuid)) => builder.append_value(uuid)?,
1627 Some(other) => {
1628 return Err(ArrowConvertError::InvalidValue(format!(
1629 "column '{}': expected Uuid value, got {:?}",
1630 col.name, other
1631 )));
1632 }
1633 }
1634 }
1635 Ok(Arc::new(builder.finish()))
1636}
1637
1638fn build_date32_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1644 let values: Vec<Option<i32>> = cells
1645 .iter()
1646 .map(|cell| match unwrap_frozen_value(*cell) {
1647 None => Ok(None),
1648 Some(Value::Date(days)) => Ok(Some(*days)),
1649 Some(Value::Null) => Ok(None),
1650 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1651 "column '{}': expected Date value, got {:?}",
1652 col.name, other
1653 ))),
1654 })
1655 .collect::<Result<Vec<Option<i32>>, ArrowConvertError>>()?;
1656 Ok(Arc::new(Date32Array::from(values)))
1657}
1658
1659fn build_time64_ns_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1661 let values: Vec<Option<i64>> = cells
1662 .iter()
1663 .map(|cell| match unwrap_frozen_value(*cell) {
1664 None => Ok(None),
1665 Some(Value::Time(nanos)) => Ok(Some(*nanos)),
1666 Some(Value::Null) => Ok(None),
1667 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1668 "column '{}': expected Time value, got {:?}",
1669 col.name, other
1670 ))),
1671 })
1672 .collect::<Result<Vec<Option<i64>>, ArrowConvertError>>()?;
1673 Ok(Arc::new(Time64NanosecondArray::from(values)))
1674}
1675
1676fn build_decimal128_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1679 let mut builder = arrow::array::Decimal128Builder::with_capacity(cells.len())
1680 .with_precision_and_scale(DECIMAL_MAX_PRECISION, DECIMAL_FIXED_SCALE as i8)?;
1681
1682 for cell in cells {
1683 match unwrap_frozen_value(*cell) {
1684 Some(Value::Decimal { scale, unscaled }) => {
1685 let rescaled = rescale_decimal(*scale, unscaled).map_err(|e| {
1686 ArrowConvertError::InvalidValue(format!("Column '{}': {e}", col.name))
1687 })?;
1688 builder.append_value(rescaled);
1689 }
1690 Some(Value::Null) | None => {
1691 builder.append_null();
1692 }
1693 Some(other) => {
1694 return Err(ArrowConvertError::InvalidValue(format!(
1695 "Column '{}': expected Decimal value, got {:?}",
1696 col.name, other
1697 )));
1698 }
1699 }
1700 }
1701 Ok(Arc::new(builder.finish()))
1702}
1703
1704fn build_varint_as_decimal128_array(
1706 col: &ColumnInfo,
1707 cells: Cells,
1708) -> Result<ArrayRef, ArrowConvertError> {
1709 use num_bigint::BigInt;
1710
1711 let mut builder = arrow::array::Decimal128Builder::with_capacity(cells.len())
1712 .with_precision_and_scale(DECIMAL_MAX_PRECISION, 0)?;
1713
1714 for cell in cells {
1715 match unwrap_frozen_value(*cell) {
1716 Some(Value::Varint(bytes)) => {
1717 if bytes.is_empty() {
1718 builder.append_value(0);
1719 } else {
1720 let bigint = BigInt::from_signed_bytes_be(bytes);
1721
1722 let max_abs = BigInt::from(10i64).pow(38u32) - BigInt::from(1i64);
1724 let abs_val = if bigint.sign() == num_bigint::Sign::Minus {
1725 -bigint.clone()
1726 } else {
1727 bigint.clone()
1728 };
1729 if abs_val > max_abs {
1730 return Err(ArrowConvertError::InvalidValue(format!(
1731 "Column '{}': varint value exceeds Decimal128(38, 0) range",
1732 col.name
1733 )));
1734 }
1735
1736 let i128_val = bigint_to_i128(&bigint).map_err(|e| {
1737 ArrowConvertError::InvalidValue(format!("Column '{}': {e}", col.name))
1738 })?;
1739 builder.append_value(i128_val);
1740 }
1741 }
1742 Some(Value::Null) | None => {
1743 builder.append_null();
1744 }
1745 Some(other) => {
1746 return Err(ArrowConvertError::InvalidValue(format!(
1747 "Column '{}': expected Varint value, got {:?}",
1748 col.name, other
1749 )));
1750 }
1751 }
1752 }
1753 Ok(Arc::new(builder.finish()))
1754}
1755
1756fn build_duration_utf8_array(
1758 col: &ColumnInfo,
1759 cells: Cells,
1760) -> Result<ArrayRef, ArrowConvertError> {
1761 let values: Vec<Option<String>> = cells
1762 .iter()
1763 .map(|cell| match unwrap_frozen_value(*cell) {
1764 None => Ok(None),
1765 Some(v @ Value::Duration { .. }) => Ok(Some(ValueFormatter::format_value(v))),
1766 Some(Value::Null) => Ok(None),
1767 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1768 "column '{}': expected Duration value, got {:?}",
1769 col.name, other
1770 ))),
1771 })
1772 .collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
1773 checked_string_offsets(&values)?;
1774 Ok(Arc::new(StringArray::from(values)))
1775}
1776
1777fn build_uuid_fixed_binary_array(
1779 col: &ColumnInfo,
1780 cells: Cells,
1781) -> Result<ArrayRef, ArrowConvertError> {
1782 let mut builder = arrow::array::FixedSizeBinaryBuilder::with_capacity(cells.len(), 16);
1783 for cell in cells {
1784 match unwrap_frozen_value(*cell) {
1785 Some(Value::Uuid(bytes)) => builder.append_value(bytes)?,
1786 Some(Value::Null) | None => builder.append_null(),
1787 Some(other) => {
1788 return Err(ArrowConvertError::InvalidValue(format!(
1789 "Column '{}': expected Uuid value, got {:?}",
1790 col.name, other
1791 )));
1792 }
1793 }
1794 }
1795 Ok(Arc::new(builder.finish()))
1796}
1797
1798fn build_inet_utf8_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1800 let values: Vec<Option<String>> = cells
1801 .iter()
1802 .map(|cell| match unwrap_frozen_value(*cell) {
1803 None => Ok(None),
1804 Some(inet @ Value::Inet(_)) => Ok(Some(ValueFormatter::format_value(inet))),
1807 Some(Value::Null) => Ok(None),
1808 Some(other) => Err(ArrowConvertError::InvalidValue(format!(
1809 "column '{}': expected Inet value, got {:?}",
1810 col.name, other
1811 ))),
1812 })
1813 .collect::<Result<Vec<Option<String>>, ArrowConvertError>>()?;
1814 checked_string_offsets(&values)?;
1815 Ok(Arc::new(StringArray::from(values)))
1816}
1817
1818fn build_list_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1819 let mut offsets: Vec<i32> = vec![0];
1821 let mut values: Vec<Option<String>> = Vec::new();
1822 let mut null_bitmap: Vec<bool> = Vec::new();
1823
1824 for cell in cells {
1825 match unwrap_frozen_value(*cell) {
1826 Some(Value::List(items)) | Some(Value::Set(items)) => {
1827 null_bitmap.push(true);
1828 for item in items {
1829 values.push(Some(ValueFormatter::format_value(item)));
1830 }
1831 offsets.push(checked_offset(values.len())?);
1832 }
1833 Some(Value::Null) | None => {
1834 null_bitmap.push(false);
1835 offsets.push(checked_offset(values.len())?);
1836 }
1837 Some(other) => {
1838 return Err(ArrowConvertError::InvalidValue(format!(
1839 "column '{}': expected List/Set value, got {:?}",
1840 col.name, other
1841 )));
1842 }
1843 }
1844 }
1845
1846 checked_string_offsets(&values)?;
1847 let values_array = Arc::new(StringArray::from(values)) as ArrayRef;
1848 let field = Arc::new(Field::new("item", ArrowDataType::Utf8, true));
1849 let offset_buffer = OffsetBuffer::new(offsets.into());
1850 let null_buffer = NullBuffer::from(null_bitmap);
1851
1852 Ok(Arc::new(ListArray::new(
1853 field,
1854 offset_buffer,
1855 values_array,
1856 Some(null_buffer),
1857 )))
1858}
1859
1860fn build_map_array(col: &ColumnInfo, cells: Cells) -> Result<ArrayRef, ArrowConvertError> {
1861 let mut offsets: Vec<i32> = vec![0];
1863 let mut keys: Vec<Option<String>> = Vec::new();
1864 let mut values: Vec<Option<String>> = Vec::new();
1865 let mut null_bitmap: Vec<bool> = Vec::new();
1866
1867 for cell in cells {
1868 match unwrap_frozen_value(*cell) {
1869 Some(Value::Map(pairs)) => {
1870 null_bitmap.push(true);
1871 for (k, v) in pairs {
1872 keys.push(Some(ValueFormatter::format_value(k)));
1873 values.push(Some(ValueFormatter::format_value(v)));
1874 }
1875 offsets.push(checked_offset(keys.len())?);
1876 }
1877 Some(Value::Null) | None => {
1878 null_bitmap.push(false);
1879 offsets.push(checked_offset(keys.len())?);
1880 }
1881 Some(other) => {
1882 return Err(ArrowConvertError::InvalidValue(format!(
1883 "column '{}': expected Map value, got {:?}",
1884 col.name, other
1885 )));
1886 }
1887 }
1888 }
1889
1890 checked_string_offsets(&keys)?;
1892 checked_string_offsets(&values)?;
1893 let key_array = Arc::new(StringArray::from(keys)) as ArrayRef;
1894 let value_array = Arc::new(StringArray::from(values)) as ArrayRef;
1895
1896 let struct_fields = Fields::from(vec![
1897 Field::new("key", ArrowDataType::Utf8, false),
1898 Field::new("value", ArrowDataType::Utf8, true),
1899 ]);
1900
1901 let entries_array = StructArray::new(struct_fields.clone(), vec![key_array, value_array], None);
1902
1903 let map_field = Arc::new(Field::new(
1904 "entries",
1905 ArrowDataType::Struct(struct_fields),
1906 false,
1907 ));
1908 let offset_buffer = OffsetBuffer::new(offsets.into());
1909 let null_buffer = NullBuffer::from(null_bitmap);
1910
1911 Ok(Arc::new(MapArray::new(
1912 map_field,
1913 offset_buffer,
1914 entries_array,
1915 Some(null_buffer),
1916 false,
1917 )))
1918}
1919
1920#[cfg(test)]
1925mod tests {
1926 use super::*;
1927 use crate::query::{ColumnInfo, QueryRow};
1928 use crate::schema::CqlType;
1929 use crate::types::{DataType, Value};
1930 use crate::RowKey;
1931 use arrow::array::{Array, Int32Array};
1932
1933 fn col(name: &str, data_type: DataType, cql_type: Option<CqlType>) -> ColumnInfo {
1935 ColumnInfo {
1936 name: name.to_string(),
1937 data_type,
1938 nullable: true,
1939 position: 0,
1940 table_name: None,
1941 cql_type,
1942 }
1943 }
1944
1945 fn row_one(name: &str, value: Value) -> QueryRow {
1947 let mut values: HashMap<Arc<str>, Value> = HashMap::new();
1948 values.insert(name.into(), value);
1949 QueryRow {
1950 values,
1951 key: RowKey::new(Vec::new()),
1952 metadata: Default::default(),
1953 cell_metadata: None,
1954 }
1955 }
1956
1957 fn row_absent() -> QueryRow {
1959 QueryRow {
1960 values: HashMap::new(),
1961 key: RowKey::new(Vec::new()),
1962 metadata: Default::default(),
1963 cell_metadata: None,
1964 }
1965 }
1966
1967 fn is_invalid_value(res: Result<arrow::record_batch::RecordBatch, ArrowConvertError>) -> bool {
1968 matches!(res, Err(ArrowConvertError::InvalidValue(_)))
1969 }
1970
1971 #[test]
1974 fn typed_scalar_type_mismatch_is_error() {
1975 let columns = vec![col("d", DataType::Timestamp, Some(CqlType::Date))];
1976 let rows = vec![row_one("d", Value::Text("not-a-date".into()))];
1977 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
1978 }
1979
1980 #[test]
1983 fn flat_builder_type_mismatch_is_error() {
1984 let columns = vec![col("n", DataType::Integer, None)];
1985 let rows = vec![row_one("n", Value::Text("nope".into()))];
1986 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
1987 }
1988
1989 #[test]
1992 fn collection_expected_list_got_scalar_is_error() {
1993 let columns = vec![col(
1994 "l",
1995 DataType::List,
1996 Some(CqlType::List(Box::new(CqlType::Int))),
1997 )];
1998 let rows = vec![row_one("l", Value::Integer(5))];
1999 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2000 }
2001
2002 #[test]
2005 fn collection_mistyped_element_is_error() {
2006 let columns = vec![col(
2007 "l",
2008 DataType::List,
2009 Some(CqlType::List(Box::new(CqlType::Int))),
2010 )];
2011 let rows = vec![row_one(
2012 "l",
2013 Value::List(vec![Value::Integer(1), Value::Text("bad".into())]),
2014 )];
2015 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2016 }
2017
2018 #[test]
2021 fn collection_expected_map_got_scalar_is_error() {
2022 let columns = vec![col(
2023 "m",
2024 DataType::Map,
2025 Some(CqlType::Map(
2026 Box::new(CqlType::Text),
2027 Box::new(CqlType::Int),
2028 )),
2029 )];
2030 let rows = vec![row_one("m", Value::Integer(7))];
2031 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2032 }
2033
2034 #[test]
2037 fn null_and_absent_still_build_ok() {
2038 let columns = vec![col("n", DataType::Integer, None)];
2039 let rows = vec![row_one("n", Value::Null), row_absent()];
2040 let batch = rows_to_record_batch(&columns, &rows).expect("null/absent must build");
2041 assert_eq!(batch.num_rows(), 2);
2042 assert_eq!(batch.column(0).null_count(), 2);
2043 }
2044
2045 #[test]
2047 fn correctly_typed_value_builds_ok() {
2048 let columns = vec![col("n", DataType::Integer, None)];
2049 let rows = vec![row_one("n", Value::Integer(42))];
2050 let batch = rows_to_record_batch(&columns, &rows).expect("well-typed value must build");
2051 let arr = batch
2052 .column(0)
2053 .as_any()
2054 .downcast_ref::<Int32Array>()
2055 .expect("Int32Array");
2056 assert_eq!(arr.value(0), 42);
2057 assert_eq!(arr.null_count(), 0);
2058 }
2059
2060 #[test]
2064 fn decimal_scale_above_fixed_is_error() {
2065 let columns = vec![col("d", DataType::Blob, Some(CqlType::Decimal))];
2066 let unscaled = num_bigint::BigInt::from(123_456_789_012i64).to_signed_bytes_be();
2068 let rows = vec![row_one(
2069 "d",
2070 Value::Decimal {
2071 scale: 12,
2072 unscaled,
2073 },
2074 )];
2075 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2076 }
2077
2078 #[test]
2081 fn decimal_scale_within_fixed_builds_ok() {
2082 use arrow::array::Decimal128Array;
2083 let columns = vec![col("d", DataType::Blob, Some(CqlType::Decimal))];
2084 let unscaled = num_bigint::BigInt::from(123_456i64).to_signed_bytes_be();
2086 let rows = vec![row_one("d", Value::Decimal { scale: 3, unscaled })];
2087 let batch = rows_to_record_batch(&columns, &rows).expect("in-range decimal must build");
2088 let arr = batch
2089 .column(0)
2090 .as_any()
2091 .downcast_ref::<Decimal128Array>()
2092 .expect("Decimal128Array");
2093 assert_eq!(arr.value(0), 123_456_000_000i128);
2094 assert_eq!(arr.null_count(), 0);
2095 }
2096
2097 #[test]
2100 fn decimal_null_and_absent_still_null() {
2101 let columns = vec![col("d", DataType::Blob, Some(CqlType::Decimal))];
2102 let rows = vec![row_one("d", Value::Null), row_absent()];
2103 let batch = rows_to_record_batch(&columns, &rows).expect("null/absent decimal must build");
2104 assert_eq!(batch.num_rows(), 2);
2105 assert_eq!(batch.column(0).null_count(), 2);
2106 }
2107
2108 #[test]
2114 fn float32_column_accepts_wide_float_value() {
2115 let flat = vec![col("h", DataType::Float32, None)];
2117 let rows = vec![row_one("h", Value::Float(1.84f32 as f64))];
2118 let batch = rows_to_record_batch(&flat, &rows).expect("wide float must narrow, not error");
2119 let arr = batch
2120 .column(0)
2121 .as_any()
2122 .downcast_ref::<Float32Array>()
2123 .expect("Float32Array");
2124 assert_eq!(arr.value(0), 1.84f32);
2125 assert_eq!(arr.null_count(), 0);
2126
2127 let typed = vec![col("h", DataType::Float32, Some(CqlType::Float))];
2129 let rows = vec![row_one("h", Value::Float(1.84f32 as f64))];
2130 let batch =
2131 rows_to_record_batch(&typed, &rows).expect("wide float (typed) must narrow, not error");
2132 let arr = batch
2133 .column(0)
2134 .as_any()
2135 .downcast_ref::<Float32Array>()
2136 .expect("Float32Array");
2137 assert_eq!(arr.value(0), 1.84f32);
2138 }
2139
2140 #[test]
2143 fn tuple_expected_tuple_got_scalar_is_error() {
2144 let columns = vec![col(
2145 "t",
2146 DataType::Text,
2147 Some(CqlType::Tuple(vec![CqlType::Int, CqlType::Text])),
2148 )];
2149 let rows = vec![row_one("t", Value::Text("not-a-tuple".into()))];
2150 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2151 }
2152
2153 #[test]
2156 fn tuple_null_and_absent_still_build_ok() {
2157 let columns = vec![col(
2158 "t",
2159 DataType::Text,
2160 Some(CqlType::Tuple(vec![CqlType::Int, CqlType::Text])),
2161 )];
2162 let rows = vec![row_one("t", Value::Null), row_absent()];
2163 let batch = rows_to_record_batch(&columns, &rows).expect("null/absent tuple must build");
2164 assert_eq!(batch.num_rows(), 2);
2165 assert_eq!(batch.column(0).null_count(), 2);
2166 }
2167
2168 #[test]
2171 fn udt_expected_udt_got_scalar_is_error() {
2172 let columns = vec![col(
2173 "u",
2174 DataType::Text,
2175 Some(CqlType::Udt(
2176 "my_type".into(),
2177 vec![("a".into(), CqlType::Int), ("b".into(), CqlType::Text)],
2178 )),
2179 )];
2180 let rows = vec![row_one("u", Value::Integer(9))];
2181 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2182 }
2183
2184 #[test]
2187 fn udt_null_and_absent_still_build_ok() {
2188 let columns = vec![col(
2189 "u",
2190 DataType::Text,
2191 Some(CqlType::Udt(
2192 "my_type".into(),
2193 vec![("a".into(), CqlType::Int), ("b".into(), CqlType::Text)],
2194 )),
2195 )];
2196 let rows = vec![row_one("u", Value::Null), row_absent()];
2197 let batch = rows_to_record_batch(&columns, &rows).expect("null/absent UDT must build");
2198 assert_eq!(batch.num_rows(), 2);
2199 assert_eq!(batch.column(0).null_count(), 2);
2200 }
2201
2202 #[test]
2206 fn empty_field_udt_expected_udt_got_scalar_is_error() {
2207 let columns = vec![col(
2208 "u",
2209 DataType::Text,
2210 Some(CqlType::Udt("unresolved".into(), vec![])),
2211 )];
2212 let rows = vec![row_one("u", Value::Integer(9))];
2213 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2214 }
2215
2216 #[test]
2219 fn empty_field_tuple_expected_tuple_got_scalar_is_error() {
2220 let columns = vec![col("t", DataType::Text, Some(CqlType::Tuple(vec![])))];
2221 let rows = vec![row_one("t", Value::Text("nope".into()))];
2222 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2223 }
2224
2225 #[test]
2229 fn authoritative_text_column_type_mismatch_is_error() {
2230 for cql in [CqlType::Text, CqlType::Ascii, CqlType::Varchar] {
2231 let columns = vec![col("s", DataType::Text, Some(cql))];
2232 let rows = vec![row_one("s", Value::Integer(1))];
2233 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2234 }
2235 }
2236
2237 #[test]
2240 fn authoritative_text_column_rejects_json() {
2241 let columns = vec![col("s", DataType::Text, Some(CqlType::Text))];
2242 let rows = vec![row_one(
2243 "s",
2244 Value::Json(Box::new(serde_json::json!({"a": 1}))),
2245 )];
2246 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2247 }
2248
2249 #[test]
2253 fn frozen_wrapped_scalar_values_build_ok() {
2254 let text_cols = vec![col(
2256 "s",
2257 DataType::Text,
2258 Some(CqlType::Frozen(Box::new(CqlType::Text))),
2259 )];
2260 let text_rows = vec![row_one(
2261 "s",
2262 Value::Frozen(Box::new(Value::Text("hi".into()))),
2263 )];
2264 let batch =
2265 rows_to_record_batch(&text_cols, &text_rows).expect("frozen<text> value must build");
2266 let arr = batch
2267 .column(0)
2268 .as_any()
2269 .downcast_ref::<StringArray>()
2270 .expect("StringArray");
2271 assert_eq!(arr.value(0), "hi");
2272
2273 let date_cols = vec![col(
2275 "d",
2276 DataType::Integer,
2277 Some(CqlType::Frozen(Box::new(CqlType::Date))),
2278 )];
2279 let date_rows = vec![row_one("d", Value::Frozen(Box::new(Value::Date(19_000))))];
2280 let batch =
2281 rows_to_record_batch(&date_cols, &date_rows).expect("frozen<date> value must build");
2282 assert_eq!(batch.num_rows(), 1);
2283 assert_eq!(batch.column(0).null_count(), 0);
2284 }
2285
2286 #[test]
2289 fn authoritative_text_column_builds_ok() {
2290 let columns = vec![col("s", DataType::Text, Some(CqlType::Text))];
2291 let rows = vec![
2292 row_one("s", Value::Text("hi".into())),
2293 row_one("s", Value::Null),
2294 row_absent(),
2295 ];
2296 let batch = rows_to_record_batch(&columns, &rows).expect("well-typed text must build");
2297 let arr = batch
2298 .column(0)
2299 .as_any()
2300 .downcast_ref::<StringArray>()
2301 .expect("StringArray");
2302 assert_eq!(arr.value(0), "hi");
2303 assert_eq!(arr.null_count(), 2);
2304 }
2305
2306 #[test]
2309 fn authoritative_int_column_rejects_date() {
2310 let columns = vec![col("n", DataType::Integer, Some(CqlType::Int))];
2311 let rows = vec![row_one("n", Value::Date(19_000))];
2312 assert!(is_invalid_value(rows_to_record_batch(&columns, &rows)));
2313 }
2314
2315 #[test]
2318 fn opaque_int_column_accepts_date() {
2319 let columns = vec![col("n", DataType::Integer, None)];
2320 let rows = vec![row_one("n", Value::Date(19_000))];
2321 let batch = rows_to_record_batch(&columns, &rows).expect("opaque int accepts Date");
2322 let arr = batch
2323 .column(0)
2324 .as_any()
2325 .downcast_ref::<Int32Array>()
2326 .expect("Int32Array");
2327 assert_eq!(arr.value(0), 19_000);
2328 }
2329
2330 #[test]
2334 fn authoritative_bigint_counter_reject_mismatch() {
2335 let bigint_time = vec![col("b", DataType::BigInt, Some(CqlType::BigInt))];
2336 assert!(is_invalid_value(rows_to_record_batch(
2337 &bigint_time,
2338 &[row_one("b", Value::Time(123))]
2339 )));
2340
2341 let counter_time = vec![col("c", DataType::BigInt, Some(CqlType::Counter))];
2342 assert!(is_invalid_value(rows_to_record_batch(
2343 &counter_time,
2344 &[row_one("c", Value::Time(123))]
2345 )));
2346
2347 let bigint_counter = vec![col("b", DataType::BigInt, Some(CqlType::BigInt))];
2348 assert!(is_invalid_value(rows_to_record_batch(
2349 &bigint_counter,
2350 &[row_one("b", Value::Counter(7))]
2351 )));
2352 }
2353
2354 #[test]
2356 fn authoritative_counter_column_accepts_counter() {
2357 let columns = vec![col("c", DataType::BigInt, Some(CqlType::Counter))];
2358 let rows = vec![row_one("c", Value::Counter(42))];
2359 let batch = rows_to_record_batch(&columns, &rows).expect("counter accepts Counter");
2360 assert_eq!(batch.num_rows(), 1);
2361 assert_eq!(batch.column(0).null_count(), 0);
2362 }
2363
2364 #[test]
2371 fn checked_offset_past_i32_max_is_error() {
2372 assert_eq!(
2374 super::checked_offset(i32::MAX as usize).ok(),
2375 Some(i32::MAX)
2376 );
2377 assert!(matches!(
2379 super::checked_offset(i32::MAX as usize + 1),
2380 Err(ArrowConvertError::InvalidValue(_))
2381 ));
2382 }
2383
2384 #[test]
2387 fn checked_offset_normal_sizes_are_identity() {
2388 assert_eq!(super::checked_offset(0).ok(), Some(0));
2389 assert_eq!(super::checked_offset(1).ok(), Some(1));
2390 assert_eq!(super::checked_offset(1_000_000).ok(), Some(1_000_000));
2391 }
2392
2393 #[test]
2401 fn checked_value_bytes_past_i32_max_is_error() {
2402 assert!(super::checked_value_bytes(i32::MAX as usize).is_ok());
2404 assert!(matches!(
2406 super::checked_value_bytes(i32::MAX as usize + 1),
2407 Err(ArrowConvertError::InvalidValue(_))
2408 ));
2409 assert!(super::checked_value_bytes(0).is_ok());
2410 assert!(super::checked_value_bytes(1_000_000).is_ok());
2411 }
2412
2413 #[test]
2421 fn typed_blob_builder_over_i32_max_fails_closed_without_2gib_clone() {
2422 const CHUNK: usize = 16 * 1024 * 1024; const N: usize = 128; let big = Value::blob(vec![0u8; CHUNK]);
2425 let refs: Vec<Option<&Value>> = (0..N).map(|_| Some(&big)).collect();
2426 let err = super::build_typed_value_array(&CqlType::Blob, &refs);
2427 assert!(
2428 matches!(err, Err(ArrowConvertError::InvalidValue(_))),
2429 "Blob arm must fail closed at the i32 offset ceiling"
2430 );
2431 }
2432
2433 #[test]
2438 fn typed_text_builder_over_i32_max_fails_closed_without_2gib_clone() {
2439 const CHUNK: usize = 16 * 1024 * 1024; const N: usize = 128; let big = Value::text("a".repeat(CHUNK));
2442 let refs: Vec<Option<&Value>> = (0..N).map(|_| Some(&big)).collect();
2443 let err = super::build_typed_value_array(&CqlType::Text, &refs);
2444 assert!(
2445 matches!(err, Err(ArrowConvertError::InvalidValue(_))),
2446 "Text arm must fail closed at the i32 offset ceiling"
2447 );
2448 }
2449
2450 #[test]
2461 fn opaque_text_fallback_over_i32_max_fails_closed_without_2gib_clone() {
2462 use std::borrow::Cow;
2463 const CHUNK: usize = 16 * 1024 * 1024; const N: usize = 128; let big = "a".repeat(CHUNK);
2466 let refs: Vec<Option<Cow<str>>> =
2469 (0..N).map(|_| Some(Cow::Borrowed(big.as_str()))).collect();
2470 let total: usize = refs.iter().flatten().map(|s| s.len()).sum();
2471 assert_eq!(total, i32::MAX as usize + 1, "test must cross i32::MAX");
2472 assert!(
2473 matches!(
2474 super::checked_string_offsets(&refs),
2475 Err(ArrowConvertError::InvalidValue(_))
2476 ),
2477 "opaque untyped Text fallback must fail closed at the i32 offset ceiling"
2478 );
2479 }
2480
2481 #[test]
2485 fn opaque_text_fallback_preserves_raw_text_verbatim() {
2486 use arrow::array::StringArray;
2487 let cols = vec![col("o", DataType::Text, None)];
2488 let rows = vec![
2489 row_one("o", Value::Text("verbatim".into())),
2490 row_one("o", Value::Null),
2491 ];
2492 let batch = rows_to_record_batch(&cols, &rows).expect("opaque text must build");
2493 let arr = batch
2494 .column(0)
2495 .as_any()
2496 .downcast_ref::<StringArray>()
2497 .expect("Utf8 array");
2498 assert_eq!(arr.value(0), "verbatim");
2499 assert!(arr.is_null(1));
2500 }
2501
2502 #[test]
2509 fn scalar_binary_cumulative_bytes_over_i32_max_is_typed_error() {
2510 const CHUNK: usize = 16 * 1024 * 1024; const N: usize = 128; let buf = vec![0u8; CHUNK];
2513 let refs: Vec<Option<&[u8]>> = (0..N).map(|_| Some(buf.as_slice())).collect();
2514 let total: usize = refs.iter().flatten().map(|b| b.len()).sum();
2516 assert_eq!(total, i32::MAX as usize + 1, "test must cross i32::MAX");
2517 assert!(matches!(
2519 super::checked_binary_offsets(&refs),
2520 Err(ArrowConvertError::InvalidValue(_))
2521 ));
2522 }
2523
2524 #[test]
2531 fn scalar_string_cumulative_bytes_over_i32_max_is_typed_error() {
2532 let ok = vec![Some("a".to_string()), None, Some("bc".to_string())];
2538 assert!(super::checked_string_offsets(&ok).is_ok());
2539 assert!(matches!(
2541 super::checked_value_bytes(i32::MAX as usize + 42),
2542 Err(ArrowConvertError::InvalidValue(_))
2543 ));
2544 }
2545
2546 #[test]
2550 fn normal_scalar_text_and_blob_still_build_through_byte_guard() {
2551 let text_cols = vec![col("t", DataType::Text, Some(CqlType::Text))];
2552 let text_rows = vec![
2553 row_one("t", Value::Text("hello".into())),
2554 row_one("t", Value::Null),
2555 ];
2556 let batch = rows_to_record_batch(&text_cols, &text_rows).expect("text must build");
2557 assert_eq!(batch.num_rows(), 2);
2558
2559 let blob_cols = vec![col("b", DataType::Blob, Some(CqlType::Blob))];
2560 let blob_rows = vec![row_one("b", Value::blob(vec![1, 2, 3, 4]))];
2561 let batch = rows_to_record_batch(&blob_cols, &blob_rows).expect("blob must build");
2562 assert_eq!(batch.num_rows(), 1);
2563 }
2564
2565 #[test]
2568 fn normal_collections_still_build_through_checked_offsets() {
2569 let list_cols = vec![col(
2570 "l",
2571 DataType::List,
2572 Some(CqlType::List(Box::new(CqlType::Int))),
2573 )];
2574 let list_rows = vec![
2575 row_one("l", Value::List(vec![Value::Integer(1), Value::Integer(2)])),
2576 row_one("l", Value::Null),
2577 ];
2578 let batch = rows_to_record_batch(&list_cols, &list_rows).expect("list must build");
2579 assert_eq!(batch.num_rows(), 2);
2580
2581 let map_cols = vec![col(
2582 "m",
2583 DataType::Map,
2584 Some(CqlType::Map(
2585 Box::new(CqlType::Text),
2586 Box::new(CqlType::Int),
2587 )),
2588 )];
2589 let map_rows = vec![row_one(
2590 "m",
2591 Value::Map(vec![(Value::Text("k".into()), Value::Integer(9))]),
2592 )];
2593 let batch = rows_to_record_batch(&map_cols, &map_rows).expect("map must build");
2594 assert_eq!(batch.num_rows(), 1);
2595 }
2596}