sqlframe 1.1.3__py3-none-any.whl

sqlframe/base/types.py

@@ -0,0 +1,418 @@
+# This code is based on code from Apache Spark under the license found in the LICENSE file located in the 'sqlframe' folder.
+
+from __future__ import annotations
+
+import typing as t
+from decimal import Decimal
+
+from sqlframe.base.exceptions import RowError
+
+
+class DataType:
+    def __repr__(self) -> str:
+        return self.__class__.__name__ + "()"
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
+    def __eq__(self, other: t.Any) -> bool:
+        return isinstance(other, self.__class__) and self.__dict__ == other.__dict__
+
+    def __ne__(self, other: t.Any) -> bool:
+        return not self.__eq__(other)
+
+    def __str__(self) -> str:
+        return self.typeName()
+
+    @classmethod
+    def typeName(cls) -> str:
+        return cls.__name__[:-4].lower()
+
+    def simpleString(self) -> str:
+        return str(self)
+
+    def jsonValue(self) -> t.Union[str, t.Dict[str, t.Any]]:
+        return str(self)
+
+
+class DataTypeWithLength(DataType):
+    def __init__(self, length: int):
+        self.length = length
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.length})"
+
+    def __str__(self) -> str:
+        return f"{self.typeName()}({self.length})"
+
+
+class StringType(DataType):
+    pass
+
+
+class CharType(DataTypeWithLength):
+    pass
+
+
+class VarcharType(DataTypeWithLength):
+    pass
+
+
+class BinaryType(DataType):
+    pass
+
+
+class BooleanType(DataType):
+    pass
+
+
+class DateType(DataType):
+    pass
+
+
+class TimestampType(DataType):
+    pass
+
+
+class TimestampNTZType(DataType):
+    @classmethod
+    def typeName(cls) -> str:
+        return "timestamp_ntz"
+
+
+class DecimalType(DataType):
+    def __init__(self, precision: int = 10, scale: int = 0):
+        self.precision = precision
+        self.scale = scale
+
+    def simpleString(self) -> str:
+        return f"decimal({self.precision}, {self.scale})"
+
+    def jsonValue(self) -> str:
+        return f"decimal({self.precision}, {self.scale})"
+
+    def __repr__(self) -> str:
+        return f"DecimalType({self.precision}, {self.scale})"
+
+
+class DoubleType(DataType):
+    pass
+
+
+class FloatType(DataType):
+    pass
+
+
+class ByteType(DataType):
+    def __str__(self) -> str:
+        return "tinyint"
+
+
+class IntegerType(DataType):
+    def __str__(self) -> str:
+        return "int"
+
+
+class LongType(DataType):
+    def __str__(self) -> str:
+        return "bigint"
+
+
+class ShortType(DataType):
+    def __str__(self) -> str:
+        return "smallint"
+
+
+class ArrayType(DataType):
+    def __init__(self, elementType: DataType, containsNull: bool = True):
+        self.elementType = elementType
+        self.containsNull = containsNull
+
+    def __repr__(self) -> str:
+        return f"ArrayType({self.elementType, str(self.containsNull)}"
+
+    def simpleString(self) -> str:
+        return f"array<{self.elementType.simpleString()}>"
+
+    def jsonValue(self) -> t.Dict[str, t.Any]:
+        return {
+            "type": self.typeName(),
+            "elementType": self.elementType.jsonValue(),
+            "containsNull": self.containsNull,
+        }
+
+
+class MapType(DataType):
+    def __init__(self, keyType: DataType, valueType: DataType, valueContainsNull: bool = True):
+        self.keyType = keyType
+        self.valueType = valueType
+        self.valueContainsNull = valueContainsNull
+
+    def __repr__(self) -> str:
+        return f"MapType({self.keyType}, {self.valueType}, {str(self.valueContainsNull)})"
+
+    def simpleString(self) -> str:
+        return f"map<{self.keyType.simpleString()}, {self.valueType.simpleString()}>"
+
+    def jsonValue(self) -> t.Dict[str, t.Any]:
+        return {
+            "type": self.typeName(),
+            "keyType": self.keyType.jsonValue(),
+            "valueType": self.valueType.jsonValue(),
+            "valueContainsNull": self.valueContainsNull,
+        }
+
+
+class StructField(DataType):
+    def __init__(
+        self,
+        name: str,
+        dataType: DataType,
+        nullable: bool = True,
+        metadata: t.Optional[t.Dict[str, t.Any]] = None,
+    ):
+        self.name = name
+        self.dataType = dataType
+        self.nullable = nullable
+        self.metadata = metadata or {}
+
+    def __repr__(self) -> str:
+        return f"StructField('{self.name}', {self.dataType}, {str(self.nullable)})"
+
+    def simpleString(self) -> str:
+        return f"{self.name}:{self.dataType.simpleString()}"
+
+    def jsonValue(self) -> t.Dict[str, t.Any]:
+        return {
+            "name": self.name,
+            "type": self.dataType.jsonValue(),
+            "nullable": self.nullable,
+            "metadata": self.metadata,
+        }
+
+
+class StructType(DataType):
+    def __init__(self, fields: t.Optional[t.List[StructField]] = None):
+        if not fields:
+            self.fields = []
+            self.names = []
+        else:
+            self.fields = fields
+            self.names = [f.name for f in fields]
+
+    def __iter__(self) -> t.Iterator[StructField]:
+        return iter(self.fields)
+
+    def __len__(self) -> int:
+        return len(self.fields)
+
+    def __repr__(self) -> str:
+        return f"StructType({', '.join(str(field) for field in self)})"
+
+    def simpleString(self) -> str:
+        return f"struct<{', '.join(x.simpleString() for x in self)}>"
+
+    def jsonValue(self) -> t.Dict[str, t.Any]:
+        return {"type": self.typeName(), "fields": [x.jsonValue() for x in self]}
+
+    def fieldNames(self) -> t.List[str]:
+        return list(self.names)
+
+
+def _create_row(
+    fields: t.Union[Row, t.List[str]], values: t.Union[t.Tuple[t.Any, ...], t.List[t.Any]]
+) -> Row:
+    row = Row(*values)
+    row.__fields__ = fields
+    return row
+
+
+class Row(tuple):
+    """
+    A row in :class:`DataFrame`.
+    The fields in it can be accessed:
+
+    * like attributes (``row.key``)
+    * like dictionary values (``row[key]``)
+
+    ``key in row`` will search through row keys.
+
+    Row can be used to create a row object by using named arguments.
+    It is not allowed to omit a named argument to represent that the value is
+    None or missing. This should be explicitly set to None in this case.
+
+    .. versionchanged:: 3.0.0
+        Rows created from named arguments no longer have
+        field names sorted alphabetically and will be ordered in the position as
+        entered.
+
+    Examples
+    --------
+    >>> from pyspark.sql import Row
+    >>> row = Row(name="Alice", age=11)
+    >>> row
+    Row(name='Alice', age=11)
+    >>> row['name'], row['age']
+    ('Alice', 11)
+    >>> row.name, row.age
+    ('Alice', 11)
+    >>> 'name' in row
+    True
+    >>> 'wrong_key' in row
+    False
+
+    Row also can be used to create another Row like class, then it
+    could be used to create Row objects, such as
+
+    >>> Person = Row("name", "age")
+    >>> Person
+    <Row('name', 'age')>
+    >>> 'name' in Person
+    True
+    >>> 'wrong_key' in Person
+    False
+    >>> Person("Alice", 11)
+    Row(name='Alice', age=11)
+
+    This form can also be used to create rows as tuple values, i.e. with unnamed
+    fields.
+
+    >>> row1 = Row("Alice", 11)
+    >>> row2 = Row(name="Alice", age=11)
+    >>> row1 == row2
+    True
+    """
+
+    @t.overload
+    def __new__(cls, *args: str) -> Row: ...
+
+    @t.overload
+    def __new__(cls, **kwargs: t.Any) -> Row: ...
+
+    def __new__(cls, *args: t.Optional[str], **kwargs: t.Optional[t.Any]) -> Row:
+        if args and kwargs:
+            raise RowError("Cannot use both args and kwargs to create Row")
+        if kwargs:
+            # create row objects
+            # psycopg2 returns Decimal type for numeric while PySpark returns float so we convert to float
+            row = tuple.__new__(
+                cls, [float(x) if isinstance(x, Decimal) else x for x in kwargs.values()]
+            )
+            row.__fields__ = list(kwargs.keys())
+            return row
+        else:
+            # create row class or objects
+            return tuple.__new__(cls, args)
+
+    def asDict(self, recursive: bool = False) -> t.Dict[str, t.Any]:
+        """
+        Return as a dict
+
+        Parameters
+        ----------
+        recursive : bool, optional
+            turns the nested Rows to dict (default: False).
+
+        Notes
+        -----
+        If a row contains duplicate field names, e.g., the rows of a join
+        between two :class:`DataFrame` that both have the fields of same names,
+        one of the duplicate fields will be selected by ``asDict``. ``__getitem__``
+        will also return one of the duplicate fields, however returned value might
+        be different to ``asDict``.
+
+        Examples
+        --------
+        >>> from pyspark.sql import Row
+        >>> Row(name="Alice", age=11).asDict() == {'name': 'Alice', 'age': 11}
+        True
+        >>> row = Row(key=1, value=Row(name='a', age=2))
+        >>> row.asDict() == {'key': 1, 'value': Row(name='a', age=2)}
+        True
+        >>> row.asDict(True) == {'key': 1, 'value': {'name': 'a', 'age': 2}}
+        True
+        """
+        if not hasattr(self, "__fields__"):
+            raise RowError("Cannot convert a Row class into dict")
+
+        if recursive:
+
+            def conv(obj: t.Any) -> t.Any:
+                if isinstance(obj, Row):
+                    return obj.asDict(True)
+                elif isinstance(obj, list):
+                    return [conv(o) for o in obj]
+                elif isinstance(obj, dict):
+                    return dict((k, conv(v)) for k, v in obj.items())
+                else:
+                    return obj
+
+            return dict(zip(self.__fields__, (conv(o) for o in self)))
+        else:
+            return dict(zip(self.__fields__, self))
+
+    def __contains__(self, item: t.Any) -> bool:
+        if hasattr(self, "__fields__"):
+            return item in self.__fields__
+        else:
+            return super(Row, self).__contains__(item)
+
+    # let object acts like class
+    def __call__(self, *args: t.Any) -> Row:
+        """create new Row object"""
+        if len(args) > len(self):
+            raise RowError(
+                "Cannot create Row with fields %s. Expected %d values but got %s instead"
+                % (self, len(self), args)
+            )
+
+        return _create_row(self, args)
+
+    def __getitem__(self, item: t.Any) -> t.Any:
+        if isinstance(item, (int, slice)):
+            return super(Row, self).__getitem__(item)
+        try:
+            # it will be slow when it has many fields,
+            # but this will not be used in normal cases
+            idx = self.__fields__.index(item)
+            return super(Row, self).__getitem__(idx)
+        except IndexError:
+            raise KeyError(item)
+        except ValueError:
+            raise RowError(item)
+
+    def __getattr__(self, item: str) -> t.Any:
+        if item.startswith("__"):
+            raise AttributeError(item)
+        try:
+            # it will be slow when it has many fields,
+            # but this will not be used in normal cases
+            idx = self.__fields__.index(item)
+            return self[idx]
+        except IndexError:
+            raise AttributeError(item)
+        except ValueError:
+            raise AttributeError(item)
+
+    def __setattr__(self, key: t.Any, value: t.Any) -> None:
+        if key != "__fields__":
+            raise RuntimeError("Row is read-only")
+        self.__dict__[key] = value
+
+    def __reduce__(
+        self,
+    ) -> t.Union[str, t.Tuple[t.Any, ...]]:
+        """Returns a tuple so Python knows how to pickle Row."""
+        if hasattr(self, "__fields__"):
+            return (_create_row, (self.__fields__, tuple(self)))
+        else:
+            return tuple.__reduce__(self)
+
+    def __repr__(self) -> str:
+        """Printable representation of Row used in Python REPL."""
+        if hasattr(self, "__fields__"):
+            return "Row(%s)" % ", ".join(
+                "%s=%r" % (k, v) for k, v in zip(self.__fields__, tuple(self))
+            )
+        else:
+            return "<Row(%s)>" % ", ".join(repr(field) for field in self)

