import datetime import unittest from sqlglot.dataframe.sql import functions as F from sqlglot.dataframe.sql.window import Window class TestDataframeColumn(unittest.TestCase): def test_eq(self): self.assertEqual("cola = 1", (F.col("cola") == 1).sql()) def test_neq(self): self.assertEqual("cola <> 1", (F.col("cola") != 1).sql()) def test_gt(self): self.assertEqual("cola > 1", (F.col("cola") > 1).sql()) def test_lt(self): self.assertEqual("cola < 1", (F.col("cola") < 1).sql()) def test_le(self): self.assertEqual("cola <= 1", (F.col("cola") <= 1).sql()) def test_ge(self): self.assertEqual("cola >= 1", (F.col("cola") >= 1).sql()) def test_and(self): self.assertEqual( "cola = colb AND colc = cold", ((F.col("cola") == F.col("colb")) & (F.col("colc") == F.col("cold"))).sql(), ) def test_or(self): self.assertEqual( "cola = colb OR colc = cold", ((F.col("cola") == F.col("colb")) | (F.col("colc") == F.col("cold"))).sql(), ) def test_mod(self): self.assertEqual("cola % 2", (F.col("cola") % 2).sql()) def test_add(self): self.assertEqual("cola + 1", (F.col("cola") + 1).sql()) def test_sub(self): self.assertEqual("cola - 1", (F.col("cola") - 1).sql()) def test_mul(self): self.assertEqual("cola * 2", (F.col("cola") * 2).sql()) def test_div(self): self.assertEqual("cola / 2", (F.col("cola") / 2).sql()) def test_radd(self): self.assertEqual("1 + cola", (1 + F.col("cola")).sql()) def test_rsub(self): self.assertEqual("1 - cola", (1 - F.col("cola")).sql()) def test_rmul(self): self.assertEqual("1 * cola", (1 * F.col("cola")).sql()) def test_rdiv(self): self.assertEqual("1 / cola", (1 / F.col("cola")).sql()) def test_pow(self): self.assertEqual("POWER(cola, 2)", (F.col("cola") ** 2).sql()) def test_rpow(self): self.assertEqual("POWER(2, cola)", (2 ** F.col("cola")).sql()) def test_invert(self): self.assertEqual("NOT cola", (~F.col("cola")).sql()) def test_startswith(self): self.assertEqual("STARTSWITH(cola, 'test')", F.col("cola").startswith("test").sql()) def test_endswith(self): self.assertEqual("ENDSWITH(cola, 'test')", F.col("cola").endswith("test").sql()) def test_rlike(self): self.assertEqual("cola RLIKE 'foo'", F.col("cola").rlike("foo").sql()) def test_like(self): self.assertEqual("cola LIKE 'foo%'", F.col("cola").like("foo%").sql()) def test_ilike(self): self.assertEqual("cola ILIKE 'foo%'", F.col("cola").ilike("foo%").sql()) def test_substring(self): self.assertEqual("SUBSTRING(cola, 2, 3)", F.col("cola").substr(2, 3).sql()) def test_isin(self): self.assertEqual("cola IN (1, 2, 3)", F.col("cola").isin([1, 2, 3]).sql()) self.assertEqual("cola IN (1, 2, 3)", F.col("cola").isin(1, 2, 3).sql()) def test_asc(self): self.assertEqual("cola ASC", F.col("cola").asc().sql()) def test_desc(self): self.assertEqual("cola DESC", F.col("cola").desc().sql()) def test_asc_nulls_first(self): self.assertEqual("cola ASC", F.col("cola").asc_nulls_first().sql()) def test_asc_nulls_last(self): self.assertEqual("cola ASC NULLS LAST", F.col("cola").asc_nulls_last().sql()) def test_desc_nulls_first(self): self.assertEqual("cola DESC NULLS FIRST", F.col("cola").desc_nulls_first().sql()) def test_desc_nulls_last(self): self.assertEqual("cola DESC", F.col("cola").desc_nulls_last().sql()) def test_when_otherwise(self): self.assertEqual("CASE WHEN cola = 1 THEN 2 END", F.when(F.col("cola") == 1, 2).sql()) self.assertEqual( "CASE WHEN cola = 1 THEN 2 END", F.col("cola").when(F.col("cola") == 1, 2).sql() ) self.assertEqual( "CASE WHEN cola = 1 THEN 2 WHEN colb = 2 THEN 3 END", (F.when(F.col("cola") == 1, 2).when(F.col("colb") == 2, 3)).sql(), ) self.assertEqual( "CASE WHEN cola = 1 THEN 2 WHEN colb = 2 THEN 3 END", F.col("cola").when(F.col("cola") == 1, 2).when(F.col("colb") == 2, 3).sql(), ) self.assertEqual( "CASE WHEN cola = 1 THEN 2 WHEN colb = 2 THEN 3 ELSE 4 END", F.when(F.col("cola") == 1, 2).when(F.col("colb") == 2, 3).otherwise(4).sql(), ) def test_is_null(self): self.assertEqual("cola IS NULL", F.col("cola").isNull().sql()) def test_is_not_null(self): self.assertEqual("NOT cola IS NULL", F.col("cola").isNotNull().sql()) def test_cast(self): self.assertEqual("CAST(cola AS INT)", F.col("cola").cast("INT").sql()) def test_alias(self): self.assertEqual("cola AS new_name", F.col("cola").alias("new_name").sql()) def test_between(self): self.assertEqual("cola BETWEEN 1 AND 3", F.col("cola").between(1, 3).sql()) self.assertEqual("cola BETWEEN 10.1 AND 12.1", F.col("cola").between(10.1, 12.1).sql()) self.assertEqual( "cola BETWEEN CAST('2022-01-01' AS DATE) AND CAST('2022-03-01' AS DATE)", F.col("cola").between(datetime.date(2022, 1, 1), datetime.date(2022, 3, 1)).sql(), ) self.assertEqual( "cola BETWEEN CAST('2022-01-01 01:01:01+00:00' AS TIMESTAMP) " "AND CAST('2022-03-01 01:01:01+00:00' AS TIMESTAMP)", F.col("cola") .between(datetime.datetime(2022, 1, 1, 1, 1, 1), datetime.datetime(2022, 3, 1, 1, 1, 1)) .sql(), ) def test_over(self): over_rows = F.sum("cola").over( Window.partitionBy("colb").orderBy("colc").rowsBetween(1, Window.unboundedFollowing) ) self.assertEqual( "SUM(cola) OVER (PARTITION BY colb ORDER BY colc ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING)", over_rows.sql(), ) over_range = F.sum("cola").over( Window.partitionBy("colb").orderBy("colc").rangeBetween(1, Window.unboundedFollowing) ) self.assertEqual( "SUM(cola) OVER (PARTITION BY colb ORDER BY colc RANGE BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING)", over_range.sql(), )