Fix sub-fields accessing of generated structs (#4683)

* Correct subfield access logical when calling ITEM

Signed-off-by: Yuanchun Shen <[email protected]>

* Add explain and integration tests

Signed-off-by: Yuanchun Shen <[email protected]>

---------

Signed-off-by: Yuanchun Shen <[email protected]>
(cherry picked from commit 3a42c51)

Author: yuancu

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java

@@ -1462,4 +1462,27 @@ public void testTopKThenSortExplain() throws IOException {
                 + "| sort age "
                 + "| fields age"));
   }
+
+  @Test
+  public void testGeoIpPushedInAgg() throws IOException {
+    // This explain IT verifies that externally registered UDF can be properly pushed down
+    assertYamlEqualsIgnoreId(
+        loadExpectedPlan("udf_geoip_in_agg_pushed.yaml"),
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval info = geoip('my-datasource', host) | stats count() by info.city",
+                TEST_INDEX_WEBLOGS)));
+  }
+
+  @Test
+  public void testInternalItemAccessOnStructs() throws IOException {
+    String expected = loadExpectedPlan("access_struct_subfield_with_item.yaml");
+    assertYamlEqualsIgnoreId(
+        expected,
+        explainQueryYaml(
+            String.format(
+                "source=%s | eval info = geoip('dummy-datasource', host) | fields host, info,"
+                    + " info.dummy_sub_field",
+                TEST_INDEX_WEBLOGS)));
+  }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteGeoIpFunctionsIT.java

@@ -55,4 +55,54 @@ public void testGeoIpEnrichmentWithIpFieldAsInput() throws IOException {
         rows("10.0.0.1", Map.of("country", "USA")),
         rows("fd12:2345:6789:1:a1b2:c3d4:e5f6:789a", Map.of("country", "India")));
   }
+
+  @Test
+  public void testGeoIpInAggregation() throws IOException {
+    JSONObject result1 =
+        executeQuery(
+            String.format(
+                "source=%s | where method='POST' | eval info = geoip('%s', host) | eval"
+                    + " date=DATE('2020-12-10') | stats count() by info.city, method, span(date,"
+                    + " 1month) as month",
+                TEST_INDEX_WEBLOGS, DATASOURCE_NAME));
+    verifySchema(
+        result1,
+        schema("count()", "bigint"),
+        schema("month", "date"),
+        schema("info.city", "string"),
+        schema("method", "string"));
+    verifyDataRows(
+        result1,
+        rows(1, "2020-12-01", "Seattle", "POST"),
+        rows(1, "2020-12-01", "Bengaluru", "POST"));
+
+    // This case is pushed down into DSL with scripts
+    JSONObject result2 =
+        executeQuery(
+            String.format(
+                "source=%s | where method='POST' | eval info = geoip('%s', host) | stats count() by"
+                    + " info.city",
+                TEST_INDEX_WEBLOGS, DATASOURCE_NAME));
+    verifySchema(result2, schema("count()", "bigint"), schema("info.city", "string"));
+    verifyDataRows(result2, rows(1, "Seattle"), rows(1, "Bengaluru"));
+  }
+
+  @Test
+  public void testGeoIpEnrichmentAccessingSubField() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | where method='POST' | eval info = geoip('%s', host) | fields host,"
+                    + " info, info.country",
+                TEST_INDEX_WEBLOGS, DATASOURCE_NAME));
+    verifySchema(
+        result, schema("host", "ip"), schema("info", "struct"), schema("info.country", "string"));
+    verifyDataRows(
+        result,
+        rows("10.0.0.1", Map.of("country", "USA", "city", "Seattle"), "USA"),
+        rows(
+            "fd12:2345:6789:1:a1b2:c3d4:e5f6:789a",
+            Map.of("country", "India", "city", "Bengaluru"),
+            "India"));
+  }
 }

integ-test/src/test/resources/expectedOutput/calcite/access_struct_subfield_with_item.yaml

@@ -0,0 +1,8 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(host=[$0], info=[GEOIP('dummy-datasource':VARCHAR, $0)], info.dummy_sub_field=[ITEM(GEOIP('dummy-datasource':VARCHAR, $0), 'dummy_sub_field')])
+        CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_weblogs]])
+  physical: |
+    EnumerableCalc(expr#0=[{inputs}], expr#1=['dummy-datasource':VARCHAR], expr#2=[GEOIP($t1, $t0)], expr#3=['dummy_sub_field'], expr#4=[ITEM($t2, $t3)], host=[$t0], $f1=[$t2], $f2=[$t4])
+      CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_weblogs]], PushDownContext=[[PROJECT->[host], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["host"],"excludes":[]}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])

integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/access_struct_subfield_with_item.yaml

@@ -0,0 +1,9 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(host=[$0], info=[GEOIP('dummy-datasource':VARCHAR, $0)], info.dummy_sub_field=[ITEM(GEOIP('dummy-datasource':VARCHAR, $0), 'dummy_sub_field')])
+        CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_weblogs]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableCalc(expr#0..11=[{inputs}], expr#12=['dummy-datasource':VARCHAR], expr#13=[GEOIP($t12, $t0)], expr#14=['dummy_sub_field'], expr#15=[ITEM($t13, $t14)], host=[$t0], info=[$t13], info.dummy_sub_field=[$t15])
+        CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_weblogs]])