continue

Author: ruslandoga

lib/ch/row_binary.ex

@@ -129,6 +129,10 @@ defmodule Ch.RowBinary do
     {t, Enum.map(ts, &encoding_type/1)}
   end
 
+  defp encoding_type({:variant = v, ts}) do
+    {v, Enum.map(ts, &encoding_type/1)}
+  end
+
   defp encoding_type({:map = m, kt, vt}) do
     {m, encoding_type(kt), encoding_type(vt)}
   end
@@ -302,6 +306,12 @@ defmodule Ch.RowBinary do
     Enum.map(types, fn type -> encode(type, nil) end)
   end
 
+  def encode({:variant, _types}, nil), do: 255
+
+  def encode({:variant, types}, value) do
+    try_encode_variant(types, 0, value)
+  end
+
   def encode(:datetime, %NaiveDateTime{} = datetime) do
     <<NaiveDateTime.diff(datetime, @epoch_naive_datetime)::32-little>>
   end
@@ -453,6 +463,20 @@ defmodule Ch.RowBinary do
 
   defp encode_many_kv([] = done, _key_type, _value_type), do: done
 
+  defp try_encode_variant([type | types], idx, value) do
+    try do
+      encode(type, value)
+    else
+      encoded -> [idx | encoded]
+    rescue
+      _e -> try_encode_variant(types, idx + 1, value)
+    end
+  end
+
+  defp try_encode_variant([], _idx, value) do
+    raise ArgumentError, "no matching type found for encoding #{inspect(value)} as Variant"
+  end
+
   @compile {:inline, d: 1}
 
   defp d(?0), do: 0
@@ -606,8 +630,8 @@ defmodule Ch.RowBinary do
     {t, Enum.map(ts, &decoding_type/1)}
   end
 
-  defp decoding_type({:variant = t, ts}) do
-    {t, Enum.map(ts, &decoding_type/1)}
+  defp decoding_type({:variant = v, ts}) do
+    {v, Enum.map(ts, &decoding_type/1)}
   end
 
   defp decoding_type({:map = m, kt, vt}) do

test/ch/variant_test.exs

@@ -1,19 +1,83 @@
 defmodule Ch.VariantTest do
   use ExUnit.Case
 
+  # https://clickhouse.com/docs/sql-reference/data-types/variant
+
   @moduletag :variant
 
   setup do
     conn = start_supervised!({Ch, database: Ch.Test.database()})
     {:ok, conn: conn}
   end
 
-  test "it works", %{conn: conn} do
+  test "basic", %{conn: conn} do
     assert Ch.query!(conn, "select null::Variant(UInt64, String, Array(UInt64))").rows == [[nil]]
     assert Ch.query!(conn, "select [1]::Variant(UInt64, String, Array(UInt64))").rows == [[[1]]]
     assert Ch.query!(conn, "select 0::Variant(UInt64, String, Array(UInt64))").rows == [[0]]
 
     assert Ch.query!(conn, "select 'Hello, World!'::Variant(UInt64, String, Array(UInt64))").rows ==
              [["Hello, World!"]]
   end
+
+  test "with a table", %{conn: conn} do
+    # https://clickhouse.com/docs/sql-reference/data-types/variant#creating-variant
+    Ch.query!(conn, """
+    CREATE TABLE variant_test (v Variant(UInt64, String, Array(UInt64))) ENGINE = Memory;
+    """)
+
+    on_exit(fn -> Ch.Test.query("DROP TABLE variant_test", [], database: Ch.Test.database()) end)
+
+    Ch.query!(
+      conn,
+      "INSERT INTO variant_test VALUES (NULL), (42), ('Hello, World!'), ([1, 2, 3]);"
+    )
+
+    assert Ch.query!(conn, "SELECT v FROM variant_test").rows == [
+             [nil],
+             [42],
+             ["Hello, World!"],
+             [[1, 2, 3]]
+           ]
+
+    # https://clickhouse.com/docs/sql-reference/data-types/variant#reading-variant-nested-types-as-subcolumns
+    assert Ch.query!(conn, "SELECT v, v.String, v.UInt64, v.`Array(UInt64)` FROM variant_test;").rows ==
+             [
+               [nil, nil, nil, []],
+               [42, nil, 42, []],
+               ["Hello, World!", "Hello, World!", nil, []],
+               [[1, 2, 3], nil, nil, [1, 2, 3]]
+             ]
+
+    assert Ch.query!(
+             conn,
+             "SELECT v, variantElement(v, 'String'), variantElement(v, 'UInt64'), variantElement(v, 'Array(UInt64)') FROM variant_test;"
+           ).rows == [
+             [nil, nil, nil, []],
+             [42, nil, 42, []],
+             ["Hello, World!", "Hello, World!", nil, []],
+             [[1, 2, 3], nil, nil, [1, 2, 3]]
+           ]
+  end
+
+  test "rowbinary", %{conn: conn} do
+    Ch.query!(conn, """
+    CREATE TABLE variant_test (v Variant(UInt64, String, Array(UInt64))) ENGINE = Memory;
+    """)
+
+    on_exit(fn -> Ch.Test.query("DROP TABLE variant_test", [], database: Ch.Test.database()) end)
+
+    Ch.query!(
+      conn,
+      "INSERT INTO variant_test FORMAT RowBinary",
+      [[nil], [42], ["Hello, World!"], [[1, 2, 3]]],
+      types: ["Variant(UInt64, String, Array(UInt64))"]
+    )
+
+    assert Ch.query!(conn, "SELECT v FROM variant_test").rows == [
+             [nil],
+             [42],
+             ["Hello, World!"],
+             [[1, 2, 3]]
+           ]
+  end
 end