Python是数据科学家最流行的编程语言之一——这是有充分理由的。Python包索引(PyPI)承载了大量令人印象深刻的数据科学库包,比如NumPy、SciPy、自然语言工具包、Pandas和Matplotlib。大量可用的高质量分析库及其庞大的开发人员社区使Python成为许多数据科学家的容易选择。

出于性能原因,这些库中的许多都是用C和c++实现的,但是提供了外部函数接口(FFIs)或Python绑定,这样您就可以从Python调用这些函数。这些“低级”语言实现用于减轻Python的一些常见问题,特别是执行时间和内存消耗。绑定执行时间和内存消耗简化了可伸缩性,这对于降低成本至关重要。如果我们可以编写performant代码来完成数据科学任务,那么与Python集成就是一个主要的优势。

数据科学和恶意软件分析的交叉不仅需要快速的执行时间,还需要有效地使用共享资源以实现可扩展性。可伸缩性对于“大数据”问题至关重要,比如有效地处理多个平台上的数百万个可执行文件的数据。在现代处理器上获得良好的性能需要并行性(通常通过多个线程),但高效的执行时间和内存使用也是必要的。对于这种问题,平衡本地系统资源可能会很困难,而正确实现多线程系统则更加困难。C和c++本身不提供线程安全。虽然存在特定于平台的外部库,但显然开发人员有责任维护线程安全。

解析恶意软件本身就很危险。恶意软件经常以意想不到的方式操纵文件格式数据结构,从而导致分析工具失败。一个相对常见的Python解析陷阱是由于缺乏强大的类型安全性造成的。当期望字节数组时,Python无缘无故地接受None值,这很容易导致普遍混乱,而不会在代码中添加None检查。这些与“duck typing”相关的假设通常会导致失败。

进入Rust。Rust语言提出了许多与上述潜在问题的理想解决方案相一致的主张:与C和c++相比,执行时间和内存消耗,以及提供广泛的线程安全性。Rust语言提供了额外的有益特性,比如强大的内存安全保证和没有运行时开销。没有运行时开销简化了Rust代码与其他语言(包括Python)的集成。在这篇博客中,我们对Rust做了一个简短的测试,看看这种炒作是否值得。

数据科学是一个非常广泛的领域,有太多的应用在一篇博客文章中讨论。一个简单的数据科学任务的例子是计算字节序列的信息熵。以比特为单位计算熵的一般公式是(参见维基百科;信息熵):

H(X)=-i Px(xi) log2Px(xi)

为了计算随机变量X的熵,我们首先计算每个可能的字节值(xi)的出现次数,然后除以出现的总次数来计算特定值xi的出现概率(Px(xi))。然后我们计算特定值xi的发生概率(Px(xi))和所谓的自我信息(log2Px(xi))的加权和的负数。因为我们是以位为单位计算熵,所以我们使用log2(以2为基数表示位)。

让我们对Rust进行测试,看看它是如何对纯Python,甚至上面提到的一些非常流行的Python库执行熵计算的。这是对数据科学应用程序performant Rust的简单评估,而不是对Python或可用的优秀库的批评。在这些测试中,我们将从Rust代码中生成一个可以从Python导入的自定义C库。所有测试都在Ubuntu 18.04上运行。

我们从一个简单的纯Python函数(在entropy.py中)开始,只使用标准库数学模块来计算bytearray的熵。此函数未进行优化,并为修改和性能度量提供基线。

import math
def compute_entropy_pure_python(data):
"""Compute entropy on bytearray `data`."""
counts = * 256
entropy = 0.0
length = len(data)
for byte in data:
counts += 1
for count in counts:
if count != 0:
probability = float(count) / length
entropy -= probability * math.log(probability, 2)
return entropy

