pydiverse-common 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

pydiverse/common/dtypes.py

@@ -49,20 +49,20 @@ class Dtype:
     @staticmethod
     def from_sql(sql_type) -> "Dtype":
         """Convert a SQL type to a Dtype."""
-        import sqlalchemy as sqa
+        import sqlalchemy as sa
 
-        if isinstance(sql_type, sqa.SmallInteger):
+        if isinstance(sql_type, sa.SmallInteger):
             return Int16()
-        if isinstance(sql_type, sqa.BigInteger):
+        if isinstance(sql_type, sa.BigInteger):
             return Int64()
-        if isinstance(sql_type, sqa.Integer):
+        if isinstance(sql_type, sa.Integer):
             return Int32()
-        if isinstance(sql_type, sqa.Float):
+        if isinstance(sql_type, sa.Float):
             precision = sql_type.precision or 53
             if precision <= 24:
                 return Float32()
             return Float64()
-        if isinstance(sql_type, sqa.Numeric | sqa.DECIMAL):
+        if isinstance(sql_type, sa.Numeric | sa.DECIMAL):
             # Just to be safe, we always use FLOAT64 for fixpoint numbers.
             # Databases are obsessed about fixpoint. However, in dataframes, it
             # is more common to just work with double precision floating point.
@@ -70,21 +70,21 @@ class Dtype:
             # Decimal to Float64 whenever it cannot guarantee semantic correctness
             # otherwise.
             return Float64()
-        if isinstance(sql_type, sqa.String):
+        if isinstance(sql_type, sa.String):
             return String()
-        if isinstance(sql_type, sqa.Boolean):
+        if isinstance(sql_type, sa.Boolean):
             return Bool()
-        if isinstance(sql_type, sqa.Date):
+        if isinstance(sql_type, sa.Date):
             return Date()
-        if isinstance(sql_type, sqa.Time):
+        if isinstance(sql_type, sa.Time):
             return Time()
-        if isinstance(sql_type, sqa.DateTime):
+        if isinstance(sql_type, sa.DateTime):
             return Datetime()
-        if isinstance(sql_type, sqa.Interval):
+        if isinstance(sql_type, sa.Interval):
             return Duration()
-        if isinstance(sql_type, sqa.ARRAY):
+        if isinstance(sql_type, sa.ARRAY):
             return List(Dtype.from_sql(sql_type.item_type))
-        if isinstance(sql_type, sqa.types.NullType):
+        if isinstance(sql_type, sa.types.NullType):
             return NullType()
 
         raise TypeError
@@ -184,7 +184,7 @@ class Dtype:
             raise TypeError
         if pa.types.is_decimal(arrow_type):
             # We don't recommend using Decimal in dataframes, but we support it.
-            return Decimal()
+            return Decimal(arrow_type.precision, arrow_type.scale)
         if pa.types.is_string(arrow_type):
             return String()
         if pa.types.is_boolean(arrow_type):
@@ -201,6 +201,11 @@ class Dtype:
             return NullType()
         if pa.types.is_list(arrow_type):
             return List(Dtype.from_arrow(arrow_type.value_type))
+        if pa.types.is_dictionary(arrow_type):
+            raise RuntimeError(
+                "Most likely this is an Enum type. But metadata about categories is "
+                "only in the pyarrow field and not in the pyarrow dtype"
+            )
         raise TypeError
 
     @staticmethod
@@ -212,6 +217,8 @@ class Dtype:
             return List(Dtype.from_polars(polars_type.inner))
         if isinstance(polars_type, pl.Enum):
             return Enum(*polars_type.categories)
