From 4bfd89c9160f623ff6f40da4af174630673ba5ac Mon Sep 17 00:00:00 2001 From: Heng Qian Date: Fri, 8 Aug 2025 12:44:45 +0800 Subject: [PATCH 1/3] Support script push down on text field Signed-off-by: Heng Qian --- .../sql/calcite/remote/CalciteExplainIT.java | 14 +++ .../calcite/explain_script_push_on_text.json | 6 + .../calcite/explain_text_like_function.json | 2 +- .../opensearch/request/PredicateAnalyzer.java | 7 +- .../storage/script/CalciteScriptEngine.java | 79 +++++-------- .../aggregation/CalciteAggregationScript.java | 6 +- .../storage/script/core/CalciteScript.java | 10 +- .../script/filter/CalciteFilterScript.java | 6 +- .../request/PredicateAnalyzerTest.java | 106 +++++++++++++++--- 9 files changed, 151 insertions(+), 85 deletions(-) create mode 100644 integ-test/src/test/resources/expectedOutput/calcite/explain_script_push_on_text.json diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java index 2c9ba3107f7..96930929daa 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java @@ -69,6 +69,7 @@ public void supportPushDownSortMergeJoin() throws IOException { } // Only for Calcite + @Ignore("We've supported script push down on text field") @Test public void supportPartialPushDown() throws IOException { Assume.assumeTrue("This test is only for push down enabled", isPushdownEnabled()); @@ -82,6 +83,7 @@ public void supportPartialPushDown() throws IOException { } // Only for Calcite + @Ignore("We've supported script push down on text field") @Test public void supportPartialPushDown_NoPushIfAllFailed() throws IOException { Assume.assumeTrue("This test is only for push down enabled", isPushdownEnabled()); @@ -94,6 +96,7 @@ public void supportPartialPushDown_NoPushIfAllFailed() throws IOException { assertJsonEqualsIgnoreId(expected, result); } + @Ignore("We've supported script push down on text field") @Test public void supportPartialPushDownScript() throws IOException { Assume.assumeTrue("This test is only for push down enabled", isPushdownEnabled()); @@ -134,6 +137,17 @@ public void testExplainWithReverse() throws IOException { assertTrue(result.contains("dir0=[DESC]")); } + // Only for Calcite + @Test + public void supportPushDownScriptOnTextField() throws IOException { + String result = + explainQueryToString( + "explain source=opensearch-sql_test_index_account | where length(address) > 0 | eval" + + " address_length = length(address) | stats count() by address_length"); + String expected = loadFromFile("expectedOutput/calcite/explain_script_push_on_text.json"); + assertJsonEqualsIgnoreId(expected, result); + } + /** * Executes the PPL query and returns the result as a string with windows-style line breaks * replaced with Unix-style ones. diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_script_push_on_text.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_script_push_on_text.json new file mode 100644 index 00000000000..5cd4c2505bf --- /dev/null +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_script_push_on_text.json @@ -0,0 +1,6 @@ +{ + "calcite": { + "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(count()=[$1], address_length=[$0])\n LogicalAggregate(group=[{0}], count()=[COUNT()])\n LogicalProject(address_length=[CHAR_LENGTH($2)])\n LogicalFilter(condition=[>(CHAR_LENGTH($2), 0)])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n", + "physical": "EnumerableLimit(fetch=[10000])\n EnumerableCalc(expr#0..1=[{inputs}], count()=[$t1], address_length=[$t0])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[address], SCRIPT->>(CHAR_LENGTH($0), 0), AGGREGATION->rel#:LogicalAggregate.NONE.[](input=RelSubset#,group={0},count()=COUNT())], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":0,\"timeout\":\"1m\",\"query\":{\"script\":{\"script\":{\"source\":\"{\\\"langType\\\":\\\"calcite\\\",\\\"script\\\":\\\"rO0ABXNyABFqYXZhLnV0aWwuQ29sbFNlcleOq7Y6G6gRAwABSQADdGFneHAAAAADdwQAAAAGdAAHcm93VHlwZXQAlnsKICAiZmllbGRzIjogWwogICAgewogICAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgInByZWNpc2lvbiI6IC0xLAogICAgICAibmFtZSI6ICJhZGRyZXNzIgogICAgfQogIF0sCiAgIm51bGxhYmxlIjogZmFsc2UKfXQABGV4cHJ0AbN7CiAgIm9wIjogewogICAgIm5hbWUiOiAiPiIsCiAgICAia2luZCI6ICJHUkVBVEVSX1RIQU4iLAogICAgInN5bnRheCI6ICJCSU5BUlkiCiAgfSwKICAib3BlcmFuZHMiOiBbCiAgICB7CiAgICAgICJvcCI6IHsKICAgICAgICAibmFtZSI6ICJDSEFSX0xFTkdUSCIsCiAgICAgICAgImtpbmQiOiAiQ0hBUl9MRU5HVEgiLAogICAgICAgICJzeW50YXgiOiAiRlVOQ1RJT04iCiAgICAgIH0sCiAgICAgICJvcGVyYW5kcyI6IFsKICAgICAgICB7CiAgICAgICAgICAiaW5wdXQiOiAwLAogICAgICAgICAgIm5hbWUiOiAiJDAiCiAgICAgICAgfQogICAgICBdCiAgICB9LAogICAgewogICAgICAibGl0ZXJhbCI6IDAsCiAgICAgICJ0eXBlIjogewogICAgICAgICJ0eXBlIjogIklOVEVHRVIiLAogICAgICAgICJudWxsYWJsZSI6IGZhbHNlCiAgICAgIH0KICAgIH0KICBdCn10AApmaWVsZFR5cGVzc3IAEWphdmEudXRpbC5IYXNoTWFwBQfawcMWYNEDAAJGAApsb2FkRmFjdG9ySQAJdGhyZXNob2xkeHA/QAAAAAAADHcIAAAAEAAAAAF0AAdhZGRyZXNzc3IAOm9yZy5vcGVuc2VhcmNoLnNxbC5vcGVuc2VhcmNoLmRhdGEudHlwZS5PcGVuU2VhcmNoVGV4dFR5cGWtg6OTBOMxRAIAAUwABmZpZWxkc3QAD0xqYXZhL3V0aWwvTWFwO3hyADpvcmcub3BlbnNlYXJjaC5zcWwub3BlbnNlYXJjaC5kYXRhLnR5cGUuT3BlblNlYXJjaERhdGFUeXBlwmO8ygL6BTUCAANMAAxleHByQ29yZVR5cGV0ACtMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByQ29yZVR5cGU7TAALbWFwcGluZ1R5cGV0AEhMb3JnL29wZW5zZWFyY2gvc3FsL29wZW5zZWFyY2gvZGF0YS90eXBlL09wZW5TZWFyY2hEYXRhVHlwZSRNYXBwaW5nVHlwZTtMAApwcm9wZXJ0aWVzcQB+AAt4cH5yAClvcmcub3BlbnNlYXJjaC5zcWwuZGF0YS50eXBlLkV4cHJDb3JlVHlwZQAAAAAAAAAAEgAAeHIADmphdmEubGFuZy5FbnVtAAAAAAAAAAASAAB4cHQAB1VOS05PV05+cgBGb3JnLm9wZW5zZWFyY2guc3FsLm9wZW5zZWFyY2guZGF0YS50eXBlLk9wZW5TZWFyY2hEYXRhVHlwZSRNYXBwaW5nVHlwZQAAAAAAAAAAEgAAeHEAfgARdAAEVGV4dHNyADxzaGFkZWQuY29tLmdvb2dsZS5jb21tb24uY29sbGVjdC5JbW11dGFibGVNYXAkU2VyaWFsaXplZEZvcm0AAAAAAAAAAAIAAkwABGtleXN0ABJMamF2YS9sYW5nL09iamVjdDtMAAZ2YWx1ZXNxAH4AGHhwdXIAE1tMamF2YS5sYW5nLk9iamVjdDuQzlifEHMpbAIAAHhwAAAAAHVxAH4AGgAAAABzcQB+AAAAAAADdwQAAAAAeHh4\\\"}\",\"lang\":\"opensearch_compounded_script\",\"params\":{\"utcTimestamp\":*}},\"boost\":1.0}},\"_source\":{\"includes\":[\"address\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}],\"aggregations\":{\"composite_buckets\":{\"composite\":{\"size\":1000,\"sources\":[{\"address_length\":{\"terms\":{\"script\":{\"source\":\"{\\\"langType\\\":\\\"calcite\\\",\\\"script\\\":\\\"rO0ABXNyABFqYXZhLnV0aWwuQ29sbFNlcleOq7Y6G6gRAwABSQADdGFneHAAAAADdwQAAAAGdAAHcm93VHlwZXQAlnsKICAiZmllbGRzIjogWwogICAgewogICAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgInByZWNpc2lvbiI6IC0xLAogICAgICAibmFtZSI6ICJhZGRyZXNzIgogICAgfQogIF0sCiAgIm51bGxhYmxlIjogZmFsc2UKfXQABGV4cHJ0AKZ7CiAgIm9wIjogewogICAgIm5hbWUiOiAiQ0hBUl9MRU5HVEgiLAogICAgImtpbmQiOiAiQ0hBUl9MRU5HVEgiLAogICAgInN5bnRheCI6ICJGVU5DVElPTiIKICB9LAogICJvcGVyYW5kcyI6IFsKICAgIHsKICAgICAgImlucHV0IjogMCwKICAgICAgIm5hbWUiOiAiJDAiCiAgICB9CiAgXQp9dAAKZmllbGRUeXBlc3NyABdqYXZhLnV0aWwuTGlua2VkSGFzaE1hcDTATlwQbMD7AgABWgALYWNjZXNzT3JkZXJ4cgARamF2YS51dGlsLkhhc2hNYXAFB9rBwxZg0QMAAkYACmxvYWRGYWN0b3JJAAl0aHJlc2hvbGR4cD9AAAAAAAAMdwgAAAAQAAAAC3QADmFjY291bnRfbnVtYmVyfnIAKW9yZy5vcGVuc2VhcmNoLnNxbC5kYXRhLnR5cGUuRXhwckNvcmVUeXBlAAAAAAAAAAASAAB4cgAOamF2YS5sYW5nLkVudW0AAAAAAAAAABIAAHhwdAAETE9OR3QACWZpcnN0bmFtZXNyADpvcmcub3BlbnNlYXJjaC5zcWwub3BlbnNlYXJjaC5kYXRhLnR5cGUuT3BlblNlYXJjaFRleHRUeXBlrYOjkwTjMUQCAAFMAAZmaWVsZHN0AA9MamF2YS91dGlsL01hcDt4cgA6b3JnLm9wZW5zZWFyY2guc3FsLm9wZW5zZWFyY2guZGF0YS50eXBlLk9wZW5TZWFyY2hEYXRhVHlwZcJjvMoC+gU1AgADTAAMZXhwckNvcmVUeXBldAArTG9yZy9vcGVuc2VhcmNoL3NxbC9kYXRhL3R5cGUvRXhwckNvcmVUeXBlO0wAC21hcHBpbmdUeXBldABITG9yZy9vcGVuc2VhcmNoL3NxbC9vcGVuc2VhcmNoL2RhdGEvdHlwZS9PcGVuU2VhcmNoRGF0YVR5cGUkTWFwcGluZ1R5cGU7TAAKcHJvcGVydGllc3EAfgAReHB+cQB+AAt0AAdVTktOT1dOfnIARm9yZy5vcGVuc2VhcmNoLnNxbC5vcGVuc2VhcmNoLmRhdGEudHlwZS5PcGVuU2VhcmNoRGF0YVR5cGUkTWFwcGluZ1R5cGUAAAAAAAAAABIAAHhxAH4ADHQABFRleHRzcgA8c2hhZGVkLmNvbS5nb29nbGUuY29tbW9uLmNvbGxlY3QuSW1tdXRhYmxlTWFwJFNlcmlhbGl6ZWRGb3JtAAAAAAAAAAACAAJMAARrZXlzdAASTGphdmEvbGFuZy9PYmplY3Q7TAAGdmFsdWVzcQB+ABx4cHVyABNbTGphdmEubGFuZy5PYmplY3Q7kM5YnxBzKWwCAAB4cAAAAAB1cQB+AB4AAAAAc3EAfgAAAAAAA3cEAAAAAnQAB2tleXdvcmRzcQB+ABJ+cQB+AAt0AAZTVFJJTkd+cQB+ABh0AAdLZXl3b3JkcQB+AB14dAAHYWRkcmVzc3NxAH4AEHEAfgAWcQB+ABlxAH4AHXNxAH4AAAAAAAN3BAAAAAB4dAAHYmFsYW5jZXEAfgANdAAGZ2VuZGVyc3EAfgAQcQB+ABZxAH4AGXEAfgAdc3EAfgAAAAAAA3cEAAAAAnEAfgAicQB+ACN4dAAEY2l0eXNxAH4AEHEAfgAWcQB+ABlxAH4AHXNxAH4AAAAAAAN3BAAAAAJxAH4AInEAfgAjeHQACGVtcGxveWVyc3EAfgAQcQB+ABZxAH4AGXEAfgAdc3EAfgAAAAAAA3cEAAAAAnEAfgAicQB+ACN4dAAFc3RhdGVzcQB+ABBxAH4AFnEAfgAZcQB+AB1zcQB+AAAAAAADdwQAAAACcQB+ACJxAH4AI3h0AANhZ2VxAH4ADXQABWVtYWlsc3EAfgAQcQB+ABZxAH4AGXEAfgAdc3EAfgAAAAAAA3cEAAAAAnEAfgAicQB+ACN4dAAIbGFzdG5hbWVzcQB+ABBxAH4AFnEAfgAZcQB+AB1zcQB+AAAAAAADdwQAAAACcQB+ACJxAH4AI3h4AHg=\\\"}\",\"lang\":\"opensearch_compounded_script\",\"params\":{\"utcTimestamp\":*}},\"missing_bucket\":true,\"missing_order\":\"first\",\"order\":\"asc\"}}}]},\"aggregations\":{\"count()\":{\"value_count\":{\"field\":\"_index\"}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n" + } +} diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_text_like_function.json b/integ-test/src/test/resources/expectedOutput/calcite/explain_text_like_function.json index cd3f2f48d10..4a729956ec2 100644 --- a/integ-test/src/test/resources/expectedOutput/calcite/explain_text_like_function.json +++ b/integ-test/src/test/resources/expectedOutput/calcite/explain_text_like_function.json @@ -1,6 +1,6 @@ { "calcite": { "logical": "LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])\n LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10])\n LogicalFilter(condition=[ILIKE($2, '%Holmes%':VARCHAR, '\\')])\n CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])\n", - "physical": "EnumerableLimit(fetch=[10000])\n EnumerableCalc(expr#0..10=[{inputs}], expr#11=['%Holmes%':VARCHAR], expr#12=['\\'], expr#13=[ILIKE($t2, $t11, $t12)], proj#0..10=[{exprs}], $condition=[$t13])\n CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname]], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"timeout\":\"1m\",\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])\n" + "physical": "CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname], SCRIPT->ILIKE($2, '%Holmes%':VARCHAR, '\\'), LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={\"from\":0,\"size\":10000,\"timeout\":\"1m\",\"query\":{\"script\":{\"script\":{\"source\":\"{\\\"langType\\\":\\\"calcite\\\",\\\"script\\\":\\\"rO0ABXNyABFqYXZhLnV0aWwuQ29sbFNlcleOq7Y6G6gRAwABSQADdGFneHAAAAADdwQAAAAGdAAHcm93VHlwZXQEj3sKICAiZmllbGRzIjogWwogICAgewogICAgICAidHlwZSI6ICJCSUdJTlQiLAogICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAibmFtZSI6ICJhY2NvdW50X251bWJlciIKICAgIH0sCiAgICB7CiAgICAgICJ0eXBlIjogIlZBUkNIQVIiLAogICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAicHJlY2lzaW9uIjogLTEsCiAgICAgICJuYW1lIjogImZpcnN0bmFtZSIKICAgIH0sCiAgICB7CiAgICAgICJ0eXBlIjogIlZBUkNIQVIiLAogICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAicHJlY2lzaW9uIjogLTEsCiAgICAgICJuYW1lIjogImFkZHJlc3MiCiAgICB9LAogICAgewogICAgICAidHlwZSI6ICJCSUdJTlQiLAogICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAibmFtZSI6ICJiYWxhbmNlIgogICAgfSwKICAgIHsKICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICJudWxsYWJsZSI6IHRydWUsCiAgICAgICJwcmVjaXNpb24iOiAtMSwKICAgICAgIm5hbWUiOiAiZ2VuZGVyIgogICAgfSwKICAgIHsKICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICJudWxsYWJsZSI6IHRydWUsCiAgICAgICJwcmVjaXNpb24iOiAtMSwKICAgICAgIm5hbWUiOiAiY2l0eSIKICAgIH0sCiAgICB7CiAgICAgICJ0eXBlIjogIlZBUkNIQVIiLAogICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAicHJlY2lzaW9uIjogLTEsCiAgICAgICJuYW1lIjogImVtcGxveWVyIgogICAgfSwKICAgIHsKICAgICAgInR5cGUiOiAiVkFSQ0hBUiIsCiAgICAgICJudWxsYWJsZSI6IHRydWUsCiAgICAgICJwcmVjaXNpb24iOiAtMSwKICAgICAgIm5hbWUiOiAic3RhdGUiCiAgICB9LAogICAgewogICAgICAidHlwZSI6ICJCSUdJTlQiLAogICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAibmFtZSI6ICJhZ2UiCiAgICB9LAogICAgewogICAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICAgIm51bGxhYmxlIjogdHJ1ZSwKICAgICAgInByZWNpc2lvbiI6IC0xLAogICAgICAibmFtZSI6ICJlbWFpbCIKICAgIH0sCiAgICB7CiAgICAgICJ0eXBlIjogIlZBUkNIQVIiLAogICAgICAibnVsbGFibGUiOiB0cnVlLAogICAgICAicHJlY2lzaW9uIjogLTEsCiAgICAgICJuYW1lIjogImxhc3RuYW1lIgogICAgfQogIF0sCiAgIm51bGxhYmxlIjogZmFsc2UKfXQABGV4cHJ0Aa57CiAgIm9wIjogewogICAgIm5hbWUiOiAiSUxJS0UiLAogICAgImtpbmQiOiAiTElLRSIsCiAgICAic3ludGF4IjogIlNQRUNJQUwiCiAgfSwKICAib3BlcmFuZHMiOiBbCiAgICB7CiAgICAgICJpbnB1dCI6IDIsCiAgICAgICJuYW1lIjogIiQyIgogICAgfSwKICAgIHsKICAgICAgImxpdGVyYWwiOiAiJUhvbG1lcyUiLAogICAgICAidHlwZSI6IHsKICAgICAgICAidHlwZSI6ICJWQVJDSEFSIiwKICAgICAgICAibnVsbGFibGUiOiBmYWxzZSwKICAgICAgICAicHJlY2lzaW9uIjogLTEKICAgICAgfQogICAgfSwKICAgIHsKICAgICAgImxpdGVyYWwiOiAiXFwiLAogICAgICAidHlwZSI6IHsKICAgICAgICAidHlwZSI6ICJDSEFSIiwKICAgICAgICAibnVsbGFibGUiOiBmYWxzZSwKICAgICAgICAicHJlY2lzaW9uIjogMQogICAgICB9CiAgICB9CiAgXQp9dAAKZmllbGRUeXBlc3NyABFqYXZhLnV0aWwuSGFzaE1hcAUH2sHDFmDRAwACRgAKbG9hZEZhY3RvckkACXRocmVzaG9sZHhwP0AAAAAAAAx3CAAAABAAAAALdAAOYWNjb3VudF9udW1iZXJ+cgApb3JnLm9wZW5zZWFyY2guc3FsLmRhdGEudHlwZS5FeHByQ29yZVR5cGUAAAAAAAAAABIAAHhyAA5qYXZhLmxhbmcuRW51bQAAAAAAAAAAEgAAeHB0AARMT05HdAAJZmlyc3RuYW1lc3IAOm9yZy5vcGVuc2VhcmNoLnNxbC5vcGVuc2VhcmNoLmRhdGEudHlwZS5PcGVuU2VhcmNoVGV4dFR5cGWtg6OTBOMxRAIAAUwABmZpZWxkc3QAD0xqYXZhL3V0aWwvTWFwO3hyADpvcmcub3BlbnNlYXJjaC5zcWwub3BlbnNlYXJjaC5kYXRhLnR5cGUuT3BlblNlYXJjaERhdGFUeXBlwmO8ygL6BTUCAANMAAxleHByQ29yZVR5cGV0ACtMb3JnL29wZW5zZWFyY2gvc3FsL2RhdGEvdHlwZS9FeHByQ29yZVR5cGU7TAALbWFwcGluZ1R5cGV0AEhMb3JnL29wZW5zZWFyY2gvc3FsL29wZW5zZWFyY2gvZGF0YS90eXBlL09wZW5TZWFyY2hEYXRhVHlwZSRNYXBwaW5nVHlwZTtMAApwcm9wZXJ0aWVzcQB+ABB4cH5xAH4ACnQAB1VOS05PV05+cgBGb3JnLm9wZW5zZWFyY2guc3FsLm9wZW5zZWFyY2guZGF0YS50eXBlLk9wZW5TZWFyY2hEYXRhVHlwZSRNYXBwaW5nVHlwZQAAAAAAAAAAEgAAeHEAfgALdAAEVGV4dHNyADxzaGFkZWQuY29tLmdvb2dsZS5jb21tb24uY29sbGVjdC5JbW11dGFibGVNYXAkU2VyaWFsaXplZEZvcm0AAAAAAAAAAAIAAkwABGtleXN0ABJMamF2YS9sYW5nL09iamVjdDtMAAZ2YWx1ZXNxAH4AG3hwdXIAE1tMamF2YS5sYW5nLk9iamVjdDuQzlifEHMpbAIAAHhwAAAAAHVxAH4AHQAAAABzcQB+AAAAAAADdwQAAAACdAAHa2V5d29yZHNxAH4AEX5xAH4ACnQABlNUUklOR35xAH4AF3QAB0tleXdvcmRxAH4AHHh0AAdhZGRyZXNzc3EAfgAPcQB+ABVxAH4AGHEAfgAcc3EAfgAAAAAAA3cEAAAAAHh0AAdiYWxhbmNlcQB+AAx0AAZnZW5kZXJzcQB+AA9xAH4AFXEAfgAYcQB+ABxzcQB+AAAAAAADdwQAAAACcQB+ACFxAH4AInh0AARjaXR5c3EAfgAPcQB+ABVxAH4AGHEAfgAcc3EAfgAAAAAAA3cEAAAAAnEAfgAhcQB+ACJ4dAAIZW1wbG95ZXJzcQB+AA9xAH4AFXEAfgAYcQB+ABxzcQB+AAAAAAADdwQAAAACcQB+ACFxAH4AInh0AAVzdGF0ZXNxAH4AD3EAfgAVcQB+ABhxAH4AHHNxAH4AAAAAAAN3BAAAAAJxAH4AIXEAfgAieHQAA2FnZXEAfgAMdAAFZW1haWxzcQB+AA9xAH4AFXEAfgAYcQB+ABxzcQB+AAAAAAADdwQAAAACcQB+ACFxAH4AInh0AAhsYXN0bmFtZXNxAH4AD3EAfgAVcQB+ABhxAH4AHHNxAH4AAAAAAAN3BAAAAAJxAH4AIXEAfgAieHh4\\\"}\",\"lang\":\"opensearch_compounded_script\",\"params\":{\"utcTimestamp\":*}},\"boost\":1.0}},\"_source\":{\"includes\":[\"account_number\",\"firstname\",\"address\",\"balance\",\"gender\",\"city\",\"employer\",\"state\",\"age\",\"email\",\"lastname\"],\"excludes\":[]},\"sort\":[{\"_doc\":{\"order\":\"asc\"}}]}, requestedTotalSize=10000, pageSize=null, startFrom=0)])\n" } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PredicateAnalyzer.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PredicateAnalyzer.java index 58f5b4cf585..7d76796e01a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PredicateAnalyzer.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/request/PredicateAnalyzer.java @@ -91,7 +91,6 @@ import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.type.OpenSearchTextType; -import org.opensearch.sql.opensearch.storage.script.CalciteScriptEngine.ReferenceFieldVisitor; import org.opensearch.sql.opensearch.storage.script.CalciteScriptEngine.UnsupportedScriptException; import org.opensearch.sql.opensearch.storage.script.CompoundedScriptEngine.ScriptEngineType; import org.opensearch.sql.opensearch.storage.script.StringUtils; @@ -627,7 +626,7 @@ private QueryExpression andOr(RexCall call) { && call.getOperands().size() == 2 && (call.getOperands().get(0).getKind() == SqlKind.IS_NULL || call.getOperands().get(1).getKind() == SqlKind.IS_NULL)) { - throw new UnsupportedScriptException( + throw new PredicateAnalyzerException( "DSL will evaluate both branches of OR with isNUll, prevent push-down to avoid NPE"); } @@ -1357,10 +1356,6 @@ public ScriptQueryExpression( RelDataType rowType, Map fieldTypes, RelOptCluster cluster) { - ReferenceFieldVisitor validator = new ReferenceFieldVisitor(rowType, fieldTypes, true); - // Dry run visitInputRef to make sure the input reference ExprType is valid for script - // pushdown - validator.visitEach(List.of(rexNode)); RelJsonSerializer serializer = new RelJsonSerializer(cluster); this.code = SerializationWrapper.wrapWithLangType( diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java index 9b3bbd13b75..b8c6d188401 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java @@ -40,7 +40,6 @@ import java.util.Map; import java.util.Set; import java.util.function.BiFunction; -import java.util.function.Supplier; import lombok.RequiredArgsConstructor; import org.apache.calcite.DataContext; import org.apache.calcite.adapter.enumerable.EnumUtils; @@ -59,33 +58,31 @@ import org.apache.calcite.linq4j.tree.MethodCallExpression; import org.apache.calcite.linq4j.tree.MethodDeclaration; import org.apache.calcite.linq4j.tree.ParameterExpression; +import org.apache.calcite.linq4j.tree.Types; import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexExecutable; -import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexNode; import org.apache.calcite.rex.RexProgram; import org.apache.calcite.rex.RexProgramBuilder; -import org.apache.calcite.rex.RexVisitorImpl; import org.apache.calcite.schema.SchemaPlus; import org.apache.calcite.sql.validate.SqlConformance; import org.apache.calcite.sql.validate.SqlConformanceEnum; import org.apache.calcite.util.BuiltInMethod; import org.apache.calcite.util.Util; -import org.apache.commons.lang3.StringUtils; -import org.apache.commons.lang3.tuple.Pair; import org.checkerframework.checker.nullness.qual.Nullable; import org.opensearch.index.fielddata.ScriptDocValues; import org.opensearch.script.AggregationScript; import org.opensearch.script.FilterScript; import org.opensearch.script.ScriptContext; import org.opensearch.script.ScriptEngine; +import org.opensearch.search.lookup.LeafSearchLookup; import org.opensearch.sql.data.model.ExprTimestampValue; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; -import org.opensearch.sql.opensearch.request.PredicateAnalyzer.NamedFieldExpression; +import org.opensearch.sql.opensearch.data.type.OpenSearchTextType; import org.opensearch.sql.opensearch.storage.script.aggregation.CalciteAggregationScriptFactory; import org.opensearch.sql.opensearch.storage.script.filter.CalciteFilterScriptFactory; import org.opensearch.sql.opensearch.storage.serde.RelJsonSerializer; @@ -187,19 +184,25 @@ public ScriptInputGetter( @Override public org.apache.calcite.linq4j.tree.Expression field( BlockBuilder list, int index, @Nullable Type storageType) { - Pair refTypePair = - getValidatedReferenceNameAndType(rowType, index, fieldTypes); + String fieldName = rowType.getFieldList().get(index).getName(); + ExprType exprType = fieldTypes.get(fieldName); + String referenceField = OpenSearchTextType.toKeywordSubField(fieldName, exprType); MethodCallExpression fieldValueExpr = - Expressions.call( - DataContext.ROOT, - BuiltInMethod.DATA_CONTEXT_GET.method, - Expressions.constant(refTypePair.getKey())); + // Have to invoke `getFromSource` if the field is the text without keyword or struct + (referenceField == null || exprType == ExprCoreType.STRUCT) + ? Expressions.call( + EnumUtils.convert(DataContext.ROOT, ScriptDataContext.class), + Types.lookupMethod(ScriptDataContext.class, "getFromSource", String.class), + Expressions.constant(fieldName)) + : Expressions.call( + DataContext.ROOT, + BuiltInMethod.DATA_CONTEXT_GET.method, + Expressions.constant(referenceField)); if (storageType == null) { final RelDataType fieldType = rowType.getFieldList().get(index).getType(); storageType = ((JavaTypeFactory) typeFactory).getJavaClass(fieldType); } - return EnumUtils.convert( - tryConvertDocValue(fieldValueExpr, refTypePair.getValue()), storageType); + return EnumUtils.convert(tryConvertDocValue(fieldValueExpr, exprType), storageType); } /** @@ -214,32 +217,13 @@ private Expression tryConvertDocValue(Expression docValueExpr, ExprType exprType } } - public static class ReferenceFieldVisitor extends RexVisitorImpl> { - - private final RelDataType rowType; - private final Map fieldTypes; - - public ReferenceFieldVisitor( - RelDataType rowType, Map fieldTypes, boolean deep) { - super(deep); - this.rowType = rowType; - this.fieldTypes = fieldTypes; - } - - @Override - public Pair visitInputRef(RexInputRef inputRef) { - return getValidatedReferenceNameAndType(rowType, inputRef.getIndex(), fieldTypes); - } - } - public static class ScriptDataContext implements DataContext { - private final Supplier>> docProvider; + private final LeafSearchLookup lookup; private final Map params; - public ScriptDataContext( - Supplier>> docProvider, Map params) { - this.docProvider = docProvider; + public ScriptDataContext(LeafSearchLookup lookup, Map params) { + this.lookup = lookup; this.params = params; } @@ -264,7 +248,7 @@ public Object get(String name) { if (Variable.UTC_TIMESTAMP.camelName.equals(name)) return params.get(Variable.UTC_TIMESTAMP.camelName); - ScriptDocValues docValue = docProvider.get().get(name); + ScriptDocValues docValue = this.lookup.doc().get(name); if (docValue == null || docValue.isEmpty()) { return null; // No way to differentiate null and missing from doc value } @@ -278,6 +262,10 @@ public Object get(String name) { } return value; } + + public Object getFromSource(String name) { + return this.lookup.source().get(name); + } } /** @@ -337,21 +325,4 @@ public static String translate( return code; } - - private static Pair getValidatedReferenceNameAndType( - RelDataType rowType, int index, Map fieldTypes) { - String fieldName = rowType.getFieldList().get(index).getName(); - ExprType exprType = fieldTypes.get(fieldName); - if (exprType == ExprCoreType.STRUCT) { - throw new UnsupportedScriptException( - "Script query does not support fields of struct type: " + fieldName); - } - NamedFieldExpression expression = new NamedFieldExpression(fieldName, exprType); - String referenceField = expression.getReferenceForTermQuery(); - if (StringUtils.isEmpty(referenceField)) { - throw new UnsupportedScriptException( - "Field name cannot be empty for expression: " + expression.getRootName()); - } - return Pair.of(referenceField, exprType); - } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java index a04ec8f94a8..54d6d254b90 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java @@ -18,6 +18,7 @@ import org.apache.calcite.rel.type.RelDataType; import org.apache.lucene.index.LeafReaderContext; import org.opensearch.script.AggregationScript; +import org.opensearch.search.lookup.LeafSearchLookup; import org.opensearch.search.lookup.SearchLookup; import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.data.model.ExprValueUtils; @@ -32,6 +33,8 @@ class CalciteAggregationScript extends AggregationScript { /** Calcite Script. */ private final CalciteScript calciteScript; + private final LeafSearchLookup searchLookup; + private final RelDataType type; public CalciteAggregationScript( @@ -42,12 +45,13 @@ public CalciteAggregationScript( Map params) { super(params, lookup, context); this.calciteScript = new CalciteScript(function, params); + this.searchLookup = lookup.getLeafSearchLookup(context); this.type = type; } @Override public Object execute() { - Object value = calciteScript.execute(this::getDoc)[0]; + Object value = calciteScript.execute(this.searchLookup)[0]; ExprType exprType = OpenSearchTypeFactory.convertRelDataTypeToExprType(type); // See logic in {@link ExpressionAggregationScript::execute} return switch ((ExprCoreType) exprType) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java index 090bb535e15..f4368558cf6 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java @@ -8,11 +8,10 @@ import java.security.AccessController; import java.security.PrivilegedAction; import java.util.Map; -import java.util.function.Supplier; import lombok.EqualsAndHashCode; import org.apache.calcite.DataContext; import org.apache.calcite.linq4j.function.Function1; -import org.opensearch.index.fielddata.ScriptDocValues; +import org.opensearch.search.lookup.LeafSearchLookup; import org.opensearch.sql.opensearch.storage.script.CalciteScriptEngine.ScriptDataContext; /** @@ -36,12 +35,11 @@ public CalciteScript(Function1 function, Map>> docProvider) { + public Object[] execute(LeafSearchLookup lookup) { return AccessController.doPrivileged( - (PrivilegedAction) - () -> function.apply(new ScriptDataContext(docProvider, params))); + (PrivilegedAction) () -> function.apply(new ScriptDataContext(lookup, params))); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java index 988117b8c9f..ce94d14ad0a 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java @@ -11,6 +11,7 @@ import org.apache.calcite.linq4j.function.Function1; import org.apache.lucene.index.LeafReaderContext; import org.opensearch.script.FilterScript; +import org.opensearch.search.lookup.LeafSearchLookup; import org.opensearch.search.lookup.SearchLookup; import org.opensearch.sql.opensearch.storage.script.core.CalciteScript; @@ -24,6 +25,8 @@ class CalciteFilterScript extends FilterScript { /** Calcite Script. */ private final CalciteScript calciteScript; + private final LeafSearchLookup searchLookup; + public CalciteFilterScript( Function1 function, SearchLookup lookup, @@ -31,11 +34,12 @@ public CalciteFilterScript( Map params) { super(params, lookup, context); this.calciteScript = new CalciteScript(function, params); + this.searchLookup = lookup.getLeafSearchLookup(context); } @Override public boolean execute() { - Object result = calciteScript.execute(this::getDoc)[0]; + Object result = calciteScript.execute(this.searchLookup)[0]; // The result should be type of BOOLEAN_NULLABLE. Treat it as false if null return result != null && (boolean) result; } diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java index 5a2a7544998..76cb3f74892 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java @@ -8,12 +8,14 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; import com.google.common.collect.ImmutableList; import java.math.BigDecimal; import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.function.Consumer; import org.apache.calcite.plan.RelOptCluster; import org.apache.calcite.plan.volcano.VolcanoPlanner; import org.apache.calcite.rel.type.RelDataType; @@ -24,10 +26,14 @@ import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; +import org.apache.calcite.runtime.Hook; import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.type.SqlTypeFactoryImpl; import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.Holder; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.ExistsQueryBuilder; import org.opensearch.index.query.MatchBoolPrefixQueryBuilder; @@ -48,6 +54,9 @@ import org.opensearch.sql.opensearch.data.type.OpenSearchDataType; import org.opensearch.sql.opensearch.data.type.OpenSearchDataType.MappingType; import org.opensearch.sql.opensearch.request.PredicateAnalyzer.ExpressionNotAnalyzableException; +import org.opensearch.sql.opensearch.request.PredicateAnalyzer.QueryExpression; +import org.opensearch.sql.opensearch.storage.script.CalciteScriptEngine.UnsupportedScriptException; +import org.opensearch.sql.opensearch.storage.serde.SerializationWrapper; public class PredicateAnalyzerTest { final RelDataTypeFactory typeFactory = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT); @@ -592,14 +601,27 @@ void likeFunction_keywordField_generatesWildcardQuery() throws ExpressionNotAnal } @Test - void likeFunction_textField_throwsException() throws ExpressionNotAnalyzableException { + void likeFunction_textField_scriptPushDown() throws ExpressionNotAnalyzableException { RexInputRef field3 = builder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 2); List arguments = Arrays.asList(field3, builder.makeLiteral("%Hi%")); RexNode call = PPLFuncImpTable.INSTANCE.resolve(builder, "like", arguments.toArray(new RexNode[0])); - assertThrows( - ExpressionNotAnalyzableException.class, - () -> PredicateAnalyzer.analyze(call, schema, fieldTypes)); + + final RelDataType rowType = + builder + .getTypeFactory() + .builder() + .kind(StructKind.FULLY_QUALIFIED) + .add("a", builder.getTypeFactory().createSqlType(SqlTypeName.BIGINT)) + .add("b", builder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR)) + .build(); + Hook.CURRENT_TIME.addThread((Consumer>) h -> h.set(0L)); + QueryExpression expression = + PredicateAnalyzer.analyzeExpression(call, schema, fieldTypes, rowType, cluster); + assert (expression + .builder() + .toString() + .contains("\"lang\" : \"opensearch_compounded_script\"")); } @Test @@ -688,7 +710,7 @@ void equals_generatesTermQuery_TextWithKeyword() throws ExpressionNotAnalyzableE } @Test - void equals_throwException_TextWithoutKeyword() { + void equals_scriptPushDown_TextWithoutKeyword() throws ExpressionNotAnalyzableException { final RelDataType rowType = builder .getTypeFactory() @@ -701,15 +723,12 @@ void equals_throwException_TextWithoutKeyword() { final RexInputRef field3 = builder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 2); RexNode call = builder.makeCall(SqlStdOperatorTable.EQUALS, field3, stringLiteral); - ExpressionNotAnalyzableException exception = - assertThrows( - ExpressionNotAnalyzableException.class, - () -> PredicateAnalyzer.analyze(call, schema, fieldTypes, rowType, cluster)); - assertEquals("Can't convert =($2, 'Hi')", exception.getMessage()); + QueryBuilder builder = PredicateAnalyzer.analyze(call, schema, fieldTypes, rowType, cluster); + assert (builder.toString().contains("\"lang\" : \"opensearch_compounded_script\"")); } @Test - void isNullOr_throwException() { + void isNullOr_ScriptPushDown() throws ExpressionNotAnalyzableException { final RelDataType rowType = builder .getTypeFactory() @@ -720,10 +739,65 @@ void isNullOr_throwException() { .build(); // PPL IS_EMPTY is translated to OR(IS_NULL(arg), IS_EMPTY(arg)) RexNode call = PPLFuncImpTable.INSTANCE.resolve(builder, BuiltinFunctionName.IS_EMPTY, field2); - ExpressionNotAnalyzableException exception = - assertThrows( - ExpressionNotAnalyzableException.class, - () -> PredicateAnalyzer.analyzeExpression(call, schema, fieldTypes, rowType, cluster)); - assertEquals("Can't convert OR(IS NULL($1), IS EMPTY($1))", exception.getMessage()); + Hook.CURRENT_TIME.addThread((Consumer>) h -> h.set(0L)); + QueryExpression expression = + PredicateAnalyzer.analyzeExpression(call, schema, fieldTypes, rowType, cluster); + assert (expression + .builder() + .toString() + .contains("\"lang\" : \"opensearch_compounded_script\"")); + } + + @Test + void verify_partial_pushdown() throws ExpressionNotAnalyzableException { + RexNode call1 = builder.makeCall(SqlStdOperatorTable.EQUALS, field1, numericLiteral); + RexNode call2 = builder.makeCall(SqlStdOperatorTable.IS_EMPTY, field2); + try (MockedStatic mockedSerializationWrapper = + Mockito.mockStatic(SerializationWrapper.class)) { + mockedSerializationWrapper + .when(() -> SerializationWrapper.wrapWithLangType(any(), any())) + .thenThrow(new UnsupportedScriptException("")); + + // Partial push down part of and + RexNode andCall = builder.makeCall(SqlStdOperatorTable.AND, List.of(call1, call2)); + QueryExpression result = + PredicateAnalyzer.analyzeExpression(andCall, schema, fieldTypes, null, null); + + QueryBuilder resultBuilder = result.builder(); + assertInstanceOf(BoolQueryBuilder.class, resultBuilder); + assertEquals( + """ + { + "bool" : { + "must" : [ + { + "term" : { + "a" : { + "value" : 12, + "boost" : 1.0 + } + } + } + ], + "adjust_pure_negative" : true, + "boost" : 1.0 + } + }""", + resultBuilder.toString()); + + List unAnalyzableNodes = result.getUnAnalyzableNodes(); + assertEquals(1, unAnalyzableNodes.size()); + assertEquals(call2, unAnalyzableNodes.getFirst()); + + // Don't push down the whole condition if part of `or` cannot be pushed down + RexNode orCall = builder.makeCall(SqlStdOperatorTable.OR, List.of(call1, call2)); + ExpressionNotAnalyzableException exception = + assertThrows( + ExpressionNotAnalyzableException.class, + () -> { + PredicateAnalyzer.analyzeExpression(orCall, schema, fieldTypes, null, null); + }); + assertEquals("Can't convert OR(=($0, 12), IS EMPTY($1))", exception.getMessage()); + } } } From 25c5151e88af48eac5b878461369ef652f65350d Mon Sep 17 00:00:00 2001 From: Heng Qian Date: Fri, 8 Aug 2025 14:39:06 +0800 Subject: [PATCH 2/3] Fix IT Signed-off-by: Heng Qian --- .../sql/calcite/remote/CalciteExplainIT.java | 1 + .../storage/script/CalciteScriptEngine.java | 17 +++++++++++------ .../aggregation/CalciteAggregationScript.java | 8 ++++---- .../storage/script/core/CalciteScript.java | 11 +++++++---- .../script/filter/CalciteFilterScript.java | 9 +++++---- 5 files changed, 28 insertions(+), 18 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java index 96930929daa..c7f483978ea 100644 --- a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java @@ -140,6 +140,7 @@ public void testExplainWithReverse() throws IOException { // Only for Calcite @Test public void supportPushDownScriptOnTextField() throws IOException { + Assume.assumeTrue("This test is only for push down enabled", isPushdownEnabled()); String result = explainQueryToString( "explain source=opensearch-sql_test_index_account | where length(address) > 0 | eval" diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java index b8c6d188401..6a8426749ae 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/CalciteScriptEngine.java @@ -78,7 +78,7 @@ import org.opensearch.script.FilterScript; import org.opensearch.script.ScriptContext; import org.opensearch.script.ScriptEngine; -import org.opensearch.search.lookup.LeafSearchLookup; +import org.opensearch.search.lookup.SourceLookup; import org.opensearch.sql.data.model.ExprTimestampValue; import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; @@ -219,11 +219,16 @@ private Expression tryConvertDocValue(Expression docValueExpr, ExprType exprType public static class ScriptDataContext implements DataContext { - private final LeafSearchLookup lookup; + private final Map> docProvider; + private final SourceLookup sourceLookup; private final Map params; - public ScriptDataContext(LeafSearchLookup lookup, Map params) { - this.lookup = lookup; + public ScriptDataContext( + Map> docProvider, + SourceLookup sourceLookup, + Map params) { + this.docProvider = docProvider; + this.sourceLookup = sourceLookup; this.params = params; } @@ -248,7 +253,7 @@ public Object get(String name) { if (Variable.UTC_TIMESTAMP.camelName.equals(name)) return params.get(Variable.UTC_TIMESTAMP.camelName); - ScriptDocValues docValue = this.lookup.doc().get(name); + ScriptDocValues docValue = this.docProvider.get(name); if (docValue == null || docValue.isEmpty()) { return null; // No way to differentiate null and missing from doc value } @@ -264,7 +269,7 @@ public Object get(String name) { } public Object getFromSource(String name) { - return this.lookup.source().get(name); + return this.sourceLookup.get(name); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java index 54d6d254b90..8809513a2ad 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/aggregation/CalciteAggregationScript.java @@ -18,8 +18,8 @@ import org.apache.calcite.rel.type.RelDataType; import org.apache.lucene.index.LeafReaderContext; import org.opensearch.script.AggregationScript; -import org.opensearch.search.lookup.LeafSearchLookup; import org.opensearch.search.lookup.SearchLookup; +import org.opensearch.search.lookup.SourceLookup; import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.data.type.ExprCoreType; @@ -33,7 +33,7 @@ class CalciteAggregationScript extends AggregationScript { /** Calcite Script. */ private final CalciteScript calciteScript; - private final LeafSearchLookup searchLookup; + private final SourceLookup sourceLookup; private final RelDataType type; @@ -45,13 +45,13 @@ public CalciteAggregationScript( Map params) { super(params, lookup, context); this.calciteScript = new CalciteScript(function, params); - this.searchLookup = lookup.getLeafSearchLookup(context); + this.sourceLookup = lookup.getLeafSearchLookup(context).source(); this.type = type; } @Override public Object execute() { - Object value = calciteScript.execute(this.searchLookup)[0]; + Object value = calciteScript.execute(this.getDoc(), this.sourceLookup)[0]; ExprType exprType = OpenSearchTypeFactory.convertRelDataTypeToExprType(type); // See logic in {@link ExpressionAggregationScript::execute} return switch ((ExprCoreType) exprType) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java index f4368558cf6..4568fddc8ed 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/core/CalciteScript.java @@ -11,7 +11,8 @@ import lombok.EqualsAndHashCode; import org.apache.calcite.DataContext; import org.apache.calcite.linq4j.function.Function1; -import org.opensearch.search.lookup.LeafSearchLookup; +import org.opensearch.index.fielddata.ScriptDocValues; +import org.opensearch.search.lookup.SourceLookup; import org.opensearch.sql.opensearch.storage.script.CalciteScriptEngine.ScriptDataContext; /** @@ -35,11 +36,13 @@ public CalciteScript(Function1 function, Map> docProvider, SourceLookup sourceLookup) { return AccessController.doPrivileged( - (PrivilegedAction) () -> function.apply(new ScriptDataContext(lookup, params))); + (PrivilegedAction) + () -> function.apply(new ScriptDataContext(docProvider, sourceLookup, params))); } } diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java index ce94d14ad0a..90f5af66481 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/CalciteFilterScript.java @@ -11,8 +11,8 @@ import org.apache.calcite.linq4j.function.Function1; import org.apache.lucene.index.LeafReaderContext; import org.opensearch.script.FilterScript; -import org.opensearch.search.lookup.LeafSearchLookup; import org.opensearch.search.lookup.SearchLookup; +import org.opensearch.search.lookup.SourceLookup; import org.opensearch.sql.opensearch.storage.script.core.CalciteScript; /** @@ -25,7 +25,7 @@ class CalciteFilterScript extends FilterScript { /** Calcite Script. */ private final CalciteScript calciteScript; - private final LeafSearchLookup searchLookup; + private final SourceLookup sourceLookup; public CalciteFilterScript( Function1 function, @@ -34,12 +34,13 @@ public CalciteFilterScript( Map params) { super(params, lookup, context); this.calciteScript = new CalciteScript(function, params); - this.searchLookup = lookup.getLeafSearchLookup(context); + // TODO: we'd better get source from the leafLookup of super once it's available + this.sourceLookup = lookup.getLeafSearchLookup(context).source(); } @Override public boolean execute() { - Object result = calciteScript.execute(this.searchLookup)[0]; + Object result = calciteScript.execute(this.getDoc(), this.sourceLookup)[0]; // The result should be type of BOOLEAN_NULLABLE. Treat it as false if null return result != null && (boolean) result; } From fa129be1f4ef8513560af08d9ade4f16b5e377ac Mon Sep 17 00:00:00 2001 From: Heng Qian Date: Mon, 11 Aug 2025 17:46:05 +0800 Subject: [PATCH 3/3] Add UT for struct type push down Signed-off-by: Heng Qian --- .../request/PredicateAnalyzerTest.java | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java index 76cb3f74892..e3e810d58c9 100644 --- a/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java +++ b/opensearch/src/test/java/org/opensearch/sql/opensearch/request/PredicateAnalyzerTest.java @@ -48,6 +48,7 @@ import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.index.query.TermsQueryBuilder; import org.opensearch.index.query.WildcardQueryBuilder; +import org.opensearch.sql.data.type.ExprCoreType; import org.opensearch.sql.data.type.ExprType; import org.opensearch.sql.expression.function.BuiltinFunctionName; import org.opensearch.sql.expression.function.PPLFuncImpTable; @@ -723,10 +724,35 @@ void equals_scriptPushDown_TextWithoutKeyword() throws ExpressionNotAnalyzableEx final RexInputRef field3 = builder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR), 2); RexNode call = builder.makeCall(SqlStdOperatorTable.EQUALS, field3, stringLiteral); + Hook.CURRENT_TIME.addThread((Consumer>) h -> h.set(0L)); QueryBuilder builder = PredicateAnalyzer.analyze(call, schema, fieldTypes, rowType, cluster); assert (builder.toString().contains("\"lang\" : \"opensearch_compounded_script\"")); } + @Test + void equals_scriptPushDown_Struct() throws ExpressionNotAnalyzableException { + final RelDataType mapType = + typeFactory.createMapType( + typeFactory.createSqlType(SqlTypeName.VARCHAR), + typeFactory.createSqlType(SqlTypeName.VARCHAR)); + final RelDataType rowType = + builder + .getTypeFactory() + .builder() + .kind(StructKind.FULLY_QUALIFIED) + .add("d", mapType) + .build(); + final RexInputRef field4 = builder.makeInputRef(mapType, 3); + final Map newFieldTypes = + Map.of("d", OpenSearchDataType.of(ExprCoreType.STRUCT)); + final List newSchema = List.of("d"); + RexNode call = builder.makeCall(SqlStdOperatorTable.IS_NOT_NULL, field4); + Hook.CURRENT_TIME.addThread((Consumer>) h -> h.set(0L)); + QueryBuilder builder = + PredicateAnalyzer.analyze(call, newSchema, newFieldTypes, rowType, cluster); + assert (builder.toString().contains("\"lang\" : \"opensearch_compounded_script\"")); + } + @Test void isNullOr_ScriptPushDown() throws ExpressionNotAnalyzableException { final RelDataType rowType =