pixeltable 0.1.0__py3-none-any.whl → 0.2.4__py3-none-any.whl

pixeltable/utils/hf_datasets.py

@@ -0,0 +1,157 @@
+import datasets
+from typing import Union, Optional, List, Dict, Any
+import pixeltable.type_system as ts
+from pixeltable import exceptions as excs
+import math
+import logging
+import pixeltable
+import random
+
+_logger = logging.getLogger(__name__)
+
+# use 100MB as the batch size limit for loading a huggingface dataset into pixeltable.
+# The primary goal is to bound memory use, regardless of dataset size.
+# Second goal is to limit overhead. 100MB is presumed to be reasonable for a lot of storage systems.
+_K_BATCH_SIZE_BYTES = 100_000_000
+
+# note, there are many more types. we allow overrides in the schema_override parameter
+# to handle cases where the appropriate type is not yet mapped, or to override this mapping.
+# https://huggingface.co/docs/datasets/v2.17.0/en/package_reference/main_classes#datasets.Value
+_hf_to_pxt: Dict[str, ts.ColumnType] = {
+    'int32': ts.IntType(nullable=True),  # pixeltable widens to big int
+    'int64': ts.IntType(nullable=True),
+    'bool': ts.BoolType(nullable=True),
+    'float32': ts.FloatType(nullable=True),
+    'string': ts.StringType(nullable=True),
+    'timestamp[s]': ts.TimestampType(nullable=True),
+    'timestamp[ms]': ts.TimestampType(nullable=True),  # HF dataset iterator converts timestamps to datetime.datetime
+}
+
+def _to_pixeltable_type(
+    feature_type: Union[datasets.ClassLabel, datasets.Value, datasets.Sequence],
+) -> Optional[ts.ColumnType]:
+    """Convert a huggingface feature type to a pixeltable ColumnType if one is defined."""
+    if isinstance(feature_type, datasets.ClassLabel):
+        # enum, example: ClassLabel(names=['neg', 'pos'], id=None)
+        return ts.StringType(nullable=True)
+    elif isinstance(feature_type, datasets.Value):
+        # example: Value(dtype='int64', id=None)
+        return _hf_to_pxt.get(feature_type.dtype, None)
+    elif isinstance(feature_type, datasets.Sequence):
+        # example: cohere wiki. Sequence(feature=Value(dtype='float32', id=None), length=-1, id=None)
+        dtype = _to_pixeltable_type(feature_type.feature)
+        length = feature_type.length if feature_type.length != -1 else None
+        return ts.ArrayType(shape=(length,), dtype=dtype)
+    else:
+        return None
+
+def _get_hf_schema(dataset: Union[datasets.Dataset, datasets.DatasetDict]) -> datasets.Features:
+    """Get the schema of a huggingface dataset as a dictionary."""
+    first_dataset = dataset if isinstance(dataset, datasets.Dataset) else next(iter(dataset.values()))
+    return first_dataset.features
+
+def huggingface_schema_to_pixeltable_schema(
+    hf_dataset: Union[datasets.Dataset, datasets.DatasetDict],
+) -> Dict[str, Optional[ts.ColumnType]]:
+    """Generate a pixeltable schema from a huggingface dataset schema.
+    Columns without a known mapping are mapped to None
+    """
+    hf_schema = _get_hf_schema(hf_dataset)
+    pixeltable_schema = {
+        column_name: _to_pixeltable_type(feature_type) for column_name, feature_type in hf_schema.items()
+    }
+    return pixeltable_schema
+
+def import_huggingface_dataset(
+    cl: 'pixeltable.Client',
+    table_path: str,
+    dataset: Union[datasets.Dataset, datasets.DatasetDict],
+    *,
+    column_name_for_split: Optional[str],
+    schema_override: Optional[Dict[str, Any]],
+    **kwargs,
+) -> 'pixeltable.InsertableTable':
+    """See `pixeltable.Client.import_huggingface_dataset` for documentation"""
+    if table_path in cl.list_tables():
+        raise excs.Error(f'table {table_path} already exists')
+
+    if not isinstance(dataset, (datasets.Dataset, datasets.DatasetDict)):
+        raise excs.Error(f'`type(dataset)` must be `datasets.Dataset` or `datasets.DatasetDict`. Got {type(dataset)=}')
+
+    if isinstance(dataset, datasets.Dataset):
+        # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
+        raw_name = dataset.split._name
+        split_name = raw_name.split('[')[0] if raw_name is not None else None
+        dataset_dict = {split_name: dataset}
+    else:
+        dataset_dict = dataset
+
+    pixeltable_schema = huggingface_schema_to_pixeltable_schema(dataset)
+    if schema_override is not None:
+        pixeltable_schema.update(schema_override)
+
+    if column_name_for_split is not None:
+        if column_name_for_split in pixeltable_schema:
+            raise excs.Error(
+                f'Column name `{column_name_for_split}` already exists in dataset schema; provide a different `column_name_for_split`'
+            )
+        pixeltable_schema[column_name_for_split] = ts.StringType(nullable=True)
+
+    for field, column_type in pixeltable_schema.items():
+        if column_type is None:
+            raise excs.Error(f'Could not infer pixeltable type for feature `{field}` in huggingface dataset')
+
+    if isinstance(dataset, datasets.Dataset):
+        # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
+        raw_name = dataset.split._name
+        split_name = raw_name.split('[')[0] if raw_name is not None else None
+        dataset_dict = {split_name: dataset}
+    elif isinstance(dataset, datasets.DatasetDict):
+        dataset_dict = dataset
+    else:
+        raise excs.Error(f'`type(dataset)` must be `datasets.Dataset` or `datasets.DatasetDict`. Got {type(dataset)=}')
+
+    # extract all class labels from the dataset to translate category ints to strings
+    hf_schema = _get_hf_schema(dataset)
+    categorical_features = {
+        feature_name: feature_type.names
+        for (feature_name, feature_type) in hf_schema.items()
+        if isinstance(feature_type, datasets.ClassLabel)
+    }
+
+    try:
+        # random tmp name
+        tmp_name = f'{table_path}_tmp_{random.randint(0, 100000000)}'
+        tab = cl.create_table(tmp_name, pixeltable_schema, **kwargs)
+
+        def _translate_row(row: Dict[str, Any], split_name: str) -> Dict[str, Any]:
+            output_row = row.copy()
+            # map all class labels to strings
+            for field, values in categorical_features.items():
+                output_row[field] = values[row[field]]
+            # add split name to row
+            if column_name_for_split is not None:
+                output_row[column_name_for_split] = split_name
+            return output_row
+
+        for split_name, split_dataset in dataset_dict.items():
+            num_batches = split_dataset.size_in_bytes / _K_BATCH_SIZE_BYTES
+            tuples_per_batch = math.ceil(split_dataset.num_rows / num_batches)
+            assert tuples_per_batch > 0
+
+            batch = []
+            for row in split_dataset:
+                batch.append(_translate_row(row, split_name))
+                if len(batch) >= tuples_per_batch:
+                    tab.insert(batch)
+                    batch = []
+            # last batch
+            if len(batch) > 0:
+                tab.insert(batch)
+
+    except Exception as e:
+        _logger.error(f'Error while inserting dataset into table: {tmp_name}')
+        raise e
+
+    cl.move(tmp_name, table_path)
+    return cl.get_table(table_path)

