diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 18e5499..38ed1d1 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -35,11 +35,3 @@ jobs:
         run: |
           gcloud auth configure-docker
           docker push gcr.io/kf-feast/spark-py:v3.0.1
-      - name: Build Docker (with extra dependencies)
-        run: | 
-          cd k8s_spark
-          docker build -t gcr.io/kf-feast/spark-py-extra:v3.0.1 --build-arg BASE_IMAGE=gcr.io/kf-feast/spark-py-base/spark-py:v3.0.1 .
-      - name: Push Docker (with extra dependencies)
-        run: |
-          gcloud auth configure-docker
-          docker push gcr.io/kf-feast/spark-py-extra:v3.0.1
diff --git a/aws_spark/Dockerfile b/aws_spark/Dockerfile
index 5911086..d1ffb7e 100644
--- a/aws_spark/Dockerfile
+++ b/aws_spark/Dockerfile
@@ -4,10 +4,12 @@ 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/k8s_spark/.Dockerfile.swp b/k8s_spark/.Dockerfile.swp
deleted file mode 100644
index 81490df..0000000
Binary files a/k8s_spark/.Dockerfile.swp and /dev/null differ
diff --git a/k8s_spark/Dockerfile b/k8s_spark/Dockerfile
deleted file mode 100644
index 5d4bfd0..0000000
--- a/k8s_spark/Dockerfile
+++ /dev/null
@@ -1,6 +0,0 @@
-ARG BASE_IMAGE=gcr.io/kf-feast/spark-py-base:v3.0.1
-
-FROM ${BASE_IMAGE}
-
-USER root
-ADD https://repo1.maven.org/maven2/com/linkedin/sparktfrecord/spark-tfrecord_2.12/0.3.0/spark-tfrecord_2.12-0.3.0.jar /opt/spark/jars