Convert column children computation from lazy to eager (rapidsai#20953)

This change makes the children match the data and mask with eager computation, which in turn makes it easier to reason about when data accesses can happen and will simplify the rewriting of other internals around pylibcudf and buffers.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Matthew Roeschke (https://github.com/mroeschke)
  - Bradley Dice (https://github.com/bdice)

URL: rapidsai#20953

Author: vyasr

python/cudf/cudf/core/column/categorical.py

@@ -132,11 +132,11 @@ def _process_values_for_isin(
         # Convert values to categorical dtype like self
         return self, column.as_column(values, dtype=self.dtype)
 
-    @property
-    def children(self) -> tuple[NumericalColumn]:
+    def _recompute_children(self) -> None:
         if self.offset == 0 and self.size == self.base_size:
-            return super().children  # type: ignore[return-value]
-        if self._children is None:
+            # Optimization: for non-sliced columns, children == base_children (just references)
+            self._children = self.base_children  # type: ignore[assignment]
+        else:
             # Pass along size, offset, null_count from __init__
             # which doesn't necessarily match the attributes of plc_column
             # Use base_children[0]'s plc_column as it has the actual codes data
@@ -148,6 +148,9 @@ def children(self) -> tuple[NumericalColumn]:
                 exposed=False,
             )
             self._children = (child,)
+
+    @property
+    def children(self) -> tuple[NumericalColumn]:
         return self._children
 
     @property

python/cudf/cudf/core/column/column.py

@@ -327,6 +327,7 @@ def __init__(
         # triggering the destruction of the exposed buffers.
         self._exposed_buffers: set[Buffer] = set()
         self._recompute_data()
+        self._recompute_children()
 
     def _get_children_from_pylibcudf_column(
         self,
@@ -500,6 +501,7 @@ def set_base_mask(self, value: None | Buffer) -> None:
         )
 
         self._clear_cache()
+        self._recompute_children()
 
     def _clear_cache(self) -> None:
         self._distinct_count.clear()
@@ -557,23 +559,8 @@ def base_children(self) -> tuple[ColumnBase, ...]:
 
     @property
     def children(self) -> tuple[ColumnBase, ...]:
-        if self.offset == 0 and self.size == self.base_size:
-            self._children = self.base_children
-        if self._children is None:
-            if not self.base_children:
-                self._children = ()
-            else:
-                # Compute children from the column view (children factoring self.size)
-                children = ColumnBase.from_pylibcudf(
-                    self.plc_column.copy()
-                ).base_children
-                dtypes = (
-                    base_child.dtype for base_child in self.base_children
-                )
-                self._children = tuple(
-                    child._with_type_metadata(dtype)
-                    for child, dtype in zip(children, dtypes, strict=True)
-                )
+        """Return the offset-aware children columns."""
+        assert self._children is not None
         return self._children
 
     def set_base_children(self, value: tuple[ColumnBase, ...]) -> None:
@@ -584,7 +571,6 @@ def set_base_children(self, value: tuple[ColumnBase, ...]) -> None:
         if any(not isinstance(child, ColumnBase) for child in value):
             raise TypeError("All children must be Columns.")
 
-        self._children = None
         self._base_children = value
 
     def _recompute_data(self) -> None:
@@ -600,6 +586,21 @@ def _recompute_data(self) -> None:
             end = start + self.size * self.dtype.itemsize
             self._data = self.base_data[start:end]
 
+    def _recompute_children(self) -> None:
+        """Recompute the offset-aware children columns."""
+        # Check for empty base_children first to avoid accessing properties on empty collections
+        if not self.base_children:
+            self._children = ()
+        elif self.offset == 0 and self.size == self.base_size:
+            # Optimization: for non-sliced columns, children == base_children (just references)
+            self._children = self.base_children
+        else:
+            # Slice each base child using the parent's offset and size
+            self._children = tuple(
+                base_child.slice(self.offset, self.offset + self.size)
+                for base_child in self.base_children
+            )
+
     def _mimic_inplace(
         self, other_col: Self, inplace: bool = False
     ) -> None | Self:
@@ -614,8 +615,8 @@ def _mimic_inplace(
             self.plc_column = other_col.plc_column
             self._base_children = other_col._base_children
             self._recompute_data()
+            self._recompute_children()
             self._mask = None
-            self._children = None
             self._clear_cache()
             return None
         else:

python/cudf/cudf/core/column/lists.py

@@ -76,6 +76,40 @@ def _get_children_from_pylibcudf_column(
             children[1]._with_type_metadata(dtype.element_type),
         )
 
+    def _recompute_children(self) -> None:
+        """Recompute the offset-aware children columns with proper type metadata."""
+        if not self.base_children:
+            self._children = ()
+        elif self.offset == 0 and self.size == self.base_size:
+            # Optimization: for non-sliced columns, children == base_children
+            self._children = self.base_children  # type: ignore[assignment]
+        else:
+            # List columns have special structure:
+            # - Child 0 (offsets): size = parent.size + 1, slice from [offset, offset+size+1)
+            # - Child 1 (values): slice based on the offset values
+            offsets_child = self.base_children[0]
+            values_child = self.base_children[1]
+
+            # Slice offsets: get size+1 offsets starting from self.offset
+            sliced_offsets = offsets_child.slice(
+                self.offset, self.offset + self.size + 1
+            )
+
+            # Get the range of values to slice: from first offset to last offset
+            # The offsets tell us where in the values array each list starts/ends
+            first_offset = offsets_child.element_indexing(self.offset)
+            last_offset = offsets_child.element_indexing(
+                self.offset + self.size
+            )
+
+            # Slice the values child to only include the values we need
+            sliced_values = values_child.slice(first_offset, last_offset)
+
+            # Adjust offsets to be relative to the sliced values (subtract first_offset)
+            adjusted_offsets = sliced_offsets - first_offset
+
+            self._children = (adjusted_offsets, sliced_values)  # type: ignore[assignment]
+
     def _prep_pandas_compat_repr(self) -> StringColumn | Self:
         """
         Preprocess Column to be compatible with pandas repr, namely handling nulls.

python/cudf/cudf/core/column/string.py

@@ -211,6 +211,33 @@ def _recompute_data(self) -> None:
             assert self.base_data is not None
             self._data = self.base_data[self.start_offset : self.end_offset]
 
+    def _recompute_children(self) -> None:
+        if not self.base_children:
+            self._children = ()
+        elif (
+            self.offset == 0
+            and len(self.base_children) > 0
+            and self.size == self.base_children[0].size - 1
+        ):
+            # Optimization: for non-sliced columns, children == base_children
+            self._children = self.base_children  # type: ignore[assignment]
+        else:
+            # String columns have one child: the offsets column
+            # For a string column with size N, the offsets child has size N+1
+
+            # Step 1: Slice the offsets column to get the correct size
+            offsets_child = self.base_children[0]
+            sliced_offsets = offsets_child.slice(
+                self.offset, self.offset + self.size + 1
+            )
+
+            # Step 2: Adjust the offsets to be relative to the sliced data
+            # by subtracting the first offset value (start_offset)
+            chars_offset = self.start_offset
+            adjusted_offsets = sliced_offsets - chars_offset
+
+            self._children = (adjusted_offsets,)  # type: ignore[assignment]
+
     def all(self, skipna: bool = True) -> bool:
         if skipna and self.null_count == self.size:
             return True

python/cudf/cudf/core/column/struct.py

@@ -79,6 +79,20 @@ def _get_children_from_pylibcudf_column(
             )
         )
 
+    def _recompute_children(self) -> None:
+        """Recompute the offset-aware children columns with proper type metadata."""
+        if not self.base_children:
+            self._children = ()
+        elif self.offset == 0 and self.size == self.base_size:
+            # Optimization: for non-sliced columns, children == base_children
+            self._children = self.base_children  # type: ignore[assignment]
+        else:
+            # Slice each child using the parent's offset and size
+            self._children = tuple(  # type: ignore[assignment]
+                base_child.slice(self.offset, self.offset + self.size)
+                for base_child in self.base_children
+            )
+
     def _prep_pandas_compat_repr(self) -> StringColumn | Self:
         """
         Preprocess Column to be compatible with pandas repr, namely handling nulls.