pinecone_datasets.dataset

  1import glob
  2import sys
  3import os
  4import itertools
  5import time
  6import json
  7import asyncio
  8import warnings
  9from urllib.parse import urlparse
 10from dataclasses import dataclass
 11from importlib.metadata import version
 12
 13import gcsfs
 14import s3fs
 15import pandas as pd
 16from tqdm.auto import tqdm
 17import pyarrow.parquet as pq
 18from pydantic import ValidationError
 19from typing import Any, Generator, Iterator, List, Union, Dict, Optional, Tuple
 20
 21from pinecone_datasets import cfg
 22from pinecone_datasets.catalog import DatasetMetadata
 23from pinecone_datasets.fs import get_cloud_fs, LocalFileSystem
 24
 25import pinecone as pc
 26from pinecone import Index
 27
 28
 29class DatasetInitializationError(Exception):
 30    long_message = """
 31    This dataset was not initialized from path, but from memory, e.g. Dataset.from_pandas(...)
 32    Therefore this dataset cannot be reloaded from path, or use methods that require a path.
 33    If you want to reload a dataset from path, please use the `from_path` method and pass a valid path.
 34    """
 35
 36    def __init__(self, message=long_message):
 37        self.message = message
 38        super().__init__(self.message)
 39
 40
 41# TODO: import from Client
 42@dataclass
 43class UpsertResponse:
 44    upserted_count: int
 45
 46
 47def iter_pandas_dataframe_slices(
 48    df: pd.DataFrame, batch_size, return_indexes
 49) -> Generator[List[Dict[str, Any]], None, None]:
 50    for i in range(0, len(df), batch_size):
 51        if return_indexes:
 52            yield (i, df.iloc[i : i + batch_size].to_dict(orient="records"))
 53        else:
 54            yield df.iloc[i : i + batch_size].to_dict(orient="records")
 55
 56
 57def iter_pandas_dataframe_single(
 58    df: pd.DataFrame,
 59) -> Generator[Dict[str, Any], None, None]:
 60    for i in range(0, len(df), 1):
 61        yield df.iloc[i : i + 1].to_dict(orient="records")[0]
 62
 63
 64class Dataset(object):
 65    @classmethod
 66    def from_path(cls, dataset_path, **kwargs):
 67        """
 68        Create a Dataset object from local or cloud storage
 69        Args:
 70            dataset_path (str): a path to a local or cloud storage path containing a valid dataset.
 71
 72        Returns:
 73            Dataset: a Dataset object
 74        """
 75        return cls(dataset_path=dataset_path, **kwargs)
 76
 77    @classmethod
 78    def from_catalog(cls, dataset_id, catalog_base_path: str = "", **kwargs):
 79        """
 80        Load a dataset from Pinecone's Datasets catalog, or from your own endpoint.
 81
 82        Args:
 83            dataset_id (str): the id of the dataset to load within a catalog
 84            catalog_base_path (str): the catalog's base path. Defaults to DATASETS_CATALOG_BASEPATH environment variable.
 85                                     If neither are set, will use Pinecone's public catalog.
 86
 87        Returns:
 88            Dataset: a Dataset object
 89        """
 90        catalog_base_path = (
 91            catalog_base_path
 92            if catalog_base_path
 93            else os.environ.get("DATASETS_CATALOG_BASEPATH", cfg.Storage.endpoint)
 94        )
 95        dataset_path = os.path.join(catalog_base_path, f"{dataset_id}")
 96        return cls(dataset_path=dataset_path, **kwargs)
 97
 98    @classmethod
 99    def from_pandas(
100        cls,
101        documents: pd.DataFrame,
102        metadata: DatasetMetadata,
103        documents_column_mapping: Optional[Dict] = None,
104        queries: Optional[pd.DataFrame] = None,
105        queries_column_mapping: Optional[Dict] = None,
106        **kwargs,
107    ) -> "Dataset":
108        """
109        Create a Dataset object from a pandas DataFrame
110
111        Args:
112            documents (pd.DataFrame): a pandas DataFrame containing the documents
113            documents_column_mapping (Dict): a dictionary mapping the columns of the documents DataFrame to the Pinecone Datasets Schema
114            queries (pd.DataFrame): a pandas DataFrame containing the queries
115            queries_column_mapping (Dict): a dictionary mapping the columns of the queries DataFrame to the Pinecone Datasets Schema
116
117        Keyword Args:
118            kwargs (Dict): additional arguments to pass to the fsspec constructor
119
120        Returns:
121            Dataset: a Dataset object
122        """
123        clazz = cls(dataset_path=None, **kwargs)
124        clazz._documents = cls._read_pandas_dataframe(
125            documents, documents_column_mapping, cfg.Schema.Names.documents
126        )
127        clazz._queries = cls._read_pandas_dataframe(
128            queries, queries_column_mapping, cfg.Schema.Names.queries
129        )
130        clazz._metadata = metadata
131        return clazz
132
133    @staticmethod
134    def _read_pandas_dataframe(
135        df: pd.DataFrame,
136        column_mapping: Dict[str, str],
137        schema: List[Tuple[str, bool, Any]],
138    ) -> pd.DataFrame:
139        """
140        Reads a pandas DataFrame and validates it against a schema.
141
142        Args:
143            df (pd.DataFrame): the pandas DataFrame to read
144            column_mapping (Dict[str, str]): a dictionary mapping the columns of the DataFrame to the Pinecone Datasets Schema (col_name, pinecone_name)
145            schema (List[Tuple[str, bool]]): the schema to validate against (column_name, is_nullable)
146
147        Returns:
148            pd.DataFrame: the validated, renamed DataFrame
149        """
150        if df is None or df.empty:
151            return pd.DataFrame(columns=[column_name for column_name, _, _ in schema])
152        else:
153            if column_mapping is not None:
154                df.rename(columns=column_mapping, inplace=True)
155            for column_name, is_nullable, null_value in schema:
156                if column_name not in df.columns and not is_nullable:
157                    raise ValueError(
158                        f"error, file is not matching Pinecone Datasets Schmea: {column_name} not found"
159                    )
160                elif column_name not in df.columns and is_nullable:
161                    df[column_name] = null_value
162            return df[[column_name for column_name, _, _ in schema]]
163
164    def __init__(
165        self,
166        dataset_path: str,
167        **kwargs,
168    ) -> None:
169        """
170        Dataset class to load and query datasets from the Pinecone Datasets catalog.
171        See `from_path` and `from_dataset_id` for examples on how to load a dataset.
172
173        Examples:
174            ```python
175            from pinecone_datasets import Dataset
176            dataset = Dataset.from_dataset_id("dataset_name")
177            # or
178            dataset = Dataset.from_path("gs://my-bucket/my-dataset")
179
180            for doc in dataset.iter_documents(batch_size=100):
181                index.upsert(doc)
182            for query in dataset.iter_queries(batch_size):
183                results = index.search(query)
184                # do something with the results
185            # or
186            dataset.documents # returns a pandas/polars DataFrame
187            dataset.queries # returns a pandas/polars DataFrame
188            ```
189
190        """
191        self._config = cfg
192        if dataset_path is not None:
193            endpoint = urlparse(dataset_path)._replace(path="").geturl()
194            self._fs = get_cloud_fs(endpoint, **kwargs)
195            self._dataset_path = dataset_path
196            if not self._fs.exists(self._dataset_path):
197                raise FileNotFoundError(
198                    "Dataset does not exist. Please check the path or dataset_id"
199                )
200        else:
201            self._fs = None
202            self._dataset_path = None
203        self._documents = None
204        self._queries = None
205        self._metadata = None
206        self._pinecone_client = None
207
208    def _is_datatype_exists(self, data_type: str) -> bool:
209        if not self._fs:
210            raise DatasetInitializationError()
211        return self._fs.exists(os.path.join(self._dataset_path, data_type))
212
213    @staticmethod
214    def _convert_metadata_from_dict_to_json(metadata: Optional[dict]) -> str:
215        if pd.isna(metadata):
216            return None
217        if metadata and not isinstance(metadata, dict):
218            raise TypeError(
219                f"metadata must be a dict but its {type(metadata)} meta = {metadata}"
220            )
221        return json.dumps(metadata, ensure_ascii=False)
222
223    @staticmethod
224    def _convert_metadata_from_json_to_dict(metadata: Optional[str]) -> dict:
225        if metadata is None:
226            return None
227        if not isinstance(metadata, str):
228            if isinstance(metadata, dict):
229                return metadata
230            else:
231                raise TypeError("metadata must be a string or dict")
232        return json.loads(metadata)
233
234    def _safe_read_from_path(self, data_type: str) -> pd.DataFrame:
235        if not self._fs:
236            raise DatasetInitializationError()
237
238        read_path_str = os.path.join(self._dataset_path, data_type, "*.parquet")
239        read_path = self._fs.glob(read_path_str)
240        if self._is_datatype_exists(data_type):
241            dataset = pq.ParquetDataset(read_path, filesystem=self._fs)
242            dataset_schema_names = dataset.schema.names
243            columns_to_null = []
244            columns_not_null = []
245            for column_name, is_nullable, null_value in getattr(
246                self._config.Schema.Names, data_type
247            ):
248                if column_name not in dataset_schema_names and not is_nullable:
249                    raise ValueError(
250                        f"error, file is not matching Pinecone Datasets Schmea: {column_name} not found"
251                    )
252                elif column_name not in dataset_schema_names and is_nullable:
253                    columns_to_null.append((column_name, null_value))
254                else:
255                    columns_not_null.append(column_name)
256            try:
257                # TODO: use of the columns_not_null and columns_to_null is only a workaround for proper schema validation and versioning
258                df = dataset.read_pandas(columns=columns_not_null).to_pandas()
259
260                # metadta supposed to be a dict [if legacy] or string
261                if data_type == "documents":
262                    df["metadata"] = df["metadata"].apply(
263                        self._convert_metadata_from_json_to_dict
264                    )
265                elif data_type == "queries":
266                    df["filter"] = df["filter"].apply(
267                        self._convert_metadata_from_json_to_dict
268                    )
269
270                for column_name, null_value in columns_to_null:
271                    df[column_name] = null_value
272                return df
273            # TODO: add more specific error handling, explain what is wrong
274            except Exception as e:
275                print("error, no exception: {}".format(e), file=sys.stderr)
276                raise (e)
277        else:
278            warnings.warn(
279                "WARNING: No data found at: {}. Returning empty DF".format(
280                    read_path_str
281                ),
282                UserWarning,
283                stacklevel=0,
284            )
285            return pd.DataFrame(
286                columns=[
287                    col[0] for col in getattr(self._config.Schema.Names, data_type)
288                ]
289            )
290
291    def _load_metadata(self) -> DatasetMetadata:
292        if not self._fs:
293            raise DatasetInitializationError()
294
295        with self._fs.open(
296            os.path.join(self._dataset_path, "metadata.json"), "rb"
297        ) as f:
298            metadata = json.load(f)
299        try:
300            out = DatasetMetadata(**metadata)
301            return out
302        # TODO: add more specific error handling, explain what is wrong
303        except ValidationError as e:
304            raise e
305
306    def __getitem__(self, key: str):
307        if key in ["documents", "queries"]:
308            return getattr(self, key)
309        else:
310            raise KeyError("Dataset does not have key: {}".format(key))
311
312    def __len__(self) -> int:
313        return self.documents.shape[0]
314
315    @property
316    def documents(self) -> pd.DataFrame:
317        if self._documents is None:
318            self._documents = self._safe_read_from_path("documents")
319        return self._documents
320
321    def iter_documents(
322        self, batch_size: int = 1, return_indexes=False
323    ) -> Iterator[List[Dict[str, Any]]]:
324        """
325        Iterates over the documents in the dataset.
326
327        Args:
328            batch_size (int, optional): The batch size to use for the iterator. Defaults to 1.
329
330        Returns:
331            Iterator[List[Dict[str, Any]]]: An iterator over the documents in the dataset.
332
333        Examples:
334            for batch in dataset.iter_documents(batch_size=100):
335                index.upsert(batch)
336        """
337        if isinstance(batch_size, int) and batch_size > 0:
338            return iter_pandas_dataframe_slices(
339                df=self.documents[self._config.Schema.documents_select_columns].dropna(
340                    axis=1, how="all"
341                ),
342                batch_size=batch_size,
343                return_indexes=return_indexes,
344            )
345        else:
346            raise ValueError("batch_size must be greater than 0")
347
348    @property
349    def queries(self) -> pd.DataFrame:
350        if self._queries is None:
351            self._queries = self._safe_read_from_path("queries")
352        return self._queries
353
354    def iter_queries(self) -> Iterator[Dict[str, Any]]:
355        """
356        Iterates over the queries in the dataset.
357
358        Returns:
359            Iterator[Dict[str, Any]]: An iterator over the queries in the dataset.
360
361        Examples:
362            for query in dataset.iter_queries():
363                results = index.query(**query)
364                # do something with the results
365        """
366        return iter_pandas_dataframe_single(
367            self.queries[self._config.Schema.queries_select_columns]
368        )
369
370    @property
371    def metadata(self) -> DatasetMetadata:
372        if self._metadata is None:
373            self._metadata = self._load_metadata()
374        return self._metadata
375
376    def head(self, n: int = 5) -> pd.DataFrame:
377        return self.documents.head(n)
378
379    def to_path(self, dataset_path: str, **kwargs):
380        """
381        Saves the dataset to a local or cloud storage path.
382        """
383        fs = get_cloud_fs(dataset_path, **kwargs)
384
385        # save documents
386        documents_path = os.path.join(dataset_path, "documents")
387        fs.makedirs(documents_path, exist_ok=True)
388
389        documents_metadta_copy = self.documents["metadata"].copy()
390        try:
391            self.documents["metadata"] = self.documents["metadata"].apply(
392                self._convert_metadata_from_dict_to_json
393            )
394            self.documents.to_parquet(
395                os.path.join(documents_path, "part-0.parquet"),
396                engine="pyarrow",
397                index=False,
398                filesystem=fs,
399            )
400        finally:
401            self.documents["metadata"] = documents_metadta_copy
402        # save queries
403        if not self.queries.empty:
404            queries_path = os.path.join(dataset_path, "queries")
405            fs.makedirs(queries_path, exist_ok=True)
406            queries_filter_copy = self.queries["filter"].copy()
407            try:
408                self.queries["filter"] = self.queries["filter"].apply(
409                    self._convert_metadata_from_dict_to_json
410                )
411                self.queries.to_parquet(
412                    os.path.join(queries_path, "part-0.parquet"),
413                    engine="pyarrow",
414                    index=False,
415                    filesystem=fs,
416                )
417            finally:
418                self.queries["filter"] = queries_filter_copy
419        else:
420            warnings.warn("Queries are empty, not saving queries")
421
422        # save metadata
423        with fs.open(os.path.join(dataset_path, "metadata.json"), "w") as f:
424            json.dump(self.metadata.dict(), f)
425
426    def to_catalog(
427        self,
428        dataset_id: str,
429        catalog_base_path: str = "",
430        **kwargs,
431    ):
432        """
433        Saves the dataset to the public catalog.
434        """
435
436        # TODO: duplicated code
437
438        catalog_base_path = (
439            catalog_base_path
440            if catalog_base_path
441            else os.environ.get("DATASETS_CATALOG_BASEPATH", cfg.Storage.endpoint)
442        )
443        dataset_path = os.path.join(catalog_base_path, f"{dataset_id}")
444        self.to_path(dataset_path, **kwargs)
445
446    def _upsert_to_index(
447        self, index_name: str, namespace: str, batch_size: int, show_progress: bool
448    ):
449        pinecone_index = Index(index_name=index_name)
450
451        res = pinecone_index.upsert_from_dataframe(
452            self.documents[self._config.Schema.documents_select_columns].dropna(
453                axis=1, how="all"
454            ),
455            namespace=namespace,
456            batch_size=batch_size,
457            show_progress=show_progress,
458        )
459        return {"upserted_count": res.upserted_count}
460
461    def _set_pinecone_index(
462        self,
463        api_key: Optional[str] = None,
464        environment: Optional[str] = None,
465        **kwargs,
466    ) -> None:
467        pc.init(api_key=api_key, environment=environment, **kwargs)
468        self._pinecone_client = pc
469
470    def _create_index(
471        self,
472        index_name: str,
473        api_key: Optional[str] = None,
474        environment: Optional[str] = None,
475        **kwargs,
476    ) -> Index:
477        self._set_pinecone_index(api_key=api_key, environment=environment)
478        pinecone_index_list = self._pinecone_client.list_indexes()
479
480        if index_name in pinecone_index_list:
481            raise ValueError(
482                f"index {index_name} already exists, Pinecone Datasets can only be upserted to a new indexe"
483            )
484        else:
485            # create index
486            print("creating index")
487            try:
488                self._pinecone_client.create_index(
489                    name=index_name,
490                    dimension=self.metadata.dense_model.dimension,
491                    **kwargs,
492                )
493                print("index created")
494                return True
495            except Exception as e:
496                print(f"error creating index: {e}")
497                return False
498
499    def to_pinecone_index(
500        self,
501        index_name: str,
502        namespace: Optional[str] = "",
503        should_create_index: bool = True,
504        batch_size: int = 100,
505        show_progress: bool = True,
506        api_key: Optional[str] = None,
507        environment: Optional[str] = None,
508        **kwargs,
509    ):
510        """
511        Saves the dataset to a Pinecone index.
512
513        this function will look for two environment variables:
514        - PINECONE_API_KEY
515        - PINECONE_ENVIRONMENT
516
517        Then, it will init a Pinecone Client and will perform an upsert to the index.
518        The upsert will be using async batches to increase performance.
519
520        Args:
521            index_name (str): the name of the index to upsert to
522            namespace (str, optional): the namespace to use for the upsert. Defaults to "".
523            batch_size (int, optional): the batch size to use for the upsert. Defaults to 100.
524            show_progress (bool, optional): whether to show a progress bar while upserting. Defaults to True.
525
526        Keyword Args:
527            kwargs (Dict): additional arguments to pass to the Pinecone Client constructor when creating the index.
528            see available parameters here: https://docs.pinecone.io/reference/create_index
529
530
531        Returns:
532            UpsertResponse: an object containing the upserted_count
533
534        Examples:
535            ```python
536            result = dataset.to_pinecone_index(index_name="my_index")
537            ```
538        """
539        if should_create_index:
540            if not self._create_index(
541                index_name, api_key=api_key, environment=environment, **kwargs
542            ):
543                raise RuntimeError("index creation failed")
544        else:
545            self._set_pinecone_index(api_key=api_key, environment=environment, **kwargs)
546
547        return self._upsert_to_index(
548            index_name=index_name,
549            namespace=namespace,
550            batch_size=batch_size,
551            show_progress=show_progress,
552        )