Viết hàm thêm cho PostgreSQL: Chú voi bay

Vừa rồi tôi đã viết bài Dùng Python viết hàm xử lý dữ liệu dưới tầng database cho PostgreSQL. Sau khi chơi với Python một chút, tôi tự hỏi, có thể tăng tốc độ thực thi thêm nữa không. Thế nên hôm nay tôi nghịch thêm vài cách khác nhau, để gắn thêm tên lửa vào đít chú voi PostgreSQL.

Dumbo Picture credit: Walt Disney

Bây giờ tôi sẽ chuyển đổi code kia sang Cython và Rust.

Cython

Cython là một sự mở rộng của Python nhằm giúp code Python của bạn trở nên biên dịch được và tận hưởng tốc độ của một chương trình viết bằng ngôn ngữ biên dịch như C. Cython mang đến hai thứ:

  • Cú pháp pha trộn giữa Python và C.
  • Công cụ để dịch code Python của bạn thành C và biên dịch thành thư viện nhị phân.

Cython

Lưu ý, để việc biên dịch thực sự có ý nghĩa, code của bạn phải được viết bằng cú pháp mới của Cython.

Đoạn code của tôi được chuyển sang phong cách Cython như sau (file ctools.pyx):

# cython: language_level=3

import string
import random
import cython


@cython.returns(Py_UNICODE)
def scramble_email(Py_UNICODE* email):
    # No need to handle the case that "@" is missing, because
    # email address was validated before.
    head, tail = email.split('@')
    # No need to handle the case that email is empty.

    # Pick 3 random letters
    c1, c2, c3 = random.sample(string.ascii_lowercase + string.digits, 3)
    # Choose a random middle position
    i = random.randint(1, max(1, len(head) - 2))

    return f'{c1}{head[:i]}{c2}{head[i:]}{c3}@{tail}'

Biên dịch file ctools.pyx thành một file *.so:

cythonize -ai ctools.pyx

Lệnh trên sẽ sinh ra một file ctools.cpython-38-x86_64-linux-gnu.so, tôi sẽ copy file này vào nơi mà Python và plugin 'plpython' của PostgreSQL tìm thấy được để import.

Để xác định nơi mà plpython nhìn thấy, ta có thể chạy Python lên và soi nội dung của biến sys.path:

>>> import sys
>>> sys.path
['', '/usr/lib/python38.zip', '/usr/lib/python3.8', '/usr/lib/python3.8/lib-dynload', '/home/quan/.local/lib/python3.8/site-packages', '/usr/local/lib/python3.8/dist-packages', '/usr/lib/python3/dist-packages']

Trong danh sách đường dẫn trên, ta chỉ quan tâm đến các đường dẫn bắt đầu bằng "/usr", vì các đường dẫn còn lại là của "ta", PostgreSQL, với tư cách là một process khởi động bởi hệ thống, sẽ không nhìn thấy chúng. Trong các đường dẫn hệ thống (bắt đầu bằng "/usr"), ta sẽ chọn lưu vào cái nào bắt đầu bằng "/usr/local/", vì theo quy ước của Linux, các thư mục "/usr/" mà không có "local" sẽ dành cho phần mềm nào được cài bởi "package manager" (như apt, dnf, pacman...), còn "/usr/local/" sẽ dành cho người quản trị tự ý cài phần mềm riêng. Cụ thể, tôi sẽ copy bằng lệnh sau:

sudo cp ctools.cpython-38-x86_64-linux-gnu.so /usr/local/lib/python3.8/dist-packages

Sau khi đã copy file *.so, ta sẽ định nghĩa lại hàm trong PL/Python như sau:

CREATE OR REPLACE FUNCTION scramble_email(email text)
    RETURNS text
    AS $$

    import ctools
    return ctools.scramble_email(email)
$$ LANGUAGE plpython3u;

Mẹo: Thay vì dán code kia vào màn hình pgcli, ta có thể lưu vào file functions.sql rồi dùng lệnh đặc biệt "\i" của psql nạp lên.

