diff --git a/src/Engine/IO/DrillDown/KeyPathLeafCollector.cs b/src/Engine/IO/DrillDown/KeyPathLeafCollector.cs new file mode 100644 index 0000000..b58fd39 --- /dev/null +++ b/src/Engine/IO/DrillDown/KeyPathLeafCollector.cs @@ -0,0 +1,157 @@ +using System.Buffers; +using System.Text.Json; +using Refedle.Engine.IO.Json; +using Refedle.Engine.Types; + +namespace Refedle.Engine.IO.DrillDown; + +/// +/// Collects leaf row(s) and schema observations at the end of a KeyPath descent, plus the byte-level +/// value lookup used to descend object-key segments. Extracted from so +/// the traversal control flow depends only one-way on collection: KeyPathTraverser calls into this +/// type, never the reverse. All value slicing delegates to . +/// +internal static class KeyPathLeafCollector +{ + internal static void CollectLeafRows( + JsonRawBytes leafBytes, + string posHash, + string colName, + byte[] colNameUtf8, + List rows, + List keyOrder, + HashSet keySet, + Dictionary columnTypes, + Dictionary keyObservedCount) + { + var reader = new Utf8JsonReader(leafBytes.Span); + if (!reader.Read()) + { + return; + } + + if (reader.TokenType == JsonTokenType.StartObject) + { + rows.Add(new FocusedTableRow(leafBytes, posHash)); + var observedKeys = new HashSet(StringComparer.Ordinal); + SchemaScanner.ScanObject(leafBytes.Span, keyOrder, keySet, columnTypes, observedKeys); + SchemaScanner.IncrementObservationCounts(observedKeys, keyObservedCount); + return; + } + + if (reader.TokenType == JsonTokenType.StartArray) + { + CollectArrayLeafRows(leafBytes, posHash, rows, keyOrder, keySet, columnTypes, keyObservedCount); + return; + } + + // Primitive leaf (including null) — synthesize a single-key object so + // JsonObjectCellExtractor can extract it without modification. + // Note: ScanObject is NOT called here, so no type inference is performed; + // the synthesized column always receives ColumnType.Text (Phase 2 limitation). + var synthBytes = SynthesizeObject(colNameUtf8, leafBytes.Span); + rows.Add(new FocusedTableRow(synthBytes, posHash)); + SchemaScanner.RegisterKeyIfNew(colName, keyOrder, keySet); + SchemaScanner.IncrementObservationCounts([colName], keyObservedCount); + } + + internal static void CollectArrayLeafRows( + JsonRawBytes leafBytes, + string posHash, + List rows, + List keyOrder, + HashSet keySet, + Dictionary columnTypes, + Dictionary keyObservedCount) + { + var reader = new Utf8JsonReader(leafBytes.Span); + if (!reader.Read() || reader.TokenType != JsonTokenType.StartArray) + { + return; + } + + var elementIndex = 0; + while (reader.Read()) + { + if (reader.TokenType == JsonTokenType.EndArray) + { + break; + } + + if (reader.CurrentDepth != 1) + { + continue; + } + + var isObjectElement = reader.TokenType == JsonTokenType.StartObject; + var elementBytes = JsonByteExtractor.ExtractValueBytes(ref reader, leafBytes); + var elementHash = $"{posHash}:{elementIndex}"; + + if (isObjectElement) + { + rows.Add(new FocusedTableRow(elementBytes, elementHash)); + var observedKeys = new HashSet(StringComparer.Ordinal); + SchemaScanner.ScanObject(elementBytes.Span, keyOrder, keySet, columnTypes, observedKeys); + SchemaScanner.IncrementObservationCounts(observedKeys, keyObservedCount); + elementIndex++; + continue; + } + + // Primitive element (including null) — synthesize {"value": element}. + var synthBytes = SynthesizeObject("value"u8, elementBytes.Span); + rows.Add(new FocusedTableRow(synthBytes, elementHash)); + SchemaScanner.RegisterKeyIfNew("value", keyOrder, keySet); + SchemaScanner.IncrementObservationCounts(["value"], keyObservedCount); + elementIndex++; + } + } + + internal static JsonRawBytes? FindValueByKey(JsonRawBytes objectBytes, string key) + { + var reader = new Utf8JsonReader(objectBytes.Span); + if (!reader.Read() || reader.TokenType != JsonTokenType.StartObject) + { + return null; + } + + while (reader.Read()) + { + if (reader.TokenType == JsonTokenType.EndObject) + { + return null; + } + + if (reader.TokenType != JsonTokenType.PropertyName) + { + continue; + } + + if (!reader.ValueTextEquals(key)) + { + reader.Skip(); + continue; + } + + if (!reader.Read()) + { + return null; + } + + return JsonByteExtractor.ExtractValueBytes(ref reader, objectBytes); + } + + return null; + } + + private static JsonRawBytes SynthesizeObject(ReadOnlySpan keyUtf8, ReadOnlySpan valueBytes) + { + var buffer = new ArrayBufferWriter(); + using var writer = new Utf8JsonWriter(buffer); + writer.WriteStartObject(); + writer.WritePropertyName(keyUtf8); + writer.WriteRawValue(valueBytes, skipInputValidation: true); + writer.WriteEndObject(); + writer.Flush(); + return buffer.WrittenMemory; + } +} diff --git a/src/Engine/IO/DrillDown/KeyPathTraverser.cs b/src/Engine/IO/DrillDown/KeyPathTraverser.cs index 91ed3f7..dceb2f0 100644 --- a/src/Engine/IO/DrillDown/KeyPathTraverser.cs +++ b/src/Engine/IO/DrillDown/KeyPathTraverser.cs @@ -1,4 +1,3 @@ -using System.Buffers; using System.Text.Json; using Refedle.Engine.IO.Json; using Refedle.Engine.Types; @@ -6,10 +5,11 @@ namespace Refedle.Engine.IO.DrillDown; /// -/// Stateless helpers that traverse a KeyPath through a single record's bytes and collect the -/// leaf row(s) reached, accumulating schema information along the way. Shared implementation -/// detail of , split out to keep both classes under the -/// project's per-class line limit. +/// Traverses a KeyPath through a single record's bytes with an explicit-stack DFS, collecting leaf +/// rows via . Descent depth is bounded by the heap, not the call +/// stack, so a keyPath whose length is driven by untrusted input cannot overflow the stack. Leaf +/// collection and value lookup live in to keep both classes under +/// the per-class line limit and the dependency one-way (this class calls the collector, never back). /// internal static class KeyPathTraverser { @@ -30,9 +30,23 @@ public static void ExtractRows( Dictionary columnTypes, Dictionary keyObservedCount) { - TraverseKeyPath( - recordBytes, keyPath, 0, posHash, colName, colNameUtf8, - rows, keyOrder, keySet, columnTypes, keyObservedCount); + Stack stack = []; + stack.Push(TraversalFrame.Descend(recordBytes, 0, posHash)); + while (stack.TryPop(out var frame)) + { + var (next, deferred) = ProcessFrame( + frame, keyPath, colName, colNameUtf8, + rows, keyOrder, keySet, columnTypes, keyObservedCount); + if (deferred is { } d) + { + stack.Push(d); + } + + if (next is { } n) + { + stack.Push(n); + } + } } /// @@ -53,51 +67,32 @@ public static string LastKeySegment(IReadOnlyList keyPath) return "value"; } - private static void TraverseKeyPath( - JsonRawBytes currentBytes, - IReadOnlyList keyPath, - int segmentIndex, - string posHash, - string colName, - byte[] colNameUtf8, - List rows, - List keyOrder, - HashSet keySet, - Dictionary columnTypes, - Dictionary keyObservedCount) - { - if (segmentIndex == keyPath.Count) - { - CollectLeafRows(currentBytes, posHash, colName, colNameUtf8, rows, keyOrder, keySet, columnTypes, keyObservedCount); - return; - } - - var segment = keyPath[segmentIndex]; - - if (segment.Kind == KeyPathSegmentKind.Index) - { - TraverseIndexSegment( - currentBytes, keyPath, segmentIndex, posHash, colName, colNameUtf8, - rows, keyOrder, keySet, columnTypes, keyObservedCount); - return; - } - - var valueBytes = FindValueByKey(currentBytes, segment.Value); - if (valueBytes is null) - { - return; // Key absent, or current value is not an object — skip record silently. - } + private enum FrameKind { Descend, ContinueArray } - TraverseKeyPath( - valueBytes.Value, keyPath, segmentIndex + 1, posHash, colName, colNameUtf8, - rows, keyOrder, keySet, columnTypes, keyObservedCount); + /// + /// Pending descent work. applies the segment at + /// to . + /// resumes an index segment's array scan from ; it is returned as the + /// deferred frame after each element so only O(depth) frames stay live, never one per sibling. + /// + private readonly record struct TraversalFrame( + FrameKind Kind, + JsonRawBytes Bytes, + int SegmentIndex, + int ElementIndex, + string PosHash, + JsonReaderState ReaderState) + { + public static TraversalFrame Descend(JsonRawBytes bytes, int segmentIndex, string posHash) => + new(FrameKind.Descend, bytes, segmentIndex, 0, posHash, default); } - private static void TraverseIndexSegment( - JsonRawBytes currentBytes, + // Returns the frame to descend next (its subtree processed first) and the frame to resume + // afterward (an array scan continuation). The caller pushes deferred, then next, so the LIFO + // stack finishes next's subtree before resuming deferred — preserving forward DFS order. + private static (TraversalFrame? next, TraversalFrame? deferred) ProcessFrame( + TraversalFrame frame, IReadOnlyList keyPath, - int segmentIndex, - string posHash, string colName, byte[] colNameUtf8, List rows, @@ -106,193 +101,89 @@ private static void TraverseIndexSegment( Dictionary columnTypes, Dictionary keyObservedCount) { - var reader = new Utf8JsonReader(currentBytes.Span); - if (!reader.Read() || reader.TokenType != JsonTokenType.StartArray) - { - return; // Wrong type at this path position — skip record silently. - } - - if (segmentIndex == keyPath.Count - 1) + if (frame.Kind == FrameKind.ContinueArray) { - // A trailing index segment expands the same array that would be reached by - // selecting it directly as the leaf (e.g. "tags" and "tags[0]" must produce - // identical output, including the "value" column for primitive elements). - CollectArrayLeafRows(currentBytes, posHash, rows, keyOrder, keySet, columnTypes, keyObservedCount); - return; - } - - var elementIndex = 0; - while (reader.Read()) - { - if (reader.TokenType == JsonTokenType.EndArray) - { - break; - } - - if (reader.CurrentDepth != 1) - { - continue; - } - - var elementBytes = ExtractElementBytes(ref reader, currentBytes); - TraverseKeyPath( - elementBytes, keyPath, segmentIndex + 1, $"{posHash}:{elementIndex}", colName, colNameUtf8, - rows, keyOrder, keySet, columnTypes, keyObservedCount); - elementIndex++; + var reader = new Utf8JsonReader(frame.Bytes.Span, isFinalBlock: true, frame.ReaderState); + return ScanOneArrayElement(ref reader, frame.Bytes, frame.SegmentIndex, frame.ElementIndex, frame.PosHash); } - } - private static void CollectLeafRows( - JsonRawBytes leafBytes, - string posHash, - string colName, - byte[] colNameUtf8, - List rows, - List keyOrder, - HashSet keySet, - Dictionary columnTypes, - Dictionary keyObservedCount) - { - var reader = new Utf8JsonReader(leafBytes.Span); - if (!reader.Read()) + if (frame.SegmentIndex == keyPath.Count) { - return; + KeyPathLeafCollector.CollectLeafRows( + frame.Bytes, frame.PosHash, colName, colNameUtf8, rows, keyOrder, keySet, columnTypes, keyObservedCount); + return (null, null); } - if (reader.TokenType == JsonTokenType.StartObject) + var segment = keyPath[frame.SegmentIndex]; + if (segment.Kind == KeyPathSegmentKind.Index) { - rows.Add(new FocusedTableRow(leafBytes, posHash)); - var observedKeys = new HashSet(StringComparer.Ordinal); - SchemaScanner.ScanObject(leafBytes.Span, keyOrder, keySet, columnTypes, observedKeys); - SchemaScanner.IncrementObservationCounts(observedKeys, keyObservedCount); - return; + return ExpandIndexSegment(frame, keyPath, rows, keyOrder, keySet, columnTypes, keyObservedCount); } - if (reader.TokenType == JsonTokenType.StartArray) + var valueBytes = KeyPathLeafCollector.FindValueByKey(frame.Bytes, segment.Value); + if (valueBytes is null) { - CollectArrayLeafRows(leafBytes, posHash, rows, keyOrder, keySet, columnTypes, keyObservedCount); - return; + return (null, null); // Key absent, or current value is not an object — skip record silently. } - // Primitive leaf (including null) — synthesize a single-key object so - // JsonObjectCellExtractor can extract it without modification. - // Note: ScanObject is NOT called here, so no type inference is performed; - // the synthesized column always receives ColumnType.Text (Phase 2 limitation). - var synthBytes = SynthesizeObject(colNameUtf8, leafBytes.Span); - rows.Add(new FocusedTableRow(synthBytes, posHash)); - SchemaScanner.RegisterKeyIfNew(colName, keyOrder, keySet); - SchemaScanner.IncrementObservationCounts([colName], keyObservedCount); + return (TraversalFrame.Descend(valueBytes.Value, frame.SegmentIndex + 1, frame.PosHash), null); } - private static void CollectArrayLeafRows( - JsonRawBytes leafBytes, - string posHash, + private static (TraversalFrame? next, TraversalFrame? deferred) ExpandIndexSegment( + TraversalFrame frame, + IReadOnlyList keyPath, List rows, List keyOrder, HashSet keySet, Dictionary columnTypes, Dictionary keyObservedCount) { - var reader = new Utf8JsonReader(leafBytes.Span); + var reader = new Utf8JsonReader(frame.Bytes.Span); if (!reader.Read() || reader.TokenType != JsonTokenType.StartArray) { - return; - } - - var elementIndex = 0; - while (reader.Read()) - { - if (reader.TokenType == JsonTokenType.EndArray) - { - break; - } - - if (reader.CurrentDepth != 1) - { - continue; - } - - var isObjectElement = reader.TokenType == JsonTokenType.StartObject; - var elementBytes = ExtractElementBytes(ref reader, leafBytes); - var elementHash = $"{posHash}:{elementIndex}"; - - if (isObjectElement) - { - rows.Add(new FocusedTableRow(elementBytes, elementHash)); - var observedKeys = new HashSet(StringComparer.Ordinal); - SchemaScanner.ScanObject(elementBytes.Span, keyOrder, keySet, columnTypes, observedKeys); - SchemaScanner.IncrementObservationCounts(observedKeys, keyObservedCount); - elementIndex++; - continue; - } - - // Primitive element (including null) — synthesize {"value": element}. - var synthBytes = SynthesizeObject("value"u8, elementBytes.Span); - rows.Add(new FocusedTableRow(synthBytes, elementHash)); - SchemaScanner.RegisterKeyIfNew("value", keyOrder, keySet); - SchemaScanner.IncrementObservationCounts(["value"], keyObservedCount); - elementIndex++; + return (null, null); // Wrong type at this path position — skip record silently. } - } - private static JsonRawBytes ExtractElementBytes(ref Utf8JsonReader reader, JsonRawBytes containingBytes) - { - if (reader.TokenType is JsonTokenType.StartObject or JsonTokenType.StartArray) + if (frame.SegmentIndex == keyPath.Count - 1) { - return JsonByteExtractor.ExtractNestedBytes(ref reader, containingBytes); + // A trailing index segment expands the same array that would be reached by selecting + // it directly as the leaf (e.g. "tags" and "tags[0]" must produce identical output). + KeyPathLeafCollector.CollectArrayLeafRows(frame.Bytes, frame.PosHash, rows, keyOrder, keySet, columnTypes, keyObservedCount); + return (null, null); } - var start = (int)reader.TokenStartIndex; - var end = (int)reader.BytesConsumed; - return containingBytes.Slice(start, end - start); + return ScanOneArrayElement(ref reader, frame.Bytes, frame.SegmentIndex, 0, frame.PosHash); } - private static JsonRawBytes? FindValueByKey(JsonRawBytes objectBytes, string key) + // Reads the next depth-1 element and returns it as the next frame to descend, plus the + // ContinueArray continuation for the remaining siblings as the deferred frame. + private static (TraversalFrame? next, TraversalFrame? deferred) ScanOneArrayElement( + ref Utf8JsonReader reader, + JsonRawBytes arrayBytes, + int segmentIndex, + int elementIndex, + string posHash) { - var reader = new Utf8JsonReader(objectBytes.Span); - if (!reader.Read() || reader.TokenType != JsonTokenType.StartObject) - { - return null; - } - while (reader.Read()) { - if (reader.TokenType == JsonTokenType.EndObject) - { - return null; - } - - if (reader.TokenType != JsonTokenType.PropertyName) + if (reader.TokenType == JsonTokenType.EndArray) { - continue; + return (null, null); } - if (!reader.ValueTextEquals(key)) + if (reader.CurrentDepth != 1) { - reader.Skip(); continue; } - if (!reader.Read()) - { - return null; - } - - return ExtractElementBytes(ref reader, objectBytes); + var elementBytes = JsonByteExtractor.ExtractValueBytes(ref reader, arrayBytes); + var remainder = arrayBytes.Slice((int)reader.BytesConsumed); + var next = TraversalFrame.Descend(elementBytes, segmentIndex + 1, $"{posHash}:{elementIndex}"); + var deferred = new TraversalFrame( + FrameKind.ContinueArray, remainder, segmentIndex, elementIndex + 1, posHash, reader.CurrentState); + return (next, deferred); } - return null; - } - - private static JsonRawBytes SynthesizeObject(ReadOnlySpan keyUtf8, ReadOnlySpan valueBytes) - { - var buffer = new ArrayBufferWriter(); - using var writer = new Utf8JsonWriter(buffer); - writer.WriteStartObject(); - writer.WritePropertyName(keyUtf8); - writer.WriteRawValue(valueBytes, skipInputValidation: true); - writer.WriteEndObject(); - writer.Flush(); - return buffer.WrittenMemory; + return (null, null); } } diff --git a/src/Engine/IO/Json/JsonByteExtractor.cs b/src/Engine/IO/Json/JsonByteExtractor.cs index f3dbea7..c6dc210 100644 --- a/src/Engine/IO/Json/JsonByteExtractor.cs +++ b/src/Engine/IO/Json/JsonByteExtractor.cs @@ -44,6 +44,23 @@ public static JsonRawBytes ExtractNestedBytes( return rawJson.Slice(startPosition, endPosition - startPosition); } + /// + /// Returns the raw bytes of the value at the reader's current token. Nested structures + /// (Object/Array) are read through ; primitives are sliced + /// directly from between the token start and consumed offset. + /// + public static JsonRawBytes ExtractValueBytes(ref Utf8JsonReader reader, JsonRawBytes containingBytes) + { + if (reader.TokenType is JsonTokenType.StartObject or JsonTokenType.StartArray) + { + return ExtractNestedBytes(ref reader, containingBytes); + } + + var start = (int)reader.TokenStartIndex; + var end = (int)reader.BytesConsumed; + return containingBytes.Slice(start, end - start); + } + /// /// Counts the top-level properties of a JSON object by tracking brace/bracket depth. /// The reader must be positioned at a token; on return diff --git a/tests/Refedle.Tests/Engine/IO/DrillDown/KeyPathTraverserTests.cs b/tests/Refedle.Tests/Engine/IO/DrillDown/KeyPathTraverserTests.cs index f897634..a484b09 100644 --- a/tests/Refedle.Tests/Engine/IO/DrillDown/KeyPathTraverserTests.cs +++ b/tests/Refedle.Tests/Engine/IO/DrillDown/KeyPathTraverserTests.cs @@ -20,6 +20,61 @@ private static JsonElement GetProperty(JsonRawBytes bytes, string propertyName) return doc.RootElement.GetProperty(propertyName).Clone(); } + // Builds {"k0":{"k1":...{"k{depth-1}":"leafValue"}...}} — `depth` single-key objects, + // exercising key-segment descent to a primitive leaf. + private static string BuildNestedObjects(int depth, string leafValue) + { + var sb = new StringBuilder(); + for (var i = 0; i < depth; i++) + { + sb.Append("{\"k").Append(i).Append("\":"); + } + + sb.Append('"').Append(leafValue).Append('"'); + sb.Append(new string('}', depth)); + return sb.ToString(); + } + + // Builds [[["leafValue"]]] — `depth` single-element arrays; each "[0]" descends one level. + private static string BuildNestedArrays(int depth, string leafValue) + { + var sb = new StringBuilder(); + for (var i = 0; i < depth; i++) + { + sb.Append('['); + } + + sb.Append('"').Append(leafValue).Append('"'); + for (var i = 0; i < depth; i++) + { + sb.Append(']'); + } + + return sb.ToString(); + } + + private static List BuildKeySegments(int depth) + { + List segments = []; + for (var i = 0; i < depth; i++) + { + segments.Add(Key($"k{i}")); + } + + return segments; + } + + private static List BuildIndexSegments(int depth) + { + List segments = []; + for (var i = 0; i < depth; i++) + { + segments.Add(Index("[0]")); + } + + return segments; + } + private static TraverseResult Traverse( JsonRawBytes recordBytes, IReadOnlyList keyPath, string posHash = "1") { @@ -273,6 +328,65 @@ public void ExtractRows_PathEndingInIndexSegment_ProducesSameRowsAsPathWithoutIt resultWithoutIndex.KeyOrder.Should().Equal(resultWithIndex.KeyOrder); } + [Fact] + public void ExtractRows_DeeplyNestedKeySegments_ReachesLeafWithExpectedAggregate() + { + // Arrange — depth 25 stays under Utf8JsonReader's default MaxDepth (64); the recursive + // implementation must reach the leaf. Pins behavior the iterative rewrite must preserve. + const int depth = 25; + var bytes = Bytes(BuildNestedObjects(depth, "leafKey")); + var keyPath = BuildKeySegments(depth); + + // Act + var result = Traverse(bytes, keyPath); + + // Assert + result.Rows.Should().HaveCount(1); + result.Rows[0].HashValue.Should().Be("1"); + GetProperty(result.Rows[0].Bytes, $"k{depth - 1}").GetString().Should().Be("leafKey"); + result.KeyOrder.Should().Equal($"k{depth - 1}"); + result.KeyObservedCount[$"k{depth - 1}"].Should().Be(1); + } + + [Fact] + public void ExtractRows_DeeplyNestedIndexSegments_ReachesLeafWithExpectedAggregate() + { + // Arrange — each "[0]" descends one array level; the leaf hash accumulates ":0" per + // segment. The iterative rewrite must reproduce this hash and aggregation exactly. + const int depth = 25; + var expectedLeafHash = "1" + string.Concat(Enumerable.Repeat(":0", depth)); + var bytes = Bytes(BuildNestedArrays(depth, "leafIndex")); + var keyPath = BuildIndexSegments(depth); + + // Act + var result = Traverse(bytes, keyPath); + + // Assert + result.Rows.Should().HaveCount(1); + result.Rows[0].HashValue.Should().Be(expectedLeafHash); + GetProperty(result.Rows[0].Bytes, "value").GetString().Should().Be("leafIndex"); + result.KeyOrder.Should().Equal("value"); + result.KeyObservedCount["value"].Should().Be(1); + } + + [Fact] + public void ExtractRows_BranchingIndexSegments_VisitsLeavesInForwardDfsOrder() + { + // Arrange — orders[*].items[*].id: two index levels, each with more than one sibling, + // pins the iterative DFS to fully visit element 0's subtree before element 1's. + var bytes = Bytes("""{"orders":[{"items":[{"id":"a1"},{"id":"a2"}]},{"items":[{"id":"b1"}]}]}"""); + IReadOnlyList keyPath = [Key("orders"), Index("[0]"), Key("items"), Index("[0]"), Key("id")]; + + // Act + var result = Traverse(bytes, keyPath); + + // Assert + result.Rows.Should().HaveCount(3); + result.Rows.Select(r => r.HashValue).Should().Equal("1:0:0", "1:0:1", "1:1:0"); + result.Rows.Select(r => GetProperty(r.Bytes, "id").GetString()).Should().Equal("a1", "a2", "b1"); + result.KeyOrder.Should().Equal("id"); + } + [Fact] public void LastKeySegment_PathEndingInKeySegment_ReturnsThatSegmentValue() {