pixeltable/utils/media_store.py

@@ -0,0 +1,76 @@
+import glob
+import os
+import re
+import shutil
+import uuid
+from typing import Optional, List, Tuple, Dict
+from pathlib import Path
+from collections import defaultdict
+from uuid import UUID
+
+from pixeltable.env import Env
+
+
+class MediaStore:
+    """
+    Utilities to manage media files stored in Env.media_dir
+
+    Media file names are a composite of: table id, column id, version, uuid:
+    the table id/column id/version are redundant but useful for identifying all files for a table
+    or all files created for a particular version of a table
+    """
+    pattern = re.compile(r'([0-9a-fA-F]+)_(\d+)_(\d+)_([0-9a-fA-F]+)')  # tbl_id, col_id, version, uuid
+
+    @classmethod
+    def prepare_media_path(cls, tbl_id: UUID, col_id: int, version: int, ext: Optional[str] = None) -> Path:
+        """
+        Construct a new, unique Path name for a persisted media file, and create the parent directory
+        for the new Path if it does not already exist. The Path will reside in
+        the environment's media_dir.
+        """
+        id_hex = uuid.uuid4().hex
+        parent = Env.get().media_dir / tbl_id.hex / id_hex[0:2] / id_hex[0:4]
+        parent.mkdir(parents=True, exist_ok=True)
+        return parent / f'{tbl_id.hex}_{col_id}_{version}_{id_hex}{ext or ""}'
+
+    @classmethod
+    def delete(cls, tbl_id: UUID, version: Optional[int] = None) -> None:
+        """Delete all files belonging to tbl_id. If version is not None, delete
+        only those files belonging to the specified version."""
+        assert tbl_id is not None
+        if version is None:
+            # Remove the entire folder for this table id.
+            path = Env.get().media_dir / tbl_id.hex
+            if path.exists():
+                shutil.rmtree(path)
+        else:
+            # Remove only the elements for the specified version.
+            paths = glob.glob(str(Env.get().media_dir / tbl_id.hex) + f'/**/{tbl_id.hex}_*_{version}_*', recursive=True)
+            for path in paths:
+                os.remove(path)
+
+    @classmethod
+    def count(cls, tbl_id: UUID) -> int:
+        """
+        Return number of files for given tbl_id.
+        """
+        paths = glob.glob(str(Env.get().media_dir / tbl_id.hex) + f'/**/{tbl_id.hex}_*', recursive=True)
+        return len(paths)
+
+    @classmethod
+    def stats(cls) -> List[Tuple[int, int, int, int]]:
+        paths = glob.glob(str(Env.get().media_dir) + "/**", recursive=True)
+        # key: (tbl_id, col_id), value: (num_files, size)
+        d: Dict[Tuple[UUID, int], List[int]] = defaultdict(lambda: [0, 0])
+        for p in paths:
+            if not os.path.isdir(p):
+                matched = re.match(cls.pattern, Path(p).name)
+                assert matched is not None
+                tbl_id, col_id = UUID(hex=matched[1]), int(matched[2])
+                file_info = os.stat(p)
+                t = d[(tbl_id, col_id)]
+                t[0] += 1
+                t[1] += file_info.st_size
+        result = [(tbl_id, col_id, num_files, size) for (tbl_id, col_id), (num_files, size) in d.items()]
+        result.sort(key=lambda e: e[3], reverse=True)
+        return result

