You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by si...@apache.org on 2022/11/08 20:28:44 UTC

[pinot] branch master updated: Add Support for IP Address Function (#9501)

This is an automated email from the ASF dual-hosted git repository.

siddteotia pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git


The following commit(s) were added to refs/heads/master by this push:
     new e3f2835686 Add Support for IP Address Function (#9501)
e3f2835686 is described below

commit e3f2835686e3619ee0febae4ba19823125440e4a
Author: Sabrina Zhao <yi...@linkedin.com>
AuthorDate: Tue Nov 8 12:28:36 2022 -0800

    Add Support for IP Address Function (#9501)
    
    * Added subnet functions/unit tests, need to add more tests
    
    * cleaned up code
    
    * style checker
    
    * Added more unit tests, handled boolean type for literal
    
    * added comments on implemenation detail
    
    * helper fns for min/max
    
    * clearer exception
    
    * added e2e tests on sample data, before lib
    
    * added lib to pom
    
    * using lib func
    
    * handle invalid ipprefix
    
    * handled validation using lib, added more tests to trigger func in filter
    
    * rebase, trigger test
    
    * fixed dependency
    
    * requestutils bool handled
    
    * got rid off inhouse
    
    * added more tests that trigger exception
    
    * style
    
    * dependency
    
    * resolve dependency
    
    * new test cases
---
 .../LiteralOnlyBrokerRequestTest.java              |  85 ++++++++
 pinot-common/pom.xml                               |   4 +
 .../common/function/scalar/IpAddressFunctions.java |  80 +++++++
 .../pinot/sql/parsers/CalciteSqlCompilerTest.java  | 140 ++++++++++++
 .../pinot-common-jdk8/pom.xml                      |   4 +
 .../pinot/queries/IsSubnetOfQueriesTest.java       | 236 +++++++++++++++++++++
 pom.xml                                            |   6 +
 7 files changed, 555 insertions(+)

diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java
index 0aafb16188..28c4ef5466 100644
--- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java
@@ -142,6 +142,22 @@ public class LiteralOnlyBrokerRequestTest {
         .compileToPinotQuery("SELECT count(*) from myTable where regexpReplace(col1, \"b(..)\", \"X$1Y\", 10 , 1, "
             + "\"m\")  = "
             + "\"fooXarYXazY\"")));
+    Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(
+        CalciteSqlParser.compileToPinotQuery("select isSubnetOf('1.2.3.128/0', '192.168.5.1') from mytable")));
+    Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery(
+        "select isSubnetOf('1.2.3.128/0', rtrim('192.168.5.1      ')) from mytable")));
+    Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery(
+        "select isSubnetOf('123:db8:85a3::8a2e:370:7334/72', '124:db8:85a3::8a2e:370:7334') from mytable")));
+    Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(
+        CalciteSqlParser.compileToPinotQuery("select isSubnetOf('1.2.3.128/0', foo) from mytable")));
+    Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery(
+        "select count(*) from mytable where isSubnetOf('7890:db8:113::8a2e:370:7334/127', ltrim('   "
+            + "7890:db8:113::8a2e:370:7336'))")));
+    Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery(
+        "select count(*) from mytable where isSubnetOf('7890:db8:113::8a2e:370:7334/127', "
+            + "'7890:db8:113::8a2e:370:7336')")));
+    Assert.assertFalse(BaseBrokerRequestHandler.isLiteralOnlyQuery(
+        CalciteSqlParser.compileToPinotQuery("select count(*) from mytable where isSubnetOf(foo, bar)")));
   }
 
   @Test
@@ -157,6 +173,8 @@ public class LiteralOnlyBrokerRequestTest {
         "SELECT toUtf8('hello!') AS encoded, " + "fromUtf8(toUtf8('hello!')) AS decoded")));
     Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery(
         "SELECT toBase64(toUtf8('hello!')) AS encoded, " + "fromBase64('aGVsbG8h') AS decoded")));
+    Assert.assertTrue(BaseBrokerRequestHandler.isLiteralOnlyQuery(CalciteSqlParser.compileToPinotQuery(
+        "select isSubnetOf('1.2.3.128/0', '192.168.5.1') AS booleanCol from mytable")));
   }
 
   @Test
