AWS Step Function - Associating inputs and outputs across multiple Map states

Question

Say I have a simple input:

{
  "input_array": ["blue", "green", "orange"]
}

Which I am passing into a state machine with a couple of Map states:

{
  "StartAt": "Step One",
  "States": {
    "Step One": {
      "Type": "Map",
      "ItemsPath": "$.input_array",
      "ResultPath": "$.step_one_output",
      "Next": "Step Two",
      "Iterator": {
        "StartAt": "Find fruit matching colour",
        "States": {
          "FindFruitMatchingColour": {
            "Type": "Task",
            "Resource": "[some resource]",
            "End": true
          }
        }
      }
    },
    "Step Two": {
      "Type": "Map",
      "ItemsPath": "$.step_one_output",
      "ResultPath": "$.step_two_output",
      "Next": "End",
      "Iterator": {
        "StartAt": "Do Thing With Colour And Its Fruit",
        "States": {
          "DoThingWithColourAndFruitCombination": {
            "Type": "Task",
            "Resource": "[some resource]",
            "End": true
          }
        }
      }
    },
    "End": {
      "Type": "Succeed"
    }
  }
}

What I'd like to do is associate the input of step one with its corresponding output and pass them together into step two. For example:

[
  {"colour": "blue", "fruit": "blueberry"}, 
  {"colour": "green", "fruit": "apple"}, 
  {"colour": "orange", "fruit": "orange"}
]

It feels like a fairly common use case.

I've tried a few ways to solve this, such as using

{
  "ResultSelector": {
    "fruit.$": "$",
    "colour.$": "$$.Map.Item"
  }
}

But I cannot seem to access the individual iterations in ResultSelector, it appears to be in the context of the overall output array. I've also explored having the two steps part of the same Iterator, but again struggle to associate the individual inputs to the outputs of step one

Alternatively, I've considered having a separate step to zip together the input and the output of step one together before passing to step two. But I'm concerned about not being able to rely on the order of the output of step one and mixing up colours and fruits.

Many thanks for any guidance provided, even if to say it's not supported. I can always have the resource return both the fruit and the colour, but this is not its current behaviour and I'd like to delegate this to the state machine if possible.

P.S - this is my first SO issue. Don't hesitate to ask for more detail.

Answer 1

Based on Justin's help, I think I have come up with the following solution:

{
  "StartAt": "Step One",
  "States": {
    "Step One": {
      "Type": "Map",
      "ItemsPath": "$.input_array",
      "ResultPath": "$.step_one_output",
      "Next": "Step Two",
      "Iterator": {
        "StartAt": "Find fruit matching colour",
        "States": {
          "FindFruitMatchingColour": {
            "Type": "Task",
            "Resource": "[some resource]",
            "End": true
          }
        }
      }
    },
    "Step Two": {
      "Type": "Map",
      "ItemsPath": "$.step_one_output",
      "ResultPath": "$.step_two_output",
      "Parameters": {
        "colour.$": "States.ArrayGetItem($.input_array, $$.Map.Item.Index)",
        "fruit.$": "$$.Map.Item.Value",
      },
      "Next": "End",
      "Iterator": {
        "StartAt": "Do Thing With Colour And Its Fruit",
        "States": {
          "DoThingWithColourAndFruitCombination": {
            "Type": "Task",
            "Resource": "[some resource]",
            "End": true
          }
        }
      }
    },
    "End": {
      "Type": "Succeed"
    }
  }
}

Specifically using the output of step on as the input to a Map task step two, and also selecting the appropriate colour from the original input_array based on the index of the current iteration using the Parameters key in 'Step Two'