pixeltable/utils/parquet.py

@@ -0,0 +1,167 @@
+import io
+import json
+import logging
+from collections import deque
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import pyarrow as pa
+import pyarrow.parquet
+
+import pixeltable.type_system as ts
+from pixeltable.utils.arrow import iter_tuples, to_arrow_schema, to_pixeltable_schema
+from pixeltable.utils.transactional_directory import transactional_directory
+import pixeltable.exceptions as exc
+
+import random
+
+_logger = logging.getLogger(__name__)
+
+
+def _write_batch(value_batch : Dict[str, deque], schema : pa.Schema, output_path : Path) -> None:
+    pydict = {}
+    for field in schema:
+        if isinstance(field.type, pa.FixedShapeTensorType):
+            stacked_arr = np.stack(value_batch[field.name])
+            pydict[field.name] = pa.FixedShapeTensorArray.from_numpy_ndarray(stacked_arr)
+        else:
+            pydict[field.name] = value_batch[field.name]
+
+    tab = pa.Table.from_pydict(pydict, schema=schema)
+    pa.parquet.write_table(tab, output_path)
+
+def save_parquet(df: 'pixeltable.DataFrame', dest_path: Path, partition_size_bytes: int = 100_000_000) -> None:
+    """
+        Internal method to stream dataframe data to parquet format.
+        Does not materialize the dataset to memory.
+
+        It preserves pixeltable type metadata in a json file, which would otherwise
+        not be available in the parquet format.
+
+        Images are stored inline in a compressed format in their parquet file.
+
+        Args:
+            df : dataframe to save.
+            dest_path : path to directory to save the parquet files to.
+            partition_size_bytes : maximum target size for each chunk. Default 100_000_000 bytes.
+    """
+    column_names = df.get_column_names()
+    column_types = df.get_column_types()
+    type_dict = {k: v.as_dict() for k, v in zip(column_names, column_types)}
+    arrow_schema = to_arrow_schema(dict(zip(column_names, column_types)))
+
+    # store the changes atomically
+    with transactional_directory(dest_path) as temp_path:
+        # dump metadata json file so we can inspect what was the source of the parquet file later on.
+        json.dump(df._as_dict(), (temp_path / '.pixeltable.json').open('w')) # pylint: disable=protected-access
+        json.dump(type_dict, (temp_path / '.pixeltable.column_types.json').open('w')) # keep type metadata
+
+        batch_num = 0
+        current_value_batch : Dict[str, deque] = {k:deque() for k in column_names}
+        current_byte_estimate = 0
+
+        for data_row in df._exec(): # pylint: disable=protected-access
+            for (col_name, col_type, e) in zip(column_names, column_types, df._select_list_exprs): # pylint: disable=protected-access
+                val = data_row[e.slot_idx]
+                if val is None:
+                    current_value_batch[col_name].append(val)
+                    continue
+
+                assert val is not None
+                if col_type.is_image_type():
+                    # images get inlined into the parquet file
+                    if data_row.file_paths is not None and data_row.file_paths[e.slot_idx] is not None:
+                        # if there is a file, read directly to preserve information
+                        with open(data_row.file_paths[e.slot_idx], 'rb') as f:
+                            val = f.read()
+                    elif isinstance(val, PIL.Image.Image):
+                        # if no file available, eg. bc it is computed, convert to png
+                        buf = io.BytesIO()
+                        val.save(buf, format='PNG')
+                        val = buf.getvalue()
+                    else:
+                        assert False, f'unknown image type {type(val)}'
+                    length = len(val)
+                elif col_type.is_string_type():
+                    length = len(val)
+                elif col_type.is_video_type():
+                    if data_row.file_paths is not None and data_row.file_paths[e.slot_idx] is not None:
+                        val = data_row.file_paths[e.slot_idx]
+                    else:
+                        assert False, f'unknown video type {type(val)}'
+                    length = len(val)
+                elif col_type.is_json_type():
+                    val = json.dumps(val)
+                    length = len(val)
+                elif col_type.is_array_type():
+                    length = val.nbytes
+                elif col_type.is_int_type():
+                    length = 8
+                elif col_type.is_float_type():
+                    length = 8
+                elif col_type.is_bool_type():
+                    length = 1
+                elif col_type.is_timestamp_type():
+                    length = 8
+                else:
+                    assert False, f'unknown type {col_type} for {col_name}'
+
+                current_value_batch[col_name].append(val)
+                current_byte_estimate += length
+            if current_byte_estimate > partition_size_bytes:
+                assert batch_num < 100_000, 'wrote too many parquet files, unclear ordering'
+                _write_batch(current_value_batch, arrow_schema,  temp_path / f'part-{batch_num:05d}.parquet')
+                batch_num += 1
+                current_value_batch = {k:deque() for k in column_names}
+                current_byte_estimate = 0
+
+        _write_batch(current_value_batch, arrow_schema, temp_path / f'part-{batch_num:05d}.parquet')
+
+
+def parquet_schema_to_pixeltable_schema(parquet_path: str) -> Dict[str, Optional[ts.ColumnType]]:
+    """Generate a default pixeltable schema for the given parquet file. Returns None for unknown types."""
+
+    input_path = Path(parquet_path).expanduser()
+    parquet_dataset = pa.parquet.ParquetDataset(input_path)
+    return to_pixeltable_schema(parquet_dataset.schema)
+
+
+def import_parquet(
+    cl: 'pixeltable.Client',
+    table_path: str,
+    *,
+    parquet_path: str,
+    schema_override: Optional[Dict[str, ts.ColumnType]],
+    **kwargs,
+) -> 'catalog.InsertableTable':
+    """See `pixeltable.Client.import_parquet` for documentation"""
+    input_path = Path(parquet_path).expanduser()
+    parquet_dataset = pa.parquet.ParquetDataset(input_path)
+
+    schema = parquet_schema_to_pixeltable_schema(parquet_path)
+    if schema_override is None:
+        schema_override = {}
+
+    schema.update(schema_override)
+    for k, v in schema.items():
+        if v is None:
+            raise exc.Error(f'Could not infer pixeltable type for column {k} from parquet file')
+
+    if table_path in cl.list_tables():
+        raise exc.Error(f'Table {table_path} already exists')
+
+    try:
+        tmp_name = f'{table_path}_tmp_{random.randint(0, 100000000)}'
+        tab = cl.create_table(tmp_name, schema, **kwargs)
+        for fragment in parquet_dataset.fragments:
+            for batch in fragment.to_batches():
+                dict_batch = list(iter_tuples(batch))
+                tab.insert(dict_batch)
+    except Exception as e:
+        _logger.error(f'Error while inserting Parquet file into table: {e}')
+        raise e
+
+    cl.move(tmp_name, table_path)
+    return cl.get_table(table_path)

