pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/io/table_data_conduit.py

@@ -8,9 +8,11 @@ import urllib.parse
 import urllib.request
 from dataclasses import dataclass, field, fields
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, cast
 
+import numpy as np
 import pandas as pd
+import PIL
 from pyarrow.parquet import ParquetDataset
 
 import pixeltable as pxt
@@ -47,16 +49,16 @@ class TableDataConduitFormat(str, enum.Enum):
 
 @dataclass
 class TableDataConduit:
-    source: TableDataSource
-    source_format: Optional[str] = None
-    source_column_map: Optional[dict[str, str]] = None
+    source: 'TableDataSource'
+    source_format: str | None = None
+    source_column_map: dict[str, str] | None = None
     if_row_exists: Literal['update', 'ignore', 'error'] = 'error'
-    pxt_schema: Optional[dict[str, Any]] = None
-    src_schema_overrides: Optional[dict[str, Any]] = None
-    src_schema: Optional[dict[str, Any]] = None
-    pxt_pk: Optional[list[str]] = None
-    src_pk: Optional[list[str]] = None
-    valid_rows: Optional[RowData] = None
+    pxt_schema: dict[str, ts.ColumnType] | None = None
+    src_schema_overrides: dict[str, ts.ColumnType] | None = None
+    src_schema: dict[str, ts.ColumnType] | None = None
+    pxt_pk: list[str] | None = None
+    src_pk: list[str] | None = None
+    valid_rows: RowData | None = None
     extra_fields: dict[str, Any] = field(default_factory=dict)
 
     reqd_col_names: set[str] = field(default_factory=set)
@@ -87,7 +89,7 @@ class TableDataConduit:
         for name, coltype in self.pxt_schema.items():
             self.pxt_schema[name] = ts.ColumnType.normalize_type(coltype)
 
-    def infer_schema(self) -> dict[str, Any]:
+    def infer_schema(self) -> dict[str, ts.ColumnType]:
         raise NotImplementedError
 
     def valid_row_batch(self) -> Iterator[RowData]:
@@ -101,7 +103,7 @@ class TableDataConduit:
     def add_table_info(self, table: pxt.Table) -> None:
         """Add information about the table into which we are inserting data"""
         assert isinstance(table, pxt.Table)
-        self.pxt_schema = table._schema
+        self.pxt_schema = table._get_schema()
         self.pxt_pk = table._tbl_version.get().primary_key
         for col in table._tbl_version_path.columns():
             if col.is_required_for_insert:
@@ -137,7 +139,7 @@ class DFTableDataConduit(TableDataConduit):
         t.pxt_df = tds.source
         return t
 
-    def infer_schema(self) -> dict[str, Any]:
+    def infer_schema(self) -> dict[str, ts.ColumnType]:
         self.pxt_schema = self.pxt_df.schema
         self.pxt_pk = self.src_pk
         return self.pxt_schema
@@ -149,7 +151,7 @@ class DFTableDataConduit(TableDataConduit):
 
 
 class RowDataTableDataConduit(TableDataConduit):
-    raw_rows: Optional[RowData] = None
+    raw_rows: RowData | None = None
     disable_mapping: bool = True
     batch_count: int = 0
 
@@ -168,7 +170,7 @@ class RowDataTableDataConduit(TableDataConduit):
         t.batch_count = 0
         return t
 
-    def infer_schema(self) -> dict[str, Any]:
+    def infer_schema(self) -> dict[str, ts.ColumnType]:
         from .datarows import _infer_schema_from_rows
 
         if self.source_column_map is None:
@@ -239,7 +241,7 @@ class PandasTableDataConduit(TableDataConduit):
         t.batch_count = 0
         return t
 
-    def infer_schema_part1(self) -> tuple[dict[str, Any], list[str]]:
+    def infer_schema_part1(self) -> tuple[dict[str, ts.ColumnType], list[str]]:
         """Return inferred schema, inferred primary key, and source column map"""
         if self.source_column_map is None:
             if self.src_schema_overrides is None:
@@ -252,7 +254,7 @@ class PandasTableDataConduit(TableDataConduit):
         else:
             raise NotImplementedError()
 
-    def infer_schema(self) -> dict[str, Any]:
+    def infer_schema(self) -> dict[str, ts.ColumnType]:
         self.pxt_schema, self.pxt_pk = self.infer_schema_part1()
         self.normalize_pxt_schema_types()
         _df_check_primary_key_values(self.pd_df, self.src_pk)
