Introduced the Column Table structure

Author: jprakash-db

src/databricks/sql/client.py

@@ -25,7 +25,7 @@
     ParamEscaper,
     inject_parameters,
     transform_paramstyle,
-    ArrowQueue,
+    ColumnTable,
     ColumnQueue
 )
 from databricks.sql.parameters.native import (
@@ -1152,10 +1152,10 @@ def _convert_columnar_table(self, table):
         column_names = [c[0] for c in self.description]
         ResultRow = Row(*column_names)
         result = []
-        for row_index in range(len(table[0])):
+        for row_index in range(table.num_rows):
             curr_row = []
-            for col_index in range(len(table)):
-                curr_row.append(table[col_index][row_index])
+            for col_index in range(table.num_columns):
+                curr_row.append(table.get_item(col_index,  row_index))
             result.append(ResultRow(*curr_row))
 
         return result
@@ -1235,11 +1235,11 @@ def merge_columnar(self, result1, result2):
         :return:
         """
 
-        if len(result1) != len(result2):
-            raise ValueError("The number of columns in both results must be the same")
+        if result1.column_names != result2.column_names:
+            raise ValueError("The columns in the results don't match")
 
-        merged_result = [result1[i] + result2[i] for i in range(len(result1))]
-        return merged_result
+        merged_result = [result1.column_table[i] + result2.column_table[i] for i in range(result1.num_columns)]
+        return ColumnTable(merged_result, result1.column_names)
 
     def fetchmany_columnar(self, size: int):
         """
@@ -1250,8 +1250,8 @@ def fetchmany_columnar(self, size: int):
             raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
 
         results = self.results.next_n_rows(size)
-        n_remaining_rows = size - len(results[0])
-        self._next_row_index += len(results[0])
+        n_remaining_rows = size - results.num_rows
+        self._next_row_index += results.num_rows
 
         while (
                 n_remaining_rows > 0
@@ -1261,8 +1261,8 @@ def fetchmany_columnar(self, size: int):
             self._fill_results_buffer()
             partial_results = self.results.next_n_rows(n_remaining_rows)
             results = self.merge_columnar(results, partial_results)
-            n_remaining_rows -= len(partial_results[0])
-            self._next_row_index += len(partial_results[0])
+            n_remaining_rows -= partial_results.num_rows
+            self._next_row_index += partial_results.num_rows
 
         return results
 
@@ -1282,13 +1282,13 @@ def fetchall_arrow(self) -> "pyarrow.Table":
     def fetchall_columnar(self):
         """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
         results = self.results.remaining_rows()
-        self._next_row_index += len(results[0])
+        self._next_row_index += results.num_rows
 
         while not self.has_been_closed_server_side and self.has_more_rows:
             self._fill_results_buffer()
             partial_results = self.results.remaining_rows()
             results = self.merge_columnar(results, partial_results)
-            self._next_row_index += len(partial_results[0])
+            self._next_row_index += partial_results.num_rows
 
         return results
 

src/databricks/sql/utils.py

@@ -88,7 +88,7 @@ def build_queue(
                 column_table, description
             )
 
-            return ColumnQueue(converted_column_table, column_names)
+            return ColumnQueue(ColumnTable(converted_column_table, column_names))
         elif row_set_type == TSparkRowSetType.URL_BASED_SET:
             return CloudFetchQueue(
                 schema_bytes=arrow_schema_bytes,
@@ -102,27 +102,47 @@ def build_queue(
         else:
             raise AssertionError("Row set type is not valid")
 
+class ColumnTable:
+    def __init__(self, column_table, column_names):
+        self.column_table = column_table
+        self.column_names = column_names
+
+    @property
+    def num_rows(self):
+        if len(self.column_table) == 0:
+            return 0
+        else:
+            return len(self.column_table[0])
+
+    @property
+    def num_columns(self):
+        return len(self.column_names)
+
+    def get_item(self, col_index, row_index):
+        return self.column_table[col_index][row_index]
+
+    def slice(self, curr_index, length):
+        sliced_column_table = [column[curr_index : curr_index + length] for column in self.column_table]
+        return ColumnTable(sliced_column_table, self.column_names)
+
+
 class ColumnQueue(ResultSetQueue):
-    def __init__(self, columnar_table, column_names):
-        self.columnar_table = columnar_table
+    def __init__(self, column_table: ColumnTable):
+        self.column_table = column_table
         self.cur_row_index = 0
-        self.n_valid_rows = len(columnar_table[0])
-        self.column_names = column_names
+        self.n_valid_rows = column_table.num_rows
 
     def next_n_rows(self, num_rows):
         length = min(num_rows, self.n_valid_rows - self.cur_row_index)
-        # Slicing using the default python slice
-        next_data = [
-            column[self.cur_row_index : self.cur_row_index + length]
-            for column in self.columnar_table
-        ]
-        self.cur_row_index += length
-        return next_data
+
+        slice = self.column_table.slice(self.cur_row_index, length)
+        self.cur_row_index += slice.num_rows
+        return slice
 
     def remaining_rows(self):
-        next_data = [column[self.cur_row_index :] for column in self.columnar_table]
-        self.cur_row_index += len(next_data[0])
-        return next_data
+        slice = self.column_table.slice(self.cur_row_index, self.n_valid_rows - self.cur_row_index)
+        self.cur_row_index += slice.num_rows
+        return slice
 
 
 class ArrowQueue(ResultSetQueue):