Support IN by integers

GitHub: fix #21 Reported by yongxianggao-chanjet. Thanks!!!
pgroonga · Sep 4, 2016 · 4fe5028 · 4fe5028
1 parent 1c8e7fb
commit 4fe5028
Show file tree

Hide file tree

Showing 7 changed files with 279 additions and 5 deletions.
diff --git a/expected/compare/integer/single/in/bitmapscan.out b/expected/compare/integer/single/in/bitmapscan.out
@@ -0,0 +1,29 @@
+CREATE TABLE ids (
+  id integer
+);
+INSERT INTO ids VALUES (2);
+INSERT INTO ids VALUES (7);
+INSERT INTO ids VALUES (6);
+INSERT INTO ids VALUES (4);
+INSERT INTO ids VALUES (5);
+INSERT INTO ids VALUES (8);
+INSERT INTO ids VALUES (1);
+INSERT INTO ids VALUES (10);
+INSERT INTO ids VALUES (3);
+INSERT INTO ids VALUES (9);
+CREATE INDEX pgroonga_index ON ids USING pgroonga (id);
+SET enable_seqscan = off;
+SET enable_indexscan = off;
+SET enable_bitmapscan = on;
+SELECT id
+  FROM ids
+ WHERE id IN (6, 1, 7)
+ ORDER BY id ASC;
+ id 
+----
+  1
+  6
+  7
+(3 rows)
+
+DROP TABLE ids;
diff --git a/expected/compare/integer/single/in/indexscan.out b/expected/compare/integer/single/in/indexscan.out
@@ -0,0 +1,29 @@
+CREATE TABLE ids (
+  id integer
+);
+INSERT INTO ids VALUES (2);
+INSERT INTO ids VALUES (7);
+INSERT INTO ids VALUES (6);
+INSERT INTO ids VALUES (4);
+INSERT INTO ids VALUES (5);
+INSERT INTO ids VALUES (8);
+INSERT INTO ids VALUES (1);
+INSERT INTO ids VALUES (10);
+INSERT INTO ids VALUES (3);
+INSERT INTO ids VALUES (9);
+CREATE INDEX pgroonga_index ON ids USING pgroonga (id);
+SET enable_seqscan = off;
+SET enable_indexscan = on;
+SET enable_bitmapscan = off;
+SELECT id
+  FROM ids
+ WHERE id IN (6, 1, 7)
+ ORDER BY id ASC;
+ id 
+----
+  1
+  6
+  7
+(3 rows)
+
+DROP TABLE ids;
diff --git a/expected/compare/integer/single/in/seqscan.out b/expected/compare/integer/single/in/seqscan.out
@@ -0,0 +1,28 @@
+CREATE TABLE ids (
+  id integer
+);
+INSERT INTO ids VALUES (2);
+INSERT INTO ids VALUES (7);
+INSERT INTO ids VALUES (6);
+INSERT INTO ids VALUES (4);
+INSERT INTO ids VALUES (5);
+INSERT INTO ids VALUES (8);
+INSERT INTO ids VALUES (1);
+INSERT INTO ids VALUES (10);
+INSERT INTO ids VALUES (3);
+INSERT INTO ids VALUES (9);
+SET enable_seqscan = on;
+SET enable_indexscan = off;
+SET enable_bitmapscan = off;
+SELECT id
+  FROM ids
+ WHERE id IN (6, 1, 7)
+ ORDER BY id ASC;
+ id 
+----
+  1
+  6
+  7
+(3 rows)
+
+DROP TABLE ids;
diff --git a/sql/compare/integer/single/in/bitmapscan.sql b/sql/compare/integer/single/in/bitmapscan.sql
@@ -0,0 +1,27 @@
+CREATE TABLE ids (
+  id integer
+);
+
+INSERT INTO ids VALUES (2);
+INSERT INTO ids VALUES (7);
+INSERT INTO ids VALUES (6);
+INSERT INTO ids VALUES (4);
+INSERT INTO ids VALUES (5);
+INSERT INTO ids VALUES (8);
+INSERT INTO ids VALUES (1);
+INSERT INTO ids VALUES (10);
+INSERT INTO ids VALUES (3);
+INSERT INTO ids VALUES (9);
+
+CREATE INDEX pgroonga_index ON ids USING pgroonga (id);
+
+SET enable_seqscan = off;
+SET enable_indexscan = off;
+SET enable_bitmapscan = on;
+
+SELECT id
+  FROM ids
+ WHERE id IN (6, 1, 7)
+ ORDER BY id ASC;
+
+DROP TABLE ids;
diff --git a/sql/compare/integer/single/in/indexscan.sql b/sql/compare/integer/single/in/indexscan.sql
@@ -0,0 +1,27 @@
+CREATE TABLE ids (
+  id integer
+);
+
+INSERT INTO ids VALUES (2);
+INSERT INTO ids VALUES (7);
+INSERT INTO ids VALUES (6);
+INSERT INTO ids VALUES (4);
+INSERT INTO ids VALUES (5);
+INSERT INTO ids VALUES (8);
+INSERT INTO ids VALUES (1);
+INSERT INTO ids VALUES (10);
+INSERT INTO ids VALUES (3);
+INSERT INTO ids VALUES (9);
+
+CREATE INDEX pgroonga_index ON ids USING pgroonga (id);
+
+SET enable_seqscan = off;
+SET enable_indexscan = on;
+SET enable_bitmapscan = off;
+
+SELECT id
+  FROM ids
+ WHERE id IN (6, 1, 7)
+ ORDER BY id ASC;
+
+DROP TABLE ids;
diff --git a/sql/compare/integer/single/in/seqscan.sql b/sql/compare/integer/single/in/seqscan.sql
@@ -0,0 +1,25 @@
+CREATE TABLE ids (
+  id integer
+);
+
+INSERT INTO ids VALUES (2);
+INSERT INTO ids VALUES (7);
+INSERT INTO ids VALUES (6);
+INSERT INTO ids VALUES (4);
+INSERT INTO ids VALUES (5);
+INSERT INTO ids VALUES (8);
+INSERT INTO ids VALUES (1);
+INSERT INTO ids VALUES (10);
+INSERT INTO ids VALUES (3);
+INSERT INTO ids VALUES (9);
+
+SET enable_seqscan = on;
+SET enable_indexscan = off;
+SET enable_bitmapscan = off;
+
+SELECT id
+  FROM ids
+ WHERE id IN (6, 1, 7)
+ ORDER BY id ASC;
+
+DROP TABLE ids;
diff --git a/src/pgroonga.c b/src/pgroonga.c
@@ -2245,6 +2245,72 @@ pgroonga_beginscan(PG_FUNCTION_ARGS)
 	PG_RETURN_POINTER(scan);
 }
 
