Column#apply_window_function supports :group_keys

ranguba · Apr 25, 2017 · 4ea68c5 · 4ea68c5
1 parent 3445bef
commit 4ea68c5
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 1 deletion.
diff --git a/ext/groonga/rb-grn-data-column.c b/ext/groonga/rb-grn-data-column.c
@@ -93,6 +93,7 @@ rb_grn_data_column_apply_window_function (int argc, VALUE *argv, VALUE self)
     grn_obj *window_function_call = NULL;
     VALUE rb_options;
     VALUE rb_sort_keys;
+    VALUE rb_group_keys;
     VALUE rb_builder;
     VALUE rb_window_function_call;
 
@@ -106,6 +107,7 @@ rb_grn_data_column_apply_window_function (int argc, VALUE *argv, VALUE self)
     rb_scan_args(argc, argv, "01", &rb_options);
     rb_grn_scan_options(rb_options,
                         "sort_keys", &rb_sort_keys,
+                        "group_keys", &rb_group_keys,
                         NULL);
 
     if (!NIL_P(rb_sort_keys)) {
@@ -126,6 +128,24 @@ rb_grn_data_column_apply_window_function (int argc, VALUE *argv, VALUE self)
                                     rb_table);
     }
 
+    if (!NIL_P(rb_group_keys)) {
+        VALUE rb_table;
+
+        if (!RVAL2CBOOL(rb_obj_is_kind_of(rb_group_keys, rb_cArray)))
+            rb_raise(rb_eArgError, ":group_keys should be an array of key: <%s>",
+                     rb_grn_inspect(rb_group_keys));
+
+        definition.n_group_keys = RARRAY_LEN(rb_group_keys);
+        definition.group_keys = ALLOCA_N(grn_table_sort_key,
+                                         definition.n_group_keys);
+        rb_table = GRNOBJECT2RVAL(Qnil, context, table, GRN_FALSE);
+        rb_grn_table_sort_keys_fill(context,
+                                    definition.group_keys,
+                                    definition.n_group_keys,
+                                    rb_group_keys,
+                                    rb_table);
+    }
+
     rb_builder = rb_grn_record_expression_builder_new(rb_table, Qnil);
     rb_window_function_call =
         rb_grn_record_expression_builder_build(rb_builder);

diff --git a/test/test-data-column.rb b/test/test-data-column.rb
@@ -1,4 +1,4 @@
-# Copyright (C) 2016  Kouhei Sutou <kou@clear-code.com>
+# Copyright (C) 2016-2017  Kouhei Sutou <kou@clear-code.com>
 #
 # This library is free software; you can redistribute it and/or
 # modify it under the terms of the GNU Lesser General Public
@@ -49,5 +49,47 @@ def test_no_argument
                    ],
                    comments.collect {|comment| [comment.id, comment.nth]})
     end
+
+    def test_group_keys
+      Groonga::Schema.define do |schema|
+        schema.create_table("Comments") do |table|
+          table.uint32("nth")
+          table.short_text("category")
+        end
+      end
+      comments = Groonga["Comments"]
+      nth = Groonga["Comments.nth"]
+
+      3.times do
+        comments.add(:category => "a")
+      end
+      3.times do
+        comments.add(:category => "b")
+      end
+
+      options = {
+        :sort_keys => [["_id", "desc"]],
+        :group_keys => ["category"],
+      }
+      nth.apply_window_function(options) do |record|
+        record.call("record_number")
+      end
+      values = comments.collect do |comment|
+        [
+          comment.id,
+          comment.nth,
+          comment.category,
+        ]
+      end
+      assert_equal([
+                     [1,  3, "a"],
+                     [2,  2, "a"],
+                     [3,  1, "a"],
+                     [4,  3, "b"],
+                     [5,  2, "b"],
+                     [6,  1, "b"],
+                   ],
+                   values)
+    end
   end
 end