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

pixeltable/type_system.py

@@ -6,9 +6,8 @@ import enum
 import json
 import typing
 import urllib.parse
-from copy import copy
 from pathlib import Path
-from typing import Any, Optional, Tuple, Dict, Callable, List, Union
+from typing import Any, Optional, Tuple, Dict, Callable, List, Union, Sequence, Mapping
 
 import PIL.Image
 import av
@@ -240,19 +239,38 @@ class ColumnType:
 
     @classmethod
     def from_python_type(cls, t: type) -> Optional[ColumnType]:
-        if t in _python_type_to_column_type:
-            return _python_type_to_column_type[t]
-        elif isinstance(t, typing._UnionGenericAlias) and t.__args__[1] is type(None):
-            # `t` is a type of the form Optional[T] (equivalently, Union[T, None]).
-            # We treat it as the underlying type but with nullable=True.
-            if t.__args__[0] in _python_type_to_column_type:
-                underlying = copy(_python_type_to_column_type[t.__args__[0]])
-                underlying.nullable = True
-                return underlying
-
+        if typing.get_origin(t) is typing.Union:
+            union_args = typing.get_args(t)
+            if union_args[1] is type(None):
+                # `t` is a type of the form Optional[T] (equivalently, Union[T, None]).
+                # We treat it as the underlying type but with nullable=True.
+                underlying = cls.from_python_type(union_args[0])
+                if underlying is not None:
+                    underlying.nullable = True
+                    return underlying
+        else:
+            # Discard type parameters to ensure that parameterized types such as `list[T]`
+            # are correctly mapped to Pixeltable types.
+            base = typing.get_origin(t)
+            if base is None:
+                # No type parameters; the base type is just `t` itself
+                base = t
+            if base is str:
+                return StringType()
+            if base is int:
+                return IntType()
+            if base is float:
+                return FloatType()
+            if base is bool:
+                return BoolType()
+            if base is datetime.date or base is datetime.datetime:
+                return TimestampType()
+            if issubclass(base, Sequence) or issubclass(base, Mapping):
+                return JsonType()
+            if issubclass(base, PIL.Image.Image):
+                return ImageType()
         return None
 
-
     def validate_literal(self, val: Any) -> None:
         """Raise TypeError if val is not a valid literal for this type"""
         if val is None:
@@ -383,10 +401,6 @@ class ColumnType:
             return sql.VARBINARY
         assert False
 
