diff --git a/android/module.go b/android/module.go
index e190e8b..66a5f60 100644
--- a/android/module.go
+++ b/android/module.go
@@ -3163,6 +3163,7 @@
 		symlinkTarget:         "",
 		executable:            executable,
 		effectiveLicenseFiles: &licenseFiles,
+		partition:             fullInstallPath.partition,
 	}
 	m.packagingSpecs = append(m.packagingSpecs, spec)
 	return spec
@@ -3280,6 +3281,7 @@
 		srcPath:          nil,
 		symlinkTarget:    relPath,
 		executable:       false,
+		partition:        fullInstallPath.partition,
 	})
 
 	return fullInstallPath
@@ -3320,6 +3322,7 @@
 		srcPath:          nil,
 		symlinkTarget:    absPath,
 		executable:       false,
+		partition:        fullInstallPath.partition,
 	})
 
 	return fullInstallPath
diff --git a/android/packaging.go b/android/packaging.go
index e3a0b54..ecd84a2 100644
--- a/android/packaging.go
+++ b/android/packaging.go
@@ -40,6 +40,8 @@
 	executable bool
 
 	effectiveLicenseFiles *Paths
+
+	partition string
 }
 
 // Get file name of installed package
@@ -67,6 +69,10 @@
 	return *p.effectiveLicenseFiles
 }
 
+func (p *PackagingSpec) Partition() string {
+	return p.partition
+}
+
 type PackageModule interface {
 	Module
 	packagingBase() *PackagingBase
@@ -76,11 +82,14 @@
 	// be copied to a zip in CopyDepsToZip, `depTag` should implement PackagingItem marker interface.
 	AddDeps(ctx BottomUpMutatorContext, depTag blueprint.DependencyTag)
 
+	// GatherPackagingSpecs gathers PackagingSpecs of transitive dependencies.
+	GatherPackagingSpecs(ctx ModuleContext) map[string]PackagingSpec
+
 	// CopyDepsToZip zips the built artifacts of the dependencies into the given zip file and
 	// returns zip entries in it. This is expected to be called in GenerateAndroidBuildActions,
 	// followed by a build rule that unzips it and creates the final output (img, zip, tar.gz,
 	// etc.) from the extracted files
-	CopyDepsToZip(ctx ModuleContext, zipOut WritablePath) []string
+	CopyDepsToZip(ctx ModuleContext, specs map[string]PackagingSpec, zipOut WritablePath) []string
 }
 
 // PackagingBase provides basic functionality for packaging dependencies. A module is expected to
@@ -211,7 +220,7 @@
 	}
 }
 
-// Returns transitive PackagingSpecs from deps
+// See PackageModule.GatherPackagingSpecs
 func (p *PackagingBase) GatherPackagingSpecs(ctx ModuleContext) map[string]PackagingSpec {
 	m := make(map[string]PackagingSpec)
 	ctx.VisitDirectDeps(func(child Module) {
@@ -229,10 +238,10 @@
 
 // CopySpecsToDir is a helper that will add commands to the rule builder to copy the PackagingSpec
 // entries into the specified directory.
-func (p *PackagingBase) CopySpecsToDir(ctx ModuleContext, builder *RuleBuilder, m map[string]PackagingSpec, dir ModuleOutPath) (entries []string) {
+func (p *PackagingBase) CopySpecsToDir(ctx ModuleContext, builder *RuleBuilder, specs map[string]PackagingSpec, dir ModuleOutPath) (entries []string) {
 	seenDir := make(map[string]bool)
-	for _, k := range SortedStringKeys(m) {
-		ps := m[k]
+	for _, k := range SortedStringKeys(specs) {
+		ps := specs[k]
 		destPath := dir.Join(ctx, ps.relPathInPackage).String()
 		destDir := filepath.Dir(destPath)
 		entries = append(entries, ps.relPathInPackage)
@@ -254,14 +263,13 @@
 }
 
 // See PackageModule.CopyDepsToZip
-func (p *PackagingBase) CopyDepsToZip(ctx ModuleContext, zipOut WritablePath) (entries []string) {
-	m := p.GatherPackagingSpecs(ctx)
+func (p *PackagingBase) CopyDepsToZip(ctx ModuleContext, specs map[string]PackagingSpec, zipOut WritablePath) (entries []string) {
 	builder := NewRuleBuilder(pctx, ctx)
 
 	dir := PathForModuleOut(ctx, ".zip")
 	builder.Command().Text("rm").Flag("-rf").Text(dir.String())
 	builder.Command().Text("mkdir").Flag("-p").Text(dir.String())
-	entries = p.CopySpecsToDir(ctx, builder, m, dir)
+	entries = p.CopySpecsToDir(ctx, builder, specs, dir)
 
 	builder.Command().
 		BuiltTool("soong_zip").
diff --git a/android/packaging_test.go b/android/packaging_test.go
index ff7446c..91ac1f3 100644
--- a/android/packaging_test.go
+++ b/android/packaging_test.go
@@ -95,7 +95,7 @@
 
 func (m *packageTestModule) GenerateAndroidBuildActions(ctx ModuleContext) {
 	zipFile := PathForModuleOut(ctx, "myzip.zip")
-	m.entries = m.CopyDepsToZip(ctx, zipFile)
+	m.entries = m.CopyDepsToZip(ctx, m.GatherPackagingSpecs(ctx), zipFile)
 }
 
 func runPackagingTest(t *testing.T, multitarget bool, bp string, expected []string) {
diff --git a/filesystem/filesystem.go b/filesystem/filesystem.go
index 0796258..9cf7bbd 100644
--- a/filesystem/filesystem.go
+++ b/filesystem/filesystem.go
@@ -226,7 +226,7 @@
 
 func (f *filesystem) buildImageUsingBuildImage(ctx android.ModuleContext) android.OutputPath {
 	depsZipFile := android.PathForModuleOut(ctx, "deps.zip").OutputPath
-	f.CopyDepsToZip(ctx, depsZipFile)
+	f.CopyDepsToZip(ctx, f.GatherPackagingSpecs(ctx), depsZipFile)
 
 	builder := android.NewRuleBuilder(pctx, ctx)
 	depsBase := proptools.StringDefault(f.properties.Base_dir, ".")
@@ -345,7 +345,7 @@
 	}
 
 	depsZipFile := android.PathForModuleOut(ctx, "deps.zip").OutputPath
-	f.CopyDepsToZip(ctx, depsZipFile)
+	f.CopyDepsToZip(ctx, f.GatherPackagingSpecs(ctx), depsZipFile)
 
 	builder := android.NewRuleBuilder(pctx, ctx)
 	depsBase := proptools.StringDefault(f.properties.Base_dir, ".")
diff --git a/snapshot/host_snapshot.go b/snapshot/host_snapshot.go
index 75b5e88..9793218 100644
--- a/snapshot/host_snapshot.go
+++ b/snapshot/host_snapshot.go
@@ -151,7 +151,7 @@
 
 	f.installDir = android.PathForModuleInstall(ctx)
 
-	f.CopyDepsToZip(ctx, depsZipFile)
+	f.CopyDepsToZip(ctx, f.GatherPackagingSpecs(ctx), depsZipFile)
 
 	builder := android.NewRuleBuilder(pctx, ctx)
 	builder.Command().