polars_core/series/
into.rs

1#[cfg(any(
2    feature = "dtype-datetime",
3    feature = "dtype-date",
4    feature = "dtype-duration",
5    feature = "dtype-time"
6))]
7use polars_compute::cast::cast_default;
8use polars_compute::cast::cast_unchecked;
9
10use crate::prelude::*;
11
12impl Series {
13    /// Returns a reference to the Arrow ArrayRef
14    #[inline]
15    pub fn array_ref(&self, chunk_idx: usize) -> &ArrayRef {
16        &self.chunks()[chunk_idx] as &ArrayRef
17    }
18
19    /// Convert a chunk in the Series to the correct Arrow type.
20    /// This conversion is needed because polars doesn't use a
21    /// 1 on 1 mapping for logical/categoricals, etc.
22    pub fn to_arrow(&self, chunk_idx: usize, compat_level: CompatLevel) -> ArrayRef {
23        ToArrowConverter {
24            compat_level,
25            #[cfg(feature = "dtype-categorical")]
26            categorical_converter: {
27                let mut categorical_converter =
28                    crate::series::categorical_to_arrow::CategoricalToArrowConverter {
29                        converters: Default::default(),
30                        persist_remap: false,
31                        output_keys_only: false,
32                    };
33
34                categorical_converter.initialize(self.dtype());
35
36                categorical_converter
37            },
38        }
39        .array_to_arrow(self.chunks().get(chunk_idx).unwrap().as_ref(), self.dtype())
40    }
41}
42
43pub struct ToArrowConverter {
44    pub compat_level: CompatLevel,
45    #[cfg(feature = "dtype-categorical")]
46    pub categorical_converter: crate::series::categorical_to_arrow::CategoricalToArrowConverter,
47}
48
49impl ToArrowConverter {
50    pub fn array_to_arrow(&mut self, array: &dyn Array, dtype: &DataType) -> Box<dyn Array> {
51        match dtype {
52            // make sure that we recursively apply all logical types.
53            #[cfg(feature = "dtype-struct")]
54            DataType::Struct(fields) => {
55                use arrow::array::StructArray;
56
57                let arr: &StructArray = array.as_any().downcast_ref().unwrap();
58                let values = arr
59                    .values()
60                    .iter()
61                    .zip(fields.iter())
62                    .map(|(values, field)| self.array_to_arrow(values.as_ref(), field.dtype()))
63                    .collect::<Vec<_>>();
64
65                StructArray::new(
66                    ArrowDataType::Struct(
67                        fields
68                            .iter()
69                            .map(|x| x.name())
70                            .zip(values.iter().map(|x| x.dtype()))
71                            .map(|(name, dtype)| ArrowField::new(name.clone(), dtype.clone(), true))
72                            .collect(),
73                    ),
74                    arr.len(),
75                    values,
76                    arr.validity().cloned(),
77                )
78                .boxed()
79            },
80            DataType::List(inner) => {
81                let arr: &ListArray<i64> = array.as_any().downcast_ref().unwrap();
82                let new_values = self.array_to_arrow(arr.values().as_ref(), inner);
83
84                let arr = ListArray::<i64>::new(
85                    ListArray::<i64>::default_datatype(new_values.dtype().clone()),
86                    arr.offsets().clone(),
87                    new_values,
88                    arr.validity().cloned(),
89                );
90                Box::new(arr)
91            },
92            #[cfg(feature = "dtype-array")]
93            DataType::Array(inner, width) => {
94                use arrow::array::FixedSizeListArray;
95
96                let arr: &FixedSizeListArray = array.as_any().downcast_ref().unwrap();
97                let new_values = self.array_to_arrow(arr.values().as_ref(), inner);
98
99                let arr = FixedSizeListArray::new(
100                    FixedSizeListArray::default_datatype(new_values.dtype().clone(), *width),
101                    arr.len(),
102                    new_values,
103                    arr.validity().cloned(),
104                );
105                Box::new(arr)
106            },
107            #[cfg(feature = "dtype-categorical")]
108            DataType::Categorical(_, _) | DataType::Enum(_, _) => self
109                .categorical_converter
110                .array_to_arrow(array, dtype, self.compat_level),
111            #[cfg(feature = "dtype-date")]
112            DataType::Date => {
113                cast_default(array, &DataType::Date.to_arrow(self.compat_level)).unwrap()
114            },
115            #[cfg(feature = "dtype-datetime")]
116            DataType::Datetime(_, _) => {
117                cast_default(array, &dtype.to_arrow(self.compat_level)).unwrap()
118            },
119            #[cfg(feature = "dtype-duration")]
120            DataType::Duration(_) => {
121                cast_default(array, &dtype.to_arrow(self.compat_level)).unwrap()
122            },
123            #[cfg(feature = "dtype-time")]
124            DataType::Time => {
125                cast_default(array, &DataType::Time.to_arrow(self.compat_level)).unwrap()
126            },
127            #[cfg(feature = "dtype-decimal")]
128            DataType::Decimal(_, _) => array
129                .as_any()
130                .downcast_ref::<arrow::array::PrimitiveArray<i128>>()
131                .unwrap()
132                .clone()
133                .to(dtype.to_arrow(CompatLevel::newest()))
134                .to_boxed(),
135            #[cfg(feature = "object")]
136            DataType::Object(_) => {
137                use crate::chunked_array::object::builder::object_series_to_arrow_array;
138                object_series_to_arrow_array(&unsafe {
139                    Series::from_chunks_and_dtype_unchecked(
140                        PlSmallStr::EMPTY,
141                        vec![array.to_boxed()],
142                        dtype,
143                    )
144                })
145            },
146            DataType::String => {
147                if self.compat_level.0 >= 1 {
148                    array.to_boxed()
149                } else {
150                    cast_unchecked(array, &ArrowDataType::LargeUtf8).unwrap()
151                }
152            },
153            DataType::Binary => {
154                if self.compat_level.0 >= 1 {
155                    array.to_boxed()
156                } else {
157                    cast_unchecked(array, &ArrowDataType::LargeBinary).unwrap()
158                }
159            },
160            _ => {
161                assert!(!dtype.is_logical());
162                array.to_boxed()
163            },
164        }
165    }
166}
polars_core/series/into.rs

polars_core/series/
into.rs
