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

hamilton/plugins/ibis_extensions.py

@@ -0,0 +1,93 @@
+# 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.
+
+from typing import Any
+
+from hamilton import registry
+
+try:
+    import ibis
+    import ibis.expr.types as ir
+except ImportError as e:
+    raise NotImplementedError("Ibis is not installed.") from e
+
+from hamilton.data_quality import base, default_validators
+
+DATAFRAME_TYPE = ir.Table
+COLUMN_TYPE = ir.Column
+
+
+@registry.get_column.register(ir.Table)
+def get_column_ibis(df: ir.Table, column_name: str) -> ir.Column:
+    return df[column_name]
+
+
+def register_types():
+    """Function to register the types for this extension."""
+    registry.register_types("ibis", DATAFRAME_TYPE, COLUMN_TYPE)
+
+
+@registry.fill_with_scalar.register(ir.Table)
+def fill_with_scalar_ibis(df: ir.Table, column_name: str, scalar_value: Any) -> ir.Table:
+    raise NotImplementedError(
+        "As of Hamilton version 1.5.1, filling Ibis columns with scalars isn't supported. This feature is actively being developed"
+    )
+
+
+register_types()
+
+
+class SchemaValidatorIbis(base.DataValidator):
+    def __init__(self, schema: ibis.expr.schema.Schema, importance: str):
+        """`schema` is an ordered mapping"""
+        super(SchemaValidatorIbis, self).__init__(importance)
+        self.schema = schema
+
+    def name(self) -> str:
+        return f"{self.arg()}_validator"
+
+    @classmethod
+    def arg(cls) -> str:
+        return "schema"
+
+    @classmethod
+    def applies_to(cls, datatype: type[type]) -> bool:
+        return issubclass(datatype, ir.Table)
+
+    def description(self) -> str:
+        return "Validates that ibis Table expression matches schema."
+
+    def validate(self, data: ir.Table) -> base.ValidationResult:
+        passes = data.schema().equals(self.schema)
+        if passes:
+            message = "Data passes Ibis schema check"
+            diagnostics = {"schema": self.schema.fields}
+        else:
+            message = "Data failed Ibis schema check"
+            diagnostics = {"schema": self.schema.fields}
+        return base.ValidationResult(
+            passes=passes,
+            message=message,
+            diagnostics=diagnostics,
+        )
+
+
+def register_validators():
+    default_validators.AVAILABLE_DEFAULT_VALIDATORS.append(SchemaValidatorIbis)
+
+
+register_validators()