pixeltable 0.0.0__py3-none-any.whl

pixeltable/io/hf_datasets.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+
+import logging
+import math
+import random
+import typing
+from typing import Union, Optional, Any
+
+import pixeltable
+import pixeltable.type_system as ts
+from pixeltable import exceptions as excs
+
+if typing.TYPE_CHECKING:
+    import datasets
+
+_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."""
+    import datasets
+
+    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."""
+    import datasets
+
+    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(
+    table_path: str,
+    dataset: Union[datasets.Dataset, datasets.DatasetDict],
+    *,
+    column_name_for_split: Optional[str] = None,
+    schema_override: Optional[dict[str, Any]] = None,
+    **kwargs,
+) -> 'pixeltable.InsertableTable':
+    """Create a new `Table` from a Huggingface dataset, or dataset dict with multiple splits.
+        Requires datasets library to be installed.
+
+    Args:
+        path_str: Path to the table.
+        dataset: Huggingface datasets.Dataset or datasets.DatasetDict to insert into the table.
+        column_name_for_split: column name to use for split information. If None, no split information will be stored.
+        schema_override: Optional dictionary mapping column names to column type to override the corresponding defaults from
+        `pixeltable.utils.hf_datasets.huggingface_schema_to_pixeltable_schema`. The column type should be a pixeltable ColumnType.
+        For example, {'col_vid': VideoType()}, rather than {'col_vid': StringType()}.
+
+        kwargs: Additional arguments to pass to `create_table`.
+
+    Returns:
+        The newly created table. The table will have loaded the data from the dataset.
+    """
+    import datasets
+    import pixeltable as pxt
+
+    if table_path in pxt.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 = pxt.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
+
+    pxt.move(tmp_name, table_path)
+    return pxt.get_table(table_path)

pixeltable/io/pandas.py

