polars_core/frame/
upstream_traits.rs

1use std::ops::{Index, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};
2
3use arrow::record_batch::RecordBatchT;
4
5use crate::prelude::*;
6
7impl TryExtend<RecordBatchT<Box<dyn Array>>> for DataFrame {
8    fn try_extend<I: IntoIterator<Item = RecordBatchT<Box<dyn Array>>>>(
9        &mut self,
10        iter: I,
11    ) -> PolarsResult<()> {
12        for record_batch in iter {
13            self.append_record_batch(record_batch)?;
14        }
15
16        Ok(())
17    }
18}
19
20impl TryExtend<PolarsResult<RecordBatchT<Box<dyn Array>>>> for DataFrame {
21    fn try_extend<I: IntoIterator<Item = PolarsResult<RecordBatchT<Box<dyn Array>>>>>(
22        &mut self,
23        iter: I,
24    ) -> PolarsResult<()> {
25        for record_batch in iter {
26            self.append_record_batch(record_batch?)?;
27        }
28
29        Ok(())
30    }
31}
32
33impl Index<usize> for DataFrame {
34    type Output = Column;
35
36    fn index(&self, index: usize) -> &Self::Output {
37        &self.columns()[index]
38    }
39}
40
41macro_rules! impl_ranges {
42    ($range_type:ty) => {
43        impl Index<$range_type> for DataFrame {
44            type Output = [Column];
45
46            fn index(&self, index: $range_type) -> &Self::Output {
47                &self.columns()[index]
48            }
49        }
50    };
51}
52
53impl_ranges!(Range<usize>);
54impl_ranges!(RangeInclusive<usize>);
55impl_ranges!(RangeFrom<usize>);
56impl_ranges!(RangeTo<usize>);
57impl_ranges!(RangeToInclusive<usize>);
58impl_ranges!(RangeFull);
59
60// we don't implement Borrow<str> or AsRef<str> as upstream crates may add impl of trait for usize.
61impl Index<&str> for DataFrame {
62    type Output = Column;
63
64    fn index(&self, index: &str) -> &Self::Output {
65        self.column(index).unwrap()
66    }
67}