Я следую руководству по плагинам Polars — неправильные прогнозы ветвей, и там говорится, что есть более быстрый способ реализовать следующий код:
#[polars_expr(output_type=Int64)]
fn sum_i64(inputs: &[Series]) -> PolarsResult {
let left: &Int64Chunked = inputs[0].i64()?;
let right: &Int64Chunked = inputs[1].i64()?;
// Note: there's a faster way of summing two columns, see
// section 7.
let out: Int64Chunked = broadcast_binary_elementwise(
left,
right,
|left: Option, right: Option| match (left, right) {
(Some(left), Some(right)) => Some(left + right),
_ => None,
},
);
Ok(out.into_series())
}
избегая ветвления при проверке того, является ли значение нулевым или нет. Однако я не могу найти способ сложить Option вместе, не проверяя, являются ли они нулевыми.
Я сравниваю код, и он выглядит значительно медленнее, чем простое сложение двух полярных столбцов (примерно в 2-3 раза медленнее). Как я могу ускорить этот код, чтобы избежать ветвления.
Кроме того, есть ли что-то, что я делаю неправильно, когда пытаюсь реализовать (a+b).sum() как функцию плагина Polars, поскольку кажется, что это значительно медленнее
Я провожу следующие тесты на Python, чтобы сравнить скорости:
python код
import polars as pl
import time
import minimal_plugin as mp
FILE = "out.parquet"
def make_df():
B = 1 PolarsResult {
let left: &Int64Chunked = inputs[0].i64()?;
let right: &Int64Chunked = inputs[1].i64()?;
let mut total = 0i64;
left.iter().zip(right.iter()).for_each(|(v, w)| {
if let (Some(v), Some(w)) = (v, w) {
total += v + w;
}
});
Ok(Series::new(PlSmallStr::EMPTY, vec![total]))
}
вывод
Number of cols: 67108865
shape: (1, 1)
┌─────────────┐
│ a │
│ --- │
│ i64 │
╞═════════════╡
│ 53366459392 │
└─────────────┘
--- 0.3663475513458252 seconds ---
shape: (1, 1)
┌─────────────┐
│ a │
│ --- │
│ i64 │
╞═════════════╡
│ 53366459392 │
└─────────────┘
--- 0.8730909824371338 seconds ---
shape: (1, 1)
┌─────────────┐
│ a │
│ --- │
│ i64 │
╞═════════════╡
│ 53366459392 │
└─────────────┘
--- 12.611181497573853 seconds ---
Подробнее здесь: https://stackoverflow.com/questions/797 ... -branching
