pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/utils/__init__.py

@@ -1,3 +1,9 @@
+import hashlib
+import urllib.parse
+import urllib.request
+from pathlib import Path
+
+
 def print_perf_counter_delta(delta: float) -> str:
     """Prints a performance counter delta in a human-readable format.
 
@@ -15,3 +21,37 @@ def print_perf_counter_delta(delta: float) -> str:
         return f'{delta * 1e3:.2f} ms'
     else:
         return f'{delta:.2f} s'
+
+
+def sha256sum(path: Path | str) -> str:
+    """
+    Compute the SHA256 hash of a file.
+    """
+    if isinstance(path, str):
+        path = Path(path)
+
+    h = hashlib.sha256()
+    with open(path, 'rb') as file:
+        while chunk := file.read(h.block_size):
+            h.update(chunk)
+
+    return h.hexdigest()
+
+
+def parse_local_file_path(file_or_url: str) -> Path | None:
+    """
+    Parses a string that may be either a URL or a local file path.
+
+    If the string is a local file path or a file-scheme URL (file://), then a Path object will be returned.
+    Otherwise, None will be returned.
+    """
+    parsed = urllib.parse.urlparse(file_or_url)
+    if len(parsed.scheme) <= 1:
+        # We're using `urlparse` to help distinguish file paths from URLs. If there is no scheme, then it's a file path.
+        # If there's a single-character scheme, we also interpret this as a file path; this insures that drive letters
+        # on Windows pathnames are correctly handled.
+        return Path(file_or_url).absolute()
+    elif parsed.scheme == 'file':
+        return Path(urllib.parse.unquote(urllib.request.url2pathname(parsed.path)))
+    else:
+        return None

pixeltable/utils/arrow.py

@@ -1,35 +1,49 @@
-import logging
-from typing import Any, Iterator, Optional, Union
+import datetime
+import io
+import json
+import uuid
+from typing import TYPE_CHECKING, Any, Iterator, cast
 
 import numpy as np
+import PIL.Image
 import pyarrow as pa
-import datetime
 
+import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
-from pixeltable.env import Env
-
-_tz_def = Env().get().default_time_zone
 
-_logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    import pixeltable as pxt
 
-_pa_to_pt: dict[pa.DataType, ts.ColumnType] = {
+PA_TO_PXT_TYPES: dict[pa.DataType, ts.ColumnType] = {
     pa.string(): ts.StringType(nullable=True),
+    pa.large_string(): ts.StringType(nullable=True),
+    pa.timestamp('us', tz='UTC'): ts.TimestampType(nullable=True),
     pa.bool_(): ts.BoolType(nullable=True),
-    pa.uint8(): ts.IntType(nullable=True),
     pa.int8(): ts.IntType(nullable=True),
-    pa.uint32(): ts.IntType(nullable=True),
-    pa.uint64(): ts.IntType(nullable=True),
+    pa.int16(): ts.IntType(nullable=True),
     pa.int32(): ts.IntType(nullable=True),
     pa.int64(): ts.IntType(nullable=True),
+    pa.uint8(): ts.IntType(nullable=True),
+    pa.uint16(): ts.IntType(nullable=True),
+    pa.uint32(): ts.IntType(nullable=True),
+    pa.uint64(): ts.IntType(nullable=True),
     pa.float32(): ts.FloatType(nullable=True),
+    pa.float64(): ts.FloatType(nullable=True),
+    pa.date32(): ts.DateType(nullable=True),
+    pa.date64(): ts.DateType(nullable=True),
+    pa.uuid(): ts.UUIDType(nullable=True),
+    pa.binary(): ts.BinaryType(nullable=True),
 }
 
-_pt_to_pa: dict[type[ts.ColumnType], pa.DataType] = {
+PXT_TO_PA_TYPES: dict[type[ts.ColumnType], pa.DataType] = {
     ts.StringType: pa.string(),
-    ts.TimestampType: pa.timestamp('us', tz=datetime.timezone.utc),  # postgres timestamp is microseconds
+    ts.TimestampType: pa.timestamp('us', tz='UTC'),  # postgres timestamp is microseconds
+    ts.DateType: pa.date32(),  # This could be date64
+    ts.UUIDType: pa.uuid(),
     ts.BoolType: pa.bool_(),
     ts.IntType: pa.int64(),
     ts.FloatType: pa.float32(),
+    ts.BinaryType: pa.binary(),
     ts.JsonType: pa.string(),  # TODO(orm) pa.struct() is possible
     ts.ImageType: pa.binary(),  # inline image
     ts.AudioType: pa.string(),  # path
@@ -38,48 +52,152 @@ _pt_to_pa: dict[type[ts.ColumnType], pa.DataType] = {
 }
 
 
-def to_pixeltable_type(arrow_type: pa.DataType) -> Optional[ts.ColumnType]:
+def to_pixeltable_type(arrow_type: pa.DataType, nullable: bool) -> ts.ColumnType | None:
     """Convert a pyarrow DataType to a pixeltable ColumnType if one is defined.
     Returns None if no conversion is currently implemented.
     """
     if isinstance(arrow_type, pa.TimestampType):
-        return ts.TimestampType(nullable=True)
-    elif arrow_type in _pa_to_pt:
-        return _pa_to_pt[arrow_type]
+        return ts.TimestampType(nullable=nullable)
+    elif arrow_type in PA_TO_PXT_TYPES:
+        pt = PA_TO_PXT_TYPES[arrow_type]
+        return pt.copy(nullable=nullable) if pt is not None else None
     elif isinstance(arrow_type, pa.FixedShapeTensorType):
-        dtype = to_pixeltable_type(arrow_type.value_type)
+        dtype = to_pixeltable_type(arrow_type.value_type, nullable)
         if dtype is None:
             return None
-        return ts.ArrayType(shape=arrow_type.shape, dtype=dtype)
+        return ts.ArrayType(shape=tuple(arrow_type.shape), dtype=dtype, nullable=nullable)
     else:
         return None
 
 
-def to_arrow_type(pixeltable_type: ts.ColumnType) -> Optional[pa.DataType]:
+def to_arrow_type(pixeltable_type: ts.ColumnType) -> pa.DataType | None:
     """Convert a pixeltable DataType to a pyarrow datatype if one is defined.
     Returns None if no conversion is currently implemented.
     """
-    if pixeltable_type.__class__ in _pt_to_pa:
-        return _pt_to_pa[pixeltable_type.__class__]
+    if pixeltable_type.__class__ in PXT_TO_PA_TYPES:
+        return PXT_TO_PA_TYPES[pixeltable_type.__class__]
     elif isinstance(pixeltable_type, ts.ArrayType):
-        return pa.fixed_shape_tensor(pa.from_numpy_dtype(pixeltable_type.numpy_dtype()), pixeltable_type.shape)
+        return pa.fixed_shape_tensor(pa.from_numpy_dtype(pixeltable_type.dtype), pixeltable_type.shape)
     else:
         return None
 
 
-def to_pixeltable_schema(arrow_schema: pa.Schema) -> dict[str, ts.ColumnType]:
-    return {field.name: to_pixeltable_type(field.type) for field in arrow_schema}
+def to_pxt_schema(
+    arrow_schema: pa.Schema, schema_overrides: dict[str, Any], primary_key: list[str]
+) -> dict[str, ts.ColumnType]:
+    """Convert a pyarrow Schema to a schema using pyarrow names and pixeltable types."""
+    pxt_schema = {
+        field.name: to_pixeltable_type(field.type, field.name not in primary_key)
+        if field.name not in schema_overrides
+        else schema_overrides[field.name]
+        for field in arrow_schema
+    }
+    return pxt_schema
 
 
 def to_arrow_schema(pixeltable_schema: dict[str, Any]) -> pa.Schema:
-    return pa.schema((name, to_arrow_type(typ)) for name, typ in pixeltable_schema.items())  # type: ignore[misc]
+    return pa.schema((name, to_arrow_type(typ)) for name, typ in pixeltable_schema.items())
+
+
+def _to_record_batch(column_vals: dict[str, list[Any]], schema: pa.Schema) -> pa.RecordBatch:
+    import pyarrow as pa
+
+    pa_arrays: list[pa.Array] = []
+    for field in schema:
+        if isinstance(field.type, pa.FixedShapeTensorType):
+            stacked_arr = np.stack(column_vals[field.name])
+            pa_arrays.append(pa.FixedShapeTensorArray.from_numpy_ndarray(stacked_arr))
+        else:
+            pa_array = cast(pa.Array, pa.array(column_vals[field.name]))
+            pa_arrays.append(pa_array)
+    return pa.RecordBatch.from_arrays(pa_arrays, schema=schema)
+
+
+def to_record_batches(query: 'pxt.Query', batch_size_bytes: int) -> Iterator[pa.RecordBatch]:
+    arrow_schema = to_arrow_schema(query.schema)
+    batch_columns: dict[str, list[Any]] = {k: [] for k in query.schema}
+    current_byte_estimate = 0
+    num_batch_rows = 0
+
+    # TODO: in order to avoid having to deal with ExprEvalError here, ResultSet should be an iterator
+    # over _exec()
+    try:
+        for data_row in query._exec():
+            num_batch_rows += 1
+            for (col_name, col_type), e in zip(query.schema.items(), query._select_list_exprs):
+                val = data_row[e.slot_idx]
+                val_size_bytes: int
+                if val is None:
+                    batch_columns[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[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):
+                        # no file available: save as png
+                        buf = io.BytesIO()
+                        val.save(buf, format='png')
+                        val = buf.getvalue()
+                    else:
+                        raise excs.Error(f'unknown image type {type(val)}')
+                    val_size_bytes = len(val)
+                elif col_type.is_string_type():
+                    val_size_bytes = len(val)
+                elif col_type.is_uuid_type():
+                    # pa.uuid() uses fixed_size_binary(16) as storage type
+                    val = val.bytes  # Convert UUID to 16-byte binary for arrow
+                    val_size_bytes = len(val)
+                elif col_type.is_binary_type():
+                    val_size_bytes = len(val)
+                elif col_type.is_media_type():
+                    assert data_row.file_paths[e.slot_idx] is not None
+                    val = data_row.file_paths[e.slot_idx]
+                    val_size_bytes = len(val)
+                elif col_type.is_json_type():
+                    val = json.dumps(val)
+                    val_size_bytes = len(val)
+                elif col_type.is_array_type():
+                    val_size_bytes = val.nbytes
+                elif col_type.is_int_type() or col_type.is_float_type():
+                    val_size_bytes = 8
+                elif col_type.is_bool_type():
+                    val_size_bytes = 1
+                elif col_type.is_date_type():
+                    val_size_bytes = 4
+                elif col_type.is_timestamp_type():
+                    val = val.astimezone(datetime.timezone.utc)
+                    val_size_bytes = 8
+                else:
+                    raise excs.Error(f'unknown type {col_type} for {col_name}')
 
+                batch_columns[col_name].append(val)
+                current_byte_estimate += val_size_bytes
 
-def to_pydict(batch: pa.RecordBatch) -> dict[str, Union[list, np.ndarray]]:
+            if current_byte_estimate > batch_size_bytes and num_batch_rows > 0:
+                record_batch = _to_record_batch(batch_columns, arrow_schema)
+                yield record_batch
+                batch_columns = {k: [] for k in query.schema}
+                current_byte_estimate = 0
+                num_batch_rows = 0
+
+    except excs.ExprEvalError as e:
+        query._raise_expr_eval_err(e)
+
+    if num_batch_rows > 0:
+        record_batch = _to_record_batch(batch_columns, arrow_schema)
+        yield record_batch
+
+
+def to_pydict(batch: pa.Table | pa.RecordBatch) -> dict[str, list | np.ndarray]:
     """Convert a RecordBatch to a dictionary of lists, unlike pa.lib.RecordBatch.to_pydict,
     this function will not convert numpy arrays to lists, and will preserve the original numpy dtype.
     """
-    out: dict[str, Union[list, np.ndarray]] = {}
+    out: dict[str, list | np.ndarray] = {}
     for k, name in enumerate(batch.schema.names):
         col = batch.column(k)
         if isinstance(col.type, pa.FixedShapeTensorType):
@@ -92,7 +210,7 @@ def to_pydict(batch: pa.RecordBatch) -> dict[str, Union[list, np.ndarray]]:
     return out
 
 
-def iter_tuples(batch: pa.RecordBatch) -> Iterator[dict[str, Any]]:
+def iter_tuples(batch: pa.Table | pa.RecordBatch) -> Iterator[dict[str, Any]]:
     """Convert a RecordBatch to an iterator of dictionaries. also works with pa.Table and pa.RowGroup"""
     pydict = to_pydict(batch)
     assert len(pydict) > 0, 'empty record batch'
@@ -102,3 +220,57 @@ def iter_tuples(batch: pa.RecordBatch) -> Iterator[dict[str, Any]]:
 
     for i in range(batch_size):
         yield {col_name: values[i] for col_name, values in pydict.items()}
+
+
+def _ar_val_to_pxt_val(val: Any, pxt_type: ts.ColumnType) -> Any:
+    """Convert a value to insertable format"""
+    if val is None:
+        return None
+    if pxt_type.is_float_type():
+        return float(val)
+    elif pxt_type.is_int_type():
+        return int(val)
+    elif pxt_type.is_bool_type():
+        return bool(val)
+    elif pxt_type.is_string_type():
+        return str(val)
+    elif pxt_type.is_uuid_type():
+        if isinstance(val, uuid.UUID):
+            return val
+        if isinstance(val, bytes):
+            return uuid.UUID(bytes=val)
+        return uuid.UUID(val)
+    elif pxt_type.is_binary_type():
+        assert isinstance(val, bytes)
+        return val
+    elif pxt_type.is_date_type():
+        if isinstance(val, str):
+            return datetime.date.fromisoformat(val)
+        if isinstance(val, datetime.date):
+            return val
+    elif pxt_type.is_timestamp_type():
+        if isinstance(val, str):
+            return datetime.datetime.fromisoformat(val)
+        if isinstance(val, datetime.datetime):
+            return val
+    elif pxt_type.is_array_type():
+        return pxt_type.create_literal(val)
+    raise ValueError(f'Unsupported type {pxt_type} for value {val}')
+
+
+def iter_tuples2(
+    batch: pa.Table | pa.RecordBatch, col_mapping: dict[str, str] | None, schema: dict[str, ts.ColumnType]
+) -> Iterator[dict[str, Any]]:
+    """Convert a RecordBatch to an iterator of dictionaries. also works with pa.Table and pa.RowGroup"""
+    pydict = to_pydict(batch)
+    assert len(pydict) > 0, 'empty record batch'
+    for _, v in pydict.items():
+        batch_size = len(v)
+        break
+
+    for i in range(batch_size):
+        # Convert a row to insertable format
+        yield {
+            (pxt_name := col_mapping.get(col_name, col_name)): _ar_val_to_pxt_val(values[i], schema[pxt_name])
+            for col_name, values in pydict.items()
+        }

pixeltable/utils/av.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from fractions import Fraction
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Iterator
+
+import av
+import av.stream
+import PIL.Image
+from typing_extensions import Self
+
+from pixeltable.env import Env
+
+# format -> (codec, extension)
+AUDIO_FORMATS: dict[str, tuple[str, str]] = {
+    'wav': ('pcm_s16le', 'wav'),
+    'mp3': ('libmp3lame', 'mp3'),
+    'flac': ('flac', 'flac'),
+    'mp4': ('aac', 'm4a'),
+}
+
+
+def get_metadata(path: str) -> dict:
+    with av.open(path) as container:
+        assert isinstance(container, av.container.InputContainer)
+        streams_info = [__get_stream_metadata(stream) for stream in container.streams]
+        result = {
+            'bit_exact': getattr(container, 'bit_exact', False),
+            'bit_rate': container.bit_rate,
+            'size': container.size,
+            'metadata': container.metadata,
+            'streams': streams_info,
+        }
+    return result
+
+
+def __get_stream_metadata(stream: av.stream.Stream) -> dict:
+    if stream.type not in ('audio', 'video'):
+        return {'type': stream.type}  # Currently unsupported
+
+    codec_context = stream.codec_context
+    codec_context_md: dict[str, Any] = {
+        'name': codec_context.name,
+        'codec_tag': codec_context.codec_tag.encode('unicode-escape').decode('utf-8'),
+        'profile': codec_context.profile,
+    }
+    metadata = {
+        'type': stream.type,
+        'duration': stream.duration,
+        'time_base': float(stream.time_base) if stream.time_base is not None else None,
+        'duration_seconds': float(stream.duration * stream.time_base)
+        if stream.duration is not None and stream.time_base is not None
+        else None,
+        'frames': stream.frames,
+        'metadata': stream.metadata,
+        'codec_context': codec_context_md,
+    }
+
+    if stream.type == 'audio':
+        # Additional metadata for audio
+        channels = getattr(stream.codec_context, 'channels', None)
+        codec_context_md['channels'] = int(channels) if channels is not None else None
+    else:
+        assert stream.type == 'video'
+        assert isinstance(stream, av.video.stream.VideoStream)
+        # Additional metadata for video
+        codec_context_md['pix_fmt'] = getattr(stream.codec_context, 'pix_fmt', None)
+        metadata.update(
+            **{
+                'width': stream.width,
+                'height': stream.height,
+                'frames': stream.frames,
+                'average_rate': float(stream.average_rate) if stream.average_rate is not None else None,
+                'base_rate': float(stream.base_rate) if stream.base_rate is not None else None,
+                'guessed_rate': float(stream.guessed_rate) if stream.guessed_rate is not None else None,
+            }
+        )
+
+    return metadata
+
+
+def get_video_duration(path: str) -> float | None:
+    """Return video duration in seconds."""
+    with av.open(path) as container:
+        video_stream = container.streams.video[0]
+        if video_stream is None:
+            return None
+        if video_stream.duration is not None:
+            return float(video_stream.duration * video_stream.time_base)
+
+        # if duration is not in the header, look for it in the last packet
+        last_pts: int | None = None
+        for packet in container.demux(video_stream):
+            if packet.pts is not None:
+                last_pts = packet.pts
+        if last_pts is not None:
+            return float(last_pts * video_stream.time_base)
+
+        return None
+
+
+def has_audio_stream(path: str) -> bool:
+    """Check if video has audio stream using PyAV."""
+    md = get_metadata(path)
+    return any(stream['type'] == 'audio' for stream in md['streams'])
+
+
+def ffmpeg_clip_cmd(
+    input_path: str,
+    output_path: str,
+    start_time: float,
+    duration: float | None = None,
+    fast: bool = True,
+    video_encoder: str | None = None,
+    video_encoder_args: dict[str, Any] | None = None,
+) -> list[str]:
+    cmd = ['ffmpeg']
+    if fast:
+        # fast: -ss before -i
+        cmd.extend(
+            [
+                '-ss',
+                str(start_time),
+                '-i',
+                input_path,
+                '-map',
+                '0',  # Copy all streams from input
+                '-c',
+                'copy',  # Stream copy (no re-encoding)
+            ]
+        )
+    else:
+        if video_encoder is None:
+            video_encoder = Env.get().default_video_encoder
+
+        # accurate: -ss after -i
+        cmd.extend(
+            [
+                '-i',
+                input_path,
+                '-ss',
+                str(start_time),
+                '-map',
+                '0',  # Copy all streams from input
+                '-c:a',
+                'copy',  # audio copy
+                '-c:s',
+                'copy',  # subtitle copy
+                '-c:v',
+                video_encoder,  # re-encode video
+            ]
+        )
+        if video_encoder_args is not None:
+            for k, v in video_encoder_args.items():
+                cmd.extend([f'-{k}', str(v)])
+
+    if duration is not None:
+        cmd.extend(['-t', str(duration)])
+    cmd.extend(['-loglevel', 'error', output_path])
+    return cmd
+
+
+def ffmpeg_segment_cmd(
+    input_path: str,
+    output_pattern: str,
+    segment_duration: float | None = None,
+    segment_times: list[float] | None = None,
+    video_encoder: str | None = None,
+    video_encoder_args: dict[str, Any] | None = None,
+) -> list[str]:
+    """Commandline for frame-accurate segmentation"""
+    assert (segment_duration is None) != (segment_times is None)
+    if video_encoder is None:
+        video_encoder = Env.get().default_video_encoder
+
+    cmd = [
+        'ffmpeg',
+        '-i',
+        input_path,
+        '-map',
+        '0',  # Copy all streams from input
+        '-c:a',
+        'copy',  # don't re-encode audio
+        '-c:v',
+        video_encoder,  # re-encode video
+    ]
+    if video_encoder_args is not None:
+        for k, v in video_encoder_args.items():
+            cmd.extend([f'-{k}', str(v)])
+    cmd.extend(['-f', 'segment'])
+
+    # -force_key_frames needs to precede -f segment
+    if segment_duration is not None:
+        cmd.extend(
+            [
+                '-force_key_frames',
+                f'expr:gte(t,n_forced*{segment_duration})',  # Force keyframe at each segment boundary
+                '-f',
+                'segment',
+                '-segment_time',
+                str(segment_duration),
+            ]
+        )
+    else:
+        assert segment_times is not None
+        times_str = ','.join([str(t) for t in segment_times])
+        cmd.extend(['-force_key_frames', times_str, '-f', 'segment', '-segment_times', times_str])
+
+    cmd.extend(
+        [
+            '-reset_timestamps',
+            '1',  # Reset timestamps for each segment
+            '-loglevel',
+            'error',  # Only show errors
+            output_pattern,
+        ]
+    )
+    return cmd
+
+
+class VideoFrames:
+    """
+    Context manager for iterating over video frames at a specified frame rate.
+
+    Args:
+        path: Path to the video file
+        fps: Number of frames to extract per second. If None or 0.0, extracts all frames.
+    """
+
+    path: Path
+    fps: float
+    container: av.container.input.InputContainer | None
+    video_framerate: Fraction | None
+    video_time_base: Fraction | None
+    video_start_time: int | None
+
+    @dataclass
+    class Item:
+        frame_idx: int
+        pts: int
+        dts: int
+        time: float
+        is_corrupt: bool
+        key_frame: bool
+        pict_type: int
+        interlaced_frame: bool
+        frame: PIL.Image.Image
+
+    def __init__(self, path: Path, fps: float | None = None) -> None:
+        self.path = path
+        self.fps = 0.0 if fps is None else fps
+        self.container = None
+        self.video_framerate = None
+        self.video_time_base = None
+        self.video_start_time = None
+
+    def __enter__(self) -> Self:
+        self.container = av.open(self.path)
+        stream = self.container.streams.video[0]
+        self.video_framerate = stream.average_rate
+        self.video_time_base = stream.time_base
+        self.video_start_time = stream.start_time or 0
+        return self
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> None:
+        # Clean up
+        if self.container:
+            self.container.close()
+
+    def __iter__(self) -> Iterator[Item]:
+        num_returned = 0
+        frame_idx = -1
+        while True:
+            try:
+                frame = next(self.container.decode(video=0))
+            except (StopIteration, EOFError):
+                return
+
+            frame_idx += 1
+            if self.fps == 0.0 or (num_returned <= frame.time * self.fps):
+                img = frame.to_image()
+                assert isinstance(img, PIL.Image.Image)
+                yield VideoFrames.Item(
+                    frame_idx=frame_idx,
+                    pts=frame.pts,
+                    dts=frame.dts,
+                    time=frame.time,
+                    is_corrupt=frame.is_corrupt,
+                    key_frame=frame.key_frame,
+                    pict_type=frame.pict_type,
+                    interlaced_frame=frame.interlaced_frame,
+                    frame=img,
+                )
+                num_returned += 1