feat(cubesql): WHERE SQL push down (#7808)

* feat(cubesql): `WHERE` SQL push down * Remove TODO as we need to allocate names for segments
cube-js · Feb 24, 2024 · 98b5709 · 98b5709
1 parent 60aad90
commit 98b5709
Show file tree

Hide file tree

Showing 26 changed files with 558 additions and 56 deletions.
diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts
@@ -1233,6 +1233,7 @@ class ApiGateway {
       ...query,
       measures: (query.measures || []).map(m => (typeof m === 'string' ? this.parseMemberExpression(m) : m)),
       dimensions: (query.dimensions || []).map(m => (typeof m === 'string' ? this.parseMemberExpression(m) : m)),
+      segments: (query.segments || []).map(m => (typeof m === 'string' ? this.parseMemberExpression(m) : m)),
     };
   }
 

diff --git a/packages/cubejs-api-gateway/src/query.js b/packages/cubejs-api-gateway/src/query.js
@@ -105,7 +105,7 @@ const querySchema = Joi.object().keys({
     Joi.object().pattern(id, Joi.valid('asc', 'desc')),
     Joi.array().items(Joi.array().min(2).ordered(id, Joi.valid('asc', 'desc')))
   ),
-  segments: Joi.array().items(id),
+  segments: Joi.array().items(Joi.alternatives(id, memberExpression)),
   timezone: Joi.string(),
   limit: Joi.number().integer().min(1),
   offset: Joi.number().integer().min(0),

diff --git a/packages/cubejs-api-gateway/src/types/query.ts b/packages/cubejs-api-gateway/src/types/query.ts
@@ -64,7 +64,7 @@ interface Query {
   dimensions?: (Member | TimeMember | MemberExpression)[];
   filters?: (QueryFilter | LogicalAndFilter | LogicalOrFilter)[];
   timeDimensions?: QueryTimeDimension[];
-  segments?: Member[];
+  segments?: (Member | MemberExpression)[];
   limit?: null | number;
   offset?: number;
   total?: boolean;

diff --git a/packages/cubejs-backend-native/src/transport.rs b/packages/cubejs-backend-native/src/transport.rs
@@ -297,8 +297,8 @@ impl TransportService for NodeBridgeTransport {
         )
         .await?;
 
-        if let Some(error) = response.get("error") {
-            if let Some(stack) = response.get("stack") {
+        if let Some(error) = response.get("error").and_then(|e| e.as_str()) {
+            if let Some(stack) = response.get("stack").and_then(|e| e.as_str()) {
                 return Err(CubeError::user(format!(
                     "Error during SQL generation: {}\n{}",
                     error, stack

diff --git a/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts b/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts
@@ -104,6 +104,7 @@ export class DatabricksQuery extends BaseQuery {
     templates.functions.LTRIM = 'LTRIM({{ args|reverse|join(", ") }})';
     templates.functions.RTRIM = 'RTRIM({{ args|reverse|join(", ") }})';
     templates.functions.DATEDIFF = 'DATEDIFF({{ date_part }}, DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}))';
+    templates.expressions.timestamp_literal = 'from_utc_timestamp(\'{{ value }}\', \'UTC\')';
     return templates;
   }
 }
diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js
@@ -2500,6 +2500,7 @@ export class BaseQuery {
       statements: {
         select: 'SELECT {{ select_concat | map(attribute=\'aliased\') | join(\', \') }} \n' +
           'FROM (\n  {{ from }}\n) AS {{ from_alias }} \n' +
+          '{% if filter %} WHERE {{ filter }}{% endif %}' +
           '{% if group_by %} GROUP BY {{ group_by | map(attribute=\'index\') | join(\', \') }}{% endif %}' +
           '{% if order_by %} ORDER BY {{ order_by | map(attribute=\'expr\') | join(\', \') }}{% endif %}' +
           '{% if limit %}\nLIMIT {{ limit }}{% endif %}' +

diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseSegment.ts b/packages/cubejs-schema-compiler/src/adapter/BaseSegment.ts
@@ -1,12 +1,34 @@
 import type { BaseQuery } from './BaseQuery';
 
 export class BaseSegment {
+  public readonly expression: any;
+
+  public readonly expressionCubeName: any;
+
+  public readonly expressionName: any;
+
+  public readonly isMemberExpression: boolean = false;
+
   public constructor(
     protected readonly query: BaseQuery,
-    public readonly segment: string
-  ) {}
+    public readonly segment: string | any
+  ) {
+    if (segment.expression) {
+      this.expression = segment.expression;
+      this.expressionCubeName = segment.cubeName;
+      this.expressionName = segment.expressionName || `${segment.cubeName}.${segment.name}`;
+      this.isMemberExpression = !!segment.definition;
+    }
+  }
 
   public filterToWhere() {
+    return this.segmentSql();
+  }
+
+  public segmentSql() {
+    if (this.expression) {
+      return this.query.evaluateSymbolSql(this.expressionCubeName, this.expressionName, this.definition(), 'segment');
+    }
     return this.query.segmentSql(this);
   }
 
@@ -19,6 +41,11 @@ export class BaseSegment {
   }
 
   public definition() {
+    if (this.expression) {
+      return {
+        sql: this.expression
+      };
+    }
     return this.segmentDefinition();
   }
 
@@ -27,6 +54,9 @@ export class BaseSegment {
   }
 
   public cube() {
+    if (this.expression) {
+      return this.query.cubeEvaluator.cubeFromPath(this.expressionCubeName);
+    }
     return this.query.cubeEvaluator.cubeFromPath(this.segment);
   }
 
@@ -35,14 +65,16 @@ export class BaseSegment {
   }
 
   public path() {
+    if (this.expression) {
+      return null;
+    }
     return this.query.cubeEvaluator.parsePath('segments', this.segment);
   }
 
   public expressionPath() {
-    // TODO expression support
-    // if (this.expression) {
-    //   return `expr:${this.expression.expressionName}`;
-    // }
+    if (this.expression) {
+      return `expr:${this.expression.expressionName}`;
+    }
     return this.query.cubeEvaluator.pathFromArray(this.path());
   }
 }
diff --git a/packages/cubejs-schema-compiler/src/adapter/BigqueryQuery.ts b/packages/cubejs-schema-compiler/src/adapter/BigqueryQuery.ts
@@ -160,6 +160,7 @@ export class BigqueryQuery extends BaseQuery {
     templates.expressions.binary = '{% if op == \'%\' %}MOD({{ left }}, {{ right }}){% else %}({{ left }} {{ op }} {{ right }}){% endif %}';
     templates.expressions.interval = 'INTERVAL {{ interval }}';
     templates.expressions.extract = 'EXTRACT({% if date_part == \'DOW\' %}DAYOFWEEK{% elif date_part == \'DOY\' %}DAYOFYEAR{% else %}{{ date_part }}{% endif %} FROM {{ expr }})';
+    templates.expressions.timestamp_literal = 'TIMESTAMP(\'{{ value }}\')';
     return templates;
   }
 }
diff --git a/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.ts b/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.ts
@@ -58,6 +58,7 @@ export class PostgresQuery extends BaseQuery {
     templates.functions.DATEDIFF = 'CASE WHEN LOWER(\'{{ date_part }}\') IN (\'year\', \'quarter\', \'month\') THEN (EXTRACT(YEAR FROM AGE(DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }}))) * 12 + EXTRACT(MONTH FROM AGE(DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }})))) / CASE LOWER(\'{{ date_part }}\') WHEN \'year\' THEN 12 WHEN \'quarter\' THEN 3 WHEN \'month\' THEN 1 END ELSE EXTRACT(EPOCH FROM DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}) - DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }})) / EXTRACT(EPOCH FROM \'1 {{ date_part }}\'::interval) END::bigint';
     templates.expressions.interval = 'INTERVAL \'{{ interval }}\'';
     templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})';
+    templates.expressions.timestamp_literal = 'timestamptz \'{{ value }}\'';
 
     return templates;
   }

diff --git a/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts b/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts
@@ -121,6 +121,7 @@ export class PrestodbQuery extends BaseQuery {
       '{% if limit %}\nLIMIT {{ limit }}{% endif %}';
     templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})';
     templates.expressions.interval = 'INTERVAL \'{{ num }}\' {{ date_part }}';
+    templates.expressions.timestamp_literal = 'from_iso8601_timestamp(\'{{ value }}\')';
     return templates;
   }
 }
diff --git a/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts b/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts
@@ -62,6 +62,7 @@ export class SnowflakeQuery extends BaseQuery {
     templates.functions.BTRIM = 'TRIM({{ args_concat }})';
     templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})';
     templates.expressions.interval = 'INTERVAL \'{{ interval }}\'';
+    templates.expressions.timestamp_literal = '\'{{ value }}\'::timestamp_tz';
     return templates;
   }
 }
diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs
@@ -10,7 +10,7 @@ use crate::{
     },
     CubeError,
 };
-use chrono::{Days, NaiveDate};
+use chrono::{Days, NaiveDate, SecondsFormat, TimeZone, Utc};
 use cubeclient::models::V1LoadRequestQuery;
 use datafusion::{
     error::{DataFusionError, Result},
@@ -165,6 +165,48 @@ lazy_static! {
     static ref DATE_PART_REGEX: Regex = Regex::new("^[A-Za-z_ ]+$").unwrap();
 }
 
+macro_rules! generate_sql_for_timestamp {
+    (@generic $value:ident, $value_block:expr, $sql_generator:expr, $sql_query:expr) => {
+        if let Some($value) = $value {
+            let value = $value_block.to_rfc3339_opts(SecondsFormat::Millis, true);
+            (
+                $sql_generator
+                    .get_sql_templates()
+                    .timestamp_literal_expr(value)
+                    .map_err(|e| {
+                        DataFusionError::Internal(format!(
+                            "Can't generate SQL for timestamp: {}",
+                            e
+                        ))
+                    })?,
+                $sql_query,
+            )
+        } else {
+            ("NULL".to_string(), $sql_query)
+        }
+    };
+    ($value:ident, timestamp, $sql_generator:expr, $sql_query:expr) => {
+        generate_sql_for_timestamp!(
+            @generic $value, { Utc.timestamp_opt($value as i64, 0).unwrap() }, $sql_generator, $sql_query
+        )
+    };
+    ($value:ident, timestamp_millis_opt, $sql_generator:expr, $sql_query:expr) => {
+        generate_sql_for_timestamp!(
+            @generic $value, { Utc.timestamp_millis_opt($value as i64).unwrap() }, $sql_generator, $sql_query
+        )
+    };
+    ($value:ident, timestamp_micros, $sql_generator:expr, $sql_query:expr) => {
+        generate_sql_for_timestamp!(
+            @generic $value, { Utc.timestamp_micros($value as i64).unwrap() }, $sql_generator, $sql_query
+        )
+    };
+    ($value:ident, $method:ident, $sql_generator:expr, $sql_query:expr) => {
+        generate_sql_for_timestamp!(
+            @generic $value, { Utc.$method($value as i64) }, $sql_generator, $sql_query
+        )
+    };
+}
+
 impl CubeScanWrapperNode {
     pub async fn generate_sql(
         &self,
@@ -310,7 +352,7 @@ impl CubeScanWrapperNode {
                         window_expr,
                         from,
                         joins: _joins,
-                        filter_expr: _filter_expr,
+                        filter_expr,
                         having_expr: _having_expr,
                         limit,
                         offset,
@@ -443,6 +485,20 @@ impl CubeScanWrapperNode {
                             )
                             .await?;
 
+                            let (filter, sql) = Self::generate_column_expr(
+                                plan.clone(),
+                                schema.clone(),
+                                filter_expr.clone(),
+                                sql,
+                                generator.clone(),
+                                &column_remapping,
+                                &mut next_remapping,
+                                alias.clone(),
+                                can_rename_columns,
+                                ungrouped_scan_node.clone(),
+                            )
+                            .await?;
+
                             let (window, sql) = Self::generate_column_expr(
                                 plan.clone(),
                                 schema.clone(),
@@ -548,6 +604,17 @@ impl CubeScanWrapperNode {
                                         })
                                         .collect::<Result<_>>()?,
                                 );
+                                load_request.segments = Some(
+                                    filter
+                                        .iter()
+                                        .map(|m| {
+                                            Self::ungrouped_member_def(
+                                                m,
+                                                &ungrouped_scan_node.used_cubes,
+                                            )
+                                        })
+                                        .collect::<Result<_>>()?,
+                                );
                                 if !order_expr.is_empty() {
                                     load_request.order = Some(
                                         order_expr
@@ -651,7 +718,16 @@ impl CubeScanWrapperNode {
                                         aggregate,
                                         // TODO
                                         from_alias.unwrap_or("".to_string()),
-                                        None,
+                                        if !filter.is_empty() {
+                                            Some(
+                                                filter
+                                                    .iter()
+                                                    .map(|f| f.expr.to_string())
+                                                    .join(" AND "),
+                                            )
+                                        } else {
+                                            None
+                                        },
                                         None,
                                         order,
                                         limit,
@@ -1284,10 +1360,33 @@ impl CubeScanWrapperNode {
                             }
                         }
                         // ScalarValue::Date64(_) => {}
-                        // ScalarValue::TimestampSecond(_, _) => {}
-                        // ScalarValue::TimestampMillisecond(_, _) => {}
-                        // ScalarValue::TimestampMicrosecond(_, _) => {}
-                        // ScalarValue::TimestampNanosecond(_, _) => {}
+                        ScalarValue::TimestampSecond(s, _) => {
+                            generate_sql_for_timestamp!(s, timestamp, sql_generator, sql_query)
+                        }
+                        ScalarValue::TimestampMillisecond(ms, None) => {
+                            generate_sql_for_timestamp!(
+                                ms,
+                                timestamp_millis_opt,
+                                sql_generator,
+                                sql_query
+                            )
+                        }
+                        ScalarValue::TimestampMicrosecond(ms, None) => {
+                            generate_sql_for_timestamp!(
+                                ms,
+                                timestamp_micros,
+                                sql_generator,
+                                sql_query
+                            )
+                        }
+                        ScalarValue::TimestampNanosecond(nanoseconds, None) => {
+                            generate_sql_for_timestamp!(
+                                nanoseconds,
+                                timestamp_nanos,
+                                sql_generator,
+                                sql_query
+                            )
+                        }
                         ScalarValue::IntervalYearMonth(x) => {
                             if let Some(x) = x {
                                 let (num, date_part) = (x, "MONTH");