feat: update test file format to support aggregate functions

substrait-io · Nov 8, 2024 · d4cf096 · d4cf096
1 parent 752aa63
commit d4cf096
Show file tree

Hide file tree

Showing 14 changed files with 9,496 additions and 5,852 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -23,9 +23,8 @@ repos:
   - id: flake8
 - repo: local
   hooks:
-  - id: check-substrait-extensions
-    name: Check Substrait extensions
-    entry: pytest tests/test_extensions.py::test_read_substrait_extensions
+  - id: check-substrait-extensions_coverage
+    name: Check Substrait extensions and test coverage
+    entry: pytest tests/test_extensions.py::test_substrait_extension_coverage
     language: python
     pass_filenames: false
-
diff --git a/grammar/FuncTestCaseLexer.g4 b/grammar/FuncTestCaseLexer.g4
@@ -10,6 +10,7 @@ Whitespace    : [ \t\n\r]+ -> channel(HIDDEN) ;
 
 TripleHash: '###';
 SubstraitScalarTest: 'SUBSTRAIT_SCALAR_TEST';
+SubstraitAggregateTest: 'SUBSTRAIT_AGGREGATE_TEST';
 SubstraitInclude: 'SUBSTRAIT_INCLUDE';
 
 FormatVersion
@@ -20,6 +21,7 @@ DescriptionLine
     : '# ' ~[\r\n]* '\r'? '\n'
     ;
 
+Define: 'DEFINE';
 ErrorResult: '<!ERROR>';
 UndefineResult: '<!UNDEFINED>';
 Overflow: 'OVERFLOW';
