tellaro-query-language 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

{tellaro_query_language-0.2.0.dist-info → tellaro_query_language-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tellaro-query-language
-Version: 0.2.0
+Version: 0.2.2
 Summary: A flexible, human-friendly query language for searching and filtering structured data
 Home-page: https://github.com/tellaro/tellaro-query-language
 License: MIT
@@ -24,6 +24,7 @@ Requires-Dist: opensearch-dsl (>=2.1.0,<3.0.0) ; extra == "opensearch"
 Requires-Dist: opensearch-py (>=2.4.2,<3.0.0) ; extra == "opensearch"
 Requires-Dist: pyparsing (>=3.2.1,<4.0.0)
 Requires-Dist: setuptools (>=80.0.0,<81.0.0)
+Requires-Dist: urllib3 (>=2.5.0,<3.0.0)
 Project-URL: Documentation, https://github.com/tellaro/tellaro-query-language/tree/main/docs
 Project-URL: Repository, https://github.com/tellaro/tellaro-query-language
 Description-Content-Type: text/markdown
@@ -217,6 +218,28 @@ poetry run tests
 
 **Note**: The development setup uses `python-dotenv` to load OpenSearch credentials from `.env` files for integration testing. This is NOT required when using TQL as a package - see the [Package Usage Guide](docs/package-usage-guide.md) for production configuration patterns.
 
+### TQL Playground
+
+The repository includes an interactive web playground for testing TQL queries:
+
+```bash
+# Navigate to the playground directory
+cd playground
+
+# Start with Docker (recommended)
+docker-compose up
+
+# Or start with OpenSearch included
+docker-compose --profile opensearch up
+```
+
+Access the playground at:
+- Frontend: http://localhost:5173
+- API: http://localhost:8000
+- API Docs: http://localhost:8000/docs
+
+The playground uses your local TQL source code, so any changes you make are immediately reflected. See [playground/README.md](playground/README.md) for more details.
+
 ### File Operations
 
 ```python

{tellaro_query_language-0.2.0.dist-info → tellaro_query_language-0.2.2.dist-info}/RECORD

@@ -4,53 +4,53 @@ tql/cache/__init__.py,sha256=GIzIEMZUZEYJj72sAhuVLEG-OJEKUG2srUWNM3Ix-T8,213
 tql/cache/base.py,sha256=0b-8uyh3JltayGmXQI45snTqsM5sQu9u0KcNvZIRa-I,687
 tql/cache/memory.py,sha256=ibcmQSAxNvqCy6DksbU7gLu6UArYp1u3fW-oLubxtV0,2056
 tql/cache/redis.py,sha256=ZU_IsVDvpSYpNvPfnZ4iulJDODpEGx3c4dkXLzPzPVc,2309
-tql/core.py,sha256=fctqzGO1bt6_6pH1-j8mROXTdzPzRfW-v59cYMbLKI4,41947
+tql/core.py,sha256=bMPrcuutY-1yvC-4M7w2y1JxNitMyBSpxPfg8ohjO60,48406
 tql/core_components/README.md,sha256=Rm7w4UHdQ0vPBEFybE5b62IOvSA5Nzq2GRvtBHOapmc,3068
 tql/core_components/__init__.py,sha256=v8BBybPlqV7dkVY9mw1mblvqyAFJZ7Pf_bEc-jAL7FI,643
 tql/core_components/file_operations.py,sha256=Jr0kkxz_OP2KHOAsIr7KMtYe_lbu8LuBUySt2LQbjJw,3925
-tql/core_components/opensearch_operations.py,sha256=lmS8VdwnqkbqgYv2f9tygD-ixODWLUZ_d1H9SHyV_sI,38268
-tql/core_components/stats_operations.py,sha256=zAfDhVOFFPMrRIMw6Qtjxbobbdi7ao_HuHBCcVc3BGY,7579
+tql/core_components/opensearch_operations.py,sha256=KvmK1FnkGZFBjBysH_sDjzIRnyUcNn7wzLzuRr1rBlg,54264
+tql/core_components/stats_operations.py,sha256=aqTGAqIFvR6EkSbJEd0qft8Ldy8uiTrK2XI9o5bZUOs,8014
 tql/core_components/validation_operations.py,sha256=_VPXh0HABBjsXF99jFT7B6-5QAPsADOCy6poinGrxeE,22454
-tql/evaluator.py,sha256=kjqf2rkjdfPhrX2gyNtiBpdnBpob3LFkTuvplLtc-Ww,15542
+tql/evaluator.py,sha256=_JYr-wK3F1wvBoNGIBiAEaP6Ot1g2qxZ4lOjPdOqvDk,17698
 tql/evaluator_components/README.md,sha256=c59yf2au34yPhrru7JWgGop_ORteB6w5vfMhsac8j3k,3882
 tql/evaluator_components/__init__.py,sha256=DourRUSYXWPnCghBFj7W0YfMeymT3X8YTDCwnLIyP1c,535
 tql/evaluator_components/field_access.py,sha256=BuXvL9jlv4H77neT70Vh7_qokmzs-d4EbSDA2FB1IT0,6435
-tql/evaluator_components/special_expressions.py,sha256=rC1klZRA6cXLeS14xBi5LYYMi0V1qU3yg6_B4etRRBs,12139
-tql/evaluator_components/value_comparison.py,sha256=A-RJ_52FLvsMO0ciEhJEdfhfFsoE4c3viNk7qYMFILI,15636
+tql/evaluator_components/special_expressions.py,sha256=K6M5pW4Re2kEqxfxj9sc7I_M1tU3pn6LKJ2AfjHeciA,12917
+tql/evaluator_components/value_comparison.py,sha256=pL7-hxdNbzJ53DrTSiDdd7KYbVLChuNwFRLjG7P_1KM,17939
 tql/exceptions.py,sha256=hatIixXci6p57J9RrkfdvmKM_2i-JKb8ViL2kU4z7a8,5550
 tql/geoip_normalizer.py,sha256=tvie-5xevJEeLp2KmjoXDjYdND8AvyVE7lCO8qgUzGY,10486
-tql/mutator_analyzer.py,sha256=nCP64sVmEVV_zhpACfuTVTcuRV2NmypGbkWgS2taWOs,38039
-tql/mutators/__init__.py,sha256=zuzv6OMu2kVveDI85mu-6FA_-CZHsJEnIBIgtsjfXgY,6798
+tql/mutator_analyzer.py,sha256=OzI7t3C4H0IJOonpywE5LWz2cm5Dco5xnp2RTQOiSWg,55638
+tql/mutators/__init__.py,sha256=eTK8sRw4KXXnTZTn5ETIqwcaIek5rSUIVyZsxTwNNHA,6966
 tql/mutators/base.py,sha256=4Ze_x1sTO11OILXfcF2XN7ttyHcZ4gwn96UXFMMaC6M,2523
-tql/mutators/dns.py,sha256=Rmn6Kv-BmhfhMSYJ4GrkNmYms5rmpcxuG6leNxJspOQ,13387
+tql/mutators/dns.py,sha256=1IKgHolFLRMR4TOgK0AiLjz5vDtFiqO328mVF4Vzk3s,14428
 tql/mutators/encoding.py,sha256=yt12BJrHAIJfBesP8VOSfVlvJqB1yOmEeT_8QDPvNN8,7985
 tql/mutators/geo.py,sha256=fFQSg_Li3KjFKS3TI26yDzrDpWsmC3MfmgcsxYoQMgM,14507
 tql/mutators/list.py,sha256=949ZrKKhL4INkH2Od8bq7Ey80kFX_23PEfRKueG82cU,7084
 tql/mutators/network.py,sha256=1lZpmKt1GoTfNxiXUmSXkTwJIzPQZnQEgU7ojpBSm3A,5458
-tql/mutators/security.py,sha256=BE5x1IFCjq85UD22Q2sgSue_xXwhCk0umvf91a1XqJY,8202
-tql/mutators/string.py,sha256=8CzqZ2G5vpDeaEQ_jSA3NRidfUqfOrkUT9o2cmUUAEA,6288
+tql/mutators/security.py,sha256=XyWuPxgpCi-igHKmkbP0_0V-evOc3FG2a1igY8rQRX4,9256
+tql/mutators/string.py,sha256=0pIqIcpbRjUvGexsVV2lREPrJAAlIA2INhjvssLEKxw,9180
 tql/opensearch.py,sha256=J7LhfVJfaXEWtyZqVDqNZaeIbIcUYJr2cQtfKzdyIhM,3362
 tql/opensearch_components/README.md,sha256=gt-qLmmach8Kh7-QwLZmoAxxIL79XIG1EDqJum8PMZE,3756
 tql/opensearch_components/__init__.py,sha256=_zIZY8Fns7mkEcY6w2p9FNRBXtEmmPFFJEcFRfrVyXA,514
-tql/opensearch_components/field_mapping.py,sha256=N4r7VkzNeXjIhNDt2cfnd1LbkvaS_9O298PIDAcR_Hw,17569
-tql/opensearch_components/lucene_converter.py,sha256=ke1cMCeGVf5UorjwhN3y9x-GO-AHJ7FmbuTPSNMxYJo,12570
-tql/opensearch_components/query_converter.py,sha256=QodPaRJf5fI_DpMzRGINFIdpmcZZ_3Umlw1h33zsrZE,36858
-tql/opensearch_mappings.py,sha256=zJCCdMrxK7mswrkxd5LiOhunQ9GIJNZdhktVoGXgVgk,11529
-tql/opensearch_stats.py,sha256=h-hA2EZ-sc4S1zxr7jaInmonfrgTXoGBbb9sOYurdFE,17823
-tql/parser.py,sha256=MnY1Gxm8vl2Uw2tU2CyFB26uE3TSeUx6i7APtlwh7Fk,75727
+tql/opensearch_components/field_mapping.py,sha256=fj388cKVyDXLJKi8giSiGHL9zg4cFRzy0VJ6nIsppSo,18102
+tql/opensearch_components/lucene_converter.py,sha256=OvYTZHNBktPGow1fsVm4TMlvxHSmWrnqo42lFZNxXTo,13175
+tql/opensearch_components/query_converter.py,sha256=vLoBqv7W3ntqUH6hcuT4PDJkGkAGSQCxMvAWC482c0g,41971
+tql/opensearch_mappings.py,sha256=sVLlQlE3eGD7iNNZ_m4F4j5GVzQAJhZyCqDKYRhLRh8,11531
+tql/opensearch_stats.py,sha256=aMV__jtlfogGBnFucsNPazORro2mYTz_C_w9uxOqsMI,24384
+tql/parser.py,sha256=9kewX4IbBL3W5hbq9Xhi4BGrQ4QaoWqz9AJV0Yuf9YA,78665
 tql/parser_components/README.md,sha256=lvQX72ckq2zyotGs8QIHHCIFqaA7bOHwkP44wU8Zoiw,2322
 tql/parser_components/__init__.py,sha256=zBwHBMPJyHSBbaOojf6qTrJYjJg5A6tPUE8nHFdRiQs,521
-tql/parser_components/ast_builder.py,sha256=-pbcYhZNoRm0AnjmJRAAlXLCAwHfauchTpX_6KO0plE,6793
+tql/parser_components/ast_builder.py,sha256=erHoeKAMzobswoRIXB9xcsZbzQ5-2ZwaYfQgRWoUAa8,9653
 tql/parser_components/error_analyzer.py,sha256=qlCD9vKyW73aeKQYI33P1OjIWSJ3LPd08wuN9cis2fU,4012
-tql/parser_components/field_extractor.py,sha256=TumeuUo2c5gPYVbTPsmU43C3TJFC8chAAWERu5v_Q3c,4182
-tql/parser_components/grammar.py,sha256=lSvjABvEBaH29-ad-_UGD4WmofdNwC_pO2OKQJ_It-U,19309
-tql/post_processor.py,sha256=-vA2wgbuLij2FVnj5I9HDHtw5bKj9Cu3EE9mtoeSWk8,28859
+tql/parser_components/field_extractor.py,sha256=eUEkmiYWX2OexanFqhHeX8hcIkRlfIcgMB667e0HRYs,4629
+tql/parser_components/grammar.py,sha256=h58RBshZHXgbP1EmNwmf7dny-fgVloNg-qN4Rivross,20599
+tql/post_processor.py,sha256=MZOJzuWTL2qdvu-AUNMryYF2D-piv8rYH5vCcrLt5-A,50069
 tql/scripts.py,sha256=VOr5vCjIvKlW36kwvJx7JGFIRM16IJZlbJcWlBexBtk,3728
-tql/stats_evaluator.py,sha256=lOEbICFuP0krZZqEjREz37xlpm35_P6eRgkHVgJLNI4,15703
+tql/stats_evaluator.py,sha256=OQZuNLwLHAtWrwAh3utdtr1fQt3tftCs6L-1G1NQCGQ,22318
 tql/stats_transformer.py,sha256=MT-4rDWZSySgn4Fuq9H0c-mvwFYLM6FqWpPv2rHX-rE,7588
 tql/validators.py,sha256=e9MlX-zQ_O3M8YP8vXyMjKU8iiJMTh6mMK0iv0_4gTY,3771
-tellaro_query_language-0.2.0.dist-info/LICENSE,sha256=zRhQ85LnW55fWgAjQctckwQ67DX5Jmt64lq343ThZFU,1063
-tellaro_query_language-0.2.0.dist-info/METADATA,sha256=c3rHLr71I1LgwoDt0rypKbIScua6k2SVa_ozNHpiOg4,15109
-tellaro_query_language-0.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-tellaro_query_language-0.2.0.dist-info/entry_points.txt,sha256=H43APfGBMsZkKsUCnFTaqprQPW-Kce2yz2qsBL3dZrw,164
-tellaro_query_language-0.2.0.dist-info/RECORD,,
+tellaro_query_language-0.2.2.dist-info/LICENSE,sha256=zRhQ85LnW55fWgAjQctckwQ67DX5Jmt64lq343ThZFU,1063
+tellaro_query_language-0.2.2.dist-info/METADATA,sha256=QsCXKY_0aHeMorc4PepJ84ViZbTK53suxILmr868Lkk,15740
+tellaro_query_language-0.2.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+tellaro_query_language-0.2.2.dist-info/entry_points.txt,sha256=H43APfGBMsZkKsUCnFTaqprQPW-Kce2yz2qsBL3dZrw,164
+tellaro_query_language-0.2.2.dist-info/RECORD,,

tql/core.py

@@ -212,8 +212,10 @@ class TQL:
         if query_type == "stats_expr":
             # This is a pure stats query like "| stats count()"
             stats_result = self.stats(data, query)
+            # Get viz_hint from the parsed AST
+            viz_hint = ast.get("viz_hint")
             # Convert to execute_opensearch format
-            result["stats"] = self._convert_stats_result(stats_result)
+            result["stats"] = self._convert_stats_result(stats_result, viz_hint)
             result["total"] = len(records)
             return result
 
@@ -232,7 +234,9 @@ class TQL:
 
             # Apply stats to filtered data
             stats_result = self.stats_evaluator.evaluate_stats(filtered_records, stats_ast)
-            result["stats"] = self._convert_stats_result(stats_result)
+            # Get viz_hint from the stats AST
+            viz_hint = stats_ast.get("viz_hint")
+            result["stats"] = self._convert_stats_result(stats_result, viz_hint)
             result["total"] = len(filtered_records)
 
             # Include filtered documents if size > 0
@@ -241,40 +245,63 @@ class TQL:
 
             return result
 
-        # Handle regular filter queries
-        matched_records = []
-        has_enrichments = False
+        # Handle regular filter queries with mutator analysis
+        # Analyze the query for mutators
+        from .mutator_analyzer import MutatorAnalyzer
+        from .post_processor import QueryPostProcessor
+
+        analyzer = MutatorAnalyzer(field_mappings=self.field_mappings)
+        analysis_result = analyzer.analyze_ast(ast, context="in_memory")
 
+        # Use the optimized AST for evaluation
+        optimized_ast = analysis_result.optimized_ast
+
+        # First pass: collect all matching records using optimized AST
+        matched_records = []
         for record in records:
-            # Check if record matches
-            if self.evaluator._evaluate_node(ast, record, self._simple_mappings):
-                # Apply any mutators to enrich the record
-                enriched_record = self._apply_mutators_to_record(ast, record)
-                matched_records.append(enriched_record)
-                # Check if enrichments were added
-                if not has_enrichments and enriched_record is not record:
-                    has_enrichments = True
+            # Check if record matches using the optimized AST (without array operators)
+            if self.evaluator._evaluate_node(optimized_ast, record, self._simple_mappings):
+                matched_records.append(record)
+
+        # Apply post-processing if needed
+        if analysis_result.post_processing_requirements:
+            processor = QueryPostProcessor()
+
+            # Apply mutators/enrichments
+            processed_records = processor.process_results(
+                matched_records, analysis_result.post_processing_requirements, track_enrichments=save_enrichment
+            )
+
+            # Apply filters (for array operators like any/all/none)
+            filtered_records = processor.filter_results(processed_records, analysis_result.post_processing_requirements)
+
+            matched_records = filtered_records
+            result["post_processing_applied"] = True
+
+            # Add post-processing stats
+            result["post_processing_stats"] = {
+                "documents_retrieved": len(processed_records),
+                "documents_returned": len(filtered_records),
+                "documents_filtered": len(processed_records) - len(filtered_records),
+            }
 
         # Set result data
         result["total"] = len(matched_records)
         if size > 0:
             result["results"] = matched_records[:size]
 
-        # Check if post-processing (mutators) were applied
-        if has_enrichments:
-            result["post_processing_applied"] = True
+        # Update health status based on analysis
+        result["health_status"] = analysis_result.health_status
+        result["health_reasons"] = [
+            reason.get("reason", reason.get("description", "")) for reason in analysis_result.health_reasons
+        ]
 
         # Save enrichments if requested
-        if save_enrichment and has_enrichments and source_file:
-            # For file sources, update all records (not just matches)
-            all_enriched = []
-            for record in records:
-                enriched_record = self._apply_mutators_to_record(ast, record)
-                all_enriched.append(enriched_record)
-
-            # Save based on file type
-            if source_file.lower().endswith(".json"):
-                self.file_ops.save_enrichments_to_json(source_file, all_enriched)
+        if save_enrichment and result.get("post_processing_applied") and source_file:
+            # For file sources, we would need to re-process all records with enrichments
+            # This is a complex operation that would require applying mutators to all records
+            # For now, we'll skip this functionality in the in-memory implementation
+            pass
 
         return result
 
@@ -370,6 +397,38 @@ class TQL:
                 elif node_type == "unary_op":
                     stats["logical_operators"].add("not")
                     traverse_ast(node.get("operand"), depth + 1)
+                elif node_type == "query_with_stats":
+                    # Traverse into the filter part to find mutators and fields
+                    filter_node = node.get("filter")
+                    if filter_node:
+                        traverse_ast(filter_node, depth + 1)
+                    # Also traverse the stats part
+                    stats_node = node.get("stats")
+                    if stats_node:
+                        traverse_ast(stats_node, depth + 1)
+                elif node_type == "stats_expr":
+                    # Check aggregations for any fields or mutators
+                    aggregations = node.get("aggregations", [])
+                    for agg in aggregations:
+                        if isinstance(agg, dict):
+                            field = agg.get("field")
+                            if field and field != "*":
+                                stats["fields"].add(field)
+                            # Check for field mutators in aggregations
+                            if agg.get("field_mutators"):
+                                stats["has_mutators"] = True
+                elif node_type == "geo_expr":
+                    # Geo expressions always have mutators
+                    field = node.get("field")
+                    if field:
+                        stats["fields"].add(field)
+                    stats["has_mutators"] = True
+                elif node_type == "nslookup_expr":
+                    # NSLookup expressions always have mutators
+                    field = node.get("field")
+                    if field:
+                        stats["fields"].add(field)
+                    stats["has_mutators"] = True
 
         traverse_ast(ast)
 
@@ -450,7 +509,7 @@ class TQL:
         index: Optional[str] = None,
         query: Optional[str] = None,
         size: int = 500,
-        search_after: Optional[List[Any]] = None,
+        from_: int = 0,
         sort: Optional[List[Dict[str, Any]]] = None,
         timestamp_field: str = "@timestamp",
         time_range: Optional[Dict[str, str]] = None,
@@ -506,24 +565,6 @@ class TQL:
         # Remove parameters that the new implementation doesn't understand
         filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ["save_enrichment", "opensearch_client"]}
 
-        # Check for backward compatibility - if from_ is in kwargs, it's old usage
-        if "from_" in kwargs:
-            # Old style pagination - convert to search_after with default sort
-            from_ = kwargs.pop("from_")
-            if from_ > 0 and not sort:
-                # Need a sort for proper pagination
-                sort = [{"_id": "asc"}]
-            # Note: We can't directly convert from_ to search_after without previous results
-            # This is a limitation of moving to search_after
-            if from_ > 0:
-                import warnings
-
-                warnings.warn(
-                    "from_ parameter is deprecated. Use search_after with sort for efficient pagination.",
-                    DeprecationWarning,
-                    stacklevel=2,
-                )
-
         # Add the supported parameters
         filtered_kwargs.update(
             {
@@ -532,7 +573,7 @@ class TQL:
                 "scan_all": scan_all,
                 "scroll_size": scroll_size,
                 "scroll_timeout": scroll_timeout,
-                "search_after": search_after,
+                "from_": from_,
                 "sort": sort,
             }
         )
@@ -1001,22 +1042,140 @@ class TQL:
         Returns:
             Enriched record (may be same as input if no enrichments)
         """
-        # For now, return the original record
-        # TODO: Implement mutator application for enrichment  # noqa: W0511
-        return record
+        # Check if we need to apply mutators
+        if not self._has_output_mutators(ast):
+            return record
+
+        # Deep copy to avoid modifying original
+        import copy
+
+        enriched_record = copy.deepcopy(record)
+
+        # Apply mutators from AST nodes
+        self._apply_node_mutators(ast, enriched_record)
+
+        return enriched_record
+
+    def _has_output_mutators(self, ast: Dict[str, Any]) -> bool:
+        """Check if AST contains mutators that should transform output.
+
+        Args:
+            ast: Query AST
+
+        Returns:
+            True if output mutators are present
+        """
+        if isinstance(ast, dict):
+            node_type = ast.get("type")
+
+            # Check for field mutators with exists operator (output transformation)
+            if node_type == "comparison" and ast.get("operator") == "exists" and ast.get("field_mutators"):
+                return True
+
+            # Recursively check child nodes
+            if node_type == "logical_op":
+                left = ast.get("left", {})
+                right = ast.get("right", {})
+                return self._has_output_mutators(left) or self._has_output_mutators(right)
+            elif node_type == "unary_op":
+                operand = ast.get("operand", {})
+                return self._has_output_mutators(operand)
+
+        return False
+
+    def _apply_node_mutators(self, ast: Dict[str, Any], record: Dict[str, Any]) -> None:
+        """Apply mutators from AST nodes to the record.
+
+        Args:
+            ast: Query AST
+            record: Record to modify (in-place)
+        """
+        if not isinstance(ast, dict):
+            return
+
+        node_type = ast.get("type")
+
+        # Apply mutators for exists operator (output transformation)
+        if node_type == "comparison" and ast.get("operator") == "exists" and ast.get("field_mutators"):
+            field_name = ast["field"]
+            field_mutators = ast["field_mutators"]
+
+            # Get field value
+            field_value = self._get_nested_field(record, field_name)
+
+            if field_value is not None:
+                # Apply mutators
+                from .mutators import apply_mutators
+
+                mutated_value = apply_mutators(field_value, field_mutators, field_name, record)
+
+                # Update record with mutated value
+                self._set_nested_field(record, field_name, mutated_value)
+
+        # Recursively process child nodes
+        elif node_type == "logical_op":
+            self._apply_node_mutators(ast.get("left", {}), record)
+            self._apply_node_mutators(ast.get("right", {}), record)
+        elif node_type == "unary_op":
+            self._apply_node_mutators(ast.get("operand", {}), record)
+
+    def _get_nested_field(self, record: Dict[str, Any], field_path: str) -> Any:
+        """Get value from nested field path.
+
+        Args:
+            record: Record dictionary
+            field_path: Dot-separated field path
+
+        Returns:
+            Field value or None if not found
+        """
+        parts = field_path.split(".")
+        current = record
+
+        for part in parts:
+            if isinstance(current, dict) and part in current:
+                current = current[part]
+            else:
+                return None
+
+        return current
+
+    def _set_nested_field(self, record: Dict[str, Any], field_path: str, value: Any) -> None:
+        """Set value in nested field path.
+
+        Args:
+            record: Record dictionary to modify
+            field_path: Dot-separated field path
+            value: Value to set
+        """
+        parts = field_path.split(".")
+        current = record
 
-    def _convert_stats_result(self, stats_result: Dict[str, Any]) -> Dict[str, Any]:
+        # Navigate to parent of target field
+        for part in parts[:-1]:
+            if part not in current:
+                current[part] = {}
+            current = current[part]
+
+        # Set the value
+        if len(parts) > 0:
+            current[parts[-1]] = value
+
+    def _convert_stats_result(self, stats_result: Dict[str, Any], viz_hint: Optional[str] = None) -> Dict[str, Any]:
         """Convert stats result from query() format to execute_opensearch format.
 
         Args:
             stats_result: Result from stats() or query_stats() method
+            viz_hint: Optional visualization hint from the query
 
         Returns:
             Stats result in execute_opensearch format
         """
         # Map the stats evaluator format to execute_opensearch format
+        result = {}
+
         if stats_result.get("type") == "simple_aggregation":
-            return {
+            result = {
                 "type": "stats",
                 "operation": stats_result["function"],
                 "field": stats_result["field"],
@@ -1024,10 +1183,22 @@ class TQL:
             }
         elif stats_result.get("type") == "multiple_aggregations":
             # For multiple aggregations, return the results dict
-            return {"type": "stats_multiple", "results": stats_result["results"]}
+            result = {"type": "stats_multiple", "results": stats_result["results"]}
         elif stats_result.get("type") == "grouped_aggregation":
             # For grouped aggregations
-            return {"type": "stats_grouped", "group_by": stats_result["group_by"], "results": stats_result["results"]}
+            result = {
+                "type": "stats_grouped",
+                "group_by": stats_result["group_by"],
+                "results": stats_result["results"],
+                "operation": stats_result.get("function", "count"),
+                "field": stats_result.get("field", "*"),
+            }
         else:
             # Return as-is if format is unknown
-            return stats_result
+            result = stats_result
+
+        # Add visualization hint if provided
+        if viz_hint:
+            result["viz_hint"] = viz_hint
+
+        return result