-    @abc.abstractmethod
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        assert False, f'Have not implemented {self.__class__.__name__} to Arrow'
- 
     @staticmethod
     def no_conversion(v: Any) -> Any:
         """
@@ -413,9 +427,6 @@ class InvalidType(ColumnType):
     def to_sa_type(self) -> Any:
         assert False
 
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        assert False
-
     def print_value(self, val: Any) -> str:
         assert False
 
@@ -442,10 +453,6 @@ class StringType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.String
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        return pa.string()
 
     def print_value(self, val: Any) -> str:
         return f"'{val}'"
@@ -454,6 +461,13 @@ class StringType(ColumnType):
         if not isinstance(val, str):
             raise TypeError(f'Expected string, got {val.__class__.__name__}')
 
+    def _create_literal(self, val: Any) -> Any:
+        # Replace null byte within python string with space to avoid issues with Postgres.
+        # Use a space to avoid merging words.
+        # TODO(orm): this will also be an issue with JSON inputs, would space still be a good replacement?
+        if isinstance(val, str) and '\x00' in val:
+            return val.replace('\x00', ' ')
+        return val
 
 class IntType(ColumnType):
     def __init__(self, nullable: bool = False):
@@ -464,10 +478,6 @@ class IntType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.BigInteger
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        return pa.int64() # to be consistent with bigint above
 
     def _validate_literal(self, val: Any) -> None:
         if not isinstance(val, int):
@@ -483,10 +493,6 @@ class FloatType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.Float
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa
-        return pa.float32()
 
     def _validate_literal(self, val: Any) -> None:
         if not isinstance(val, float):
@@ -506,10 +512,6 @@ class BoolType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.Boolean
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        return pa.bool_()
 
     def _validate_literal(self, val: Any) -> None:
         if not isinstance(val, bool):
@@ -529,10 +531,6 @@ class TimestampType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.TIMESTAMP
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        return pa.timestamp('us') # postgres timestamp is microseconds
 
     def _validate_literal(self, val: Any) -> None:
         if not isinstance(val, datetime.datetime) and not isinstance(val, datetime.date):
@@ -570,10 +568,6 @@ class JsonType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.dialects.postgresql.JSONB
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        return pa.string() # TODO: weight advantage of pa.struct type.
 
     def print_value(self, val: Any) -> str:
         val_type = self.infer_literal_type(val)
@@ -669,7 +663,9 @@ class ArrayType(ColumnType):
 
     def _create_literal(self, val: Any) -> Any:
         if isinstance(val, (list,tuple)):
-            return np.array(val)
+            # map python float to whichever numpy float is
+            # declared for this type, rather than assume float64
+            return np.array(val, dtype=self.numpy_dtype())
         return val
 
     def to_sql(self) -> str:
@@ -677,12 +673,6 @@ class ArrayType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.LargeBinary
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        if any([n is None for n in self.shape]):
-            raise TypeError(f'Cannot convert array with unknown shape to Arrow')        
-        return pa.fixed_shape_tensor(pa.from_numpy_dtype(self.numpy_dtype()), self.shape)
 
     def numpy_dtype(self) -> np.dtype:
         if self.dtype == self.Type.INT:
@@ -788,10 +778,6 @@ class ImageType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.String
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        return pa.binary()
 
     def _validate_literal(self, val: Any) -> None:
         if isinstance(val, PIL.Image.Image):
@@ -815,10 +801,6 @@ class VideoType(ColumnType):
 
     def to_sa_type(self) -> str:
         return sql.String
-    
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa # pylint: disable=import-outside-toplevel
-        return pa.string()
 
     def _validate_literal(self, val: Any) -> None:
         self._validate_file_path(val)
@@ -854,10 +836,6 @@ class AudioType(ColumnType):
     def to_sa_type(self) -> str:
         return sql.String
 
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa  # pylint: disable=import-outside-toplevel
-        return pa.string()
-
     def _validate_literal(self, val: Any) -> None:
         self._validate_file_path(val)
 
@@ -901,10 +879,6 @@ class DocumentType(ColumnType):
     def to_sa_type(self) -> str:
         return sql.String
 
-    def to_arrow_type(self) -> 'pyarrow.DataType':
-        import pyarrow as pa  # pylint: disable=import-outside-toplevel
-        return pa.string()
-
     def _validate_literal(self, val: Any) -> None:
         self._validate_file_path(val)
 
@@ -919,20 +893,3 @@ class DocumentType(ColumnType):
                     raise excs.Error(f'Not a recognized document format: {val}')
             except Exception as e:
                 raise excs.Error(f'Not a recognized document format: {val}') from None
-
-
-# A dictionary mapping various Python types to their respective ColumnTypes.
-# This can be used to infer Pixeltable ColumnTypes from type hints on Python
-# functions. (Since Python functions do not necessarily have type hints, this
-# should always be an optional/convenience inference.)
-_python_type_to_column_type: dict[type, ColumnType] = {
-    str: StringType(),
-    int: IntType(),
-    float: FloatType(),
-    bool: BoolType(),
-    datetime.datetime: TimestampType(),
-    datetime.date: TimestampType(),
-    list: JsonType(),
-    dict: JsonType(),
-    PIL.Image.Image: ImageType()
-}

pixeltable/utils/arrow.py

@@ -0,0 +1,98 @@
+import logging
+from typing import Any, Dict, Iterable, Iterator, Optional
+
+import pyarrow as pa
+
+import pixeltable.type_system as ts
+
+_logger = logging.getLogger(__name__)
+
+_pa_to_pt: Dict[pa.DataType, ts.ColumnType] = {
+    pa.string(): ts.StringType(nullable=True),
+    pa.timestamp('us'): 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.int32(): ts.IntType(nullable=True),
+    pa.int64(): ts.IntType(nullable=True),
+    pa.float32(): ts.FloatType(nullable=True),
+}
+
+_pt_to_pa: Dict[ts.ColumnType, pa.DataType] = {
+    ts.StringType: pa.string(),
+    ts.TimestampType: pa.timestamp('us'),  # postgres timestamp is microseconds
+    ts.BoolType: pa.bool_(),
+    ts.IntType: pa.int64(),
+    ts.FloatType: pa.float32(),
+    ts.JsonType: pa.string(),  # TODO(orm) pa.struct() is possible
+    ts.ImageType: pa.binary(),  # inline image
+    ts.AudioType: pa.string(),  # path
+    ts.VideoType: pa.string(),  # path
+    ts.DocumentType: pa.string(),  # path
+}
+
+
+def to_pixeltable_type(arrow_type: pa.DataType) -> Optional[ts.ColumnType]:
+    """Convert a pyarrow DataType to a pixeltable ColumnType if one is defined.
+    Returns None if no conversion is currently implemented.
+    """
+    if arrow_type in _pa_to_pt:
+        return _pa_to_pt[arrow_type]
+    elif isinstance(arrow_type, pa.FixedShapeTensorType):
+        dtype = to_pixeltable_type(arrow_type.value_type)
+        if dtype is None:
+            return None
+        return ts.ArrayType(shape=arrow_type.shape, dtype=dtype)
+    else:
+        return None
+
+
+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__]
+    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 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())
+
+
+def to_pydict(batch: pa.RecordBatch) -> Dict[str, Iterable[Any]]:
+    """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 = {}
+    for k, name in enumerate(batch.schema.names):
+        col = batch.column(k)
+        if isinstance(col.type, pa.FixedShapeTensorType):
+            # treat array columns as numpy arrays to easily preserve numpy type
+            out[name] = col.to_numpy(zero_copy_only=False)
+        else:
+            # for the rest, use pydict to preserve python types
+            out[name] = col.to_pylist()
+
+    return out
+
+
+def iter_tuples(batch: 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'
+    for _, v in pydict.items():
+        batch_size = len(v)
+        break
+
+    for i in range(batch_size):
+        yield {col_name: values[i] for col_name, values in pydict.items()}

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/parquet.py

@@ -1,14 +1,24 @@
+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 io
-import pyarrow.parquet as pq
 import pyarrow as pa
-import numpy as np
-from pathlib import Path
-from collections import deque
+import pyarrow.parquet
 
-from typing import List, Tuple, Any, Dict
+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 = {}
@@ -20,18 +30,18 @@ def _write_batch(value_batch : Dict[str, deque], schema : pa.Schema, output_path
             pydict[field.name] = value_batch[field.name]
 
     tab = pa.Table.from_pydict(pydict, schema=schema)
-    pq.write_table(tab, output_path)
+    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 
+        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.
@@ -39,10 +49,9 @@ def save_parquet(df: 'pixeltable.DataFrame', dest_path: Path, partition_size_byt
     """
     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)))
 
-    arrow_schema = pa.schema([pa.field(name, column_types[i].to_arrow_type()) for i, name in enumerate(column_names)])
     # 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.
@@ -65,8 +74,9 @@ def save_parquet(df: 'pixeltable.DataFrame', dest_path: Path, partition_size_byt
                     # 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
-                        val = open(data_row.file_paths[e.slot_idx], 'rb').read()
-                    elif isinstance(val, PIL.Image.Image): 
+                        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')
@@ -109,18 +119,49 @@ def save_parquet(df: 'pixeltable.DataFrame', dest_path: Path, partition_size_byt
 
         _write_batch(current_value_batch, arrow_schema, temp_path / f'part-{batch_num:05d}.parquet')
 
-def get_part_metadata(path : Path) -> List[Tuple[str, int]]:
-    """
-        Args:
-            path: path to directory containing parquet files.
-        Returns:
-            A list of (file_name, num_rows) tuples for the parquet files in file name order.
-    """
-    parts = sorted([f for f in path.iterdir() if f.suffix == '.parquet'])
-    rows_per_file = {}
-
-    for part in parts:
-        parquet_file = pq.ParquetFile(str(part))
-        rows_per_file[part] = parquet_file.metadata.num_rows
 
-    return [(file, num_rows) for file, num_rows in rows_per_file.items()]
+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)