Adds Support for Join Operations

pydruid/query.py

@@ -21,6 +21,7 @@
 from pydruid.utils.dimensions import build_dimension
 from pydruid.utils.filters import Filter
 from pydruid.utils.having import Having
+from pydruid.utils.joins import Join
 from pydruid.utils.postaggregator import Postaggregator
 from pydruid.utils.query_utils import UnicodeWriter
 
@@ -241,13 +242,16 @@ def parse_datasource(datasource, query_type):
                 and all([isinstance(x, str) for x in datasource])
             )
             or isinstance(datasource, dict)
+            or isinstance(datasource, Join)
         ):
             raise ValueError(
                 "Datasource definition not valid. Must be string or "
-                "dict or list of strings"
+                "dict or list of strings or Join"
             )
         if isinstance(datasource, str):
             return datasource
+        elif isinstance(datasource, Join):
+            return Join.build_join(datasource)
         else:
             return {"type": "union", "dataSources": datasource}
 

pydruid/utils/joins.py

@@ -0,0 +1,69 @@
+#
+# Copyright 2013 Metamarkets Group Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+LEFT = "LEFT"
+INNER = "INNER"
+
+
+class Join:
+    """ A base class for building a join operation """
+
+    def __init__(self, left, right, right_prefix, condition, join_type=INNER):
+        """ A class for building a right join """
+        self.left = left
+        self.right = right
+        self.right_prefix = right_prefix
+        self.condition = condition
+        self.join_type = join_type
+
+    @property
+    def join(self):
+        """ Gets the join as a dictionary """
+        return {
+            "left": self.left,
+            "right": self.right,
+            "rightPrefix": self.right_prefix,
+            "condition": self.condition,
+            "joinType": self.join_type,
+            "type": "join",
+        }
+
+    @staticmethod
+    def build_join(join):
+        """ Builds the join dictionary """
+        return join.join
+
+
+class LeftJoin(Join):
+    """ A class for building a left join operation """
+
+    def __init__(self, left, right, right_prefix, condition):
+        super(LeftJoin, self).__init__(left, right, right_prefix, condition, LEFT)
+
+
+class InnerJoin(Join):
+    """ A class for building a left join operation """
+
+    def __init__(self, left, right, right_prefix, condition):
+        super(InnerJoin, self).__init__(left, right, right_prefix, condition, INNER)
+
+
+class CrossJoin(Join):
+    """ A class for building a cross join operation """
+
+    def __init__(self, left, right, right_prefix):
+        super(CrossJoin, self).__init__(left, right, right_prefix, "1 = 1", INNER)

tests/test_query.py

@@ -16,14 +16,13 @@
 #
 
 import csv
-import os
 
 import pandas
 import pytest
 from pandas.testing import assert_frame_equal
 
 from pydruid.query import Query, QueryBuilder
-from pydruid.utils import aggregators, filters, having, postaggregator
+from pydruid.utils import aggregators, filters, having, postaggregator, joins
 
 
 def create_query_with_results():
@@ -226,6 +225,52 @@ def test_build_subquery(self):
         # then
         assert subquery_dict == expected_query_dict
 
+    def test_build_join_query(self):
+        """ Tests building a query with a join """
+        expected_query_dict = {
+            "queryType": "groupBy",
+            "dataSource": {
+                "type": "join",
+                "joinType": "INNER",
+                "left": {
+                    "type": "query",
+                    "query": {
+                        "dataSource": "things",
+                        "dimensions": ["id", "a"],
+                        "queryType": "groupBy",
+                    },
+                },
+                "right": {
+                    "type": "query",
+                    "query": {
+                        "dataSource": "other_things",
+                        "dimensions": ["id", "b"],
+                        "queryType": "groupBy",
+                    },
+                },
+                "rightPrefix": "other",
+                "condition": "id = other_id",
+            },
+            "dimensions": ["id"],
+        }
+
+        builder = QueryBuilder()
+
+        left = builder.subquery({"datasource": "things", "dimensions": ["id", "a"]})
+
+        right = builder.subquery(
+            {"datasource": "other_things", "dimensions": ["id", "b"]}
+        )
+
+        joined = joins.InnerJoin(
+            left=left, right=right, right_prefix="other", condition="id = other_id"
+        )
+
+        query = builder.groupby({"datasource": joined, "dimensions": ["id"]})
+
+        # then
+        assert query.query_dict == expected_query_dict
+
 
 class TestQuery:
     def test_export_tsv(self, tmpdir):

tests/utils/test_joins.py

@@ -0,0 +1,65 @@
+# -*- coding: UTF-8 -*-
+"""
+Tests the pydruid.utils.joins module
+"""
+
+from pydruid.utils import joins
+
+
+class TestJoin:
+    """ Test cases the pydruid.util.join module """
+
+    def test_join(self):
+        join = joins.Join("some", "other", "other_", "a = other_b", join_type="INNER")
+        actual = joins.Join.build_join(join)
+        expected = {
+            "type": "join",
+            "joinType": "INNER",
+            "condition": "a = other_b",
+            "left": "some",
+            "right": "other",
+            "rightPrefix": "other_",
+        }
+        assert actual == expected
+
+    def test_inner_join(self):
+        """ Tests InnerJoin """
+        join = joins.InnerJoin("some", "other", "other_", "a = other_b")
+        actual = joins.Join.build_join(join)
+        expected = {
+            "type": "join",
+            "joinType": "INNER",
+            "condition": "a = other_b",
+            "left": "some",
+            "right": "other",
+            "rightPrefix": "other_",
+        }
+        assert actual == expected
+
+    def test_left_join(self):
+        """ Tests LeftJoin """
+        join = joins.LeftJoin("some", "other", "other_", "a = other_b")
+        actual = joins.Join.build_join(join)
+        expected = {
+            "type": "join",
+            "joinType": "LEFT",
+            "condition": "a = other_b",
+            "left": "some",
+            "right": "other",
+            "rightPrefix": "other_",
+        }
+        assert actual == expected
+
+    def test_cross_join(self):
+        """ Tests CrossJoin """
+        join = joins.CrossJoin("some", "other", "other_")
+        actual = joins.Join.build_join(join)
+        expected = {
+            "type": "join",
+            "joinType": "INNER",  # Converts to inner join with condition 1 = 1
+            "condition": "1 = 1",
+            "left": "some",
+            "right": "other",
+            "rightPrefix": "other_",
+        }
+        assert actual == expected