fix(databricks-jdbc-driver): Rolling window & count_distinct_approx (HLL) (#8323)

Author: ovr

packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts

@@ -33,6 +33,14 @@ export class DatabricksQuery extends BaseQuery {
     return `hll_union_agg(${sql})`;
   }
 
+  public hllCardinality(sql: string): string {
+    return `hll_sketch_estimate(${sql})`;
+  }
+
+  public hllCardinalityMerge(sql: string): string {
+    return `hll_sketch_estimate(hll_union_agg(${sql}))`;
+  }
+
   public countDistinctApprox(sql: string) {
     return `approx_count_distinct(${sql})`;
   }

packages/cubejs-jdbc-driver/src/JDBCDriver.ts

@@ -25,7 +25,7 @@ import path from 'path';
 import { DriverOptionsInterface, SupportedDrivers } from './supported-drivers';
 // eslint-disable-next-line @typescript-eslint/no-unused-vars
 import { JDBCDriverConfiguration } from './types';
-import { QueryStream, nextFn, Row } from './QueryStream';
+import { QueryStream, nextFn, Row, transformRow } from './QueryStream';
 
 const DriverManager = require('@cubejs-backend/jdbc/lib/drivermanager');
 const Connection = require('@cubejs-backend/jdbc/lib/connection');
@@ -246,6 +246,7 @@ export class JDBCDriver extends BaseDriver {
 
   protected async queryPromised(query: string, cancelObj: any, options: any) {
     options = options || {};
+
     try {
       const conn = await this.pool.acquire();
       try {
@@ -337,11 +338,18 @@ export class JDBCDriver extends BaseDriver {
     await setQueryTimeout(600);
     const executeQueryAsync = promisify(statement.execute.bind(statement));
     const resultSet = await executeQueryAsync(query);
-    const toObjArrayAsync =
-      resultSet.toObjArray && promisify(resultSet.toObjArray.bind(resultSet)) ||
-      (() => Promise.resolve(resultSet));
 
-    return toObjArrayAsync();
+    if (resultSet.toObjArray) {
+      const result: any = await (promisify(resultSet.toObjArray.bind(resultSet)))();
+
+      for (const [key, row] of Object.entries(result)) {
+        result[key] = transformRow(row);
+      }
+
+      return result;
+    }
+
+    return resultSet;
   }
 
   public async release() {

packages/cubejs-jdbc-driver/src/QueryStream.ts

@@ -9,6 +9,18 @@ export type nextFn = () => {
   value: Row,
 };
 
+export function transformRow(row: any) {
+  // eslint-disable-next-line no-restricted-syntax
+  for (const [name, field] of Object.entries(row)) {
+    // console.log({ name, field });
+    if (field instanceof Int8Array) {
+      row[name] = Buffer.from(field).toString('base64');
+    }
+  }
+
+  return row;
+}
+
 export class QueryStream extends Readable {
   private next: null | nextFn;
 
@@ -23,18 +35,6 @@ export class QueryStream extends Readable {
     this.next = nextFn;
   }
 
-  protected transformRow(row: any) {
-    // eslint-disable-next-line no-restricted-syntax
-    for (const [name, field] of Object.entries(row)) {
-      // console.log({ name, field });
-      if (field instanceof Int8Array) {
-        row[name] = Buffer.from(field).toString('base64');
-      }
-    }
-
-    return row;
-  }
-
   /**
    * @override
    */
@@ -44,7 +44,7 @@ export class QueryStream extends Readable {
         if (this.next) {
           const row = this.next();
           if (row.value) {
-            this.push(this.transformRow(row.value));
+            this.push(transformRow(row.value));
           }
           if (row.done) {
             this.push(null);

packages/cubejs-testing-drivers/fixtures/_schemas.json

@@ -339,6 +339,30 @@
           "rolling_window": {
             "trailing": "2 month"
           }
+        },
+        {
+          "name": "rollingCountApproxBy2Day",
+          "type": "count_distinct_approx",
+          "sql": "order_id",
+          "rolling_window": {
+            "trailing": "2 day"
+          }
+        },
+        {
+          "name": "rollingCountApproxBy2Week",
+          "type": "count_distinct_approx",
+          "sql": "order_id",
+          "rolling_window": {
+            "trailing": "2 week"
+          }
+        },
+        {
+          "name": "rollingCountApproxBy2Month",
+          "type": "count_distinct_approx",
+          "sql": "order_id",
+          "rolling_window": {
+            "trailing": "2 month"
+          }
         }
       ]
     }

packages/cubejs-testing-drivers/src/tests/testQueries.ts

@@ -1535,6 +1535,50 @@ export function testQueries(type: string, { includeIncrementalSchemaSuite, exten
       expect(response.rawData()).toMatchSnapshot();
     });
 
+    if (includeHLLSuite) {
+      execute('querying BigECommerce: rolling count_distinct_approx window by 2 day', async () => {
+        const response = await client.load({
+          measures: [
+            'BigECommerce.rollingCountApproxBy2Day',
+          ],
+          timeDimensions: [{
+            dimension: 'BigECommerce.orderDate',
+            granularity: 'month',
+            dateRange: ['2020-01-01', '2020-12-31'],
+          }],
+        });
+        expect(response.rawData()).toMatchSnapshot();
+      });
+
+      execute('querying BigECommerce: rolling count_distinct_approx window by 2 week', async () => {
+        const response = await client.load({
+          measures: [
+            'BigECommerce.rollingCountApproxBy2Week',
+          ],
+          timeDimensions: [{
+            dimension: 'BigECommerce.orderDate',
+            granularity: 'month',
+            dateRange: ['2020-01-01', '2020-12-31'],
+          }],
+        });
+        expect(response.rawData()).toMatchSnapshot();
+      });
+
+      execute('querying BigECommerce: rolling count_distinct_approx window by 2 month', async () => {
+        const response = await client.load({
+          measures: [
+            'BigECommerce.rollingCountApproxBy2Month',
+          ],
+          timeDimensions: [{
+            dimension: 'BigECommerce.orderDate',
+            granularity: 'month',
+            dateRange: ['2020-01-01', '2020-12-31'],
+          }],
+        });
+        expect(response.rawData()).toMatchSnapshot();
+      });
+    }
+
     execute('querying BigECommerce: totalProfitYearAgo', async () => {
       const response = await client.load({
         measures: [