fix: route in-memory shuffle partitions to remote fetch when not present locally

When using in-memory shuffle (shuffle_location: memory), shuffle partitions are
stored in the executor's local InMemoryShuffleManager. Previously, the shuffle
reader incorrectly assumed all memory:// paths could be read locally.

In distributed mode with multiple executors, a shuffle reader on executor A may
need to read a partition that was written by executor B. The partition would
have a memory:// path, but only exist in executor B's shuffle manager. This
caused 'Shuffle partition not found in memory' errors.

This fix modifies split_partition_locations() to check if a memory:// partition
actually exists in the local shuffle manager before categorizing it for local
read. If the partition doesn't exist locally, it's routed to the remote fetch
path which uses the Arrow Flight service. The Flight service already correctly
handles memory:// paths by reading from the executor's local shuffle manager.

Fixes spiceai/spiceai#9290

Author: lukekim

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -425,8 +425,19 @@ fn split_partition_locations(
 
     for loc in partition_locations {
         if check_is_memory_location(&loc) {
-            // Memory locations are always read locally
-            result.memory.push(loc);
+            // Memory locations should only be read locally if they exist in this executor's
+            // shuffle manager. If the partition doesn't exist locally, it means the partition
+            // was written by another executor and we need to fetch it remotely via Flight.
+            if check_is_local_memory_location(&loc) {
+                result.memory.push(loc);
+            } else {
+                // Partition is in memory on another executor, fetch remotely
+                debug!(
+                    "Memory partition {} not found locally, will fetch remotely from executor {}",
+                    loc.path, loc.executor_meta.id
+                );
+                result.remote.push(loc);
+            }
         } else if check_is_object_store_location(&loc) {
             // Object store locations are handled via the runtime_env's registered object stores
             result.object_store.push(loc);
@@ -598,6 +609,16 @@ fn check_is_memory_location(location: &PartitionLocation) -> bool {
     location.path.starts_with("memory://")
 }
 
+/// Check if a memory:// partition actually exists in the local shuffle manager.
+/// This is used to determine whether to read the partition locally or fetch it remotely.
+fn check_is_local_memory_location(location: &PartitionLocation) -> bool {
+    if let Some(key) = location.path.strip_prefix("memory://") {
+        global_shuffle_manager().contains_partition(key)
+    } else {
+        false
+    }
+}
+
 /// Partition reader Trait, different partition reader can have
 #[async_trait]
 trait PartitionReader: Send + Sync + Clone {