pixeltable 0.3.3__py3-none-any.whl → 0.3.5__py3-none-any.whl

pixeltable/type_system.py

@@ -8,12 +8,9 @@ import json
 import typing
 import urllib.parse
 import urllib.request
-from pathlib import Path
+from typing import Any, Iterable, Literal, Mapping, Optional, Sequence, Union
 
-from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
-from typing import Any, Iterable, Mapping, Optional, Sequence, Union
-
-import av  # type: ignore
+import av
 import jsonschema
 import jsonschema.protocols
 import jsonschema.validators
@@ -24,6 +21,9 @@ import sqlalchemy as sql
 from typing_extensions import _AnnotatedAlias
 
 import pixeltable.exceptions as excs
+from pixeltable.utils import parse_local_file_path
+
+from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
 
 
 class ColumnType:
@@ -47,8 +47,8 @@ class ColumnType:
         @classmethod
         def supertype(
             cls,
-            type1: 'ColumnType.Type',
-            type2: 'ColumnType.Type',
+            type1: Optional['ColumnType.Type'],
+            type2: Optional['ColumnType.Type'],
             # we need to pass this in because we can't easily append it as a class member
             common_supertypes: dict[tuple['ColumnType.Type', 'ColumnType.Type'], 'ColumnType.Type'],
         ) -> Optional['ColumnType.Type']:
@@ -93,6 +93,9 @@ class ColumnType:
         self._type = t
         self._nullable = nullable
 
+    def has_supertype(self) -> bool:
+        return True
+
     @property
     def nullable(self) -> bool:
         return self._nullable
@@ -213,9 +216,9 @@ class ColumnType:
             return self.copy(nullable=(self.nullable or other.nullable))
 
         if self.is_invalid_type():
-            return other
+            return other.copy(nullable=(self.nullable or other.nullable))
         if other.is_invalid_type():
-            return self
+            return self.copy(nullable=(self.nullable or other.nullable))
 
         if self.is_scalar_type() and other.is_scalar_type():
             t = self.Type.supertype(self._type, other._type, self.common_supertypes)
@@ -271,8 +274,10 @@ class ColumnType:
                 inferred_type = val_type
             else:
                 inferred_type = inferred_type.supertype(val_type)
-                if inferred_type is None:
-                    return None
+            if inferred_type is None:
+                return None
+            if not inferred_type.has_supertype():
+                return inferred_type
         return inferred_type
 
     @classmethod