@@ -325,10 +327,13 @@ class JsonTableDataConduit(TableDataConduit):
 
 
 class HFTableDataConduit(TableDataConduit):
-    hf_ds: Optional[Union[datasets.Dataset, datasets.DatasetDict]] = None
-    column_name_for_split: Optional[str] = None
+    """
+    TODO:
+    - use set_format('arrow') and convert ChunkedArrays to PIL.Image.Image instead of going through numpy, which is slow
+    """
+
+    column_name_for_split: str | None = None
     categorical_features: dict[str, dict[int, str]]
-    hf_schema: dict[str, Any] = None
     dataset_dict: dict[str, datasets.Dataset] = None
     hf_schema_source: dict[str, Any] = None
 
@@ -340,9 +345,19 @@ class HFTableDataConduit(TableDataConduit):
         import datasets
 
         assert isinstance(tds.source, (datasets.Dataset, datasets.DatasetDict))
-        t.hf_ds = tds.source
         if 'column_name_for_split' in t.extra_fields:
             t.column_name_for_split = t.extra_fields['column_name_for_split']
+
+        # make sure we get numpy arrays for arrays, not Python lists
+        source = tds.source.with_format(type='numpy')
+        if isinstance(source, datasets.Dataset):
+            # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
+            raw_name = source.split._name
+            split_name = raw_name.split('[')[0] if raw_name is not None else None
+            t.dataset_dict = {split_name: source}
+        else:
+            assert isinstance(source, datasets.DatasetDict)
+            t.dataset_dict = source
         return t
 
     @classmethod
@@ -356,13 +371,13 @@ class HFTableDataConduit(TableDataConduit):
         except ImportError:
             return False
 
-    def infer_schema_part1(self) -> tuple[dict[str, Any], list[str]]:
+    def infer_schema_part1(self) -> tuple[dict[str, ts.ColumnType], list[str]]:
         from pixeltable.io.hf_datasets import _get_hf_schema, huggingface_schema_to_pxt_schema
 
         if self.source_column_map is None:
             if self.src_schema_overrides is None:
                 self.src_schema_overrides = {}
-            self.hf_schema_source = _get_hf_schema(self.hf_ds)
+            self.hf_schema_source = _get_hf_schema(self.source)
             self.src_schema = huggingface_schema_to_pxt_schema(
                 self.hf_schema_source, self.src_schema_overrides, self.src_pk
             )
@@ -397,15 +412,6 @@ class HFTableDataConduit(TableDataConduit):
     def prepare_insert(self) -> None:
         import datasets
 
