diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 38ed1d1..fefe70c 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -20,7 +20,7 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-maven-
       - name: Set up Cloud SDK
-        uses: google-github-actions/setup-gcloud@master
+        uses: google-github-actions/setup-gcloud@v0
         with:
           project_id: ${{ secrets.GCP_PROJECT_ID }}
           service_account_key: ${{ secrets.GCP_SA_KEY }}
@@ -29,9 +29,8 @@ jobs:
         run: bash build.sh
       - name: Build Docker
         run: | 
-          cd aws_spark
-          docker build -t gcr.io/kf-feast/spark-py:v3.0.1 --build-arg BASE_IMAGE=gcr.io/kf-feast/spark-py-base/spark-py:v3.0.1 .
+          docker build -t gcr.io/kf-feast/spark-py:v3.0.1-py39 .
       - name: Push Docker
         run: |
           gcloud auth configure-docker
-          docker push gcr.io/kf-feast/spark-py:v3.0.1
+          docker push gcr.io/kf-feast/spark-py:v3.0.1-py39
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..2699bdd
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,23 @@
+ARG BASE_IMAGE=gcr.io/kf-feast/spark-py-base/spark-py:v3.0.1
+
+FROM ${BASE_IMAGE}
+
+ARG HADOOP_AWS_VERSION=3.2.1
+ARG AWS_JAVA_SDK_VERSION=1.11.874
+ARG TFRECORD_VERSION=0.3.0
+ARG GCS_CONNECTOR_VERSION=2.0.1
+ARG BQ_CONNECTOR_VERSION=0.18.1
+
+USER root
+# Add HADOOP_AWS_JAR and AWS_JAVA_SDK
+ADD https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/${HADOOP_AWS_VERSION}/hadoop-aws-${HADOOP_AWS_VERSION}.jar /opt/spark/jars
+ADD https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/${AWS_JAVA_SDK_VERSION}/aws-java-sdk-bundle-${AWS_JAVA_SDK_VERSION}.jar /opt/spark/jars
+ADD https://repo1.maven.org/maven2/com/linkedin/sparktfrecord/spark-tfrecord_2.12/${TFRECORD_VERSION}/spark-tfrecord_2.12-${TFRECORD_VERSION}.jar /opt/spark/jars
+ADD https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop2-${GCS_CONNECTOR_VERSION}.jar /opt/spark/jars
+ADD https://repo1.maven.org/maven2/com/google/cloud/spark/spark-bigquery-with-dependencies_2.12/${BQ_CONNECTOR_VERSION}/spark-bigquery-with-dependencies_2.12-${BQ_CONNECTOR_VERSION}.jar /opt/spark/jars
+ADD https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-azure/3.3.0/hadoop-azure-3.3.0.jar /opt/spark/jars
+ADD https://repo1.maven.org/maven2/com/microsoft/azure/azure-storage/8.6.5/azure-storage-8.6.5.jar /opt/spark/jars
+ADD https://repo1.maven.org/maven2/org/eclipse/jetty/jetty-util-ajax/9.4.26.v20200117/jetty-util-ajax-9.4.26.v20200117.jar /opt/spark/jars
+ADD https://repo1.maven.org/maven2/org/eclipse/jetty/jetty-util/9.4.26.v20200117/jetty-util-9.4.26.v20200117.jar /opt/spark/jars
+
+ENTRYPOINT [ "/opt/entrypoint.sh" ]
diff --git a/README.md b/README.md
index 9f01716..03f6987 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,11 @@
-This builds a container image for Spark 3.0.1 with two tweaks:
+This builds a container image for Spark 3.0.1 with several tweaks:
 
 1. a patch that makes it backwards compatible with spark k8s operator `v1beta2-1.1.2-2.4.5`
 2. added AWS libraries
+3. added TFRecord support
+4. added GCS libraries
+5. added Azure blob storage libraries
+
+To build from scratch, first run `build.sh` to build the base image then build the Dockerfile in this repository.
 
 The image is publicly available as `gcr.io/kf-feast/spark-py:v3.0.1`
diff --git a/aws_spark/Dockerfile b/aws_spark/Dockerfile
deleted file mode 100644
index d1ffb7e..0000000
--- a/aws_spark/Dockerfile
+++ /dev/null
@@ -1,15 +0,0 @@
-ARG BASE_IMAGE=seedjeffwan/spark:v2.4.5
-
-FROM ${BASE_IMAGE}
-
-ARG HADOOP_AWS_VERSION=3.2.1
-ARG AWS_JAVA_SDK_VERSION=1.11.874
-ARG TFRECORD_VERSION=0.3.0
-
-USER root
-# Add HADOOP_AWS_JAR and AWS_JAVA_SDK
-ADD https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/${HADOOP_AWS_VERSION}/hadoop-aws-${HADOOP_AWS_VERSION}.jar /opt/spark/jars
-ADD https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-bundle/${AWS_JAVA_SDK_VERSION}/aws-java-sdk-bundle-${AWS_JAVA_SDK_VERSION}.jar /opt/spark/jars
-ADD https://repo1.maven.org/maven2/com/linkedin/sparktfrecord/spark-tfrecord_2.12/${TFRECORD_VERSION}/spark-tfrecord_2.12-${TFRECORD_VERSION}.jar /opt/spark/jars
-
-ENTRYPOINT [ "/opt/entrypoint.sh" ]
diff --git a/build.sh b/build.sh
old mode 100644
new mode 100755
index 72bcd95..b00bf74
--- a/build.sh
+++ b/build.sh
@@ -2,13 +2,11 @@
 
 set -euxo pipefail
 
-git clone https://github.com/apache/spark.git
+git clone --branch v3.0.1 --single-branch --depth 1 https://github.com/apache/spark.git
 cd spark
 
-git checkout v3.0.1
-
 patch -p1 <../p01.patch
 
 ./build/mvn -Pscala-2.12 -Pkubernetes -DskipTests -Phadoop-3.2 -Dhadoop.version=3.2.1 --no-transfer-progress clean package
 
-./bin/docker-image-tool.sh -r gcr.io/kf-feast/spark-py-base -t v3.0.1 -p resource-managers/kubernetes/docker/src/main/dockerfiles/spark/bindings/python/Dockerfile build
+./bin/docker-image-tool.sh -r gcr.io/kf-feast/spark-py-base -t v3.0.1-py39 -p ../python.Dockerfile build
diff --git a/python.Dockerfile b/python.Dockerfile
new file mode 100644
index 0000000..443ff91
--- /dev/null
+++ b/python.Dockerfile
@@ -0,0 +1,42 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+ARG base_img
+
+FROM $base_img
+WORKDIR /
+
+# Reset to root to run installation tasks
+USER 0
+
+RUN mkdir ${SPARK_HOME}/python
+# TODO: Investigate running both pip and pip3 via virtualenvs
+RUN apt-get update && \
+    apt install -y python3 python3-pip && \
+    pip install --upgrade pip setuptools && \
+    # Removed the .cache to save space
+    rm -r /root/.cache && rm -rf /var/cache/apt/*
+
+COPY python/pyspark ${SPARK_HOME}/python/pyspark
+COPY python/lib ${SPARK_HOME}/python/lib
+
+WORKDIR /opt/spark/work-dir
+ENTRYPOINT [ "/opt/entrypoint.sh" ]
+
+# Specify the User that the actual main process will run as
+ARG spark_uid=185
+USER ${spark_uid}