En ocasiones es necesario pasar a una función de transformación información en la que buscar datos necesarios para la transformación y el resultado.

En estos casos, se pueden pasar dichos datos como List o idealmente como HashMap si vas a conocer el valor de la clave por la que buscar, ya que así será mucho más rápido.

HashMap map2 = new HashMap<>();

Dataset<Row> ds = sparkSession.sql(select * from db.tableName);
ds.collectAsList().forEach((Row row) -> {

String id = row.get(1).toString();
if (map2.containsKey(id)) {
	((ArrayList) map2.get(id)).add(row);
	} else {
		List<Row> ls = new ArrayList<Row>();
		ls.add(row);
		map2.put(id, ls);
	}
});

En este caso lo que se va a guardar dentro del HashMap es un listado de los Row que tienen el mismo Id