postgresml · levkk · Aug 31, 2022 · Aug 30, 2022 · Aug 30, 2022 · Aug 30, 2022
diff --git a/.editorconfig b/.editorconfig
@@ -5,7 +5,7 @@ end_of_line = lf
 insert_final_newline = true
 charset = utf-8
 
-[*.py]
+[*.py, *.rs]
 indent_style = space
 indent_size = 4
 

diff --git a/pgml-extension/.dockerignore b/pgml-extension/.dockerignore
@@ -158,3 +158,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+pgml_rust/target/
diff --git a/pgml-extension/Dockerfile b/pgml-extension/Dockerfile
@@ -1,10 +1,18 @@
+# FROM rust:1-bullseye AS rust_builder
+# COPY pgml_rust /pgml_rust
+# WORKDIR /pgml_rust
+# RUN apt-get update && apt-get install -y postgresql-13 libpq-dev cmake libclang-dev
+# RUN cargo install cargo-pgx
+# RUN cargo pgx init
+# RUN cargo pgx package
+
 FROM debian:bullseye-slim
 MAINTAINER docker@postgresml.com
 
 RUN apt-get update
 ARG DEBIAN_FRONTEND=noninteractive
 ENV TZ=Etc/UTC
-RUN apt-get install -y postgresql-plpython3-13 python3 python3-pip postgresql-13 tzdata sudo cmake libpq-dev
+RUN apt-get install -y postgresql-plpython3-13 python3 python3-pip postgresql-13 tzdata sudo cmake libpq-dev libclang-dev
 
 # Cache this, quicker
 RUN pip3 install xgboost sklearn diptest torch lightgbm transformers datasets sentencepiece sacremoses sacrebleu rouge

