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().clone(), x.dtype()))
70                            .zip(values.iter().map(|x| x.dtype()))
71                            .map(|((name, dtype), converted_arrow_dtype)| {
72                                create_arrow_field(
73                                    name,
74                                    dtype,
75                                    converted_arrow_dtype,
76                                    self.compat_level,
77                                )
78                            })
79                            .collect(),
80                    ),
81                    arr.len(),
82                    values,
83                    arr.validity().cloned(),
84                )
85                .boxed()
86            },
87            DataType::List(inner) => {
88                let arr: &ListArray<i64> = array.as_any().downcast_ref().unwrap();
89                let new_values = self.array_to_arrow(arr.values().as_ref(), inner);
90
91                let arr = ListArray::<i64>::new(
92                    ArrowDataType::LargeList(Box::new(create_arrow_field(
93                        LIST_VALUES_NAME,
94                        inner.as_ref(),
95                        new_values.dtype(),
96                        self.compat_level,
97                    ))),
98                    arr.offsets().clone(),
99                    new_values,
100                    arr.validity().cloned(),
101                );
102                Box::new(arr)
103            },
104            #[cfg(feature = "dtype-array")]
105            DataType::Array(inner, width) => {
106                use arrow::array::FixedSizeListArray;
107
108                let arr: &FixedSizeListArray = array.as_any().downcast_ref().unwrap();
109                let new_values = self.array_to_arrow(arr.values().as_ref(), inner);
110
111                let arr = FixedSizeListArray::new(
112                    ArrowDataType::FixedSizeList(
113                        Box::new(create_arrow_field(
114                            LIST_VALUES_NAME,
115                            inner.as_ref(),
116                            new_values.dtype(),
117                            self.compat_level,
118                        )),
119                        *width,
120                    ),
121                    arr.len(),
122                    new_values,
123                    arr.validity().cloned(),
124                );
125                Box::new(arr)
126            },
127            #[cfg(feature = "dtype-categorical")]
128            DataType::Categorical(_, _) | DataType::Enum(_, _) => self
129                .categorical_converter
130                .array_to_arrow(array, dtype, self.compat_level),
131            #[cfg(feature = "dtype-date")]
132            DataType::Date => {
133                cast_default(array, &DataType::Date.to_arrow(self.compat_level)).unwrap()
134            },
135            #[cfg(feature = "dtype-datetime")]
136            DataType::Datetime(_, _) => {
137                cast_default(array, &dtype.to_arrow(self.compat_level)).unwrap()
138            },
139            #[cfg(feature = "dtype-duration")]
140            DataType::Duration(_) => {
141                cast_default(array, &dtype.to_arrow(self.compat_level)).unwrap()
142            },
143            #[cfg(feature = "dtype-time")]
144            DataType::Time => {
145                cast_default(array, &DataType::Time.to_arrow(self.compat_level)).unwrap()
146            },
147            #[cfg(feature = "dtype-decimal")]
148            DataType::Decimal(_, _) => array
149                .as_any()
150                .downcast_ref::<arrow::array::PrimitiveArray<i128>>()
151                .unwrap()
152                .clone()
153                .to(dtype.to_arrow(CompatLevel::newest()))
154                .to_boxed(),
155            #[cfg(feature = "object")]
156            DataType::Object(_) => {
157                use crate::chunked_array::object::builder::object_series_to_arrow_array;
158                object_series_to_arrow_array(&unsafe {
159                    Series::from_chunks_and_dtype_unchecked(
160                        PlSmallStr::EMPTY,
161                        vec![array.to_boxed()],
162                        dtype,
163                    )
164                })
165            },
166            DataType::String => {
167                if self.compat_level.0 >= 1 {
168                    array.to_boxed()
169                } else {
170                    cast_unchecked(array, &ArrowDataType::LargeUtf8).unwrap()
171                }
172            },
173            DataType::Binary => {
174                if self.compat_level.0 >= 1 {
175                    array.to_boxed()
176                } else {
177                    cast_unchecked(array, &ArrowDataType::LargeBinary).unwrap()
178                }
179            },
180            _ => {
181                assert!(!dtype.is_logical());
182                array.to_boxed()
183            },
184        }
185    }
186}
187
188fn create_arrow_field(
189    name: PlSmallStr,
190    dtype: &DataType,
191    arrow_dtype: &ArrowDataType,
192    compat_level: CompatLevel,
193) -> ArrowField {
194    match (dtype, arrow_dtype) {
195        #[cfg(feature = "dtype-categorical")]
196        (DataType::Categorical(..) | DataType::Enum(..), ArrowDataType::Dictionary(_, _, _)) => {
197            // Sets _PL_ metadata
198            dtype.to_arrow_field(name, compat_level)
199        },
200        _ => ArrowField::new(name, arrow_dtype.clone(), true),
201    }
202}
polars_core/series/into.rs

polars_core/series/
into.rs
