def METHOD_NAME(statement):
    """Get entities in statement that are referenced as columns.
    Examples::
        >> get_column_entities(query(A)) == {A}
        >> get_column_entities(query(A.field)) == {A}
        >> get_column_entities(query(A, B)) == {A, B})
    Does not include eager loaded entities.
    """
    def _entities_in_statement(statement):
        try:
            entities = (cd.get("entity") for cd in statement.column_descriptions)
            return {e for e in entities if e is not None}
        except AttributeError:
            return set()
    entities = _entities_in_statement(statement)
    # TODO(gj): currently walking way more than we have to. Probably
    # some points in the tree where we can safely call it good for that
    # branch and continue on to more fruitful pastures.
    for child in statement.get_children():
        entities |= METHOD_NAME(child)
    return entities