@@ -102,3 +104,7 @@ NullLiteral: 'null';
 StringLiteral
     : '\'' ('\\' . | '\'\'' | ~['\\])* '\''
     ;
+
+ColumnName
+    : 'COL' Int
+    ;
diff --git a/grammar/FuncTestCaseParser.g4 b/grammar/FuncTestCaseParser.g4
@@ -15,7 +15,7 @@ header
     ;
 
 version
-    : TripleHash SubstraitScalarTest Colon FormatVersion
+    : TripleHash (SubstraitScalarTest | SubstraitAggregateTest) Colon FormatVersion
     ;
 
 include
@@ -31,7 +31,8 @@ testCase
     ;
 
 testGroup
-    : testGroupDescription (testCase)+
+    : testGroupDescription (testCase)+                          #scalarFuncTestGroup
+    | testGroupDescription (aggFuncTestCase)+                   #aggregateFuncTestGroup
     ;
 
 arguments
@@ -58,6 +59,63 @@ argument
     | intervalDayArg
     ;
 
+aggFuncTestCase
+    : aggFuncCall ( OBracket func_options CBracket )? Eq result
+    ;
+
+aggFuncCall
+    : tableData funcName=Identifier OParen qualifiedAggregateFuncArgs CParen            #multiArgAggregateFuncCall
+    | tableRows functName=Identifier OParen aggregateFuncArgs CParen                    #compactAggregateFuncCall
+    | functName=Identifier OParen dataColumn CParen                                     #singleArgAggregateFuncCall
+    ;
+
+tableData
+    : Define tableName=Identifier OParen dataType (Comma dataType)* CParen Eq tableRows
+    ;
+
+tableRows
+    : OParen (literalValueList (Comma literalValueList)*)? CParen
+    ;
+
+dataColumn
+    : literalValueList DoubleColon dataType
+    ;
+
+literalValueList
+    : OParen (literal (Comma literal)*)? CParen
+    ;
+
+literal
+    : NullLiteral
+    | numericLiteral
+    | BooleanLiteral
+    | StringLiteral
+    | DateLiteral
+    | TimeLiteral
+    | TimestampLiteral
+    | TimestampTzLiteral
+    | IntervalYearLiteral
+    | IntervalDayLiteral
+    ;
+
+qualifiedAggregateFuncArgs
+    : qualifiedAggregateFuncArg (Comma qualifiedAggregateFuncArg)*
+    ;
+
+aggregateFuncArgs
+    : aggregateFuncArg (Comma aggregateFuncArg)*
+    ;
+
+qualifiedAggregateFuncArg
+    : tableName=Identifier Dot ColumnName
+    | argument
+    ;
+
+aggregateFuncArg
+    : ColumnName DoubleColon dataType
+    | argument
+    ;
+
 numericLiteral
     : DecimalLiteral | IntegerLiteral | floatLiteral
     ;
@@ -66,7 +124,7 @@ floatLiteral
     : FloatLiteral | NaN
     ;
 
-nullArg: NullLiteral DoubleColon datatype;
+nullArg: NullLiteral DoubleColon dataType;
 
 intArg: IntegerLiteral DoubleColon (I8 | I16 | I32 | I64);
 
@@ -77,11 +135,11 @@ decimalArg
     ;
 
 booleanArg
-    : BooleanLiteral DoubleColon Bool
+    : BooleanLiteral DoubleColon (Bool | Boolean)
     ;
 
 stringArg
-    : StringLiteral DoubleColon Str
+    : StringLiteral DoubleColon (Str | String)
     ;
 
 dateArg
@@ -93,19 +151,19 @@ timeArg
     ;
 
 timestampArg
-    : TimestampLiteral DoubleColon Ts
+    : TimestampLiteral DoubleColon (Ts | Timestamp)
     ;
 
 timestampTzArg
-    : TimestampTzLiteral DoubleColon TsTZ
+    : TimestampTzLiteral DoubleColon (TsTZ | Timestamp_TZ)
     ;
 
 intervalYearArg
-    : IntervalYearLiteral DoubleColon IYear
+    : IntervalYearLiteral DoubleColon (IYear | Interval_Year)
     ;
 
 intervalDayArg
-    : IntervalDayLiteral DoubleColon IDay
+    : IntervalDayLiteral DoubleColon (IDay | Interval_Day)
     ;
 
 intervalYearLiteral
@@ -126,53 +184,56 @@ timeInterval
     | fractionalSeconds=IntegerLiteral FractionalSecondSuffix
     ;
 
-datatype
+dataType
     : scalarType
     | parameterizedType
     ;
 
 scalarType
-  : Bool #Boolean
-  | I8 #i8
-  | I16 #i16
-  | I32 #i32
-  | I64 #i64
-  | FP32 #fp32
-  | FP64 #fp64
-  | Str #string
-  | Binary #binary
-  | Ts #timestamp
-  | TsTZ #timestampTz
-  | Date #date
-  | Time #time
-  | IDay #intervalDay
-  | IYear #intervalYear
-  | UUID #uuid
-  | UserDefined Identifier #userDefined
+  : (Bool | Boolean)        #boolean
+  | I8                      #i8
+  | I16                     #i16
+  | I32                     #i32
+  | I64                     #i64
+  | FP32                    #fp32
+  | FP64                    #fp64
+  | (Str | String)          #string
+  | (Binary | VBin)         #binary
+  | (Ts | Timestamp)        #timestamp
+  | (TsTZ | Timestamp_TZ)   #timestampTz
+  | Date                    #date
+  | Time                    #time
+  | (IDay | Interval_Year)  #intervalDay
+  | (IYear | Interval_Day)  #intervalYear
+  | UUID                    #uuid
+  | UserDefined Identifier  #userDefined
   ;
 
 fixedCharType
-    : FChar isnull=QMark? OAngleBracket len=numericParameter CAngleBracket #fixedChar
+    : (FChar | FixedChar) isnull=QMark? OAngleBracket len=numericParameter CAngleBracket #fixedChar
     ;
 
 varCharType
-    : VChar isnull=QMark? OAngleBracket len=numericParameter CAngleBracket #varChar
+    : (VChar | VarChar) isnull=QMark? OAngleBracket len=numericParameter CAngleBracket #varChar
     ;
 
 fixedBinaryType
-    : FBin isnull=QMark? OAngleBracket len=numericParameter CAngleBracket #fixedBinary
+    : (FBin | FixedBinary) isnull=QMark? OAngleBracket len=numericParameter CAngleBracket #fixedBinary
     ;
 
 decimalType
-    : Dec isnull=QMark? (OAngleBracket precision=numericParameter Comma scale=numericParameter CAngleBracket)?  #decimal
+    : (Dec | Decimal) isnull=QMark?
+        (OAngleBracket precision=numericParameter Comma scale=numericParameter CAngleBracket)?  #decimal
     ;
 
 precisionTimestampType
-    : PTs isnull=QMark? OAngleBracket precision=numericParameter CAngleBracket #precisionTimestamp
+    : (PTs | Precision_Timestamp) isnull=QMark?
+        OAngleBracket precision=numericParameter CAngleBracket #precisionTimestamp
     ;
 
 precisionTimestampTZType
-    : PTsTZ isnull=QMark? OAngleBracket precision=numericParameter CAngleBracket #precisionTimestampTZ
+    : (PTsTZ | Precision_Timestamp_TZ) isnull=QMark?
+        OAngleBracket precision=numericParameter CAngleBracket #precisionTimestampTZ
     ;
 
 parameterizedType

diff --git a/tests/cases/arithmetic/max.test b/tests/cases/arithmetic/max.test
@@ -0,0 +1,18 @@
+### SUBSTRAIT_AGGREGATE_TEST: v1.0
+### SUBSTRAIT_INCLUDE: '/extensions/functions_arithmetic.yaml'
+
+# basic: Basic examples without any special cases
+max((20, -3, 1, -10, 0, 5)::i8) = 20::i8
+max((-32768, 32767, 20000, -30000)::i16) = 32767::i16
+max((-214748648, 214748647, 21470048, 4000000)::i32) = 214748647::i32
+max((2000000000, -3217908979, 629000000, -100000000, 0, 987654321)::i64) = 2000000000::i64
+max((2.5, 0, 5.0, -2.5, -7.5)::fp32) = 5.0::fp32
+max((1.5e+308, 1.5e+10, -1.5e+8, -1.5e+7, -1.5e+70)::fp64) = 1.5e+308::fp64
+
+# null_handling: Examples with null as input or output
+max((Null, Null, Null)::i16) = Null::i16
+max(()::i16) = Null::i16
+max((2000000000, Null, 629000000, -100000000, Null, 987654321)::i64) = 2000000000::i64
+max((Null, inf)::fp64) = inf::fp64
+max((Null, -inf, -1.5e+8, -1.5e+7, -1.5e+70)::fp64) = -1.5e+7::fp64
+max((1.5e+308, 1.5e+10, Null, -1.5e+7, Null)::fp64) = 1.5e+308::fp64