"Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)"
{'login': 'eliasws', 'id': 19492473, 'node_id': 'MDQ6VXNlcjE5NDkyNDcz', 'avatar_url': 'https://avatars.githubusercontent.com/u/19492473?v=4', 'gravatar_id': '', 'url': 'https://api.github.com/users/eliasws', 'html_url': 'https://github.com/eliasws', 'followers_url': 'https://api.github.com/users/eliasws/followers', 'following_url': 'https://api.github.com/users/eliasws/following{/other_user}', 'gists_url': 'https://api.github.com/users/eliasws/gists{/gist_id}', 'starred_url': 'https://api.github.com/users/eliasws/starred{/owner}{/repo}', 'subscriptions_url': 'https://api.github.com/users/eliasws/subscriptions', 'organizations_url': 'https://api.github.com/users/eliasws/orgs', 'repos_url': 'https://api.github.com/users/eliasws/repos', 'events_url': 'https://api.github.com/users/eliasws/events{/privacy}', 'received_events_url': 'https://api.github.com/users/eliasws/received_events', 'type': 'User', 'site_admin': False}
[{'id': 1935892857, 'node_id': 'MDU6TGFiZWwxOTM1ODkyODU3', 'url': 'https://api.github.com/repos/huggingface/datasets/labels/bug', 'name': 'bug', 'color': 'd73a4a', 'default': True, 'description': "Something isn't working"}]
[]
#3323 
## Describe the bug
Passing a correctly shaped Numpy-Array to get_nearest_examples leads to the Exception

"Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)"

Probably the reason for this is a wrongly converted assertion.

1.15.1:

`assert len(query.shape) == 1 or (len(query.shape) == 2 and query.shape[0] == 1)`

1.16.1:

```
 if len(query.shape) != 1 or (len(query.shape) == 2 and query.shape[0] != 1):
            raise ValueError("Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)")
```

## Steps to reproduce the bug

follow the steps described here: https://huggingface.co/course/chapter5/6?fw=tf

```python
 question_embedding.shape # (1, 768)

 scores, samples = embeddings_dataset.get_nearest_examples(
    "embeddings", question_embedding, k=5 # Error
)

# "Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)"
```

## Expected results
Should work without exception

## Actual results
Throws exception

## Environment info
- `datasets` version: 1.15.1
- Platform: Darwin-20.6.0-x86_64-i386-64bit
- Python version: 3.7.12
- PyArrow version: 6.0.

{'url': 'https://api.github.com/repos/huggingface/datasets/issues/3327/reactions', 'total_count': 0, '+1': 0, '-1': 0, 'laugh': 0, 'hooray': 0, 'confused': 0, 'heart': 0, 'rocket': 0, 'eyes': 0}