diff --git a/init/init.cpp b/init/init.cpp
index 17498da..5b0b0dd 100644
--- a/init/init.cpp
+++ b/init/init.cpp
@@ -315,8 +315,7 @@
         if (apex_info_list.has_value()) {
             std::vector<std::string> subcontext_apexes;
             for (const auto& info : apex_info_list->getApexInfo()) {
-                if (info.hasPreinstalledModulePath() &&
-                    subcontext->PathMatchesSubcontext(info.getPreinstalledModulePath())) {
+                if (subcontext->PartitionMatchesSubcontext(info.getPartition())) {
                     subcontext_apexes.push_back(info.getModuleName());
                 }
             }
diff --git a/init/subcontext.cpp b/init/subcontext.cpp
index 6a095fb..3fe448f 100644
--- a/init/subcontext.cpp
+++ b/init/subcontext.cpp
@@ -263,6 +263,10 @@
     return false;
 }
 
+bool Subcontext::PartitionMatchesSubcontext(const std::string& partition) const {
+    return std::find(partitions_.begin(), partitions_.end(), partition) != partitions_.end();
+}
+
 void Subcontext::SetApexList(std::vector<std::string>&& apex_list) {
     apex_list_ = std::move(apex_list);
 }
@@ -352,12 +356,13 @@
     }
 
     if (SelinuxGetVendorAndroidVersion() >= __ANDROID_API_P__) {
-        subcontext.reset(
-                new Subcontext(std::vector<std::string>{"/vendor", "/odm"}, kVendorContext));
+        subcontext.reset(new Subcontext(std::vector<std::string>{"/vendor", "/odm"},
+                                        std::vector<std::string>{"VENDOR", "ODM"}, kVendorContext));
     }
 }
+
 void InitializeHostSubcontext(std::vector<std::string> vendor_prefixes) {
-    subcontext.reset(new Subcontext(vendor_prefixes, kVendorContext, /*host=*/true));
+    subcontext.reset(new Subcontext(vendor_prefixes, {}, kVendorContext, /*host=*/true));
 }
 
 Subcontext* GetSubcontext() {
diff --git a/init/subcontext.h b/init/subcontext.h
index 93ebace..23c4a24 100644
--- a/init/subcontext.h
+++ b/init/subcontext.h
@@ -36,8 +36,10 @@
 
 class Subcontext {
   public:
-    Subcontext(std::vector<std::string> path_prefixes, std::string_view context, bool host = false)
+    Subcontext(std::vector<std::string> path_prefixes, std::vector<std::string> partitions,
+               std::string_view context, bool host = false)
         : path_prefixes_(std::move(path_prefixes)),
+          partitions_(std::move(partitions)),
           context_(context.begin(), context.end()),
           pid_(0) {
         if (!host) {
@@ -49,6 +51,7 @@
     Result<std::vector<std::string>> ExpandArgs(const std::vector<std::string>& args);
     void Restart();
     bool PathMatchesSubcontext(const std::string& path) const;
+    bool PartitionMatchesSubcontext(const std::string& partition) const;
     void SetApexList(std::vector<std::string>&& apex_list);
 
     const std::string& context() const { return context_; }
@@ -59,6 +62,7 @@
     Result<SubcontextReply> TransmitMessage(const SubcontextCommand& subcontext_command);
 
     std::vector<std::string> path_prefixes_;
+    std::vector<std::string> partitions_;
     std::vector<std::string> apex_list_;
     std::string context_;
     pid_t pid_;
diff --git a/init/subcontext_benchmark.cpp b/init/subcontext_benchmark.cpp
index ccef2f3..172ee31 100644
--- a/init/subcontext_benchmark.cpp
+++ b/init/subcontext_benchmark.cpp
@@ -33,7 +33,7 @@
         return;
     }
 
-    auto subcontext = Subcontext({"path"}, context);
+    auto subcontext = Subcontext({"path"}, {"partition"}, context);
     free(context);
 
     while (state.KeepRunning()) {
diff --git a/init/subcontext_test.cpp b/init/subcontext_test.cpp
index da1f455..85a2f2a 100644
--- a/init/subcontext_test.cpp
+++ b/init/subcontext_test.cpp
@@ -41,7 +41,7 @@
 
 template <typename F>
 void RunTest(F&& test_function) {
-    auto subcontext = Subcontext({"dummy_path"}, kTestContext);
+    auto subcontext = Subcontext({"dummy_path"}, {"dummy_partition"}, kTestContext);
     ASSERT_NE(0, subcontext.pid());
 
     test_function(subcontext);
@@ -177,6 +177,19 @@
     });
 }
 
+TEST(subcontext, PartitionMatchesSubcontext) {
+    RunTest([](auto& subcontext) {
+        static auto& existent_partition = "dummy_partition";
+        static auto& non_existent_partition = "not_dummy_partition";
+
+        auto existent_result = subcontext.PartitionMatchesSubcontext(existent_partition);
+        auto non_existent_result = subcontext.PartitionMatchesSubcontext(non_existent_partition);
+
+        ASSERT_TRUE(existent_result);
+        ASSERT_FALSE(non_existent_result);
+    });
+}
+
 BuiltinFunctionMap BuildTestFunctionMap() {
     // For CheckDifferentPid
     auto do_return_pids_as_error = [](const BuiltinArguments& args) -> Result<void> {