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

pixeltable/io/lancedb.py

@@ -0,0 +1,3 @@
+from pixeltable.utils.lancedb import export_lancedb
+
+__all__ = ['export_lancedb']

pixeltable/io/pandas.py

@@ -1,18 +1,26 @@
-from typing import Any, Optional, Union
+import os
+import uuid
+from typing import Any
 
 import numpy as np
 import pandas as pd
+from pandas._typing import DtypeObj  # For pandas dtype type hints
+from pandas.api.types import is_datetime64_any_dtype, is_extension_array_dtype
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
+from pixeltable.env import Env
 
 
 def import_pandas(
-    tbl_name: str, df: pd.DataFrame, *, schema_overrides: Optional[dict[str, pxt.ColumnType]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    tbl_name: str,
+    df: pd.DataFrame,
+    *,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
-    comment: str = ''
+    comment: str = '',
 ) -> pxt.Table:
     """Creates a new base table from a Pandas
     [`DataFrame`](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html), with the
@@ -36,26 +44,24 @@ def import_pandas(
     Returns:
         A handle to the newly created [`Table`][pixeltable.Table].
     """
-    if schema_overrides is None:
-        schema_overrides = {}
-    if primary_key is None:
-        primary_key = []
-    elif isinstance(primary_key, str):
-        primary_key = [primary_key]
-
-    schema, pxt_pk = __df_to_pxt_schema(df, schema_overrides, primary_key)
-    tbl_rows = (dict(__df_row_to_pxt_row(row, schema)) for row in df.itertuples())
-    table = pxt.create_table(tbl_name, schema, primary_key=pxt_pk, num_retained_versions=num_retained_versions, comment=comment)
-    table.insert(tbl_rows)
-    return table
+    return pxt.create_table(
+        tbl_name,
+        source=df,
+        schema_overrides=schema_overrides,
+        primary_key=primary_key,
+        num_retained_versions=num_retained_versions,
+        comment=comment,
+    )
 
 
 def import_csv(
-    tbl_name: str, filepath_or_buffer, schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    tbl_name: str,
+    filepath_or_buffer: str | os.PathLike,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
     comment: str = '',
-    **kwargs
+    **kwargs: Any,
 ) -> pxt.Table:
     """
     Creates a new base table from a csv file. This is a convenience method and is equivalent
@@ -66,16 +72,26 @@ def import_csv(
     Returns:
         A handle to the newly created [`Table`][pixeltable.Table].
     """
-    df = pd.read_csv(filepath_or_buffer, **kwargs)
-    return import_pandas(tbl_name, df, schema_overrides=schema_overrides, primary_key=primary_key, num_retained_versions=num_retained_versions, comment=comment)
+    return pxt.create_table(
+        tbl_name,
+        source=filepath_or_buffer,
+        schema_overrides=schema_overrides,
+        primary_key=primary_key,
+        num_retained_versions=num_retained_versions,
+        comment=comment,
+        extra_args=kwargs,
+    )
 
 
 def import_excel(
-    tbl_name: str, io, *args, schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    tbl_name: str,
+    io: str | os.PathLike,
+    *,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
     comment: str = '',
-    **kwargs
+    **kwargs: Any,
 ) -> pxt.Table:
     """
     Creates a new base table from an Excel (.xlsx) file. This is a convenience method and is
@@ -86,97 +102,77 @@ def import_excel(
     Returns:
         A handle to the newly created [`Table`][pixeltable.Table].
     """
-    df = pd.read_excel(io, *args, **kwargs)
-    return import_pandas(tbl_name, df, schema_overrides=schema_overrides, primary_key=primary_key, num_retained_versions=num_retained_versions, comment=comment)
+    return pxt.create_table(
+        tbl_name,
+        source=io,
+        schema_overrides=schema_overrides,
+        primary_key=primary_key,
+        num_retained_versions=num_retained_versions,
+        comment=comment,
+        extra_args=kwargs,
+    )
+
+
+def _df_check_primary_key_values(df: pd.DataFrame, primary_key: list[str]) -> None:
+    for pd_name in primary_key:
+        # This can be faster for large DataFrames
+        has_nulls = df[pd_name].count() < len(df)
+        if has_nulls:
+            raise excs.Error(f'Primary key column `{pd_name}` cannot contain null values.')
 
 
-def __df_to_pxt_schema(
-    df: pd.DataFrame, schema_overrides: dict[str, pxt.ColumnType], primary_key: list[str]
-) -> tuple[dict[str, pxt.ColumnType], list[str]]:
+def df_infer_schema(
+    df: pd.DataFrame, schema_overrides: dict[str, ts.ColumnType], primary_key: list[str]
+) -> dict[str, ts.ColumnType]:
     """
     Infers a Pixeltable schema from a Pandas DataFrame.
 
     Returns:
         A tuple containing a Pixeltable schema and a list of primary key column names.
     """
-    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`.'
-            )
-    for pd_name in primary_key:
-        if pd_name not in df.columns:
-            raise excs.Error(f'Primary key column `{pd_name}` does not exist in the given `DataFrame`.')
-
-    schema: dict[str, pxt.ColumnType] = {}
-    col_mapping: dict[str, str] = {}  # Maps Pandas column names to Pixeltable column names
-
+    pd_schema: dict[str, ts.ColumnType] = {}
     for pd_name, pd_dtype in zip(df.columns, df.dtypes):
         if pd_name in schema_overrides:
+            assert isinstance(schema_overrides[pd_name], ts.ColumnType)
             pxt_type = schema_overrides[pd_name]
         else:
-            # This complicated-looking condition is necessary because we cannot safely call `pd.isna()` on
-            # general objects, so we need to check for nulls in the specific cases where we might expect them.
-            # isinstance(val, float) will check for NaN values in float columns *as well as* floats appearing
-            # in object columns (where Pandas uses NaN as a general null).
-            # np.issubdtype(pd_dtype, np.datetime64) checks for NaT values specifically in datetime columns.
-            has_na = any(
-                (isinstance(val, float) or np.issubdtype(pd_dtype, np.datetime64)) and pd.isna(val)
-                for val in df[pd_name]
-            )
-            if has_na and pd_name in primary_key:
-                raise excs.Error(f'Primary key column `{pd_name}` cannot contain null values.')
-            pxt_type = __np_dtype_to_pxt_type(pd_dtype, df[pd_name], pd_name not in primary_key)
-        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
-        col_mapping[pd_name] = pxt_name
-
-    pxt_pk = [col_mapping[pk] for pk in primary_key]
-    return schema, pxt_pk
-
-
-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
-
+            pxt_type = __pd_coltype_to_pxt_type(pd_dtype, df[pd_name], pd_name not in primary_key)
+        pd_schema[pd_name] = pxt_type
 
-def __np_dtype_to_pxt_type(np_dtype: np.dtype, data_col: pd.Series, nullable: bool) -> pxt.ColumnType:
-    """
-    Infers a Pixeltable type based on a Numpy dtype.
-    """
-    if np.issubdtype(np_dtype, np.integer):
-        return pxt.IntType(nullable=nullable)
+    return pd_schema
 
-    if np.issubdtype(np_dtype, np.floating):
-        return pxt.FloatType(nullable=nullable)
 
-    if np.issubdtype(np_dtype, np.bool_):
-        return pxt.BoolType(nullable=nullable)
+def __pd_dtype_to_pxt_type(pd_dtype: DtypeObj, nullable: bool) -> ts.ColumnType | None:
+    """
+    Determines a pixeltable ColumnType from a pandas dtype
 
-    if np.issubdtype(np_dtype, np.character):
-        return pxt.StringType(nullable=nullable)
+    Args:
+        pd_dtype: A pandas dtype object
 
-    if np.issubdtype(np_dtype, np.datetime64):
-        return pxt.TimestampType(nullable=nullable)
+    Returns:
+        ts.ColumnType: A pixeltable ColumnType
+    """
+    # Pandas extension arrays / types (Int64, boolean, string[pyarrow], etc.) are not directly
+    # compatible with NumPy dtypes
+    # The timezone-aware datetime64[ns, tz=] dtype is a pandas extension dtype
+    if is_datetime64_any_dtype(pd_dtype):
+        return ts.TimestampType(nullable=nullable)
+    if is_extension_array_dtype(pd_dtype):
+        return None
+    # Most other pandas dtypes are directly NumPy compatible
+    assert isinstance(pd_dtype, np.dtype)
+    return ts.ColumnType.from_np_dtype(pd_dtype, nullable)
+
+
+def __pd_coltype_to_pxt_type(pd_dtype: DtypeObj, data_col: pd.Series, nullable: bool) -> ts.ColumnType:
+    """
+    Infers a Pixeltable type based on a pandas dtype.
+    """
+    pxttype = __pd_dtype_to_pxt_type(pd_dtype, nullable)
+    if pxttype is not None:
+        return pxttype
 
-    if np_dtype == np.object_:
+    if pd_dtype == np.object_:
         # The `object_` dtype can mean all sorts of things; see if we can infer the Pixeltable type
         # based on the actual data in `data_col`.
         # First drop any null values (they don't contribute to type inference).
@@ -184,40 +180,65 @@ def __np_dtype_to_pxt_type(np_dtype: np.dtype, data_col: pd.Series, nullable: bo
 
         if len(data_col) == 0:
             # No non-null values; default to FloatType (the Pandas type of an all-NaN column)
-            return pxt.FloatType(nullable=nullable)
+            return ts.FloatType(nullable=nullable)
 
-        inferred_type = pxt.ColumnType.infer_common_literal_type(data_col)
+        inferred_type = ts.ColumnType.infer_common_literal_type(data_col)
         if inferred_type is None:
             # Fallback on StringType if everything else fails
-            return pxt.StringType(nullable=nullable)
+            return ts.StringType(nullable=nullable)
         else:
             return inferred_type.copy(nullable=nullable)
 
-    raise excs.Error(f'Could not infer Pixeltable type of column: {data_col.name} (dtype: {np_dtype})')
+    raise excs.Error(f'Could not infer Pixeltable type of column: {data_col.name} (dtype: {pd_dtype})')
 
 
-def __df_row_to_pxt_row(row: tuple[Any, ...], schema: dict[str, pxt.ColumnType]) -> dict[str, Any]:
-    rows = {}
+def _df_row_to_pxt_row(
+    row: tuple[Any, ...], schema: dict[str, ts.ColumnType], col_mapping: dict[str, str] | None
+) -> dict[str, Any]:
+    """Convert a row to insertable format"""
+    pxt_row: dict[str, Any] = {}
     for val, (col_name, pxt_type) in zip(row[1:], schema.items()):
+        pxt_name = col_mapping.get(col_name, col_name)
+        nval: Any
         if pxt_type.is_float_type():
-            val = float(val)
+            nval = 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
+            nval = None
         elif pxt_type.is_int_type():
-            val = int(val)
+            nval = int(val)
         elif pxt_type.is_bool_type():
-            val = bool(val)
+            nval = bool(val)
         elif pxt_type.is_string_type():
-            val = str(val)
+            nval = str(val)
+        elif pxt_type.is_date_type():
+            if pd.isnull(val):
+                # pandas has the bespoke 'NaT' valud for a missing timestamp
+                # This is not supported by postgres, and must be converted to None
+                nval = None
+            else:
+                nval = pd.Timestamp(val).date()
         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
+                # pandas has the bespoke 'NaT' value for a missing timestamp
+                # This is not supported by postgres, and must be converted to None
+                nval = None
             else:
-                val = pd.Timestamp(val).to_pydatetime()
-        rows[col_name] = val
-    return rows
+                tval = pd.Timestamp(val)
+                # pandas supports tz-aware and naive timestamps.
+                if tval.tz is None:
+                    nval = pd.Timestamp(tval).tz_localize(tz=Env.get().default_time_zone)
+                else:
+                    nval = tval.astimezone(Env.get().default_time_zone)
+        elif pxt_type.is_uuid_type():
+            if pd.isnull(val):
+                nval = None
+            elif isinstance(val, uuid.UUID):
+                nval = val
+            else:
+                nval = uuid.UUID(val)
+        else:
+            nval = val
+        pxt_row[pxt_name] = nval
+    return pxt_row

pixeltable/io/parquet.py

@@ -1,168 +1,78 @@
 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 Any, Optional, Union
-
-import numpy as np
-import PIL.Image
-import datetime
+from typing import Any
 
 import pixeltable as pxt
-from pixeltable.env import Env
-import pixeltable.exceptions as exc
-import pixeltable.type_system as ts
+import pixeltable.exceptions as excs
+from pixeltable.catalog import Catalog
 from pixeltable.utils.transactional_directory import transactional_directory
 
 if typing.TYPE_CHECKING:
-    import pyarrow as pa
     import pixeltable as pxt
 
-_logger = logging.getLogger(__name__)
-
-
-def _write_batch(value_batch: dict[str, deque], schema: pa.Schema, output_path: Path) -> None:
-    import pyarrow as pa
-    from pyarrow import parquet
-
-    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)
-    parquet.write_table(tab, str(output_path))
+_logger = logging.getLogger('pixeltable')
 
 
 def export_parquet(
-            table_or_df: Union[pxt.Table, pxt.DataFrame],
-            parquet_path: Path,
-            partition_size_bytes: int = 100_000_000,
-            inline_images: bool = False
-            ) -> None:
+    table_or_query: pxt.Table | pxt.Query,
+    parquet_path: Path,
+    partition_size_bytes: int = 100_000_000,
+    inline_images: bool = False,
+) -> None:
     """
-    Exports a dataframe's data to one or more Parquet files. Requires pyarrow to be installed.
+    Exports a Query's data to one or more Parquet files. Requires pyarrow to be installed.
 
     It additionally writes the pixeltable metadata in a json file, which would otherwise
     not be available in the parquet format.
 
     Args:
-        table_or_df : Table or Dataframe to export.
+        table_or_query : Table or Query to export.
         parquet_path : Path to directory to write the parquet files to.
         partition_size_bytes : The maximum target size for each chunk. Default 100_000_000 bytes.
         inline_images : If True, images are stored inline in the parquet file. This is useful
                         for small images, to be imported as pytorch dataset. But can be inefficient
                         for large images, and cannot be imported into pixeltable.
-                        If False, will raise an error if the Dataframe has any image column.
+                        If False, will raise an error if the Query has any image column.
                         Default False.
     """
-    from pixeltable.utils.arrow import to_arrow_schema
+    import pyarrow as pa
 
-    df: pxt.DataFrame
-    if isinstance(table_or_df, pxt.catalog.Table):
-        df = table_or_df._df()
-    else:
-        df = table_or_df
+    from pixeltable.utils.arrow import to_record_batches
 
-    type_dict = {k: v.as_dict() for k, v in df.schema.items()}
-    arrow_schema = to_arrow_schema(df.schema)
+    query: pxt.Query
+    if isinstance(table_or_query, pxt.catalog.Table):
+        query = table_or_query.select()
+    else:
+        query = table_or_query
 
-    if not inline_images and any(col_type.is_image_type() for col_type in df.schema.values()):
-        raise exc.Error('Cannot export Dataframe with image columns when inline_images is False')
+    if not inline_images and any(col_type.is_image_type() for col_type in query.schema.values()):
+        raise excs.Error('Cannot export Query with image columns when inline_images is False')
 
     # store the changes atomically
     with transactional_directory(parquet_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'))
+        json.dump(query.as_dict(), (temp_path / '.pixeltable.json').open('w'))
+        type_dict = {k: v.as_dict() for k, v in query.schema.items()}
         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 df.schema.keys()}
-        current_byte_estimate = 0
-
-        for data_row in df._exec():
-            for (col_name, col_type), e in zip(df.schema.items(), df._select_list_exprs):
-                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():
-                    val = val.astimezone(datetime.timezone.utc)
-                    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')
+        with Catalog.get().begin_xact(for_write=False):
+            for record_batch in to_record_batches(query, partition_size_bytes):
+                output_path = temp_path / f'part-{batch_num:05d}.parquet'
+                arrow_tbl = pa.Table.from_batches([record_batch])
+                pa.parquet.write_table(arrow_tbl, str(output_path))
                 batch_num += 1
-                current_value_batch = {k: deque() for k in df.schema.keys()}
-                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."""
-    from pyarrow import parquet
-
-    from pixeltable.utils.arrow import to_pixeltable_schema
-
-    input_path = Path(parquet_path).expanduser()
-    parquet_dataset = parquet.ParquetDataset(str(input_path))
-    return to_pixeltable_schema(parquet_dataset.schema)
 
 
 def import_parquet(
     table: str,
     *,
     parquet_path: str,
-    schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     **kwargs: Any,
 ) -> pxt.Table:
     """Creates a new base table from a Parquet file or set of files. Requires pyarrow to be installed.
@@ -174,41 +84,18 @@ def import_parquet(
             name `name` will be given type `type`, instead of being inferred from the Parquet dataset. The keys in
             `schema_overrides` should be the column names of the Parquet dataset (whether or not they are valid
             Pixeltable identifiers).
+        primary_key: The primary key of the table (see [`create_table()`][pixeltable.create_table]).
         kwargs: Additional arguments to pass to `create_table`.
 
     Returns:
         A handle to the newly created table.
     """
-    from pyarrow import parquet
-
-    import pixeltable as pxt
-    from pixeltable.utils.arrow import iter_tuples
-
-    input_path = Path(parquet_path).expanduser()
-    parquet_dataset = parquet.ParquetDataset(str(input_path))
-
-    schema = parquet_schema_to_pixeltable_schema(parquet_path)
-    if schema_overrides is None:
-        schema_overrides = {}
-
-    schema.update(schema_overrides)
-    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 in pxt.list_tables():
-        raise exc.Error(f'Table {table} already exists')
-
-    try:
-        tmp_name = f'{table}_tmp_{random.randint(0, 100000000)}'
-        tab = pxt.create_table(tmp_name, schema, **kwargs)
-        for fragment in parquet_dataset.fragments:  # type: ignore[attr-defined]
-            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)
-    return pxt.get_table(table)
+    value = kwargs.pop('source_format', None)
+    return pxt.create_table(
+        table,
+        source=parquet_path,
+        source_format=value,
+        schema_overrides=schema_overrides,
+        primary_key=primary_key,
+        extra_args=kwargs,
+    )