-        if isinstance(self.source, datasets.Dataset):
-            # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
-            raw_name = self.source.split._name
-            split_name = raw_name.split('[')[0] if raw_name is not None else None
-            self.dataset_dict = {split_name: self.source}
-        else:
-            assert isinstance(self.source, datasets.DatasetDict)
-            self.dataset_dict = self.source
-
         # extract all class labels from the dataset to translate category ints to strings
         self.categorical_features = {
             feature_name: feature_type.names
@@ -416,26 +422,44 @@ class HFTableDataConduit(TableDataConduit):
             self.source_column_map = {}
         self.check_source_columns_are_insertable(self.hf_schema_source.keys())
 
-    def _translate_row(self, row: dict[str, Any], split_name: str) -> dict[str, Any]:
+    def _translate_row(self, row: dict[str, Any], split_name: str, features: datasets.Features) -> dict[str, Any]:
         output_row: dict[str, Any] = {}
         for col_name, val in row.items():
             # translate category ints to strings
             new_val = self.categorical_features[col_name][val] if col_name in self.categorical_features else val
             mapped_col_name = self.source_column_map.get(col_name, col_name)
 
-            # Convert values to the appropriate type if needed
-            try:
-                checked_val = self.pxt_schema[mapped_col_name].create_literal(new_val)
-            except TypeError as e:
-                msg = str(e)
-                raise excs.Error(f'Error in column {col_name}: {msg[0].lower() + msg[1:]}\nRow: {row}') from e
-            output_row[mapped_col_name] = checked_val
+            new_val = self._translate_val(new_val, features[col_name])
+            output_row[mapped_col_name] = new_val
 
         # add split name to output row
         if self.column_name_for_split is not None:
             output_row[self.column_name_for_split] = split_name
         return output_row
 
+    def _translate_val(self, val: Any, feature: datasets.Feature) -> Any:
+        """Convert numpy scalars to Python types and images to PIL.Image.Image"""
+        import datasets
+
+        if isinstance(feature, datasets.Value):
+            if isinstance(val, (np.generic, np.ndarray)):
+                # a scalar, which we want as a standard Python type
+                assert np.ndim(val) == 0
+                return val.item()
+            else:
+                # a standard Python object
+                return val
+        elif isinstance(feature, datasets.Sequence):
+            assert np.ndim(val) > 0
+            return val
+        elif isinstance(feature, datasets.Image):
+            return PIL.Image.fromarray(val)
+        elif isinstance(feature, dict):
+            assert isinstance(val, dict)
+            return {k: self._translate_val(v, feature[k]) for k, v in val.items()}
+        else:
+            return val
+
     def valid_row_batch(self) -> Iterator[RowData]:
         for split_name, split_dataset in self.dataset_dict.items():
             num_batches = split_dataset.size_in_bytes / self._K_BATCH_SIZE_BYTES
@@ -444,7 +468,7 @@ class HFTableDataConduit(TableDataConduit):
 
             batch = []
             for row in split_dataset:
-                batch.append(self._translate_row(row, split_name))
+                batch.append(self._translate_row(row, split_name, split_dataset.features))
                 if len(batch) >= tuples_per_batch:
                     yield batch
                     batch = []
@@ -454,7 +478,7 @@ class HFTableDataConduit(TableDataConduit):
 
 
 class ParquetTableDataConduit(TableDataConduit):
-    pq_ds: Optional[ParquetDataset] = None
+    pq_ds: ParquetDataset | None = None
 
     @classmethod
     def from_tds(cls, tds: TableDataConduit) -> 'ParquetTableDataConduit':
@@ -469,13 +493,13 @@ class ParquetTableDataConduit(TableDataConduit):
         t.pq_ds = parquet.ParquetDataset(str(input_path))
         return t
 
-    def infer_schema_part1(self) -> tuple[dict[str, Any], list[str]]:
-        from pixeltable.utils.arrow import ar_infer_schema
+    def infer_schema_part1(self) -> tuple[dict[str, ts.ColumnType], list[str]]:
+        from pixeltable.utils.arrow import to_pxt_schema
 
         if self.source_column_map is None:
             if self.src_schema_overrides is None:
                 self.src_schema_overrides = {}
-            self.src_schema = ar_infer_schema(self.pq_ds.schema, self.src_schema_overrides, self.src_pk)
+            self.src_schema = to_pxt_schema(self.pq_ds.schema, self.src_schema_overrides, self.src_pk)
             inferred_schema, inferred_pk, self.source_column_map = normalize_schema_names(
                 self.src_schema, self.src_pk, self.src_schema_overrides
             )
@@ -483,7 +507,7 @@ class ParquetTableDataConduit(TableDataConduit):
         else:
             raise NotImplementedError()
 
-    def infer_schema(self) -> dict[str, Any]:
+    def infer_schema(self) -> dict[str, ts.ColumnType]:
         self.pxt_schema, self.pxt_pk = self.infer_schema_part1()
         self.normalize_pxt_schema_types()
         self.prepare_insert()
@@ -504,7 +528,7 @@ class ParquetTableDataConduit(TableDataConduit):
         from pixeltable.utils.arrow import iter_tuples2
 
         try:
-            for fragment in self.pq_ds.fragments:  # type: ignore[attr-defined]
+            for fragment in self.pq_ds.fragments:
                 for batch in fragment.to_batches():
                     dict_batch = list(iter_tuples2(batch, self.source_column_map, self.pxt_schema))
                     self.total_rows += len(dict_batch)

pixeltable/io/utils.py

@@ -1,5 +1,5 @@
 from keyword import iskeyword as is_python_keyword
-from typing import Any, Optional, Union
+from typing import Any
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
@@ -21,7 +21,7 @@ def normalize_pxt_col_name(name: str) -> str:
     return id
 
 
-def normalize_primary_key_parameter(primary_key: Optional[Union[str, list[str]]] = None) -> list[str]:
+def normalize_primary_key_parameter(primary_key: str | list[str] | None = None) -> list[str]:
     if primary_key is None:
         primary_key = []
     elif isinstance(primary_key, str):
@@ -40,7 +40,7 @@ def normalize_schema_names(
     primary_key: list[str],
     schema_overrides: dict[str, Any],
     require_valid_pxt_column_names: bool = False,
-) -> tuple[dict[str, Any], list[str], Optional[dict[str, str]]]:
+) -> tuple[dict[str, Any], list[str], dict[str, str] | None]:
     """
     Convert all names in the input schema from source names to valid Pixeltable identifiers
     - Ensure that all names are unique.

pixeltable/iterators/__init__.py

@@ -1,3 +1,4 @@
+"""Iterators for splitting media and documents into components."""
 # ruff: noqa: F401
 
 from .audio import AudioSplitter
@@ -5,7 +6,7 @@ from .base import ComponentIterator
 from .document import DocumentSplitter
 from .image import TileIterator
 from .string import StringSplitter
-from .video import FrameIterator
+from .video import FrameIterator, VideoSplitter
 
 __default_dir = {symbol for symbol in dir() if not symbol.startswith('_')}
 __removed_symbols = {'base', 'document', 'video'}

pixeltable/iterators/audio.py

@@ -1,12 +1,12 @@
 import logging
-import uuid
 from fractions import Fraction
 from pathlib import Path
-from typing import Any, ClassVar, Optional
+from typing import Any, ClassVar
 
 import av
 
-from pixeltable import env, exceptions as excs, type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
+from pixeltable.utils.local_store import TempStore
 
 from .base import ComponentIterator
 
@@ -37,7 +37,7 @@ class AudioSplitter(ComponentIterator):
 
     # List of chunks to extract
     # Each chunk is defined by start and end presentation timestamps in audio file (int)
-    chunks_to_extract_in_pts: Optional[list[tuple[int, int]]]
+    chunks_to_extract_in_pts: list[tuple[int, int]] | None
     # next chunk to extract
     next_pos: int
 
@@ -55,12 +55,9 @@ class AudioSplitter(ComponentIterator):
     def __init__(
         self, audio: str, chunk_duration_sec: float, *, overlap_sec: float = 0.0, min_chunk_duration_sec: float = 0.0
     ):
-        if chunk_duration_sec <= 0.0:
-            raise excs.Error('chunk_duration_sec must be a positive number')
-        if chunk_duration_sec < min_chunk_duration_sec:
-            raise excs.Error('chunk_duration_sec must be at least min_chunk_duration_sec')
-        if overlap_sec >= chunk_duration_sec:
-            raise excs.Error('overlap_sec must be less than chunk_duration_sec')
+        assert chunk_duration_sec > 0.0
+        assert chunk_duration_sec >= min_chunk_duration_sec
+        assert overlap_sec < chunk_duration_sec
         audio_path = Path(audio)
         assert audio_path.exists() and audio_path.is_file()
         self.audio_path = audio_path
@@ -128,6 +125,19 @@ class AudioSplitter(ComponentIterator):
 
     @classmethod
     def output_schema(cls, *args: Any, **kwargs: Any) -> tuple[dict[str, ts.ColumnType], list[str]]:
+        param_names = ['chunk_duration_sec', 'min_chunk_duration_sec', 'overlap_sec']
+        params = dict(zip(param_names, args))
+        params.update(kwargs)
+
+        chunk_duration_sec = params['chunk_duration_sec']
+        min_chunk_duration_sec = params.get('min_chunk_duration_sec', 0.0)
+        overlap_sec = params.get('overlap_sec', 0.0)
+        if chunk_duration_sec <= 0.0:
+            raise excs.Error('chunk_duration_sec must be a positive number')
+        if chunk_duration_sec < min_chunk_duration_sec:
+            raise excs.Error('chunk_duration_sec must be at least min_chunk_duration_sec')
+        if overlap_sec >= chunk_duration_sec:
+            raise excs.Error('overlap_sec must be less than chunk_duration_sec')
         return {
             'start_time_sec': ts.FloatType(),
             'end_time_sec': ts.FloatType(),
@@ -140,7 +150,7 @@ class AudioSplitter(ComponentIterator):
         target_chunk_start, target_chunk_end = self.chunks_to_extract_in_pts[self.next_pos]
         chunk_start_pts = 0
         chunk_end_pts = 0
-        chunk_file = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}{self.audio_path.suffix}')
+        chunk_file = str(TempStore.create_path(extension=self.audio_path.suffix))
         output_container = av.open(chunk_file, mode='w')
         input_stream = self.container.streams.audio[0]
         codec_name = AudioSplitter.__codec_map.get(input_stream.codec_context.name, input_stream.codec_context.name)