Merge pull request #1083 from nikandfor/main

SpencerTorres · web-flow · commit ec8982d95cfb · 2025-08-22T19:14:07.000-07:00
feat(proto): ColArr.RowLen
diff --git a/proto/col_arr.go b/proto/col_arr.go
@@ -80,6 +80,18 @@ func (c *ColArr[T]) Infer(t ColumnType) error {
 	return nil
 }
 
+// RowLen returns i-th row array length.
+func (c ColArr[T]) RowLen(i int) int {
+	var start int
+	if i > 0 {
+		start = int(c.Offsets[i-1])
+	}
+
+	end := int(c.Offsets[i])
+
+	return end - start
+}
+
 // RowAppend appends i-th row to target and returns it.
 func (c ColArr[T]) RowAppend(i int, target []T) []T {
 	var start int
diff --git a/proto/col_arr_test.go b/proto/col_arr_test.go
@@ -74,7 +74,9 @@ func TestColArrOfStr(t *testing.T) {
 		require.Equal(t, col.Rows(), dec.Rows())
 		require.Equal(t, ColumnType("Array(String)"), dec.Type())
 		require.Equal(t, []string{"foo", "bar", "foo", "foo", "baz"}, dec.Row(0))
+		require.Equal(t, 5, dec.RowLen(0))
 		require.Equal(t, []string{"foo", "baz"}, dec.Row(1))
+		require.Equal(t, 2, dec.RowLen(1))
 	})
 	t.Run("EOF", func(t *testing.T) {
 		r := NewReader(bytes.NewReader(nil))
@@ -109,7 +111,9 @@ func TestArrOfLowCordStr(t *testing.T) {
 		require.Equal(t, col.Rows(), dec.Rows())
 		require.Equal(t, ColumnType("Array(LowCardinality(String))"), dec.Type())
 		require.Equal(t, []string{"foo", "bar", "foo", "foo", "baz"}, dec.Row(0))
+		require.Equal(t, 5, dec.RowLen(0))
 		require.Equal(t, []string{"foo", "baz"}, dec.Row(1))
+		require.Equal(t, 2, dec.RowLen(1))
 	})
 	t.Run("EOF", func(t *testing.T) {
 		r := NewReader(bytes.NewReader(nil))
