SAMZA-123: Move topic partition grouping to the AM and generalize