@@ -0,0 +1,148 @@
+from typing import Optional, Any, Iterable
+
+import numpy as np
+import pandas as pd
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+
+
+def import_pandas(
+    tbl_name: str, df: pd.DataFrame, *, schema_overrides: Optional[dict[str, pxt.ColumnType]] = None
+) -> pxt.catalog.InsertableTable:
+    """Creates a new `Table` from a Pandas `DataFrame`, with the specified name. The schema of the table
+    will be inferred from the `DataFrame`, unless `schema` is specified.
+
+    The column names of the new `Table` will be identical to those in the `DataFrame`, as long as they are valid
+    Pixeltable identifiers. If a column name is not a valid Pixeltable identifier, it will be normalized according to
+    the following procedure:
+    - first replace any non-alphanumeric characters with underscores;
+    - then, preface the result with the letter 'c' if it begins with a number or an underscore;
+    - then, if there are any duplicate column names, suffix the duplicates with '_2', '_3', etc., in column order.
+
+    Args:
+        tbl_name: The name of the table to create.
+        df: The Pandas `DataFrame`.
+        schema_overrides: If specified, then for each (name, type) pair in `schema_overrides`, the column with
+            name `name` will be given type `type`, instead of being inferred from the `DataFrame`. The keys in
+            `schema_overrides` should be the column names of the `DataFrame` (whether or not they are valid
+            Pixeltable identifiers).
+    """
+    schema = _df_to_pxt_schema(df, schema_overrides)
+    tbl_rows = (dict(_df_row_to_pxt_row(row, schema)) for row in df.itertuples())
+    table = pxt.create_table(tbl_name, schema)
+    table.insert(tbl_rows)
+    return table
+
+
+def import_csv(
+    table_path: str, filepath_or_buffer, schema_overrides: Optional[dict[str, ts.ColumnType]] = None, **kwargs
+) -> pxt.catalog.InsertableTable:
+    """
+    Creates a new `Table` from a csv file. This is a convenience method and is equivalent
+    to calling `import_pandas(table_path, pd.read_csv(filepath_or_buffer, **kwargs), schema=schema)`.
+    See the Pandas documentation for `read_csv` for more details.
+    """
+    df = pd.read_csv(filepath_or_buffer, **kwargs)
+    return import_pandas(table_path, df, schema_overrides=schema_overrides)
+
+
+def import_excel(
+    table_path: str, io, *args, schema_overrides: Optional[dict[str, ts.ColumnType]] = None, **kwargs
+) -> pxt.catalog.InsertableTable:
+    """
+    Creates a new `Table` from an excel (.xlsx) file. This is a convenience method and is equivalent
+    to calling `import_pandas(table_path, pd.read_excel(io, *args, **kwargs), schema=schema)`.
+    See the Pandas documentation for `read_excel` for more details.
+    """
+    df = pd.read_excel(io, *args, **kwargs)
+    return import_pandas(table_path, df, schema_overrides=schema_overrides)
+
+
+def _df_to_pxt_schema(
+    df: pd.DataFrame, schema_overrides: Optional[dict[str, pxt.ColumnType]]
+) -> dict[str, pxt.ColumnType]:
+    if schema_overrides is not None:
+        for pd_name in schema_overrides:
+            if pd_name not in df.columns:
+                raise excs.Error(
+                    f'Column `{pd_name}` specified in `schema_overrides` does not exist in the given `DataFrame`.'
+                )
+    schema = {}
+    for pd_name, pd_dtype in zip(df.columns, df.dtypes):
+        if schema_overrides is not None and pd_name in schema_overrides:
+            pxt_type = schema_overrides[pd_name]
+        else:
+            pxt_type = _np_dtype_to_pxt_type(pd_dtype, df[pd_name])
+        pxt_name = _normalize_pxt_col_name(pd_name)
+        # Ensure that column names are unique by appending a distinguishing suffix
+        # to any collisions
+        if pxt_name in schema:
+            n = 2
+            while f'{pxt_name}_{n}' in schema:
+                n += 1
+            pxt_name = f'{pxt_name}_{n}'
+        schema[pxt_name] = pxt_type
+    return schema
+
+
+def _normalize_pxt_col_name(pd_name: str) -> str:
+    """
+    Normalizes an arbitrary DataFrame column name into a valid Pixeltable identifier by:
+    - replacing any non-ascii or non-alphanumeric characters with an underscore _
+    - prefixing the result with the letter 'c' if it starts with an underscore or a number
+    """
+    id = ''.join(ch if ch.isascii() and ch.isalnum() else '_' for ch in pd_name)
+    if id[0].isnumeric():
+        id = f'c_{id}'
+    elif id[0] == '_':
+        id = f'c{id}'
+    assert pxt.catalog.is_valid_identifier(id), id
+    return id
+
+
+def _np_dtype_to_pxt_type(np_dtype: np.dtype, data_col: pd.Series) -> pxt.ColumnType:
+    """
+    Infers a Pixeltable type based on a Numpy dtype.
+    """
+    if np.issubdtype(np_dtype, np.integer):
+        return pxt.IntType()
+    if np.issubdtype(np_dtype, np.floating):
+        return pxt.FloatType()
+    if np.issubdtype(np_dtype, np.bool_):
+        return pxt.BoolType()
+    if np_dtype == np.object_ or np.issubdtype(np_dtype, np.character):
+        has_nan = any(isinstance(val, float) and np.isnan(val) for val in data_col)
+        return pxt.StringType(nullable=has_nan)
+    if np.issubdtype(np_dtype, np.datetime64):
+        has_nat = any(pd.isnull(val) for val in data_col)
+        return pxt.TimestampType(nullable=has_nat)
+    raise excs.Error(f'Unsupported dtype: {np_dtype}')
+
+
+def _df_row_to_pxt_row(row: tuple[Any, ...], schema: dict[str, pxt.ColumnType]) -> dict[str, Any]:
+    rows = {}
+    for val, (col_name, pxt_type) in zip(row[1:], schema.items()):
+        if pxt_type.is_float_type():
+            val = float(val)
+        elif isinstance(val, float) and np.isnan(val):
+            # pandas uses NaN for empty cells, even for types other than float;
+            # for any type but a float, convert these to None
+            val = None
+        elif pxt_type.is_int_type():
+            val = int(val)
+        elif pxt_type.is_bool_type():
+            val = bool(val)
+        elif pxt_type.is_string_type():
+            val = str(val)
+        elif pxt_type.is_timestamp_type():
+            if pd.isnull(val):
+                # pandas has the bespoke 'NaT' type for a missing timestamp; postgres is very
+                # much not-ok with it. (But if we convert it to None and then load out the
+                # table contents as a pandas DataFrame, it will correctly restore the 'NaT'!)
+                val = None
+            else:
+                val = pd.Timestamp(val).to_pydatetime()
+        rows[col_name] = val
+    return rows

