From 78f1ce9acc1e956106540a492c8d442485cd7f4c Mon Sep 17 00:00:00 2001
From: Neha Nene <nehanene@google.com>
Date: Thu, 10 Nov 2022 11:34:52 -0600
Subject: [PATCH] feat: to_char support for Oracle and Postgres (#632)

* feat: to_char support for oracle and postgres

* support user defined to_char fmt

* lint

Co-authored-by: Sriram Srinivasan <sramsrinivasan@google.com>
---
 .../query_builder/query_builder.py            | 15 +++++++++-----
 third_party/ibis/ibis_addon/operations.py     | 20 +++++++++++++++++--
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/data_validation/query_builder/query_builder.py b/data_validation/query_builder/query_builder.py
index 05a4e2f65..894b72fa1 100644
--- a/data_validation/query_builder/query_builder.py
+++ b/data_validation/query_builder/query_builder.py
@@ -306,14 +306,19 @@ def hash(config, fields):
                 how=how,
             )
 
+    @staticmethod
+    def to_char(config, fields):
+        fmt = ibis.literal(config.get("default_to_char_fmt", "FM90.099"))
+        return CalculatedField(
+            ibis.expr.api.NumericValue.to_char, config, fields, fmt=fmt
+        )
+
     @staticmethod
     def ifnull(config, fields):
-        config["default_string"] = (
-            ibis.literal("DEFAULT_REPLACEMENT_STRING")
-            if config.get("default_null_string") is None
-            else ibis.literal(config.get("default_null_string"))
+        default_null_string = ibis.literal(
+            config.get("default_null_string", "DEFAULT_REPLACEMENT_STRING")
         )
-        fields = [fields[0], config["default_string"]]
+        fields = [fields[0], default_null_string]
         return CalculatedField(
             ibis.expr.api.ValueExpr.fillna,
             config,
diff --git a/third_party/ibis/ibis_addon/operations.py b/third_party/ibis/ibis_addon/operations.py
index a7085fcfb..971ccf340 100644
--- a/third_party/ibis/ibis_addon/operations.py
+++ b/third_party/ibis/ibis_addon/operations.py
@@ -33,7 +33,7 @@
 from data_validation.clients import _raise_missing_client_error
 from ibis_bigquery.compiler import reduction as bq_reduction, BigQueryExprTranslator
 from ibis.expr.operations import Arg, Comparison, Reduction, ValueOp
-from ibis.expr.types import BinaryValue, IntegerColumn, StringValue
+from ibis.expr.types import BinaryValue, IntegerColumn, StringValue, NumericValue
 from ibis.backends.impala.compiler import ImpalaExprTranslator
 from ibis.backends.pandas import client as _pandas_client
 from ibis.backends.base_sqlalchemy.alchemy import AlchemyExprTranslator
@@ -66,6 +66,10 @@ class HashBytes(ValueOp):
     how = Arg(rlz.isin({"sha256", "farm_fingerprint"}))
     output_type = rlz.shape_like("arg", "binary")
 
+class ToChar(ValueOp):
+    arg = Arg(rlz.one_of([rlz.value(dt.Decimal), rlz.value(dt.float64)]))
+    fmt = Arg(rlz.string)
+    output_type = rlz.shape_like("arg", dt.string)
 
 class RawSQL(Comparison):
     pass
@@ -95,6 +99,10 @@ def compile_hashbytes(binary_value, how):
     return HashBytes(binary_value, how=how).to_expr()
 
 
+def compile_to_char(numeric_value, fmt):
+    return ToChar(numeric_value, fmt=fmt).to_expr()
+    
+
 def format_hash_bigquery(translator, expr):
     arg, how = expr.op().args
     compiled_arg = translator.translate(arg)
@@ -188,6 +196,11 @@ def sa_format_hashbytes_postgres(translator, expr):
     hash_func = sa.func.sha256(convert)
     return sa.func.encode(hash_func, sa.sql.literal_column("'hex'"))
 
+def sa_format_to_char(translator, expr):
+    arg, fmt = expr.op().args
+    compiled_arg = translator.translate(arg)
+    compiled_fmt = translator.translate(fmt)
+    return sa.func.to_char(compiled_arg, compiled_fmt)
 
 _pandas_client._inferable_pandas_dtypes["floating"] = _pandas_client.dt.float64
 IntegerColumn.bit_xor = ibis.expr.api._agg_function("bit_xor", BitXor, True)
@@ -195,6 +208,7 @@ def sa_format_hashbytes_postgres(translator, expr):
 StringValue.hash = compile_hash
 BinaryValue.hashbytes = compile_hashbytes
 StringValue.hashbytes = compile_hashbytes
+NumericValue.to_char = compile_to_char
 BigQueryExprTranslator._registry[BitXor] = bq_reduction("BIT_XOR")
 BigQueryExprTranslator._registry[Hash] = format_hash_bigquery
 BigQueryExprTranslator._registry[HashBytes] = format_hashbytes_bigquery
@@ -209,7 +223,9 @@ def sa_format_hashbytes_postgres(translator, expr):
 ImpalaExprTranslator._registry[HashBytes] = format_hashbytes_hive
 OracleExprTranslator._registry[RawSQL] = sa_format_raw_sql
 OracleExprTranslator._registry[HashBytes] = sa_format_hashbytes_oracle
+OracleExprTranslator._registry[ToChar] = sa_format_to_char
 TeradataExprTranslator._registry[RawSQL] = format_raw_sql
 TeradataExprTranslator._registry[HashBytes] = format_hashbytes_teradata
 PostgreSQLExprTranslator._registry[HashBytes] = sa_format_hashbytes_postgres
-PostgreSQLExprTranslator._registry[RawSQL] = sa_format_raw_sql
\ No newline at end of file
+PostgreSQLExprTranslator._registry[RawSQL] = sa_format_raw_sql
+PostgreSQLExprTranslator._registry[ToChar] = sa_format_to_char
\ No newline at end of file