lance_nd/
cast.rs

1//! Cast new rows to the types of the table, see [`Casting`].
2
3use std::sync::Arc;
4
5use arrow::array::{
6    Array, ArrayRef, AsArray, ListArray, RecordBatch, RecordBatchOptions, StructArray,
7};
8use arrow::compute::{CastOptions, cast_with_options};
9use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema};
10
11use crate::error::Result;
12
13/// How a write treats a value or attribute type that differs from the table.
14///
15/// The rules follow the `casting` argument of numpy. The table type wins:
16/// a cast changes the new rows, never the table. Only [`Casting::Promote`]
17/// can widen the table.
18#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
19pub enum Casting {
20    /// Reject the write.
21    #[default]
22    No,
23    /// Cast when no value can change, such as `float32` to `float64` or
24    /// `uint16` to `int32`. Otherwise reject the write.
25    Safe,
26    /// Also cast within a kind or to a later kind in the order bool,
27    /// unsigned, signed, float, such as `float64` to `float32`. An integer
28    /// that does not fit fails the write.
29    SameKind,
30    /// As numpy type promotion: cast as [`Casting::Safe`] when the new
31    /// values fit the table type. Otherwise widen the variable in the table
32    /// to the smallest type that holds both, see [`crate::widen()`], and
33    /// cast the new values. Attributes follow [`Casting::Safe`].
34    Promote,
35}
36
37/// `batch` with each field cast to the type of the same field in `target`.
38/// A field that the rule does not allow keeps its type, so the schema check
39/// rejects it later.
40pub(crate) fn cast_to(
41    batch: RecordBatch,
42    target: &Schema,
43    casting: Casting,
44) -> Result<RecordBatch> {
45    if casting == Casting::No {
46        return Ok(batch);
47    }
48    let mut fields = Vec::with_capacity(batch.num_columns());
49    let mut columns = Vec::with_capacity(batch.num_columns());
50    for (field, column) in batch.schema().fields().iter().zip(batch.columns()) {
51        let (field, column) = match target.field_with_name(field.name()) {
52            Ok(target) => cast_field(field, column, target, casting)?,
53            Err(_) => (field.clone(), column.clone()),
54        };
55        fields.push(field);
56        columns.push(column);
57    }
58    let schema = Schema::new_with_metadata(fields, batch.schema().metadata().clone());
59    let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
60    Ok(RecordBatch::try_new_with_options(
61        Arc::new(schema),
62        columns,
63        &options,
64    )?)
65}
66
67fn cast_field(
68    field: &FieldRef,
69    array: &ArrayRef,
70    target: &Field,
71    casting: Casting,
72) -> Result<(FieldRef, ArrayRef)> {
73    let typed = |data_type: DataType| Arc::new(field.as_ref().clone().with_data_type(data_type));
74    match (field.data_type(), target.data_type()) {
75        (from, to) if from == to => Ok((field.clone(), array.clone())),
76        (DataType::Struct(children), DataType::Struct(targets)) => {
77            let source = array.as_struct();
78            let mut fields = Vec::with_capacity(children.len());
79            let mut columns = Vec::with_capacity(children.len());
80            for (child, column) in children.iter().zip(source.columns()) {
81                let (child, column) = match targets.find(child.name()) {
82                    Some((_, target)) => cast_field(child, column, target, casting)?,
83                    None => (child.clone(), column.clone()),
84                };
85                fields.push(child);
86                columns.push(column);
87            }
88            let fields = Fields::from(fields);
89            let array = StructArray::try_new(fields.clone(), columns, source.nulls().cloned())?;
90            Ok((typed(DataType::Struct(fields)), Arc::new(array)))
91        }
92        (DataType::List(item), DataType::List(target)) => {
93            let source = array.as_list::<i32>();
94            let (item, values) = cast_field(item, source.values(), target, casting)?;
95            let array = ListArray::try_new(
96                item.clone(),
97                source.offsets().clone(),
98                values,
99                source.nulls().cloned(),
100            )?;
101            Ok((typed(DataType::List(item)), Arc::new(array)))
102        }
103        (from, to) if allowed(from, to, casting) => {
104            // Not `safe`: an integer that does not fit fails instead of becoming null.
105            let options = CastOptions {
106                safe: false,
107                ..Default::default()
108            };
109            let array = cast_with_options(array, to, &options)?;
110            Ok((typed(to.clone()), array))
111        }
112        _ => Ok((field.clone(), array.clone())),
113    }
114}
115
116/// `true` when `casting` allows a cast of values from `from` to `to`.
117pub(crate) fn allowed(from: &DataType, to: &DataType, casting: Casting) -> bool {
118    let text =
119        |t: &DataType| matches!(t, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View);
120    match casting {
121        Casting::No => false,
122        _ if text(from) && text(to) => true,
123        Casting::Safe | Casting::Promote => is_safe(from, to),
124        Casting::SameKind => {
125            is_safe(from, to)
126                || kind(from)
127                    .zip(kind(to))
128                    .is_some_and(|(from, to)| from <= to)
129        }
130    }
131}
132
133/// The smallest type that holds each value of `a` and of `b`, as numpy
134/// `promote_types`, or `None` when no number type holds both.
135pub(crate) fn promote_type(a: &DataType, b: &DataType) -> Option<DataType> {
136    use DataType::*;
137    if a == b {
138        return Some(a.clone());
139    }
140    let holds = |from: &DataType, to: &DataType| from == to || is_safe(from, to);
141    [
142        UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64, Float16, Float32, Float64,
143    ]
144    .into_iter()
145    .find(|to| holds(a, to) && holds(b, to))
146}
147
148/// The numpy kind order: bool, unsigned, signed, float.
149fn kind(data_type: &DataType) -> Option<u8> {
150    match data_type {
151        DataType::Boolean => Some(0),
152        t if t.is_unsigned_integer() => Some(1),
153        t if t.is_signed_integer() => Some(2),
154        t if t.is_floating() => Some(3),
155        _ => None,
156    }
157}
158
159/// `true` when each value of `from` has the same value in `to`, as numpy
160/// `can_cast(from, to, "safe")`.
161fn is_safe(from: &DataType, to: &DataType) -> bool {
162    let bits = |t: &DataType| t.primitive_width().map_or(0, |bytes| bytes * 8);
163    match (kind(from), kind(to)) {
164        (Some(0), Some(_)) => true,
165        (Some(f), Some(t)) if f == t => bits(to) > bits(from),
166        // An unsigned integer fits a signed integer with more bits.
167        (Some(1), Some(2)) => bits(to) > bits(from),
168        // As numpy: up to 8 bits fit float16, 16 bits float32, all float64.
169        (Some(1 | 2), Some(3)) => match bits(to) {
170            16 => bits(from) <= 8,
171            32 => bits(from) <= 16,
172            _ => true,
173        },
174        _ => false,
175    }
176}
lance_nd/cast.rs

lance_nd/
cast.rs