diff --git a/pgml-extension/pgml_rust/.cargo/config b/pgml-extension/pgml_rust/.cargo/config
@@ -0,0 +1,3 @@
+[build]
+# Postgres symbols won't be available until runtime
+rustflags = ["-C", "link-args=-Wl,-undefined,dynamic_lookup"]
diff --git a/pgml-extension/pgml_rust/.gitignore b/pgml-extension/pgml_rust/.gitignore
@@ -0,0 +1,6 @@
+.DS_Store
+.idea/
+/target
+*.iml
+**/*.rs.bk
+Cargo.lock
diff --git a/pgml-extension/pgml_rust/Cargo.toml b/pgml-extension/pgml_rust/Cargo.toml
@@ -0,0 +1,36 @@
+[package]
+name = "pgml_rust"
+version = "0.0.0"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib"]
+
+[features]
+default = ["pg13"]
+pg10 = ["pgx/pg10", "pgx-tests/pg10" ]
+pg11 = ["pgx/pg11", "pgx-tests/pg11" ]
+pg12 = ["pgx/pg12", "pgx-tests/pg12" ]
+pg13 = ["pgx/pg13", "pgx-tests/pg13" ]
+pg14 = ["pgx/pg14", "pgx-tests/pg14" ]
+pg_test = []
+
+[dependencies]
+pgx = "=0.4.5"
+xgboost = { path = "rust-xgboost" }
+rustlearn = "0.5"
+once_cell = "1"
+rand = "0.8"
+
+[dev-dependencies]
+pgx-tests = "=0.4.5"
+
+[profile.dev]
+panic = "unwind"
+lto = "thin"
+
+[profile.release]
+panic = "unwind"
+opt-level = 3
+lto = "fat"
+codegen-units = 1
diff --git a/pgml-extension/pgml_rust/README.md b/pgml-extension/pgml_rust/README.md
@@ -0,0 +1,23 @@
+# Rust meet PostgresML
+
+Here we have some POC code to use Rust for PostgresML.
+
+## Dependencies
+
+All dependencies are vendored. I downloaded XGBoost 1.62 and all its submodules. We're also using the `master` branch of `xgboost` Rust crate.
+
+If you haven't already, install:
+
+- `cmake`
+- `libclang-dev`
+
+## Local development
+
+1. `cargo install pgx`
+2. `cargo pgx run`
+3. `DROP EXTENSION IF EXISTS pgml_rust;`
+4. `CREATE EXTENSION pgml_rust;`
+5. `SELECT pgml_train('pgml.diabetes', ARRAY['age', 'sex'], 'target');`
+6. `SELECT * FROM pgml_predict(ARRAY[1, 5.0]);`
+
+Lots of todos, but still a decent PoC.
diff --git a/pgml-extension/pgml_rust/pgml_rust.control b/pgml-extension/pgml_rust/pgml_rust.control
@@ -0,0 +1,5 @@
+comment = 'pgml_rust:  Created by pgx'
+default_version = '@CARGO_VERSION@'
+module_pathname = '$libdir/pgml_rust'
+relocatable = false
+superuser = false
diff --git a/pgml-extension/pgml_rust/rust-xgboost/.gitignore b/pgml-extension/pgml_rust/rust-xgboost/.gitignore
@@ -0,0 +1,12 @@
+# Generated by Cargo
+# will have compiled files and executables
+/target/
+/examples/*/target/
+/xgboost-sys/target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
diff --git a/pgml-extension/pgml_rust/rust-xgboost/.gitmodules b/pgml-extension/pgml_rust/rust-xgboost/.gitmodules
@@ -0,0 +1,4 @@
+[submodule "xgboost-sys/xgboost"]
+	path = xgboost-sys/xgboost
+	url = https://github.com/davechallis/xgboost
+	branch = master
diff --git a/pgml-extension/pgml_rust/rust-xgboost/.travis.yml b/pgml-extension/pgml_rust/rust-xgboost/.travis.yml
@@ -0,0 +1,22 @@
+language: rust
+
+os:
+    - linux
+    - osx
+
+rust:
+    - stable
+    - nightly
+matrix:
+  allow_failures:
+    - rust: nightly
+  fast_finish: true
+
+cache: cargo
+
+script:
+  - cd xgboost-sys && cargo test --verbose --all
+  - cd .. && cargo test --verbose --all
+  - cd examples/basic && cargo run
+  - cd ../custom_objective && cargo run
+  - cd ../generalised_linear_model && cargo run
diff --git a/pgml-extension/pgml_rust/rust-xgboost/CHANGELOG.md b/pgml-extension/pgml_rust/rust-xgboost/CHANGELOG.md
@@ -0,0 +1,3 @@
+# 0.1.4 (2019-03-05)
+
+* `Booster::load_buffer` method added (thanks [jonathanstrong](https://github.com/jonathanstrong))
diff --git a/pgml-extension/pgml_rust/rust-xgboost/Cargo.toml b/pgml-extension/pgml_rust/rust-xgboost/Cargo.toml
@@ -0,0 +1,18 @@
+[package]
+name = "xgboost"
+version = "0.2.0"
+authors = ["Dave Challis <dave@suicas.net>"]
+license = "MIT"
+repository = "https://github.com/davechallis/rust-xgboost"
+homepage = "https://github.com/davechallis/rust-xgboost"
+description = "Machine learning using XGBoost"
+documentation = "https://docs.rs/xgboost"
+readme = "README.md"
+
+[dependencies]
+xgboost-sys = { path = "xgboost-sys" }
+libc = "0.2"
+derive_builder = "0.5"
+log = "0.4"
+tempfile = "3.0"
+indexmap = "1.0"
diff --git a/pgml-extension/pgml_rust/rust-xgboost/LICENSE b/pgml-extension/pgml_rust/rust-xgboost/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Dave Challis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/pgml-extension/pgml_rust/rust-xgboost/README.md b/pgml-extension/pgml_rust/rust-xgboost/README.md
@@ -0,0 +1,95 @@
+# rust-xgboost
+
+[![Travis Build Status](https://travis-ci.com/davechallis/rust-xgboost.svg?branch=master)](https://travis-ci.com/davechallis/rust-xgboost)
+[![Documentation link](https://docs.rs/xgboost/badge.svg)](https://docs.rs/xgboost/badge.svg)
+
+Rust bindings for the [XGBoost](https://xgboost.ai) gradient boosting library.
+
+* [Documentation](https://docs.rs/xgboost)
+
+Basic usage example:
+
+```rust
+extern crate xgboost;
+
+use xgboost::{parameters, DMatrix, Booster};
+
+fn main() {
+    // training matrix with 5 training examples and 3 features
+    let x_train = &[1.0, 1.0, 1.0,
+                    1.0, 1.0, 0.0,
+                    1.0, 1.0, 1.0,
+                    0.0, 0.0, 0.0,
+                    1.0, 1.0, 1.0];
+    let num_rows = 5;
+    let y_train = &[1.0, 1.0, 1.0, 0.0, 1.0];
+
+    // convert training data into XGBoost's matrix format
+    let mut dtrain = DMatrix::from_dense(x_train, num_rows).unwrap();
+
+    // set ground truth labels for the training matrix
+    dtrain.set_labels(y_train).unwrap();
+
+    // test matrix with 1 row
+    let x_test = &[0.7, 0.9, 0.6];
+    let num_rows = 1;
+    let y_test = &[1.0];
+    let mut dtest = DMatrix::from_dense(x_test, num_rows).unwrap();
+    dtest.set_labels(y_test).unwrap();
+
+    // configure objectives, metrics, etc.
+    let learning_params = parameters::learning::LearningTaskParametersBuilder::default()
+        .objective(parameters::learning::Objective::BinaryLogistic)
+        .build().unwrap();
+
+    // configure the tree-based learning model's parameters
+    let tree_params = parameters::tree::TreeBoosterParametersBuilder::default()
+            .max_depth(2)
+            .eta(1.0)
+            .build().unwrap();
+
+    // overall configuration for Booster
+    let booster_params = parameters::BoosterParametersBuilder::default()
+        .booster_type(parameters::BoosterType::Tree(tree_params))
+        .learning_params(learning_params)
+        .verbose(true)
+        .build().unwrap();
+
+    // specify datasets to evaluate against during training
+    let evaluation_sets = &[(&dtrain, "train"), (&dtest, "test")];
+
+    // overall configuration for training/evaluation
+    let params = parameters::TrainingParametersBuilder::default()
+        .dtrain(&dtrain)                         // dataset to train with
+        .boost_rounds(2)                         // number of training iterations
+        .booster_params(booster_params)          // model parameters
+        .evaluation_sets(Some(evaluation_sets)) // optional datasets to evaluate against in each iteration
+        .build().unwrap();
+
+    // train model, and print evaluation data
+    let bst = Booster::train(&params).unwrap();
+
+    println!("{:?}", bst.predict(&dtest).unwrap());
+}
+```
+
+See the [examples](https://github.com/davechallis/rust-xgboost/tree/master/examples) directory for
+more detailed examples of different features.
+
+## Status
+
+Currently in a very early stage of development, so the API is changing as usability issues occur,
+or new features are supported.
+
+Builds against XGBoost 0.81.
+
+### Platforms
+
+Tested:
+
+* Linux
+* Mac OS
+
+Unsupported:
+
+* Windows
diff --git a/pgml-extension/pgml_rust/rust-xgboost/examples/README.md b/pgml-extension/pgml_rust/rust-xgboost/examples/README.md
@@ -0,0 +1,6 @@
+# xgboost feature examples
+
+* [Basic usage](basic/src/main.rs)
+* [Custom objective and evaluation functions](custom_objective/src/main.rs)
+* [Generalised linear model](generalised_linear_model/src/main.rs)
+* [Multiclass classification](multiclass_classification/src/main.rs)
diff --git a/pgml-extension/pgml_rust/rust-xgboost/examples/basic/Cargo.toml b/pgml-extension/pgml_rust/rust-xgboost/examples/basic/Cargo.toml
@@ -0,0 +1,11 @@
+[package]
+name = "xgboost-basic-example"
+version = "0.1.0"
+authors = ["Dave Challis <dave@suicas.net>"]
+publish = false
+
+[dependencies]
+xgboost = { path = "../../" }
+sprs = "0.11"
+log = "0.4"
+env_logger = "0.5"
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		# 0.1.4 (2019-03-05)

		* `Booster::load_buffer` method added (thanks [jonathanstrong](https://github.com/jonathanstrong))