vgi_rpc/
arrow_type.rs

1//! `VgiArrow` — the bridge between idiomatic Rust types and Arrow.
2//!
3//! The proc-macro layer (`vgi-rpc-macros`) generates code that maps each
4//! RPC method parameter and return value through this trait so user
5//! handler signatures stay free of Arrow types. Cross-language wire
6//! compatibility with Python `vgi_rpc` is preserved by mirroring the
7//! Arrow `DataType` choices Python's `ArrowSerializableDataclass` uses.
8//!
9//! # Implementing for your own types
10//!
11//! Use `#[derive(VgiArrow)]` from `vgi-rpc-macros` for plain structs.
12//! Hand-implement only when the wire format must diverge from
13//! Python-canonical defaults.
14
15use std::sync::Arc;
16
17use arrow_array::{
18    builder::BinaryBuilder, Array, ArrayRef, BinaryArray, BooleanArray, FixedSizeBinaryArray,
19    Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeBinaryArray,
20    LargeStringArray, ListArray, MapArray, StringArray, UInt16Array, UInt32Array, UInt64Array,
21    UInt8Array,
22};
23use arrow_buffer::{Buffer, OffsetBuffer};
24use arrow_schema::{DataType, Field};
25
26use crate::errors::{Result, RpcError};
27
28/// Round-trip a Rust value through a single Arrow column.
29///
30/// `arrow_data_type()` returns the column's `DataType`; `nullable()`
31/// indicates whether the column accepts nulls (set by `Option<T>`).
32/// `read()` extracts a value at `idx`; `build_singleton()` builds a
33/// 1-row array carrying `value`.
34///
35/// All builtin scalar / collection / option impls in this module are
36/// `Send + Sync` and allocate at most once per call.
37pub trait VgiArrow: Sized {
38    /// The Arrow `DataType` carrying values of this Rust type.
39    fn arrow_data_type() -> DataType;
40
41    /// Whether the column should be flagged nullable. The default is
42    /// `false`; the `Option<T>` blanket impl returns `true`.
43    fn nullable() -> bool {
44        false
45    }
46
47    /// Wire-format type name surfaced via `__describe__` metadata.
48    /// Mirrors Python: `"str"`, `"int"`, `"list[int]"`, `"int | None"`.
49    fn describe_name() -> String;
50
51    /// Pull this value out of `arr` at row `idx`. Errors with a
52    /// `RpcError::type_error` if `arr`'s concrete type doesn't match
53    /// `Self::arrow_data_type()`.
54    fn read(arr: &dyn Array, idx: usize) -> Result<Self>;
55
56    /// Build a 1-row `ArrayRef` containing `value`.
57    fn build_singleton(value: Self) -> Result<ArrayRef>;
58}
59
60/// Helper: typed downcast or `RpcError::type_error("expected …")`.
61fn as_array<'a, A: Array + 'static>(arr: &'a dyn Array, expected: &str) -> Result<&'a A> {
62    arr.as_any()
63        .downcast_ref::<A>()
64        .ok_or_else(|| RpcError::type_error(format!("expected {expected} array")))
65}
66
67// ---------------------------------------------------------------------------
68// Scalars
69// ---------------------------------------------------------------------------
70
71impl VgiArrow for String {
72    fn arrow_data_type() -> DataType {
73        DataType::Utf8
74    }
75    fn describe_name() -> String {
76        "str".into()
77    }
78    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
79        // Accept Utf8 directly, plus DictionaryArray<Int16|Int32, Utf8>
80        // — Python's enum-typed params arrive as dict-encoded strings.
81        if let Some(s) = arr.as_any().downcast_ref::<StringArray>() {
82            return Ok(s.value(idx).to_string());
83        }
84        if let Some(d) = arr
85            .as_any()
86            .downcast_ref::<arrow_array::DictionaryArray<arrow_array::types::Int16Type>>()
87        {
88            let key = d.keys().value(idx);
89            let values = as_array::<StringArray>(d.values().as_ref(), "Utf8 (dict values)")?;
90            return Ok(values.value(key as usize).to_string());
91        }
92        if let Some(d) = arr
93            .as_any()
94            .downcast_ref::<arrow_array::DictionaryArray<arrow_array::types::Int32Type>>()
95        {
96            let key = d.keys().value(idx);
97            let values = as_array::<StringArray>(d.values().as_ref(), "Utf8 (dict values)")?;
98            return Ok(values.value(key as usize).to_string());
99        }
100        Err(RpcError::type_error(
101            "expected Utf8 (or DictionaryArray<Int16|Int32, Utf8>) array",
102        ))
103    }
104    fn build_singleton(value: Self) -> Result<ArrayRef> {
105        Ok(Arc::new(StringArray::from(vec![value])))
106    }
107}
108
109impl VgiArrow for i64 {
110    fn arrow_data_type() -> DataType {
111        DataType::Int64
112    }
113    fn describe_name() -> String {
114        "int".into()
115    }
116    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
117        if let Some(a) = arr.as_any().downcast_ref::<Int64Array>() {
118            return Ok(a.value(idx));
119        }
120        if let Some(a) = arr.as_any().downcast_ref::<Int32Array>() {
121            return Ok(a.value(idx) as i64);
122        }
123        Err(RpcError::type_error("expected Int64/Int32 array"))
124    }
125    fn build_singleton(value: Self) -> Result<ArrayRef> {
126        Ok(Arc::new(Int64Array::from(vec![value])))
127    }
128}
129
130impl VgiArrow for i32 {
131    fn arrow_data_type() -> DataType {
132        DataType::Int32
133    }
134    fn describe_name() -> String {
135        "int".into()
136    }
137    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
138        if let Some(a) = arr.as_any().downcast_ref::<Int32Array>() {
139            return Ok(a.value(idx));
140        }
141        if let Some(a) = arr.as_any().downcast_ref::<Int64Array>() {
142            return Ok(a.value(idx) as i32);
143        }
144        Err(RpcError::type_error("expected Int32/Int64 array"))
145    }
146    fn build_singleton(value: Self) -> Result<ArrayRef> {
147        Ok(Arc::new(Int32Array::from(vec![value])))
148    }
149}
150
151// Smaller / unsigned integer widths. Python's `Annotated[int, ArrowType(pa.int8())]`
152// shows up on the wire as `Int8` etc.; we expose them as the matching
153// Rust primitive so user signatures are natural.
154macro_rules! impl_int_vgi {
155    ($t:ty, $arr:ty, $dt:expr) => {
156        impl VgiArrow for $t {
157            fn arrow_data_type() -> DataType {
158                $dt
159            }
160            fn describe_name() -> String {
161                "int".into()
162            }
163            fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
164                Ok(as_array::<$arr>(arr, stringify!($t))?.value(idx))
165            }
166            fn build_singleton(value: Self) -> Result<ArrayRef> {
167                Ok(Arc::new(<$arr>::from(vec![value])))
168            }
169        }
170    };
171}
172impl_int_vgi!(i8, Int8Array, DataType::Int8);
173impl_int_vgi!(i16, Int16Array, DataType::Int16);
174impl_int_vgi!(u8, UInt8Array, DataType::UInt8);
175impl_int_vgi!(u16, UInt16Array, DataType::UInt16);
176impl_int_vgi!(u32, UInt32Array, DataType::UInt32);
177impl_int_vgi!(u64, UInt64Array, DataType::UInt64);
178
179impl VgiArrow for f64 {
180    fn arrow_data_type() -> DataType {
181        DataType::Float64
182    }
183    fn describe_name() -> String {
184        "float".into()
185    }
186    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
187        if let Some(a) = arr.as_any().downcast_ref::<Float64Array>() {
188            return Ok(a.value(idx));
189        }
190        if let Some(a) = arr.as_any().downcast_ref::<Float32Array>() {
191            return Ok(a.value(idx) as f64);
192        }
193        Err(RpcError::type_error("expected Float64/Float32 array"))
194    }
195    fn build_singleton(value: Self) -> Result<ArrayRef> {
196        Ok(Arc::new(Float64Array::from(vec![value])))
197    }
198}
199
200impl VgiArrow for f32 {
201    fn arrow_data_type() -> DataType {
202        DataType::Float32
203    }
204    fn describe_name() -> String {
205        "float".into()
206    }
207    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
208        if let Some(a) = arr.as_any().downcast_ref::<Float32Array>() {
209            return Ok(a.value(idx));
210        }
211        if let Some(a) = arr.as_any().downcast_ref::<Float64Array>() {
212            return Ok(a.value(idx) as f32);
213        }
214        Err(RpcError::type_error("expected Float32/Float64 array"))
215    }
216    fn build_singleton(value: Self) -> Result<ArrayRef> {
217        Ok(Arc::new(Float32Array::from(vec![value])))
218    }
219}
220
221impl VgiArrow for bool {
222    fn arrow_data_type() -> DataType {
223        DataType::Boolean
224    }
225    fn describe_name() -> String {
226        "bool".into()
227    }
228    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
229        Ok(as_array::<BooleanArray>(arr, "Boolean")?.value(idx))
230    }
231    fn build_singleton(value: Self) -> Result<ArrayRef> {
232        Ok(Arc::new(BooleanArray::from(vec![value])))
233    }
234}
235
236// ---------------------------------------------------------------------------
237// Bytes (Binary), kept distinct from Vec<u8> via a newtype-style wrapper.
238// ---------------------------------------------------------------------------
239
240/// Newtype indicating a `Vec<u8>` should be carried as Arrow `Binary`,
241/// not `List<UInt8>`. Use this in handler signatures where the wire
242/// type should be `bytes` rather than a list of bytes.
243///
244/// `#[derive(VgiArrow)]` does not auto-pick between the two — there is
245/// no idiomatic Rust signal that `Vec<u8>` in a struct field means
246/// "blob" rather than "byte list", so users opt in via this wrapper.
247#[derive(Clone, Debug, PartialEq, Eq)]
248pub struct Bytes(pub Vec<u8>);
249
250impl From<Vec<u8>> for Bytes {
251    fn from(v: Vec<u8>) -> Self {
252        Self(v)
253    }
254}
255
256impl From<Bytes> for Vec<u8> {
257    fn from(b: Bytes) -> Self {
258        b.0
259    }
260}
261
262impl VgiArrow for Bytes {
263    fn arrow_data_type() -> DataType {
264        DataType::Binary
265    }
266    fn describe_name() -> String {
267        "bytes".into()
268    }
269    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
270        Ok(Bytes(
271            as_array::<BinaryArray>(arr, "Binary")?.value(idx).to_vec(),
272        ))
273    }
274    fn build_singleton(value: Self) -> Result<ArrayRef> {
275        let mut b = BinaryBuilder::new();
276        b.append_value(value.0);
277        Ok(Arc::new(b.finish()))
278    }
279}
280
281// ---------------------------------------------------------------------------
282// Wide-binary / wide-string newtypes.
283// ---------------------------------------------------------------------------
284
285/// `LargeUtf8` (64-bit-offset string array) wire type. Stored as a
286/// regular `String` in user code; the wrapper just tags the wire shape.
287#[derive(Clone, Debug, PartialEq, Eq)]
288pub struct LargeString(pub String);
289
290impl From<String> for LargeString {
291    fn from(s: String) -> Self {
292        Self(s)
293    }
294}
295impl From<LargeString> for String {
296    fn from(s: LargeString) -> Self {
297        s.0
298    }
299}
300
301impl VgiArrow for LargeString {
302    fn arrow_data_type() -> DataType {
303        DataType::LargeUtf8
304    }
305    fn describe_name() -> String {
306        "str".into()
307    }
308    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
309        Ok(LargeString(
310            as_array::<LargeStringArray>(arr, "LargeUtf8")?
311                .value(idx)
312                .to_string(),
313        ))
314    }
315    fn build_singleton(value: Self) -> Result<ArrayRef> {
316        Ok(Arc::new(LargeStringArray::from(vec![value.0])))
317    }
318}
319
320/// `LargeBinary` wire type. See [`LargeString`].
321#[derive(Clone, Debug, PartialEq, Eq)]
322pub struct LargeBytes(pub Vec<u8>);
323
324impl From<Vec<u8>> for LargeBytes {
325    fn from(v: Vec<u8>) -> Self {
326        Self(v)
327    }
328}
329impl From<LargeBytes> for Vec<u8> {
330    fn from(b: LargeBytes) -> Self {
331        b.0
332    }
333}
334
335impl VgiArrow for LargeBytes {
336    fn arrow_data_type() -> DataType {
337        DataType::LargeBinary
338    }
339    fn describe_name() -> String {
340        "bytes".into()
341    }
342    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
343        Ok(LargeBytes(
344            as_array::<LargeBinaryArray>(arr, "LargeBinary")?
345                .value(idx)
346                .to_vec(),
347        ))
348    }
349    fn build_singleton(value: Self) -> Result<ArrayRef> {
350        let len = i64::try_from(value.0.len())
351            .map_err(|_| RpcError::value_error("large_binary value exceeds i64 offsets"))?;
352        let offsets = OffsetBuffer::new(vec![0_i64, len].into());
353        // Take ownership of the Vec allocation instead of copying the payload
354        // into a builder-owned Arrow buffer.
355        let arr = LargeBinaryArray::new(offsets, Buffer::from_vec(value.0), None);
356        Ok(Arc::new(arr))
357    }
358}
359
360/// Zero-copy `LargeBinary` wire value for payload-oriented handlers.
361///
362/// Unlike [`LargeBytes`], `read` retains a slice of the inbound Arrow buffer
363/// and `build_singleton` transfers that same buffer into the response array.
364/// Cloning this type is also zero-copy. Use [`AsRef::as_ref`] or dereference it
365/// to inspect the bytes, and [`to_vec`](Self::to_vec) only when owned mutable
366/// storage is actually required.
367#[derive(Clone, Debug, PartialEq, Eq)]
368pub struct LargeBytesBuffer(pub Buffer);
369
370impl LargeBytesBuffer {
371    /// Copy the payload into an owned `Vec<u8>`.
372    pub fn to_vec(&self) -> Vec<u8> {
373        self.0.as_slice().to_vec()
374    }
375}
376
377impl From<Vec<u8>> for LargeBytesBuffer {
378    fn from(value: Vec<u8>) -> Self {
379        Self(Buffer::from_vec(value))
380    }
381}
382
383impl From<Buffer> for LargeBytesBuffer {
384    fn from(value: Buffer) -> Self {
385        Self(value)
386    }
387}
388
389impl AsRef<[u8]> for LargeBytesBuffer {
390    fn as_ref(&self) -> &[u8] {
391        self.0.as_slice()
392    }
393}
394
395impl std::ops::Deref for LargeBytesBuffer {
396    type Target = [u8];
397
398    fn deref(&self) -> &Self::Target {
399        self.as_ref()
400    }
401}
402
403impl VgiArrow for LargeBytesBuffer {
404    fn arrow_data_type() -> DataType {
405        DataType::LargeBinary
406    }
407
408    fn describe_name() -> String {
409        "bytes".into()
410    }
411
412    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
413        let arr = as_array::<LargeBinaryArray>(arr, "LargeBinary")?;
414        let offsets = arr.value_offsets();
415        let start = usize::try_from(offsets[idx])
416            .map_err(|_| RpcError::value_error("large_binary offset is negative"))?;
417        let end = usize::try_from(offsets[idx + 1])
418            .map_err(|_| RpcError::value_error("large_binary offset is negative"))?;
419        Ok(Self(arr.values().slice_with_length(start, end - start)))
420    }
421
422    fn build_singleton(value: Self) -> Result<ArrayRef> {
423        let len = i64::try_from(value.0.len())
424            .map_err(|_| RpcError::value_error("large_binary value exceeds i64 offsets"))?;
425        let offsets = OffsetBuffer::new(vec![0_i64, len].into());
426        Ok(Arc::new(LargeBinaryArray::new(offsets, value.0, None)))
427    }
428}
429
430/// `FixedSizeBinary(N)` wire type carried as `[u8; N]`. The const
431/// generic encodes the width so the schema is fully determined.
432#[derive(Clone, Debug, PartialEq, Eq)]
433pub struct FixedBinary<const N: usize>(pub [u8; N]);
434
435impl<const N: usize> From<[u8; N]> for FixedBinary<N> {
436    fn from(b: [u8; N]) -> Self {
437        Self(b)
438    }
439}
440
441impl<const N: usize> VgiArrow for FixedBinary<N> {
442    fn arrow_data_type() -> DataType {
443        DataType::FixedSizeBinary(N as i32)
444    }
445    fn describe_name() -> String {
446        "bytes".into()
447    }
448    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
449        let a = as_array::<FixedSizeBinaryArray>(arr, "FixedSizeBinary")?;
450        let raw = a.value(idx);
451        if raw.len() != N {
452            return Err(RpcError::type_error(format!(
453                "FixedSizeBinary width mismatch: expected {N}, got {}",
454                raw.len()
455            )));
456        }
457        let mut out = [0u8; N];
458        out.copy_from_slice(raw);
459        Ok(FixedBinary(out))
460    }
461    fn build_singleton(value: Self) -> Result<ArrayRef> {
462        let arr = FixedSizeBinaryArray::try_from_iter([value.0.as_slice()].into_iter())
463            .map_err(RpcError::from)?;
464        Ok(Arc::new(arr))
465    }
466}
467
468/// Dictionary-encoded `Utf8` (`Dictionary(Int16, Utf8)`) wire type.
469/// On the user-facing side it's just a `String`; the newtype controls
470/// the schema choice.
471#[derive(Clone, Debug, PartialEq, Eq)]
472pub struct DictString(pub String);
473
474impl From<String> for DictString {
475    fn from(s: String) -> Self {
476        Self(s)
477    }
478}
479impl From<DictString> for String {
480    fn from(s: DictString) -> Self {
481        s.0
482    }
483}
484
485impl VgiArrow for DictString {
486    fn arrow_data_type() -> DataType {
487        DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8))
488    }
489    fn describe_name() -> String {
490        "str".into()
491    }
492    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
493        // Reuse the `String` reader which already accepts both plain
494        // Utf8 and DictionaryArray<Int16|Int32, Utf8>.
495        Ok(DictString(<String as VgiArrow>::read(arr, idx)?))
496    }
497    fn build_singleton(value: Self) -> Result<ArrayRef> {
498        use arrow_array::builder::StringDictionaryBuilder;
499        use arrow_array::types::Int16Type;
500        let mut b = StringDictionaryBuilder::<Int16Type>::new();
501        b.append_value(&value.0);
502        Ok(Arc::new(b.finish()))
503    }
504}
505
506// ---------------------------------------------------------------------------
507// Date / time / duration / decimal — chrono + rust_decimal backed.
508// ---------------------------------------------------------------------------
509
510use arrow_array::{
511    Date32Array, Decimal128Array, DurationMicrosecondArray, Time64MicrosecondArray,
512    TimestampMicrosecondArray,
513};
514
515const DATE32_EPOCH: chrono::NaiveDate = match chrono::NaiveDate::from_ymd_opt(1970, 1, 1) {
516    Some(d) => d,
517    None => panic!("epoch"),
518};
519
520impl VgiArrow for chrono::NaiveDate {
521    fn arrow_data_type() -> DataType {
522        DataType::Date32
523    }
524    fn describe_name() -> String {
525        "date".into()
526    }
527    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
528        let days = as_array::<Date32Array>(arr, "Date32")?.value(idx);
529        DATE32_EPOCH
530            .checked_add_signed(chrono::Duration::days(days as i64))
531            .ok_or_else(|| RpcError::value_error("date32 out of range"))
532    }
533    fn build_singleton(value: Self) -> Result<ArrayRef> {
534        let days = (value - DATE32_EPOCH).num_days() as i32;
535        Ok(Arc::new(Date32Array::from(vec![days])))
536    }
537}
538
539impl VgiArrow for chrono::NaiveDateTime {
540    fn arrow_data_type() -> DataType {
541        DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None)
542    }
543    fn describe_name() -> String {
544        "datetime".into()
545    }
546    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
547        let micros = as_array::<TimestampMicrosecondArray>(arr, "Timestamp(us)")?.value(idx);
548        chrono::DateTime::from_timestamp_micros(micros)
549            .map(|dt| dt.naive_utc())
550            .ok_or_else(|| RpcError::value_error("timestamp out of range"))
551    }
552    fn build_singleton(value: Self) -> Result<ArrayRef> {
553        let micros = value.and_utc().timestamp_micros();
554        Ok(Arc::new(TimestampMicrosecondArray::from(vec![micros])))
555    }
556}
557
558/// UTC-tagged timestamp wire type (`Timestamp(us, tz="UTC")`). User
559/// holds a `chrono::DateTime<Utc>`.
560#[derive(Clone, Debug, PartialEq, Eq)]
561pub struct UtcTimestamp(pub chrono::DateTime<chrono::Utc>);
562
563impl From<chrono::DateTime<chrono::Utc>> for UtcTimestamp {
564    fn from(d: chrono::DateTime<chrono::Utc>) -> Self {
565        Self(d)
566    }
567}
568
569impl VgiArrow for UtcTimestamp {
570    fn arrow_data_type() -> DataType {
571        DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("UTC".into()))
572    }
573    fn describe_name() -> String {
574        "datetime".into()
575    }
576    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
577        let micros = as_array::<TimestampMicrosecondArray>(arr, "Timestamp(us, UTC)")?.value(idx);
578        chrono::DateTime::<chrono::Utc>::from_timestamp_micros(micros)
579            .map(UtcTimestamp)
580            .ok_or_else(|| RpcError::value_error("UTC timestamp out of range"))
581    }
582    fn build_singleton(value: Self) -> Result<ArrayRef> {
583        let micros = value.0.timestamp_micros();
584        let arr = TimestampMicrosecondArray::from(vec![micros]).with_timezone("UTC");
585        Ok(Arc::new(arr))
586    }
587}
588
589impl VgiArrow for chrono::NaiveTime {
590    fn arrow_data_type() -> DataType {
591        DataType::Time64(arrow_schema::TimeUnit::Microsecond)
592    }
593    fn describe_name() -> String {
594        "time".into()
595    }
596    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
597        let micros = as_array::<Time64MicrosecondArray>(arr, "Time64(us)")?.value(idx);
598        let secs = (micros / 1_000_000) as u32;
599        let nanos = ((micros % 1_000_000) * 1_000) as u32;
600        chrono::NaiveTime::from_num_seconds_from_midnight_opt(secs, nanos)
601            .ok_or_else(|| RpcError::value_error("time-of-day out of range"))
602    }
603    fn build_singleton(value: Self) -> Result<ArrayRef> {
604        use chrono::Timelike;
605        let micros = (value.num_seconds_from_midnight() as i64) * 1_000_000
606            + (value.nanosecond() as i64) / 1_000;
607        Ok(Arc::new(Time64MicrosecondArray::from(vec![micros])))
608    }
609}
610
611impl VgiArrow for chrono::Duration {
612    fn arrow_data_type() -> DataType {
613        DataType::Duration(arrow_schema::TimeUnit::Microsecond)
614    }
615    fn describe_name() -> String {
616        "duration".into()
617    }
618    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
619        let micros = as_array::<DurationMicrosecondArray>(arr, "Duration(us)")?.value(idx);
620        Ok(chrono::Duration::microseconds(micros))
621    }
622    fn build_singleton(value: Self) -> Result<ArrayRef> {
623        let micros = value.num_microseconds().ok_or_else(|| {
624            RpcError::value_error("duration overflows microsecond representation")
625        })?;
626        Ok(Arc::new(DurationMicrosecondArray::from(vec![micros])))
627    }
628}
629
630/// Decimal128 with precision 20, scale 4 — matches the conformance
631/// schema. Other (precision, scale) combinations use additional
632/// newtypes if needed.
633#[derive(Clone, Copy, Debug, PartialEq, Eq)]
634pub struct Decimal20_4(pub rust_decimal::Decimal);
635
636impl From<rust_decimal::Decimal> for Decimal20_4 {
637    fn from(d: rust_decimal::Decimal) -> Self {
638        Self(d)
639    }
640}
641
642impl VgiArrow for Decimal20_4 {
643    fn arrow_data_type() -> DataType {
644        DataType::Decimal128(20, 4)
645    }
646    fn describe_name() -> String {
647        "Decimal".into()
648    }
649    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
650        let raw = as_array::<Decimal128Array>(arr, "Decimal128")?.value(idx);
651        // Decimal128 carries the unscaled integer; scale is in the type.
652        let mut d = rust_decimal::Decimal::from_i128_with_scale(raw, 4);
653        d.normalize_assign();
654        Ok(Decimal20_4(d))
655    }
656    fn build_singleton(value: Self) -> Result<ArrayRef> {
657        let mut d = value.0;
658        d.rescale(4);
659        let raw = d.mantissa();
660        let arr = Decimal128Array::from(vec![raw])
661            .with_precision_and_scale(20, 4)
662            .map_err(RpcError::from)?;
663        Ok(Arc::new(arr))
664    }
665}
666
667// ---------------------------------------------------------------------------
668// Option<T> — wraps any VgiArrow with nullable=true.
669// ---------------------------------------------------------------------------
670
671impl<T> VgiArrow for Option<T>
672where
673    T: VgiArrow,
674{
675    fn arrow_data_type() -> DataType {
676        T::arrow_data_type()
677    }
678    fn nullable() -> bool {
679        true
680    }
681    fn describe_name() -> String {
682        format!("{} | None", T::describe_name())
683    }
684    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
685        if arr.is_null(idx) {
686            Ok(None)
687        } else {
688            Ok(Some(T::read(arr, idx)?))
689        }
690    }
691    fn build_singleton(value: Self) -> Result<ArrayRef> {
692        match value {
693            Some(v) => T::build_singleton(v),
694            None => build_null_singleton::<T>(),
695        }
696    }
697}
698
699fn build_null_singleton<T: VgiArrow>() -> Result<ArrayRef> {
700    use arrow_array::array::new_null_array;
701    Ok(new_null_array(&T::arrow_data_type(), 1))
702}
703
704// ---------------------------------------------------------------------------
705// Vec<T> — list types.
706//
707// Handled as `List<inner>` for arbitrary VgiArrow inner types. Common
708// scalar inners (i64, i32, f64, f32, bool, String) get fast-path
709// builders; everything else falls back to a generic per-row push that
710// goes through `T::build_singleton`.
711// ---------------------------------------------------------------------------
712
713impl<T> VgiArrow for Vec<T>
714where
715    T: VgiArrow,
716{
717    fn arrow_data_type() -> DataType {
718        DataType::List(Arc::new(Field::new("item", T::arrow_data_type(), true)))
719    }
720    fn describe_name() -> String {
721        format!("list[{}]", T::describe_name())
722    }
723    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
724        let la = as_array::<ListArray>(arr, "List")?;
725        let inner = la.value(idx);
726        let len = inner.len();
727        let mut out = Vec::with_capacity(len);
728        for i in 0..len {
729            out.push(T::read(inner.as_ref(), i)?);
730        }
731        Ok(out)
732    }
733    fn build_singleton(values: Self) -> Result<ArrayRef> {
734        // Generic path: build each element as a 1-row array via
735        // T::build_singleton, concat them into the list's inner array,
736        // and wrap in a ListArray with a single (0..len) offset pair.
737        // No specialization in V1 — fast-path scalar builders can come
738        // later when `min_specialization` stabilizes.
739        let len = values.len();
740        let mut singletons: Vec<ArrayRef> = Vec::with_capacity(len);
741        for v in values {
742            singletons.push(T::build_singleton(v)?);
743        }
744        let refs: Vec<&dyn Array> = singletons.iter().map(|a| a.as_ref()).collect();
745        let inner = if refs.is_empty() {
746            arrow_array::array::new_empty_array(&T::arrow_data_type())
747        } else {
748            arrow_select::concat::concat(&refs).map_err(RpcError::from)?
749        };
750        let offsets = arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(vec![
751            0i32, len as i32,
752        ]));
753        let field = Arc::new(Field::new("item", T::arrow_data_type(), true));
754        Ok(Arc::new(ListArray::new(field, offsets, inner, None)))
755    }
756}
757
758// ---------------------------------------------------------------------------
759// Map: Vec<(K, V)>
760//
761// Mirrors the Python wire layout for `dict[K, V]`:
762// `Map(entries{key: K, value: V (nullable)})` — the pyarrow/canonical spelling.
763// Only string-keyed maps are supported in V1 because that's what the
764// Python canonical's dataclass introspection emits.
765// ---------------------------------------------------------------------------
766
767/// `Vec<(String, V)>` — wire format `Map<Utf8, V>` (Python canonical).
768impl<V> VgiArrow for Vec<(String, V)>
769where
770    V: VgiArrow,
771{
772    fn arrow_data_type() -> DataType {
773        // The entries struct children are `key`/`value` — the spelling pyarrow's
774        // `pa.map_()` produces and therefore what the canonical Python protocol,
775        // the C++ extension and the Go worker all carry. arrow-rs's own
776        // `MapBuilder` defaults to `keys`/`values`, which is where the earlier
777        // mismatch came from; the read path below is positional, so only the
778        // advertised names change.
779        let entries = Field::new(
780            "entries",
781            DataType::Struct(
782                vec![
783                    Field::new("key", DataType::Utf8, false),
784                    Field::new("value", V::arrow_data_type(), true),
785                ]
786                .into(),
787            ),
788            false,
789        );
790        DataType::Map(Arc::new(entries), false)
791    }
792    fn describe_name() -> String {
793        format!("dict[str, {}]", V::describe_name())
794    }
795    fn read(arr: &dyn Array, idx: usize) -> Result<Self> {
796        let m = as_array::<MapArray>(arr, "Map")?;
797        let entry = m.value(idx);
798        let keys = as_array::<StringArray>(entry.column(0).as_ref(), "Map.keys (Utf8)")?;
799        let values = entry.column(1);
800        let mut out = Vec::with_capacity(keys.len());
801        for i in 0..keys.len() {
802            let v = V::read(values.as_ref(), i)?;
803            out.push((keys.value(i).to_string(), v));
804        }
805        Ok(out)
806    }
807    fn build_singleton(entries: Self) -> Result<ArrayRef> {
808        use arrow_array::array::new_empty_array;
809        let len = entries.len();
810        let (keys, values): (Vec<String>, Vec<V>) = entries.into_iter().unzip();
811        let key_arr = Arc::new(StringArray::from(keys)) as ArrayRef;
812        let value_arr: ArrayRef = if values.is_empty() {
813            new_empty_array(&V::arrow_data_type())
814        } else {
815            let mut singletons: Vec<ArrayRef> = Vec::with_capacity(values.len());
816            for v in values {
817                singletons.push(V::build_singleton(v)?);
818            }
819            let refs: Vec<&dyn Array> = singletons.iter().map(|a| a.as_ref()).collect();
820            arrow_select::concat::concat(&refs).map_err(RpcError::from)?
821        };
822        // Must match `arrow_data_type()` above exactly, or `RecordBatch::try_new`
823        // rejects the array as mismatching its own schema.
824        let entries_struct = arrow_array::StructArray::from(vec![
825            (Arc::new(Field::new("key", DataType::Utf8, false)), key_arr),
826            (
827                Arc::new(Field::new("value", V::arrow_data_type(), true)),
828                value_arr,
829            ),
830        ]);
831        let offsets = arrow_buffer::OffsetBuffer::new(arrow_buffer::ScalarBuffer::from(vec![
832            0i32, len as i32,
833        ]));
834        let entries_field = Arc::new(Field::new(
835            "entries",
836            entries_struct.data_type().clone(),
837            false,
838        ));
839        Ok(Arc::new(MapArray::new(
840            entries_field,
841            offsets,
842            entries_struct,
843            None,
844            false,
845        )))
846    }
847}
848
849#[cfg(test)]
850mod tests {
851    use super::*;
852    use arrow_array::RecordBatch;
853    use arrow_schema::Schema;
854
855    fn round_trip<T: VgiArrow + std::fmt::Debug + PartialEq>(value: T) -> T {
856        let arr = T::build_singleton(value).expect("build_singleton");
857        let schema = Arc::new(Schema::new(vec![Field::new(
858            "v",
859            T::arrow_data_type(),
860            T::nullable(),
861        )]));
862        let batch = RecordBatch::try_new(schema, vec![arr]).unwrap();
863        T::read(batch.column(0).as_ref(), 0).expect("read")
864    }
865
866    #[test]
867    fn roundtrip_string() {
868        assert_eq!(round_trip("hello".to_string()), "hello".to_string());
869    }
870
871    #[test]
872    fn roundtrip_i64() {
873        assert_eq!(round_trip(42i64), 42);
874        assert_eq!(round_trip(-1i64), -1);
875    }
876
877    #[test]
878    fn roundtrip_i32() {
879        assert_eq!(round_trip(7i32), 7);
880    }
881
882    #[test]
883    fn roundtrip_f64() {
884        assert_eq!(round_trip(1.5f64), 1.5);
885    }
886
887    #[test]
888    fn roundtrip_f32() {
889        assert_eq!(round_trip(2.5f32), 2.5);
890    }
891
892    #[test]
893    fn roundtrip_bool() {
894        assert!(round_trip(true));
895        assert!(!round_trip(false));
896    }
897
898    #[test]
899    fn roundtrip_bytes() {
900        assert_eq!(
901            round_trip(Bytes(vec![1, 2, 3, 4, 5])),
902            Bytes(vec![1, 2, 3, 4, 5])
903        );
904    }
905
906    #[test]
907    fn large_bytes_build_reuses_vec_allocation() {
908        let payload = vec![1_u8; 1024];
909        let payload_ptr = payload.as_ptr();
910        let arr = LargeBytes::build_singleton(LargeBytes(payload)).unwrap();
911        let arr = arr.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
912
913        assert_eq!(arr.value(0).as_ptr(), payload_ptr);
914    }
915
916    #[test]
917    fn large_bytes_buffer_roundtrip_reuses_arrow_value_buffer() {
918        let input = LargeBinaryArray::from_iter_values([
919            b"prefix".as_slice(),
920            b"the payload remains in its Arrow buffer".as_slice(),
921        ]);
922        let input_ptr = input.value(1).as_ptr();
923
924        let value = LargeBytesBuffer::read(&input, 1).unwrap();
925        assert_eq!(value.as_ref(), input.value(1));
926        assert_eq!(value.as_ref().as_ptr(), input_ptr);
927
928        let output = LargeBytesBuffer::build_singleton(value).unwrap();
929        let output = output.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
930        assert_eq!(output.value(0), input.value(1));
931        assert_eq!(output.value(0).as_ptr(), input_ptr);
932    }
933
934    #[test]
935    fn roundtrip_option_some() {
936        assert_eq!(round_trip(Some(123i64)), Some(123));
937    }
938
939    #[test]
940    fn roundtrip_option_none() {
941        assert_eq!(round_trip::<Option<i64>>(None), None);
942    }
943
944    #[test]
945    fn roundtrip_option_string() {
946        assert_eq!(
947            round_trip(Some("hello".to_string())),
948            Some("hello".to_string())
949        );
950        assert_eq!(round_trip::<Option<String>>(None), None);
951    }
952
953    #[test]
954    fn roundtrip_vec_i64() {
955        assert_eq!(round_trip(vec![1i64, 2, 3, 4, 5]), vec![1, 2, 3, 4, 5]);
956        let empty: Vec<i64> = Vec::new();
957        assert_eq!(round_trip(empty.clone()), empty);
958    }
959
960    #[test]
961    fn roundtrip_vec_string() {
962        assert_eq!(
963            round_trip(vec!["a".to_string(), "b".to_string()]),
964            vec!["a".to_string(), "b".to_string()]
965        );
966    }
967
968    #[test]
969    fn roundtrip_vec_f64() {
970        assert_eq!(round_trip(vec![1.0f64, 2.5]), vec![1.0, 2.5]);
971    }
972
973    #[test]
974    fn roundtrip_vec_bool() {
975        assert_eq!(round_trip(vec![true, false, true]), vec![true, false, true]);
976    }
977
978    #[test]
979    fn roundtrip_vec_vec_i64() {
980        let v = vec![vec![1i64, 2], vec![3], vec![]];
981        assert_eq!(round_trip(v.clone()), v);
982    }
983
984    #[test]
985    fn roundtrip_map_str_i64() {
986        let m = vec![("a".to_string(), 1i64), ("b".into(), 2)];
987        assert_eq!(round_trip(m.clone()), m);
988    }
989
990    #[test]
991    fn roundtrip_map_str_str() {
992        let m = vec![
993            ("k1".to_string(), "v1".to_string()),
994            ("k2".into(), "v2".into()),
995        ];
996        assert_eq!(round_trip(m.clone()), m);
997    }
998
999    #[test]
1000    fn describe_names_match_python() {
1001        assert_eq!(<String as VgiArrow>::describe_name(), "str");
1002        assert_eq!(<i64 as VgiArrow>::describe_name(), "int");
1003        assert_eq!(<f64 as VgiArrow>::describe_name(), "float");
1004        assert_eq!(<bool as VgiArrow>::describe_name(), "bool");
1005        assert_eq!(<Bytes as VgiArrow>::describe_name(), "bytes");
1006        assert_eq!(<Option<String> as VgiArrow>::describe_name(), "str | None");
1007        assert_eq!(<Vec<i64> as VgiArrow>::describe_name(), "list[int]");
1008        assert_eq!(
1009            <Vec<Vec<i64>> as VgiArrow>::describe_name(),
1010            "list[list[int]]"
1011        );
1012        assert_eq!(
1013            <Vec<(String, i64)> as VgiArrow>::describe_name(),
1014            "dict[str, int]"
1015        );
1016    }
1017
1018    #[test]
1019    fn nullable_flag_only_set_for_option() {
1020        assert!(!<i64 as VgiArrow>::nullable());
1021        assert!(!<String as VgiArrow>::nullable());
1022        assert!(<Option<i64> as VgiArrow>::nullable());
1023        assert!(<Option<Vec<i64>> as VgiArrow>::nullable());
1024    }
1025}
vgi_rpc/arrow_type.rs

vgi_rpc/
arrow_type.rs
