feat: support simple "_row_key BETWEEN xxx AND yyy" query

datafusion-contrib · Mar 9, 2022 · 4c9760b · 4c9760b
1 parent d36b3e0
commit 4c9760b
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -29,7 +29,7 @@ ctx.sql("SELECT \"_row_key\", pressure, \"_timestamp\" FROM weather_balloons whe
 ### SQL
 - ✅ select by `"_row_key" =`
 - ✅ select by `"_row_key" IN`
-- [ ] select by `"_row_key" BETWEEN`
+- ✅ select by `"_row_key" BETWEEN`
 - [ ] select by composite row keys (via `table_partition_cols` and `table_partition_separator`)
 - [ ] Projection pushdown
 - [ ] Predicate push down ([Value range](https://cloud.google.com/bigtable/docs/using-filters#value-range))

diff --git a/src/datasource.rs b/src/datasource.rs
@@ -445,6 +445,18 @@ mod tests {
             "+-------------------------------+----------+-------------------------+",
         ];
         assert_batches_eq!(expected, &batches);
+
+        batches = ctx.sql("SELECT \"_row_key\", pressure, \"_timestamp\" FROM weather_balloons where \"_row_key\" BETWEEN 'us-west2#3698#2021-03-05-1200' AND 'us-west2#3698#2021-03-05-1202' ORDER BY \"_row_key\"").await?.collect().await?;
+        expected = vec![
+            "+-------------------------------+----------+-------------------------+",
+            "| _row_key                      | pressure | _timestamp              |",
+            "+-------------------------------+----------+-------------------------+",
+            "| us-west2#3698#2021-03-05-1200 | 94558    | 2021-03-05 12:00:05.100 |",
+            "| us-west2#3698#2021-03-05-1201 | 94122    | 2021-03-05 12:01:05.200 |",
+            "| us-west2#3698#2021-03-05-1202 | 95992    | 2021-03-05 12:02:05.300 |",
+            "+-------------------------------+----------+-------------------------+",
+        ];
+        assert_batches_eq!(expected, &batches);
         Ok(())
     }
 }
diff --git a/src/datasource/composer.rs b/src/datasource/composer.rs
@@ -81,6 +81,40 @@ pub fn compose(
                 }
                 _ => (),
             },
+            Expr::Between {
+                expr,
+                negated,
+                low,
+                high,
+            } => match expr.as_ref() {
+                Expr::Column(col) => {
+                    if datasource.table_partition_cols.contains(&col.name) {
+                        if negated.to_owned() {
+                            return Err(DataFusionError::Execution(
+                                "_row_key filter: NOT IN is not supported".to_owned(),
+                            ));
+                        }
+                        match low.as_ref() {
+                            Expr::Literal(ScalarValue::Utf8(Some(low_key))) => {
+                                match high.as_ref() {
+                                    Expr::Literal(ScalarValue::Utf8(Some(high_key))) => row_ranges
+                                        .push(RowRange {
+                                            start_key: Some(StartKey::StartKeyClosed(
+                                                low_key.clone().into_bytes(),
+                                            )),
+                                            end_key: Some(EndKey::EndKeyClosed(
+                                                high_key.clone().into_bytes(),
+                                            )),
+                                        }),
+                                    _ => (),
+                                }
+                            }
+                            _ => (),
+                        }
+                    }
+                }
+                _ => (),
+            },
             _ => (),
         }
     }