@@ -292,26 +297,24 @@ class ColumnType:
                 designations will be allowed regardless.
         """
         origin = typing.get_origin(t)
+        type_args = typing.get_args(t)
         if origin is typing.Union:
             # Check if `t` has the form Optional[T].
-            union_args = typing.get_args(t)
-            if len(union_args) == 2 and type(None) in union_args:
+            if len(type_args) == 2 and type(None) in type_args:
                 # `t` is a type of the form Optional[T] (equivalently, Union[T, None] or Union[None, T]).
                 # We treat it as the underlying type but with nullable=True.
-                underlying_py_type = union_args[0] if union_args[1] is type(None) else union_args[1]
+                underlying_py_type = type_args[0] if type_args[1] is type(None) else type_args[1]
                 underlying = cls.from_python_type(underlying_py_type, allow_builtin_types=allow_builtin_types)
                 if underlying is not None:
                     return underlying.copy(nullable=True)
         elif origin is Required:
-            required_args = typing.get_args(t)
-            assert len(required_args) == 1
+            assert len(type_args) == 1
             return cls.from_python_type(
-                required_args[0], nullable_default=False, allow_builtin_types=allow_builtin_types
-            )
+                type_args[0], nullable_default=False, allow_builtin_types=allow_builtin_types
+            ).copy(nullable=False)
         elif origin is typing.Annotated:
-            annotated_args = typing.get_args(t)
-            origin = annotated_args[0]
-            parameters = annotated_args[1]
+            origin = type_args[0]
+            parameters = type_args[1]
             if isinstance(parameters, ColumnType):
                 return parameters.copy(nullable=nullable_default)
         else:
@@ -323,6 +326,11 @@ class ColumnType:
             if isinstance(t, type) and issubclass(t, _PxtType):
                 return t.as_col_type(nullable=nullable_default)
             elif allow_builtin_types:
+                if t is Literal and len(type_args) > 0:
+                    literal_type = cls.infer_common_literal_type(type_args)
+                    if literal_type is None:
+                        return None
+                    return literal_type.copy(nullable=(literal_type.nullable or nullable_default))
                 if t is str:
                     return StringType(nullable=nullable_default)
                 if t is int:
@@ -335,7 +343,7 @@ class ColumnType:
                     return TimestampType(nullable=nullable_default)
                 if t is PIL.Image.Image:
                     return ImageType(nullable=nullable_default)
-                if issubclass(t, Sequence) or issubclass(t, Mapping) or issubclass(t, pydantic.BaseModel):
+                if isinstance(t, type) and issubclass(t, (Sequence, Mapping, pydantic.BaseModel)):
                     return JsonType(nullable=nullable_default)
         return None
 
@@ -394,12 +402,9 @@ class ColumnType:
     def _validate_file_path(self, val: Any) -> None:
         """Raises TypeError if not a valid local file path or not a path/byte sequence"""
         if isinstance(val, str):
-            parsed = urllib.parse.urlparse(val)
-            if parsed.scheme != '' and parsed.scheme != 'file':
-                return
-            path = Path(urllib.parse.unquote(urllib.request.url2pathname(parsed.path)))
-            if not path.is_file():
-                raise TypeError(f'File not found: {str(path)}')
+            path = parse_local_file_path(val)
+            if path is not None and not path.is_file():
+                raise TypeError(f'File not found: {path}')
         else:
             if not isinstance(val, bytes):
                 raise TypeError(f'expected file path or bytes, got {type(val)}')
@@ -492,7 +497,7 @@ class InvalidType(ColumnType):
         super().__init__(self.Type.INVALID, nullable=nullable)
 
     def to_sa_type(self) -> sql.types.TypeEngine:
-        assert False
+        return sql.types.NullType()
 
     def print_value(self, val: Any) -> str:
         return str(val)
@@ -505,6 +510,9 @@ class StringType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.STRING, nullable=nullable)
 
+    def has_supertype(self):
+        return not self.nullable
+
     def to_sa_type(self) -> sql.types.TypeEngine:
         return sql.String()
 
@@ -588,6 +596,9 @@ class TimestampType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.TIMESTAMP, nullable=nullable)
 
+    def has_supertype(self):
+        return not self.nullable
+
     def to_sa_type(self) -> sql.types.TypeEngine:
         return sql.TIMESTAMP(timezone=True)
 
@@ -598,6 +609,8 @@ class TimestampType(ColumnType):
     def _create_literal(self, val: Any) -> Any:
         if isinstance(val, str):
             return datetime.datetime.fromisoformat(val)
+        if isinstance(val, datetime.datetime):
+            return val
         return val
 
 
@@ -648,6 +661,10 @@ class JsonType(ColumnType):
         return val_type.print_value(val)
 
     def _validate_literal(self, val: Any) -> None:
+        if isinstance(val, tuple):
+            val = list(val)
+        if isinstance(val, pydantic.BaseModel):
+            val = val.model_dump()
         if not self.__is_valid_json(val):
             raise TypeError(f'That literal is not a valid Pixeltable JSON object: {val}')
         if self.__validator is not None:
@@ -815,14 +832,20 @@ class ArrayType(ColumnType):
         return hash((self._type, self.nullable, self.shape, self.dtype))
 
     def supertype(self, other: ColumnType) -> Optional[ArrayType]:
+        basic_supertype = super().supertype(other)
+        if basic_supertype is not None:
+            assert isinstance(basic_supertype, ArrayType)
+            return basic_supertype
+
         if not isinstance(other, ArrayType):
             return None
+
         super_dtype = self.Type.supertype(self.dtype, other.dtype, self.common_supertypes)
         if super_dtype is None:
             # if the dtypes are incompatible, then the supertype is a fully general array
             return ArrayType(nullable=(self.nullable or other.nullable))
         super_shape: Optional[tuple[Optional[int], ...]]
-        if len(self.shape) != len(other.shape):
+        if self.shape is None or other.shape is None or len(self.shape) != len(other.shape):
             super_shape = None
         else:
             super_shape = tuple(n1 if n1 == n2 else None for n1, n2 in zip(self.shape, other.shape))
@@ -851,23 +874,39 @@ class ArrayType(ColumnType):
         dtype = None if d['dtype'] is None else cls.make_type(cls.Type(d['dtype']))
         return cls(shape, dtype, nullable=d['nullable'])
 
+    @classmethod
+    def from_np_dtype(cls, dtype: np.dtype, nullable: bool) -> Optional[ColumnType]:
+        """
+        Return pixeltable type corresponding to a given simple numpy dtype
+        """
+        if np.issubdtype(dtype, np.integer):
+            return IntType(nullable=nullable)
+
+        if np.issubdtype(dtype, np.floating):
+            return FloatType(nullable=nullable)
+
+        if dtype == np.bool_:
+            return BoolType(nullable=nullable)
+
+        if np.issubdtype(dtype, np.str_):
+            return StringType(nullable=nullable)
+
+        if np.issubdtype(dtype, np.character):
+            return StringType(nullable=nullable)
+
+        if np.issubdtype(dtype, np.datetime64):
+            return TimestampType(nullable=nullable)
+
+        return None
+
     @classmethod
     def from_literal(cls, val: np.ndarray, nullable: bool = False) -> Optional[ArrayType]:
         # determine our dtype
         assert isinstance(val, np.ndarray)
-        dtype: ColumnType
-        if np.issubdtype(val.dtype, np.integer):
-            dtype = IntType()
-        elif np.issubdtype(val.dtype, np.floating):
-            dtype = FloatType()
-        elif val.dtype == np.bool_:
-            dtype = BoolType()
-        elif np.issubdtype(val.dtype, np.str_):
-            # Note that this includes NumPy types like '<U1' -- arrays of single Unicode characters
-            dtype = StringType()
-        else:
+        pxttype: Optional[ColumnType] = cls.from_np_dtype(val.dtype, nullable)
+        if pxttype == None:
             return None
-        return cls(val.shape, dtype=dtype, nullable=nullable)
+        return cls(val.shape, dtype=pxttype, nullable=nullable)
 
     def is_valid_literal(self, val: np.ndarray) -> bool:
         if not isinstance(val, np.ndarray):
@@ -990,8 +1029,14 @@ class ImageType(ColumnType):
         return hash((self._type, self.nullable, self.size, self.mode))
 
     def supertype(self, other: ColumnType) -> Optional[ImageType]:
+        basic_supertype = super().supertype(other)
+        if basic_supertype is not None:
+            assert isinstance(basic_supertype, ImageType)
+            return basic_supertype
+
         if not isinstance(other, ImageType):
             return None
+
         width = self.width if self.width == other.width else None
         height = self.height if self.height == other.height else None
         mode = self.mode if self.mode == other.mode else None

pixeltable/utils/__init__.py

@@ -1,3 +1,10 @@
+import hashlib
+import urllib.parse
+import urllib.request
+from pathlib import Path
+from typing import Optional, Union
+
+
 def print_perf_counter_delta(delta: float) -> str:
     """Prints a performance counter delta in a human-readable format.
 