@@ -324,6 +342,73 @@ public class LiteralOnlyBrokerRequestTest {
     brokerResponse = requestHandler.handleRequest(request, null, requestStats);
     Assert.assertTrue(
         brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException"));
+
+    request = JsonUtils.stringToJsonNode(
+        "{\"sql\":\"SELECT isSubnetOf('2001:db8:85a3::8a2e:370:7334/62', '2001:0db8:85a3:0003:ffff:ffff:ffff:ffff')"
+            + " as booleanCol\"}");
+    requestStats = Tracing.getTracer().createRequestScope();
+    brokerResponse = requestHandler.handleRequest(request, null, requestStats);
+    resultTable = brokerResponse.getResultTable();
+    dataSchema = resultTable.getDataSchema();
+    rows = resultTable.getRows();
+    Assert.assertEquals(dataSchema.getColumnName(0), "booleanCol");
+    Assert.assertEquals(dataSchema.getColumnDataType(0), DataSchema.ColumnDataType.BOOLEAN);
+    Assert.assertEquals(rows.size(), 1);
+    Assert.assertEquals(rows.get(0).length, 1);
+    Assert.assertTrue((boolean) rows.get(0)[0]);
+    Assert.assertEquals(brokerResponse.getTotalDocs(), 0);
+
+    // first argument must be in prefix format
+    request = JsonUtils.stringToJsonNode(
+        "{\"sql\":\"SELECT isSubnetOf('2001:db8:85a3::8a2e:370:7334', '2001:0db8:85a3:0003:ffff:ffff:ffff:ffff') as"
+            + " booleanCol\"}");
+    requestStats = Tracing.getTracer().createRequestScope();
+    brokerResponse = requestHandler.handleRequest(request, null, requestStats);
+    Assert.assertTrue(
+        brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException"));
+
+    // first argument must be in prefix format
+    request = JsonUtils.stringToJsonNode(
+        "{\"sql\":\"SELECT isSubnetOf('105.25.245.115', '105.25.245.115') as" + " booleanCol\"}");
+    requestStats = Tracing.getTracer().createRequestScope();
+    brokerResponse = requestHandler.handleRequest(request, null, requestStats);
+    Assert.assertTrue(
+        brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException"));
+
+    // second argument should not be a prefix
+    request = JsonUtils.stringToJsonNode(
+        "{\"sql\":\"SELECT isSubnetOf('10.3.168.0/22', '3.175.47.239/26') as" + " booleanCol\"}");
+    requestStats = Tracing.getTracer().createRequestScope();
+    brokerResponse = requestHandler.handleRequest(request, null, requestStats);
+    Assert.assertTrue(
+        brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException"));
+
+    // second argument should not be a prefix
+    request = JsonUtils.stringToJsonNode(
+        "{\"sql\":\"SELECT isSubnetOf('5f3f:bfdb:1bbe:a824:6bf9:0fbb:d358:1889/64', "
+            + "'4275:386f:b2b5:0664:04aa:d7bd:0589:6909/64') as"
+            + " booleanCol\"}");
+    requestStats = Tracing.getTracer().createRequestScope();
+    brokerResponse = requestHandler.handleRequest(request, null, requestStats);
+    Assert.assertTrue(
+        brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException"));
+
+    // invalid prefix length
+    request = JsonUtils.stringToJsonNode(
+        "{\"sql\":\"SELECT isSubnetOf('2001:4801:7825:103:be76:4eff::/129', '2001:4801:7825:103:be76:4eff::') as"
+            + " booleanCol\"}");
+    requestStats = Tracing.getTracer().createRequestScope();
+    brokerResponse = requestHandler.handleRequest(request, null, requestStats);
+    Assert.assertTrue(
+        brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException"));
+
+    // invalid prefix length
+    request = JsonUtils.stringToJsonNode(
+        "{\"sql\":\"SELECT isSubnetOf('170.189.0.175/33', '170.189.0.175') as" + " booleanCol\"}");
+    requestStats = Tracing.getTracer().createRequestScope();
+    brokerResponse = requestHandler.handleRequest(request, null, requestStats);
+    Assert.assertTrue(
+        brokerResponse.getProcessingExceptions().get(0).getMessage().contains("IllegalArgumentException"));
   }
 
   /** Tests for EXPLAIN PLAN for literal only queries. */
diff --git a/pinot-common/pom.xml b/pinot-common/pom.xml
index 7af3c3fdea..cba89837e2 100644
--- a/pinot-common/pom.xml
+++ b/pinot-common/pom.xml
@@ -446,6 +446,10 @@
       <artifactId>jbcrypt</artifactId>
       <version>0.4</version>
     </dependency>
+    <dependency>
+      <groupId>com.github.seancfoley</groupId>
+      <artifactId>ipaddress</artifactId>
+    </dependency>
   </dependencies>
   <profiles>
     <profile>
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/IpAddressFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/IpAddressFunctions.java
new file mode 100644
index 0000000000..83b73056d8
--- /dev/null
+++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/IpAddressFunctions.java
@@ -0,0 +1,80 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.pinot.common.function.scalar;
+
+import inet.ipaddr.AddressStringException;
+import inet.ipaddr.IPAddress;
+import inet.ipaddr.IPAddressString;
+import inet.ipaddr.PrefixLenException;
+import org.apache.pinot.spi.annotations.ScalarFunction;
+
+
+/**
+ * Inbuilt IP related transform functions
+ *
+ * Functions added:
+ * isSubnetOf(String ipPrefix, String ipAddress) --> boolean
+ *
+ * Functions to add:
+ * ipPrefix(String ipAddress, int prefixBits) -> String ipPrefix
+ * ipSubnetMin(String ipPrefix) -> String ipMin
+ * ipSubnetMax(String ipPrefix) -> String ipMax
+ */
+public class IpAddressFunctions {
+
+  private IpAddressFunctions() {
+  }
+
+  /**
+   * Validates IP prefix prefixStr and returns IPAddress if validated
+   */
+  private static IPAddress getPrefix(String prefixStr) {
+    IPAddress prefixAddr = getAddress(prefixStr);
+    if (!prefixAddr.isPrefixed()) {
+      throw new IllegalArgumentException("IP Address " + prefixStr + " should be prefixed.");
+    }
+    try {
+      return prefixAddr.toPrefixBlock();
+    } catch (PrefixLenException e) {
+      throw e;
+    }
+  }
+
+  /**
+   * Validates IP address ipString and returns IPAddress if validated
+   */
+  private static IPAddress getAddress(String ipString) {
+    try {
+      return new IPAddressString(ipString).toAddress();
+    } catch (AddressStringException e) {
+      throw new IllegalArgumentException("Invalid IP Address format for " + ipString);
+    }
+  }
+
+  @ScalarFunction
+  public static boolean isSubnetOf(String ipPrefix, String ipAddress) {
+    IPAddress prefix = getPrefix(ipPrefix);
+    IPAddress ip = getAddress(ipAddress);
+    if (ip.isPrefixed()) {
+      throw new IllegalArgumentException("IP Address " + ipAddress + " should not be prefixed.");
+    }
+    return prefix.contains(ip);
+  }
+}
diff --git a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
index d4afdc71b3..5ffe1b29be 100644
--- a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
@@ -360,6 +360,39 @@ public class CalciteSqlCompilerTest {
       Assert.assertEquals(eqOperands.get(0).getIdentifier().getName(), "d");
       Assert.assertEquals(eqOperands.get(1).getLiteral(), Literal.boolValue(true));
     }
+
+    {
+      PinotQuery pinotQuery =
+          CalciteSqlParser.compileToPinotQuery("select * from vegetable where isSubnetOf('192.168.0.1/24', foo)");
+      Function func = pinotQuery.getFilterExpression().getFunctionCall();
+      Assert.assertEquals(func.getOperator(), FilterKind.EQUALS.name());
+      List<Expression> operands = func.getOperands();
+      Assert.assertEquals(operands.size(), 2);
+      Assert.assertEquals(operands.get(0).getFunctionCall().getOperator(), "issubnetof");
+      Assert.assertEquals(operands.get(1).getLiteral(), Literal.boolValue(true));
+    }
+
+    {
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(
+          "select * from vegetable where isSubnetOf('192.168.0.1/24', foo)=true AND isSubnetOf('192.168.0.1/24', "
+              + "foo)");
+      Function func = pinotQuery.getFilterExpression().getFunctionCall();
+      Assert.assertEquals(func.getOperator(), FilterKind.AND.name());
+      List<Expression> operands = func.getOperands();
+      Assert.assertEquals(operands.size(), 2);
+      Assert.assertEquals(operands.get(0).getFunctionCall().getOperator(), FilterKind.EQUALS.name());
+      Assert.assertEquals(operands.get(1).getFunctionCall().getOperator(), FilterKind.EQUALS.name());
+
+      List<Expression> lhs = operands.get(0).getFunctionCall().getOperands();
+      Assert.assertEquals(lhs.size(), 2);
+      Assert.assertEquals(lhs.get(0).getFunctionCall().getOperator(), "issubnetof");
+      Assert.assertEquals(lhs.get(1).getLiteral(), Literal.boolValue(true));
+
+      List<Expression> rhs = operands.get(1).getFunctionCall().getOperands();
+      Assert.assertEquals(rhs.size(), 2);
+      Assert.assertEquals(rhs.get(0).getFunctionCall().getOperator(), "issubnetof");
+      Assert.assertEquals(rhs.get(1).getLiteral(), Literal.boolValue(true));
+    }
   }
 
   @Test
@@ -1981,6 +2014,113 @@ public class CalciteSqlCompilerTest {
     }
     Assert.assertNotNull(expectedError);
     Assert.assertTrue(expectedError instanceof SqlCompilationException);
+
+    query = "select isSubnetOf('192.168.0.1/24', '192.168.0.225') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    boolean result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('192.168.0.1/24', '192.168.0.1') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('130.191.23.32/27', '130.191.23.40') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('130.191.23.32/26', '130.192.23.33') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertFalse(result);
+
+    query = "select isSubnetOf('153.87.199.160/28', '153.87.199.166') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('2001:4800:7825:103::/64', '2001:4800:7825:103::2050') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('130.191.23.32/26', '130.191.23.33') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('2001:4801:7825:103:be76:4efe::/96', '2001:4801:7825:103:be76:4efe::e15') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('122.152.15.0/26', '122.152.15.28') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('96.141.228.254/26', '96.141.228.254') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('3.175.47.128/26', '3.175.48.178') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertFalse(result);
+
+    query = "select isSubnetOf('192.168.0.1/24', '192.168.0.0') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('10.3.168.0/22', '10.3.168.123') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('10.3.168.0/22', '10.3.171.255') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('10.3.168.0/22', '1.2.3.1') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertFalse(result);
+
+    query = "select isSubnetOf('1.2.3.128/1', '127.255.255.255') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('1.2.3.128/0', '192.168.5.1') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query =
+        "select isSubnetOf('2001:db8:85a3::8a2e:370:7334/62', '2001:0db8:85a3:0003:ffff:ffff:ffff:ffff') from "
+            + "mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
+
+    query = "select isSubnetOf('123:db8:85a3::8a2e:370:7334/72', '124:db8:85a3::8a2e:370:7334') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertFalse(result);
+
+    query = "select isSubnetOf('7890:db8:113::8a2e:370:7334/127', '7890:db8:113::8a2e:370:7336') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertFalse(result);
+
+    query = "select isSubnetOf('7890:db8:113::8a2e:370:7334/127', '7890:db8:113::8a2e:370:7335') from mytable";
+    pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
+    result = pinotQuery.getSelectList().get(0).getLiteral().getBoolValue();
+    Assert.assertTrue(result);
   }
 
   @Test
diff --git a/pinot-connectors/prestodb-pinot-dependencies/pinot-common-jdk8/pom.xml b/pinot-connectors/prestodb-pinot-dependencies/pinot-common-jdk8/pom.xml
index 350bd5af51..b1e89f3c80 100644
--- a/pinot-connectors/prestodb-pinot-dependencies/pinot-common-jdk8/pom.xml
+++ b/pinot-connectors/prestodb-pinot-dependencies/pinot-common-jdk8/pom.xml
@@ -430,5 +430,9 @@
       <artifactId>jbcrypt</artifactId>
       <version>0.4</version>
     </dependency>
+    <dependency>
+      <groupId>com.github.seancfoley</groupId>
+      <artifactId>ipaddress</artifactId>
+    </dependency>
   </dependencies>
 </project>
diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/IsSubnetOfQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/IsSubnetOfQueriesTest.java
new file mode 100644
index 0000000000..6337147b0a
--- /dev/null
+++ b/pinot-core/src/test/java/org/apache/pinot/queries/IsSubnetOfQueriesTest.java
@@ -0,0 +1,236 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.queries;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.common.response.broker.BrokerResponseNative;
+import org.apache.pinot.common.response.broker.ResultTable;
+import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.spi.ImmutableSegment;
+import org.apache.pinot.segment.spi.IndexSegment;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+/**
+ * Queries test for Subnet Containment for IP Address queries.
+ */
+public class IsSubnetOfQueriesTest extends BaseQueriesTest {
+  private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "IsSubnetOfQueriesTest");
+  private static final String RAW_TABLE_NAME = "testTable";
+  private static final String SEGMENT_NAME = "testSegment";
+
+  private static final String IPv4_PREFIX_COLUMN_STRING = "IPv4prefixColumn";
+  private static final String IPv6_PREFIX_COLUMN_STRING = "IPv6prefixColumn";
+  private static final String IPv4_ADDRESS_COLUMN = "IPv4AddressColumn";
+  private static final String IPv6_ADDRESS_COLUMN = "IPv6AddressColumn";
+  // columns storing hardcoded expected results:
+  // isSubnetOf(IPv4prefixColumn, IPv4AddressColumn), isSubnetOf(IPv6prefixColumn, IPv6AddressColumn)
+  private static final String IPv4_CONTAINS_COLUMN = "IPv4ContainsColumn";
+  private static final String IPv6_CONTAINS_COLUMN = "IPv6ContainsColumn";
+
+  private static final String DEFAULT_IPv4_PREFIX = "1.2.3.128/26";
+  private static final String DEFAULT_IPv6_PREFIX = "64:fa9b::17/64";
+  private static final String DEFAULT_IPv4_ADDRESS = "1.2.3.129";
+  private static final String DEFAULT_IPv6_ADDRESS = "64:ffff::17";
+  private static final boolean DEFAULT_IPv4_CONTAINS = true;
+  private static final boolean DEFAULT_IPv6_CONTAINS = false;
+
+  private static final Schema SCHEMA =
+      new Schema.SchemaBuilder().addSingleValueDimension(IPv4_PREFIX_COLUMN_STRING, FieldSpec.DataType.STRING)
+          .addSingleValueDimension(IPv6_PREFIX_COLUMN_STRING, FieldSpec.DataType.STRING)
+          .addSingleValueDimension(IPv4_ADDRESS_COLUMN, FieldSpec.DataType.STRING)
+          .addSingleValueDimension(IPv6_ADDRESS_COLUMN, FieldSpec.DataType.STRING)
+          .addSingleValueDimension(IPv4_CONTAINS_COLUMN, FieldSpec.DataType.BOOLEAN)
+          .addSingleValueDimension(IPv6_CONTAINS_COLUMN, FieldSpec.DataType.BOOLEAN).build();
+  private static final TableConfig TABLE_CONFIG =
+      new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build();
+  private IndexSegment _indexSegment;
+  private List<IndexSegment> _indexSegments;
+  private long _expectedNumberIpv4Contains = 0L;
+  private long _expectedNumberIpv6Contains = 0L;
+
+  @BeforeClass
+  public void setUp()
+      throws Exception {
+    FileUtils.deleteDirectory(INDEX_DIR);
+
+    List<GenericRow> records = new ArrayList<>();
+
+    // add IPv4 test cases
+    addIPv4Row(records, "105.25.245.115/27", "105.25.245.98", true);
+    addIPv4Row(records, "122.152.0.204/28", "122.152.0.198", true);
+    addIPv4Row(records, "130.191.23.32/26", "130.191.23.33", true);
+    addIPv4Row(records, "122.152.15.0/26", "122.152.15.28", true);
+    addIPv4Row(records, "96.141.228.254/26", "96.141.228.254", true);
+    addIPv4Row(records, "3.175.47.128/26", "3.175.47.178", true);
+
+    addIPv4Row(records, "105.25.245.115/27", "105.25.245.0", false);
+    addIPv4Row(records, "122.152.0.204/28", "122.152.0.254", false);
+    addIPv4Row(records, "130.191.23.32/26", "130.192.23.33", false);
+    addIPv4Row(records, "122.152.15.0/26", "122.152.0.63", false);
+    addIPv4Row(records, "96.141.228.254/26", "96.141.227.254", false);
+    addIPv4Row(records, "3.175.47.128/26", "3.175.48.178", false);
+
+    addIPv4Row(records, "10.3.168.0/22", "1.2.3.1", false);
+    addIPv4Row(records, "1.2.3.128/26", "1.2.5.1", false);
+    addIPv4Row(records, "1.2.3.128/26", "1.1.3.1", false);
+
+    // add IPv6 test cases
+    addIPv6Row(records, "2001:4800:7825:103::/64", "2001:4800:7825:103::2050", true);
+    addIPv6Row(records, "2001:4801:7825:103:be76:4efe::/96", "2001:4801:7825:103:be76:4efe::e15", true);
+    addIPv6Row(records, "2001:db8:85a3::8a2e:370:7334/62", "2001:0db8:85a3:0003:ffff:ffff:ffff:ffff", true);
+    addIPv6Row(records, "7890:db8:113::8a2e:370:7334/127", "7890:db8:113::8a2e:370:7336", false);
+    addIPv6Row(records, "64:ff9b::17/64", "64:ffff::17", false);
+    addIPv6Row(records, "123:db8:85a3::8a2e:370:7334/72", "124:db8:85a3::8a2e:370:7334", false);
+
+    SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA);
+    segmentGeneratorConfig.setTableName(RAW_TABLE_NAME);
+    segmentGeneratorConfig.setSegmentName(SEGMENT_NAME);
+    segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath());
+
+    SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+    driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records));
+    driver.build();
+
+    ImmutableSegment immutableSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap);
+    _indexSegment = immutableSegment;
+    _indexSegments = Arrays.asList(immutableSegment, immutableSegment);
+  }
+
+  @Test
+  public void testIsSubnetOf() {
+    // called in select
+    String query = String.format(
+        "select isSubnetOf(%s, %s) as IPv4Result, isSubnetOf(%s, %s) as IPv6Result, %s, %s from %s limit 100",
+        IPv4_PREFIX_COLUMN_STRING, IPv4_ADDRESS_COLUMN, IPv6_PREFIX_COLUMN_STRING, IPv6_ADDRESS_COLUMN,
+        IPv4_CONTAINS_COLUMN, IPv6_CONTAINS_COLUMN, RAW_TABLE_NAME);
+    BrokerResponseNative brokerResponse = getBrokerResponse(query);
+    ResultTable resultTable = brokerResponse.getResultTable();
+    DataSchema dataSchema = resultTable.getDataSchema();
+    assertEquals(dataSchema.getColumnDataTypes(), new DataSchema.ColumnDataType[]{
+        DataSchema.ColumnDataType.BOOLEAN, DataSchema.ColumnDataType.BOOLEAN, DataSchema.ColumnDataType.BOOLEAN,
+        DataSchema.ColumnDataType.BOOLEAN
+    });
+    List<Object[]> rows = resultTable.getRows();
+    for (int i = 0; i < rows.size(); i++) {
+      Object[] row = rows.get(i);
+      boolean iPv4Result = (boolean) row[0];
+      boolean iPv6Result = (boolean) row[1];
+      boolean expectedIPv4Result = (boolean) row[2];
+      boolean expectedIPv6Result = (boolean) row[3];
+      assertEquals(iPv4Result, expectedIPv4Result);
+      assertEquals(iPv6Result, expectedIPv6Result);
+    }
+
+    // called in filter
+    query = String.format("select count(*) from %s where isSubnetOf(%s, %s)", RAW_TABLE_NAME, IPv4_PREFIX_COLUMN_STRING,
+        IPv4_ADDRESS_COLUMN);
+    brokerResponse = getBrokerResponse(query);
+    resultTable = brokerResponse.getResultTable();
+    rows = resultTable.getRows();
+    assertEquals(rows.size(), 1);
+    assertEquals(rows.get(0)[0], _expectedNumberIpv4Contains * 4);
+
+    query = String.format("select count(*) from %s where isSubnetOf(%s, %s)", RAW_TABLE_NAME, IPv6_PREFIX_COLUMN_STRING,
+        IPv6_ADDRESS_COLUMN);
+    brokerResponse = getBrokerResponse(query);
+    resultTable = brokerResponse.getResultTable();
+    rows = resultTable.getRows();
+    assertEquals(rows.size(), 1);
+    assertEquals(rows.get(0)[0], _expectedNumberIpv6Contains * 4);
+  }
+
+  private void addIPv4Row(List<GenericRow> records, String prefix, String address, boolean expectedBool) {
+    if (expectedBool) {
+      _expectedNumberIpv4Contains += 1;
+    }
+    if (DEFAULT_IPv6_CONTAINS) {
+      _expectedNumberIpv6Contains += 1;
+    }
+    GenericRow record = new GenericRow();
+    record.putValue(IPv4_PREFIX_COLUMN_STRING, prefix);
+    record.putValue(IPv4_ADDRESS_COLUMN, address);
+    record.putValue(IPv4_CONTAINS_COLUMN, expectedBool);
+
+    record.putValue(IPv6_PREFIX_COLUMN_STRING, DEFAULT_IPv6_PREFIX);
+    record.putValue(IPv6_ADDRESS_COLUMN, DEFAULT_IPv6_ADDRESS);
+    record.putValue(IPv6_CONTAINS_COLUMN, DEFAULT_IPv6_CONTAINS);
+    records.add(record);
+  }
+
+  private void addIPv6Row(List<GenericRow> records, String prefix, String address, boolean expectedBool) {
+    if (expectedBool) {
+      _expectedNumberIpv6Contains += 1;
+    }
+    if (DEFAULT_IPv4_CONTAINS) {
+      _expectedNumberIpv4Contains += 1;
+    }
+    GenericRow record = new GenericRow();
+    record.putValue(IPv6_PREFIX_COLUMN_STRING, prefix);
+    record.putValue(IPv6_ADDRESS_COLUMN, address);
+    record.putValue(IPv6_CONTAINS_COLUMN, expectedBool);
+
+    record.putValue(IPv4_PREFIX_COLUMN_STRING, DEFAULT_IPv4_PREFIX);
+    record.putValue(IPv4_ADDRESS_COLUMN, DEFAULT_IPv4_ADDRESS);
+    record.putValue(IPv4_CONTAINS_COLUMN, DEFAULT_IPv4_CONTAINS);
+    records.add(record);
+  }
+
+  @Override
+  protected String getFilter() {
+    return null;
+  }
+
+  @Override
+  protected IndexSegment getIndexSegment() {
+    return _indexSegment;
+  }
+
+  @Override
+  protected List<IndexSegment> getIndexSegments() {
+    return _indexSegments;
+  }
+
+  @AfterClass
+  public void tearDown()
+      throws IOException {
+    _indexSegment.destroy();
+    FileUtils.deleteDirectory(INDEX_DIR);
+  }
+}
diff --git a/pom.xml b/pom.xml
index 79fb0b283e..0a937a5abd 100644
--- a/pom.xml
+++ b/pom.xml
@@ -1198,6 +1198,12 @@
         <artifactId>h3</artifactId>
         <version>${h3.version}</version>
       </dependency>
+
+      <dependency>
+        <groupId>com.github.seancfoley</groupId>
+        <artifactId>ipaddress</artifactId>
+        <version>5.3.4</version>
+      </dependency>
     </dependencies>
   </dependencyManagement>
   <build>


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org