Add support for various polymorphic scenarios in sort_query

- Rename query_entities to get_query_entities
- Add tests for get_query_entities