apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/plugins/plotly_extensions.py

@@ -0,0 +1,150 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+import pathlib
+from collections.abc import Collection
+from typing import IO, Any
+
+try:
+    import plotly.graph_objects
+except ImportError as e:
+    raise NotImplementedError("Plotly is not installed.") from e
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataSaver
+
+
+@dataclasses.dataclass
+class PlotlyStaticWriter(DataSaver):
+    """Write Plotly figure as static image format
+    ref: https://plotly.com/python/static-image-export/
+    """
+
+    path: str | pathlib.Path | IO
+    format: str | None = None
+    width: int | None = None
+    height: int | None = None
+    scale: int | float | None = None
+    validate: bool = True
+    engine: str = "auto"
+
+    def _get_saving_kwargs(self) -> dict:
+        kwargs = {}
+        if self.format is not None:
+            kwargs["format"] = self.format
+        if self.width is not None:
+            kwargs["width"] = self.width
+        if self.height is not None:
+            kwargs["height"] = self.height
+        if self.scale is not None:
+            kwargs["scale"] = self.scale
+        if self.validate is not None:
+            kwargs["validate"] = self.validate
+        if self.engine is not None:
+            kwargs["engine"] = self.engine
+
+        return kwargs
+
+    def save_data(self, data: plotly.graph_objects.Figure) -> dict[str, Any]:
+        data.write_image(file=self.path, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [plotly.graph_objects.Figure]
+
+    @classmethod
+    def name(cls) -> str:
+        return "plotly"
+
+
+@dataclasses.dataclass
+class PlotlyInteractiveWriter(DataSaver):
+    """Write Plotly figure as interactive HTML + JS
+    ref: https://plotly.com/python/interactive-html-export/
+    """
+
+    path: str | pathlib.Path | IO
+    config: dict | None = None
+    auto_play: bool = True
+    include_plotlyjs: bool | str = (
+        True  # or "cdn", "directory", "require", "False", "other string .js"
+    )
+    include_mathjax: bool | str = False  # "cdn", "string .js"
+    post_script: str | list[str] | None = None
+    full_html: bool = True
+    animation_opts: dict | None = None
+    default_width: int | float | str = "100%"
+    default_height: int | float | str = "100%"
+    validate: bool = True
+    auto_open: bool = True
+    div_id: str | None = None
+
+    def _get_saving_kwargs(self) -> dict:
+        kwargs = {}
+        if self.config is not None:
+            kwargs["config"] = self.config
+        if self.auto_play is not None:
+            kwargs["auto_play"] = self.auto_play
+        if self.include_plotlyjs is not None:
+            kwargs["include_plotlyjs"] = self.include_plotlyjs
+        if self.include_mathjax is not None:
+            kwargs["include_mathjax"] = self.include_mathjax
+        if self.post_script is not None:
+            kwargs["post_script"] = self.post_script
+        if self.full_html is not None:
+            kwargs["full_html"] = self.full_html
+        if self.animation_opts is not None:
+            kwargs["animation_opts"] = self.animation_opts
+        if self.default_width is not None:
+            kwargs["default_width"] = self.default_width
+        if self.default_height is not None:
+            kwargs["default_height"] = self.default_height
+        if self.validate is not None:
+            kwargs["validate"] = self.validate
+        if self.auto_open is not None:
+            kwargs["auto_open"] = self.auto_open
+        if self.div_id is not None:
+            kwargs["div_id"] = self.div_id
+        return kwargs
+
+    def save_data(self, data: plotly.graph_objects.Figure) -> dict[str, Any]:
+        data.write_html(file=self.path, **self._get_saving_kwargs())
+        return utils.get_file_metadata(self.path)
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [plotly.graph_objects.Figure]
+
+    @classmethod
+    def name(cls) -> str:
+        return "html"
+
+
+def register_data_loaders():
+    for loader in [
+        PlotlyStaticWriter,
+        PlotlyInteractiveWriter,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()
+
+COLUMN_FRIENDLY_DF_TYPE = False

hamilton/plugins/polars_extensions.py

@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import logging
+from typing import Any
+
+from packaging import version
+
+from hamilton import registry
+
+logger = logging.getLogger(__name__)
+
+try:
+    from xlsxwriter.workbook import Workbook
+except ImportError:
+    Workbook = type
+
+try:
+    import polars as pl
+except ImportError as e:
+    raise NotImplementedError("Polars is not installed.") from e
+
+pl_version = version.Version(pl.__version__)
+if pl_version < version.Version("1.0.0"):
+    from hamilton.plugins.polars_pre_1_0_0_extension import register_data_loaders
+
+    logger.warning(
+        "Using pre-1.0.0 Polars integration -- we will stop supporting this in Hamilton 2.0, so please upgrade your version of polars! "
+        f"Current version: {pl_version}, minimum required version: 1.0.0."
+    )
+else:
+    from hamilton.plugins.polars_post_1_0_0_extensions import register_data_loaders
+
+register_data_loaders()
+
+DATAFRAME_TYPE = pl.DataFrame
+COLUMN_TYPE = pl.Series
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("polars", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+@registry.get_column.register(pl.DataFrame)
+def get_column_polars(df: pl.DataFrame, column_name: str) -> pl.Series:
+    return df[column_name]
+
+
+@registry.fill_with_scalar.register(pl.DataFrame)
+def fill_with_scalar_polars(df: pl.DataFrame, column_name: str, scalar_value: Any) -> pl.DataFrame:
+    if not isinstance(scalar_value, pl.Series):
+        scalar_value = [scalar_value]
+    return df.with_columns(pl.Series(name=column_name, values=scalar_value))
+
+
+register_types()

hamilton/plugins/polars_implementations.py

@@ -0,0 +1,25 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import logging
+
+from hamilton.plugins.h_polars import PolarsDataFrameResult
+
+logger = logging.getLogger(__name__)
+logger.warning("This module is deprecated. Please use hamilton.plugins.h_polars instead.")
+
+__all__ = ["PolarsDataFrameResult"]

hamilton/plugins/polars_lazyframe_extensions.py

@@ -0,0 +1,302 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+from collections.abc import Collection, Mapping, Sequence
+from io import BytesIO
+from pathlib import Path
+from typing import (
+    Any,
+    BinaryIO,
+    TextIO,
+)
+
+try:
+    import polars as pl
+except ImportError as e:
+    raise NotImplementedError("Polars is not installed.") from e
+
+
+# for polars <0.16.0 we need to determine whether type_aliases exist.
+has_alias = False
+if hasattr(pl, "type_aliases"):
+    has_alias = True
+
+# for polars 0.18.0 we need to check what to do.
+if has_alias and hasattr(pl.type_aliases, "CsvEncoding"):
+    from polars.type_aliases import CsvEncoding
+else:
+    CsvEncoding = type
+
+# import these types to make type hinting work
+from polars.datatypes import DataType, DataTypeClass  # noqa: F401
+
+from hamilton import registry
+from hamilton.io import utils
+from hamilton.io.data_adapters import DataLoader
+
+DATAFRAME_TYPE = pl.LazyFrame
+COLUMN_TYPE = pl.Expr
+# COLUMN_FRIENDLY_DF_TYPE = False
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("polars_lazyframe", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+@registry.get_column.register(pl.LazyFrame)
+def get_column_polars_lazyframe(df: pl.LazyFrame, column_name: str) -> pl.Expr:
+    # TODO: figure out if we can validate this here already or need to wait to the end
+    # when query.collect() resolves the lazy frame
+    # df.collect_schema().names() gives a list of names but it can be expensive
+    # https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.columns.html
+    # https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.collect_schema.html#polars.LazyFrame.collect_schema
+    return pl.col(column_name)
+
+
+@registry.fill_with_scalar.register(pl.LazyFrame)
+def fill_with_scalar_polars_lazyframe(
+    df: pl.LazyFrame, column_name: str, scalar_value: Any
+) -> pl.LazyFrame:
+    if not isinstance(scalar_value, pl.Expr):
+        scalar_value = pl.lit(scalar_value)
+    return df.with_columns(scalar_value.alias(column_name))
+
+
+register_types()
+
+
+@dataclasses.dataclass
+class PolarsScanCSVReader(DataLoader):
+    """Class specifically to handle loading CSV files with Polars.
+    Should map to https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.read_csv.html
+    """
+
+    file: str | TextIO | BytesIO | Path | BinaryIO | bytes
+    # kwargs:
+    has_header: bool = True
+    columns: Sequence[int] | Sequence[str] = None
+    new_columns: Sequence[str] = None
+    separator: str = ","
+    comment_char: str = None
+    quote_char: str = '"'
+    skip_rows: int = 0
+    dtypes: Mapping[str, Any] | Sequence[Any] = None
+    null_values: str | Sequence[str] | dict[str, str] = None
+    missing_utf8_is_empty_string: bool = False
+    ignore_errors: bool = False
+    try_parse_dates: bool = False
+    n_threads: int = None
+    infer_schema_length: int = 100
+    batch_size: int = 8192
+    n_rows: int = None
+    encoding: CsvEncoding | str = "utf8"
+    low_memory: bool = False
+    rechunk: bool = True
+    use_pyarrow: bool = False
+    storage_options: dict[str, Any] = None
+    skip_rows_after_header: int = 0
+    row_count_name: str = None
+    row_count_offset: int = 0
+    eol_char: str = "\n"
+    raise_if_empty: bool = True
+
+    def _get_loading_kwargs(self):
+        kwargs = {}
+        if self.has_header is not None:
+            kwargs["has_header"] = self.has_header
+        if self.columns is not None:
+            kwargs["columns"] = self.columns
+        if self.new_columns is not None:
+            kwargs["new_columns"] = self.new_columns
+        if self.separator is not None:
+            kwargs["separator"] = self.separator
+        if self.comment_char is not None:
+            kwargs["comment_char"] = self.comment_char
+        if self.quote_char is not None:
+            kwargs["quote_char"] = self.quote_char
+        if self.skip_rows is not None:
+            kwargs["skip_rows"] = self.skip_rows
+        if self.dtypes is not None:
+            kwargs["dtypes"] = self.dtypes
+        if self.null_values is not None:
+            kwargs["null_values"] = self.null_values
+        if self.missing_utf8_is_empty_string is not None:
+            kwargs["missing_utf8_is_empty_string"] = self.missing_utf8_is_empty_string
+        if self.ignore_errors is not None:
+            kwargs["ignore_errors"] = self.ignore_errors
+        if self.try_parse_dates is not None:
+            kwargs["try_parse_dates"] = self.try_parse_dates
+        if self.n_threads is not None:
+            kwargs["n_threads"] = self.n_threads
+        if self.infer_schema_length is not None:
+            kwargs["infer_schema_length"] = self.infer_schema_length
+        if self.n_rows is not None:
+            kwargs["n_rows"] = self.n_rows
+        if self.encoding is not None:
+            kwargs["encoding"] = self.encoding
+        if self.low_memory is not None:
+            kwargs["low_memory"] = self.low_memory
+        if self.rechunk is not None:
+            kwargs["rechunk"] = self.rechunk
+        if self.storage_options is not None:
+            kwargs["storage_options"] = self.storage_options
+        if self.skip_rows_after_header is not None:
+            kwargs["skip_rows_after_header"] = self.skip_rows_after_header
+        if self.row_count_name is not None:
+            kwargs["row_count_name"] = self.row_count_name
+        if self.row_count_offset is not None:
+            kwargs["row_count_offset"] = self.row_count_offset
+        if self.eol_char is not None:
+            kwargs["eol_char"] = self.eol_char
+        if self.raise_if_empty is not None:
+            kwargs["raise_if_empty"] = self.raise_if_empty
+        return kwargs
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [DATAFRAME_TYPE]
+
+    def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]:
+        df = pl.scan_csv(self.file, **self._get_loading_kwargs())
+
+        metadata = utils.get_file_and_dataframe_metadata(self.file, df)
+        return df, metadata
+
+    @classmethod
+    def name(cls) -> str:
+        return "csv"
+
+
+@dataclasses.dataclass
+class PolarsScanParquetReader(DataLoader):
+    """Class specifically to handle loading parquet files with polars
+    Should map to https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.read_parquet.html
+    """
+
+    file: str | TextIO | BytesIO | Path | BinaryIO | bytes
+    # kwargs:
+    columns: list[int] | list[str] = None
+    n_rows: int = None
+    use_pyarrow: bool = False
+    memory_map: bool = True
+    storage_options: dict[str, Any] = None
+    parallel: Any = "auto"
+    row_count_name: str = None
+    row_count_offset: int = 0
+    low_memory: bool = False
+    use_statistics: bool = True
+    rechunk: bool = True
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [DATAFRAME_TYPE]
+
+    def _get_loading_kwargs(self):
+        kwargs = {}
+        if self.columns is not None:
+            kwargs["columns"] = self.columns
+        if self.n_rows is not None:
+            kwargs["n_rows"] = self.n_rows
+        if self.storage_options is not None:
+            kwargs["storage_options"] = self.storage_options
+        if self.parallel is not None:
+            kwargs["parallel"] = self.parallel
+        if self.row_count_name is not None:
+            kwargs["row_count_name"] = self.row_count_name
+        if self.row_count_offset is not None:
+            kwargs["row_count_offset"] = self.row_count_offset
+        if self.low_memory is not None:
+            kwargs["low_memory"] = self.low_memory
+        if self.use_statistics is not None:
+            kwargs["use_statistics"] = self.use_statistics
+        if self.rechunk is not None:
+            kwargs["rechunk"] = self.rechunk
+        return kwargs
+
+    def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]:
+        df = pl.scan_parquet(self.file, **self._get_loading_kwargs())
+        metadata = utils.get_file_and_dataframe_metadata(self.file, df)
+        return df, metadata
+
+    @classmethod
+    def name(cls) -> str:
+        return "parquet"
+
+
+@dataclasses.dataclass
+class PolarsScanFeatherReader(DataLoader):
+    """
+    Class specifically to handle loading Feather/Arrow IPC files with Polars.
+    Should map to https://pola-rs.github.io/polars/py-polars/html/reference/api/polars.read_ipc.html
+    """
+
+    source: str | BinaryIO | BytesIO | Path | bytes
+    # kwargs:
+    columns: list[str] | list[int] | None = None
+    n_rows: int | None = None
+    use_pyarrow: bool = False
+    memory_map: bool = True
+    storage_options: dict[str, Any] | None = None
+    row_count_name: str | None = None
+    row_count_offset: int = 0
+    rechunk: bool = True
+
+    @classmethod
+    def applicable_types(cls) -> Collection[type]:
+        return [DATAFRAME_TYPE]
+
+    def _get_loading_kwargs(self):
+        kwargs = {}
+        if self.columns is not None:
+            kwargs["columns"] = self.columns
+        if self.n_rows is not None:
+            kwargs["n_rows"] = self.n_rows
+        if self.memory_map is not None:
+            kwargs["memory_map"] = self.memory_map
+        if self.storage_options is not None:
+            kwargs["storage_options"] = self.storage_options
+        if self.row_count_name is not None:
+            kwargs["row_count_name"] = self.row_count_name
+        if self.row_count_offset is not None:
+            kwargs["row_count_offset"] = self.row_count_offset
+        if self.rechunk is not None:
+            kwargs["rechunk"] = self.rechunk
+        return kwargs
+
+    def load_data(self, type_: type) -> tuple[DATAFRAME_TYPE, dict[str, Any]]:
+        df = pl.scan_ipc(self.source, **self._get_loading_kwargs())
+        metadata = utils.get_file_metadata(self.source)
+        return df, metadata
+
+    @classmethod
+    def name(cls) -> str:
+        return "feather"
+
+
+def register_data_loaders():
+    """Function to register the data loaders for this extension."""
+    for loader in [
+        PolarsScanCSVReader,
+        PolarsScanParquetReader,
+        PolarsScanFeatherReader,
+    ]:
+        registry.register_adapter(loader)
+
+
+register_data_loaders()