Tu código R puede ser 100 veces más lento de lo que debería
.png?width=50&name=image%20(1).png)
Durante el último año me he dedicado a practicar y aprender a escribir código de alto rendimiento sin interrumpir el flujo de trabajo del equipo con el que estoy trabajando.Esta práctica, también conocida como Programación consciente del rendimiento (en inglés: Performance-Aware Programming) nos ha permitido aportar un inmenso valor a nuestros clientes con paquetes de R como orbweaver y faucet.
Sin embargo, para llevar esto a la práctica no es necesario llevar un problema o algoritmo a un paquete completo, podemos integrarlo directamente en nuestro flujo de trabajo habitual. En este artículo explicaré cómo programar con un sombrero de rendimiento puesto.
Este artículo utiliza un ejemplo en R, pero aplica igualmente a Python o cualquier otro lenguaje.
El reto
Trabajando en un reporte, llegó la necesidad de traducir semanas ISO a meses ISO, según el estándar ISO 8601 (https://es.wikipedia.org/wiki/ISO_8601). Por simplicidad, optamos por utilizar un tabla de traducción, donde tenemos el año ISO, la semana ISO y el mes ISO como valores predefinidos. Esto nos permite filtrar esta tabla para obtener los meses ISO con base en el año y la semana.
Con base en esta idea, definimos una tabla que se ve más o menos así:
isodates <- tibble::tribble(
~isoyear, ~isomonth, ~isoweek,
2024, 1, 1 ,
2024, 1, 2 ,
2024, 1, 3 ,
2024, 1, 4 ,
2024, 2, 5 ,
2024, 2, 6 ,
2024, 2, 7 ,
# Aquí van más valores
2024, 12, 49,
2024, 12, 50,
2024, 12, 51,
2024, 12, 52
)Después se creó una función que filtra esta tabla para obtener el mes ISO.
get_iso_month_r <- function(year, week) {
isodates |>
dplyr::filter(isoyear == year & isoweek == week) |>
dplyr::pull(isomonth)
}En teoría, esto debe funcionar perfectamente. Entonces, intentamos utilizarla en una expresión utilizando dplyr y después de algunos minutos de esperar. Nos dimos cuenta que esta forma de escribir el código no sería viable.
mock_dataset |>
dplyr::rowwise() |>
dplyr::mutate(isomonth = get_iso_month_r(year, week))Solucionar el problema de rendimiento
La primera solución que se nos vino a la cabeza fue: escribamos esto en Rust. Rust nos permitiría vectorizar la operación y minimizar el gasto de nuestra función.
Entonces empezamos a escribir la solución en Rust. Para poder hacer esto incluimos la función rextendr::rust_function directamente en nuestro código de R, lo cual nos generó una función optimizada, compilada y lista para usar.
rextendr::rust_function(
r"{
fn get_iso_month_rs(years: &[f64], weeks: &[f64]) -> Vec<i32> {
years.into_iter().zip(weeks.into_iter()).map(|(&year, &week)| {
match (year as i32, week as i32) {
(2024, 1..=4 ) => 1,
(2024, 5..=8 ) => 2,
(2024, 9..=13 ) => 3,
(2024, 14..=17) => 4,
(2024, 18..=22) => 5,
(2024, 23..=26) => 6,
(2024, 27..=30) => 7,
(2024, 31..=35) => 8,
(2024, 36..=39) => 9,
(2024, 40..=44) => 10,
(2024, 45..=48) => 11,
(2024, 49..=52) => 12,
_ => panic!("Fecha invalida")
}
}).collect()
}
}",
# El perfil "release" es para utilizar la mayor optimización
profile = "release"
)Actualizamos nuestra expresión con dplyr para utilizar la nueva función en Rust y corrió en tan solo unos cuantos milisegundos.
mock_dataset |>
dplyr::mutate(isomonth = get_iso_month_rs(year, week))¿Por qué no escribirlo en R primero?
Escribir una versión vectorizada de nuestra función original no hubiera sido mucho más trabajo que escribirla en Rust. El código es incluso muy similar. La razón por la que optamos escribir nuestra función de primeras en Rust se debe a que es más sencillo entender los costos de correr bucles en un lenguaje compilado como Rust a entender los costos de vectorizar código en R. Además, no implicaba ningún esfuerzo adicional más que saber un poco de Rust.
Comparación de rendimiento
Para hacer este artículo un poco más interesante escribimos la versión vectorizada en R utilizando dplyr::case_when y la comparamos con la versión de Rust para ver cuánto rendimiento ganamos.
# Versión vectorizada en R
get_iso_month_r_2 <- function(year, week) {
dplyr::case_when(
year == 2024 & week %in% 1:4 ~ 1,
year == 2024 & week %in% 5:8 ~ 2,
year == 2024 & week %in% 9:13 ~ 3,
year == 2024 & week %in% 14:17 ~ 4,
year == 2024 & week %in% 18:22 ~ 5,
year == 2024 & week %in% 23:26 ~ 6,
year == 2024 & week %in% 27:30 ~ 7,
year == 2024 & week %in% 31:35 ~ 8,
year == 2024 & week %in% 36:39 ~ 9,
year == 2024 & week %in% 40:44 ~ 10,
year == 2024 & week %in% 45:48 ~ 11,
year == 2024 & week %in% 49:52 ~ 12
)
}Para correr este benchmark utilizamos el paquete microbenchmark y creamos un dataset inicialmente con 100.000 observaciones.
mock_dataset <- data.frame(
year = 2024,
week = round(runif(100000, min = 1, max = 52), 0)
)
microbenchmark::microbenchmark(
times = 100,
"add_isomonth_r_vectorized" = {
mock_dataset |>
dplyr::mutate(isomonth = get_iso_month_r_2(year, week))
},
"add_isomonth_rs" = {
mock_dataset |>
dplyr::mutate(isomonth = get_iso_month_rs(year, week))
}
)Los resultados fueron impresionantes. La versión de Rust fue 20 veces más rápida que la versión en R.
Unit: milliseconds
expr min lq mean median uq max neval
add_isomonth_r_vectorized 58.809796 59.765772 72.863495 60.942076 68.028255 198.090077 100
add_isomonth_rs 3.283077 3.527616 3.843486 3.738133 3.844333 6.574269 100Si aumentamos el número de observaciones de 100.000 a 1.000.000 el espacio entre nuestra función en R y nuestra función en Rust aumentó de 20x a 40x.
Unit: milliseconds
expr min lq mean median uq max neval
add_isomonth_r_vectorized 599.49971 825.10239 834.11362 836.73390 844.02839 881.37060 100
add_isomonth_rs 21.52827 21.79659 21.95012 21.86705 21.98735 23.65436 100Por pura diversión, agregamos una versión paralela de nuestra función en Rust. Para esto utilizamos rayon. Para esto solo tendremos que cambiar una línea de código en nuestra función de Rust. Pasaremos de utilizar into_iter a into_par_iter.
rextendr::rust_function(
r"{
fn get_iso_month_rs_par(years: &[f64], weeks: &[f64]) -> Vec<i32> {
use rayon::prelude::*;
years.into_par_iter().zip(weeks.into_par_iter()).map(|(&year, &week)| {
match (year as i32, week as i32) {
(2024, 1..=4 ) => 1,
(2024, 5..=8 ) => 2,
(2024, 9..=13 ) => 3,
(2024, 14..=17) => 4,
(2024, 18..=22) => 5,
(2024, 23..=26) => 6,
(2024, 27..=30) => 7,
(2024, 31..=35) => 8,
(2024, 36..=39) => 9,
(2024, 40..=44) => 10,
(2024, 45..=48) => 11,
(2024, 49..=52) => 12,
_ => panic!("Fecha invalida")
}
}).collect()
}
}",
dependencies = list(rayon = "1.10"),
profile = "release"
)Agregamos esta función a nuestro benchmark (utilizando 1.000.000 de observaciones) y lo corremos para ver los resultados. La máquina en la que corremos este benchmark tiene solo 4 hilos de procesamiento por lo cual el rendimiento está sujeto a la CPU en la que corre.
Unit: milliseconds
expr min lq mean median uq max neval
add_isomonth_r_vectorized 572.39804 579.536659 637.9863 589.355421 614.49464 1075.81790 100
add_isomonth_rs 21.97963 22.159766 23.5626 22.348087 22.89876 46.58920 100
add_isomonth_rs_par 8.84868 9.014089 10.0714 9.180416 10.07903 20.37571 100Ahora vemos que el rendimiento de nuestra versión paralela es más o menos 2x de nuestra versión iterativa. Por lo cual nuestro rendimiento comparado a la versión en R (con solo dplyr) es de más o menos 80x.
Conclusión
Escribir código R de alto rendimiento no tiene que ser una tarea tediosa. Mejorar el rendimiento de un flujo de datos en R por 10x o incluso 100x es posible y típicamente se puede hacer sin dedicar un enorme esfuerzo de optimización. En particular cuando nos enfrentamos a tareas que se han de repetir miles o millones de veces, por ejemplo en una transformación de datos, puede añadir valor recordar que operaciones pueden beneficiarse de vectorizar en un lenguaje como Rust o C++. Con solo un cambio relativamente menor, como el de arriba puedes ahorrar millones de ciclos de tu CPU y horas de tu valioso tiempo y el de tu equipo.
Nota: En un siguiente blog post compararemos esta solución con DuckDB para mostrar que con solo R y un poquito de Rust, C o C++ podemos obtener un excelente rendimiento sin cambiar nuestra forma de trabajar.
