Fix nth_value function

trinodb · Nov 5, 2014 · a461eee · a461eee
1 parent dc1c6ef
commit a461eee
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 14 deletions.
diff --git a/presto-docs/src/main/sphinx/release/release-0.81.rst b/presto-docs/src/main/sphinx/release/release-0.81.rst
@@ -7,3 +7,8 @@ Hive Changes
 
 * Fix ORC predicate pushdown.
 * Fix column selection in RCFile.
+
+General Changes
+---------------
+
+* Fix :func:`nth_value` function.
diff --git a/presto-main/src/main/java/com/facebook/presto/operator/window/NthValueFunction.java b/presto-main/src/main/java/com/facebook/presto/operator/window/NthValueFunction.java
@@ -102,23 +102,23 @@ public void processRow(BlockBuilder output, boolean newPeerGroup, int peerGroupC
     {
         if (pagesIndex.isNull(offsetChannel, currentPosition)) {
             output.appendNull();
-            return;
         }
-
-        int offset = Ints.checkedCast(pagesIndex.getLong(offsetChannel, currentPosition));
-        checkCondition(offset >= 1, INVALID_FUNCTION_ARGUMENT, "Offset must be at least 1");
-
-        // offset is base 1
-        int valuePosition = partitionStartPosition + offset - 1;
-
-        // if the value is out of range, result is null
-        if (valuePosition >= partitionStartPosition + partitionRowCount) {
-            output.appendNull();
-            return;
+        else {
+            int offset = Ints.checkedCast(pagesIndex.getLong(offsetChannel, currentPosition));
+            checkCondition(offset >= 1, INVALID_FUNCTION_ARGUMENT, "Offset must be at least 1");
+
+            // offset is base 1
+            int valuePosition = partitionStartPosition + offset - 1;
+
+            // if the value is out of range, result is null
+            if (valuePosition >= partitionStartPosition + partitionRowCount) {
+                output.appendNull();
+            }
+            else {
+                pagesIndex.appendTo(valueChannel, valuePosition, output);
+            }
         }
 
-        pagesIndex.appendTo(valueChannel, valuePosition, output);
-
         currentPosition++;
     }
 }
diff --git a/presto-main/src/test/java/com/facebook/presto/operator/window/TestWindowFunctions.java b/presto-main/src/test/java/com/facebook/presto/operator/window/TestWindowFunctions.java
@@ -721,4 +721,65 @@ public void testLeadFunction()
                         .row(null, null, -1)
                         .build(), queryRunner);
     }
+
+    @Test
+    public void testNthValue()
+    {
+        // Test the function with a constant offset
+        assertWindowQuery("nth_value(orderkey, 2) OVER (PARTITION BY orderstatus ORDER BY orderkey)",
+                resultBuilder(TEST_SESSION, BIGINT, VARCHAR, VARCHAR)
+                        .row(3, "F", 5)
+                        .row(5, "F", 5)
+                        .row(6, "F", 5)
+                        .row(33, "F", 5)
+                        .row(1, "O", 2)
+                        .row(2, "O", 2)
+                        .row(4, "O", 2)
+                        .row(7, "O", 2)
+                        .row(32, "O", 2)
+                        .row(34, "O", 2)
+                        .build(), queryRunner);
+        assertWindowQueryWithNulls("nth_value(orderkey, 2) OVER (PARTITION BY orderstatus ORDER BY orderkey)",
+                resultBuilder(TEST_SESSION, BIGINT, VARCHAR, VARCHAR)
+                        .row(3, "F", 5)
+                        .row(5, "F", 5)
+                        .row(null, "F", 5)
+                        .row(null, "F", 5)
+                        .row(34, "O", null)
+                        .row(null, "O", null)
+                        .row(1, null, 7)
+                        .row(7, null, 7)
+                        .row(null, null, 7)
+                        .row(null, null, 7)
+                        .build(), queryRunner);
+
+        // Test the function with a variable offset
+        assertWindowQuery("nth_value(orderkey, orderkey) OVER (PARTITION BY orderstatus ORDER BY orderkey)",
+                resultBuilder(TEST_SESSION, BIGINT, VARCHAR, VARCHAR)
+                        .row(3, "F", 6)
+                        .row(5, "F", null)
+                        .row(6, "F", null)
+                        .row(33, "F", null)
+                        .row(1, "O", 1)
+                        .row(2, "O", 2)
+                        .row(4, "O", 7)
+                        .row(7, "O", null)
+                        .row(32, "O", null)
+                        .row(34, "O", null)
+                        .build(), queryRunner);
+        assertWindowQueryWithNulls(
+                "nth_value(orderkey, orderkey) OVER (PARTITION BY orderstatus ORDER BY orderkey DESC NULLS FIRST)",
+                resultBuilder(TEST_SESSION, BIGINT, VARCHAR, VARCHAR)
+                        .row(3, "F", 5)
+                        .row(5, "F", null)
+                        .row(null, "F", null)
+                        .row(null, "F", null)
+                        .row(34, "O", null)
+                        .row(null, "O", null)
+                        .row(1, null, null)
+                        .row(7, null, null)
+                        .row(null, null, null)
+                        .row(null, null, null)
+                        .build(), queryRunner);
+    }
 }