From db4bef41b1a652d275fb5730c6b04a0a18e72cf3 Mon Sep 17 00:00:00 2001
From: EJ Song <51077614+sezruby@users.noreply.github.com>
Date: Tue, 4 Jan 2022 14:49:32 -0800
Subject: [PATCH] Add python binding for ZOrderCoveringIndexConfig (#519)

---
 python/hyperspace/__init__.py                 |  4 +--
 python/hyperspace/hyperspace.py               | 30 ++++++++++++++++---
 python/hyperspace/indexconfig.py              | 15 ++++++++++
 .../hyperspace/tests/test_indexmanagement.py  |  7 +++++
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/python/hyperspace/__init__.py b/python/hyperspace/__init__.py
index cb9599cd..f0d8708f 100644
--- a/python/hyperspace/__init__.py
+++ b/python/hyperspace/__init__.py
@@ -1,6 +1,6 @@
 from .hyperspace import Hyperspace
-from .indexconfig import CoveringIndexConfig, IndexConfig
+from .indexconfig import CoveringIndexConfig, ZOrderCoveringIndexConfig, IndexConfig
 
 __all__ = [
-    'Hyperspace', 'CoveringIndexConfig', 'IndexConfig'
+    'Hyperspace', 'CoveringIndexConfig', 'ZOrderCoveringIndexConfig', 'IndexConfig'
 ]
diff --git a/python/hyperspace/hyperspace.py b/python/hyperspace/hyperspace.py
index e5c1173e..13c2e931 100644
--- a/python/hyperspace/hyperspace.py
+++ b/python/hyperspace/hyperspace.py
@@ -21,9 +21,9 @@ class Hyperspace:
 
     def _getJavaCoveringIndexConfig(self, index_config):
         """
-        Constructs CoveringIndexConfig Java object from python wrapper CoveringIndexConfig object.
-        :param index_config: CoveringIndexConfig java object
-        :return: CoveringIndexConfig python object
+        Constructs CoveringIndexConfig Java object from python wrapper object.
+        :param index_config: CoveringIndexConfig python object
+        :return: CoveringIndexConfig java object
 
         >>> _getJavaCoveringIndexConfig(idx_config)
         """
@@ -33,6 +33,21 @@ class Hyperspace:
             self.jvm.java.lang.String(index_config.indexName), indexed_columns, included_columns)
         return _jindexConfig
 
+    def _getJavaZOrderCoveringIndexConfig(self, index_config):
+        """
+        Constructs ZOrderCoveringIndexConfig Java object from python wrapper object.
+        :param index_config: ZOrderCoveringIndexConfig python object
+        :return: ZOrderCoveringIndexConfig java object
+
+        >>> _getJavaZOrderCoveringIndexConfig(idx_config)
+        """
+        indexed_columns = self._getScalaSeqFromList(index_config.indexedColumns)
+        included_columns = self._getScalaSeqFromList(index_config.includedColumns)
+        index_name = self.jvm.java.lang.String(index_config.indexName)
+        _jZOrderIndexConfig = self.jvm.com.microsoft.hyperspace.index.zordercovering. \
+            ZOrderCoveringIndexConfig(index_name, indexed_columns, included_columns)
+        return _jZOrderIndexConfig
+
     def _getScalaSeqFromList(self, list):
         """
         Constructs scala sequence from Java's List object.
@@ -71,7 +86,14 @@ class Hyperspace:
         >>> df = spark.read.parquet("./sample.parquet").toDF("c1", "c2", "c3")
         >>> hyperspace.createIndex(df, indexConfig)
         """
-        self.hyperspace.createIndex(dataFrame._jdf, self._getJavaCoveringIndexConfig(indexConfig))
+        if isinstance(indexConfig, CoveringIndexConfig):
+            self.hyperspace.createIndex(dataFrame._jdf, \
+                self._getJavaCoveringIndexConfig(indexConfig))
+        elif isinstance(indexConfig, ZOrderCoveringIndexConfig):
+            self.hyperspace.createIndex(dataFrame._jdf, \
+                self._getJavaZOrderCoveringIndexConfig(indexConfig))
+        else:
+            raise Exception("Invalid index config type: " + type(indexConfig).__name__)
 
     def deleteIndex(self, indexName):
         """
diff --git a/python/hyperspace/indexconfig.py b/python/hyperspace/indexconfig.py
index 6093d36c..98d61b4b 100644
--- a/python/hyperspace/indexconfig.py
+++ b/python/hyperspace/indexconfig.py
@@ -13,4 +13,19 @@ class CoveringIndexConfig:
         self.indexedColumns = indexedColumns
         self.includedColumns = includedColumns
 
+class ZOrderCoveringIndexConfig:
+    def __init__(self, indexName, indexedColumns, includedColumns):
+        """
+        Initializes ZOrderCoveringIndexConfig object.
+        :param indexName: index name
+        :param indexedColumns: indexed columns
+        :param includedColumns: included columns
+        :return: ZOrderCoveringIndexConfig object
+
+        >>> idxConfig = ZOrderCoveringIndexConfig("indexName", ["c1"], ["c2","c3"])
+        """
+        self.indexName = indexName
+        self.indexedColumns = indexedColumns
+        self.includedColumns = includedColumns
+
 IndexConfig = CoveringIndexConfig
diff --git a/python/hyperspace/tests/test_indexmanagement.py b/python/hyperspace/tests/test_indexmanagement.py
index 30e75991..23ff414e 100644
--- a/python/hyperspace/tests/test_indexmanagement.py
+++ b/python/hyperspace/tests/test_indexmanagement.py
@@ -34,6 +34,13 @@ class HyperspaceIndexManagementTests(HyperspaceTestCase):
         self.hyperspace.createIndex(self.df, idx_config)
         self.assertEqual(self.hyperspace.indexes().filter("""name = "idx1" """).count(), 1)
 
+    def test_index_create_zorder(self):
+        idx_config = ZOrderCoveringIndexConfig('zidx1', ['name'], ['age'])
+        self.hyperspace.createIndex(self.df, idx_config)
+        self.assertEqual(self.hyperspace.indexes().filter("""name = "zidx1" """).count(), 1)
+        self.assertEqual(self.hyperspace.index("zidx1"). \
+          filter("""kind = "ZOrderCoveringIndex" """).count(), 1)
+
     def test_index_delete(self):
         idx_config = IndexConfig('idx2', ['name'], ['age'])
         self.hyperspace.createIndex(self.df, idx_config)