plpython on cython

9043 dòng được xử lý trong 0,35 giây, so với 0,36 giây của code Python thuần. Có cải thiện, nhưng không nhiều lắm.

Rust

Rust là một ngôn ngữ biên dịch ở tầng thấp ngang C. Đây là một ngôn ngữ mới, được yêu thích bởi các lập trình viên Python khi cần một ngôn ngữ biên dịch bổ sung cho kho vũ khí của mình. Là một ngôn ngữ thông dịch, kiểu động (dynamic type) nên Python có bất lợi về tốc độ chạy, bởi vậy các lập trình viên Python sẽ nhặt thêm một ngôn ngữ biên dịch để có thể "chinh chiến trên nhiều mặt trận", áp dụng khi tốc độ chạy là điều quan trọng hàng đầu. Khi chọn một ngôn ngữ biên dịch, tôi chọn Rust vì nó mới, áp dụng được kĩ thuật lập trình hiện đại, cách quản lý bộ nhớ tân tiến, nó lại có nhiều chỗ mà cách viết khá giống Python (khía cạnh functional programming).

Rust

Khi dùng ngôn ngữ biên dịch như Rust thì ta không thể viết hàm thêm trong một thủ tục, mà phải viết luôn một extension cho PostgreSQL và cung cấp hàm mới qua đường đó.

Tôi sẽ theo framework pgx để viết extension mới này. Sau đây là code Rust của tôi (file lib.rs):

use std::cmp::max;
use std::iter::repeat_with;

use pgx::*;
use fastrand;
use itertools::Itertools;

pg_module_magic!();


#[pg_extern]
fn scramble_email(email: &str) -> String {
    // No need to handle the case that "@" is missing, because
    // email address was validated before.
    let (head, tail) = email.split("@").next_tuple().unwrap();

    // Pick 3 random letters
    let (a, b, c) = repeat_with(|| fastrand::alphanumeric().to_lowercase())
                    .next_tuple().unwrap();
    // Choose a random middle position.
    // We know that our email addresses are short enough to fit u8.
    let upto = max(2, head.chars().count() - 1) as u8;
    // fastrand is crazy with inclusive range (m..=n),
    // so we use open range.
    let i = fastrand::u8(1..upto) as usize;
    format!("{}{}{}{}{}@{}", a, &head[..i], b, &head[i..], c, tail)
}

Có thể thấy, thuật toán i xì code bên Python, cũng tách ra, chập vào, cũng dùng format chuỗi. Tuy nhiên, có vài khác biệt nhỏ:

  • fastrand::alphanumeric() trả về cả chữ in hoa (A - Z), nên tôi phải chuyển đổi sang in thường.
  • fastrand::u8 bên Rust không gộp chặn trên như randint bên Python.
  • Python dùng hàm len() để tính chiều dài chuỗi, còn Rust dùng "method" .len() hoặc .count() (sự khác nhau của .len().count() có liên quan đến vấn đề UTF-8).