@@ -15,3 +22,37 @@ def print_perf_counter_delta(delta: float) -> str:
         return f'{delta * 1e3:.2f} ms'
     else:
         return f'{delta:.2f} s'
+
+
+def sha256sum(path: Union[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) -> Optional[Path]:
+    """
+    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

@@ -6,8 +6,10 @@ import pyarrow as pa
 
 import pixeltable.type_system as ts
 
-_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=datetime.timezone.utc): ts.TimestampType(nullable=True),
     pa.bool_(): ts.BoolType(nullable=True),
     pa.uint8(): ts.IntType(nullable=True),
     pa.int8(): ts.IntType(nullable=True),
@@ -16,9 +18,10 @@ _pa_to_pt: dict[pa.DataType, ts.ColumnType] = {
     pa.int32(): ts.IntType(nullable=True),
     pa.int64(): ts.IntType(nullable=True),
     pa.float32(): ts.FloatType(nullable=True),
+    pa.float64(): ts.FloatType(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.BoolType: pa.bool_(),
@@ -32,19 +35,20 @@ _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) -> Optional[ts.ColumnType]:
     """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)
     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=arrow_type.shape, dtype=dtype, nullable=nullable)
     else:
         return None
 
@@ -53,16 +57,25 @@ def to_arrow_type(pixeltable_type: ts.ColumnType) -> Optional[pa.DataType]:
     """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)
     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 ar_infer_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."""
+    ar_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 ar_schema
 
 
 def to_arrow_schema(pixeltable_schema: dict[str, Any]) -> pa.Schema:
@@ -96,3 +109,23 @@ def iter_tuples(batch: Union[pa.Table, pa.RecordBatch]) -> Iterator[dict[str, An
 
     for i in range(batch_size):
         yield {col_name: values[i] for col_name, values in pydict.items()}
+
+
+def iter_tuples2(
+    batch: Union[pa.Table, pa.RecordBatch], col_mapping: Optional[dict[str, str]], 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_name if col_mapping is None else col_mapping[col_name]): schema[pxt_name].create_literal(
+                values[i]
+            )
+            for col_name, values in pydict.items()
+        }

pixeltable/utils/formatter.py

@@ -6,7 +6,7 @@ import logging
 import mimetypes
 from typing import Any, Callable, Optional
 
-import av  # type: ignore[import-untyped]
+import av
 import numpy as np
 import PIL
 import PIL.Image as Image

pixeltable/utils/iceberg.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+from typing import Union
+
+from pyiceberg.catalog.sql import SqlCatalog
+
+
+def sqlite_catalog(warehouse_path: Union[str, Path], name: str = 'pixeltable') -> SqlCatalog:
+    """
+    Instantiate a sqlite Iceberg catalog at the specified path. If no catalog exists, one will be created.
+    """
+    if isinstance(warehouse_path, str):
+        warehouse_path = Path(warehouse_path)
+    warehouse_path.mkdir(exist_ok=True)
+    return SqlCatalog(name, uri=f'sqlite:///{warehouse_path}/catalog.db', warehouse=f'file://{warehouse_path}')

pixeltable/utils/media_store.py

@@ -30,7 +30,7 @@ class MediaStore:
         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 = Env.get().media_dir / tbl_id.hex / id_hex[:2] / id_hex[:4]
         parent.mkdir(parents=True, exist_ok=True)
         return parent / f'{tbl_id.hex}_{col_id}_{version}_{id_hex}{ext or ""}'