GH-46629: [Python] Add options to DatasetFactory.inspect

python/pyarrow/_dataset.pyx

@@ -3151,9 +3151,19 @@ cdef class DatasetFactory(_Weakrefable):
  schemas.append(pyarrow_wrap_schema(s))
  return schemas
 
- def inspect(self):
+ def inspect(self, *, promote_options="default", fragments=None):
  """
- Inspect all data fragments and return a common Schema.
+ Inspect data fragments and return a common Schema.
+
+ Parameters
+ ----------
+ promote_options : str, default "default"
+ Control how to unify types. Accepts strings "default" and "permissive".
+ Default: types must match exactly, except nulls can be merged with other types.
+ Permissive: types are promoted when possible.
+ fragments : int, default None
+ How many fragments should be inspected to infer the unified schema.
+ Use ``None`` to inspect all fragments.
 
  Returns
  -------
@@ -3162,6 +3172,17 @@ cdef class DatasetFactory(_Weakrefable):
  cdef:
  CInspectOptions options
  CResult[shared_ptr[CSchema]] result
+
+ options.field_merge_options = _parse_field_merge_options(promote_options)
+
+ if fragments is None:
+ options.fragments = -1 # InspectOptions::kInspectAllFragments
+ elif isinstance(fragments, int) and fragments >= 0:
+ options.fragments = fragments
+ else:
+ raise ValueError(
+ f"Fragment count must be a non-negative int or None; got {fragments!r}")
+
  with nogil:
  result = self.factory.Inspect(options)
  return pyarrow_wrap_schema(GetResultValue(result))

python/pyarrow/includes/libarrow_dataset.pxd

@@ -156,6 +156,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
 
  cdef cppclass CInspectOptions "arrow::dataset::InspectOptions":
  int fragments
+ CField.CMergeOptions field_merge_options
 
  cdef cppclass CFinishOptions "arrow::dataset::FinishOptions":
  shared_ptr[CSchema] schema

python/pyarrow/lib.pxd

@@ -678,6 +678,8 @@ cdef shared_ptr[const CKeyValueMetadata] pyarrow_unwrap_metadata(
 cdef object pyarrow_wrap_metadata(
  const shared_ptr[const CKeyValueMetadata]& meta)
 
+cdef CField.CMergeOptions _parse_field_merge_options(str promote_options) except *
+
 #
 # Public Cython API for 3rd party code
 #

python/pyarrow/table.pxi

@@ -6307,12 +6307,9 @@ def concat_tables(tables, MemoryPool memory_pool=None, str promote_options="none
  for table in tables:
  c_tables.push_back(table.sp_table)
 
- if promote_options == "permissive":
- options.field_merge_options = CField.CMergeOptions.Permissive()
- elif promote_options in {"default", "none"}:
- options.field_merge_options = CField.CMergeOptions.Defaults()
- else:
- raise ValueError(f"Invalid promote options: {promote_options}")
+ options.field_merge_options = _parse_field_merge_options(
+ "default" if promote_options == "none" else promote_options
+ )
 
  with nogil:
  options.unify_schemas = promote_options != "none"

python/pyarrow/tests/test_dataset.py

@@ -147,6 +147,42 @@ def mockfs():
  return mockfs
 
 
+@pytest.fixture
+def promotable_mockfs():
+ """
+ Creates a _MockFileSystem with two parquet files that have promotable schemas.
+ - file1: value: int8, dictionary: dictionary<int8, string>
+ - file2: value: uint16, dictionary: dictionary<int16, string>
+ """
+ mockfs = fs._MockFileSystem()
+
+ table1 = pa.table({
+ 'value': pa.array([1, 2], type=pa.int8()),
+ 'dictionary': pa.DictionaryArray.from_arrays(
+ pa.array([0, 1], type=pa.int8()),
+ ['a', 'b'],
+ ),
+ })
+ table2 = pa.table({
+ 'value': pa.array([3, 4], type=pa.uint16()),
+ 'dictionary': pa.DictionaryArray.from_arrays(
+ pa.array([1, 0], type=pa.int16()),
+ ['d', 'c'],
+ ),
+ })
+
+ mockfs.create_dir('subdir/zzz')
+ path1 = 'subdir/zzz/file1.parquet'
+ with mockfs.open_output_stream(path1) as out:
+ pq.write_table(table1, out)
+
+ path2 = 'subdir/zzz/file2.parquet'
+ with mockfs.open_output_stream(path2) as out:
+ pq.write_table(table2, out)
+
+ return mockfs, path1, path2
+
+
 @pytest.fixture
 def open_logging_fs(monkeypatch):
  from pyarrow.fs import LocalFileSystem, PyFileSystem
@@ -451,6 +487,70 @@ def test_dataset(dataset, dataset_reader):
  assert_dataset_fragment_convenience_methods(dataset)
 
 
+def test_dataset_factory_inspect_schema_promotion(promotable_mockfs):
+ mockfs, path1, path2 = promotable_mockfs
+ factory = ds.FileSystemDatasetFactory(
+ mockfs, [path1, path2], ds.ParquetFileFormat()
+ )
+
+ with pytest.raises(
+ pa.ArrowTypeError,
+ match='Unable to merge: Field value has incompatible types: int8 vs uint16',
+ ):
+ factory.inspect()
+
+ schema = factory.inspect(promote_options='permissive')
+ expected_schema = pa.schema([
+ # Promoted from int8 and uint16
+ pa.field('value', pa.int32()),
+ # Dictionary index type promoted from int8 and int16
+ pa.field('dictionary', pa.dictionary(pa.int16(), pa.string())),
+ ])
+ assert schema.equals(expected_schema)
+ dataset = factory.finish(schema)
+ table = dataset.to_table()
+ table.validate(full=True)
+ expected_table = pa.table({
+ 'value': pa.chunked_array([[1, 2], [3, 4]], type=pa.int32()),
+ 'dictionary': pa.chunked_array([
+ pa.DictionaryArray.from_arrays(
+ pa.array([0, 1], type=pa.int16()),
+ ['a', 'b']
+ ),
+ pa.DictionaryArray.from_arrays(
+ pa.array([1, 0], type=pa.int16()),
+ ['d', 'c']
+ )
+ ]),
+ })
+ assert table.equals(expected_table)
+
+ # Inspecting only one fragment should not promote the 'value' field
+ inspected_schema_one_frag = factory.inspect(
+ promote_options='permissive', fragments=1)
+ assert inspected_schema_one_frag.field('value').type != pa.int32()
+
+
+def test_dataset_factory_inspect_bad_params(promotable_mockfs):
+ mockfs, path1, path2 = promotable_mockfs
+ factory = ds.FileSystemDatasetFactory(
+ mockfs, [path1, path2], ds.ParquetFileFormat()
+ )
+
+ with pytest.raises(ValueError, match='Invalid promote_options: bad_option'):
+ factory.inspect(promote_options='bad_option')
+
+ with pytest.raises(
+ ValueError, match="Fragment count must be a non-negative int or None; got 'one'"
+ ):
+ factory.inspect(fragments="one")
+
+ with pytest.raises(
+ ValueError, match='Fragment count must be a non-negative int or None; got -1'
+ ):
+ factory.inspect(fragments=-1)
+
+
 @pytest.mark.parquet
 def test_scanner_options(dataset):
  scanner = dataset.to_batches(fragment_readahead=16, batch_readahead=8)

python/pyarrow/tests/test_table.py

@@ -1963,7 +1963,7 @@ def test_concat_tables_permissive():
 def test_concat_tables_invalid_option():
  t = pa.Table.from_arrays([list(range(10))], names=('a',))
 
- with pytest.raises(ValueError, match="Invalid promote options: invalid"):
+ with pytest.raises(ValueError, match="Invalid promote_options: invalid"):
  pa.concat_tables([t, t], promote_options="invalid")
 
 

python/pyarrow/types.pxi

@@ -3676,6 +3676,19 @@ cdef class Schema(_Weakrefable):
  return pyarrow_wrap_schema(result)
 
 
+cdef CField.CMergeOptions _parse_field_merge_options(str promote_options) except *:
+ """
+ Returns MergeOptions::Permissive() or MergeOptions::Defaults() based on the value
+ of `promote_options`.
+ """
+ if promote_options == "permissive":
+ return CField.CMergeOptions.Permissive()
+ elif promote_options == "default":
+ return CField.CMergeOptions.Defaults()
+ else:
+ raise ValueError(f"Invalid promote_options: {promote_options}")
+
+
 def unify_schemas(schemas, *, promote_options="default"):
  """
  Unify schemas by merging fields by name.
@@ -3720,12 +3733,7 @@ def unify_schemas(schemas, *, promote_options="default"):
  raise TypeError(f"Expected Schema, got {type(schema)}")
  c_schemas.push_back(pyarrow_unwrap_schema(schema))
 
- if promote_options == "default":
- c_options = CField.CMergeOptions.Defaults()
- elif promote_options == "permissive":
- c_options = CField.CMergeOptions.Permissive()
- else:
- raise ValueError(f"Invalid merge mode: {promote_options}")
+ c_options = _parse_field_merge_options(promote_options)
 
  return pyarrow_wrap_schema(
  GetResultValue(UnifySchemas(c_schemas, c_options)))