"Shape of query is incorrect, it has to be either a 1D array or 2D (1, N)"
CONTRIBUTOR
#3323 
2021-11-26T16:27:22Z
https://github.com/huggingface/datasets/issues/3327#issuecomment-980104741
https://api.github.com/repos/huggingface/datasets/issues/3327
IC_kwDODunzps46azYl
https://api.github.com/repos/huggingface/datasets/issues/comments/980104741/reactions
2021-11-26T16:27:22Z
https://api.github.com/repos/huggingface/datasets/issues/comments/980104741
https://avatars.githubusercontent.com/u/19492473?v=4
https://api.github.com/users/eliasws/events{/privacy}
https://api.github.com/users/eliasws/followers
https://api.github.com/users/eliasws/following{/other_user}
https://api.github.com/users/eliasws/gists{/gist_id}
https://github.com/eliasws
eliasws
MDQ6VXNlcjE5NDkyNDcz
https://api.github.com/users/eliasws/orgs
https://api.github.com/users/eliasws/received_events
https://api.github.com/users/eliasws/repos
https://api.github.com/users/eliasws/starred{/owner}{/repo}
https://api.github.com/users/eliasws/subscriptions
User
https://api.github.com/users/eliasws
public
## 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.