+        if isinstance(polars_type, pl.Decimal):
+            return Decimal(polars_type.precision, polars_type.scale)
 
         return {
             pl.Int64: Int64(),
@@ -224,7 +231,6 @@ class Dtype:
             pl.UInt8: UInt8(),
             pl.Float64: Float64(),
             pl.Float32: Float32(),
-            pl.Decimal: Decimal(),
             pl.Utf8: String(),
             pl.Boolean: Bool(),
             pl.Datetime: Datetime(),
@@ -236,29 +242,28 @@ class Dtype:
 
     def to_sql(self):
         """Convert this Dtype to a SQL type."""
-        import sqlalchemy as sqa
+        import sqlalchemy as sa
 
         return {
-            Int(): sqa.BigInteger(),  # we default to 64 bit
-            Int8(): sqa.SmallInteger(),
-            Int16(): sqa.SmallInteger(),
-            Int32(): sqa.Integer(),
-            Int64(): sqa.BigInteger(),
-            UInt8(): sqa.SmallInteger(),
-            UInt16(): sqa.Integer(),
-            UInt32(): sqa.BigInteger(),
-            UInt64(): sqa.BigInteger(),
-            Float(): sqa.Float(53),  # we default to 64 bit
-            Float32(): sqa.Float(24),
-            Float64(): sqa.Float(53),
-            Decimal(): sqa.DECIMAL(),
-            String(): sqa.String(),
-            Bool(): sqa.Boolean(),
-            Date(): sqa.Date(),
-            Time(): sqa.Time(),
-            Datetime(): sqa.DateTime(),
-            Duration(): sqa.Interval(),
-            NullType(): sqa.types.NullType(),
+            Int(): sa.BigInteger(),  # we default to 64 bit
+            Int8(): sa.SmallInteger(),
+            Int16(): sa.SmallInteger(),
+            Int32(): sa.Integer(),
+            Int64(): sa.BigInteger(),
+            UInt8(): sa.SmallInteger(),
+            UInt16(): sa.Integer(),
+            UInt32(): sa.BigInteger(),
+            UInt64(): sa.BigInteger(),
+            Float(): sa.Float(53),  # we default to 64 bit
+            Float32(): sa.Float(24),
+            Float64(): sa.Float(53),
+            String(): sa.String(),
+            Bool(): sa.Boolean(),
+            Date(): sa.Date(),
+            Time(): sa.Time(),
+            Datetime(): sa.DateTime(),
+            Duration(): sa.Interval(),
+            NullType(): sa.types.NullType(),
         }[self]
 
     def to_pandas(self, backend: PandasBackend = PandasBackend.ARROW):
@@ -268,7 +273,7 @@ class Dtype:
         if backend == PandasBackend.NUMPY:
             return self.to_pandas_nullable(backend)
         if backend == PandasBackend.ARROW:
-            if self == String():
+            if isinstance(self, String) or isinstance(self, Enum):
                 return pd.StringDtype(storage="pyarrow")
             return pd.ArrowDtype(self.to_arrow())
 
@@ -316,7 +321,7 @@ class Dtype:
             Float(): pd.Float64Dtype(),  # we default to 64 bit
             Float32(): pd.Float32Dtype(),
             Float64(): pd.Float64Dtype(),
-            Decimal(): pd.Float64Dtype(),  # NumericDtype is
+            Decimal(): pd.Float64Dtype(),  # NumericDtype exists but is not used
             String(): pd.StringDtype(),
             Bool(): pd.BooleanDtype(),
             Date(): "datetime64[s]",
@@ -329,9 +334,6 @@ class Dtype:
         """Convert this Dtype to a PyArrow type."""
         import pyarrow as pa
 
-        if isinstance(self, Enum):
-            return pa.string()
-
         return {
             Int(): pa.int64(),  # we default to 64 bit
             Int8(): pa.int8(),
@@ -345,7 +347,6 @@ class Dtype:
             Float(): pa.float64(),  # we default to 64 bit
             Float32(): pa.float32(),
             Float64(): pa.float64(),
-            Decimal(): pa.decimal128(35, 10),  # Arbitrary precision
             String(): pa.string(),
             Bool(): pa.bool_(),
             Date(): pa.date32(),
@@ -355,6 +356,12 @@ class Dtype:
             NullType(): pa.null(),
         }[self]
 
+    def to_arrow_field(self, name: str, nullable: bool = True):
+        """Convert this Dtype to a PyArrow Field."""
+        import pyarrow as pa
+
+        return pa.field(name, self.to_arrow(), nullable=nullable)
+
     def to_polars(self: "Dtype"):
         """Convert this Dtype to a Polars type."""
         import polars as pl
@@ -372,8 +379,6 @@ class Dtype:
             Float(): pl.Float64,  # we default to 64 bit
             Float64(): pl.Float64,
             Float32(): pl.Float32,
-            Decimal(): pl.Decimal(scale=10),  # Arbitrary precision
-            String(): pl.Utf8,
             Bool(): pl.Boolean,
             Datetime(): pl.Datetime("us"),
             Duration(): pl.Duration("us"),
@@ -395,7 +400,43 @@ class Float64(Float): ...
 class Float32(Float): ...
 
 
-class Decimal(Float): ...
+class Decimal(Float):
+    def __init__(self, precision: int | None = None, scale: int | None = None):
+        """
+        Initialize a Decimal Dtype.
+
+        Default is Decimal(31,10) which is the highest precision that works with DB2.
+        If you like to save memory, Decimal(15,6) will get you quite far as well.
+
+        :param precision: total number of digits in the number
+            If not specified, it is assumed to be 31.
+        :param scale: number of digits after the decimal point
+            If not specified, it is assumed to be (precision//3+1).
+        """
+        self.precision = precision or 31
+        self.scale = scale or (self.precision // 3 + 1)
+
+    def to_sql(self):
+        import sqlalchemy as sa
+
+        return sa.Numeric(self.precision, self.scale)
+
+    def to_polars(self):
+        import polars as pl
+
+        return pl.Decimal(self.precision, self.scale)
+
+    def to_arrow(self):
+        import pyarrow as pa
+
+        if self.precision > 38:
+            return pa.decimal256(self.precision, self.scale)
+        elif self.precision > 18:
+            return pa.decimal128(self.precision, self.scale)
+        elif self.precision > 9:
+            return pa.decimal64(self.precision, self.scale)
+        else:
+            return pa.decimal32(self.precision, self.scale)
 
 
 class Int(Dtype):
@@ -428,7 +469,32 @@ class UInt16(Int): ...
 class UInt8(Int): ...
 
 
-class String(Dtype): ...
+class String(Dtype):
+    def __init__(self, max_length: int | None = None):
+        """
+        Initialize a String Dtype.
+
+        :param max_length: maximum length of string
+            This length will only be used for specifying fixed length strings in SQL.
+            Thus, the meaning of characters vs. bytes is dependent on the SQL dialect.
+        """
+        self.max_length = max_length
+
+    def to_sql(self):
+        """Convert this Dtype to a SQL type."""
+        import sqlalchemy as sa
+
+        return sa.String(length=self.max_length)
+
+    def to_polars(self):
+        import polars as pl
+
+        return pl.Utf8
+
+    def to_arrow(self):
+        import pyarrow as pa
+
+        return pa.string()
 
 
 class Bool(Dtype): ...
@@ -463,9 +529,9 @@ class List(Dtype):
         return f"List[{repr(self.inner)}]"
 
     def to_sql(self):
-        import sqlalchemy as sqa
+        import sqlalchemy as sa
 
-        return sqa.ARRAY(self.inner.to_sql())
+        return sa.ARRAY(self.inner.to_sql())
 
     def to_polars(self):
         import polars as pl
@@ -483,6 +549,9 @@ class Enum(String):
         if not all(isinstance(c, str) for c in categories):
             raise TypeError("arguments for `Enum` must have type `str`")
         self.categories = list(categories)
+        self.max_length = (
+            max([len(c) for c in categories]) if len(categories) > 0 else None
+        )
 
     def __eq__(self, rhs):
         return isinstance(rhs, Enum) and self.categories == rhs.categories
@@ -498,14 +567,23 @@ class Enum(String):
 
         return pl.Enum(self.categories)
 
-    def to_sql(self):
-        import sqlalchemy as sqa
-
-        return sqa.String()
-
     def to_arrow(self):
         import pyarrow as pa
 
-        # There is also pa.dictionary(), which seems to be kind of similar to an enum.
-        # Maybe it is better to convert to this.
+        # enum categories can only be maintained in pyarrow field (see to_arrow_field)
         return pa.string()
+
+    def to_arrow_field(self, name: str, nullable: bool = True):
+        """Convert this Dtype to a PyArrow Field."""
+        import pyarrow as pa
+
+        # try to mimic what polars does
+        return pa.field(
+            name,
+            pa.dictionary(pa.uint32(), pa.large_string()),
+            nullable=nullable,
+            metadata={
+                # the key might change with polars versions
+                "_PL_ENUM_VALUES2": "".join([f"{len(c)};{c}" for c in self.categories])
+            },
+        )

pydiverse/common/util/hashing.py

@@ -52,11 +52,14 @@ def hash_polars_dataframe(df: pl.DataFrame, use_init_repr=False) -> str:
                 ]:
                     df = df.with_columns(
                         pl.col(struct_col_name).struct.rename_fields(
-                            [stable_hash(struct_col_name, struct_field_name)]
+                            [stable_hash(struct_col_name, struct_dtype)]
                         )
-                        for struct_col_name, struct_field_name in struct_cols_and_dtypes
+                        for struct_col_name, struct_dtype in struct_cols_and_dtypes
                     ).unnest(
-                        struct_col_name for struct_col_name, _ in struct_cols_and_dtypes
+                        [
+                            struct_col_name
+                            for struct_col_name, _ in struct_cols_and_dtypes
+                        ]
                     )
                 return df
 

{pydiverse_common-0.3.8.dist-info → pydiverse_common-0.3.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydiverse-common
-Version: 0.3.8
+Version: 0.3.10
 Summary: Common functionality shared between pydiverse libraries
 Author: QuantCo, Inc.
 Author-email: Martin Trautmann <windiana@users.sf.net>, Finn Rudolph <finn.rudolph@t-online.de>

{pydiverse_common-0.3.8.dist-info → pydiverse_common-0.3.10.dist-info}/RECORD

@@ -1,5 +1,5 @@
 pydiverse/common/__init__.py,sha256=J7b4iStFyaEMYre_jdlZ4l_8dLyrMWCIpQdsMQcB8aI,806
-pydiverse/common/dtypes.py,sha256=LYZKaKYq_4uI4kUhoaCTTo5j1SRurswIOfN11Bkz25A,15986
+pydiverse/common/dtypes.py,sha256=ekeH6Ygnfy6sTsMVG4CcEdGHH_C6Z550tkdkaMQNIyY,18782
 pydiverse/common/testing.py,sha256=FcivI5wn0X3gzJhwnysKvCOgjSTTXaN6FtSFJ72jfSg,341
 pydiverse/common/version.py,sha256=1IU_m4r76_Qq0u-Tyo2_bERZFOkh0ZFueVzDqcCfLO0,336
 pydiverse/common/errors/__init__.py,sha256=FNeEfVbUa23b9sHkFsmxHYhY6sRgjaZysPQmlovpJrI,262
@@ -8,10 +8,10 @@ pydiverse/common/util/computation_tracing.py,sha256=HeXRHRUI8vxpzQ27Xcpa0StndSTP
 pydiverse/common/util/deep_map.py,sha256=JtY5ViWMMelOiLzPF7ZjzruCfB-bETISGxCk37qETxg,2540
 pydiverse/common/util/deep_merge.py,sha256=bV5p5_lsC-9nFah28EiEyG2h6U3Z5AuTqSooxOgCHN0,1929
 pydiverse/common/util/disposable.py,sha256=4XoGz70YRWA9TAqnUBvRCTAdsOGBviFN0gzxU7veY9o,993
-pydiverse/common/util/hashing.py,sha256=8Z1NybJ_zd3ONpn5annHGjowwArWkd2ZkCtlb3dtz_Q,4576
+pydiverse/common/util/hashing.py,sha256=EofnKULVKXv-S9kry0mOqHU5bxPGomCtr6XfYqIhGgc,4650
 pydiverse/common/util/import_.py,sha256=K7dSgz4YyrqEvqhoOzbwgD7D8HScMoO5XoSWtjbaoUs,4056
 pydiverse/common/util/structlog.py,sha256=xxhauxMuyxcKXTVg1MiPTkuvPBj8Zcr4o_v8Bq59Nig,3778
-pydiverse_common-0.3.8.dist-info/METADATA,sha256=ptAGp299BY9NSaM-XEaojLzhL_KVc0SEY-MFqqqAwL0,3399
-pydiverse_common-0.3.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydiverse_common-0.3.8.dist-info/licenses/LICENSE,sha256=AcE6SDVuAq6v9ZLE_8eOCe_NvSE0rAPR3NR7lSowYh4,1517
-pydiverse_common-0.3.8.dist-info/RECORD,,
+pydiverse_common-0.3.10.dist-info/METADATA,sha256=8xmLOqUJZbuETqz49ve_W8NNVkHGyaeOwT2exOhk0uY,3400
+pydiverse_common-0.3.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydiverse_common-0.3.10.dist-info/licenses/LICENSE,sha256=AcE6SDVuAq6v9ZLE_8eOCe_NvSE0rAPR3NR7lSowYh4,1517
+pydiverse_common-0.3.10.dist-info/RECORD,,