pixeltable/utils/pytorch.py

@@ -0,0 +1,91 @@
+import io
+import pyarrow as pa
+import pyarrow.parquet
+import torch
+import torch.utils.data
+from pathlib import Path
+import PIL.Image
+import json
+from typing import Dict, Iterator, Any
+import datetime
+
+from pixeltable.type_system import ColumnType
+import numpy as np
+
+class PixeltablePytorchDataset(torch.utils.data.IterableDataset):
+    """
+    PyTorch dataset interface for pixeltable data.
+    NB. This class must inherit from torch.utils.data.IterableDataset for it
+    to work with torch.utils.data.DataLoader.
+    """
+    def __init__(
+        self,
+        path: Path,
+        image_format: str,
+    ):
+        """
+        Args:
+            path: path to directory containing parquet files
+            image_format: 'np' or 'pt'. 'np' is RGB uint8 array,
+                        'pt' is result of torchvision.transforms.ToTensor()
+        """
+        super().__init__()
+
+        self.path = path
+        self.image_format = image_format
+        assert image_format in ["np", "pt"]
+        column_type_path = path / '.pixeltable.column_types.json'
+        assert column_type_path.exists(), f"missing {column_type_path}"
+        with column_type_path.open() as f:
+            column_types = json.load(f)
+        self.column_types = {k: ColumnType.from_dict(v) for k, v in column_types.items()}
+        self.part_metadata = pa.parquet.ParquetDataset(path).files
+
+    def _unmarshall(self, k: str, v: Any) -> Any:
+        if self.column_types[k].is_image_type():
+            assert isinstance(v, bytes)
+            im = PIL.Image.open(io.BytesIO(v))
+            arr = np.array(im) # will copy data to guarantee "WRITEABLE" flag assertion below.
+            assert arr.flags["WRITEABLE"]
+
+            if self.image_format == "np":
+                return arr
+
+            assert self.image_format == "pt"
+            import torchvision  # pylint: disable = import-outside-toplevel
+
+            # use arr instead of im in ToTensor() to guarantee array input
+            # to torch.from_numpy is writable. Using im is a suspected cause of
+            # https://github.com/pixeltable/pixeltable/issues/69
+            return torchvision.transforms.ToTensor()(arr)
+        elif self.column_types[k].is_json_type():
+            assert isinstance(v, str)
+            return json.loads(v)
+        elif self.column_types[k].is_array_type():
+            assert isinstance(v, np.ndarray)
+            if not v.flags["WRITEABLE"]:
+                v = v.copy()
+            assert v.flags["WRITEABLE"]
+            return v
+        elif self.column_types[k].is_timestamp_type():
+            # pytorch default collation only supports numeric types
+            assert isinstance(v, datetime.datetime)
+            return v.timestamp()
+        else:
+            assert not isinstance(v, np.ndarray) # all array outputs should be handled above
+            return v
+
+    def __iter__(self) -> Iterator[Dict[str, Any]]:
+        import pixeltable.utils.arrow as arrow
+        worker_info = torch.utils.data.get_worker_info()
+
+        if worker_info is None:
+            part_list = range(len(self.part_metadata))
+        else:
+            part_list = [ i for i in part_list if (i % worker_info.num_workers) == worker_info.id ]
+
+        for part_no in part_list:
+            pqf = pa.parquet.ParquetFile(self.part_metadata[part_no])
+            for batch in pqf.iter_batches():
+                for tup in arrow.iter_tuples(batch):
+                    yield {k: self._unmarshall(k, v) for k, v in tup.items()}

