From 33fc15627ce35825b15c2270c3373720a0c229c5 Mon Sep 17 00:00:00 2001
From: Tanner Clary <tannerclary@google.com>
Date: Mon, 4 Mar 2024 10:01:36 -0800
Subject: [PATCH] [CALCITE-6290] Incorrect return type for BigQuery TRUNC

---
 .../calcite/sql/fun/SqlLibraryOperators.java       |  8 ++++++--
 .../calcite/rel/rel2sql/RelToSqlConverterTest.java | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java b/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
index 86a72d2652f..fc5cfe90fe0 100644
--- a/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
+++ b/core/src/main/java/org/apache/calcite/sql/fun/SqlLibraryOperators.java
@@ -456,7 +456,7 @@ static RelDataType deriveTypeSplit(SqlOperatorBinding operatorBinding,
       .withName("CEIL_BIG_QUERY")
       .withReturnTypeInference(ReturnTypes.ARG0_EXCEPT_INTEGER_NULLABLE);
 
-  /** The "FLOOR(value)" function. Identical to the stadnard <code>FLOOR</code> function
+  /** The "FLOOR(value)" function. Identical to the standard <code>FLOOR</code> function
    * except the return type should be a double if the operand is an integer. */
   @LibraryOperator(libraries = {BIG_QUERY})
   public static final SqlFunction FLOOR_BIG_QUERY = new SqlFloorFunction(SqlKind.FLOOR)
@@ -2161,9 +2161,13 @@ private static RelDataType deriveTypeMapFromEntries(SqlOperatorBinding opBinding
   public static final SqlFunction POW =
       SqlStdOperatorTable.POWER.withName("POW");
 
+
+  /** The "TRUNC(numeric1 [, integer2])" function. Identical to the standard <code>TRUNCATE</code>
+  * function except the return type should be a double if numeric1 is an integer. */
   @LibraryOperator(libraries = {BIG_QUERY})
   public static final SqlFunction TRUNC =
-      SqlStdOperatorTable.TRUNCATE.withName("TRUNC");
+      SqlStdOperatorTable.TRUNCATE.withName("TRUNC")
+        .withReturnTypeInference(ReturnTypes.ARG0_EXCEPT_INTEGER_NULLABLE);
 
   /** Infix "::" cast operator used by PostgreSQL, for example
    * {@code '100'::INTEGER}. */
diff --git a/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java b/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java
index 1de16cb69ee..312815ded6a 100644
--- a/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java
+++ b/core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java
@@ -592,6 +592,20 @@ private static String toSql(RelNode root, SqlDialect dialect,
         + "created_thing\nFROM foodmart.product");
   }
 
+    /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6290">[CALCITE-6290]
+   * Incorrect return type for BigQuery TRUNC</a>. */
+  @Test void testBigQueryTruncPreservesCast() {
+    final String query = "SELECT CAST(TRUNC(3) AS BIGINT) "
+        + "as created_thing\n FROM `foodmart`.`product`";
+    final SqlParser.Config parserConfig =
+        BigQuerySqlDialect.DEFAULT.configureParser(SqlParser.config());
+    final Sql sql = fixture()
+        .withBigQuery().withLibrary(SqlLibrary.BIG_QUERY).parserConfig(parserConfig);
+    sql.withSql(query).ok("SELECT CAST(TRUNC(3) AS INT64) AS "
+        + "created_thing\nFROM foodmart.product");
+  }
+
   @Test void testSelectLiteralAgg() {
     final Function<RelBuilder, RelNode> relFn = b -> b
         .scan("EMP")