enhance: [reader] add snapshot-backed read entrypoint

[congqixia]

Add MilvusReadApp and a spark-submit wrapper for client and snapshot reads, and prefer client-created snapshots with fallback only when snapshot RPCs are unavailable.

Reuse snapshot planning for V2/V3 segments, expose real row_id/timestamp plus $segment_id/$row_offset metadata, and make Arrow string conversion handle JSON stored as binary.

Update Milvus proto for snapshot RPCs and add focused tests for the read app, snapshot planner helpers, and Arrow VarBinary StringType conversion.

Also: reader/backfill metadata columns are now named $segment_id and $row_offset instead of segment_id and row_offset.

[sre-ci-robot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: congqixia
To complete the pull request process, please assign xiaofan-luan after the PR has been reviewed.
You can assign the PR to them by writing /assign @xiaofan-luan in a comment when ready.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[liliu-z]

MilvusDataSource.scala:540 In snapshot mode, MilvusTable.schema() (line 348) returns the snapshot-inferred schema directly, skipping the extra-column append logic at lines 391–419. The inferred schema never includes $segment_id/$row_offset. When pruneColumns() appends these names from milvus.extra.columns (lines 524–533), line 540 calls .find(...).get on None, throwing NoSuchElementException. This is a guaranteed crash on the newly introduced MilvusReadApp --mode snapshot --extra-columns code path.

Fix: Have MilvusTable.schema() append extra columns in the snapshot branch before returning (mirroring lines 391–419), or change line 540 to filter out names not found in the schema.

[liliu-z]

MilvusSnapshotReader.scala:645 Three things conspire to produce nulls for V3 system fields in snapshot mode: (1) inferSchema() uses toSparkSchema(includeSystemFields=true) producing field names "RowID"/"Timestamp" from snapshot JSON; (2) toProtobufSchemaBytes() (line 645) strips RowID/Timestamp from protobuf schema bytes; (3) V3 reader's fieldNameToId (MilvusLoonPartitionReader.scala:64–65) hardcodes only lowercase "row_id" → 0, "timestamp" → 1. Since the protobuf schema lacks system fields, fieldNameToId never maps "RowID" → 0, so columns are excluded and return null. Not a regression (was null before), but the PR claims to "expose real row_id/timestamp." The V2 reader is NOT affected due to its compatible casing.

[liliu-z]

MilvusClient.scala:887 The check message.contains("unimplemented") on line 887 matches any exception in the cause chain containing this word, potentially triggering a silent fallback to legacy mode for unrelated errors (e.g., "Feature X is unimplemented for storage backend Y"). The gRPC status code check on lines 874–881 is correct and sufficient for standard gRPC errors. This is a plausibility-based concern rather than a proven false-positive case, but the fix is trivial.

[liliu-z]

MilvusDataSource.scala:1026 Both MilvusDataSource.scala:1026–1042 and MilvusReadApp.scala:292–316 read entire files into memory with no size limit. A misconfigured snapshot path pointing to a large file could cause OOM on the driver.

[liliu-z]

MilvusDataSource.scala:977 Nearly identical bucket-scoped S3A configuration code exists in MilvusDataSource.scala:977–1024 and MilvusReadApp.scala:318–353. This duplication increases maintenance burden and divergence risk.

[liliu-z]

src/main/scala/operations/backfill/BackfillApp.scala:219 The check at line 219 is i + 1 < args.length but does not verify !args(i + 1).startsWith("--"). So --mode --snapshot-path silently accepts "--snapshot-path" as the value of mode instead of raising a missing-value error. This turns a simple CLI typo into a confusing downstream failure. In contrast, MilvusReadApp.parseArgs (line 102) correctly checks for this.

[liliu-z]

src/main/scala/operations/backfill/README.md:248 Lines 248-253 reference segment_id / row_offset without the $ prefix. After this PR, the actual column names are $segment_id / $row_offset. Users following the README will get incorrect column references.

[liliu-z]

At line 670, the inner catch calls dropSnapshot(...) and ignores the Try result. If the drop fails, there is no warning logged, making it impossible to diagnose orphaned snapshots.

Fix: Log a warning when the cleanup dropSnapshot call fails.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/MilvusDataSource.scala:0 When the connector detects a legacy server that doesn't support snapshots and falls back, it drops filter pushdown entirely. This is a performance regression for all users on older Milvus versions — queries that previously pushed filters down to the server now scan the full dataset on the Spark side.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/SchemaUtil.scala:0 When field IDs are missing and the code falls back to idx + 1, the generated IDs can collide with the hard-coded RowID (fieldID=0) and Timestamp (fieldID=1) system fields. If a user field occupies fieldID=0, the system RowID field is silently discarded, causing data misalignment at read time. Use idx + 100 or another offset that avoids system-reserved IDs.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/MilvusOption.scala:0 MilvusOption is serialized as part of the InputPartition closure and shipped to Spark executors. This means access keys, secret keys, and tokens are embedded in the serialized task payload, which may be logged, persisted to disk, or visible in the Spark UI. Sensitive credentials must be stripped before serialization or fetched independently on the executor side.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/ArrowConverter.scala:0 ArrowConverter performs unchecked casts that will throw ClassCastException at runtime if the actual Arrow vector type doesn't match expectations. Additionally, BinaryVector values (unsigned bytes 0-255) are mapped to Spark's ByteType (signed -128 to 127), silently reinterpreting values above 127 as negative numbers. This is a silent data corruption issue.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/ArrowConverter.scala:0 Spark's BinaryType is already byte[]. Wrapping it in ArrayType produces Array[byte[]], which doesn't correspond to any valid Milvus type. This branch is either dead code (never reached) or produces incorrect results if reached. It should be deleted.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/SchemaUtil.scala:0 SchemaUtil uses case-sensitive existingNames.contains("RowID") to check for system field conflicts, but MilvusLoonPartitionReader's alias mapping is case-insensitive. If a user defines a field named row_id (lowercase), it bypasses the guard, and the system appends a duplicate RowID field to the Arrow Schema. This can cause two fields with different IDs but semantically the same name, leading to ambiguous reads.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/MilvusClient.scala:1129 At lines 1129-1144, the retry interceptor calls super.onClose(status, trailers) before initiating the retry. This fires the close event on the downstream listener prematurely, meaning the application sees a failure notification and then a new RPC silently starts behind its back. This violates gRPC's listener state machine contract and can cause double-processing or incorrect error handling.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/MilvusDataSource.scala:713 When dropClientReadSnapshot fails, it retries up to 3 times in a tight loop with no Thread.sleep, backoff, or jitter. During transient network issues, these immediate retries are ineffective and add unnecessary load. For distributed cleanup operations, even a simple fixed delay would significantly improve resilience.

[liliu-z]

src/main/scala/zilliztech/spark/milvus/MilvusPackedV2PartitionReader.scala:109 When a requested column is not present in the Column Group, the reader silently skips it. If this happens due to an upstream planning error, the query returns incomplete data with no indication of the missing column. This makes debugging extremely difficult in production.

[liliu-z]

The interceptor wraps a single ClientCall from next.newCall(...) at line 1104. On retry, executeCall() recursively invokes super.start(...) on that same call from inside the listener's onClose (line 1157). gRPC's ClientCall is single-use — the second start() will throw IllegalStateException("call already started"). Combined with the removal of the pre-retry super.onClose(...), the interceptor also never delivers the original failure to the caller, so the first transient UNAVAILABLE swallows the result silently. This is worse than having no retry interceptor. Either delete the custom interceptor and use gRPC's built-in retryPolicy via ManagedChannelBuilder.enableRetry() + service config, or rewrite to allocate a fresh next.newCall(...) per attempt.

[liliu-z]

The test at line 41 expects BinaryVector code 100 to map to ArrayType(ByteType), but the implementation at MilvusSnapshotReader.scala:625 returns BinaryType. The implementation matches the PR description and DataTypeUtilTest. This test will fail in CI. The commit 0689668 ("Address comment") appears to have flipped the expectation in the wrong direction.

[liliu-z]

Line 37 filters lowercase aliases (row_id/timestamp) but not the CamelCase canonical names RowID/Timestamp. If a user collection has a field named RowID with fieldID 100, the mapping produces Map(0L → "RowID", 100L → "RowID"); inverting to fieldNameToId is order-dependent and non-deterministic. The same issue exists in V3 at MilvusLoonPartitionReader.scala:64.

[liliu-z]

Same issue as the V2 reader: line 64 handles lowercase aliases but not the literal canonical names RowID/Timestamp. If Milvus allows those as user field names, system-field access becomes ambiguous.

[liliu-z]

The vector-search branch at line 271 returns rows from vectorSearchResults.next() without setting _lastReturnedRowOffset, which defaults to -1L. The wrapper at MilvusPartitionReaderFactory.scala:97 reads underlyingReader.lastReturnedRowOffset, so $row_offset is always -1 for vector search results. Either carry the segment-local row offset alongside the row+distance pair, or reject the combination with a clear error.

[liliu-z]

The V3 path emits p.partitionName (string name like "_default") while the V2 path emits p.partitionID.toString (numeric ID as string). A snapshot read can mix V2 + V3 segments, so users selecting the partition extra column get heterogeneous values from the same column depending on segment format. Both paths must agree — either both emit the partition name or both emit the ID.

[liliu-z]

src/main/scala/MilvusDataSource.scala:0 Snapshot cleanup is tied to SparkListenerSQLExecutionEnd only. SIGKILL, OOM, or yarn application -kill leaves orphan snapshot records on the Milvus server. The driver-death case cannot be fully solved in Spark client code, but the PR needs either a documented Milvus-side TTL/sweeper guarantee or a user-visible warning that snapshots may accumulate after unclean driver exits.

[liliu-z]

The V3 wrapper reads underlyingReader.lastReturnedRowOffset (real segment-relative offset accounting for filter-skipped rows). The V2 wrapper uses a wrapper-local var rowOffset counter that increments per get() — a post-filter sequential index, not the segment-relative offset. If downstream code uses $row_offset to re-order rows for backfill (the documented use case), V2-segment rows will be misordered relative to V3-segment rows in the same collection.

[liliu-z]

The fallback that checks for substring matches in UNKNOWN-status error messages is fragile across proxy versions. It would not match if the message is ever rephrased (e.g., "method not registered" instead of the current wording).

[liliu-z]

src/main/scala/MilvusReadApp.scala:0 Passing S3 credentials via --s3-access-key and --s3-secret-key CLI arguments makes them visible to every process on the host via ps or /proc. The CLI app should also accept AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY environment variables or a credentials file as alternatives. At minimum, document "use --use-iam in production" prominently.

[liliu-z]

Sets fs.s3a.bucket.$bucket.path.style.access=true and fs.s3a.bucket.$bucket.connection.ssl.enabled=false unconditionally, overriding any explicit global Hadoop config the user expects to inherit. The fs.s3a.impl check preserves existing values; these two keys should follow the same pattern or be configurable.

[liliu-z]

A valid empty collection triggers IllegalArgumentException("Snapshot mode has no StorageV3 manifests or StorageV2 segments"). This should return an empty DataFrame (Array.empty[InputPartition]). The guard is redundant for safety since parseSnapshotMetadata already fails loudly on unparseable data. Distinguish "zero segments" (return empty) from "unparseable snapshot" (fail loud) and remove this overly aggressive guard.

[liliu-z]

src/main/scala/MilvusPartitionReaderFactory.scala:94 The partition extra column is emitted as a raw Java String instead of UTF8String.fromString(...). Spark's InternalRow contract requires UTF8String for string-type columns; a raw String will cause a ClassCastException or silent data corruption at read time. The same issue exists in the V3 wrapper at line 168. Wrap the value with UTF8String.fromString(...).

[liliu-z]

src/main/scala/MilvusDataSource.scala:384 When the Milvus collection already contains fields named row_id or timestamp, MilvusDataSource.schema() (lines 384-410) will produce a schema with duplicate column names — one from the system fields and one from the user fields. This causes ambiguous column resolution at query time. Apply the same conflict/dedup policy that SchemaUtil already uses on the Arrow side.

[liliu-z]

src/main/scala/MilvusScanBuilder.scala:0 The pruneColumns implementation derives field IDs positionally (index < 2 → fieldID = 0/1), which assumes system fields (RowID, Timestamp) are always the first two columns. When an external schema omits system fields, the reader will silently project RowID/Timestamp instead of the requested user fields, returning wrong column data. Source field IDs from milvusCollection.schema.fields(...).fieldID instead of position.

[liliu-z]

Both overloads of isSnapshotMode use .contains to check whether the snapshot manifest key is present. Setting milvus.snapshot.manifests="" will trigger snapshot mode even though no manifests are specified, which is fragile. Filter out empty/blank values before treating the option as present.

[liliu-z]

When an InterruptedException is caught during the retry sleep at line 713, the interrupt flag is cleared. The method then throws without restoring Thread.currentThread().interrupt(), which violates the Java interrupt contract and can cause callers to miss the interrupt signal. Catch InterruptedException, restore the interrupt flag, and stop retrying.

[liliu-z]

The gRPC channel builder calls .enableRetry() but no defaultServiceConfig is installed, so the retry policy is never actually applied. Either install a proper service config with retry parameters, or remove the misleading .enableRetry() call.