pixeltable/utils/s3.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+
+def get_client() -> Any:
+    import boto3
+    import botocore
+    try:
+        boto3.Session().get_credentials().get_frozen_credentials()
+        return boto3.client('s3')  # credentials are available
+    except AttributeError:
+        # No credentials available, use unsigned mode
+        config = botocore.config.Config(signature_version=botocore.UNSIGNED)
+        return boto3.client('s3', config=config)

pixeltable/utils/sql.py

@@ -0,0 +1,17 @@
+import logging
+
+import sqlalchemy as sql
+
+
+def log_stmt(logger: logging.Logger, stmt) -> None:
+    logger.debug(f'executing {str(stmt.compile(dialect=sql.dialects.postgresql.dialect()))}')
+
+def log_explain(logger: logging.Logger, stmt: sql.sql.ClauseElement, conn: sql.engine.Connection) -> None:
+    try:
+        # don't set dialect=Env.get().engine.dialect: x % y turns into x %% y, which results in a syntax error
+        stmt_str = str(stmt.compile(compile_kwargs={'literal_binds': True}))
+        explain_result = conn.execute(sql.text(f'EXPLAIN {stmt_str}'))
+        explain_str = '\n'.join([str(row) for row in explain_result])
+        logger.debug(f'SqlScanNode explain:\n{explain_str}')
+    except Exception as e:
+        logger.warning(f'EXPLAIN failed')

