diff --git a/build_kzip.bash b/build_kzip.bash
index 0018ea9..9564723 100755
--- a/build_kzip.bash
+++ b/build_kzip.bash
@@ -7,14 +7,16 @@
 #   BUILD_NUMBER          build number, used to generate unique ID (will use UUID if not set)
 #   DIST_DIR              where the resulting all.kzip will be placed
 #   KYTHE_KZIP_ENCODING   proto or json (proto is default)
+#   KYTHE_JAVA_SOURCE_BATCH_SIZE maximum number of the Java source files in a compilation unit
 #   OUT_DIR               output directory (out if not specified})
 #   TARGET_BUILD_VARIANT  variant, e.g., `userdebug`
 #   TARGET_PRODUCT        target device name, e.g., 'aosp_blueline'
 #   XREF_CORPUS           source code repository URI, e.g., 'android.googlesource.com/platform/superproject'
 
 : ${BUILD_NUMBER:=$(uuidgen)}
+: ${KYTHE_JAVA_SOURCE_BATCH_SIZE:=500}
 : ${KYTHE_KZIP_ENCODING:=proto}
-export KYTHE_KZIP_ENCODING
+export KYTHE_JAVA_SOURCE_BATCH_SIZE KYTHE_KZIP_ENCODING
 
 # The extraction might fail for some source files, so run with -k and then check that
 # sufficiently many files were generated.