Package overhaul: update dependencies, improve load time, and more (#48)

## Problem

A lot of dependencies are far out of date, preventing this package from
being run on modern jupyter notebooks and creating incompatibilities
with updated SDK libraries. Also, this package loads very slowly despite
not doing very much.

It wasn't my intention to change everything in one PR, but since the
tests weren't runnable I couldn't assess the impact of the dependency
changes without updating and adding tests. And once you start doing that
you naturally want to modify other parts of the code.

## Solution

- Upgrade all dependencies in pyproject.toml and actions in .github
workflows
- Tried running tests to assess what broke after upgrading deps, and
found myself struggling because the test coverage and organization isn't
great. For example, there were tests under "unit tests" which can't run
without S3 credentials for a private bucket (which I don't have).
Problems like these, and a lack of specific testing for the
`load_dataset` and `list_datasets` functions, meant I definitely needed
to expand testing.

Added tests:
- For `load_dataset` and `list_datasets` functions which are the main
way people use this package from the examples repository.
- Added tests for working with local catalogs since that seemed like an
intended behavior of this package but wasn't clearly tested
- Added tests for uploading datasets to a google storage bucket using
service account credentials, which seems important since that is how we
would need to update or add additional datasets to the public set.

Refactoring (non-breaking):
- I ended up refactoring and adding a lot of tests to scope down the
responsibility of the Dataset class by extracting groups of
functionality into smaller more focused classes: `DatasetFSWriter`,
`DatasetFSReader`. This seemed faster than trying to reason about
everything smashed together into one giant `Dataset` class.
- Use lazy-loading for heavy dependencies such as `pandas`, `gcsfs`, and
`s3fs`. With these changes, importing the `load_dataset` function is now
about 8x faster, with import time dropping from a measured ~1.849s down
to 0.230s.
- Incorporated Dave Rigby's suggestion to use the `fs.glob` function
when listing datasets in a Catalog; this significantly cuts down the
number of network calls needed to build the catalog list.
- When calling `load_dataset`, skip loading all metadata. In the past
the entire catalog metadata was loaded to check if a dataset exists
before trying to build the `Dataset` object, which is unnecessary and
just adds a ton of overhead. If you try to load a dataset that doesn't
exist, the error message is already pretty clear.

Breaking changes:
- Removed `to_pinecone_index`. Having this in creates a coupling between
this package and the SDK package that is going to be a perpetual thorn
in our side when it comes to maintaining docs and examples. The SDK will
continue marching forward whereas most of the other logic in here
related to uploading and downloading from buckets into dataframes should
not change much. There's really no need for it in here, so I am removing
it.
- `Dataset.to_catalog` and `Dataset.from_catalog` now error and tells
you to use `Catalog.save_dataset` and `Catalog.load_dataset`. I'm not
aware of anyone using these legacy methods, but it was a change that
felt right to do to make the Catalog class mostly responsible for
"where" things are saved to while the Dataset class is responsible for
"what" things are saved. Now you can easily download a dataset and save
it to local, for example, which wasn't something you could easily reason
about when the dataset itself was responsible for writing to a catalog.

Despite all these changes, little should change for most users of the
package aside from dependency updates. `load_dataset` and
`list_datasets` are still the same (although much more performant) and
they are by far the most used.

## Usage

Most people who just want to load a demo dataset will do something like
this

```python
from pinecone import Pinecone
from pinecone_dataset import load_dataset

ds = load_dataset('dataset_name')

pc = Pinecone(api_key='key')
index = pc.Index(host='host')
index.upsert_from_dataframe(df=ds.documents, batch_size=100)
```

## Type of Change

- [x] Bug fix (non-breaking change which fixes an issue)
- [x] Breaking change (fix or feature that would cause existing
functionality to not work as expected)

## Testing load performance

To investigate load performance, I used a python feature called
`importtime` and a cool package I added as a dev dependency called
`tuna` for visualizing the outputs.

```
poetry run python3 -X importtime -c "from pinecone_datasets import load_dataset" 2> load_times.log
poetry run tuna load_times.log
```

## Testing

Besides the tests you see added here, I did some manual testing in a
notebook setting using the `1.0.0.dev3` build that I built from this
branch.

Author: jhamon

.github/workflows/PR.yml

@@ -3,6 +3,7 @@ name: CD
 on:
   workflow_dispatch:
 
+
 jobs:
 
   release:
@@ -17,7 +18,7 @@ jobs:
     - name: Install Poetry
       uses: snok/install-poetry@v1
       with:
-        version: 1.3.2
+        version: 1.5.0
 
     - name: Set Version
       run: echo "VERSION=$(poetry version -s)" >> $GITHUB_ENV

.github/workflows/cd.yml

@@ -1,41 +1,89 @@
 name: CI
 
-on: push
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+      branches:
+        - main
 
 jobs:
-  run-tests:
-    name: Run tests
+  linting:
+    name: Run lint and type checking
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
-        python-version: [3.8, 3.9, '3.10', 3.11]
-
+        python-version: ['3.10']
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
 
     - name: Install Poetry
       uses: snok/install-poetry@v1
       with:
-        version: 1.3.2
+        version: 1.5.0
     - name: install dependencies
       run: poetry install --with dev --all-extras
+    
     - name: Run Black Check
       run: poetry run black --check .
+    
     - name: Run mypy check
       run: poetry run mypy .
-    - name: Run pytest
+
+  run-tests:
+    name: Run tests
+    needs: linting
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [3.9, '3.10', 3.11, 3.12, 3.13]
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install Poetry
+      uses: snok/install-poetry@v1
+      with:
+        version: 1.5.0
+    - name: install dependencies
+      run: poetry install --with dev --all-extras
+
+    - name: Run pytest (unit tests)
       env:
         PY_VERSION: ${{ matrix.python-version }}
-        AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
-        AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        # AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+        # AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
       run: poetry run pytest -n 4 --html=report.html --cov pinecone_datasets tests/unit
+    
     - name: upload pytest report.html
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       if: always()
       with:
         name: dataset-pytest-report-py${{ matrix.python-version }}
-        path: report.html
+        path: report.html
+
+    - name: Write google service account credentials to a file
+      id: prepare-google-credentials
+      shell: bash
+      run: |
+        secrets_file="$(mktemp)"
+        echo "$GCS_SERVICE_ACCOUNT_CREDS_BASE64" | base64 -d > $secrets_file
+        echo "google_credentials_file=$secrets_file" >> $GITHUB_OUTPUT
+      env:
+        GCS_SERVICE_ACCOUNT_CREDS_BASE64: '${{ secrets.GCS_SERVICE_ACCOUNT_CREDS_BASE64 }}'
+
+    - name: Run pytest (integration tests)
+      run: poetry run pytest tests/integration
+      env:
+        GOOGLE_APPLICATION_CREDENTIALS: ${{ steps.prepare-google-credentials.outputs.google_credentials_file }}

.github/workflows/ci.yml

@@ -12,14 +12,14 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: '3.10'
       - name: Install Poetry
         uses: snok/install-poetry@v1
         with:
-          version: 1.3.2
+          version: 1.5.0
 
       - run: poetry install --with dev --all-extras
       # ADJUST THIS: build your documentation into docs/.

.github/workflows/docs.yml

@@ -0,0 +1,172 @@
+# Pinecone Datasets
+
+### Supported storage options
+
+pinecone_datasets can load datasets from Google Cloud storage, Amazon S3, and local files.
+
+By default, the `load_dataset` and `list_datasets` packages will pull from Pinecone's public GCS bucket at `gs://pinecone-datasets-dev`, but you can interact with catalogs stored in other locations.
+
+```python
+from pinecone_datasets import Catalog
+
+# Local catalog
+catalog = Catalog(base_path="/path/to/local/catalog")
+catalog.list_datasets()
+
+# Google Cloud
+catalog = Catalog(base_path="gs://bucket-name")
+
+# S3 catalog
+s3_catalog = Catalog(base_path="s3://bucket-name")
+```
+
+If you are using Amazon S3 or Google Cloud to access private buckets, you can use environment variables to configure your credentials. For example, if you set a base_path starting with "gs://", the `gcsfs` package will attempt to find credentials by looking in cache locations used by `gcloud auth login` or reading environment variables such as `GOOGLE_APPLICATION_CREDENTIALS`.
+
+## Adding a new dataset to the public datasets repo
+
+Note: Only Pinecone employees with access to the bucket can complete this step.
+
+Prerequisites:
+
+1. Install google cloud CLI
+2. Authenticate with `gcloud auth login`
+
+```python
+from pinecone_datasets import Catalog, Dataset, DatasetMetadata, DenseModelMetadata
+
+# 1. Prepare pandas dataframes containing your embeddings
+documents_df = ...
+queries_df = ...
+
+# 2. Create metadata to describe the dataset
+import datatime
+metadata = DatasetMetadata(
+    name="new-dataset-name",
+    created_at=datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
+    documents=len(documents_df),
+    queries=len(queries_df),
+    dense_model=DenseModelMetadata(
+        name="ada2",
+        dimension=2,
+    ),
+)
+
+# 3. Take all this, and instantiate a Dataset
+ds = Dataset.from_pandas(
+    documents=documents_df,
+    queries=queries_df,
+    metadata=metadata
+)
+
+# 4. Save to catalog (requires gcloud auth step above)
+catalog = Catalog(base_path="gs://pinecone-datasets-dev")
+catalog.save_dataset(ds)
+```
+
+Afterwards, verify the new dataset appears in list function and can be used
+
+```python
+from pinecone_datasets import list_datasets, load_dataset
+
+list_datasets(as_df=True)
+
+ds = load_dataset("new-dataset-name")
+ds.documents
+ds.head()
+```
+
+### Expected dataset structure
+
+The package expects data to be laid out with the following directory structure:
+
+    ├── my-subdir                     # path to where all datasets
+    │   ├── my-dataset                # name of dataset
+    │   │   ├── metadata.json         # dataset metadata (optional, only for listed)
+    │   │   ├── documents             # datasets documents
+    │   │   │   ├── file1.parquet      
+    │   │   │   └── file2.parquet      
+    │   │   ├── queries               # dataset queries
+    │   │   │   ├── file1.parquet  
+    │   │   │   └── file2.parquet   
+    └── ...
+
+The data schema is expected to be as follows:
+
+- `documents` directory contains parquet files with the following schema:
+    - Mandatory: `id: str, values: list[float]`
+    - Optional: `sparse_values: Dict: indices: List[int], values: List[float]`, `metadata: Dict`, `blob: dict`
+        - note: blob is a dict that can contain any data, it is not returned when iterating over the dataset and is inteded to be used for storing additional data that is not part of the dataset schema. however, it is sometime useful to store additional data in the dataset, for example, a document text. In future version this may become a first class citizen in the dataset schema.
+- `queries` directory contains parquet files with the following schema:
+    - Mandatory: `vector: list[float], top_k: int`
+    - Optional: `sparse_vector: Dict: indices: List[int], values: List[float]`, `filter: Dict`
+        - note: filter is a dict that contain pinecone filters, for more information see [here](https://docs.pinecone.io/docs/metadata-filtering)
+
+in addition, a metadata file is expected to be in the dataset directory, for example: `s3://my-bucket/my-dataset/metadata.json`
+
+```python
+from pinecone_datasets.catalog import DatasetMetadata
+
+meta = DatasetMetadata(
+    name="test_dataset",
+    created_at="2023-02-17 14:17:01.481785",
+    documents=2,
+    queries=2,
+    source="manual",
+    bucket="LOCAL",
+    task="unittests",
+    dense_model={"name": "bert", "dimension": 3},
+    sparse_model={"name": "bm25"},
+)
+```
+
+full metadata schema can be found in `pinecone_datasets.dataset_metadata.DatasetMetadata.schema`
+
+### The 'blob' column
+
+Pinecone dataset ship with a blob column which is inteneded to be used for storing additional data that is not part of the dataset schema. however, it is sometime useful to store additional data in the dataset, for example, a document text. We added a utility function to move data from the blob column to the metadata column. This is useful for example when upserting a dataset to an index and want to use the metadata to store text data.
+
+```python
+from pinecone_datasets import import_documents_keys_from_blob_to_metadata
+
+new_dataset = import_documents_keys_from_blob_to_metadata(dataset, keys=["text"])
+```
+
+## Usage saving
+
+You can save your dataset to a catalog managed by you or to a local path or a remote path (GCS or S3). 
+
+### Saving a dataset to a Catalog
+
+To set you own catalog endpoint, set the environment variable `DATASETS_CATALOG_BASEPATH` to your bucket. Note that pinecone uses the default authentication method for the storage type (gcsfs for GCS and s3fs for S3).
+
+After this environment variable is set you can save your dataset to the catalog using the `save` function
+
+```python
+from pinecone_datasets import Dataset
+
+metadata = DatasetMetadata(**{"name": "my-dataset", ...})
+```
+
+
+### Saving to Path
+
+You can save your dataset to a local path or a remote path (GCS or S3). Note that pinecone uses the default authentication method for the storage type (gcsfs for GCS and s3fs for S3).
+
+```python
+dataset = Dataset.from_pandas(documents, queries, metadata)
+dataset.to_path("s3://my-bucket/my-subdir/my-dataset")
+```
+
+## Running tests
+
+This project is using poetry for dependency managemet. To start developing, on project root directory run:
+
+```bash
+poetry install --with dev
+```
+
+To run test locally run 
+
+```bash
+poetry run pytest test/unit --cov pinecone_datasets
+```