diff --git a/setup.sh b/setup.sh
index ee3a45a..6a0fe0a 100644
--- a/setup.sh
+++ b/setup.sh
@@ -52,14 +52,20 @@ function get_partitions () {
 }
 
 function get_parts () {
+    local zfs_install_drive
     declare parttype parts
     parttype="${1:?}"
+    zfs_install_drive="${2:-}"
     case "${parttype}" in
         zfs)
             parts="$(get_partitions | jq --raw-output '.[][] | select(.parttype=="6a85cf4d-1dd2-11b2-99a6-080020736631") | .path')" || exit 1
             ;;
         efi)
-            parts="$(get_partitions | jq --raw-output '.[][] | select(.parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b") | .path')" || exit 1
+            if [[ "${zfs_install_drive}" ]]; then
+                parts="$(get_partitions | jq --raw-output '.[][] | select(.path=="'"${zfs_install_drive}"'") | select(.parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b") | .path')" || exit 1
+            else
+                parts="$(get_partitions | jq --raw-output '.[][] | select(.parttype=="c12a7328-f81f-11d2-ba4b-00a0c93ec93b") | .path')" || exit 1
+            fi
             ;;
         *)
             >2 printf -- '%s\n' 'Unknown partition type '"'"'"${parttype}"'"'"', exiting ...'
@@ -106,10 +112,15 @@ function get_drive_id () {
 }
 
 function select_part () {
-    local parts enriched_parts enriched_parts_count part_number part_type
+    local parts enriched_parts enriched_parts_count part_number part_type zfs_install_drive zfs_install_parent
     declare part
     part_type="${1:?}" # 'efi' or 'zfs'
-    parts="$(get_parts "${part_type}")"
+    zfs_install_drive="${2:-}"
+    if [[ "${zfs_install_drive}" ]]; then
+        parts="$(get_parts "${part_type}" "${zfs_install_drive}")"
+    else
+        parts="$(get_parts "${part_type}")"
+    fi
 
     if we_have_exactly_one_part "${part_type}" "${parts}"; then
         part="${parts}"
@@ -221,7 +232,7 @@ function mount_system () {
     zfs mount "${zpool_name}"'/root/'"${zfs_arch_dataset_name}"
     zfs mount -a
 
-    local efi_part="$(select_part 'efi')"
+    local efi_part="$(select_part 'efi' "${zpool_drive:?})"
     mkdir -p '/mnt/efi'
     mount "${efi_part}" '/mnt/efi'
 }