fix: Preserve inner element types in PyArrow schema inference and optimize JSON nested list detection

soooojinlee · claude · soooojinlee · commit e558c814c5d0 · 2026-04-02T00:24:07.000+09:00
- Add _parse_pa_type_str() to reconstruct PyArrow types from type strings
  for VALUE_LIST/VALUE_SET, avoiding lossy round-trip through placeholder
- Optimize proto_json nested list detection: only scan with any() when
  first element is None, avoiding O(n) scan for flat lists
- Add warning log for unrecognized PyArrow type strings

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
Signed-off-by: soojin &lt;soojin@dable.io&gt;
diff --git a/sdk/python/feast/infra/offline_stores/offline_utils.py b/sdk/python/feast/infra/offline_stores/offline_utils.py
@@ -1,3 +1,4 @@
+import logging
 import uuid
 from dataclasses import asdict, dataclass
 from datetime import datetime, timedelta, timezone
@@ -21,6 +22,7 @@
 from feast.repo_config import RepoConfig
 from feast.type_map import feast_value_type_to_pa
 from feast.utils import _get_requested_feature_views_to_features_dict, to_naive_utc
+from feast.value_type import ValueType
 
 DEFAULT_ENTITY_DF_EVENT_TIMESTAMP_COL = "event_timestamp"
 
@@ -241,6 +243,37 @@ def get_offline_store_from_config(offline_store_config: Any) -> OfflineStore:
     return offline_store_class()
 
 
+_PA_BASIC_TYPES = {
+    "int32": pa.int32(),
+    "int64": pa.int64(),
+    "double": pa.float64(),
+    "float": pa.float32(),
+    "string": pa.string(),
+    "binary": pa.binary(),
+    "bool": pa.bool_(),
+    "large_string": pa.large_string(),
+    "null": pa.null(),
+}
+
+
+def _parse_pa_type_str(pa_type_str: str) -> pa.DataType:
+    """Parse a PyArrow type string to preserve inner element types for nested lists."""
+    pa_type_str = pa_type_str.strip()
+    if pa_type_str.startswith("list<item: ") and pa_type_str.endswith(">"):
+        inner = pa_type_str[len("list<item: ") : -1]
+        return pa.list_(_parse_pa_type_str(inner))
+    if pa_type_str in _PA_BASIC_TYPES:
+        return _PA_BASIC_TYPES[pa_type_str]
+    if pa_type_str.startswith("timestamp"):
+        return pa.timestamp("us")
+    logger = logging.getLogger(__name__)
+    logger.warning(
+        "Unrecognized PyArrow type string '%s', falling back to pa.string()",
+        pa_type_str,
+    )
+    return pa.string()
+
+
 def get_pyarrow_schema_from_batch_source(
     config: RepoConfig, batch_source: DataSource, timestamp_unit: str = "us"
 ) -> Tuple[pa.Schema, List[str]]:
@@ -250,15 +283,12 @@ def get_pyarrow_schema_from_batch_source(
     pa_schema = []
     column_names = []
     for column_name, column_type in column_names_and_types:
-        pa_schema.append(
-            (
-                column_name,
-                feast_value_type_to_pa(
-                    batch_source.source_datatype_to_feast_value_type()(column_type),
-                    timestamp_unit=timestamp_unit,
-                ),
-            )
-        )
+        value_type = batch_source.source_datatype_to_feast_value_type()(column_type)
+        if value_type in (ValueType.VALUE_LIST, ValueType.VALUE_SET):
+            pa_type = _parse_pa_type_str(column_type)
+        else:
+            pa_type = feast_value_type_to_pa(value_type, timestamp_unit=timestamp_unit)
+        pa_schema.append((column_name, pa_type))
         column_names.append(column_name)
 
     return pa.schema(pa_schema), column_names
diff --git a/sdk/python/feast/proto_json.py b/sdk/python/feast/proto_json.py
@@ -92,7 +92,9 @@ def from_json_object(
             if len(value) == 0:
                 # Clear will mark the struct as modified so it will be created even if there are no values
                 message.int64_list_val.Clear()
-            elif isinstance(value[0], list) or any(isinstance(v, list) for v in value):
+            elif isinstance(value[0], list) or (
+                value[0] is None and any(isinstance(v, list) for v in value)
+            ):
                 # Nested collection (list of lists).
                 # Check any() to handle cases where the first element is None
                 # (empty inner collections round-trip through proto as None).