如您所料,SciPy提供了一个计算熵的函数。我们将首先使用NumPy的unique函数来计算字节频率。将SciPy的熵函数的性能与其他实现进行比较有点不公平,因为SciPy实现有额外的功能来计算相对熵(Kullback-Leibler divergence)。同样,我们只是进行一个(希望不是太)缓慢的测试,以查看从Python导入的Rust编译库的性能。下面是包含在我们的entropy.py脚本中的一个基于scipy的实现。

import numpy as np
from scipy.stats import entropy as scipy_entropy
def compute_entropy_scipy_numpy(data):
"""Compute entropy on bytearray `data` with SciPy and NumPy."""
counts = np.bincount(bytearray(data), minlength=256)
return scipy_entropy(counts, base=2)

与以前的实现相比,我们现在对Rust实现的彻底性和可重复性进行了更深入的研究。我们从使用Cargo生成的默认库包开始。下面的章节描述了我们如何修改Rust包。

cargo new --lib rust_entropy
Cargo.toml

从义务货物开始。toml manifest文件,我们定义了Cargo包和库名rust_entropy_lib。我们使用板条箱上的公共cpython板条箱(v0.4.1)。io, Rust包注册表。我们还使用Rust v1.42.0,这是撰写本文时可用的最新稳定版本。

name = "rust-entropy"
version = "0.1.0"
authors = edition = "2018"
name = "rust_entropy_lib"
crate-type =
version = "0.4.1"
features =
lib.rs

Rust库的实现相当简单。与在纯Python实现中所做的一样,我们为每个可能的字节值初始化一个计数数组,并遍历数据以填充计数。为了完成计算,我们计算并返回概率的负和乘以概率的log2。

use cpython::{py_fn, py_module_initializer, PyResult, Python};
/// Compute entropy on byte array.
fn compute_entropy_pure_rust(data: &) -> f64 {
let mut counts = ;
let mut entropy = 0_f64;
let length = data.len as f64;
// collect byte counts
for &byte in data.iter {
counts += 1;
}
// make entropy calculation
for &count in counts.iter {
if count != 0 {
let probability = f64::from(count) / length;
entropy -= probability * probability.log2;
}
}
entropy
}

对于lib.rs,剩下的就是从Python调用纯Rust函数的机制。我们在lib.rs中包含了一个可识别CPython的函数(compute_entropy_cpython)来调用我们的“纯”Rust函数(compute_entropy_pure_rust)。这种设计为我们提供了维护单一纯Rust实现的好处,还提供了一个cpython友好的“包装器”。

/// Rust-CPython aware function
fn compute_entropy_cpython(_: Python, data: &) -> PyResult {
let _gil = Python::acquire_gil;
let entropy = compute_entropy_pure_rust(data);
Ok(entropy)
}
// initialize Python module and add Rust CPython aware function
py_module_initializer!(
librust_entropy_lib,
initlibrust_entropy_lib,
PyInit_rust_entropy_lib,
|py, m | {
m.add(py, "__doc__", "Entropy module implemented in Rust")?;
m.add(
py,
"compute_entropy_cpython",
py_fn!(py, compute_entropy_cpython(data: &)
)
)?;
Ok()
}
);

最后,我们通过首先导入由Rust编译的自定义动态系统库,从Python调用Rust实现(同样是在entropy.py中)。然后调用之前用py_module_initializer初始化Python模块时指定的库函数!在Rust代码中的宏。现在,我们有了一个Python模块(entropy.py),它包含调用所有熵计算实现的函数。

import rust_entropy_lib
def compute_entropy_rust_from_python(data):
"""Compute entropy on bytearray `data` with Rust."""
return rust_entropy_lib.compute_entropy_cpython(data)

我们使用Cargo在Ubuntu 18.04上构建上面的Rust库包。(此链接可能对OS X用户有帮助。)

cargo build --release

一旦构建完成,我们将生成的动态库复制并重命名到Python模块所在的目录,这样就可以从Python脚本导入它。这个由货物生成的库名为librust_entropy_lib。但是需要重命名为rust_entropy_lib。所以要在这些测试中成功导入。性能结果