pixeltable/io/parquet.py

@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import io
+import json
+import logging
+import random
+import typing
+from collections import deque
+from pathlib import Path
+from typing import Dict, Optional
+
+import PIL.Image
+import numpy as np
+
+import pixeltable.exceptions as exc
+import pixeltable.type_system as ts
+from pixeltable.utils.transactional_directory import transactional_directory
+
+if typing.TYPE_CHECKING:
+    import pixeltable as pxt
+    import pyarrow as pa
+
+_logger = logging.getLogger(__name__)
+
+
+def _write_batch(value_batch: Dict[str, deque], schema: pa.Schema, output_path: Path) -> None:
+    import pyarrow as pa
+
+    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: pxt.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.
+    """
+    from pixeltable.utils.arrow import to_arrow_schema
+
+    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."""
+    import pyarrow as pa
+    from pixeltable.utils.arrow import to_pixeltable_schema
+
+    input_path = Path(parquet_path).expanduser()
+    parquet_dataset = pa.parquet.ParquetDataset(input_path)
+    return to_pixeltable_schema(parquet_dataset.schema)
+
+
+def import_parquet(
+    table_path: str,
+    *,
+    parquet_path: str,
+    schema_override: Optional[Dict[str, ts.ColumnType]] = None,
+    **kwargs,
+) -> pxt.catalog.InsertableTable:
+    """Create a new `Table` from a Parquet file or set of files. Requires pyarrow to be installed.
+    Args:
+        path_str: Path to the table within pixeltable.
+        parquet_path: Path to an individual Parquet file or directory of Parquet files.
+        schema_override: Optional dictionary mapping column names to column type to override the default
+                        schema inferred from the Parquet file. The column type should be a pixeltable ColumnType.
+                        For example, {'col_vid': VideoType()}, rather than {'col_vid': StringType()}.
+                        Any fields not provided explicitly will map to types with `pixeltable.utils.parquet.parquet_schema_to_pixeltable_schema`
+        kwargs: Additional arguments to pass to `create_table`.
+
+    Returns:
+        The newly created table. The table will have loaded the data from the Parquet file(s).
+    """
+    import pixeltable as pxt
+    import pyarrow as pa
+    from pixeltable.utils.arrow import iter_tuples
+
+    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 pxt.list_tables():
+        raise exc.Error(f'Table {table_path} already exists')
+
+    try:
+        tmp_name = f'{table_path}_tmp_{random.randint(0, 100000000)}'
+        tab = pxt.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
+
+    pxt.move(tmp_name, table_path)
+    return pxt.get_table(table_path)

pixeltable/iterators/__init__.py

@@ -0,0 +1,3 @@
+from .base import ComponentIterator
+from .document import DocumentSplitter
+from .video import FrameIterator

pixeltable/iterators/base.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+from typing import Dict, Any, Tuple, List
+from abc import abstractmethod, ABC
+
+from pixeltable.type_system import ColumnType
+
+
+class ComponentIterator(ABC):
+    """Base class for iterators."""
+
+    @classmethod
+    @abstractmethod
+    def input_schema(cls) -> Dict[str, ColumnType]:
+        """Provide the Pixeltable types of the init() parameters
+
+        The keys need to match the names of the init() parameters. This is equivalent to the parameters_types
+        parameter of the @function decorator.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def output_schema(cls, *args: Any, **kwargs: Any) -> Tuple[Dict[str, ColumnType], List[str]]:
+        """Specify the dictionary returned by next() and a list of unstored column names
+
+        Returns:
+            a dictionary which is turned into a list of columns in the output table
+            a list of unstored column names
+        """
+        raise NotImplementedError
+
+    def __iter__(self) -> ComponentIterator:
+        return self
+
+    @abstractmethod
+    def __next__(self) -> Dict[str, Any]:
+        """Return the next element of the iterator as a dictionary or raise StopIteration"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the iterator and release all resources"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def set_pos(self, pos: int) -> None:
+        """Set the iterator position to pos"""
+        raise NotImplementedError
+
+    @classmethod
+    def create(cls, **kwargs: Any) -> tuple[type[ComponentIterator], dict[str, Any]]:
+        return cls, kwargs