Supported division pushdown

GrigoriyPA · Nov 12, 2024 · e85ffb3 · e85ffb3
1 parent 8fb382c
commit e85ffb3
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 3 deletions.
diff --git a/ydb/library/yql/providers/common/pushdown/collection.cpp b/ydb/library/yql/providers/common/pushdown/collection.cpp
@@ -360,6 +360,9 @@ bool CheckExpressionNodeForPushdown(const TExprBase& node, const TExprNode* lamb
     } else if (const auto op = node.Maybe<TCoUnaryArithmetic>(); op && settings.IsEnabled(TSettings::EFeatureFlag::UnaryOperators)) {
         return CheckExpressionNodeForPushdown(op.Cast().Arg(), lambdaArg, settings);
     } else if (const auto op = node.Maybe<TCoBinaryArithmetic>(); op && settings.IsEnabled(TSettings::EFeatureFlag::ArithmeticalExpressions)) {
+        if (!settings.IsEnabled(TSettings::EFeatureFlag::DivisionExpressions) && (op.Maybe<TCoDiv>() || op.Maybe<TCoMod>())) {
+            return false;
+        }
         return CheckExpressionNodeForPushdown(op.Cast().Left(), lambdaArg, settings) && CheckExpressionNodeForPushdown(op.Cast().Right(), lambdaArg, settings);
     }
     return false;

diff --git a/ydb/library/yql/providers/common/pushdown/settings.h b/ydb/library/yql/providers/common/pushdown/settings.h
@@ -28,14 +28,15 @@ struct TSettings {
         JustPassthroughOperators = 1 << 18, // if + coalesce + just
         InOperator = 1 << 19, // IN()
         IsDistinctOperator = 1 << 20, // IS NOT DISTINCT FROM / IS DISTINCT FROM 
+        DivisionExpressions = 1 << 21, // %, / -- NOTE: division by zero is not handled and also pushdown
 
         // Option which enables partial pushdown for sequence of OR
         // For example next predicate:
         // ($A AND $B) OR ($C AND $D)
         // May be partially pushdowned as:
         // $A OR $C
         // In case of unsupported / complicated expressions $B and $D
-        SplitOrOperator = 1 << 21
+        SplitOrOperator = 1 << 22
     };
 
     explicit TSettings(NLog::EComponent logComponent)

diff --git a/ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto b/ydb/library/yql/providers/generic/connector/api/service/protos/connector.proto
@@ -327,10 +327,11 @@ message TExpression {
             MUL = 1; // left_value * right_value
             ADD = 2; // left_value + right_value
             SUB = 3; // left_value - right_value
+            DIV = 7; // left_value / right_value
+            MOD = 8; // left_value % right_value
             BIT_AND = 4; // left_value & right_value
             BIT_OR = 5; // left_value | right_value
             BIT_XOR = 6; // left_value ^ right_value
-            // TODO: support `/` and `%`
         }
         EOperation operation = 1;
         TExpression left_value = 2;

diff --git a/ydb/library/yql/providers/generic/provider/yql_generic_predicate_pushdown.cpp b/ydb/library/yql/providers/generic/provider/yql_generic_predicate_pushdown.cpp
@@ -86,6 +86,8 @@ namespace NYql {
             MATCH_ARITHMETICAL(Sub, SUB);
             MATCH_ARITHMETICAL(Add, ADD);
             MATCH_ARITHMETICAL(Mul, MUL);
+            MATCH_ARITHMETICAL(Div, DIV);
+            MATCH_ARITHMETICAL(Mod, MOD);
 
             if (auto maybeNull = expression.Maybe<TCoNull>()) {
                 proto->mutable_null();
@@ -342,6 +344,12 @@ namespace NYql {
             case TExpression_TArithmeticalExpression::SUB:
                 operation = " - ";
                 break;
+            case TExpression_TArithmeticalExpression::DIV:
+                operation = " / ";
+                break;
+            case TExpression_TArithmeticalExpression::MOD:
+                operation = " % ";
+                break;
             case TExpression_TArithmeticalExpression::BIT_AND:
                 operation = " & ";
                 break;

diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_logical_opt.cpp
@@ -34,7 +34,7 @@ namespace {
                 // Operator features
                 EFlag::ExpressionAsPredicate | EFlag::ArithmeticalExpressions | EFlag::ImplicitConversionToInt64 |
                 EFlag::StringTypes | EFlag::LikeOperator | EFlag::DoNotCheckCompareArgumentsTypes | EFlag::InOperator |
-                EFlag::IsDistinctOperator | EFlag::JustPassthroughOperators |
+                EFlag::IsDistinctOperator | EFlag::JustPassthroughOperators | DivisionExpressions |
 
                 // Split features
                 EFlag::SplitOrOperator

diff --git a/ydb/tests/fq/yds/test_row_dispatcher.py b/ydb/tests/fq/yds/test_row_dispatcher.py
@@ -351,6 +351,8 @@ def test_filters_optional_field(self, kikimr, client):
         self.run_and_check(kikimr, client, sql + filter, data, expected, 'predicate: WHERE `flag`')
         filter = 'time * (field2 - field1) != 0'
         self.run_and_check(kikimr, client, sql + filter, data, expected, 'predicate: WHERE (`time` * (`field2` - `field1`)) <> 0')
+        filter = '(field1 % field2) / 5 = 1'
+        self.run_and_check(kikimr, client, sql + filter, data, expected, 'predicate: WHERE ((`field1` % `field2`) / 5) = 1')
         filter = ' event IS NOT DISTINCT FROM "event2"'
         self.run_and_check(kikimr, client, sql + filter, data, expected, 'predicate: WHERE `event` IS NOT DISTINCT FROM \\"event2\\"')
         filter = ' event IS DISTINCT FROM "event1"'