Merge branch 'main' into SupportOverwriteForSpark

lancedb · Dec 31, 2024 · c6dfb51 · c6dfb51
2 parents c37bfe9 + 11f6e26
commit c6dfb51
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/java/spark/src/main/java/com/lancedb/lance/spark/read/FilterPushDown.java b/java/spark/src/main/java/com/lancedb/lance/spark/read/FilterPushDown.java
@@ -36,6 +36,7 @@
 import java.sql.Date;
 import java.sql.Timestamp;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
 
@@ -89,7 +90,7 @@ public static boolean isFilterSupported(Filter filter) {
     } else if (filter instanceof EqualNullSafe) {
       return false;
     } else if (filter instanceof In) {
-      return false;
+      return true;
     } else if (filter instanceof LessThan) {
       return true;
     } else if (filter instanceof LessThanOrEqual) {
@@ -163,6 +164,13 @@ private static Optional<String> compileFilter(Filter filter) {
       Optional<String> child = compileFilter(f.child());
       if (child.isEmpty()) return child;
       return Optional.of(String.format("NOT (%s)", child.get()));
+    } else if (filter instanceof In) {
+      In in = (In) filter;
+      String values =
+          Arrays.stream(in.values())
+              .map(FilterPushDown::compileValue)
+              .collect(Collectors.joining(","));
+      return Optional.of(String.format("%s IN (%s)", in.attribute(), values));
     }
 
     return Optional.empty();

diff --git a/java/spark/src/test/java/com/lancedb/lance/spark/read/FilterPushDownTest.java b/java/spark/src/test/java/com/lancedb/lance/spark/read/FilterPushDownTest.java
@@ -82,4 +82,26 @@ public void testCompileFiltersToSqlWhereClauseWithEmptyFilters() {
     Optional<String> whereClause = FilterPushDown.compileFiltersToSqlWhereClause(filters);
     assertFalse(whereClause.isPresent());
   }
+
+  @Test
+  public void testIntegerInFilterPushDown() {
+    Object[] values = new Object[2];
+    values[0] = 500;
+    values[1] = 600;
+    Filter[] filters = new Filter[] {new GreaterThan("age", 30), new In("salary", values)};
+    Optional<String> whereClause = FilterPushDown.compileFiltersToSqlWhereClause(filters);
+    assertTrue(whereClause.isPresent());
+    assertEquals("(age > 30) AND (salary IN (500,600))", whereClause.get());
+  }
+
+  @Test
+  public void testStringInFilterPushDown() {
+    Object[] values = new Object[2];
+    values[0] = "500";
+    values[1] = "600";
+    Filter[] filters = new Filter[] {new GreaterThan("age", 30), new In("salary", values)};
+    Optional<String> whereClause = FilterPushDown.compileFiltersToSqlWhereClause(filters);
+    assertTrue(whereClause.isPresent());
+    assertEquals("(age > 30) AND (salary IN ('500','600'))", whereClause.get());
+  }
 }