cudf-polars 0.3.1

GPU execution engine for Polars using NVIDIA libcudf
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

//! Type and operator mapping between Polars and cudf.

use cudf::BinaryOp;
use cudf::types::{DataType as GpuDataType, TypeId as GpuTypeId};
use polars_core::prelude::DataType;
use polars_error::{PolarsResult, polars_bail};
use polars_plan::dsl::Operator;

/// Map a Polars `Operator` to a cudf `BinaryOp`.
pub fn map_operator(op: Operator) -> PolarsResult<BinaryOp> {
    match op {
        Operator::Plus => Ok(BinaryOp::Add),
        Operator::Minus => Ok(BinaryOp::Sub),
        Operator::Multiply => Ok(BinaryOp::Mul),
        Operator::RustDivide => Ok(BinaryOp::Div),
        Operator::TrueDivide => Ok(BinaryOp::TrueDiv),
        Operator::FloorDivide => Ok(BinaryOp::FloorDiv),
        Operator::Modulus => Ok(BinaryOp::Mod),
        Operator::Eq => Ok(BinaryOp::Equal),
        Operator::EqValidity => Ok(BinaryOp::NullEquals),
        Operator::NotEq => Ok(BinaryOp::NotEqual),
        Operator::NotEqValidity => Ok(BinaryOp::NullNotEquals),
        Operator::Lt => Ok(BinaryOp::Less),
        Operator::LtEq => Ok(BinaryOp::LessEqual),
        Operator::Gt => Ok(BinaryOp::Greater),
        Operator::GtEq => Ok(BinaryOp::GreaterEqual),
        Operator::And => Ok(BinaryOp::BitwiseAnd),
        Operator::Or => Ok(BinaryOp::BitwiseOr),
        Operator::Xor => Ok(BinaryOp::BitwiseXor),
        Operator::LogicalAnd => Ok(BinaryOp::LogicalAnd),
        Operator::LogicalOr => Ok(BinaryOp::LogicalOr),
    }
}

/// Returns true if the operator produces a boolean result.
pub fn is_comparison(op: Operator) -> bool {
    matches!(
        op,
        Operator::Eq
            | Operator::EqValidity
            | Operator::NotEq
            | Operator::NotEqValidity
            | Operator::Lt
            | Operator::LtEq
            | Operator::Gt
            | Operator::GtEq
            | Operator::LogicalAnd
            | Operator::LogicalOr
    )
}

/// Map a Polars `DataType` to a cudf `DataType`.
pub fn map_dtype(dtype: &DataType) -> PolarsResult<GpuDataType> {
    let type_id = match dtype {
        DataType::Boolean => GpuTypeId::Bool8,
        DataType::Int8 => GpuTypeId::Int8,
        DataType::Int16 => GpuTypeId::Int16,
        DataType::Int32 => GpuTypeId::Int32,
        DataType::Int64 => GpuTypeId::Int64,
        DataType::UInt8 => GpuTypeId::Uint8,
        DataType::UInt16 => GpuTypeId::Uint16,
        DataType::UInt32 => GpuTypeId::Uint32,
        DataType::UInt64 => GpuTypeId::Uint64,
        DataType::Float32 => GpuTypeId::Float32,
        DataType::Float64 => GpuTypeId::Float64,
        DataType::String => GpuTypeId::String,
        DataType::Date => GpuTypeId::TimestampDays,
        DataType::Datetime(tu, tz) => {
            if tz.is_some() {
                polars_bail!(ComputeError: "GPU engine: timezone-aware Datetime not yet supported");
            }
            use polars_core::prelude::TimeUnit;
            match tu {
                TimeUnit::Milliseconds => GpuTypeId::TimestampMilliseconds,
                TimeUnit::Microseconds => GpuTypeId::TimestampMicroseconds,
                TimeUnit::Nanoseconds => GpuTypeId::TimestampNanoseconds,
            }
        }
        DataType::Duration(tu) => {
            use polars_core::prelude::TimeUnit;
            match tu {
                TimeUnit::Milliseconds => GpuTypeId::DurationMilliseconds,
                TimeUnit::Microseconds => GpuTypeId::DurationMicroseconds,
                TimeUnit::Nanoseconds => GpuTypeId::DurationNanoseconds,
            }
        }
        _ => polars_bail!(ComputeError: "GPU engine: unsupported dtype {:?}", dtype),
    };
    Ok(GpuDataType::new(type_id))
}