Skip to main content

No project description provided

Project description

Apache License Read The Doc javadoc Pypi version Github Action stability-experimental

Join the community: Join the chat at

:heavy_exclamation_mark: This repository is still experimental. No API-compatibility is guaranteed.


Rikai is a parquet based ML data format built for working with unstructured data at scale. Processing large amounts of data for ML is never trivial, but that is especially true for images and videos often at the core of deep learning applications. We are building Rikai with two main goals:

  1. Enable ML engineers/researchers to have a seamless workflow from Feature Engineering (Spark) to Training (PyTorch/Tensorflow), from notebook to production.
  2. Enable advanced analytics capabilities to support much faster active learning, model debugging, and monitoring in production pipelines.

Current (v0.0.1) main features:

  1. Native support in Spark and PyTorch for images/videos: reduce ad-hoc type conversions when moving between ETL and training.
  2. Custom functionality for working with images and videos at scale: reduce boilerplate and low-level code currently required to process images, filter/sample videos, etc.
  3. ML-enabled SQL analytics API


  1. TensorFlow integration
  2. Versioning support built into the dataset
  3. Richer video capabilities (ffmpeg-python integration)
  4. Declarative annotation API (think vega-lite for annotating images/videos)


from import DenseMetrix
from rikai.types import Image, Box2d
from rikai import numpy as np

df = spark.createDataFrame(
        "id": 1,
        "mat": DenseMatrix(2, 2, range(4)),
        "image": Image("s3://foo/bar/1.png"),
        "annotations": [
                "label": "cat",
                "mask": np.random(size=(256,256)),
                "bbox": Box2d(xmin=1.0, ymin=2.0, xmax=3.0, ymax=4.0)


Train dataset in Pytorch

from import Dataset
from rikai.torch import DataLoader # Do not need this with Pytorch 1.8+
from torchvision import transforms as T

transform = T.Compose([
   T.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))

dataset = Dataset(
loader = DataLoader(
for batch in data_loader:
    predicts = model(

Using a ML model in Spark SQL (experiemental)

OPTIONS (min_confidence=0.3, device="gpu", batch_size=32)
USING "s3://bucket/to/yolo5_spec.yaml";

SELECT id, ML_PREDICT(yolo5, image) FROM my_dataset
WHERE split = "train" LIMIT 100;

Getting Started

Currently Rikai is maintained for Scala 2.12 and Python 3.7 and 3.8.

There are multiple ways to install Rikai:

  1. Try it using the included Dockerfile.
  2. OR install it via pip pip install rikai, with extras for gcp, pytorch/tf, and others.
  3. OR install it from source

Note: if you want to use Rikai with your own pyspark, please consult rikai documentation for tips.


The included Dockerfile creates a standalone demo image with Jupyter, Pytorch, Spark, and rikai preinstalled with notebooks for you to play with the capabilities of the rikai feature store.

To build and run the docker image from the current directory:

# Clone the repo
git clone rikai
# Build the docker image
docker build --tag rikai --network host .
# Run the image
docker run -p rikai:latest jupyter lab -ip --port 8888

If successful, the console should then print out a clickable link to JupyterLab. You can also open a browser tab and go to localhost:8888.

Install from pypi

Base rikai library can be installed with just pip install rikai. Dependencies for supporting pytorch (pytorch and torchvision), jupyter (matplotlib and jupyterlab) are all part of optional extras. Many open-source datasets also use Youtube videos so we've also added pafy and youtube-dl as optional extras as well.

For example, if you want to use pytorch in Jupyter to train models on rikai datasets in s3 containing Youtube videos you would run:

pip install rikai[pytorch,jupyter,youtube]

If you're not sure what you need and don't mind installing some extra dependencies, you can simply install everything:

pip install rikai[all]

Install from source

To build from source you'll need python as well as Scala with sbt installed:

# Clone the repo
git clone rikai
# Build the jar
sbt publishLocal
# Install python package
cd python
pip install -e . # pip install -e .[all] to install all optional extras (see "Install from pypi")

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

rikai-0.0.4.tar.gz (38.8 kB view hashes)

Uploaded Source

Built Distribution

rikai-0.0.4-py3-none-any.whl (67.3 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page