+static bool
+PGrnSearchIsInCondition(ScanKey key)
+{
+	return (key->sk_flags & SK_SEARCHARRAY &&
+			key->sk_strategy == PGrnEqualStrategyNumber);
+}
+
+static bool
+PGrnSearchBuildConditionIn(PGrnSearchData *data,
+						   ScanKey key,
+						   grn_obj *targetColumn,
+						   Form_pg_attribute attribute)
+{
+	grn_id domain;
+	unsigned char flags = 0;
+	ArrayType *values;
+	int i, n;
+
+	domain = PGrnPGTypeToGrnType(attribute->atttypid, &flags);
+	grn_obj_reinit(ctx, &(buffers->general), domain, flags);
+	values = DatumGetArrayTypeP(key->sk_argument);
+	n = ARR_DIMS(values)[0];
+
+	grn_expr_append_obj(ctx, data->expression,
+						PGrnLookup("in_values", ERROR),
+						GRN_OP_PUSH,
+						1);
+	PGrnCheck("pgroonga: IN: failed to push in_values()");
+	grn_expr_append_obj(ctx, data->expression,
+						targetColumn,
+						GRN_OP_PUSH,
+						1);
+	PGrnCheck("pgroonga: IN: failed to push target column");
+	grn_expr_append_op(ctx, data->expression, GRN_OP_GET_VALUE, 1);
+	PGrnCheck("pgroonga: IN: failed to push GET_VALUE");
+
+	for (i = 1; i <= n; i++)
+	{
+		Datum valueDatum;
+		bool isNULL;
+
+		valueDatum = array_ref(values, 1, &i, -1,
+							   attribute->attlen,
+							   attribute->attbyval,
+							   attribute->attalign,
+							   &isNULL);
+		if (isNULL)
+			return false;
+
+		PGrnConvertFromData(valueDatum,
+							attribute->atttypid,
+							&(buffers->general));
+		grn_expr_append_const(ctx,
+							  data->expression,
+							  &(buffers->general),
+							  GRN_OP_PUSH,
+							  1);
+		PGrnCheck("pgroonga: IN: failed to push a value");
+	}
+
+	grn_expr_append_op(ctx, data->expression, GRN_OP_CALL, 2 + (n - 1));
+	PGrnCheck("pgroonga: IN: failed to push CALL");
+
+	return true;
+}
+
 static void
 PGrnSearchBuildConditionLikeMatchFlush(grn_obj *expression,
 									   grn_obj *targetColumn,
@@ -2570,6 +2636,9 @@ PGrnSearchBuildCondition(Relation index,
 	targetColumn = PGrnLookupColumn(data->sourcesTable, targetColumnName, ERROR);
 	GRN_PTR_PUT(ctx, &(data->targetColumns), targetColumn);
 
+	if (PGrnSearchIsInCondition(key))
+		return PGrnSearchBuildConditionIn(data, key, targetColumn, attribute);
+
 	if (PGrnAttributeIsJSONB(attribute->atttypid))
 		return PGrnJSONBBuildSearchCondition(data, key, targetColumn);
 
@@ -2845,7 +2914,44 @@ PGrnSearch(IndexScanDesc scan)
 static void
 PGrnSort(IndexScanDesc scan)
 {
-	/* TODO */
+	PGrnScanOpaque so = (PGrnScanOpaque) scan->opaque;
+	ScanKey key;
+	TupleDesc desc;
+	Form_pg_attribute attribute;
+	const char *targetColumnName;
+	grn_table_sort_key sort_key;
+
+	if (!so->searched)
+		return;
+
+	if (scan->numberOfKeys != 1)
+		return;
+
+	key = &(scan->keyData[0]);
+	if (!PGrnSearchIsInCondition(key))
+		return;
+
+	so->sorted = grn_table_create(ctx, NULL, 0, NULL,
+								  GRN_OBJ_TABLE_NO_KEY,
+								  NULL, so->searched);
+
+	desc = RelationGetDescr(scan->indexRelation);
+	attribute = desc->attrs[key->sk_attno - 1];
+	targetColumnName = attribute->attname.data;
+	sort_key.key = grn_obj_column(ctx, so->searched,
+								  targetColumnName,
+								  strlen(targetColumnName));
+
+	sort_key.flags = GRN_TABLE_SORT_ASC;
+	sort_key.offset = 0;
+	grn_table_sort(ctx,
+				   so->searched,
+				   0,
+				   -1,
+				   so->sorted,
+				   &sort_key,
+				   1);
+	grn_obj_close(ctx, sort_key.key);
 }
 
 static void
@@ -3122,7 +3228,9 @@ PGrnIsRangeSearchable(IndexScanDesc scan)
 }
 
 static void
-PGrnEnsureCursorOpened(IndexScanDesc scan, ScanDirection dir)
+PGrnEnsureCursorOpened(IndexScanDesc scan,
+					   ScanDirection dir,
+					   bool needSort)
 {
 	PGrnScanOpaque so = (PGrnScanOpaque) scan->opaque;
 
@@ -3156,7 +3264,8 @@ PGrnEnsureCursorOpened(IndexScanDesc scan, ScanDirection dir)
 	else
 	{
 		PGrnSearch(scan);
-		PGrnSort(scan);
+		if (needSort)
+			PGrnSort(scan);
 		PGrnOpenTableCursor(scan, dir);
 	}
 }
@@ -3221,7 +3330,7 @@ pgroonga_gettuple_raw(IndexScanDesc scan,
 {
 	PGrnScanOpaque so = (PGrnScanOpaque) scan->opaque;
 
-	PGrnEnsureCursorOpened(scan, direction);
+	PGrnEnsureCursorOpened(scan, direction, true);
 
 	if (scan->kill_prior_tuple && so->currentID != GRN_ID_NIL)
 	{
@@ -3284,7 +3393,7 @@ pgroonga_getbitmap_raw(IndexScanDesc scan,
 	PGrnScanOpaque so = (PGrnScanOpaque) scan->opaque;
 	int64 nRecords = 0;
 
-	PGrnEnsureCursorOpened(scan, ForwardScanDirection);
+	PGrnEnsureCursorOpened(scan, ForwardScanDirection, false);
 
 	if (so->indexCursor)
 	{