Một số giải thích cho người mới biết đến Rust:

  • Là một ngôn ngữ biên dịch nên bạn phải khai báo kiểu dữ liệu của biến (email: &str), kiểu dữ liệu trả về bởi hàm (-> String).
  • Rust có tính năng "suy luận kiểu" nên nhiều khi bạn không cần khai báo kiểu cho biến (let upto =). Trình biên dịch Rust sẽ nhìn vào bối cảnh xung quanh để tự suy ra kiểu biến.
  • Là một ngôn ngữ thuộc họ ML, theo trường phái "functional programming" nên Rust dùng nhiều khái niệm của "functional programming" như "iterator".
  • Nếu như bên Python, email.split('@') trả về một list thì email.split("@") bên Rust trả về một "iterator", nên ta cần gọi thêm next_tuple() để bóc từng giá trị của iterator thành một tuple, phục vụ cho việc gán hai biến cùng lúc.
  • Cú pháp "tuple" (như (head, tail)) trong Rust tương ứng với việc định nghĩa một kiểu struct mới, với số lượng thành viên biết trước, nên next_tuple() không cần truyền vào số lượng phần tử (số lần cần lặp).
  • Để biểu thị một hàm nào đó chạy không thành công (lỗi), Rust không dùng khái niệm exception (như nhiều ngôn ngữ OOP quen thuộc), và dùng một kiểu dữ liệu đặc biệt là Result, vừa biểu thị trạng thái (thành công / lỗi), vừa bao bọc giá trị thực (là giá trị ta cần lấy khi hàm thực thi thành công và chi tiết về lỗi khi thất bại). Để lấy giá trị thực thì ta phải gọi unwrap(). Quan điểm "đính kèm trạng thái với giá trị trả về, không dùng exception" này tương tự với Go, nhưng không dùng giá trị đôi như Go. Cá nhân tôi thích phương pháp của Rust hơn (chỉ là sở thích cá nhân).
  • Vì biết trước dữ liệu địa chỉ email của tôi là hợp lệ, biết chắc split("@") sẽ tách ra được hai giá trị nên tôi "tự tin" gọi thẳng unwrap() mà không cần bước kiểm tra trạng thái.
  • Trong các hàm sinh số nguyên ngẫu nhiên của fastrand, tôi thấy fastrand::u8 có tốc độ cao, nên biến upto phía trên được ép kiểu về u8 để có thể dùng được với hàm này. Tôi có thể yên tâm ép kiểu từ usize về u8 vì biết chắc các địa chỉ email đã lưu là đủ ngắn để độ dài của nó không tràn ra khỏi u8.

Biên dịch code trên và tạo ra file sẵn sàng cho việc cài đặt:

cargo pgx package

Cài đặt (tên dự án của tôi là pgrust nên extension sẽ có tên pgrust):

cd target/release/pgrust-pg12/usr
sudo rsync -rvh --progress . /usr

Cài xong rồi, hãy thử sử dụng nó xem sao.

Đầu tiên, nhớ xóa hàm scramble_email tạo bởi Python, để khỏi xung đột với hàm mà ta định nghĩa trong extension Rust:

DROP FUNCTION scramble_email(email text);

Nạp extension của ta lên.

CREATE EXTENSION pgrust;

Chạy thử câu truy vấn, dùng hàm mới viết:

rust-extension

Wow, thời gian thực thi rút ngắn xuống còn 0,189 giây. Tăng tốc được 46%.

Như vậy là chuyến phiêu lưu của tôi với PostgreSQL đã đạt kết quả tốt đẹp. Không những được viết bằng ngôn ngữ yêu thích mà tôi còn chứng tỏ được công hiệu của chúng.


Lưu ý

Để giữ nội dung gọn gàng, bài viết đã bỏ qua bước đóng gói phần mềm mà copy trực tiếp file vào thư mục hệ thống. Cách làm này có rủi ro ghi đè lên file trùng tên có sẵn và có hạn chế là thiếu cơ chế gỡ cài đặt.


Cập nhật

  • 09/12/2020

    Giải thích code Rust.

  • 27/11/2020

    Viết lại code Rust tối ưu cho bài toán này hơn. Sau đây là code cũ, chạy trong 0,27 giây:

    use std::cmp;
    
    use pgx::*;
    use rand::Rng;
    use rand::distributions::Alphanumeric;
    use itertools::Itertools;
    
    pg_module_magic!();
    
    #[pg_extern]
    fn scramble_email(email: &str) -> String {
        let (head, tail) = email.split("@").next_tuple().unwrap();
        let mut rng = rand::thread_rng();
        let (a, b, c) = rng.sample_iter(&Alphanumeric)
            .map(|s| s.to_lowercase()).next_tuple().unwrap();
        let hlen = head.chars().count();
        let i = rng.gen_range(1, cmp::max(2, hlen - 1));
        format!("{}{}{}{}{}@{}", a, &head[..i], b, &head[i..], c, tail)
    }