pixeltable/utils/transactional_directory.py

@@ -0,0 +1,35 @@
+import shutil
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Generator
+
+import pixeltable.exceptions as excs
+
+
+@contextmanager
+def transactional_directory(folder_path: Path) -> Generator[Path, Any, Any]:
+    """
+    Args:
+        folder_path: path to the folder we want to create
+
+    Yields:
+        A pathlib.Path to a hidden temporary folder, which can be used to accumulate changes.
+        If everything succeeds, the changes are committed via an atomic move operation upon exiting the 'with' block (os.replace)
+        If an exception occurred, no changes are visible in the original folder.
+
+    Example:
+        folder_path = pathlib.Path("path/to/folder")
+        with transactional_folder(folder_path) as temp_folder:
+            (temp_folder / "subfolder1").mkdir()
+            (temp_folder / "subfolder2").mkdir()
+    """
+    if folder_path.exists():
+        raise excs.Error(f"Folder {folder_path} already exists")
+
+    tmp_folder = folder_path.parent / f".tmp_{folder_path.name}"
+    # Remove the temporary folder if it already exists, eg if the previous run crashed
+    shutil.rmtree(str(tmp_folder), ignore_errors=True)
+    tmp_folder.mkdir(parents=True)
+    yield tmp_folder
+    # If everything succeeds, `commit' the changes by moving the temporary folder
+    tmp_folder.rename(folder_path)

pixeltable-0.2.4.dist-info/LICENSE

@@ -0,0 +1,18 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+
+   Copyright 2023 Marcel Kornacker
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.