我们用pytest基准测试了每个函数实现的执行时间,计算熵超过100万个随机字节。所有实现都提供相同的数据。基准测试(也包括在entropy.py中)如下所示。

# ### BENCHMARKS ###
# generate some random bytes to test w/ NumPy
NUM = 1000000
VAL = np.random.randint(0, 256, size=(NUM, ), dtype=np.uint8)
def test_pure_python(benchmark):
"""Test pure Python."""
benchmark(compute_entropy_pure_python, VAL)
def test_python_scipy_numpy(benchmark):
"""Test pure Python with SciPy."""
benchmark(compute_entropy_scipy_numpy, VAL)
def test_rust(benchmark):
"""Test Rust implementation called from Python."""
benchmark(compute_entropy_rust_from_python, VAL)

最后,我们为每种计算熵的方法分别编写了简单的驱动脚本。下面是一个典型的驱动程序脚本,用于测试纯Python实现。testdata.bin文件是用于测试所有方法的1,000,000随机字节。为了简化捕获内存使用数据,所有方法都重复计算100次。

import entropy
with open('testdata.bin', 'rb') as f:
DATA = f.read
for _ in range(100):
entropy.compute_entropy_pure_python(DATA)

SciPy/NumPy和Rust实现都表现出强大的性能,很容易就比未优化的纯Python实现性能高出100倍以上。Rust版本的性能只比SciPy/NumPy稍好一点,但结果证实了我们的预期:纯Python比编译语言慢得多,用Rust编写的扩展可以与用C编写的扩展极具竞争力(甚至在这个微基准测试中击败它们)。

还有其他提高性能的方法。我们可以使用ctypes或cffi模块。我们可以添加类型提示,并使用Cython来生成一个可以从Python导入的库。所有这些选项都需要考虑特定于解决方案的权衡。

Function/ImplementationMinimum Benchmark
Execution Time (µs)compute_entropy_pure_python
294,319
compute_entropy_scipy_numpy
2,370
compute_entropy_rust_from_python

我们还使用GNU time应用程序(不要与内置的shell命令time混淆)测量了每个函数实现的内存使用情况。特别地,我们度量最大驻留集大小。

虽然纯Python和Rust实现有非常相似的最大驻留集大小,但SciPy/NumPy在这个基准测试中使用的内存明显更多,可能是由于在导入时加载到内存中的额外功能。在这两种情况下,从Python调用Rust代码都不会增加大量的内存开销。

Function/ImplementationMaximum resident
set size (KB)compute_entropy_pure_python
65,262
73,934

65,444总结

Python调用Rust的性能给我们留下了深刻的印象。在我们公认的简短评估中,我们的Rust实现性能与来自SciPy和NumPy包的底层C代码相当。Rust似乎很适合大规模高效加工。

Rust不仅在执行时间内具有性能,而且在这些测试中额外的内存开销也是最小的。执行时间和内存使用特性应该是可伸缩性的理想选择。SciPy和NumPy C FFI实现的性能当然是可以比较的,但是Rust提供了C和c++所没有的额外好处。内存和线程安全保证是吸引人的优点。

虽然C提供了类似的运行时执行改进,但它本身并不提供线程安全。外部库的存在为C提供了这种功能,但正确性的责任完全落在开发人员身上。Rust使用它的所有权模型在编译时检查线程安全问题,比如竞争条件,标准库提供了一套并发机制,包括通道、锁和引用计数智能指针。

我们不提倡任何人将SciPy或NumPy移植到Rust,因为这些包已经经过了大量优化,并具有强大的支持社区。另一方面,我们会强烈考虑将纯Python代码移植到Rust,否则在高性能库中是不可用的。对于安全领域的数据科学应用程序,Rust似乎是一个引人注目的选择,因为它的速度和安全保证。谁会介意在他们的数据科学上加点氧化铁来换取更高的速度和安全性呢?