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 as cast;
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        match self.dtype() {
24            // make sure that we recursively apply all logical types.
25            #[cfg(feature = "dtype-struct")]
26            dt @ DataType::Struct(fields) => {
27                let ca = self.struct_().unwrap();
28                let arr = ca.downcast_chunks().get(chunk_idx).unwrap();
29                let values = arr
30                    .values()
31                    .iter()
32                    .zip(fields.iter())
33                    .map(|(values, field)| {
34                        let dtype = &field.dtype;
35                        let s = unsafe {
36                            Series::from_chunks_and_dtype_unchecked(
37                                PlSmallStr::EMPTY,
38                                vec![values.clone()],
39                                &dtype.to_physical(),
40                            )
41                            .from_physical_unchecked(dtype)
42                            .unwrap()
43                        };
44                        s.to_arrow(0, compat_level)
45                    })
46                    .collect::<Vec<_>>();
47                StructArray::new(
48                    dt.to_arrow(compat_level),
49                    arr.len(),
50                    values,
51                    arr.validity().cloned(),
52                )
53                .boxed()
54            },
55            // special list branch to
56            // make sure that we recursively apply all logical types.
57            DataType::List(inner) => {
58                let ca = self.list().unwrap();
59                let arr = ca.chunks[chunk_idx].clone();
60                let arr = arr.as_any().downcast_ref::<ListArray<i64>>().unwrap();
61
62                let new_values = if let DataType::Null = &**inner {
63                    arr.values().clone()
64                } else {
65                    // We pass physical arrays and cast to logical before we convert to arrow.
66                    let s = unsafe {
67                        Series::from_chunks_and_dtype_unchecked(
68                            PlSmallStr::EMPTY,
69                            vec![arr.values().clone()],
70                            &inner.to_physical(),
71                        )
72                        .from_physical_unchecked(inner)
73                        .unwrap()
74                    };
75
76                    s.to_arrow(0, compat_level)
77                };
78
79                let dtype = ListArray::<i64>::default_datatype(inner.to_arrow(compat_level));
80                let arr = ListArray::<i64>::new(
81                    dtype,
82                    arr.offsets().clone(),
83                    new_values,
84                    arr.validity().cloned(),
85                );
86                Box::new(arr)
87            },
88            #[cfg(feature = "dtype-array")]
89            DataType::Array(inner, width) => {
90                let ca = self.array().unwrap();
91                let arr = ca.chunks[chunk_idx].clone();
92                let arr = arr.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
93
94                let new_values = if let DataType::Null = &**inner {
95                    arr.values().clone()
96                } else {
97                    let s = unsafe {
98                        Series::from_chunks_and_dtype_unchecked(
99                            PlSmallStr::EMPTY,
100                            vec![arr.values().clone()],
101                            &inner.to_physical(),
102                        )
103                        .from_physical_unchecked(inner)
104                        .unwrap()
105                    };
106
107                    s.to_arrow(0, compat_level)
108                };
109
110                let dtype =
111                    FixedSizeListArray::default_datatype(inner.to_arrow(compat_level), *width);
112                let arr =
113                    FixedSizeListArray::new(dtype, arr.len(), new_values, arr.validity().cloned());
114                Box::new(arr)
115            },
116            #[cfg(feature = "dtype-categorical")]
117            dt @ (DataType::Categorical(_, ordering) | DataType::Enum(_, ordering)) => {
118                let ca = self.categorical().unwrap();
119                let arr = ca.physical().chunks()[chunk_idx].clone();
120                // SAFETY: categoricals are always u32's.
121                let cats = unsafe { UInt32Chunked::from_chunks(PlSmallStr::EMPTY, vec![arr]) };
122
123                // SAFETY: we only take a single chunk and change nothing about the index/rev_map mapping.
124                let new = unsafe {
125                    CategoricalChunked::from_cats_and_rev_map_unchecked(
126                        cats,
127                        ca.get_rev_map().clone(),
128                        matches!(dt, DataType::Enum(_, _)),
129                        *ordering,
130                    )
131                };
132
133                new.to_arrow(compat_level, false)
134            },
135            #[cfg(feature = "dtype-date")]
136            DataType::Date => cast(
137                &*self.chunks()[chunk_idx],
138                &DataType::Date.to_arrow(compat_level),
139            )
140            .unwrap(),
141            #[cfg(feature = "dtype-datetime")]
142            DataType::Datetime(_, _) => cast(
143                &*self.chunks()[chunk_idx],
144                &self.dtype().to_arrow(compat_level),
145            )
146            .unwrap(),
147            #[cfg(feature = "dtype-duration")]
148            DataType::Duration(_) => cast(
149                &*self.chunks()[chunk_idx],
150                &self.dtype().to_arrow(compat_level),
151            )
152            .unwrap(),
153            #[cfg(feature = "dtype-time")]
154            DataType::Time => cast(
155                &*self.chunks()[chunk_idx],
156                &DataType::Time.to_arrow(compat_level),
157            )
158            .unwrap(),
159            #[cfg(feature = "dtype-decimal")]
160            DataType::Decimal(_, _) => self.decimal().unwrap().chunks()[chunk_idx]
161                .as_any()
162                .downcast_ref::<PrimitiveArray<i128>>()
163                .unwrap()
164                .clone()
165                .to(self.dtype().to_arrow(CompatLevel::newest()))
166                .to_boxed(),
167            #[cfg(feature = "object")]
168            DataType::Object(_) => {
169                use crate::chunked_array::object::builder::object_series_to_arrow_array;
170                if self.chunks().len() == 1 && chunk_idx == 0 {
171                    object_series_to_arrow_array(self)
172                } else {
173                    // we slice the series to only that chunk
174                    let offset = self.chunks()[..chunk_idx]
175                        .iter()
176                        .map(|arr| arr.len())
177                        .sum::<usize>() as i64;
178                    let len = self.chunks()[chunk_idx].len();
179                    let s = self.slice(offset, len);
180                    object_series_to_arrow_array(&s)
181                }
182            },
183            DataType::String => {
184                if compat_level.0 >= 1 {
185                    self.array_ref(chunk_idx).clone()
186                } else {
187                    let arr = self.array_ref(chunk_idx);
188                    cast_unchecked(arr.as_ref(), &ArrowDataType::LargeUtf8).unwrap()
189                }
190            },
191            DataType::Binary => {
192                if compat_level.0 >= 1 {
193                    self.array_ref(chunk_idx).clone()
194                } else {
195                    let arr = self.array_ref(chunk_idx);
196                    cast_unchecked(arr.as_ref(), &ArrowDataType::LargeBinary).unwrap()
197                }
198            },
199            _ => self.array_ref(chunk_idx).clone(),
200        }
201    }
202}
polars_core/series/into.rs

polars_core/series/
into.rs