sqlframe/base/util.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import importlib
+import typing as t
+import unicodedata
+
+from sqlglot import expressions as exp
+from sqlglot.dialects.dialect import Dialect, DialectType
+from sqlglot.schema import ensure_column_mapping as sqlglot_ensure_column_mapping
+
+if t.TYPE_CHECKING:
+    from pandas.core.frame import DataFrame as PandasDataFrame
+    from pyspark.sql.dataframe import SparkSession as PySparkSession
+
+    from sqlframe.base import types
+    from sqlframe.base._typing import OptionalPrimitiveType, SchemaInput
+    from sqlframe.base.session import _BaseSession
+    from sqlframe.base.types import StructType
+
+
+def decoded_str(value: t.Union[str, bytes]) -> str:
+    if isinstance(value, bytes):
+        return value.decode("utf-8")
+    return value
+
+
+def schema_(
+    db: exp.Identifier | str,
+    catalog: t.Optional[exp.Identifier | str] = None,
+    quoted: t.Optional[bool] = None,
+) -> exp.Table:
+    """Build a Schema.
+
+    Args:
+        db: Database name.
+        catalog: Catalog name.
+        quoted: Whether to force quotes on the schema's identifiers.
+
+    Returns:
+        The new Schema instance.
+    """
+    return exp.Table(
+        this=None,
+        db=exp.to_identifier(db, quoted=quoted) if db else None,
+        catalog=exp.to_identifier(catalog, quoted=quoted) if catalog else None,
+    )
+
+
+def to_schema(
+    sql_path: t.Union[str, exp.Table], dialect: t.Optional[DialectType] = None
+) -> exp.Table:
+    if isinstance(sql_path, exp.Table) and sql_path.this is None:
+        return sql_path
+    table = exp.to_table(
+        sql_path.copy() if isinstance(sql_path, exp.Table) else sql_path, dialect=dialect
+    )
+    table.set("catalog", table.args.get("db"))
+    table.set("db", table.args.get("this"))
+    table.set("this", None)
+    return table
+
+
+def get_column_mapping_from_schema_input(
+    schema: SchemaInput, dialect: DialectType = None
+) -> t.Dict[str, t.Optional[exp.DataType]]:
+    from sqlframe.base import types
+
+    if isinstance(schema, dict):
+        value = schema
+    elif isinstance(schema, str):
+        col_name_type_strs = [x.strip() for x in schema.split(",")]
+        if len(col_name_type_strs) == 1 and len(col_name_type_strs[0].split(" ")) == 1:
+            value = {"value": col_name_type_strs[0].strip()}
+        else:
+            value = {
+                name_type_str.split(" ")[0].strip(): name_type_str.split(" ")[1].strip()
+                for name_type_str in col_name_type_strs
+            }
+    elif isinstance(schema, types.StructType):
+        value = {struct_field.name: struct_field.dataType.simpleString() for struct_field in schema}
+    else:
+        value = {x.strip(): None for x in schema}
+    return {
+        exp.to_column(k).sql(dialect=dialect): exp.DataType.build(v, dialect=dialect)
+        if v is not None
+        else v
+        for k, v in value.items()
+    }
+    # return {x.strip(): None for x in schema}  # type: ignore
+
+
+def get_tables_from_expression_with_join(expression: exp.Select) -> t.List[exp.Table]:
+    if not expression.args.get("joins"):
+        return []
+
+    left_table = expression.args["from"].this
+    other_tables = [join.this for join in expression.args["joins"]]
+    return [left_table] + other_tables
+
+
+def to_csv(options: t.Dict[str, OptionalPrimitiveType], equality_char: str = "=") -> str:
+    return ", ".join(
+        [f"{k}{equality_char}{v}" for k, v in (options or {}).items() if v is not None]
+    )
+
+
+def ensure_column_mapping(schema: t.Union[str, StructType]) -> t.Dict:
+    if isinstance(schema, str):
+        col_name_type_strs = [x.strip() for x in schema.split(",")]
+        schema = {  # type: ignore
+            name_type_str.split(" ")[0].strip(): name_type_str.split(" ")[1].strip()
+            for name_type_str in col_name_type_strs
+        }
+    # TODO: Make a protocol with a `simpleString` attribute as what it looks for instead of the actual
+    # `StructType` object.
+    elif hasattr(schema, "simpleString"):
+        return {struct_field.name: struct_field.dataType.simpleString() for struct_field in schema}
+    return sqlglot_ensure_column_mapping(schema)  # type: ignore
+
+
+# SO: https://stackoverflow.com/questions/37513355/converting-pandas-dataframe-into-spark-dataframe-error
+def get_equivalent_spark_type(pandas_type) -> types.DataType:
+    """
+    This method will retrieve the corresponding spark type given a pandas
+    type.
+
+    Args:
+        pandas_type (str): pandas data type
+
+    Returns:
+        spark data type
+    """
+    from sqlframe.base import types
+
+    type_map = {
+        "datetime64[ns]": types.TimestampType(),
+        "int64": types.LongType(),
+        "int32": types.IntegerType(),
+        "float64": types.DoubleType(),
+        "float32": types.FloatType(),
+    }
+    return type_map.get(str(pandas_type).lower(), types.StringType())
+
+
+def pandas_to_spark_schema(pandas_df: PandasDataFrame) -> types.StructType:
+    """
+    This method will return a spark dataframe schema given a pandas dataframe.
+
+    Args:
+        pandas_df (pandas.core.frame.DataFrame): pandas DataFrame
+
+    Returns:
+        equivalent spark DataFrame schema
+    """
+    from sqlframe.base import types
+
+    columns = list([x.replace("?column?", "unknown_column") for x in pandas_df.columns])
+    d_types = list(pandas_df.dtypes)
+    p_schema = types.StructType(
+        [
+            types.StructField(column, get_equivalent_spark_type(pandas_type))
+            for column, pandas_type in zip(columns, d_types)
+        ]
+    )
+    return p_schema
+
+
+def dialect_to_string(dialect: Dialect) -> str:
+    mapping = {v: k for k, v in Dialect.classes.items()}
+    return mapping[type(dialect)]
+
+
+def get_func_from_session(
+    name: str,
+    session: t.Optional[t.Union[_BaseSession, PySparkSession]] = None,
+    fallback: bool = True,
+) -> t.Callable:
+    from sqlframe.base.session import _BaseSession
+
+    session = session if session else _BaseSession()
+
+    if isinstance(session, _BaseSession):
+        dialect_str = dialect_to_string(session.input_dialect)
+        import_path = f"sqlframe.{dialect_str}.functions"
+    else:
+        import_path = "pyspark.sql.functions"
+    try:
+        func = getattr(importlib.import_module(import_path), name)
+    except AttributeError as e:
+        if not fallback:
+            raise e
+        func = getattr(importlib.import_module("sqlframe.base.functions"), name)
+        if session.output_dialect in func.unsupported_engines:  # type: ignore
+            raise NotImplementedError(
+                f"{name} is not supported by the engine: {session.output_dialect}"  # type: ignore
+            )
+    return func
+
+
+def soundex(s):
+    if not s:
+        return ""
+
+    s = unicodedata.normalize("NFKD", s)
+    s = s.upper()
+
+    replacements = (
+        ("BFPV", "1"),
+        ("CGJKQSXZ", "2"),
+        ("DT", "3"),
+        ("L", "4"),
+        ("MN", "5"),
+        ("R", "6"),
+    )
+    result = [s[0]]
+    count = 1
+
+    # find would-be replacment for first character
+    for lset, sub in replacements:
+        if s[0] in lset:
+            last = sub
+            break
+    else:
+        last = None
+
+    for letter in s[1:]:
+        for lset, sub in replacements:
+            if letter in lset:
+                if sub != last:
+                    result.append(sub)
+                    count += 1
+                last = sub
+                break
+        else:
+            if letter != "H" and letter != "W":
+                # leave last alone if middle letter is H or W
+                last = None
+        if count == 4:
+            break
+
+    result += "0" * (4 - count)
+    return "".join(result)