Package: dbgym
https://github.com/JiaxuanYou/DBGym/assets/49870769/cc4437eb-bdb6-4285-914b-892062d7d84eOverviewDBGym is a platform designed to facilitate ML research and application on databases.
With less than5 lines of code, you can provide the path of your database, write the predictive query you want, and DBGym will output the predictions along with your database.DBGym is an ongoing research project serving the vibrant open-source community. We sincerely welcome new ideas, use cases, and pull requests to extend DBGym's capability.InstallationPrebuilt DBGym can be easily installed withpip(tested with Python 3.8 and above):pipinstalldbgymAlternatively, you can install the latest DBGym version from our code repository:gitclonehttps://github.com/JiaxuanYou/DBGym.gitcdDBGym
pipinstall-e.DBGym makes ML on Relational DBs simpleIn the current release, DBGym focuses on relational databases (databases with multiple tables), which can be as simple as a directory of CSV or Parquet files. Here we show a simple example of what DBGym can help you achieve.Example DBGym use caseYou may have the following CSV files in your directorydatabase, each representing a tabledatabase├──user.csv├──item.csv└──trans.csvFor example,user.csvmay look like this:_userxytime00.1A2021/01/0110.8A2022/02/0120.5B2023/05/0130.5N/A2024/08/01Your goal is toleverage the full database(not just theuser.csvtable) to predict the unknownN/Avalues in columnyinuser.csv.Let's pause here - the simple capability of predicting unknown column values is in fact extremely expressive. For example,It can be used for classification or regression tasksIt can be used for time series prediction tasks suppose the N/A values are associated with future time column valuesIt can be used for recommendation tasks by predicting with respect to pairs of user and item IDsIt can unify node, edge and graph-level prediction tasks on graphs by representing (heterogeneous) graphs as multiple tables...This is a non-exhaustive list. By being creative, we can do so many things with the capability of predicting unknown column values.What DBGym can do for youDBGym defines a simple declarativepredictive queryto describe the prediction task above, by specifyingtable_name.column_namein DBGym config. In the example above, the predictive query should beuser.y, indicating that you want to set up an ML prediction task for columnyin tableuser, based on the full database's information. DBGym will then automatically set up the ML prediction pipeline, including splitting the train/val/test sets, training an ML model, getting evaluation metrics for train/val/test sets, and returninguser_pred.csvwhere any missing N/A values will be replaced with DBGym predictions.To solve a given predictive query, DBGym provides a variety of ML models from tabular ML to graph ML communities, including:DBGym data APIModelTabularXGBoost, MLPGraphGCN, GraphSAGE, GIN, GATHeterogeneous graphHGCN, HGTYou can also easily create your own models and register them to DBGym. We hope to integrate more ML models into DBGym, and we sincerely appreciate your pull requests!DBGym Quick Start0 Your first DBGym experimentWith less than 5 lines of code, you can set up and carry out experiments using DBGym.fromdbgym.runimportrunfromdbgym.configimportget_config# Get default DBGym configconfig=get_config()stats=run(config)The experiment statistics and predictions returned from therun()function and are further saved in theoutputdirectory by default.1 Run DBGym experiments on RDBench datasetsFor users who want to customize DBGym experiments, we provide two ways to customize configuration: using Python and/or usingYAMLconfig file. Please refer todbgym/config.pyfor all the available configurations.Customize configs with PythonYou can easily set customized config values in the Python code, e.g., picking an RDBench dataset and writing your favorite predictive query. For example:fromdbgym.runimportrunfromdbgym.configimportget_configconfig=get_config()# point to an RDBench dataset. Will auto-download if not cachedconfig.dataset.name='rdb1-ather'# predictive query. Format: table_name.column_nameconfig.dataset.query='entry_examination.Cholesterol'stats=run(config)Customize configs with YAML fileAlternatively, you can write a customized config in a YAML file, e.g.,config_gcn.yaml, and load with DBGym:fromdbgym.runimportrunfromdbgym.configimportget_configconfig=get_config()config.merge_from_file('config_gcn.yaml')stats=run(config)where the yamlconfig_gcn.yamlcan be defined as follows:# Example Experiment Configuration for GCNtrain:epoch:200model:name:'GCN'hidden_dim:128layer:3optim:lr:0.0012 Apply DBGym to your own database datasetsYou can also easily apply DBGym to your own database datasets. To start, you can first organize your customized data as follows:dataset_path
├── dataset_1
│   ├── user.csv
│   ├── item.csv
│   └── ...
├── dataset_2
└── ...Within each CSV file, DBGym assumes that most of the table columns can be arbitrarily named, except for the key columns, where you shouldfollow the naming convention, where the primary key column for tabletableis named as_table. For example, the column names for each CSV file may look like:user.csv:_user,x1,x2,...
item.csv:_item,Feature1,Feature2,...where the special columns begin with_indicate key columns, and the remaining columns (can be arbitrarily named) are regarded as feature columns. For example,_useris the primary key column foruser.csv, which should save the unique ID for each user; other tables could also refer to user ID information, e.g., by adding_usercolumn intrans.csv, in which case_useris the foreign key column intrans.csvthat could have non-unique values.Then, you can point DBGym to the dataset directory, specify the dataset you want, and write a simple predictive query by pointing to any feature column:fromdbgym.runimportrunfromdbgym.configimportget_configconfig=get_config()# provide path to your datasetconfig.dataset.dir='dataset_path'config.dataset.name='dataset_1'config.dataset.query='target.x1'# (optional) set additional customized configsconfig.merge_from_file('config_gcn.yaml')stats=run(config)Finally, DBGym will generate experiment logs and predictions in theoutputdirectory indataset_pathfolder by default. The predictions for columnx1oftarget.csvwill be saved astarget_pred.csv.Alternatively, you can customize your output directory name by adding the following lineconfig.log_dir='output_path'3 Include customized models to DBGymYou can further easily register your customized models to DBGym and run experiments with them. This is especially helpful if you want to benchmark your proposed model with DBGym.
You may follow the instructions below to customize your own model and potentially contribute back to the DBGym repository via pull requests.To start, instead of using the simplepip install dbgymmethod, you should:gitclonehttps://github.com/JiaxuanYou/DBGym.git
pipuninstalldbgym# uninstall existing DBGympipinstall-e.# install with developer modeThen, you can customize your own model indbgym/contribute/your_model.py:fromdbgym.registerimportregisterclassYourModel(nn.Module):...your_model_type='tabular_model'or'graph_model'register(your_model_type,'your_model_name',YourModel)Next, you can use DBGym with your customized modelfromdbgym.runimportrunfromdbgym.configimportget_configconfig=get_config()config.model.name='your_model_name'stats=run(config)DBGym Feature HighlightsFigure 1: An overview of DBGym design principle.1. Unified Task Definition for Various Data Formats.To meet the requirements of diverse users, DBGym provides 3 kinds of APIs: tabular data, homogeneous graphs and heterogeneous graphs.For all these data formats, we propose a unified task definition, enabling results comparison between models for different formats of data.2. Hierarchical Datasets with Comprehensive ExperimentsVia RDBench, DBGym provides 11 datasets with a range of scales, domains, and relationships.These datasets are categorized into three groups based on their relationship complexity.Extensive experiments with 10 baselines are carried out on these datasets.3. Easy-to-use Interfaces with Robust ResultsHighly modularized pipeline.Reproducible experiment configuration.Scalable experiment management.Flexible user customization.Results reported are averaged over the same dataset and same task type (classification or regression).DBGym has enabled RDBench, a user-friendly toolkit for benchmarking ML methods on relational databases. Please refer to our paper for more details:RDBench: ML Benchmark for Relational Databases, arXiv, Zizhao Zhang*, Yi Yang*, Lutong Zou*, He Wen*, Tao Feng, Jiaxuan You.LoggingLogs and predictionsAfter running DBGym, the logs and predictions are saved to an output directory named asoutputby default. Moreover, the command line output will also indicate the path where logs and predictions are saved. For exampleLogs and predictions are saved to Datasets/output/rdb2-bank/loan.Status_MLP_42_20231024_123456Visualization with tensorboardTo see the visualization of the training process, you can go into the output directory and runtensorboard--logdir.Then visitlocalhost:6006, or any other port you specified.CitationIf you find DBGym useful or relevant to your research, please kindly cite our paper:@article{zhang2023rdbench,title={RDBench: ML Benchmark for Relational Databases},author={Zizhao, Zhang and Yi, Yang and Lutong, Zou and He, Wen and Tao, Feng and Jiaxuan, You},journal={arXiv preprint arXiv:2310.16837},year={2023}}