feat(clickhouse): support asof_join

ibis-project · Feb 13, 2023 · 7ed5143 · 7ed5143
1 parent aa3c31c
commit 7ed5143
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 0 deletions.
diff --git a/ci/schema/clickhouse.sql b/ci/schema/clickhouse.sql
@@ -52,6 +52,29 @@ INSERT INTO array_types VALUES
     ([2, NULL, 3], ['b', NULL, 'c'], NULL, 'b', 5.0, []),
     ([4, NULL, NULL, 5], ['d', NULL, NULL, 'e'], [4.0, NULL, NULL, 5.0], 'c', 6.0, [[1, 2, 3]]);
 
+CREATE OR REPLACE TABLE time_df1 (
+    time Int64,
+    value Nullable(Float64),
+    key Nullable(String)
+) ENGINE = Memory;
+INSERT INTO time_df1 VALUES
+    (1, 1.0, 'x'),
+    (20, 20.0, 'x'),
+    (30, 30.0, 'x'),
+    (40, 40.0, 'x'),
+    (50, 50.0, 'x');
+
+CREATE OR REPLACE TABLE time_df2 (
+    time Int64,
+    value Nullable(Float64),
+    key Nullable(String)
+) ENGINE = Memory;
+INSERT INTO time_df2 VALUES
+    (19, 19.0, 'x'),
+    (21, 21.0, 'x'),
+    (39, 39.0, 'x'),
+    (49, 49.0, 'x'),
+    (1000, 1000.0, 'x');
 
 CREATE OR REPLACE TABLE struct (
     abc Tuple(

diff --git a/ibis/backends/clickhouse/compiler/relations.py b/ibis/backends/clickhouse/compiler/relations.py
@@ -103,6 +103,7 @@ def _aggregation(op: ops.Aggregation, *, table, **kw):
     ops.CrossJoin: "CROSS",
     ops.LeftSemiJoin: "LEFT SEMI",
     ops.LeftAntiJoin: "LEFT ANTI",
+    ops.AsOfJoin: "LEFT ASOF",
 }
 
 

diff --git a/ibis/backends/clickhouse/tests/test_select.py b/ibis/backends/clickhouse/tests/test_select.py
@@ -23,6 +23,16 @@ def awards_players(con):
     return con.table('awards_players')
 
 
+@pytest.fixture(scope='module')
+def time_left(con):
+    return con.table('time_df1')
+
+
+@pytest.fixture(scope='module')
+def time_right(con):
+    return con.table('time_df2')
+
+
 def test_timestamp_extract_field(alltypes, snapshot):
     t = alltypes.timestamp_col
     expr = alltypes[
@@ -354,3 +364,18 @@ def test_join_with_external_table(alltypes, df):
     expected = expected.sort_values('id').reset_index(drop=True)
 
     tm.assert_frame_equal(result, expected, check_column_type=False)
+
+
+def test_asof_join(time_left, time_right):
+    expr = time_left.asof_join(
+        time_right,
+        predicates=[
+            time_left['key'] == time_right['key'],
+            time_left['time'] >= time_right['time'],
+        ],
+    )
+    result = expr.execute()
+    result['time'] = result['time_x']
+    result.drop(['time_x', 'time_y'], axis=1, inplace=True)
+    expected = pd.merge_asof(time_left.execute(), time_right.execute(), on='time')
+    tm.assert_frame_equal(result[expected.columns], expected)